mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-10 17:29:39 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3386d25022 | ||
|
|
6dca2ace8a | ||
|
|
63548739e8 | ||
|
|
3e046537f4 | ||
|
|
4a34677567 | ||
|
|
19e91cb669 | ||
|
|
5a6125021a | ||
|
|
da26887f24 | ||
|
|
1b2f7139f9 | ||
|
|
aafbc86180 | ||
|
|
3d4ca1b7c0 | ||
|
|
1180aa38bd | ||
|
|
65f65e5c7b | ||
|
|
f81a12d5e8 | ||
|
|
cf6f5fbefe | ||
|
|
908d7710db | ||
|
|
7d5b41e31f | ||
|
|
e46a0bc9f3 | ||
|
|
911483c92c | ||
|
|
4ffb9d6e3f |
No files matched your search
@@ -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.9.1" \
|
||||
"praisonai>=3.0.2" \
|
||||
"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.9.1" \
|
||||
"praisonai>=3.0.2" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.9.1" \
|
||||
"praisonai>=3.0.2" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# Profiling CLI
|
||||
|
||||
PraisonAI provides unified profiling capabilities for performance analysis and debugging, with **full visibility from ENTER to first response**.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Basic profiling
|
||||
praisonai "What is 2+2?" --profile
|
||||
|
||||
# Deep profiling with call graph, decision trace, and module breakdown
|
||||
praisonai "What is 2+2?" --profile --profile-deep
|
||||
|
||||
# JSON output format
|
||||
praisonai "What is 2+2?" --profile --profile-format json
|
||||
```
|
||||
|
||||
## Time to First Response (TTFR)
|
||||
|
||||
**TTFR** is the key user-perceived latency metric. It measures the time from when you press ENTER until the first visible output appears.
|
||||
|
||||
### What contributes to TTFR?
|
||||
|
||||
1. **CLI Entry** - Time to reach the CLI entry point
|
||||
2. **CLI Parse** - Argparse/Typer routing time
|
||||
3. **Routing** - Command routing to handler
|
||||
4. **Imports** - Loading required modules
|
||||
5. **Agent Init** - Initializing the agent
|
||||
6. **Network Request** - Sending request to LLM API
|
||||
7. **First Token** - Receiving first token (streaming)
|
||||
8. **First Output** - Rendering first visible output
|
||||
|
||||
### How to reduce TTFR
|
||||
|
||||
1. **Warm runs** - Imports are cached after first run
|
||||
2. **Streaming mode** - See output as tokens arrive
|
||||
3. **Lighter models** - Faster response times
|
||||
4. **Local models** - Eliminate network latency
|
||||
|
||||
## CLI Flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--profile` | Enable profiling with timing breakdown | Disabled |
|
||||
| `--profile-deep` | Enable deep profiling with call graph (higher overhead) | Disabled |
|
||||
| `--profile-format` | Output format: `text` or `json` | `text` |
|
||||
|
||||
## Profile Command
|
||||
|
||||
For more detailed profiling options, use the `profile` subcommand:
|
||||
|
||||
```bash
|
||||
# Profile a query
|
||||
praisonai profile query "What is 2+2?"
|
||||
|
||||
# With deep call tracing
|
||||
praisonai profile query "Hello" --deep
|
||||
|
||||
# Limit functions shown
|
||||
praisonai profile query "Test" --limit 20
|
||||
|
||||
# Save artifacts to directory
|
||||
praisonai profile query "Test" --save ./profile_results
|
||||
|
||||
# JSON output
|
||||
praisonai profile query "Test" --format json
|
||||
```
|
||||
|
||||
### Profile Query Options
|
||||
|
||||
| Option | Description | Default |
|
||||
|--------|-------------|---------|
|
||||
| `--model, -m` | Model to use | Default model |
|
||||
| `--stream/--no-stream` | Use streaming mode | No stream |
|
||||
| `--deep` | Enable deep call tracing | Disabled |
|
||||
| `--limit, -n` | Top N functions to show | 30 |
|
||||
| `--sort, -s` | Sort by: `cumulative` or `tottime` | `cumulative` |
|
||||
| `--show-callers` | Show caller functions | Disabled |
|
||||
| `--show-callees` | Show callee functions | Disabled |
|
||||
| `--save` | Save artifacts to path | None |
|
||||
| `--format, -f` | Output format: `text` or `json` | `text` |
|
||||
|
||||
## Other Profile Commands
|
||||
|
||||
```bash
|
||||
# Profile module import times
|
||||
praisonai profile imports
|
||||
|
||||
# Profile CLI startup time
|
||||
praisonai profile startup
|
||||
|
||||
# Run comprehensive profiling suite
|
||||
praisonai profile suite
|
||||
|
||||
# Create performance snapshots for comparison
|
||||
praisonai profile snapshot --baseline
|
||||
praisonai profile snapshot current --compare
|
||||
|
||||
# Configure performance optimizations
|
||||
praisonai profile optimize --show
|
||||
```
|
||||
|
||||
## Output Format
|
||||
|
||||
### Text Report
|
||||
|
||||
```
|
||||
======================================================================
|
||||
PraisonAI Profile Report
|
||||
======================================================================
|
||||
|
||||
Run ID: abc12345
|
||||
Timestamp: 2024-01-02T12:00:00.000000Z
|
||||
Method: cli_direct
|
||||
Version: 3.0.1
|
||||
|
||||
## Execution Timeline
|
||||
---------------------------------------------
|
||||
Imports : 500.00 ms
|
||||
Agent Init : 0.50 ms
|
||||
Execution : 2000.00 ms
|
||||
───────────────────────────────────────────
|
||||
⏱ Time to First Response : 2500.50 ms
|
||||
TOTAL : 2500.50 ms
|
||||
|
||||
## Timing Breakdown
|
||||
----------------------------------------
|
||||
Imports: 500.00 ms
|
||||
Agent Init: 0.50 ms
|
||||
Execution: 2000.00 ms
|
||||
─────────────────────────────────────
|
||||
TOTAL: 2500.50 ms
|
||||
|
||||
## Decision Trace (deep profile only)
|
||||
----------------------------------------
|
||||
Agent Config: default
|
||||
Model: gpt-4
|
||||
Streaming: disabled
|
||||
Profile Layer: 2
|
||||
|
||||
## Top Functions by Cumulative Time
|
||||
----------------------------------------------------------------------
|
||||
Function Calls Cumul (ms)
|
||||
----------------------------------------------------------------------
|
||||
request 1 1800.00
|
||||
send 1 1750.00
|
||||
...
|
||||
```
|
||||
|
||||
### JSON Report
|
||||
|
||||
```json
|
||||
{
|
||||
"schema_version": "1.0",
|
||||
"run_id": "abc12345",
|
||||
"timestamp": "2024-01-02T12:00:00.000000Z",
|
||||
"invocation": {
|
||||
"method": "cli_direct",
|
||||
"flags": {"layer": 1, "stream": false, "model": null},
|
||||
"praisonai_version": "3.0.1",
|
||||
"python_version": "3.12.0"
|
||||
},
|
||||
"timing": {
|
||||
"total_ms": 2500.5,
|
||||
"imports_ms": 500.0,
|
||||
"agent_init_ms": 0.5,
|
||||
"execution_ms": 2000.0
|
||||
},
|
||||
"functions": [...],
|
||||
"response_preview": "four"
|
||||
}
|
||||
```
|
||||
|
||||
## Profiling Layers
|
||||
|
||||
| Layer | Overhead | Features |
|
||||
|-------|----------|----------|
|
||||
| 0 | < 1ms | Wall-clock phases only |
|
||||
| 1 | ~5% | + cProfile function stats |
|
||||
| 2 | ~15% | + Call graph (callers/callees) |
|
||||
|
||||
- `--profile` uses Layer 1
|
||||
- `--profile --profile-deep` uses Layer 2
|
||||
|
||||
## Performance Tips
|
||||
|
||||
1. **Use Layer 0 for production monitoring** - minimal overhead
|
||||
2. **Use Layer 1 for development profiling** - good balance
|
||||
3. **Use Layer 2 only for deep debugging** - higher overhead
|
||||
|
||||
## Observational Equivalence
|
||||
|
||||
The following commands produce equivalent profile reports:
|
||||
|
||||
```bash
|
||||
# CLI direct
|
||||
praisonai "prompt" --profile
|
||||
|
||||
# Profile command
|
||||
praisonai profile query "prompt"
|
||||
```
|
||||
|
||||
Both use the same unified profiling architecture and produce identical schema output.
|
||||
@@ -0,0 +1,257 @@
|
||||
# Profiling Module
|
||||
|
||||
PraisonAI provides a unified profiling module for programmatic performance analysis.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import (
|
||||
ExecutionRequest,
|
||||
Profiler,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
# Create execution request
|
||||
request = ExecutionRequest(prompt="What is 2+2?")
|
||||
|
||||
# Configure profiler
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
# Run with profiling
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request)
|
||||
|
||||
# Access results
|
||||
print(f"Output: {result.output}")
|
||||
print(f"Total time: {report.timing.total_ms}ms")
|
||||
```
|
||||
|
||||
## ExecutionRequest
|
||||
|
||||
Immutable request descriptor for execution.
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import ExecutionRequest
|
||||
|
||||
# Basic request
|
||||
request = ExecutionRequest(prompt="Hello")
|
||||
|
||||
# With model
|
||||
request = ExecutionRequest(
|
||||
prompt="Hello",
|
||||
model="gpt-4",
|
||||
agent_name="MyAgent",
|
||||
)
|
||||
|
||||
# With streaming
|
||||
request = ExecutionRequest(
|
||||
prompt="Hello",
|
||||
stream=True,
|
||||
)
|
||||
```
|
||||
|
||||
### Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `prompt` | str | Required | The prompt to execute |
|
||||
| `agent_name` | str | "Assistant" | Name of the agent |
|
||||
| `agent_instructions` | str | None | Custom instructions |
|
||||
| `model` | str | None | Model to use |
|
||||
| `stream` | bool | False | Enable streaming |
|
||||
| `tools` | tuple | () | Tools to provide |
|
||||
|
||||
## ProfilerConfig
|
||||
|
||||
Configuration for profiling behavior.
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import ProfilerConfig
|
||||
|
||||
# Layer 0: Wall-clock only (< 1ms overhead)
|
||||
config = ProfilerConfig(layer=0)
|
||||
|
||||
# Layer 1: + cProfile stats (~5% overhead)
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
# Layer 2: + Call graph (~15% overhead)
|
||||
config = ProfilerConfig(
|
||||
layer=2,
|
||||
show_callers=True,
|
||||
show_callees=True,
|
||||
)
|
||||
|
||||
# From CLI flags
|
||||
config = ProfilerConfig.from_flags(
|
||||
profile=True,
|
||||
deep=True,
|
||||
network=False,
|
||||
)
|
||||
```
|
||||
|
||||
### Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `layer` | int | 1 | Profiling layer (0, 1, or 2) |
|
||||
| `limit` | int | 30 | Max functions to show |
|
||||
| `sort_by` | str | "cumulative" | Sort by: cumulative, time, calls |
|
||||
| `show_callers` | bool | False | Show caller functions |
|
||||
| `show_callees` | bool | False | Show callee functions |
|
||||
| `track_network` | bool | False | Track network requests |
|
||||
| `output_format` | str | "text" | Output format: text, json |
|
||||
| `save_path` | str | None | Path to save artifacts |
|
||||
|
||||
## Profiler
|
||||
|
||||
The main profiler class that wraps execution.
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import Profiler, ProfilerConfig, ExecutionRequest
|
||||
|
||||
config = ProfilerConfig(layer=1)
|
||||
profiler = Profiler(config)
|
||||
|
||||
request = ExecutionRequest(prompt="Hello")
|
||||
result, report = profiler.profile_sync(request, invocation_method="api")
|
||||
```
|
||||
|
||||
### Methods
|
||||
|
||||
#### `profile_sync(request, invocation_method="profile_command")`
|
||||
|
||||
Profile synchronous execution.
|
||||
|
||||
**Parameters:**
|
||||
- `request`: ExecutionRequest - The execution request
|
||||
- `invocation_method`: str - How this was invoked (for report)
|
||||
|
||||
**Returns:** Tuple of (ExecutionResult, ProfileReport)
|
||||
|
||||
## ProfileReport
|
||||
|
||||
The canonical profile report with timing and statistics.
|
||||
|
||||
```python
|
||||
# Access timing breakdown
|
||||
print(f"Total: {report.timing.total_ms}ms")
|
||||
print(f"Imports: {report.timing.imports_ms}ms")
|
||||
print(f"Agent Init: {report.timing.agent_init_ms}ms")
|
||||
print(f"Execution: {report.timing.execution_ms}ms")
|
||||
|
||||
# Access function stats (layer 1+)
|
||||
if report.functions:
|
||||
for func in report.functions[:5]:
|
||||
print(f"{func.name}: {func.cumulative_time_ms}ms")
|
||||
|
||||
# Access call graph (layer 2)
|
||||
if report.call_graph:
|
||||
print(f"Edges: {len(report.call_graph.edges)}")
|
||||
|
||||
# Export formats
|
||||
json_str = report.to_json()
|
||||
text_str = report.to_text()
|
||||
dict_data = report.to_dict()
|
||||
```
|
||||
|
||||
### Schema Version
|
||||
|
||||
The report uses schema version `1.0`. Schema is frozen after release - new fields are always optional with defaults.
|
||||
|
||||
### Mandatory Fields
|
||||
|
||||
- `schema_version`: str
|
||||
- `run_id`: str
|
||||
- `timestamp`: str (ISO 8601)
|
||||
- `invocation`: InvocationInfo
|
||||
- `timing`: TimingBreakdown
|
||||
- `response_preview`: str
|
||||
|
||||
### Optional Fields
|
||||
|
||||
- `functions`: List[FunctionStat] (layer 1+)
|
||||
- `call_graph`: CallGraph (layer 2)
|
||||
- `network`: List[RequestTiming] (if track_network)
|
||||
|
||||
## ExecutionResult
|
||||
|
||||
Result of execution with output and metadata.
|
||||
|
||||
```python
|
||||
# Access output
|
||||
print(result.output)
|
||||
print(result.success)
|
||||
print(result.run_id)
|
||||
|
||||
# Check for errors
|
||||
if result.error:
|
||||
print(f"Error: {result.error}")
|
||||
|
||||
# Duration
|
||||
print(f"Duration: {result.duration_ms}ms")
|
||||
```
|
||||
|
||||
## Data Bounds
|
||||
|
||||
Profile data is strictly bounded to prevent memory issues:
|
||||
|
||||
| Data Type | Maximum |
|
||||
|-----------|---------|
|
||||
| Function stats | 1000 |
|
||||
| Call graph edges | 5000 |
|
||||
| Network requests | 100 |
|
||||
|
||||
## Example: Full Profiling Workflow
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import (
|
||||
ExecutionRequest,
|
||||
Profiler,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
# Configure deep profiling
|
||||
config = ProfilerConfig(
|
||||
layer=2,
|
||||
limit=50,
|
||||
show_callers=True,
|
||||
show_callees=True,
|
||||
)
|
||||
|
||||
# Create request
|
||||
request = ExecutionRequest(
|
||||
prompt="Explain quantum computing in simple terms",
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
# Profile execution
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request)
|
||||
|
||||
# Print text report
|
||||
print(report.to_text())
|
||||
|
||||
# Save JSON report
|
||||
with open("profile.json", "w") as f:
|
||||
f.write(report.to_json())
|
||||
|
||||
# Analyze timing
|
||||
print(f"\nTiming Analysis:")
|
||||
print(f" Import overhead: {report.timing.imports_ms / report.timing.total_ms * 100:.1f}%")
|
||||
print(f" Execution time: {report.timing.execution_ms / report.timing.total_ms * 100:.1f}%")
|
||||
|
||||
# Top 5 hotspots
|
||||
if report.functions:
|
||||
print(f"\nTop 5 Hotspots:")
|
||||
for func in report.functions[:5]:
|
||||
print(f" {func.name}: {func.cumulative_time_ms:.2f}ms ({func.calls} calls)")
|
||||
```
|
||||
|
||||
## Invariants
|
||||
|
||||
The profiling module guarantees:
|
||||
|
||||
1. **Output Equivalence**: Profiled execution produces identical output to non-profiled execution
|
||||
2. **Schema Stability**: Schema version 1.0 is frozen; new fields are always optional
|
||||
3. **Bounded Data**: All profile data respects maximum bounds
|
||||
4. **Zero Overhead When Disabled**: No profiling code runs unless explicitly enabled
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
Interactive Default Tools Example
|
||||
|
||||
This example demonstrates how PraisonAI's interactive modes (TUI and prompt)
|
||||
include ACP and LSP tools by default.
|
||||
|
||||
Tool Groups:
|
||||
- ACP: acp_create_file, acp_edit_file, acp_delete_file, acp_execute_command
|
||||
- LSP: lsp_list_symbols, lsp_find_definition, lsp_find_references, lsp_get_diagnostics
|
||||
- Basic: read_file, write_file, list_files, execute_command, internet_search
|
||||
|
||||
Usage:
|
||||
python interactive_tools_example.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add paths for development
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai-agents'))
|
||||
|
||||
|
||||
def example_get_default_tools():
|
||||
"""Example: Get all default interactive tools."""
|
||||
from praisonai.cli.features.interactive_tools import (
|
||||
get_interactive_tools,
|
||||
ToolConfig,
|
||||
TOOL_GROUPS,
|
||||
)
|
||||
|
||||
print("=== Default Interactive Tools ===\n")
|
||||
|
||||
# Show tool groups
|
||||
print("Tool Groups:")
|
||||
for group, tools in TOOL_GROUPS.items():
|
||||
if group != "interactive":
|
||||
print(f" {group}: {tools}")
|
||||
print()
|
||||
|
||||
# Get all tools
|
||||
config = ToolConfig(workspace=os.getcwd())
|
||||
tools = get_interactive_tools(config=config)
|
||||
|
||||
print(f"Total tools loaded: {len(tools)}")
|
||||
print(f"Tool names: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def example_disable_acp():
|
||||
"""Example: Disable ACP tools."""
|
||||
from praisonai.cli.features.interactive_tools import (
|
||||
get_interactive_tools,
|
||||
ToolConfig,
|
||||
)
|
||||
|
||||
print("=== Disable ACP Tools (--no-acp) ===\n")
|
||||
|
||||
config = ToolConfig(workspace=os.getcwd(), enable_acp=False)
|
||||
tools = get_interactive_tools(config=config, disable=["acp"])
|
||||
|
||||
acp_tools = [t for t in tools if t.__name__.startswith("acp_")]
|
||||
print(f"ACP tools after --no-acp: {len(acp_tools)}")
|
||||
print(f"Remaining tools: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def example_disable_lsp():
|
||||
"""Example: Disable LSP tools."""
|
||||
from praisonai.cli.features.interactive_tools import (
|
||||
get_interactive_tools,
|
||||
ToolConfig,
|
||||
)
|
||||
|
||||
print("=== Disable LSP Tools (--no-lsp) ===\n")
|
||||
|
||||
config = ToolConfig(workspace=os.getcwd(), enable_lsp=False)
|
||||
tools = get_interactive_tools(config=config, disable=["lsp"])
|
||||
|
||||
lsp_tools = [t for t in tools if t.__name__.startswith("lsp_")]
|
||||
print(f"LSP tools after --no-lsp: {len(lsp_tools)}")
|
||||
print(f"Remaining tools: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def example_agent_with_default_tools():
|
||||
"""Example: Create an Agent with default interactive tools."""
|
||||
from praisonai.cli.features.interactive_tools import get_interactive_tools, ToolConfig
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("=== Agent with Default Interactive Tools ===\n")
|
||||
|
||||
# Get default tools
|
||||
config = ToolConfig(workspace=os.getcwd())
|
||||
tools = get_interactive_tools(config=config)
|
||||
|
||||
# Create agent with tools
|
||||
agent = Agent(
|
||||
name="InteractiveAgent",
|
||||
instructions="You are a helpful coding assistant with file and code analysis capabilities.",
|
||||
tools=tools,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
print(f"Agent created with {len(tools)} tools")
|
||||
print(f"Tools: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
# Simple test - list files
|
||||
print("Testing agent with list_files tool...")
|
||||
result = agent.start("List the Python files in the current directory using list_files tool.")
|
||||
print(f"Result: {result[:200]}..." if len(str(result)) > 200 else f"Result: {result}")
|
||||
print()
|
||||
|
||||
return agent
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("PraisonAI Interactive Default Tools Example")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
# Example 1: Get default tools
|
||||
example_get_default_tools()
|
||||
|
||||
# Example 2: Disable ACP
|
||||
example_disable_acp()
|
||||
|
||||
# Example 3: Disable LSP
|
||||
example_disable_lsp()
|
||||
|
||||
# Example 4: Agent with default tools (requires API key)
|
||||
if os.environ.get("OPENAI_API_KEY"):
|
||||
example_agent_with_default_tools()
|
||||
else:
|
||||
print("=== Agent Example Skipped (no OPENAI_API_KEY) ===")
|
||||
|
||||
print("=" * 60)
|
||||
print("Examples Complete")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,83 @@
|
||||
# PraisonAI Profiling Examples
|
||||
|
||||
This directory contains examples for profiling PraisonAI agent performance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install praisonai
|
||||
export OPENAI_API_KEY=your_key_here
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
### 1. Basic Profiling (`basic_profiling.py`)
|
||||
|
||||
Demonstrates programmatic profiling of a single query:
|
||||
|
||||
```bash
|
||||
python basic_profiling.py
|
||||
```
|
||||
|
||||
### 2. Suite Profiling (`suite_profiling.py`)
|
||||
|
||||
Runs a comprehensive profiling suite with multiple scenarios:
|
||||
|
||||
```bash
|
||||
python suite_profiling.py
|
||||
```
|
||||
|
||||
### 3. Optimization Demo (`optimization_example.py`)
|
||||
|
||||
Demonstrates Tier 0/1/2 performance optimizations:
|
||||
|
||||
```bash
|
||||
python optimization_example.py
|
||||
```
|
||||
|
||||
## CLI Commands
|
||||
|
||||
You can also profile directly from the command line:
|
||||
|
||||
```bash
|
||||
# Profile a query
|
||||
praisonai profile query "What is 2+2?"
|
||||
|
||||
# Profile with file grouping
|
||||
praisonai profile query "Hello" --show-files --limit 20
|
||||
|
||||
# Profile imports
|
||||
praisonai profile imports
|
||||
|
||||
# Profile startup time
|
||||
praisonai profile startup
|
||||
|
||||
# Run comprehensive suite
|
||||
praisonai profile suite --quick
|
||||
|
||||
# Create performance baseline
|
||||
praisonai profile snapshot --baseline
|
||||
|
||||
# Compare against baseline
|
||||
praisonai profile snapshot current --compare
|
||||
|
||||
# Show optimization status
|
||||
praisonai profile optimize --show
|
||||
```
|
||||
|
||||
## Output
|
||||
|
||||
- **Text output**: Human-readable timing breakdown
|
||||
- **JSON output**: Machine-readable for CI/CD (`--format json`)
|
||||
- **Artifacts**: Binary cProfile data and reports (`--save ./results`)
|
||||
- **Snapshots**: Performance baselines for regression detection
|
||||
|
||||
## Performance Optimizations
|
||||
|
||||
Enable opt-in optimizations via environment variables:
|
||||
|
||||
```bash
|
||||
export PRAISONAI_LITE_MODE=1
|
||||
export PRAISONAI_SKIP_TYPE_VALIDATION=1
|
||||
export PRAISONAI_MINIMAL_IMPORTS=1
|
||||
```
|
||||
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
Basic Profiling Example for PraisonAI
|
||||
|
||||
This example demonstrates how to use the unified profiler programmatically
|
||||
to measure agent performance.
|
||||
|
||||
Run with:
|
||||
python basic_profiling.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
from praisonai.cli.execution import (
|
||||
ExecutionRequest,
|
||||
Profiler,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Run a profiled query and display results."""
|
||||
|
||||
# Configure profiler (Layer 1 = basic profiling with function stats)
|
||||
config = ProfilerConfig(
|
||||
layer=1, # 0=minimal, 1=basic, 2=deep
|
||||
limit=15, # Show top 15 functions
|
||||
show_callers=False, # Set True for caller info (layer 2)
|
||||
show_callees=False, # Set True for callee info (layer 2)
|
||||
)
|
||||
|
||||
# Create execution request
|
||||
request = ExecutionRequest(
|
||||
prompt="What is 2+2?",
|
||||
agent_name="ProfiledAgent",
|
||||
)
|
||||
|
||||
# Run profiled query
|
||||
print("🔬 Running profiled query...")
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request)
|
||||
|
||||
# Print formatted report
|
||||
print("\n" + report.to_text())
|
||||
|
||||
# Access timing data programmatically
|
||||
print("\n📊 Timing Summary:")
|
||||
print(f" Imports: {report.timing.imports_ms:>10.2f} ms")
|
||||
print(f" Agent Init: {report.timing.agent_init_ms:>10.2f} ms")
|
||||
print(f" Execution: {report.timing.execution_ms:>10.2f} ms")
|
||||
print(f" Total: {report.timing.total_ms:>10.2f} ms")
|
||||
|
||||
# Show top 5 functions
|
||||
if report.functions:
|
||||
print("\n🔥 Top 5 Functions by Cumulative Time:")
|
||||
for i, func in enumerate(report.functions[:5], 1):
|
||||
print(f" {i}. {func.name}: {func.cumulative_time_ms:.2f}ms")
|
||||
|
||||
# Show response
|
||||
print(f"\n💬 Response: {result.output}")
|
||||
|
||||
# Export to JSON
|
||||
print("\n📄 JSON Export available via: report.to_json()")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
Performance Optimization Example for PraisonAI
|
||||
|
||||
This example demonstrates how to use the Tier 0/1/2 optimization features
|
||||
programmatically.
|
||||
|
||||
Run with:
|
||||
python optimization_example.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
|
||||
def demo_provider_cache():
|
||||
"""Demonstrate provider caching (Tier 0)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 0: Provider Cache Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
get_provider_cache,
|
||||
resolve_model_provider,
|
||||
)
|
||||
|
||||
cache = get_provider_cache()
|
||||
cache.invalidate() # Start fresh
|
||||
|
||||
# Resolve some models
|
||||
models = ['gpt-4o', 'claude-3-opus', 'gemini-pro', 'gpt-4o-mini']
|
||||
|
||||
print("\nResolving models:")
|
||||
for model in models:
|
||||
provider, name = resolve_model_provider(model)
|
||||
print(f" {model} -> {provider}/{name}")
|
||||
|
||||
# Check cache stats
|
||||
stats = cache.get_stats()
|
||||
print("\nCache Stats:")
|
||||
print(f" Hits: {stats['hits']}")
|
||||
print(f" Misses: {stats['misses']}")
|
||||
print(f" Size: {stats['size']}")
|
||||
|
||||
|
||||
def demo_lazy_importer():
|
||||
"""Demonstrate lazy imports (Tier 0)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 0: Lazy Importer Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import LazyImporter
|
||||
import time
|
||||
|
||||
LazyImporter.clear()
|
||||
|
||||
# Time lazy import
|
||||
start = time.perf_counter()
|
||||
json_module = LazyImporter.get('json')
|
||||
first_import = (time.perf_counter() - start) * 1000
|
||||
|
||||
# Second import (cached)
|
||||
start = time.perf_counter()
|
||||
json_module2 = LazyImporter.get('json')
|
||||
cached_import = (time.perf_counter() - start) * 1000
|
||||
|
||||
print(f"\nFirst import: {first_import:.3f}ms")
|
||||
print(f"Cached import: {cached_import:.3f}ms")
|
||||
print(f"Same object: {json_module is json_module2}")
|
||||
|
||||
|
||||
def demo_client_pool():
|
||||
"""Demonstrate client pooling (Tier 1)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 1: Client Pool Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import ClientPool
|
||||
|
||||
pool = ClientPool(max_size=3)
|
||||
|
||||
# Create some mock clients
|
||||
def create_client(name):
|
||||
return {'name': name, 'created': True}
|
||||
|
||||
# Get or create clients
|
||||
client1 = pool.get_or_create('openai', lambda: create_client('openai'))
|
||||
_ = pool.get_or_create('anthropic', lambda: create_client('anthropic'))
|
||||
|
||||
# Reuse existing client
|
||||
client1_again = pool.get_or_create('openai', lambda: create_client('openai_new'))
|
||||
|
||||
print(f"\nPool size: {pool.size()}")
|
||||
print(f"Client reused: {client1 is client1_again}")
|
||||
print(f"Client1 name: {client1['name']}")
|
||||
|
||||
|
||||
def demo_prewarm():
|
||||
"""Demonstrate pre-warming (Tier 1)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 1: Pre-warming Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import PrewarmManager
|
||||
import time
|
||||
|
||||
PrewarmManager.reset()
|
||||
|
||||
print(f"\nPre-warming enabled: {PrewarmManager.is_enabled()}")
|
||||
|
||||
# Enable and prewarm
|
||||
PrewarmManager.enable()
|
||||
print(f"After enable: {PrewarmManager.is_enabled()}")
|
||||
|
||||
# Prewarm OpenAI (runs in background)
|
||||
PrewarmManager.prewarm_provider('openai')
|
||||
print("Pre-warming OpenAI in background...")
|
||||
|
||||
# Wait a bit for background thread
|
||||
time.sleep(0.5)
|
||||
|
||||
print(f"OpenAI prewarmed: {PrewarmManager.is_prewarmed('openai')}")
|
||||
|
||||
PrewarmManager.reset()
|
||||
|
||||
|
||||
def demo_lite_mode():
|
||||
"""Demonstrate lite mode (Tier 2)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 2: Lite Mode Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
LiteModeConfig,
|
||||
is_lite_mode,
|
||||
)
|
||||
|
||||
# Default config
|
||||
config = LiteModeConfig()
|
||||
print(f"\nDefault lite mode: {config.enabled}")
|
||||
|
||||
# Config from env
|
||||
print("\nTo enable lite mode, set environment variables:")
|
||||
print(" export PRAISONAI_LITE_MODE=1")
|
||||
print(" export PRAISONAI_SKIP_TYPE_VALIDATION=1")
|
||||
|
||||
print(f"\nCurrent is_lite_mode(): {is_lite_mode()}")
|
||||
|
||||
|
||||
def demo_perf_snapshot():
|
||||
"""Demonstrate performance snapshots (Tier 2)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 2: Performance Snapshot Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
PerfSnapshot,
|
||||
PerfComparison,
|
||||
PerfSnapshotManager,
|
||||
)
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
# Create manager with temp directory
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
manager = PerfSnapshotManager(snapshot_dir=Path(tmpdir))
|
||||
|
||||
# Create baseline
|
||||
baseline = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='baseline',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
manager.save_baseline(baseline)
|
||||
print("\nBaseline saved")
|
||||
|
||||
# Create current snapshot (simulated improvement)
|
||||
current = PerfSnapshot(
|
||||
timestamp='2025-01-02T00:00:00Z',
|
||||
name='current',
|
||||
startup_cold_ms=95.0, # 5% faster
|
||||
startup_warm_ms=78.0,
|
||||
import_time_ms=480.0, # 4% faster
|
||||
query_time_ms=1900.0, # 5% faster
|
||||
first_token_ms=280.0,
|
||||
)
|
||||
|
||||
# Compare
|
||||
comparison = PerfComparison(baseline=baseline, current=current)
|
||||
|
||||
print(f"\nStartup improvement: {-comparison.startup_cold_diff_pct:.1f}%")
|
||||
print(f"Import improvement: {-comparison.import_time_diff_pct:.1f}%")
|
||||
print(f"Query improvement: {-comparison.query_time_diff_pct:.1f}%")
|
||||
print(f"Is regression: {comparison.is_regression()}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all demos."""
|
||||
print("=" * 60)
|
||||
print("PraisonAI Performance Optimization Demo")
|
||||
print("=" * 60)
|
||||
|
||||
demo_provider_cache()
|
||||
demo_lazy_importer()
|
||||
demo_client_pool()
|
||||
demo_prewarm()
|
||||
demo_lite_mode()
|
||||
demo_perf_snapshot()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All demos complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,227 @@
|
||||
{
|
||||
"timestamp": "2025-12-31T18:45:38.187539Z",
|
||||
"metadata": {
|
||||
"python_version": "3.12.11",
|
||||
"platform": "macOS-15.7.4-arm64-arm-64bit",
|
||||
"praisonai_version": "2.9.2",
|
||||
"timestamp": "2025-12-31T18:45:38.226558Z"
|
||||
},
|
||||
"startup": {
|
||||
"cold_ms": 80.9305000002496,
|
||||
"warm_ms": 55.266541996388696
|
||||
},
|
||||
"import_analysis": [
|
||||
{
|
||||
"module": "praisonaiagents",
|
||||
"self_us": 2419,
|
||||
"cumulative_us": 1645390,
|
||||
"self_ms": 2.419,
|
||||
"cumulative_ms": 1645.39
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.workflows",
|
||||
"self_us": 670,
|
||||
"cumulative_us": 1375919,
|
||||
"self_ms": 0.67,
|
||||
"cumulative_ms": 1375.919
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.workflows.yaml_parser",
|
||||
"self_us": 775,
|
||||
"cumulative_us": 1369435,
|
||||
"self_ms": 0.775,
|
||||
"cumulative_ms": 1369.435
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.agent.agent",
|
||||
"self_us": 19,
|
||||
"cumulative_us": 1349174,
|
||||
"self_ms": 0.019,
|
||||
"cumulative_ms": 1349.174
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.agent",
|
||||
"self_us": 537,
|
||||
"cumulative_us": 1349155,
|
||||
"self_ms": 0.537,
|
||||
"cumulative_ms": 1349.155
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.agent.agent",
|
||||
"self_us": 1453,
|
||||
"cumulative_us": 1285299,
|
||||
"self_ms": 1.453,
|
||||
"cumulative_ms": 1285.299
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.llm.openai_client",
|
||||
"self_us": 7887,
|
||||
"cumulative_us": 1123578,
|
||||
"self_ms": 7.887,
|
||||
"cumulative_ms": 1123.578
|
||||
},
|
||||
{
|
||||
"module": "openai",
|
||||
"self_us": 1083,
|
||||
"cumulative_us": 1115692,
|
||||
"self_ms": 1.083,
|
||||
"cumulative_ms": 1115.692
|
||||
},
|
||||
{
|
||||
"module": "openai.types",
|
||||
"self_us": 2606,
|
||||
"cumulative_us": 957990,
|
||||
"self_ms": 2.606,
|
||||
"cumulative_ms": 957.99
|
||||
},
|
||||
{
|
||||
"module": "openai.types.batch",
|
||||
"self_us": 1787,
|
||||
"cumulative_us": 485667,
|
||||
"self_ms": 1.787,
|
||||
"cumulative_ms": 485.667
|
||||
},
|
||||
{
|
||||
"module": "openai._models",
|
||||
"self_us": 13423,
|
||||
"cumulative_us": 450798,
|
||||
"self_ms": 13.423,
|
||||
"cumulative_ms": 450.798
|
||||
},
|
||||
{
|
||||
"module": "openai.types.eval_create_params",
|
||||
"self_us": 644,
|
||||
"cumulative_us": 393275,
|
||||
"self_ms": 0.644,
|
||||
"cumulative_ms": 393.275
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.grader_inputs_param",
|
||||
"self_us": 17,
|
||||
"cumulative_us": 359515,
|
||||
"self_ms": 0.017,
|
||||
"cumulative_ms": 359.515
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders",
|
||||
"self_us": 392,
|
||||
"cumulative_us": 359498,
|
||||
"self_ms": 0.392,
|
||||
"cumulative_ms": 359.498
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.multi_grader",
|
||||
"self_us": 3847,
|
||||
"cumulative_us": 351251,
|
||||
"self_ms": 3.847,
|
||||
"cumulative_ms": 351.251
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.label_model_grader",
|
||||
"self_us": 5450,
|
||||
"cumulative_us": 337291,
|
||||
"self_ms": 5.45,
|
||||
"cumulative_ms": 337.291
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.grader_inputs",
|
||||
"self_us": 545,
|
||||
"cumulative_us": 331841,
|
||||
"self_ms": 0.545,
|
||||
"cumulative_ms": 331.841
|
||||
},
|
||||
{
|
||||
"module": "openai.types.responses.response_input_text",
|
||||
"self_us": 21,
|
||||
"cumulative_us": 331297,
|
||||
"self_ms": 0.021,
|
||||
"cumulative_ms": 331.297
|
||||
},
|
||||
{
|
||||
"module": "openai.types.responses",
|
||||
"self_us": 1301,
|
||||
"cumulative_us": 331277,
|
||||
"self_ms": 1.301,
|
||||
"cumulative_ms": 331.277
|
||||
},
|
||||
{
|
||||
"module": "rich.logging",
|
||||
"self_us": 1066,
|
||||
"cumulative_us": 238256,
|
||||
"self_ms": 1.066,
|
||||
"cumulative_ms": 238.256
|
||||
}
|
||||
],
|
||||
"scenarios": [
|
||||
{
|
||||
"name": "simple_query",
|
||||
"config": {
|
||||
"prompt": "What is 2+2?",
|
||||
"model": "default",
|
||||
"stream": false,
|
||||
"iterations": 2
|
||||
},
|
||||
"import_time_stats": {
|
||||
"mean": 0.008833496394800022,
|
||||
"min": 0.005916997906751931,
|
||||
"max": 0.011749994882848114,
|
||||
"stdev": 0.0041245517164382375,
|
||||
"median": 0.008833496394800022
|
||||
},
|
||||
"total_time_stats": {
|
||||
"mean": 1727.4090415012324,
|
||||
"min": 1367.2264170018025,
|
||||
"max": 2087.5916660006624,
|
||||
"stdev": 509.3751524982296,
|
||||
"median": 1727.4090415012324
|
||||
},
|
||||
"first_token_stats": null,
|
||||
"raw_import_times_ms": [
|
||||
0.011749994882848114,
|
||||
0.005916997906751931
|
||||
],
|
||||
"raw_total_times_ms": [
|
||||
1367.2264170018025,
|
||||
2087.5916660006624
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "simple_stream",
|
||||
"config": {
|
||||
"prompt": "What is 2+2?",
|
||||
"model": "default",
|
||||
"stream": true,
|
||||
"iterations": 2
|
||||
},
|
||||
"import_time_stats": {
|
||||
"mean": 0.008792001608526334,
|
||||
"min": 0.006375004886649549,
|
||||
"max": 0.01120899833040312,
|
||||
"stdev": 0.0034181495442894613,
|
||||
"median": 0.008792001608526334
|
||||
},
|
||||
"total_time_stats": {
|
||||
"mean": 2321.0982290001994,
|
||||
"min": 1697.1089580038097,
|
||||
"max": 2945.087499996589,
|
||||
"stdev": 882.4540898183949,
|
||||
"median": 2321.0982290001994
|
||||
},
|
||||
"first_token_stats": {
|
||||
"mean": 2312.8511880022415,
|
||||
"min": 1681.2440420035273,
|
||||
"max": 2944.458334000956,
|
||||
"stdev": 893.2273919631452,
|
||||
"median": 2312.8511880022415
|
||||
},
|
||||
"raw_import_times_ms": [
|
||||
0.01120899833040312,
|
||||
0.006375004886649549
|
||||
],
|
||||
"raw_total_times_ms": [
|
||||
2945.087499996589,
|
||||
1697.1089580038097
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
"""
|
||||
Profile Suite Example for PraisonAI
|
||||
|
||||
This example demonstrates how to run a comprehensive profiling suite
|
||||
programmatically with multiple scenarios.
|
||||
|
||||
Run with:
|
||||
python suite_profiling.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
ScenarioConfig,
|
||||
run_profile_suite,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Run a profiling suite with custom scenarios."""
|
||||
|
||||
# Define custom scenarios
|
||||
scenarios = [
|
||||
ScenarioConfig(
|
||||
name="simple_query",
|
||||
prompt="What is 2+2?",
|
||||
stream=False,
|
||||
iterations=2,
|
||||
warmup=1,
|
||||
),
|
||||
ScenarioConfig(
|
||||
name="simple_stream",
|
||||
prompt="What is 2+2?",
|
||||
stream=True,
|
||||
iterations=2,
|
||||
warmup=1,
|
||||
),
|
||||
]
|
||||
|
||||
# Run suite with custom output directory
|
||||
print("🔬 Running Profile Suite...")
|
||||
result = run_profile_suite(
|
||||
output_dir="./profile_results",
|
||||
scenarios=scenarios,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Print summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Profile Suite Results")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n📊 Startup Times:")
|
||||
print(f" Cold Start: {result.startup_cold_ms:.2f}ms")
|
||||
print(f" Warm Start: {result.startup_warm_ms:.2f}ms")
|
||||
|
||||
if result.import_analysis:
|
||||
print("\n📦 Top Import:")
|
||||
print(f" {result.import_analysis[0]['module']}: {result.import_analysis[0]['cumulative_ms']:.2f}ms")
|
||||
|
||||
print("\n🎯 Scenario Results:")
|
||||
for scenario in result.scenarios:
|
||||
stats = scenario.get_stats(scenario.total_times)
|
||||
print(f" {scenario.name}:")
|
||||
print(f" Mean: {stats['mean']:.2f}ms")
|
||||
print(f" Min: {stats['min']:.2f}ms")
|
||||
print(f" Max: {stats['max']:.2f}ms")
|
||||
|
||||
print("\n✅ Results saved to: ./profile_results/")
|
||||
print(" - suite_results.json (machine-readable)")
|
||||
print(" - suite_report.txt (human-readable)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
Basic TUI Usage Example for PraisonAI.
|
||||
|
||||
This example demonstrates how to launch the interactive TUI.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
|
||||
def main():
|
||||
"""Launch the TUI."""
|
||||
try:
|
||||
from praisonai.cli.features.tui import run_tui
|
||||
|
||||
# Launch with default settings
|
||||
run_tui()
|
||||
|
||||
except ImportError:
|
||||
print("Textual is required for TUI. Install with:")
|
||||
print(" pip install praisonai[tui]")
|
||||
exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
Headless TUI Simulation Example for PraisonAI.
|
||||
|
||||
This example demonstrates how to run TUI simulations without
|
||||
launching the interactive UI - useful for testing and CI.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonai.cli.features.tui import (
|
||||
TuiOrchestrator,
|
||||
SimulationRunner,
|
||||
MockProvider,
|
||||
)
|
||||
from praisonai.cli.features.tui.orchestrator import OutputMode
|
||||
from praisonai.cli.features.queue import QueueConfig
|
||||
|
||||
|
||||
async def run_simulation():
|
||||
"""Run a headless simulation."""
|
||||
|
||||
# Create orchestrator with mock provider (no real API calls)
|
||||
config = QueueConfig(enable_persistence=False)
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.PRETTY,
|
||||
debug=True,
|
||||
)
|
||||
|
||||
# Create simulation runner
|
||||
runner = SimulationRunner(orchestrator, assert_mode=True)
|
||||
|
||||
# Define simulation script
|
||||
script = {
|
||||
"session_id": "sim-example",
|
||||
"model": "gpt-4o-mini",
|
||||
"steps": [
|
||||
# Navigate to different screens
|
||||
{"action": "navigate", "args": {"screen": "main"}},
|
||||
{"action": "focus", "args": {"widget": "composer"}},
|
||||
|
||||
# Change model
|
||||
{
|
||||
"action": "model",
|
||||
"args": {"model": "gpt-4"},
|
||||
"expected": {"model": "gpt-4"}
|
||||
},
|
||||
|
||||
# Navigate to queue screen
|
||||
{"action": "navigate", "args": {"screen": "queue"}},
|
||||
|
||||
# Take a snapshot
|
||||
{"action": "snapshot"},
|
||||
|
||||
# Navigate back
|
||||
{"action": "navigate", "args": {"screen": "main"}},
|
||||
|
||||
# Small delay
|
||||
{"action": "sleep", "args": {"seconds": 0.5}},
|
||||
]
|
||||
}
|
||||
|
||||
# Run the simulation
|
||||
print("Running headless TUI simulation...")
|
||||
print("=" * 50)
|
||||
|
||||
success = await runner.run_script(script)
|
||||
|
||||
# Print results
|
||||
print("=" * 50)
|
||||
summary = runner.get_summary()
|
||||
|
||||
print(f"\nSimulation {'PASSED' if success else 'FAILED'}")
|
||||
print(f" Assertions passed: {summary['assertions_passed']}")
|
||||
print(f" Assertions failed: {summary['assertions_failed']}")
|
||||
|
||||
if summary['errors']:
|
||||
print("\nErrors:")
|
||||
for error in summary['errors']:
|
||||
print(f" - {error}")
|
||||
|
||||
return success
|
||||
|
||||
|
||||
async def run_with_events():
|
||||
"""Run simulation and capture all events."""
|
||||
|
||||
config = QueueConfig(enable_persistence=False)
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
# Capture events
|
||||
events = []
|
||||
orchestrator.add_event_callback(lambda e: events.append(e))
|
||||
|
||||
await orchestrator.start(session_id="event-capture")
|
||||
|
||||
# Perform some actions
|
||||
orchestrator.set_model("gpt-4")
|
||||
orchestrator.navigate_screen("queue")
|
||||
orchestrator.set_focus("queue-panel")
|
||||
orchestrator.navigate_screen("main")
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
# Print captured events
|
||||
print("\nCaptured Events:")
|
||||
print("-" * 40)
|
||||
for event in events:
|
||||
print(f" {event.event_type.value}")
|
||||
|
||||
return events
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point."""
|
||||
print("PraisonAI TUI Headless Simulation Example")
|
||||
print("=" * 50)
|
||||
|
||||
# Run simulation
|
||||
success = asyncio.run(run_simulation())
|
||||
|
||||
# Run event capture
|
||||
print("\n" + "=" * 50)
|
||||
events = asyncio.run(run_with_events())
|
||||
|
||||
print(f"\nTotal events captured: {len(events)}")
|
||||
|
||||
return 0 if success else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
Queue Management Example for PraisonAI.
|
||||
|
||||
This example demonstrates how to use the queue system programmatically.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonai.cli.features.queue import (
|
||||
QueueManager,
|
||||
QueueConfig,
|
||||
QueuedRun,
|
||||
RunState,
|
||||
RunPriority,
|
||||
)
|
||||
|
||||
|
||||
async def main():
|
||||
"""Demonstrate queue management."""
|
||||
|
||||
print("PraisonAI Queue Management Example")
|
||||
print("=" * 50)
|
||||
|
||||
# Create queue configuration
|
||||
config = QueueConfig(
|
||||
max_concurrent_global=2,
|
||||
max_concurrent_per_agent=1,
|
||||
max_queue_size=10,
|
||||
enable_persistence=False, # Use in-memory for demo
|
||||
)
|
||||
|
||||
# Track outputs
|
||||
outputs = {}
|
||||
|
||||
async def on_output(run_id: str, chunk: str):
|
||||
if run_id not in outputs:
|
||||
outputs[run_id] = ""
|
||||
outputs[run_id] += chunk
|
||||
print(f" [{run_id[:8]}] {chunk}", end="", flush=True)
|
||||
|
||||
async def on_complete(run_id: str, run: QueuedRun):
|
||||
print(f"\n [{run_id[:8]}] Completed: {run.state.value}")
|
||||
|
||||
async def on_error(run_id: str, error: Exception):
|
||||
print(f"\n [{run_id[:8]}] Error: {error}")
|
||||
|
||||
# Create queue manager
|
||||
manager = QueueManager(
|
||||
config=config,
|
||||
on_output=on_output,
|
||||
on_complete=on_complete,
|
||||
on_error=on_error,
|
||||
)
|
||||
|
||||
# Start the manager
|
||||
await manager.start(recover=False)
|
||||
print("Queue manager started")
|
||||
|
||||
# Submit some runs with different priorities
|
||||
print("\nSubmitting runs...")
|
||||
|
||||
run1 = await manager.submit(
|
||||
input_content="Low priority task",
|
||||
agent_name="Agent1",
|
||||
priority=RunPriority.LOW,
|
||||
)
|
||||
print(f" Submitted run1: {run1} (LOW priority)")
|
||||
|
||||
run2 = await manager.submit(
|
||||
input_content="High priority task",
|
||||
agent_name="Agent2",
|
||||
priority=RunPriority.HIGH,
|
||||
)
|
||||
print(f" Submitted run2: {run2} (HIGH priority)")
|
||||
|
||||
run3 = await manager.submit(
|
||||
input_content="Normal priority task",
|
||||
agent_name="Agent1",
|
||||
priority=RunPriority.NORMAL,
|
||||
)
|
||||
print(f" Submitted run3: {run3} (NORMAL priority)")
|
||||
|
||||
# Check queue status
|
||||
print(f"\nQueue status:")
|
||||
print(f" Queued: {manager.queued_count}")
|
||||
print(f" Running: {manager.running_count}")
|
||||
|
||||
# List runs
|
||||
runs = manager.list_runs()
|
||||
print(f"\nAll runs ({len(runs)}):")
|
||||
for run in runs:
|
||||
print(f" {run.run_id[:8]}: {run.agent_name} - {run.state.value} ({run.priority.name})")
|
||||
|
||||
# Wait a bit for processing
|
||||
print("\nWaiting for processing...")
|
||||
await asyncio.sleep(2)
|
||||
|
||||
# Cancel a run
|
||||
if runs:
|
||||
cancel_id = runs[-1].run_id
|
||||
print(f"\nCancelling run: {cancel_id[:8]}")
|
||||
await manager.cancel(cancel_id)
|
||||
|
||||
# Get stats
|
||||
stats = manager.get_stats()
|
||||
print(f"\nQueue statistics:")
|
||||
print(f" Total runs: {stats.total_runs}")
|
||||
print(f" Queued: {stats.queued_count}")
|
||||
print(f" Running: {stats.running_count}")
|
||||
print(f" Succeeded: {stats.succeeded_count}")
|
||||
print(f" Failed: {stats.failed_count}")
|
||||
print(f" Cancelled: {stats.cancelled_count}")
|
||||
|
||||
# Stop the manager
|
||||
await manager.stop()
|
||||
print("\nQueue manager stopped")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("Example complete!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.5.3"
|
||||
version = "0.6.0"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
Generated
+1
-1
@@ -2422,7 +2422,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.5.3"
|
||||
version = "0.6.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -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.9.1.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.9.1.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.0.2.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v3.0.2.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -176,6 +176,11 @@ def register_commands():
|
||||
from .commands.serve import app as serve_app
|
||||
from .commands.schedule import app as schedule_app
|
||||
from .commands.run import app as run_app
|
||||
from .commands.profile import app as profile_app
|
||||
|
||||
# Import TUI and queue commands
|
||||
from .features.tui.debug import create_debug_app as create_tui_debug_app
|
||||
from .features.tui.cli import create_queue_app
|
||||
|
||||
# Register sub-apps
|
||||
app.add_typer(config_app, name="config", help="Configuration management")
|
||||
@@ -193,6 +198,15 @@ def register_commands():
|
||||
app.add_typer(serve_app, name="serve", help="API server management")
|
||||
app.add_typer(schedule_app, name="schedule", help="Scheduler management")
|
||||
app.add_typer(run_app, name="run", help="Run agents")
|
||||
app.add_typer(profile_app, name="profile", help="Performance profiling and diagnostics")
|
||||
|
||||
# Register TUI and queue commands
|
||||
tui_app = create_tui_debug_app()
|
||||
queue_app = create_queue_app()
|
||||
if tui_app:
|
||||
app.add_typer(tui_app, name="tui", help="Interactive TUI and simulation")
|
||||
if queue_app:
|
||||
app.add_typer(queue_app, name="queue", help="Queue management")
|
||||
|
||||
|
||||
# Register commands on import
|
||||
|
||||
@@ -0,0 +1,466 @@
|
||||
"""
|
||||
Profile command group for PraisonAI CLI.
|
||||
|
||||
Provides detailed cProfile-based profiling for query execution.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
|
||||
app = typer.Typer(help="Performance profiling and diagnostics")
|
||||
|
||||
|
||||
@app.command("query")
|
||||
def profile_query(
|
||||
prompt: str = typer.Argument(..., help="Prompt to profile"),
|
||||
model: Optional[str] = typer.Option(None, "--model", "-m", help="Model to use"),
|
||||
stream: bool = typer.Option(False, "--stream/--no-stream", help="Use streaming mode"),
|
||||
deep: bool = typer.Option(False, "--deep", help="Enable deep call tracing (higher overhead)"),
|
||||
limit: int = typer.Option(30, "--limit", "-n", help="Top N functions to show"),
|
||||
sort: str = typer.Option("cumulative", "--sort", "-s", help="Sort by: cumulative or tottime"),
|
||||
show_callers: bool = typer.Option(False, "--show-callers", help="Show caller functions"),
|
||||
show_callees: bool = typer.Option(False, "--show-callees", help="Show callee functions"),
|
||||
save: Optional[str] = typer.Option(None, "--save", help="Save artifacts to path (creates .prof, .txt, .json)"),
|
||||
output_format: str = typer.Option("text", "--format", "-f", help="Output format: text or json"),
|
||||
):
|
||||
"""
|
||||
Profile a query execution with detailed timing breakdown.
|
||||
|
||||
Uses the unified profiling architecture for consistent results across
|
||||
CLI direct invocation (--profile) and this command.
|
||||
|
||||
Shows per-function timing, call graphs, and latency metrics.
|
||||
|
||||
Examples:
|
||||
praisonai profile query "What is 2+2?"
|
||||
praisonai profile query "Hello" --limit 20
|
||||
praisonai profile query "Test" --deep --show-callers --show-callees
|
||||
praisonai profile query "Test" --save ./profile_results
|
||||
"""
|
||||
# Use unified profiler for consistent results
|
||||
try:
|
||||
from ..execution import ExecutionRequest, Profiler, ProfilerConfig as UnifiedProfilerConfig
|
||||
except ImportError as e:
|
||||
typer.echo(f"Error: Unified profiler not available: {e}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Warn about deep tracing overhead
|
||||
if deep:
|
||||
typer.echo("⚠️ Deep call tracing enabled - this adds significant overhead", err=True)
|
||||
|
||||
# Create unified execution request
|
||||
request = ExecutionRequest(
|
||||
prompt=prompt,
|
||||
agent_name="ProfiledAgent",
|
||||
model=model,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
# Create unified profiler config
|
||||
config = UnifiedProfilerConfig(
|
||||
layer=2 if deep else 1,
|
||||
limit=limit,
|
||||
sort_by=sort,
|
||||
show_callers=show_callers,
|
||||
show_callees=show_callees,
|
||||
output_format=output_format,
|
||||
save_path=save,
|
||||
)
|
||||
|
||||
# Run profiler
|
||||
typer.echo("🔬 Starting profiled execution...", err=True)
|
||||
|
||||
try:
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request, invocation_method="profile_command")
|
||||
except ImportError as e:
|
||||
typer.echo(f"Error: {e}", err=True)
|
||||
typer.echo("Install praisonaiagents: pip install praisonaiagents", err=True)
|
||||
raise typer.Exit(1)
|
||||
except Exception as e:
|
||||
typer.echo(f"Error during profiling: {e}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Output results
|
||||
if output_format == "json":
|
||||
typer.echo(report.to_json())
|
||||
else:
|
||||
typer.echo(report.to_text())
|
||||
|
||||
# Save artifacts if requested
|
||||
if save:
|
||||
try:
|
||||
import os
|
||||
os.makedirs(save, exist_ok=True)
|
||||
|
||||
# Save JSON report
|
||||
json_path = os.path.join(save, "profile_report.json")
|
||||
with open(json_path, 'w') as f:
|
||||
f.write(report.to_json())
|
||||
|
||||
# Save text report
|
||||
txt_path = os.path.join(save, "profile_report.txt")
|
||||
with open(txt_path, 'w') as f:
|
||||
f.write(report.to_text())
|
||||
|
||||
typer.echo("\n✅ Artifacts saved:", err=True)
|
||||
typer.echo(f" {json_path}", err=True)
|
||||
typer.echo(f" {txt_path}", err=True)
|
||||
except Exception as e:
|
||||
typer.echo(f"Warning: Failed to save artifacts: {e}", err=True)
|
||||
|
||||
|
||||
@app.command("imports")
|
||||
def profile_imports():
|
||||
"""
|
||||
Profile module import times.
|
||||
|
||||
Shows which modules take longest to import, useful for optimizing startup time.
|
||||
"""
|
||||
import subprocess
|
||||
import re
|
||||
|
||||
typer.echo("🔬 Profiling import times...", err=True)
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-X", "importtime", "-c", "import praisonaiagents"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
# Parse import times
|
||||
import_times = []
|
||||
for line in result.stderr.split('\n'):
|
||||
if 'import time:' in line:
|
||||
match = re.search(r'import time:\s+(\d+)\s+\|\s+(\d+)\s+\|\s+(.+)', line)
|
||||
if match:
|
||||
self_time = int(match.group(1))
|
||||
cumulative = int(match.group(2))
|
||||
module = match.group(3).strip()
|
||||
import_times.append((module, self_time, cumulative))
|
||||
|
||||
# Sort by cumulative time
|
||||
import_times.sort(key=lambda x: x[2], reverse=True)
|
||||
|
||||
# Display
|
||||
typer.echo("\n" + "=" * 70)
|
||||
typer.echo("Import Time Analysis")
|
||||
typer.echo("=" * 70)
|
||||
typer.echo(f"{'Module':<45} {'Self (μs)':>10} {'Cumul (μs)':>12}")
|
||||
typer.echo("-" * 70)
|
||||
|
||||
for module, self_time, cumulative in import_times[:30]:
|
||||
module_name = module
|
||||
if len(module_name) > 43:
|
||||
module_name = "..." + module_name[-40:]
|
||||
typer.echo(f"{module_name:<45} {self_time:>10} {cumulative:>12}")
|
||||
|
||||
typer.echo("-" * 70)
|
||||
|
||||
if import_times:
|
||||
total_time = max(t[2] for t in import_times)
|
||||
typer.echo(f"Total import time: {total_time / 1000:.2f} ms")
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
typer.echo("Error: Import profiling timed out", err=True)
|
||||
raise typer.Exit(1)
|
||||
except Exception as e:
|
||||
typer.echo(f"Error: {e}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
@app.command("startup")
|
||||
def profile_startup():
|
||||
"""
|
||||
Profile CLI startup time.
|
||||
|
||||
Measures time from CLI invocation to ready state.
|
||||
"""
|
||||
import time
|
||||
import subprocess
|
||||
|
||||
typer.echo("🔬 Profiling startup time...", err=True)
|
||||
|
||||
# Measure cold start (new process)
|
||||
start = time.perf_counter()
|
||||
subprocess.run(
|
||||
[sys.executable, "-c", "import praisonai; import praisonai.cli"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
cold_start = (time.perf_counter() - start) * 1000
|
||||
|
||||
# Measure warm start (imports cached)
|
||||
start = time.perf_counter()
|
||||
subprocess.run(
|
||||
[sys.executable, "-c", "import praisonai; import praisonai.cli"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
warm_start = (time.perf_counter() - start) * 1000
|
||||
|
||||
typer.echo("\n" + "=" * 50)
|
||||
typer.echo("Startup Time Analysis")
|
||||
typer.echo("=" * 50)
|
||||
typer.echo(f"Cold Start: {cold_start:>10.2f} ms")
|
||||
typer.echo(f"Warm Start: {warm_start:>10.2f} ms")
|
||||
typer.echo("=" * 50)
|
||||
|
||||
if cold_start > 1000:
|
||||
typer.echo("\n⚠️ Cold start > 1s - consider lazy imports", err=True)
|
||||
|
||||
|
||||
@app.command("suite")
|
||||
def profile_suite(
|
||||
output_dir: Optional[str] = typer.Option("/tmp/praisonai_profile_suite", "--output", "-o", help="Output directory for results"),
|
||||
iterations: int = typer.Option(3, "--iterations", "-n", help="Iterations per scenario"),
|
||||
quick: bool = typer.Option(False, "--quick", help="Quick mode (fewer iterations)"),
|
||||
):
|
||||
"""
|
||||
Run a comprehensive profiling suite.
|
||||
|
||||
Runs multiple scenarios (streaming/non-streaming, simple/complex prompts)
|
||||
and produces aggregated statistics and reports.
|
||||
|
||||
Examples:
|
||||
praisonai profile suite
|
||||
praisonai profile suite --output ./my_profile_results
|
||||
praisonai profile suite --quick
|
||||
"""
|
||||
try:
|
||||
from ..features.profiler import run_profile_suite, ScenarioConfig
|
||||
except ImportError as e:
|
||||
typer.echo(f"Error: Profiler module not available: {e}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Define scenarios
|
||||
iters = 1 if quick else iterations
|
||||
scenarios = [
|
||||
ScenarioConfig(name="simple_non_stream", prompt="hi", stream=False, iterations=iters, warmup=0 if quick else 1),
|
||||
ScenarioConfig(name="simple_stream", prompt="hi", stream=True, iterations=iters, warmup=0 if quick else 1),
|
||||
]
|
||||
|
||||
if not quick:
|
||||
scenarios.extend([
|
||||
ScenarioConfig(name="medium_non_stream", prompt="Explain Python in 2 sentences", stream=False, iterations=max(1, iters-1)),
|
||||
ScenarioConfig(name="medium_stream", prompt="Explain Python in 2 sentences", stream=True, iterations=max(1, iters-1)),
|
||||
])
|
||||
|
||||
typer.echo("🔬 Running Profile Suite...", err=True)
|
||||
typer.echo(f" Output: {output_dir}", err=True)
|
||||
typer.echo(f" Scenarios: {len(scenarios)}", err=True)
|
||||
typer.echo(f" Iterations: {iterations}", err=True)
|
||||
|
||||
try:
|
||||
result = run_profile_suite(
|
||||
output_dir=output_dir,
|
||||
scenarios=scenarios,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Print summary
|
||||
typer.echo("\n" + "=" * 60)
|
||||
typer.echo("Profile Suite Summary")
|
||||
typer.echo("=" * 60)
|
||||
typer.echo(f"Startup Cold: {result.startup_cold_ms:.2f}ms")
|
||||
typer.echo(f"Startup Warm: {result.startup_warm_ms:.2f}ms")
|
||||
|
||||
if result.import_analysis:
|
||||
typer.echo(f"\nTop Import: {result.import_analysis[0]['module']}")
|
||||
typer.echo(f" Time: {result.import_analysis[0]['cumulative_ms']:.2f}ms")
|
||||
|
||||
typer.echo("\nScenario Results:")
|
||||
for scenario in result.scenarios:
|
||||
stats = scenario.get_stats(scenario.total_times)
|
||||
typer.echo(f" {scenario.name}: {stats['mean']:.2f}ms (±{stats['stdev']:.2f}ms)")
|
||||
|
||||
typer.echo(f"\n✅ Full results saved to: {output_dir}")
|
||||
|
||||
except Exception as e:
|
||||
typer.echo(f"Error running suite: {e}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
@app.command("snapshot")
|
||||
def profile_snapshot(
|
||||
name: str = typer.Argument("current", help="Name for the snapshot"),
|
||||
baseline: bool = typer.Option(False, "--baseline", "-b", help="Save as baseline for comparison"),
|
||||
compare: bool = typer.Option(False, "--compare", "-c", help="Compare against baseline"),
|
||||
output_format: str = typer.Option("text", "--format", "-f", help="Output format: text or json"),
|
||||
):
|
||||
"""
|
||||
Create or compare performance snapshots.
|
||||
|
||||
Snapshots record timing data for baseline comparison.
|
||||
Use --baseline to save a reference point, then --compare to check for regressions.
|
||||
|
||||
Examples:
|
||||
praisonai profile snapshot --baseline
|
||||
praisonai profile snapshot current --compare
|
||||
praisonai profile snapshot v2.0 --format json
|
||||
"""
|
||||
try:
|
||||
from ..features.profiler import (
|
||||
run_profile_suite,
|
||||
ScenarioConfig,
|
||||
PerfSnapshotManager,
|
||||
create_snapshot_from_suite,
|
||||
format_comparison_report,
|
||||
)
|
||||
except ImportError as e:
|
||||
typer.echo(f"Error: Profiler module not available: {e}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
manager = PerfSnapshotManager()
|
||||
|
||||
if compare:
|
||||
# Load baseline and compare
|
||||
baseline_snap = manager.load_baseline()
|
||||
if not baseline_snap:
|
||||
typer.echo("❌ No baseline found. Create one with: praisonai profile snapshot --baseline", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
typer.echo("🔬 Running current profile for comparison...", err=True)
|
||||
|
||||
# Run quick suite
|
||||
scenarios = [
|
||||
ScenarioConfig(name="simple", prompt="hi", stream=False, iterations=2, warmup=1),
|
||||
]
|
||||
result = run_profile_suite(
|
||||
output_dir="/tmp/praisonai_snapshot_compare",
|
||||
scenarios=scenarios,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
current_snap = create_snapshot_from_suite(result, name)
|
||||
comparison = manager.compare(current_snap)
|
||||
|
||||
if output_format == "json":
|
||||
import json
|
||||
typer.echo(json.dumps(comparison.to_dict(), indent=2))
|
||||
else:
|
||||
typer.echo(format_comparison_report(comparison))
|
||||
|
||||
if comparison.is_regression():
|
||||
raise typer.Exit(1) # Exit with error on regression
|
||||
else:
|
||||
# Create new snapshot
|
||||
typer.echo(f"🔬 Creating performance snapshot '{name}'...", err=True)
|
||||
|
||||
scenarios = [
|
||||
ScenarioConfig(name="simple", prompt="hi", stream=False, iterations=3, warmup=1),
|
||||
ScenarioConfig(name="stream", prompt="hi", stream=True, iterations=2, warmup=1),
|
||||
]
|
||||
result = run_profile_suite(
|
||||
output_dir="/tmp/praisonai_snapshot",
|
||||
scenarios=scenarios,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
snapshot = create_snapshot_from_suite(result, name)
|
||||
|
||||
if baseline:
|
||||
path = manager.save_baseline(snapshot)
|
||||
typer.echo(f"\n✅ Baseline saved to: {path}")
|
||||
else:
|
||||
path = manager.save(snapshot)
|
||||
typer.echo(f"\n✅ Snapshot saved to: {path}")
|
||||
|
||||
if output_format == "json":
|
||||
import json
|
||||
typer.echo(json.dumps(snapshot.to_dict(), indent=2))
|
||||
else:
|
||||
typer.echo("\nSnapshot Summary:")
|
||||
typer.echo(f" Startup Cold: {snapshot.startup_cold_ms:.2f}ms")
|
||||
typer.echo(f" Import Time: {snapshot.import_time_ms:.2f}ms")
|
||||
typer.echo(f" Query Time: {snapshot.query_time_ms:.2f}ms")
|
||||
|
||||
|
||||
@app.command("optimize")
|
||||
def profile_optimize(
|
||||
prewarm: bool = typer.Option(False, "--prewarm", help="Enable provider pre-warming"),
|
||||
lite: bool = typer.Option(False, "--lite", help="Enable lite mode (skip type validation)"),
|
||||
show_config: bool = typer.Option(False, "--show", help="Show current optimization config"),
|
||||
):
|
||||
"""
|
||||
Configure performance optimizations.
|
||||
|
||||
All optimizations are opt-in and safe by default.
|
||||
|
||||
Examples:
|
||||
praisonai profile optimize --show
|
||||
praisonai profile optimize --prewarm
|
||||
praisonai profile optimize --lite
|
||||
"""
|
||||
try:
|
||||
from ..features.profiler import (
|
||||
PrewarmManager,
|
||||
get_lite_mode_config,
|
||||
get_provider_cache,
|
||||
)
|
||||
except ImportError as e:
|
||||
typer.echo(f"Error: Profiler module not available: {e}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
if show_config:
|
||||
typer.echo("=" * 50)
|
||||
typer.echo("Performance Optimization Status")
|
||||
typer.echo("=" * 50)
|
||||
|
||||
lite_config = get_lite_mode_config()
|
||||
typer.echo(f"\nLite Mode: {'Enabled' if lite_config.enabled else 'Disabled'}")
|
||||
if lite_config.enabled:
|
||||
typer.echo(f" Skip Type Validation: {lite_config.skip_type_validation}")
|
||||
typer.echo(f" Skip Model Validation: {lite_config.skip_model_validation}")
|
||||
typer.echo(f" Minimal Imports: {lite_config.minimal_imports}")
|
||||
|
||||
typer.echo(f"\nPre-warming: {'Enabled' if PrewarmManager.is_enabled() else 'Disabled'}")
|
||||
|
||||
cache = get_provider_cache()
|
||||
stats = cache.get_stats()
|
||||
typer.echo("\nProvider Cache:")
|
||||
typer.echo(f" Hits: {stats['hits']}")
|
||||
typer.echo(f" Misses: {stats['misses']}")
|
||||
typer.echo(f" Size: {stats['size']}")
|
||||
|
||||
typer.echo("\n" + "=" * 50)
|
||||
typer.echo("\nTo enable optimizations, set environment variables:")
|
||||
typer.echo(" PRAISONAI_LITE_MODE=1")
|
||||
typer.echo(" PRAISONAI_SKIP_TYPE_VALIDATION=1")
|
||||
typer.echo(" PRAISONAI_MINIMAL_IMPORTS=1")
|
||||
return
|
||||
|
||||
if prewarm:
|
||||
typer.echo("🔥 Enabling provider pre-warming...")
|
||||
PrewarmManager.enable()
|
||||
|
||||
import os
|
||||
if os.environ.get('OPENAI_API_KEY'):
|
||||
typer.echo(" Pre-warming OpenAI...")
|
||||
PrewarmManager.prewarm_provider('openai')
|
||||
if os.environ.get('ANTHROPIC_API_KEY'):
|
||||
typer.echo(" Pre-warming Anthropic...")
|
||||
PrewarmManager.prewarm_provider('anthropic')
|
||||
|
||||
typer.echo("✅ Pre-warming initiated (runs in background)")
|
||||
|
||||
if lite:
|
||||
typer.echo("⚡ Lite mode configuration:")
|
||||
typer.echo(" Set these environment variables to enable:")
|
||||
typer.echo(" export PRAISONAI_LITE_MODE=1")
|
||||
typer.echo(" export PRAISONAI_SKIP_TYPE_VALIDATION=1")
|
||||
typer.echo(" export PRAISONAI_MINIMAL_IMPORTS=1")
|
||||
|
||||
if not prewarm and not lite and not show_config:
|
||||
typer.echo("Use --show to see current config, or --prewarm/--lite to enable optimizations")
|
||||
|
||||
|
||||
@app.callback(invoke_without_command=True)
|
||||
def profile_callback(ctx: typer.Context):
|
||||
"""Show profile help if no subcommand."""
|
||||
if ctx.invoked_subcommand is None:
|
||||
typer.echo(ctx.get_help())
|
||||
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
PraisonAI Unified Execution Module.
|
||||
|
||||
Provides a single, unified execution primitive for all CLI modes:
|
||||
- CLI direct prompt
|
||||
- Profile command
|
||||
- TUI workers
|
||||
|
||||
Design Principles:
|
||||
1. Profiling observes, never alters execution
|
||||
2. Zero overhead when profiling disabled
|
||||
3. Sync and async are first-class citizens
|
||||
4. Fail-safe, not fail-fast
|
||||
5. Bounded resource usage
|
||||
6. Schema stability with versioning
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .request import ExecutionRequest
|
||||
from .result import ExecutionResult
|
||||
from .core import execute_sync, execute_async
|
||||
from .profiler import Profiler, ProfilerConfig, ProfileReport
|
||||
|
||||
# Lazy loading to maintain zero-overhead when not used
|
||||
_lazy_cache = {}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load execution components."""
|
||||
global _lazy_cache
|
||||
|
||||
if name in _lazy_cache:
|
||||
return _lazy_cache[name]
|
||||
|
||||
if name == "ExecutionRequest":
|
||||
from .request import ExecutionRequest
|
||||
_lazy_cache[name] = ExecutionRequest
|
||||
return ExecutionRequest
|
||||
elif name == "ExecutionResult":
|
||||
from .result import ExecutionResult
|
||||
_lazy_cache[name] = ExecutionResult
|
||||
return ExecutionResult
|
||||
elif name == "execute_sync":
|
||||
from .core import execute_sync
|
||||
_lazy_cache[name] = execute_sync
|
||||
return execute_sync
|
||||
elif name == "execute_async":
|
||||
from .core import execute_async
|
||||
_lazy_cache[name] = execute_async
|
||||
return execute_async
|
||||
elif name == "Profiler":
|
||||
from .profiler import Profiler
|
||||
_lazy_cache[name] = Profiler
|
||||
return Profiler
|
||||
elif name == "ProfilerConfig":
|
||||
from .profiler import ProfilerConfig
|
||||
_lazy_cache[name] = ProfilerConfig
|
||||
return ProfilerConfig
|
||||
elif name == "ProfileReport":
|
||||
from .profiler import ProfileReport
|
||||
_lazy_cache[name] = ProfileReport
|
||||
return ProfileReport
|
||||
elif name == "TimingBreakdown":
|
||||
from .profiler import TimingBreakdown
|
||||
_lazy_cache[name] = TimingBreakdown
|
||||
return TimingBreakdown
|
||||
elif name == "InvocationInfo":
|
||||
from .profiler import InvocationInfo
|
||||
_lazy_cache[name] = InvocationInfo
|
||||
return InvocationInfo
|
||||
elif name == "DecisionTrace":
|
||||
from .profiler import DecisionTrace
|
||||
_lazy_cache[name] = DecisionTrace
|
||||
return DecisionTrace
|
||||
elif name == "ModuleBreakdown":
|
||||
from .profiler import ModuleBreakdown
|
||||
_lazy_cache[name] = ModuleBreakdown
|
||||
return ModuleBreakdown
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Core execution
|
||||
"ExecutionRequest",
|
||||
"ExecutionResult",
|
||||
"execute_sync",
|
||||
"execute_async",
|
||||
# Profiling
|
||||
"Profiler",
|
||||
"ProfilerConfig",
|
||||
"ProfileReport",
|
||||
"TimingBreakdown",
|
||||
"InvocationInfo",
|
||||
"DecisionTrace",
|
||||
"ModuleBreakdown",
|
||||
]
|
||||
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
Core Execution Module.
|
||||
|
||||
Provides the unified execution primitives:
|
||||
- _execute_core(): Single implementation shared by all paths
|
||||
- execute_sync(): Synchronous wrapper for CLI
|
||||
- execute_async(): Asynchronous wrapper for TUI workers
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from typing import Optional, AsyncGenerator
|
||||
|
||||
from .request import ExecutionRequest
|
||||
from .result import ExecutionResult
|
||||
|
||||
|
||||
def _generate_run_id() -> str:
|
||||
"""Generate a unique run ID."""
|
||||
return str(uuid.uuid4())[:8]
|
||||
|
||||
|
||||
def _execute_core(
|
||||
request: ExecutionRequest,
|
||||
timing: Optional[dict] = None,
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
THE SINGLE EXECUTION IMPLEMENTATION.
|
||||
|
||||
All execution modes call this function.
|
||||
Profiler wraps this function, never injects into it.
|
||||
|
||||
Args:
|
||||
request: The execution request
|
||||
timing: Optional dict to record timing phases (for profiling)
|
||||
|
||||
Returns:
|
||||
ExecutionResult with output and metadata
|
||||
"""
|
||||
run_id = _generate_run_id()
|
||||
start_time = time.perf_counter()
|
||||
|
||||
if timing is not None:
|
||||
timing["core_start"] = start_time
|
||||
|
||||
try:
|
||||
# Phase 1: Import (measured inline)
|
||||
import_start = time.perf_counter()
|
||||
from praisonaiagents import Agent
|
||||
import_end = time.perf_counter()
|
||||
|
||||
if timing is not None:
|
||||
timing["imports_ms"] = (import_end - import_start) * 1000
|
||||
|
||||
# Phase 2: Agent Construction
|
||||
agent_start = time.perf_counter()
|
||||
|
||||
agent_config = {
|
||||
"name": request.agent_name,
|
||||
"verbose": False,
|
||||
}
|
||||
|
||||
if request.agent_instructions:
|
||||
agent_config["instructions"] = request.agent_instructions
|
||||
|
||||
if request.model:
|
||||
agent_config["llm"] = request.model
|
||||
|
||||
if request.tools:
|
||||
agent_config["tools"] = list(request.tools)
|
||||
|
||||
agent = Agent(**agent_config)
|
||||
agent_end = time.perf_counter()
|
||||
|
||||
if timing is not None:
|
||||
timing["agent_init_ms"] = (agent_end - agent_start) * 1000
|
||||
|
||||
# Phase 3: Execution
|
||||
exec_start = time.perf_counter()
|
||||
|
||||
if request.stream:
|
||||
# For streaming, collect all chunks
|
||||
chunks = []
|
||||
if hasattr(agent, '_start_stream'):
|
||||
for chunk in agent._start_stream(request.prompt):
|
||||
chunks.append(chunk)
|
||||
else:
|
||||
# Fallback to non-streaming
|
||||
result = agent.start(request.prompt)
|
||||
chunks = [result] if result else []
|
||||
output = "".join(chunks)
|
||||
else:
|
||||
output = agent.start(request.prompt)
|
||||
|
||||
exec_end = time.perf_counter()
|
||||
|
||||
if timing is not None:
|
||||
timing["execution_ms"] = (exec_end - exec_start) * 1000
|
||||
|
||||
end_time = time.perf_counter()
|
||||
|
||||
if timing is not None:
|
||||
timing["total_ms"] = (end_time - start_time) * 1000
|
||||
|
||||
return ExecutionResult(
|
||||
output=output or "",
|
||||
run_id=run_id,
|
||||
success=True,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
metadata={
|
||||
"model": request.model,
|
||||
"stream": request.stream,
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
end_time = time.perf_counter()
|
||||
|
||||
if timing is not None:
|
||||
timing["total_ms"] = (end_time - start_time) * 1000
|
||||
timing["error"] = str(e)
|
||||
|
||||
return ExecutionResult.from_error(
|
||||
error=str(e),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
|
||||
def execute_sync(request: ExecutionRequest) -> ExecutionResult:
|
||||
"""
|
||||
Synchronous execution wrapper.
|
||||
|
||||
Used by:
|
||||
- CLI direct prompt
|
||||
- Profile command
|
||||
|
||||
Args:
|
||||
request: The execution request
|
||||
|
||||
Returns:
|
||||
ExecutionResult with output and metadata
|
||||
"""
|
||||
return _execute_core(request)
|
||||
|
||||
|
||||
async def execute_async(request: ExecutionRequest) -> ExecutionResult:
|
||||
"""
|
||||
Asynchronous execution wrapper.
|
||||
|
||||
Used by:
|
||||
- TUI workers
|
||||
- Queue system
|
||||
|
||||
Args:
|
||||
request: The execution request
|
||||
|
||||
Returns:
|
||||
ExecutionResult with output and metadata
|
||||
"""
|
||||
# Run the synchronous core in a thread pool to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, _execute_core, request, None)
|
||||
|
||||
|
||||
async def execute_async_stream(
|
||||
request: ExecutionRequest,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Asynchronous streaming execution.
|
||||
|
||||
Yields chunks as they arrive.
|
||||
|
||||
Args:
|
||||
request: The execution request (stream=True recommended)
|
||||
|
||||
Yields:
|
||||
String chunks of the response
|
||||
"""
|
||||
try:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent_config = {
|
||||
"name": request.agent_name,
|
||||
"verbose": False,
|
||||
}
|
||||
|
||||
if request.agent_instructions:
|
||||
agent_config["instructions"] = request.agent_instructions
|
||||
|
||||
if request.model:
|
||||
agent_config["llm"] = request.model
|
||||
|
||||
if request.tools:
|
||||
agent_config["tools"] = list(request.tools)
|
||||
|
||||
agent = Agent(**agent_config)
|
||||
|
||||
if hasattr(agent, '_start_stream'):
|
||||
for chunk in agent._start_stream(request.prompt):
|
||||
yield chunk
|
||||
await asyncio.sleep(0) # Yield control
|
||||
else:
|
||||
result = agent.start(request.prompt)
|
||||
if result:
|
||||
yield result
|
||||
|
||||
except Exception as e:
|
||||
yield f"Error: {e}"
|
||||
@@ -0,0 +1,814 @@
|
||||
"""
|
||||
Profiler Module.
|
||||
|
||||
Provides external profiling that wraps execution.
|
||||
Execution code has ZERO awareness of profiling.
|
||||
|
||||
Design Principles:
|
||||
1. Profiler wraps execution, does not inject into it
|
||||
2. Layered profiling with explicit overhead
|
||||
3. Bounded resource usage
|
||||
4. Versioned, stable output schema
|
||||
"""
|
||||
|
||||
import cProfile
|
||||
import pstats
|
||||
import time
|
||||
import platform
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .request import ExecutionRequest
|
||||
from .result import ExecutionResult
|
||||
from .core import _execute_core
|
||||
|
||||
|
||||
# Schema version - MUST be incremented for breaking changes
|
||||
SCHEMA_VERSION = "1.0"
|
||||
|
||||
# Bounds for profile data
|
||||
MAX_FUNCTION_STATS = 1000
|
||||
MAX_CALL_GRAPH_EDGES = 5000
|
||||
MAX_NETWORK_REQUESTS = 100
|
||||
MAX_IMPORT_MODULES = 200
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProfilerConfig:
|
||||
"""
|
||||
Profiler configuration with explicit layer control.
|
||||
|
||||
Layers:
|
||||
- 0: Wall-clock phases only (<1ms overhead)
|
||||
- 1: + cProfile function stats (~5% overhead)
|
||||
- 2: + Call graph callers/callees (~15% overhead)
|
||||
|
||||
Network timing is orthogonal (opt-in via track_network).
|
||||
"""
|
||||
|
||||
# Layer control (0, 1, or 2)
|
||||
layer: int = 1
|
||||
|
||||
# Layer 1+ options
|
||||
limit: int = 30 # Max functions to show
|
||||
sort_by: str = "cumulative" # cumulative, time, calls
|
||||
|
||||
# Layer 2 options
|
||||
show_callers: bool = False
|
||||
show_callees: bool = False
|
||||
|
||||
# Network timing (orthogonal to layers)
|
||||
track_network: bool = False
|
||||
|
||||
# Output options
|
||||
output_format: str = "text" # text, json
|
||||
save_path: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def from_flags(
|
||||
cls,
|
||||
profile: bool = False,
|
||||
deep: bool = False,
|
||||
network: bool = False,
|
||||
) -> 'ProfilerConfig':
|
||||
"""Create config from CLI flags."""
|
||||
if deep:
|
||||
return cls(
|
||||
layer=2,
|
||||
show_callers=True,
|
||||
show_callees=True,
|
||||
track_network=network,
|
||||
)
|
||||
elif profile:
|
||||
return cls(layer=1, track_network=network)
|
||||
else:
|
||||
return cls(layer=0, track_network=network)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TimingBreakdown:
|
||||
"""
|
||||
Phase timing breakdown with full timeline visibility.
|
||||
|
||||
All fields in milliseconds.
|
||||
Timeline order: cli_entry -> cli_parse -> routing -> imports -> agent_init ->
|
||||
network_start -> first_token -> first_output -> execution_end
|
||||
"""
|
||||
|
||||
# MANDATORY fields (always present)
|
||||
total_ms: float = 0.0
|
||||
imports_ms: float = 0.0
|
||||
agent_init_ms: float = 0.0
|
||||
execution_ms: float = 0.0
|
||||
|
||||
# TIMELINE PHASES (new - for full visibility)
|
||||
cli_entry_ms: float = 0.0 # Time from ENTER to CLI entry point
|
||||
cli_parse_ms: float = 0.0 # Argparse/Typer routing time
|
||||
routing_ms: float = 0.0 # Command routing time
|
||||
network_start_ms: float = 0.0 # When network request started (relative to start)
|
||||
first_token_ms: float = 0.0 # Time to first token (streaming)
|
||||
first_output_ms: float = 0.0 # Time to First Response (TTFR) - first visible output
|
||||
|
||||
# OPTIONAL fields (may be 0)
|
||||
tool_time_ms: float = 0.0 # If tools were called
|
||||
|
||||
@property
|
||||
def time_to_first_response_ms(self) -> float:
|
||||
"""Time to First Response (TTFR) - the key user-perceived latency metric."""
|
||||
if self.first_output_ms > 0:
|
||||
return self.first_output_ms
|
||||
# Fallback: use execution end if first_output not captured
|
||||
return self.total_ms
|
||||
|
||||
def to_dict(self) -> Dict[str, float]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"total_ms": self.total_ms,
|
||||
"imports_ms": self.imports_ms,
|
||||
"agent_init_ms": self.agent_init_ms,
|
||||
"execution_ms": self.execution_ms,
|
||||
"cli_entry_ms": self.cli_entry_ms,
|
||||
"cli_parse_ms": self.cli_parse_ms,
|
||||
"routing_ms": self.routing_ms,
|
||||
"network_start_ms": self.network_start_ms,
|
||||
"first_token_ms": self.first_token_ms,
|
||||
"first_output_ms": self.first_output_ms,
|
||||
"time_to_first_response_ms": self.time_to_first_response_ms,
|
||||
"tool_time_ms": self.tool_time_ms,
|
||||
}
|
||||
|
||||
def to_timeline(self) -> List[Tuple[str, float, float]]:
|
||||
"""
|
||||
Get ordered timeline of phases.
|
||||
|
||||
Returns list of (phase_name, start_ms, duration_ms) tuples.
|
||||
"""
|
||||
timeline = []
|
||||
cursor = 0.0
|
||||
|
||||
# CLI Entry (if captured)
|
||||
if self.cli_entry_ms > 0:
|
||||
timeline.append(("CLI Entry", 0.0, self.cli_entry_ms))
|
||||
cursor = self.cli_entry_ms
|
||||
|
||||
# CLI Parse
|
||||
if self.cli_parse_ms > 0:
|
||||
timeline.append(("CLI Parse", cursor, self.cli_parse_ms))
|
||||
cursor += self.cli_parse_ms
|
||||
|
||||
# Routing
|
||||
if self.routing_ms > 0:
|
||||
timeline.append(("Routing", cursor, self.routing_ms))
|
||||
cursor += self.routing_ms
|
||||
|
||||
# Imports
|
||||
if self.imports_ms > 0:
|
||||
timeline.append(("Imports", cursor, self.imports_ms))
|
||||
cursor += self.imports_ms
|
||||
|
||||
# Agent Init
|
||||
if self.agent_init_ms > 0:
|
||||
timeline.append(("Agent Init", cursor, self.agent_init_ms))
|
||||
cursor += self.agent_init_ms
|
||||
|
||||
# Network to First Token (if streaming)
|
||||
if self.network_start_ms > 0 and self.first_token_ms > 0:
|
||||
network_duration = self.first_token_ms - self.network_start_ms
|
||||
if network_duration > 0:
|
||||
timeline.append(("Network → First Token", self.network_start_ms, network_duration))
|
||||
|
||||
# Execution
|
||||
if self.execution_ms > 0:
|
||||
timeline.append(("Execution", cursor, self.execution_ms))
|
||||
|
||||
return timeline
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionStat:
|
||||
"""Statistics for a single function."""
|
||||
|
||||
name: str
|
||||
file: str
|
||||
line: int
|
||||
calls: int
|
||||
total_time_ms: float
|
||||
cumulative_time_ms: float
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"file": self.file,
|
||||
"line": self.line,
|
||||
"calls": self.calls,
|
||||
"total_time_ms": self.total_time_ms,
|
||||
"cumulative_time_ms": self.cumulative_time_ms,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class CallGraph:
|
||||
"""Call graph data (callers and callees)."""
|
||||
|
||||
callers: Dict[str, List[str]] = field(default_factory=dict)
|
||||
callees: Dict[str, List[str]] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def edges(self) -> List[Tuple[str, str]]:
|
||||
"""Get all edges as (caller, callee) tuples."""
|
||||
edges = []
|
||||
for callee, callers in self.callers.items():
|
||||
for caller in callers:
|
||||
edges.append((caller, callee))
|
||||
return edges
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"callers": self.callers,
|
||||
"callees": self.callees,
|
||||
"edge_count": len(self.edges),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestTiming:
|
||||
"""Single HTTP request timing."""
|
||||
|
||||
url: str
|
||||
method: str
|
||||
start_time: float
|
||||
first_byte_time: float
|
||||
end_time: float
|
||||
status_code: int
|
||||
|
||||
@property
|
||||
def ttfb_ms(self) -> float:
|
||||
"""Time to first byte in milliseconds."""
|
||||
return (self.first_byte_time - self.start_time) * 1000
|
||||
|
||||
@property
|
||||
def total_ms(self) -> float:
|
||||
"""Total request time in milliseconds."""
|
||||
return (self.end_time - self.start_time) * 1000
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"url": self.url,
|
||||
"method": self.method,
|
||||
"ttfb_ms": self.ttfb_ms,
|
||||
"total_ms": self.total_ms,
|
||||
"status_code": self.status_code,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class InvocationInfo:
|
||||
"""How the profile was invoked."""
|
||||
|
||||
method: str # "cli_direct", "profile_command", "tui"
|
||||
flags: Dict[str, Any] = field(default_factory=dict)
|
||||
praisonai_version: str = ""
|
||||
python_version: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
"""Set version info if not provided."""
|
||||
if not self.python_version:
|
||||
self.python_version = platform.python_version()
|
||||
if not self.praisonai_version:
|
||||
try:
|
||||
from praisonai.version import __version__
|
||||
self.praisonai_version = __version__
|
||||
except ImportError:
|
||||
self.praisonai_version = "unknown"
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"method": self.method,
|
||||
"flags": self.flags,
|
||||
"praisonai_version": self.praisonai_version,
|
||||
"python_version": self.python_version,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class DecisionTrace:
|
||||
"""
|
||||
Decision trace for deep profile visibility.
|
||||
|
||||
Explains what decisions were made during execution.
|
||||
"""
|
||||
|
||||
agent_config: str = "" # Agent configuration chosen
|
||||
model_selected: str = "" # Model used
|
||||
streaming_mode: bool = False # Streaming vs non-streaming
|
||||
profile_layer: int = 0 # Profile layer (0/1/2)
|
||||
tools_enabled: List[str] = field(default_factory=list)
|
||||
tools_disabled: List[str] = field(default_factory=list)
|
||||
fallbacks_used: List[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"agent_config": self.agent_config,
|
||||
"model_selected": self.model_selected,
|
||||
"streaming_mode": self.streaming_mode,
|
||||
"profile_layer": self.profile_layer,
|
||||
"tools_enabled": self.tools_enabled,
|
||||
"tools_disabled": self.tools_disabled,
|
||||
"fallbacks_used": self.fallbacks_used,
|
||||
}
|
||||
|
||||
def to_text(self) -> str:
|
||||
"""Format as human-readable text."""
|
||||
lines = []
|
||||
lines.append(f" Agent Config: {self.agent_config or 'default'}")
|
||||
lines.append(f" Model: {self.model_selected or 'default'}")
|
||||
lines.append(f" Streaming: {'enabled' if self.streaming_mode else 'disabled'}")
|
||||
lines.append(f" Profile Layer: {self.profile_layer}")
|
||||
if self.tools_enabled:
|
||||
lines.append(f" Tools Enabled: {', '.join(self.tools_enabled)}")
|
||||
if self.tools_disabled:
|
||||
lines.append(f" Tools Disabled: {', '.join(self.tools_disabled)}")
|
||||
if self.fallbacks_used:
|
||||
lines.append(f" Fallbacks: {', '.join(self.fallbacks_used)}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModuleBreakdown:
|
||||
"""
|
||||
Module/file breakdown for deep profile visibility.
|
||||
|
||||
Groups functions by module category.
|
||||
"""
|
||||
|
||||
cli_modules: List[str] = field(default_factory=list)
|
||||
execution_modules: List[str] = field(default_factory=list)
|
||||
agent_modules: List[str] = field(default_factory=list)
|
||||
tool_modules: List[str] = field(default_factory=list)
|
||||
network_modules: List[str] = field(default_factory=list)
|
||||
third_party_modules: List[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, List[str]]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"cli": self.cli_modules,
|
||||
"execution": self.execution_modules,
|
||||
"agent": self.agent_modules,
|
||||
"tools": self.tool_modules,
|
||||
"network": self.network_modules,
|
||||
"third_party": self.third_party_modules,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProfileReport:
|
||||
"""
|
||||
Canonical profile report.
|
||||
|
||||
VERSION 1.0 - Schema is frozen after release.
|
||||
New fields MUST be optional with defaults.
|
||||
"""
|
||||
|
||||
# Schema version (MANDATORY, NEVER CHANGES MEANING)
|
||||
schema_version: str = SCHEMA_VERSION
|
||||
|
||||
# Identity (MANDATORY)
|
||||
run_id: str = ""
|
||||
timestamp: str = "" # ISO 8601
|
||||
|
||||
# Invocation (MANDATORY)
|
||||
invocation: InvocationInfo = field(default_factory=lambda: InvocationInfo(method="unknown"))
|
||||
|
||||
# Timing (MANDATORY)
|
||||
timing: TimingBreakdown = field(default_factory=TimingBreakdown)
|
||||
|
||||
# Function stats (OPTIONAL, present if layer >= 1)
|
||||
functions: Optional[List[FunctionStat]] = None
|
||||
|
||||
# Call graph (OPTIONAL, present if layer >= 2)
|
||||
call_graph: Optional[CallGraph] = None
|
||||
|
||||
# Network (OPTIONAL, present if track_network)
|
||||
network: Optional[List[RequestTiming]] = None
|
||||
|
||||
# Response (MANDATORY but truncated)
|
||||
response_preview: str = "" # Max 500 chars
|
||||
|
||||
# Decision trace (OPTIONAL, present if layer >= 2)
|
||||
decision_trace: Optional[DecisionTrace] = None
|
||||
|
||||
# Module breakdown (OPTIONAL, present if layer >= 2)
|
||||
module_breakdown: Optional[ModuleBreakdown] = None
|
||||
|
||||
def __post_init__(self):
|
||||
"""Set timestamp if not provided."""
|
||||
if not self.timestamp:
|
||||
self.timestamp = datetime.utcnow().isoformat() + "Z"
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
result = {
|
||||
"schema_version": self.schema_version,
|
||||
"run_id": self.run_id,
|
||||
"timestamp": self.timestamp,
|
||||
"invocation": self.invocation.to_dict(),
|
||||
"timing": self.timing.to_dict(),
|
||||
"response_preview": self.response_preview,
|
||||
}
|
||||
|
||||
if self.functions is not None:
|
||||
result["functions"] = [f.to_dict() for f in self.functions]
|
||||
|
||||
if self.call_graph is not None:
|
||||
result["call_graph"] = self.call_graph.to_dict()
|
||||
|
||||
if self.network is not None:
|
||||
result["network"] = [r.to_dict() for r in self.network]
|
||||
|
||||
if self.decision_trace is not None:
|
||||
result["decision_trace"] = self.decision_trace.to_dict()
|
||||
|
||||
if self.module_breakdown is not None:
|
||||
result["module_breakdown"] = self.module_breakdown.to_dict()
|
||||
|
||||
return result
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Convert to JSON string."""
|
||||
import json
|
||||
return json.dumps(self.to_dict(), indent=2, default=str)
|
||||
|
||||
def to_text(self) -> str:
|
||||
"""Format as human-readable text."""
|
||||
lines = []
|
||||
lines.append("=" * 70)
|
||||
lines.append("PraisonAI Profile Report")
|
||||
lines.append("=" * 70)
|
||||
lines.append("")
|
||||
|
||||
# Identity
|
||||
lines.append(f"Run ID: {self.run_id}")
|
||||
lines.append(f"Timestamp: {self.timestamp}")
|
||||
lines.append(f"Method: {self.invocation.method}")
|
||||
lines.append(f"Version: {self.invocation.praisonai_version}")
|
||||
lines.append("")
|
||||
|
||||
# TIMELINE SECTION (new - shows ENTER → First Response)
|
||||
lines.append("## Execution Timeline")
|
||||
lines.append("-" * 45)
|
||||
timeline = self.timing.to_timeline()
|
||||
if timeline:
|
||||
for phase_name, start_ms, duration_ms in timeline:
|
||||
lines.append(f" {phase_name:<25} : {duration_ms:>8.2f} ms")
|
||||
else:
|
||||
# Fallback to basic phases
|
||||
if self.timing.imports_ms > 0:
|
||||
lines.append(f" {'Imports':<25} : {self.timing.imports_ms:>8.2f} ms")
|
||||
if self.timing.agent_init_ms > 0:
|
||||
lines.append(f" {'Agent Init':<25} : {self.timing.agent_init_ms:>8.2f} ms")
|
||||
if self.timing.execution_ms > 0:
|
||||
lines.append(f" {'Execution':<25} : {self.timing.execution_ms:>8.2f} ms")
|
||||
lines.append(" " + "─" * 43)
|
||||
|
||||
# TIME TO FIRST RESPONSE (TTFR) - the key metric
|
||||
ttfr = self.timing.time_to_first_response_ms
|
||||
lines.append(f" {'⏱ Time to First Response':<25} : {ttfr:>8.2f} ms")
|
||||
lines.append(f" {'TOTAL':<25} : {self.timing.total_ms:>8.2f} ms")
|
||||
lines.append("")
|
||||
|
||||
# Timing breakdown (legacy format for compatibility)
|
||||
lines.append("## Timing Breakdown")
|
||||
lines.append("-" * 40)
|
||||
if self.timing.cli_entry_ms > 0:
|
||||
lines.append(f" CLI Entry: {self.timing.cli_entry_ms:>10.2f} ms")
|
||||
if self.timing.cli_parse_ms > 0:
|
||||
lines.append(f" CLI Parse: {self.timing.cli_parse_ms:>10.2f} ms")
|
||||
if self.timing.routing_ms > 0:
|
||||
lines.append(f" Routing: {self.timing.routing_ms:>10.2f} ms")
|
||||
lines.append(f" Imports: {self.timing.imports_ms:>10.2f} ms")
|
||||
lines.append(f" Agent Init: {self.timing.agent_init_ms:>10.2f} ms")
|
||||
lines.append(f" Execution: {self.timing.execution_ms:>10.2f} ms")
|
||||
if self.timing.network_start_ms > 0:
|
||||
lines.append(f" Network Start: {self.timing.network_start_ms:>10.2f} ms")
|
||||
if self.timing.first_token_ms > 0:
|
||||
lines.append(f" First Token: {self.timing.first_token_ms:>10.2f} ms")
|
||||
if self.timing.first_output_ms > 0:
|
||||
lines.append(f" First Output: {self.timing.first_output_ms:>10.2f} ms")
|
||||
lines.append(" ─────────────────────────────────────")
|
||||
lines.append(f" TOTAL: {self.timing.total_ms:>10.2f} ms")
|
||||
lines.append("")
|
||||
|
||||
# Decision trace (deep profile only)
|
||||
if self.decision_trace:
|
||||
lines.append("## Decision Trace")
|
||||
lines.append("-" * 40)
|
||||
lines.append(self.decision_trace.to_text())
|
||||
lines.append("")
|
||||
|
||||
# Function stats
|
||||
if self.functions:
|
||||
lines.append("## Top Functions by Cumulative Time")
|
||||
lines.append("-" * 70)
|
||||
lines.append(f"{'Function':<40} {'Calls':>8} {'Cumul (ms)':>12}")
|
||||
lines.append("-" * 70)
|
||||
for func in self.functions[:20]:
|
||||
name = func.name[:38] if len(func.name) > 38 else func.name
|
||||
lines.append(f"{name:<40} {func.calls:>8} {func.cumulative_time_ms:>12.2f}")
|
||||
lines.append("")
|
||||
|
||||
# Module breakdown (deep profile only)
|
||||
if self.module_breakdown:
|
||||
lines.append("## Module Breakdown")
|
||||
lines.append("-" * 40)
|
||||
mb = self.module_breakdown
|
||||
if mb.cli_modules:
|
||||
lines.append(f" CLI: {len(mb.cli_modules)} modules")
|
||||
if mb.execution_modules:
|
||||
lines.append(f" Execution: {len(mb.execution_modules)} modules")
|
||||
if mb.agent_modules:
|
||||
lines.append(f" Agent: {len(mb.agent_modules)} modules")
|
||||
if mb.tool_modules:
|
||||
lines.append(f" Tools: {len(mb.tool_modules)} modules")
|
||||
if mb.network_modules:
|
||||
lines.append(f" Network: {len(mb.network_modules)} modules")
|
||||
if mb.third_party_modules:
|
||||
lines.append(f" Third-party:{len(mb.third_party_modules)} modules")
|
||||
lines.append("")
|
||||
|
||||
# Call graph summary
|
||||
if self.call_graph:
|
||||
lines.append(f"## Call Graph: {len(self.call_graph.edges)} edges")
|
||||
lines.append("")
|
||||
|
||||
# Network timing
|
||||
if self.network:
|
||||
lines.append("## Network Requests")
|
||||
lines.append("-" * 70)
|
||||
for req in self.network[:10]:
|
||||
lines.append(f" {req.method} {req.url[:50]} - {req.total_ms:.2f}ms (TTFB: {req.ttfb_ms:.2f}ms)")
|
||||
lines.append("")
|
||||
|
||||
# Response preview
|
||||
if self.response_preview:
|
||||
lines.append("## Response Preview")
|
||||
lines.append("-" * 40)
|
||||
preview = self.response_preview[:200]
|
||||
if len(self.response_preview) > 200:
|
||||
preview += "..."
|
||||
lines.append(preview)
|
||||
lines.append("")
|
||||
|
||||
lines.append("=" * 70)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
class Profiler:
|
||||
"""
|
||||
External profiler that wraps execution.
|
||||
|
||||
Execution code has ZERO awareness of profiling.
|
||||
Profiler uses cProfile.runctx() or similar.
|
||||
"""
|
||||
|
||||
def __init__(self, config: ProfilerConfig):
|
||||
self.config = config
|
||||
self._cprofile: Optional[cProfile.Profile] = None
|
||||
self._timing: Dict[str, float] = {}
|
||||
self._network_requests: List[RequestTiming] = []
|
||||
|
||||
def profile_sync(
|
||||
self,
|
||||
request: ExecutionRequest,
|
||||
invocation_method: str = "profile_command",
|
||||
) -> Tuple[ExecutionResult, ProfileReport]:
|
||||
"""
|
||||
Profile synchronous execution.
|
||||
|
||||
Args:
|
||||
request: The execution request
|
||||
invocation_method: How this was invoked (for report)
|
||||
|
||||
Returns:
|
||||
Tuple of (ExecutionResult, ProfileReport)
|
||||
"""
|
||||
timing_dict: Dict[str, float] = {}
|
||||
|
||||
# Layer 0: Wall-clock timing (always)
|
||||
start = time.perf_counter()
|
||||
|
||||
# Layer 1+: cProfile
|
||||
if self.config.layer >= 1:
|
||||
self._cprofile = cProfile.Profile()
|
||||
self._cprofile.enable()
|
||||
|
||||
try:
|
||||
result = _execute_core(request, timing_dict)
|
||||
finally:
|
||||
if self._cprofile:
|
||||
self._cprofile.disable()
|
||||
|
||||
total_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
# Build report
|
||||
report = self._build_report(
|
||||
request=request,
|
||||
result=result,
|
||||
timing_dict=timing_dict,
|
||||
total_ms=total_ms,
|
||||
invocation_method=invocation_method,
|
||||
)
|
||||
|
||||
return result, report
|
||||
|
||||
def _build_report(
|
||||
self,
|
||||
request: ExecutionRequest,
|
||||
result: ExecutionResult,
|
||||
timing_dict: Dict[str, float],
|
||||
total_ms: float,
|
||||
invocation_method: str,
|
||||
) -> ProfileReport:
|
||||
"""Build the profile report from collected data."""
|
||||
|
||||
# Calculate first_output_ms (TTFR) - time from start to when output is available
|
||||
# This is imports + agent_init + execution (when output becomes available)
|
||||
imports_ms = timing_dict.get("imports_ms", 0.0)
|
||||
agent_init_ms = timing_dict.get("agent_init_ms", 0.0)
|
||||
execution_ms = timing_dict.get("execution_ms", 0.0)
|
||||
first_output_ms = imports_ms + agent_init_ms + execution_ms
|
||||
|
||||
# Timing breakdown with full timeline
|
||||
timing = TimingBreakdown(
|
||||
total_ms=total_ms,
|
||||
imports_ms=imports_ms,
|
||||
agent_init_ms=agent_init_ms,
|
||||
execution_ms=execution_ms,
|
||||
cli_entry_ms=timing_dict.get("cli_entry_ms", 0.0),
|
||||
cli_parse_ms=timing_dict.get("cli_parse_ms", 0.0),
|
||||
routing_ms=timing_dict.get("routing_ms", 0.0),
|
||||
network_start_ms=timing_dict.get("network_start_ms", 0.0),
|
||||
first_token_ms=timing_dict.get("first_token_ms", 0.0),
|
||||
first_output_ms=first_output_ms,
|
||||
)
|
||||
|
||||
# Invocation info
|
||||
invocation = InvocationInfo(
|
||||
method=invocation_method,
|
||||
flags={
|
||||
"layer": self.config.layer,
|
||||
"stream": request.stream,
|
||||
"model": request.model,
|
||||
},
|
||||
)
|
||||
|
||||
# Function stats (layer 1+)
|
||||
functions = None
|
||||
if self.config.layer >= 1 and self._cprofile:
|
||||
functions = self._extract_function_stats()
|
||||
|
||||
# Call graph (layer 2)
|
||||
call_graph = None
|
||||
if self.config.layer >= 2 and self._cprofile:
|
||||
call_graph = self._extract_call_graph()
|
||||
|
||||
# Decision trace (layer 2 - deep profile)
|
||||
decision_trace = None
|
||||
if self.config.layer >= 2:
|
||||
decision_trace = DecisionTrace(
|
||||
agent_config=request.agent_name or "default",
|
||||
model_selected=request.model or "default",
|
||||
streaming_mode=request.stream,
|
||||
profile_layer=self.config.layer,
|
||||
tools_enabled=list(request.tools) if request.tools else [],
|
||||
tools_disabled=[],
|
||||
fallbacks_used=[],
|
||||
)
|
||||
|
||||
# Module breakdown (layer 2 - deep profile)
|
||||
module_breakdown = None
|
||||
if self.config.layer >= 2 and functions:
|
||||
module_breakdown = self._extract_module_breakdown(functions)
|
||||
|
||||
# Response preview (truncated)
|
||||
response_preview = result.output[:500] if result.output else ""
|
||||
|
||||
return ProfileReport(
|
||||
run_id=result.run_id,
|
||||
invocation=invocation,
|
||||
timing=timing,
|
||||
functions=functions,
|
||||
call_graph=call_graph,
|
||||
network=self._network_requests if self._network_requests else None,
|
||||
response_preview=response_preview,
|
||||
decision_trace=decision_trace,
|
||||
module_breakdown=module_breakdown,
|
||||
)
|
||||
|
||||
def _extract_module_breakdown(self, functions: List[FunctionStat]) -> ModuleBreakdown:
|
||||
"""Extract module breakdown from function stats."""
|
||||
cli_modules = set()
|
||||
execution_modules = set()
|
||||
agent_modules = set()
|
||||
tool_modules = set()
|
||||
network_modules = set()
|
||||
third_party_modules = set()
|
||||
|
||||
for func in functions:
|
||||
file_path = func.file.lower()
|
||||
|
||||
if "praisonai/cli" in file_path:
|
||||
cli_modules.add(func.file)
|
||||
elif "praisonai" in file_path and "execution" in file_path:
|
||||
execution_modules.add(func.file)
|
||||
elif "praisonaiagents" in file_path and "agent" in file_path:
|
||||
agent_modules.add(func.file)
|
||||
elif "praisonaiagents" in file_path and "tool" in file_path:
|
||||
tool_modules.add(func.file)
|
||||
elif any(x in file_path for x in ["httpx", "httpcore", "urllib", "requests", "aiohttp"]):
|
||||
network_modules.add(func.file)
|
||||
elif "site-packages" in file_path or not file_path.startswith("/"):
|
||||
third_party_modules.add(func.file)
|
||||
|
||||
return ModuleBreakdown(
|
||||
cli_modules=list(cli_modules)[:20],
|
||||
execution_modules=list(execution_modules)[:20],
|
||||
agent_modules=list(agent_modules)[:20],
|
||||
tool_modules=list(tool_modules)[:20],
|
||||
network_modules=list(network_modules)[:20],
|
||||
third_party_modules=list(third_party_modules)[:20],
|
||||
)
|
||||
|
||||
def _extract_function_stats(self) -> List[FunctionStat]:
|
||||
"""Extract function statistics from cProfile."""
|
||||
if not self._cprofile:
|
||||
return []
|
||||
|
||||
stats = pstats.Stats(self._cprofile)
|
||||
|
||||
# Sort by cumulative time
|
||||
stats.sort_stats(self.config.sort_by)
|
||||
|
||||
# Extract stats
|
||||
function_stats = []
|
||||
for (filename, line, name), (cc, nc, tt, ct, callers) in stats.stats.items():
|
||||
if len(function_stats) >= MAX_FUNCTION_STATS:
|
||||
break
|
||||
|
||||
function_stats.append(FunctionStat(
|
||||
name=name,
|
||||
file=filename,
|
||||
line=line,
|
||||
calls=nc,
|
||||
total_time_ms=tt * 1000,
|
||||
cumulative_time_ms=ct * 1000,
|
||||
))
|
||||
|
||||
# Sort and limit
|
||||
function_stats.sort(key=lambda x: x.cumulative_time_ms, reverse=True)
|
||||
return function_stats[:self.config.limit]
|
||||
|
||||
def _extract_call_graph(self) -> CallGraph:
|
||||
"""Extract call graph from cProfile."""
|
||||
if not self._cprofile:
|
||||
return CallGraph()
|
||||
|
||||
stats = pstats.Stats(self._cprofile)
|
||||
|
||||
callers: Dict[str, List[str]] = {}
|
||||
callees: Dict[str, List[str]] = {}
|
||||
edge_count = 0
|
||||
|
||||
for (filename, line, name), (cc, nc, tt, ct, caller_dict) in stats.stats.items():
|
||||
if edge_count >= MAX_CALL_GRAPH_EDGES:
|
||||
break
|
||||
|
||||
callee_key = f"{name}:{filename}:{line}"
|
||||
|
||||
for (caller_file, caller_line, caller_name), _ in caller_dict.items():
|
||||
if edge_count >= MAX_CALL_GRAPH_EDGES:
|
||||
break
|
||||
|
||||
caller_key = f"{caller_name}:{caller_file}:{caller_line}"
|
||||
|
||||
# Add to callers
|
||||
if callee_key not in callers:
|
||||
callers[callee_key] = []
|
||||
if caller_key not in callers[callee_key]:
|
||||
callers[callee_key].append(caller_key)
|
||||
|
||||
# Add to callees
|
||||
if caller_key not in callees:
|
||||
callees[caller_key] = []
|
||||
if callee_key not in callees[caller_key]:
|
||||
callees[caller_key].append(callee_key)
|
||||
|
||||
edge_count += 1
|
||||
|
||||
return CallGraph(callers=callers, callees=callees)
|
||||
@@ -0,0 +1,85 @@
|
||||
"""
|
||||
ExecutionRequest - Immutable request for execution.
|
||||
|
||||
This is the minimal, immutable request that all execution paths use.
|
||||
CLI/TUI translate their flags into this structure.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Optional, Tuple, Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExecutionRequest:
|
||||
"""
|
||||
Minimal, immutable request for execution.
|
||||
|
||||
NOT a kitchen-sink of all CLI flags.
|
||||
CLI/TUI translate their flags into this.
|
||||
|
||||
Attributes:
|
||||
prompt: The user prompt to execute
|
||||
agent_name: Name of the agent (default: "Assistant")
|
||||
agent_instructions: Custom instructions for the agent
|
||||
model: LLM model to use (None = use default)
|
||||
tools: Tuple of tool callables (immutable)
|
||||
stream: Whether to stream the response
|
||||
|
||||
Note: This is frozen (immutable) to:
|
||||
- Prevent accidental mutation during profiling
|
||||
- Enable safe sharing across threads/coroutines
|
||||
"""
|
||||
|
||||
# Required
|
||||
prompt: str
|
||||
|
||||
# Agent identity (optional, has defaults)
|
||||
agent_name: str = "Assistant"
|
||||
agent_instructions: Optional[str] = None
|
||||
|
||||
# Model (optional, has defaults)
|
||||
model: Optional[str] = None # None = use default
|
||||
|
||||
# Tools (optional) - using tuple for immutability
|
||||
tools: Tuple[Callable[..., Any], ...] = ()
|
||||
|
||||
# Execution mode
|
||||
stream: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
"""Validate request after initialization."""
|
||||
if not self.prompt:
|
||||
raise ValueError("prompt cannot be empty")
|
||||
|
||||
def with_model(self, model: str) -> 'ExecutionRequest':
|
||||
"""Create a new request with a different model."""
|
||||
return ExecutionRequest(
|
||||
prompt=self.prompt,
|
||||
agent_name=self.agent_name,
|
||||
agent_instructions=self.agent_instructions,
|
||||
model=model,
|
||||
tools=self.tools,
|
||||
stream=self.stream,
|
||||
)
|
||||
|
||||
def with_tools(self, tools: Tuple[Callable[..., Any], ...]) -> 'ExecutionRequest':
|
||||
"""Create a new request with different tools."""
|
||||
return ExecutionRequest(
|
||||
prompt=self.prompt,
|
||||
agent_name=self.agent_name,
|
||||
agent_instructions=self.agent_instructions,
|
||||
model=self.model,
|
||||
tools=tools,
|
||||
stream=self.stream,
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"prompt": self.prompt,
|
||||
"agent_name": self.agent_name,
|
||||
"agent_instructions": self.agent_instructions,
|
||||
"model": self.model,
|
||||
"tools_count": len(self.tools),
|
||||
"stream": self.stream,
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
ExecutionResult - Result from execution.
|
||||
|
||||
Contains the output and metadata from an execution.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutionResult:
|
||||
"""
|
||||
Result from executing a request.
|
||||
|
||||
Attributes:
|
||||
output: The text output from the agent
|
||||
run_id: Unique identifier for this execution
|
||||
success: Whether execution completed successfully
|
||||
error: Error message if execution failed
|
||||
metadata: Additional metadata about the execution
|
||||
tool_calls: List of tools that were called
|
||||
start_time: When execution started (perf_counter)
|
||||
end_time: When execution ended (perf_counter)
|
||||
"""
|
||||
|
||||
# Primary output
|
||||
output: str
|
||||
|
||||
# Identity
|
||||
run_id: str = ""
|
||||
|
||||
# Status
|
||||
success: bool = True
|
||||
error: Optional[str] = None
|
||||
|
||||
# Metadata
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Tool tracking
|
||||
tool_calls: List[str] = field(default_factory=list)
|
||||
|
||||
# Timing (raw perf_counter values)
|
||||
start_time: float = 0.0
|
||||
end_time: float = 0.0
|
||||
|
||||
@property
|
||||
def duration_ms(self) -> float:
|
||||
"""Duration in milliseconds."""
|
||||
if self.start_time and self.end_time:
|
||||
return (self.end_time - self.start_time) * 1000
|
||||
return 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"output": self.output,
|
||||
"run_id": self.run_id,
|
||||
"success": self.success,
|
||||
"error": self.error,
|
||||
"metadata": self.metadata,
|
||||
"tool_calls": self.tool_calls,
|
||||
"duration_ms": self.duration_ms,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_error(cls, error: str, run_id: str = "") -> 'ExecutionResult':
|
||||
"""Create a result from an error."""
|
||||
return cls(
|
||||
output="",
|
||||
run_id=run_id,
|
||||
success=False,
|
||||
error=error,
|
||||
)
|
||||
@@ -102,6 +102,11 @@ __all__ = [
|
||||
'RuntimeConfig',
|
||||
'CodeIntelligenceRouter',
|
||||
'ActionOrchestrator',
|
||||
# Interactive tools provider (canonical source)
|
||||
'get_interactive_tools',
|
||||
'resolve_tool_groups',
|
||||
'ToolConfig',
|
||||
'TOOL_GROUPS',
|
||||
]
|
||||
|
||||
def __getattr__(name):
|
||||
@@ -247,4 +252,17 @@ def __getattr__(name):
|
||||
elif name == 'ActionOrchestrator':
|
||||
from .action_orchestrator import ActionOrchestrator
|
||||
return ActionOrchestrator
|
||||
# Interactive tools provider (canonical source)
|
||||
elif name == 'get_interactive_tools':
|
||||
from .interactive_tools import get_interactive_tools
|
||||
return get_interactive_tools
|
||||
elif name == 'resolve_tool_groups':
|
||||
from .interactive_tools import resolve_tool_groups
|
||||
return resolve_tool_groups
|
||||
elif name == 'ToolConfig':
|
||||
from .interactive_tools import ToolConfig
|
||||
return ToolConfig
|
||||
elif name == 'TOOL_GROUPS':
|
||||
from .interactive_tools import TOOL_GROUPS
|
||||
return TOOL_GROUPS
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,380 @@
|
||||
"""
|
||||
Interactive Default Tools Provider for PraisonAI.
|
||||
|
||||
This module provides the canonical source of truth for default tools
|
||||
in interactive modes (TUI and `praison "prompt"`).
|
||||
|
||||
Tool Groups:
|
||||
- `acp`: ACP-powered file operations (create, edit, delete, execute)
|
||||
- `lsp`: LSP-powered code intelligence (symbols, definitions, references)
|
||||
- `basic`: Basic file tools (read, list, search)
|
||||
- `interactive`: Union of all groups (default)
|
||||
|
||||
Usage:
|
||||
from praisonai.cli.features.interactive_tools import (
|
||||
get_interactive_tools,
|
||||
resolve_tool_groups,
|
||||
TOOL_GROUPS,
|
||||
)
|
||||
|
||||
# Get all interactive tools (default)
|
||||
tools = get_interactive_tools()
|
||||
|
||||
# Get specific groups
|
||||
tools = get_interactive_tools(groups=["acp", "basic"])
|
||||
|
||||
# Disable specific groups
|
||||
tools = get_interactive_tools(disable=["lsp"])
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, Dict, List, Optional, Set
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Tool group definitions
|
||||
TOOL_GROUPS = {
|
||||
"acp": [
|
||||
"acp_create_file",
|
||||
"acp_edit_file",
|
||||
"acp_delete_file",
|
||||
"acp_execute_command",
|
||||
],
|
||||
"lsp": [
|
||||
"lsp_list_symbols",
|
||||
"lsp_find_definition",
|
||||
"lsp_find_references",
|
||||
"lsp_get_diagnostics",
|
||||
],
|
||||
"basic": [
|
||||
"read_file",
|
||||
"write_file",
|
||||
"list_files",
|
||||
"execute_command",
|
||||
"internet_search",
|
||||
],
|
||||
}
|
||||
|
||||
# Default interactive group includes all
|
||||
TOOL_GROUPS["interactive"] = (
|
||||
TOOL_GROUPS["acp"] +
|
||||
TOOL_GROUPS["lsp"] +
|
||||
TOOL_GROUPS["basic"]
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolConfig:
|
||||
"""Configuration for tool loading."""
|
||||
workspace: str = field(default_factory=os.getcwd)
|
||||
enable_acp: bool = True
|
||||
enable_lsp: bool = True
|
||||
enable_basic: bool = True
|
||||
approval_mode: str = "auto" # auto, manual, scoped
|
||||
lsp_enabled: bool = True
|
||||
acp_enabled: bool = True
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "ToolConfig":
|
||||
"""Create config from environment variables."""
|
||||
config = cls()
|
||||
|
||||
# Check env vars for disabling groups
|
||||
disable_str = os.environ.get("PRAISON_TOOLS_DISABLE", "")
|
||||
if disable_str:
|
||||
disabled = [g.strip().lower() for g in disable_str.split(",")]
|
||||
if "acp" in disabled:
|
||||
config.enable_acp = False
|
||||
if "lsp" in disabled:
|
||||
config.enable_lsp = False
|
||||
if "basic" in disabled:
|
||||
config.enable_basic = False
|
||||
|
||||
# Check for workspace
|
||||
workspace = os.environ.get("PRAISON_WORKSPACE", "")
|
||||
if workspace:
|
||||
config.workspace = workspace
|
||||
|
||||
# Check approval mode
|
||||
approval = os.environ.get("PRAISON_APPROVAL_MODE", "")
|
||||
if approval in ("auto", "manual", "scoped"):
|
||||
config.approval_mode = approval
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def resolve_tool_groups(
|
||||
groups: Optional[List[str]] = None,
|
||||
disable: Optional[List[str]] = None,
|
||||
config: Optional[ToolConfig] = None,
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Resolve which tool names should be included based on groups and config.
|
||||
|
||||
Args:
|
||||
groups: Specific groups to include (None = all enabled)
|
||||
disable: Groups to explicitly disable
|
||||
config: Tool configuration
|
||||
|
||||
Returns:
|
||||
Set of tool names to include
|
||||
"""
|
||||
if config is None:
|
||||
config = ToolConfig.from_env()
|
||||
|
||||
disable = disable or []
|
||||
|
||||
# Start with requested groups or all enabled
|
||||
if groups:
|
||||
tool_names = set()
|
||||
for group in groups:
|
||||
if group in TOOL_GROUPS:
|
||||
tool_names.update(TOOL_GROUPS[group])
|
||||
else:
|
||||
tool_names = set()
|
||||
if config.enable_acp:
|
||||
tool_names.update(TOOL_GROUPS["acp"])
|
||||
if config.enable_lsp:
|
||||
tool_names.update(TOOL_GROUPS["lsp"])
|
||||
if config.enable_basic:
|
||||
tool_names.update(TOOL_GROUPS["basic"])
|
||||
|
||||
# Remove disabled groups
|
||||
for group in disable:
|
||||
if group in TOOL_GROUPS:
|
||||
tool_names -= set(TOOL_GROUPS[group])
|
||||
|
||||
return tool_names
|
||||
|
||||
|
||||
def _load_basic_tools() -> Dict[str, Callable]:
|
||||
"""Lazy load basic tools from praisonaiagents."""
|
||||
tools = {}
|
||||
|
||||
try:
|
||||
from praisonaiagents.tools import read_file
|
||||
tools["read_file"] = read_file
|
||||
except ImportError:
|
||||
logger.debug("read_file not available")
|
||||
|
||||
try:
|
||||
from praisonaiagents.tools import write_file
|
||||
tools["write_file"] = write_file
|
||||
except ImportError:
|
||||
logger.debug("write_file not available")
|
||||
|
||||
try:
|
||||
from praisonaiagents.tools import list_files
|
||||
tools["list_files"] = list_files
|
||||
except ImportError:
|
||||
logger.debug("list_files not available")
|
||||
|
||||
try:
|
||||
from praisonaiagents.tools import execute_command
|
||||
tools["execute_command"] = execute_command
|
||||
except ImportError:
|
||||
logger.debug("execute_command not available")
|
||||
|
||||
try:
|
||||
from praisonaiagents.tools import internet_search
|
||||
tools["internet_search"] = internet_search
|
||||
except ImportError:
|
||||
logger.debug("internet_search not available")
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def _load_acp_tools(config: ToolConfig) -> Dict[str, Callable]:
|
||||
"""Lazy load ACP tools with runtime initialization."""
|
||||
tools = {}
|
||||
|
||||
try:
|
||||
from .interactive_runtime import InteractiveRuntime, RuntimeConfig
|
||||
from .agent_tools import create_agent_centric_tools
|
||||
|
||||
# Create runtime config
|
||||
runtime_config = RuntimeConfig(
|
||||
workspace=config.workspace,
|
||||
lsp_enabled=config.lsp_enabled,
|
||||
acp_enabled=config.acp_enabled,
|
||||
approval_mode=config.approval_mode,
|
||||
)
|
||||
|
||||
# Create runtime (lazy - doesn't start until needed)
|
||||
runtime = InteractiveRuntime(runtime_config)
|
||||
|
||||
# Get all agent-centric tools
|
||||
all_tools = create_agent_centric_tools(runtime)
|
||||
|
||||
# Map by name
|
||||
for tool in all_tools:
|
||||
name = tool.__name__
|
||||
if name.startswith("acp_"):
|
||||
tools[name] = tool
|
||||
|
||||
logger.debug(f"Loaded {len(tools)} ACP tools")
|
||||
|
||||
except ImportError as e:
|
||||
logger.debug(f"ACP tools not available: {e}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error loading ACP tools: {e}")
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def _load_lsp_tools(config: ToolConfig) -> Dict[str, Callable]:
|
||||
"""Lazy load LSP tools with runtime initialization."""
|
||||
tools = {}
|
||||
|
||||
try:
|
||||
from .interactive_runtime import InteractiveRuntime, RuntimeConfig
|
||||
from .agent_tools import create_agent_centric_tools
|
||||
|
||||
# Create runtime config
|
||||
runtime_config = RuntimeConfig(
|
||||
workspace=config.workspace,
|
||||
lsp_enabled=config.lsp_enabled,
|
||||
acp_enabled=config.acp_enabled,
|
||||
approval_mode=config.approval_mode,
|
||||
)
|
||||
|
||||
# Create runtime
|
||||
runtime = InteractiveRuntime(runtime_config)
|
||||
|
||||
# Get all agent-centric tools
|
||||
all_tools = create_agent_centric_tools(runtime)
|
||||
|
||||
# Map by name
|
||||
for tool in all_tools:
|
||||
name = tool.__name__
|
||||
if name.startswith("lsp_"):
|
||||
tools[name] = tool
|
||||
|
||||
logger.debug(f"Loaded {len(tools)} LSP tools")
|
||||
|
||||
except ImportError as e:
|
||||
logger.debug(f"LSP tools not available: {e}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error loading LSP tools: {e}")
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def get_interactive_tools(
|
||||
groups: Optional[List[str]] = None,
|
||||
disable: Optional[List[str]] = None,
|
||||
config: Optional[ToolConfig] = None,
|
||||
workspace: Optional[str] = None,
|
||||
) -> List[Callable]:
|
||||
"""
|
||||
Get the default interactive tools for TUI and prompt modes.
|
||||
|
||||
This is the canonical source of truth for interactive tool defaults.
|
||||
Both `praisonai tui launch` and `praison "prompt"` should use this.
|
||||
|
||||
Args:
|
||||
groups: Specific tool groups to include (default: all enabled)
|
||||
disable: Tool groups to disable (e.g., ["acp", "lsp"])
|
||||
config: Tool configuration (created from env if not provided)
|
||||
workspace: Override workspace path
|
||||
|
||||
Returns:
|
||||
List of tool callables
|
||||
|
||||
Example:
|
||||
# Default (all tools)
|
||||
tools = get_interactive_tools()
|
||||
|
||||
# Disable ACP tools
|
||||
tools = get_interactive_tools(disable=["acp"])
|
||||
|
||||
# Only basic tools
|
||||
tools = get_interactive_tools(groups=["basic"])
|
||||
"""
|
||||
if config is None:
|
||||
config = ToolConfig.from_env()
|
||||
|
||||
if workspace:
|
||||
config.workspace = workspace
|
||||
|
||||
# Resolve which tools to include
|
||||
tool_names = resolve_tool_groups(groups, disable, config)
|
||||
|
||||
# Load tools by group
|
||||
all_tools: Dict[str, Callable] = {}
|
||||
|
||||
# Load basic tools (always available, no runtime needed)
|
||||
if config.enable_basic and not (disable and "basic" in disable):
|
||||
basic_tools = _load_basic_tools()
|
||||
all_tools.update(basic_tools)
|
||||
|
||||
# Load ACP tools (requires runtime)
|
||||
if config.enable_acp and not (disable and "acp" in disable):
|
||||
acp_tools = _load_acp_tools(config)
|
||||
all_tools.update(acp_tools)
|
||||
|
||||
# Load LSP tools (requires runtime)
|
||||
if config.enable_lsp and not (disable and "lsp" in disable):
|
||||
lsp_tools = _load_lsp_tools(config)
|
||||
all_tools.update(lsp_tools)
|
||||
|
||||
# Filter to requested tools
|
||||
result = []
|
||||
for name in sorted(tool_names): # Deterministic ordering
|
||||
if name in all_tools:
|
||||
result.append(all_tools[name])
|
||||
|
||||
logger.debug(f"Loaded {len(result)} interactive tools: {[t.__name__ for t in result]}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def merge_tool_overrides(
|
||||
default_tools: List[Callable],
|
||||
add_tools: Optional[List[Callable]] = None,
|
||||
remove_names: Optional[List[str]] = None,
|
||||
) -> List[Callable]:
|
||||
"""
|
||||
Merge tool overrides with default tools.
|
||||
|
||||
Args:
|
||||
default_tools: Base list of tools
|
||||
add_tools: Tools to add
|
||||
remove_names: Tool names to remove
|
||||
|
||||
Returns:
|
||||
Merged tool list
|
||||
"""
|
||||
# Build name -> tool map
|
||||
tool_map = {t.__name__: t for t in default_tools}
|
||||
|
||||
# Remove specified tools
|
||||
if remove_names:
|
||||
for name in remove_names:
|
||||
tool_map.pop(name, None)
|
||||
|
||||
# Add new tools
|
||||
if add_tools:
|
||||
for tool in add_tools:
|
||||
tool_map[tool.__name__] = tool
|
||||
|
||||
return list(tool_map.values())
|
||||
|
||||
|
||||
def get_tool_group_names() -> Dict[str, List[str]]:
|
||||
"""Get available tool groups and their tool names."""
|
||||
return TOOL_GROUPS.copy()
|
||||
|
||||
|
||||
def print_tool_summary(tools: List[Callable]) -> None:
|
||||
"""Print a summary of loaded tools (for debugging)."""
|
||||
acp_count = sum(1 for t in tools if t.__name__.startswith("acp_"))
|
||||
lsp_count = sum(1 for t in tools if t.__name__.startswith("lsp_"))
|
||||
basic_count = len(tools) - acp_count - lsp_count
|
||||
|
||||
print(f"Interactive tools loaded: {len(tools)} total")
|
||||
print(f" - ACP tools: {acp_count}")
|
||||
print(f" - LSP tools: {lsp_count}")
|
||||
print(f" - Basic tools: {basic_count}")
|
||||
@@ -0,0 +1,81 @@
|
||||
"""
|
||||
PraisonAI CLI Profiler Module.
|
||||
|
||||
Provides detailed cProfile-based profiling for CLI commands and agent execution.
|
||||
Supports per-file/per-function timing, call graphs, and import time analysis.
|
||||
|
||||
Includes Tier 0/1/2 optimizations:
|
||||
- Tier 0: Lazy imports, provider caching, CLI startup optimization
|
||||
- Tier 1: Connection pooling, prewarm hooks (opt-in)
|
||||
- Tier 2: Lite mode, async init, perf snapshot
|
||||
"""
|
||||
|
||||
from .core import (
|
||||
ProfilerConfig,
|
||||
ProfilerResult,
|
||||
QueryProfiler,
|
||||
run_profiled_query,
|
||||
format_profile_report,
|
||||
)
|
||||
|
||||
from .suite import (
|
||||
ScenarioConfig,
|
||||
ScenarioResult,
|
||||
SuiteResult,
|
||||
ProfileSuiteRunner,
|
||||
run_profile_suite,
|
||||
)
|
||||
|
||||
from .optimizations import (
|
||||
# Tier 0
|
||||
ProviderCache,
|
||||
get_provider_cache,
|
||||
resolve_model_provider,
|
||||
LazyImporter,
|
||||
# Tier 1
|
||||
ClientPool,
|
||||
get_client_pool,
|
||||
PrewarmManager,
|
||||
# Tier 2
|
||||
LiteModeConfig,
|
||||
get_lite_mode_config,
|
||||
is_lite_mode,
|
||||
PerfSnapshot,
|
||||
PerfComparison,
|
||||
PerfSnapshotManager,
|
||||
create_snapshot_from_suite,
|
||||
format_comparison_report,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Core profiler
|
||||
'ProfilerConfig',
|
||||
'ProfilerResult',
|
||||
'QueryProfiler',
|
||||
'run_profiled_query',
|
||||
'format_profile_report',
|
||||
# Suite runner
|
||||
'ScenarioConfig',
|
||||
'ScenarioResult',
|
||||
'SuiteResult',
|
||||
'ProfileSuiteRunner',
|
||||
'run_profile_suite',
|
||||
# Tier 0 optimizations
|
||||
'ProviderCache',
|
||||
'get_provider_cache',
|
||||
'resolve_model_provider',
|
||||
'LazyImporter',
|
||||
# Tier 1 optimizations
|
||||
'ClientPool',
|
||||
'get_client_pool',
|
||||
'PrewarmManager',
|
||||
# Tier 2 optimizations
|
||||
'LiteModeConfig',
|
||||
'get_lite_mode_config',
|
||||
'is_lite_mode',
|
||||
'PerfSnapshot',
|
||||
'PerfComparison',
|
||||
'PerfSnapshotManager',
|
||||
'create_snapshot_from_suite',
|
||||
'format_comparison_report',
|
||||
]
|
||||
@@ -0,0 +1,558 @@
|
||||
"""
|
||||
Core profiler implementation for PraisonAI CLI.
|
||||
|
||||
Provides cProfile-based profiling with detailed per-file/per-function timing,
|
||||
call graphs, and import time analysis.
|
||||
"""
|
||||
|
||||
import cProfile
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import pstats
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
|
||||
# Patterns for secrets to redact
|
||||
SECRET_PATTERNS = [
|
||||
r'sk-[a-zA-Z0-9]{20,}', # OpenAI keys
|
||||
r'sk-ant-[a-zA-Z0-9-]{20,}', # Anthropic keys
|
||||
r'AIza[a-zA-Z0-9_-]{35}', # Google API keys
|
||||
r'tvly-[a-zA-Z0-9-]{20,}', # Tavily keys
|
||||
r'[a-zA-Z0-9_-]{32,}', # Generic long tokens (be careful)
|
||||
]
|
||||
|
||||
|
||||
def redact_secrets(text: str) -> str:
|
||||
"""Redact potential secrets from text."""
|
||||
result = text
|
||||
for pattern in SECRET_PATTERNS[:4]: # Skip generic pattern for safety
|
||||
result = re.sub(pattern, '[REDACTED]', result)
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProfilerConfig:
|
||||
"""Configuration for profiler."""
|
||||
deep: bool = False # Enable deep call tracing
|
||||
limit: int = 30 # Top N functions to show
|
||||
sort_by: str = "cumulative" # cumulative or tottime
|
||||
show_files: bool = False # Group by file
|
||||
show_callers: bool = False # Show callers
|
||||
show_callees: bool = False # Show callees
|
||||
importtime: bool = False # Show import timing
|
||||
first_token: bool = False # Track time to first token
|
||||
save_path: Optional[str] = None # Path to save artifacts
|
||||
output_format: str = "text" # text or json
|
||||
stream: bool = False # Streaming mode
|
||||
|
||||
|
||||
@dataclass
|
||||
class TimingBreakdown:
|
||||
"""Timing breakdown for profiling."""
|
||||
cli_parse_ms: float = 0.0
|
||||
imports_ms: float = 0.0
|
||||
agent_construction_ms: float = 0.0
|
||||
model_init_ms: float = 0.0
|
||||
first_token_ms: float = 0.0
|
||||
total_run_ms: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionStats:
|
||||
"""Statistics for a single function."""
|
||||
name: str
|
||||
filename: str
|
||||
lineno: int
|
||||
calls: int
|
||||
tottime: float # Time in function excluding subcalls
|
||||
cumtime: float # Time in function including subcalls
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"filename": self.filename,
|
||||
"lineno": self.lineno,
|
||||
"calls": self.calls,
|
||||
"tottime_ms": self.tottime * 1000,
|
||||
"cumtime_ms": self.cumtime * 1000,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileStats:
|
||||
"""Statistics aggregated by file."""
|
||||
filename: str
|
||||
total_time: float
|
||||
functions: List[FunctionStats] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"filename": self.filename,
|
||||
"total_time_ms": self.total_time * 1000,
|
||||
"function_count": len(self.functions),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProfilerResult:
|
||||
"""Result from a profiling run."""
|
||||
prompt: str
|
||||
response: str
|
||||
timing: TimingBreakdown
|
||||
function_stats: List[FunctionStats]
|
||||
file_stats: List[FileStats]
|
||||
callers: Dict[str, List[str]] = field(default_factory=dict)
|
||||
callees: Dict[str, List[str]] = field(default_factory=dict)
|
||||
import_times: List[Tuple[str, float]] = field(default_factory=list)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"timestamp": self.timestamp,
|
||||
"metadata": self.metadata,
|
||||
"prompt": self.prompt[:200] + "..." if len(self.prompt) > 200 else self.prompt,
|
||||
"response_preview": self.response[:200] + "..." if len(self.response) > 200 else self.response,
|
||||
"timing": {
|
||||
"cli_parse_ms": self.timing.cli_parse_ms,
|
||||
"imports_ms": self.timing.imports_ms,
|
||||
"agent_construction_ms": self.timing.agent_construction_ms,
|
||||
"model_init_ms": self.timing.model_init_ms,
|
||||
"first_token_ms": self.timing.first_token_ms,
|
||||
"total_run_ms": self.timing.total_run_ms,
|
||||
},
|
||||
"top_functions": [f.to_dict() for f in self.function_stats[:20]],
|
||||
"top_files": [f.to_dict() for f in self.file_stats[:10]],
|
||||
}
|
||||
|
||||
|
||||
class QueryProfiler:
|
||||
"""
|
||||
Profiler for query execution.
|
||||
|
||||
Uses cProfile for function-level profiling and optional sys.setprofile
|
||||
for deep call tracing.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[ProfilerConfig] = None):
|
||||
self.config = config or ProfilerConfig()
|
||||
self.profiler = cProfile.Profile()
|
||||
self._first_token_time: Optional[float] = None
|
||||
self._start_time: float = 0.0
|
||||
self._deep_trace_data: Dict[str, Dict] = {}
|
||||
|
||||
def _on_first_token(self):
|
||||
"""Called when first token is received (for streaming)."""
|
||||
if self._first_token_time is None:
|
||||
self._first_token_time = time.perf_counter()
|
||||
|
||||
def _deep_trace_callback(self, frame, event, arg):
|
||||
"""Callback for sys.setprofile deep tracing."""
|
||||
if event not in ('call', 'return'):
|
||||
return
|
||||
|
||||
filename = frame.f_code.co_filename
|
||||
funcname = frame.f_code.co_name
|
||||
lineno = frame.f_lineno
|
||||
|
||||
key = f"{filename}:{funcname}:{lineno}"
|
||||
|
||||
if key not in self._deep_trace_data:
|
||||
self._deep_trace_data[key] = {
|
||||
'filename': filename,
|
||||
'funcname': funcname,
|
||||
'lineno': lineno,
|
||||
'calls': 0,
|
||||
'total_time': 0.0,
|
||||
'start_times': [],
|
||||
}
|
||||
|
||||
if event == 'call':
|
||||
self._deep_trace_data[key]['calls'] += 1
|
||||
self._deep_trace_data[key]['start_times'].append(time.perf_counter())
|
||||
elif event == 'return' and self._deep_trace_data[key]['start_times']:
|
||||
start = self._deep_trace_data[key]['start_times'].pop()
|
||||
self._deep_trace_data[key]['total_time'] += time.perf_counter() - start
|
||||
|
||||
def profile_query(
|
||||
self,
|
||||
prompt: str,
|
||||
model: Optional[str] = None,
|
||||
stream: bool = False,
|
||||
) -> ProfilerResult:
|
||||
"""
|
||||
Profile a query execution.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to execute
|
||||
model: Optional model to use
|
||||
stream: Whether to use streaming mode
|
||||
|
||||
Returns:
|
||||
ProfilerResult with detailed timing and statistics
|
||||
"""
|
||||
timing = TimingBreakdown()
|
||||
response = ""
|
||||
|
||||
# Metadata
|
||||
metadata = self._collect_metadata(model)
|
||||
|
||||
# Time CLI parsing (simulated - actual parsing happens before this)
|
||||
cli_start = time.perf_counter()
|
||||
timing.cli_parse_ms = (time.perf_counter() - cli_start) * 1000
|
||||
|
||||
# Time imports
|
||||
import_start = time.perf_counter()
|
||||
try:
|
||||
from praisonaiagents import Agent
|
||||
except ImportError:
|
||||
raise ImportError("praisonaiagents not installed")
|
||||
timing.imports_ms = (time.perf_counter() - import_start) * 1000
|
||||
|
||||
# Time agent construction
|
||||
construct_start = time.perf_counter()
|
||||
agent_config = {
|
||||
"name": "ProfilerAgent",
|
||||
"role": "Assistant",
|
||||
"goal": "Complete the task",
|
||||
"verbose": False,
|
||||
}
|
||||
if model:
|
||||
agent_config["llm"] = model
|
||||
|
||||
agent = Agent(**agent_config)
|
||||
timing.agent_construction_ms = (time.perf_counter() - construct_start) * 1000
|
||||
|
||||
# Time model initialization (first call)
|
||||
model_init_start = time.perf_counter()
|
||||
timing.model_init_ms = (time.perf_counter() - model_init_start) * 1000
|
||||
|
||||
# Profile the actual execution
|
||||
self._start_time = time.perf_counter()
|
||||
self._first_token_time = None
|
||||
|
||||
if self.config.deep:
|
||||
sys.setprofile(self._deep_trace_callback)
|
||||
|
||||
self.profiler.enable()
|
||||
|
||||
try:
|
||||
if stream and hasattr(agent, '_start_stream'):
|
||||
# Streaming mode with first token tracking
|
||||
chunks = []
|
||||
for chunk in agent._start_stream(prompt):
|
||||
if not chunks:
|
||||
self._on_first_token()
|
||||
chunks.append(chunk)
|
||||
response = ''.join(chunks)
|
||||
else:
|
||||
# Non-streaming mode
|
||||
response = agent.start(prompt)
|
||||
if response is None:
|
||||
response = ""
|
||||
finally:
|
||||
self.profiler.disable()
|
||||
if self.config.deep:
|
||||
sys.setprofile(None)
|
||||
|
||||
end_time = time.perf_counter()
|
||||
timing.total_run_ms = (end_time - self._start_time) * 1000
|
||||
|
||||
if self._first_token_time:
|
||||
timing.first_token_ms = (self._first_token_time - self._start_time) * 1000
|
||||
|
||||
# Extract statistics
|
||||
function_stats = self._extract_function_stats()
|
||||
file_stats = self._extract_file_stats(function_stats)
|
||||
callers, callees = self._extract_call_graph()
|
||||
|
||||
# Get import times if requested
|
||||
import_times = []
|
||||
if self.config.importtime:
|
||||
import_times = self._get_import_times()
|
||||
|
||||
return ProfilerResult(
|
||||
prompt=prompt,
|
||||
response=str(response),
|
||||
timing=timing,
|
||||
function_stats=function_stats,
|
||||
file_stats=file_stats,
|
||||
callers=callers if self.config.show_callers else {},
|
||||
callees=callees if self.config.show_callees else {},
|
||||
import_times=import_times,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
def _collect_metadata(self, model: Optional[str]) -> Dict[str, Any]:
|
||||
"""Collect system metadata."""
|
||||
import platform
|
||||
try:
|
||||
from praisonai.version import __version__ as praisonai_version
|
||||
except ImportError:
|
||||
praisonai_version = "unknown"
|
||||
|
||||
return {
|
||||
"python_version": platform.python_version(),
|
||||
"platform": platform.platform(),
|
||||
"praisonai_version": praisonai_version,
|
||||
"model": model or "default",
|
||||
"timestamp": datetime.utcnow().isoformat() + "Z",
|
||||
}
|
||||
|
||||
def _extract_function_stats(self) -> List[FunctionStats]:
|
||||
"""Extract function statistics from profiler."""
|
||||
stats_stream = io.StringIO()
|
||||
stats = pstats.Stats(self.profiler, stream=stats_stream)
|
||||
|
||||
# Sort by configured key
|
||||
sort_key = pstats.SortKey.CUMULATIVE if self.config.sort_by == "cumulative" else pstats.SortKey.TIME
|
||||
stats.sort_stats(sort_key)
|
||||
|
||||
function_stats = []
|
||||
for (filename, lineno, funcname), (cc, nc, tt, ct, callers) in stats.stats.items():
|
||||
function_stats.append(FunctionStats(
|
||||
name=funcname,
|
||||
filename=filename,
|
||||
lineno=lineno,
|
||||
calls=nc,
|
||||
tottime=tt,
|
||||
cumtime=ct,
|
||||
))
|
||||
|
||||
# Sort and limit
|
||||
if self.config.sort_by == "cumulative":
|
||||
function_stats.sort(key=lambda x: x.cumtime, reverse=True)
|
||||
else:
|
||||
function_stats.sort(key=lambda x: x.tottime, reverse=True)
|
||||
|
||||
return function_stats[:self.config.limit]
|
||||
|
||||
def _extract_file_stats(self, function_stats: List[FunctionStats]) -> List[FileStats]:
|
||||
"""Aggregate statistics by file."""
|
||||
file_map: Dict[str, FileStats] = {}
|
||||
|
||||
for func in function_stats:
|
||||
if func.filename not in file_map:
|
||||
file_map[func.filename] = FileStats(
|
||||
filename=func.filename,
|
||||
total_time=0.0,
|
||||
functions=[],
|
||||
)
|
||||
file_map[func.filename].total_time += func.cumtime
|
||||
file_map[func.filename].functions.append(func)
|
||||
|
||||
file_stats = list(file_map.values())
|
||||
file_stats.sort(key=lambda x: x.total_time, reverse=True)
|
||||
|
||||
return file_stats[:self.config.limit]
|
||||
|
||||
def _extract_call_graph(self) -> Tuple[Dict[str, List[str]], Dict[str, List[str]]]:
|
||||
"""Extract caller/callee relationships."""
|
||||
callers: Dict[str, List[str]] = {}
|
||||
callees: Dict[str, List[str]] = {}
|
||||
|
||||
stats = pstats.Stats(self.profiler)
|
||||
|
||||
for (filename, lineno, funcname), (cc, nc, tt, ct, caller_dict) in stats.stats.items():
|
||||
func_key = f"{funcname} ({os.path.basename(filename)}:{lineno})"
|
||||
|
||||
if caller_dict:
|
||||
callers[func_key] = []
|
||||
for (caller_file, caller_line, caller_name), _ in caller_dict.items():
|
||||
caller_key = f"{caller_name} ({os.path.basename(caller_file)}:{caller_line})"
|
||||
callers[func_key].append(caller_key)
|
||||
|
||||
if caller_key not in callees:
|
||||
callees[caller_key] = []
|
||||
callees[caller_key].append(func_key)
|
||||
|
||||
return callers, callees
|
||||
|
||||
def _get_import_times(self) -> List[Tuple[str, float]]:
|
||||
"""Get import timing by running subprocess with -X importtime."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-X", "importtime", "-c", "import praisonaiagents"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
import_times = []
|
||||
for line in result.stderr.split('\n'):
|
||||
if 'import time:' in line:
|
||||
# Parse: "import time: 123 | 456 | module_name"
|
||||
match = re.search(r'import time:\s+(\d+)\s+\|\s+(\d+)\s+\|\s+(.+)', line)
|
||||
if match:
|
||||
self_time = int(match.group(1))
|
||||
cumulative = int(match.group(2))
|
||||
module = match.group(3).strip()
|
||||
import_times.append((module, cumulative / 1000000)) # Convert to seconds
|
||||
|
||||
# Sort by time descending
|
||||
import_times.sort(key=lambda x: x[1], reverse=True)
|
||||
return import_times[:20] # Top 20
|
||||
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
def save_artifacts(self, result: ProfilerResult, base_path: str):
|
||||
"""Save profiling artifacts to files."""
|
||||
base = Path(base_path)
|
||||
base.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Save .prof binary
|
||||
prof_path = str(base) + ".prof"
|
||||
self.profiler.dump_stats(prof_path)
|
||||
|
||||
# Save .txt report
|
||||
txt_path = str(base) + ".txt"
|
||||
with open(txt_path, 'w') as f:
|
||||
f.write(format_profile_report(result, self.config))
|
||||
|
||||
# Save .json if requested
|
||||
if self.config.output_format == "json":
|
||||
json_path = str(base) + ".json"
|
||||
with open(json_path, 'w') as f:
|
||||
json.dump(result.to_dict(), f, indent=2, default=str)
|
||||
|
||||
return prof_path, txt_path
|
||||
|
||||
|
||||
def run_profiled_query(
|
||||
prompt: str,
|
||||
config: Optional[ProfilerConfig] = None,
|
||||
model: Optional[str] = None,
|
||||
) -> ProfilerResult:
|
||||
"""
|
||||
Run a profiled query.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to execute
|
||||
config: Profiler configuration
|
||||
model: Optional model to use
|
||||
|
||||
Returns:
|
||||
ProfilerResult with detailed timing and statistics
|
||||
"""
|
||||
profiler = QueryProfiler(config)
|
||||
return profiler.profile_query(prompt, model=model, stream=config.stream if config else False)
|
||||
|
||||
|
||||
def format_profile_report(result: ProfilerResult, config: Optional[ProfilerConfig] = None) -> str:
|
||||
"""
|
||||
Format a profile result as a text report.
|
||||
|
||||
Args:
|
||||
result: ProfilerResult to format
|
||||
config: Optional profiler config
|
||||
|
||||
Returns:
|
||||
Formatted text report
|
||||
"""
|
||||
config = config or ProfilerConfig()
|
||||
lines = []
|
||||
|
||||
# Header
|
||||
lines.append("=" * 70)
|
||||
lines.append("PraisonAI Profile Report")
|
||||
lines.append("=" * 70)
|
||||
lines.append("")
|
||||
|
||||
# Metadata
|
||||
lines.append("## System Information")
|
||||
lines.append(f" Timestamp: {result.timestamp}")
|
||||
lines.append(f" Python Version: {result.metadata.get('python_version', 'N/A')}")
|
||||
lines.append(f" Platform: {result.metadata.get('platform', 'N/A')}")
|
||||
lines.append(f" PraisonAI: {result.metadata.get('praisonai_version', 'N/A')}")
|
||||
lines.append(f" Model: {result.metadata.get('model', 'N/A')}")
|
||||
lines.append("")
|
||||
|
||||
# Timing breakdown
|
||||
lines.append("## Timing Breakdown")
|
||||
lines.append(f" CLI Parse: {result.timing.cli_parse_ms:>10.2f} ms")
|
||||
lines.append(f" Imports: {result.timing.imports_ms:>10.2f} ms")
|
||||
lines.append(f" Agent Construct: {result.timing.agent_construction_ms:>10.2f} ms")
|
||||
lines.append(f" Model Init: {result.timing.model_init_ms:>10.2f} ms")
|
||||
if result.timing.first_token_ms > 0:
|
||||
lines.append(f" First Token: {result.timing.first_token_ms:>10.2f} ms")
|
||||
lines.append(f" Total Run: {result.timing.total_run_ms:>10.2f} ms")
|
||||
lines.append("")
|
||||
|
||||
# Per-file timing (if requested)
|
||||
if config.show_files and result.file_stats:
|
||||
lines.append("## Per-File Timing (Top Files)")
|
||||
lines.append("-" * 70)
|
||||
lines.append(f"{'File':<50} {'Time (ms)':>15}")
|
||||
lines.append("-" * 70)
|
||||
for fs in result.file_stats[:15]:
|
||||
filename = os.path.basename(fs.filename)
|
||||
if len(filename) > 48:
|
||||
filename = "..." + filename[-45:]
|
||||
lines.append(f"{filename:<50} {fs.total_time * 1000:>15.2f}")
|
||||
lines.append("")
|
||||
|
||||
# Per-function timing
|
||||
lines.append("## Per-Function Timing (Top Functions)")
|
||||
lines.append("-" * 70)
|
||||
sort_label = "Cumulative" if config.sort_by == "cumulative" else "Total"
|
||||
lines.append(f"{'Function':<35} {'Calls':>8} {sort_label + ' (ms)':>12} {'Self (ms)':>12}")
|
||||
lines.append("-" * 70)
|
||||
for fs in result.function_stats[:config.limit]:
|
||||
funcname = fs.name
|
||||
if len(funcname) > 33:
|
||||
funcname = funcname[:30] + "..."
|
||||
lines.append(f"{funcname:<35} {fs.calls:>8} {fs.cumtime * 1000:>12.2f} {fs.tottime * 1000:>12.2f}")
|
||||
lines.append("")
|
||||
|
||||
# Callers (if requested)
|
||||
if config.show_callers and result.callers:
|
||||
lines.append("## Callers (Who called each function)")
|
||||
lines.append("-" * 70)
|
||||
for func, caller_list in list(result.callers.items())[:10]:
|
||||
lines.append(f" {func}:")
|
||||
for caller in caller_list[:5]:
|
||||
lines.append(f" <- {caller}")
|
||||
lines.append("")
|
||||
|
||||
# Callees (if requested)
|
||||
if config.show_callees and result.callees:
|
||||
lines.append("## Callees (What each function called)")
|
||||
lines.append("-" * 70)
|
||||
for func, callee_list in list(result.callees.items())[:10]:
|
||||
lines.append(f" {func}:")
|
||||
for callee in callee_list[:5]:
|
||||
lines.append(f" -> {callee}")
|
||||
lines.append("")
|
||||
|
||||
# Import times (if requested)
|
||||
if config.importtime and result.import_times:
|
||||
lines.append("## Import Times (Top Modules)")
|
||||
lines.append("-" * 70)
|
||||
lines.append(f"{'Module':<50} {'Time (ms)':>15}")
|
||||
lines.append("-" * 70)
|
||||
for module, time_sec in result.import_times[:15]:
|
||||
module_name = module
|
||||
if len(module_name) > 48:
|
||||
module_name = "..." + module_name[-45:]
|
||||
lines.append(f"{module_name:<50} {time_sec * 1000:>15.2f}")
|
||||
lines.append("")
|
||||
|
||||
# Response preview
|
||||
lines.append("## Response Preview")
|
||||
lines.append("-" * 70)
|
||||
preview = result.response[:500] + "..." if len(result.response) > 500 else result.response
|
||||
lines.append(preview)
|
||||
lines.append("")
|
||||
|
||||
lines.append("=" * 70)
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,652 @@
|
||||
"""
|
||||
Performance Optimizations for PraisonAI CLI.
|
||||
|
||||
Implements Tier 0/1/2 optimizations:
|
||||
- Tier 0: Lazy imports, provider caching, CLI startup optimization
|
||||
- Tier 1: Connection pooling, prewarm hooks (opt-in)
|
||||
- Tier 2: Lite mode, async init, perf snapshot
|
||||
|
||||
All optimizations are safe-by-default and opt-in where behavior changes.
|
||||
Multi-agent safe: no global mutable state that affects concurrent runs.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
from functools import lru_cache
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TIER 0: Safe Fast Wins (No behavior change)
|
||||
# =============================================================================
|
||||
|
||||
class ProviderCache:
|
||||
"""
|
||||
Thread-safe cache for provider/model resolution.
|
||||
|
||||
Caches resolved provider configurations to avoid repeated lookups.
|
||||
Per-process cache with clear invalidation rules.
|
||||
Multi-agent safe: uses thread-local storage for mutable state.
|
||||
"""
|
||||
|
||||
_instance = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._cache = {}
|
||||
cls._instance._cache_lock = threading.Lock()
|
||||
cls._instance._hits = 0
|
||||
cls._instance._misses = 0
|
||||
return cls._instance
|
||||
|
||||
def get(self, key: str) -> Optional[Any]:
|
||||
"""Get cached value (thread-safe)."""
|
||||
with self._cache_lock:
|
||||
if key in self._cache:
|
||||
self._hits += 1
|
||||
return self._cache[key]
|
||||
self._misses += 1
|
||||
return None
|
||||
|
||||
def set(self, key: str, value: Any, ttl_seconds: int = 300) -> None:
|
||||
"""Set cached value with TTL (thread-safe)."""
|
||||
with self._cache_lock:
|
||||
self._cache[key] = {
|
||||
'value': value,
|
||||
'expires': time.time() + ttl_seconds,
|
||||
}
|
||||
|
||||
def invalidate(self, key: Optional[str] = None) -> None:
|
||||
"""Invalidate cache entry or entire cache."""
|
||||
with self._cache_lock:
|
||||
if key:
|
||||
self._cache.pop(key, None)
|
||||
else:
|
||||
self._cache.clear()
|
||||
|
||||
def get_stats(self) -> Dict[str, int]:
|
||||
"""Get cache statistics."""
|
||||
with self._cache_lock:
|
||||
return {
|
||||
'hits': self._hits,
|
||||
'misses': self._misses,
|
||||
'size': len(self._cache),
|
||||
}
|
||||
|
||||
def cleanup_expired(self) -> int:
|
||||
"""Remove expired entries. Returns count of removed entries."""
|
||||
now = time.time()
|
||||
removed = 0
|
||||
with self._cache_lock:
|
||||
expired_keys = [
|
||||
k for k, v in self._cache.items()
|
||||
if isinstance(v, dict) and v.get('expires', float('inf')) < now
|
||||
]
|
||||
for k in expired_keys:
|
||||
del self._cache[k]
|
||||
removed += 1
|
||||
return removed
|
||||
|
||||
|
||||
# Global provider cache instance (singleton, thread-safe)
|
||||
_provider_cache = None
|
||||
|
||||
def get_provider_cache() -> ProviderCache:
|
||||
"""Get the global provider cache instance."""
|
||||
global _provider_cache
|
||||
if _provider_cache is None:
|
||||
_provider_cache = ProviderCache()
|
||||
return _provider_cache
|
||||
|
||||
|
||||
@lru_cache(maxsize=32)
|
||||
def resolve_model_provider(model: str) -> Tuple[str, str]:
|
||||
"""
|
||||
Resolve model string to (provider, model_name) tuple.
|
||||
|
||||
Cached using LRU cache for repeated lookups.
|
||||
"""
|
||||
if '/' in model:
|
||||
parts = model.split('/', 1)
|
||||
return parts[0], parts[1]
|
||||
|
||||
# Default provider mappings
|
||||
if model.startswith('gpt-') or model.startswith('o1-') or model.startswith('o3-'):
|
||||
return 'openai', model
|
||||
elif model.startswith('claude-'):
|
||||
return 'anthropic', model
|
||||
elif model.startswith('gemini-'):
|
||||
return 'google', model
|
||||
elif model.startswith('llama') or model.startswith('mistral'):
|
||||
return 'ollama', model
|
||||
|
||||
return 'openai', model # Default to OpenAI
|
||||
|
||||
|
||||
class LazyImporter:
|
||||
"""
|
||||
Lazy importer for heavy modules.
|
||||
|
||||
Defers import until first use, reducing CLI startup time.
|
||||
Thread-safe and caches imported modules.
|
||||
"""
|
||||
|
||||
_imports: Dict[str, Any] = {}
|
||||
_lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def get(cls, module_path: str, attr: Optional[str] = None) -> Any:
|
||||
"""
|
||||
Lazily import a module or attribute.
|
||||
|
||||
Args:
|
||||
module_path: Full module path (e.g., 'openai.OpenAI')
|
||||
attr: Optional attribute to get from module
|
||||
|
||||
Returns:
|
||||
Imported module or attribute
|
||||
"""
|
||||
cache_key = f"{module_path}.{attr}" if attr else module_path
|
||||
|
||||
with cls._lock:
|
||||
if cache_key in cls._imports:
|
||||
return cls._imports[cache_key]
|
||||
|
||||
# Import outside lock to avoid deadlocks
|
||||
import importlib
|
||||
|
||||
if attr:
|
||||
module = importlib.import_module(module_path)
|
||||
result = getattr(module, attr)
|
||||
else:
|
||||
result = importlib.import_module(module_path)
|
||||
|
||||
with cls._lock:
|
||||
cls._imports[cache_key] = result
|
||||
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def clear(cls) -> None:
|
||||
"""Clear import cache."""
|
||||
with cls._lock:
|
||||
cls._imports.clear()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TIER 1: Medium Effort (Guarded/opt-in)
|
||||
# =============================================================================
|
||||
|
||||
class ClientPool:
|
||||
"""
|
||||
Connection pool for API clients.
|
||||
|
||||
Reuses client instances for repeated calls to the same provider.
|
||||
Multi-agent safe: clients are isolated by key.
|
||||
Does not leak API keys between different configurations.
|
||||
"""
|
||||
|
||||
def __init__(self, max_size: int = 10):
|
||||
self._pool: Dict[str, Any] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._max_size = max_size
|
||||
self._access_times: Dict[str, float] = {}
|
||||
|
||||
def get_or_create(
|
||||
self,
|
||||
key: str,
|
||||
factory: Callable[[], Any],
|
||||
) -> Any:
|
||||
"""
|
||||
Get existing client or create new one.
|
||||
|
||||
Args:
|
||||
key: Unique key for this client configuration
|
||||
factory: Function to create new client if needed
|
||||
|
||||
Returns:
|
||||
Client instance
|
||||
"""
|
||||
with self._lock:
|
||||
if key in self._pool:
|
||||
self._access_times[key] = time.time()
|
||||
return self._pool[key]
|
||||
|
||||
# Evict oldest if at capacity
|
||||
if len(self._pool) >= self._max_size:
|
||||
oldest_key = min(self._access_times, key=self._access_times.get)
|
||||
del self._pool[oldest_key]
|
||||
del self._access_times[oldest_key]
|
||||
|
||||
# Create outside lock
|
||||
client = factory()
|
||||
|
||||
with self._lock:
|
||||
self._pool[key] = client
|
||||
self._access_times[key] = time.time()
|
||||
|
||||
return client
|
||||
|
||||
def remove(self, key: str) -> None:
|
||||
"""Remove client from pool."""
|
||||
with self._lock:
|
||||
self._pool.pop(key, None)
|
||||
self._access_times.pop(key, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all clients from pool."""
|
||||
with self._lock:
|
||||
self._pool.clear()
|
||||
self._access_times.clear()
|
||||
|
||||
def size(self) -> int:
|
||||
"""Get current pool size."""
|
||||
with self._lock:
|
||||
return len(self._pool)
|
||||
|
||||
|
||||
# Global client pool (opt-in usage)
|
||||
_client_pool: Optional[ClientPool] = None
|
||||
|
||||
def get_client_pool() -> ClientPool:
|
||||
"""Get the global client pool instance."""
|
||||
global _client_pool
|
||||
if _client_pool is None:
|
||||
_client_pool = ClientPool()
|
||||
return _client_pool
|
||||
|
||||
|
||||
class PrewarmManager:
|
||||
"""
|
||||
Manager for pre-warming provider connections.
|
||||
|
||||
OPT-IN ONLY: Must be explicitly enabled.
|
||||
Pre-initializes provider clients in background to reduce first-call latency.
|
||||
"""
|
||||
|
||||
_enabled = False
|
||||
_prewarmed: Dict[str, bool] = {}
|
||||
_lock = threading.Lock()
|
||||
_background_thread: Optional[threading.Thread] = None
|
||||
|
||||
@classmethod
|
||||
def enable(cls) -> None:
|
||||
"""Enable pre-warming (opt-in)."""
|
||||
cls._enabled = True
|
||||
|
||||
@classmethod
|
||||
def disable(cls) -> None:
|
||||
"""Disable pre-warming."""
|
||||
cls._enabled = False
|
||||
|
||||
@classmethod
|
||||
def is_enabled(cls) -> bool:
|
||||
"""Check if pre-warming is enabled."""
|
||||
return cls._enabled
|
||||
|
||||
@classmethod
|
||||
def prewarm_provider(cls, provider: str, api_key: Optional[str] = None) -> None:
|
||||
"""
|
||||
Pre-warm a provider connection in background.
|
||||
|
||||
Only runs if pre-warming is enabled.
|
||||
"""
|
||||
if not cls._enabled:
|
||||
return
|
||||
|
||||
with cls._lock:
|
||||
if provider in cls._prewarmed:
|
||||
return
|
||||
cls._prewarmed[provider] = False
|
||||
|
||||
def _prewarm():
|
||||
try:
|
||||
if provider == 'openai':
|
||||
# Just import and create client - don't make API call
|
||||
from openai import OpenAI
|
||||
key = api_key or os.environ.get('OPENAI_API_KEY')
|
||||
if key:
|
||||
_ = OpenAI(api_key=key)
|
||||
elif provider == 'anthropic':
|
||||
from anthropic import Anthropic
|
||||
key = api_key or os.environ.get('ANTHROPIC_API_KEY')
|
||||
if key:
|
||||
_ = Anthropic(api_key=key)
|
||||
|
||||
with cls._lock:
|
||||
cls._prewarmed[provider] = True
|
||||
except Exception:
|
||||
pass # Silently fail - pre-warming is optional
|
||||
|
||||
# Run in background thread
|
||||
thread = threading.Thread(target=_prewarm, daemon=True)
|
||||
thread.start()
|
||||
|
||||
@classmethod
|
||||
def is_prewarmed(cls, provider: str) -> bool:
|
||||
"""Check if provider is pre-warmed."""
|
||||
with cls._lock:
|
||||
return cls._prewarmed.get(provider, False)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""Reset pre-warm state."""
|
||||
with cls._lock:
|
||||
cls._prewarmed.clear()
|
||||
cls._enabled = False
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TIER 2: Architectural (Safe and non-breaking)
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class LiteModeConfig:
|
||||
"""
|
||||
Configuration for lite runtime mode.
|
||||
|
||||
OPT-IN ONLY: Defaults to OFF.
|
||||
Avoids expensive type/model loading when not required.
|
||||
"""
|
||||
enabled: bool = False
|
||||
skip_type_validation: bool = False
|
||||
skip_model_validation: bool = False
|
||||
minimal_imports: bool = False
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> 'LiteModeConfig':
|
||||
"""Create config from environment variables."""
|
||||
return cls(
|
||||
enabled=os.environ.get('PRAISONAI_LITE_MODE', '').lower() in ('1', 'true', 'yes'),
|
||||
skip_type_validation=os.environ.get('PRAISONAI_SKIP_TYPE_VALIDATION', '').lower() in ('1', 'true'),
|
||||
skip_model_validation=os.environ.get('PRAISONAI_SKIP_MODEL_VALIDATION', '').lower() in ('1', 'true'),
|
||||
minimal_imports=os.environ.get('PRAISONAI_MINIMAL_IMPORTS', '').lower() in ('1', 'true'),
|
||||
)
|
||||
|
||||
|
||||
# Global lite mode config
|
||||
_lite_mode_config: Optional[LiteModeConfig] = None
|
||||
|
||||
def get_lite_mode_config() -> LiteModeConfig:
|
||||
"""Get lite mode configuration."""
|
||||
global _lite_mode_config
|
||||
if _lite_mode_config is None:
|
||||
_lite_mode_config = LiteModeConfig.from_env()
|
||||
return _lite_mode_config
|
||||
|
||||
def is_lite_mode() -> bool:
|
||||
"""Check if lite mode is enabled."""
|
||||
return get_lite_mode_config().enabled
|
||||
|
||||
|
||||
@dataclass
|
||||
class PerfSnapshot:
|
||||
"""
|
||||
Performance snapshot for baseline comparison.
|
||||
|
||||
Records timing data that can be compared against later runs.
|
||||
"""
|
||||
timestamp: str
|
||||
name: str
|
||||
startup_cold_ms: float
|
||||
startup_warm_ms: float
|
||||
import_time_ms: float
|
||||
query_time_ms: float
|
||||
first_token_ms: float
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
'timestamp': self.timestamp,
|
||||
'name': self.name,
|
||||
'startup_cold_ms': self.startup_cold_ms,
|
||||
'startup_warm_ms': self.startup_warm_ms,
|
||||
'import_time_ms': self.import_time_ms,
|
||||
'query_time_ms': self.query_time_ms,
|
||||
'first_token_ms': self.first_token_ms,
|
||||
'metadata': self.metadata,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> 'PerfSnapshot':
|
||||
return cls(
|
||||
timestamp=data['timestamp'],
|
||||
name=data['name'],
|
||||
startup_cold_ms=data['startup_cold_ms'],
|
||||
startup_warm_ms=data['startup_warm_ms'],
|
||||
import_time_ms=data['import_time_ms'],
|
||||
query_time_ms=data['query_time_ms'],
|
||||
first_token_ms=data.get('first_token_ms', 0.0),
|
||||
metadata=data.get('metadata', {}),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PerfComparison:
|
||||
"""Comparison between two performance snapshots."""
|
||||
baseline: PerfSnapshot
|
||||
current: PerfSnapshot
|
||||
|
||||
@property
|
||||
def startup_cold_diff_ms(self) -> float:
|
||||
return self.current.startup_cold_ms - self.baseline.startup_cold_ms
|
||||
|
||||
@property
|
||||
def startup_cold_diff_pct(self) -> float:
|
||||
if self.baseline.startup_cold_ms == 0:
|
||||
return 0.0
|
||||
return (self.startup_cold_diff_ms / self.baseline.startup_cold_ms) * 100
|
||||
|
||||
@property
|
||||
def import_time_diff_ms(self) -> float:
|
||||
return self.current.import_time_ms - self.baseline.import_time_ms
|
||||
|
||||
@property
|
||||
def import_time_diff_pct(self) -> float:
|
||||
if self.baseline.import_time_ms == 0:
|
||||
return 0.0
|
||||
return (self.import_time_diff_ms / self.baseline.import_time_ms) * 100
|
||||
|
||||
@property
|
||||
def query_time_diff_ms(self) -> float:
|
||||
return self.current.query_time_ms - self.baseline.query_time_ms
|
||||
|
||||
@property
|
||||
def query_time_diff_pct(self) -> float:
|
||||
if self.baseline.query_time_ms == 0:
|
||||
return 0.0
|
||||
return (self.query_time_diff_ms / self.baseline.query_time_ms) * 100
|
||||
|
||||
def is_regression(self, threshold_pct: float = 10.0) -> bool:
|
||||
"""Check if there's a performance regression above threshold."""
|
||||
return (
|
||||
self.startup_cold_diff_pct > threshold_pct or
|
||||
self.import_time_diff_pct > threshold_pct or
|
||||
self.query_time_diff_pct > threshold_pct
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
'baseline': self.baseline.to_dict(),
|
||||
'current': self.current.to_dict(),
|
||||
'diffs': {
|
||||
'startup_cold_ms': self.startup_cold_diff_ms,
|
||||
'startup_cold_pct': self.startup_cold_diff_pct,
|
||||
'import_time_ms': self.import_time_diff_ms,
|
||||
'import_time_pct': self.import_time_diff_pct,
|
||||
'query_time_ms': self.query_time_diff_ms,
|
||||
'query_time_pct': self.query_time_diff_pct,
|
||||
},
|
||||
'is_regression': self.is_regression(),
|
||||
}
|
||||
|
||||
|
||||
class PerfSnapshotManager:
|
||||
"""
|
||||
Manager for performance snapshots.
|
||||
|
||||
Stores snapshots locally for baseline comparison.
|
||||
OPT-IN: User must explicitly invoke snapshot commands.
|
||||
"""
|
||||
|
||||
DEFAULT_DIR = Path.home() / '.praisonai' / 'perf_snapshots'
|
||||
|
||||
def __init__(self, snapshot_dir: Optional[Path] = None):
|
||||
self.snapshot_dir = snapshot_dir or self.DEFAULT_DIR
|
||||
self.snapshot_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def save(self, snapshot: PerfSnapshot) -> Path:
|
||||
"""Save a snapshot to disk."""
|
||||
filename = f"{snapshot.name}_{snapshot.timestamp.replace(':', '-').replace('.', '-')}.json"
|
||||
path = self.snapshot_dir / filename
|
||||
with open(path, 'w') as f:
|
||||
json.dump(snapshot.to_dict(), f, indent=2)
|
||||
return path
|
||||
|
||||
def load(self, name: str) -> Optional[PerfSnapshot]:
|
||||
"""Load the most recent snapshot with given name."""
|
||||
pattern = f"{name}_*.json"
|
||||
files = sorted(self.snapshot_dir.glob(pattern), reverse=True)
|
||||
if not files:
|
||||
return None
|
||||
|
||||
with open(files[0]) as f:
|
||||
data = json.load(f)
|
||||
return PerfSnapshot.from_dict(data)
|
||||
|
||||
def load_baseline(self, name: str = 'baseline') -> Optional[PerfSnapshot]:
|
||||
"""Load the baseline snapshot."""
|
||||
return self.load(name)
|
||||
|
||||
def save_baseline(self, snapshot: PerfSnapshot) -> Path:
|
||||
"""Save as baseline (overwrites previous baseline)."""
|
||||
snapshot.name = 'baseline'
|
||||
# Remove old baselines
|
||||
for f in self.snapshot_dir.glob('baseline_*.json'):
|
||||
f.unlink()
|
||||
return self.save(snapshot)
|
||||
|
||||
def compare(self, current: PerfSnapshot, baseline_name: str = 'baseline') -> Optional[PerfComparison]:
|
||||
"""Compare current snapshot against baseline."""
|
||||
baseline = self.load(baseline_name)
|
||||
if not baseline:
|
||||
return None
|
||||
return PerfComparison(baseline=baseline, current=current)
|
||||
|
||||
def list_snapshots(self) -> List[str]:
|
||||
"""List all saved snapshot names."""
|
||||
files = self.snapshot_dir.glob('*.json')
|
||||
names = set()
|
||||
for f in files:
|
||||
# Extract name from filename (before first underscore)
|
||||
name = f.stem.rsplit('_', 1)[0] if '_' in f.stem else f.stem
|
||||
names.add(name)
|
||||
return sorted(names)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Utility Functions
|
||||
# =============================================================================
|
||||
|
||||
def create_snapshot_from_suite(suite_result: Any, name: str = 'current') -> PerfSnapshot:
|
||||
"""
|
||||
Create a PerfSnapshot from a SuiteResult.
|
||||
|
||||
Args:
|
||||
suite_result: SuiteResult from profile suite
|
||||
name: Name for the snapshot
|
||||
|
||||
Returns:
|
||||
PerfSnapshot instance
|
||||
"""
|
||||
# Get average query time from scenarios
|
||||
query_times = []
|
||||
first_token_times = []
|
||||
for scenario in suite_result.scenarios:
|
||||
query_times.extend(scenario.total_times)
|
||||
first_token_times.extend(scenario.first_token_times)
|
||||
|
||||
avg_query_time = sum(query_times) / len(query_times) if query_times else 0.0
|
||||
avg_first_token = sum(first_token_times) / len(first_token_times) if first_token_times else 0.0
|
||||
|
||||
# Get import time from analysis
|
||||
import_time = 0.0
|
||||
if suite_result.import_analysis:
|
||||
import_time = suite_result.import_analysis[0].get('cumulative_ms', 0.0)
|
||||
|
||||
return PerfSnapshot(
|
||||
timestamp=suite_result.timestamp,
|
||||
name=name,
|
||||
startup_cold_ms=suite_result.startup_cold_ms,
|
||||
startup_warm_ms=suite_result.startup_warm_ms,
|
||||
import_time_ms=import_time,
|
||||
query_time_ms=avg_query_time,
|
||||
first_token_ms=avg_first_token,
|
||||
metadata=suite_result.metadata,
|
||||
)
|
||||
|
||||
|
||||
def format_comparison_report(comparison: PerfComparison) -> str:
|
||||
"""Format a comparison report as text."""
|
||||
lines = []
|
||||
lines.append("=" * 70)
|
||||
lines.append("Performance Comparison Report")
|
||||
lines.append("=" * 70)
|
||||
lines.append("")
|
||||
|
||||
lines.append(f"Baseline: {comparison.baseline.name} ({comparison.baseline.timestamp})")
|
||||
lines.append(f"Current: {comparison.current.name} ({comparison.current.timestamp})")
|
||||
lines.append("")
|
||||
|
||||
lines.append("-" * 70)
|
||||
lines.append(f"{'Metric':<25} {'Baseline':>12} {'Current':>12} {'Diff':>12} {'%':>8}")
|
||||
lines.append("-" * 70)
|
||||
|
||||
def format_row(name: str, baseline: float, current: float, diff: float, pct: float) -> str:
|
||||
sign = '+' if diff > 0 else ''
|
||||
pct_sign = '+' if pct > 0 else ''
|
||||
return f"{name:<25} {baseline:>12.2f} {current:>12.2f} {sign}{diff:>11.2f} {pct_sign}{pct:>7.1f}%"
|
||||
|
||||
lines.append(format_row(
|
||||
"Startup Cold (ms)",
|
||||
comparison.baseline.startup_cold_ms,
|
||||
comparison.current.startup_cold_ms,
|
||||
comparison.startup_cold_diff_ms,
|
||||
comparison.startup_cold_diff_pct,
|
||||
))
|
||||
|
||||
lines.append(format_row(
|
||||
"Import Time (ms)",
|
||||
comparison.baseline.import_time_ms,
|
||||
comparison.current.import_time_ms,
|
||||
comparison.import_time_diff_ms,
|
||||
comparison.import_time_diff_pct,
|
||||
))
|
||||
|
||||
lines.append(format_row(
|
||||
"Query Time (ms)",
|
||||
comparison.baseline.query_time_ms,
|
||||
comparison.current.query_time_ms,
|
||||
comparison.query_time_diff_ms,
|
||||
comparison.query_time_diff_pct,
|
||||
))
|
||||
|
||||
lines.append("-" * 70)
|
||||
lines.append("")
|
||||
|
||||
if comparison.is_regression():
|
||||
lines.append("⚠️ REGRESSION DETECTED (>10% slower)")
|
||||
else:
|
||||
lines.append("✅ No significant regression")
|
||||
|
||||
lines.append("=" * 70)
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,386 @@
|
||||
"""
|
||||
Profile Suite Runner for PraisonAI CLI.
|
||||
|
||||
Runs a matrix of profiling scenarios and produces aggregated reports.
|
||||
Supports cold/warm start, streaming/non-streaming, and multiple iterations.
|
||||
"""
|
||||
|
||||
import json
|
||||
import statistics
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .core import ProfilerConfig, ProfilerResult, QueryProfiler
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScenarioConfig:
|
||||
"""Configuration for a single profiling scenario."""
|
||||
name: str
|
||||
prompt: str
|
||||
model: Optional[str] = None
|
||||
stream: bool = False
|
||||
iterations: int = 3
|
||||
warmup: int = 1
|
||||
show_files: bool = True
|
||||
limit: int = 20
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScenarioResult:
|
||||
"""Result from running a scenario."""
|
||||
name: str
|
||||
config: ScenarioConfig
|
||||
results: List[ProfilerResult] = field(default_factory=list)
|
||||
|
||||
@property
|
||||
def import_times(self) -> List[float]:
|
||||
return [r.timing.imports_ms for r in self.results]
|
||||
|
||||
@property
|
||||
def total_times(self) -> List[float]:
|
||||
return [r.timing.total_run_ms for r in self.results]
|
||||
|
||||
@property
|
||||
def first_token_times(self) -> List[float]:
|
||||
return [r.timing.first_token_ms for r in self.results if r.timing.first_token_ms > 0]
|
||||
|
||||
def get_stats(self, values: List[float]) -> Dict[str, float]:
|
||||
if not values:
|
||||
return {"mean": 0, "min": 0, "max": 0, "stdev": 0, "median": 0}
|
||||
return {
|
||||
"mean": statistics.mean(values),
|
||||
"min": min(values),
|
||||
"max": max(values),
|
||||
"stdev": statistics.stdev(values) if len(values) > 1 else 0,
|
||||
"median": statistics.median(values),
|
||||
}
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"config": {
|
||||
"prompt": self.config.prompt[:50] + "..." if len(self.config.prompt) > 50 else self.config.prompt,
|
||||
"model": self.config.model or "default",
|
||||
"stream": self.config.stream,
|
||||
"iterations": self.config.iterations,
|
||||
},
|
||||
"import_time_stats": self.get_stats(self.import_times),
|
||||
"total_time_stats": self.get_stats(self.total_times),
|
||||
"first_token_stats": self.get_stats(self.first_token_times) if self.first_token_times else None,
|
||||
"raw_import_times_ms": self.import_times,
|
||||
"raw_total_times_ms": self.total_times,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SuiteResult:
|
||||
"""Result from running a full profiling suite."""
|
||||
scenarios: List[ScenarioResult] = field(default_factory=list)
|
||||
startup_cold_ms: float = 0.0
|
||||
startup_warm_ms: float = 0.0
|
||||
import_analysis: List[Dict[str, Any]] = field(default_factory=list)
|
||||
timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"timestamp": self.timestamp,
|
||||
"metadata": self.metadata,
|
||||
"startup": {
|
||||
"cold_ms": self.startup_cold_ms,
|
||||
"warm_ms": self.startup_warm_ms,
|
||||
},
|
||||
"import_analysis": self.import_analysis[:20],
|
||||
"scenarios": [s.to_dict() for s in self.scenarios],
|
||||
}
|
||||
|
||||
|
||||
class ProfileSuiteRunner:
|
||||
"""
|
||||
Runs a suite of profiling scenarios.
|
||||
|
||||
Collects comprehensive timing data across multiple runs and scenarios.
|
||||
"""
|
||||
|
||||
DEFAULT_SCENARIOS = [
|
||||
ScenarioConfig(name="simple_non_stream", prompt="hi", stream=False, iterations=3),
|
||||
ScenarioConfig(name="simple_stream", prompt="hi", stream=True, iterations=3),
|
||||
ScenarioConfig(name="medium_non_stream", prompt="Explain what Python is in 2 sentences", stream=False, iterations=2),
|
||||
ScenarioConfig(name="medium_stream", prompt="Explain what Python is in 2 sentences", stream=True, iterations=2),
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scenarios: Optional[List[ScenarioConfig]] = None,
|
||||
output_dir: Optional[str] = None,
|
||||
verbose: bool = False,
|
||||
):
|
||||
self.scenarios = scenarios or self.DEFAULT_SCENARIOS
|
||||
self.output_dir = Path(output_dir) if output_dir else Path("/tmp/praisonai_profile_suite")
|
||||
self.verbose = verbose
|
||||
|
||||
def run(self) -> SuiteResult:
|
||||
"""Run the full profiling suite."""
|
||||
result = SuiteResult()
|
||||
result.metadata = self._collect_metadata()
|
||||
|
||||
# Create output directory
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if self.verbose:
|
||||
print("🔬 Starting Profile Suite...")
|
||||
print(f" Output: {self.output_dir}")
|
||||
|
||||
# 1. Measure startup times
|
||||
if self.verbose:
|
||||
print("\n📊 Measuring startup times...")
|
||||
result.startup_cold_ms, result.startup_warm_ms = self._measure_startup()
|
||||
if self.verbose:
|
||||
print(f" Cold: {result.startup_cold_ms:.2f}ms, Warm: {result.startup_warm_ms:.2f}ms")
|
||||
|
||||
# 2. Analyze imports
|
||||
if self.verbose:
|
||||
print("\n📊 Analyzing imports...")
|
||||
result.import_analysis = self._analyze_imports()
|
||||
if self.verbose and result.import_analysis:
|
||||
print(f" Top import: {result.import_analysis[0]['module']} ({result.import_analysis[0]['cumulative_ms']:.2f}ms)")
|
||||
|
||||
# 3. Run scenarios
|
||||
for scenario in self.scenarios:
|
||||
if self.verbose:
|
||||
print(f"\n📊 Running scenario: {scenario.name}")
|
||||
scenario_result = self._run_scenario(scenario)
|
||||
result.scenarios.append(scenario_result)
|
||||
|
||||
if self.verbose:
|
||||
stats = scenario_result.get_stats(scenario_result.total_times)
|
||||
print(f" Total time: {stats['mean']:.2f}ms (±{stats['stdev']:.2f}ms)")
|
||||
|
||||
# 4. Save results
|
||||
self._save_results(result)
|
||||
|
||||
if self.verbose:
|
||||
print(f"\n✅ Suite complete. Results saved to {self.output_dir}")
|
||||
|
||||
return result
|
||||
|
||||
def _collect_metadata(self) -> Dict[str, Any]:
|
||||
"""Collect system metadata."""
|
||||
import platform
|
||||
try:
|
||||
from praisonai.version import __version__ as praisonai_version
|
||||
except ImportError:
|
||||
praisonai_version = "unknown"
|
||||
|
||||
return {
|
||||
"python_version": platform.python_version(),
|
||||
"platform": platform.platform(),
|
||||
"praisonai_version": praisonai_version,
|
||||
"timestamp": datetime.utcnow().isoformat() + "Z",
|
||||
}
|
||||
|
||||
def _measure_startup(self) -> tuple:
|
||||
"""Measure cold and warm startup times."""
|
||||
# Cold start
|
||||
start = time.perf_counter()
|
||||
subprocess.run(
|
||||
[sys.executable, "-c", "import praisonai; import praisonai.cli"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
cold_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
# Warm start
|
||||
start = time.perf_counter()
|
||||
subprocess.run(
|
||||
[sys.executable, "-c", "import praisonai; import praisonai.cli"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
warm_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
return cold_ms, warm_ms
|
||||
|
||||
def _analyze_imports(self) -> List[Dict[str, Any]]:
|
||||
"""Analyze import times."""
|
||||
import re
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-X", "importtime", "-c", "import praisonaiagents"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
import_times = []
|
||||
for line in result.stderr.split('\n'):
|
||||
if 'import time:' in line:
|
||||
match = re.search(r'import time:\s+(\d+)\s+\|\s+(\d+)\s+\|\s+(.+)', line)
|
||||
if match:
|
||||
self_time = int(match.group(1))
|
||||
cumulative = int(match.group(2))
|
||||
module = match.group(3).strip()
|
||||
import_times.append({
|
||||
"module": module,
|
||||
"self_us": self_time,
|
||||
"cumulative_us": cumulative,
|
||||
"self_ms": self_time / 1000,
|
||||
"cumulative_ms": cumulative / 1000,
|
||||
})
|
||||
|
||||
import_times.sort(key=lambda x: x["cumulative_us"], reverse=True)
|
||||
return import_times[:30]
|
||||
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
def _run_scenario(self, scenario: ScenarioConfig) -> ScenarioResult:
|
||||
"""Run a single scenario with multiple iterations."""
|
||||
result = ScenarioResult(name=scenario.name, config=scenario)
|
||||
|
||||
config = ProfilerConfig(
|
||||
stream=scenario.stream,
|
||||
show_files=scenario.show_files,
|
||||
limit=scenario.limit,
|
||||
first_token=scenario.stream,
|
||||
)
|
||||
|
||||
# Warmup runs (discarded)
|
||||
for _ in range(scenario.warmup):
|
||||
try:
|
||||
profiler = QueryProfiler(config)
|
||||
profiler.profile_query(scenario.prompt, model=scenario.model, stream=scenario.stream)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Actual runs
|
||||
for i in range(scenario.iterations):
|
||||
try:
|
||||
profiler = QueryProfiler(config)
|
||||
profile_result = profiler.profile_query(
|
||||
scenario.prompt,
|
||||
model=scenario.model,
|
||||
stream=scenario.stream,
|
||||
)
|
||||
result.results.append(profile_result)
|
||||
|
||||
if self.verbose:
|
||||
print(f" Iteration {i+1}: {profile_result.timing.total_run_ms:.2f}ms")
|
||||
|
||||
except Exception as e:
|
||||
if self.verbose:
|
||||
print(f" Iteration {i+1}: ERROR - {e}")
|
||||
|
||||
return result
|
||||
|
||||
def _save_results(self, result: SuiteResult):
|
||||
"""Save results to files."""
|
||||
# Save JSON summary
|
||||
json_path = self.output_dir / "suite_results.json"
|
||||
with open(json_path, 'w') as f:
|
||||
json.dump(result.to_dict(), f, indent=2, default=str)
|
||||
|
||||
# Save human-readable report
|
||||
report_path = self.output_dir / "suite_report.txt"
|
||||
with open(report_path, 'w') as f:
|
||||
f.write(self._format_report(result))
|
||||
|
||||
def _format_report(self, result: SuiteResult) -> str:
|
||||
"""Format a human-readable report."""
|
||||
lines = []
|
||||
lines.append("=" * 70)
|
||||
lines.append("PraisonAI Profile Suite Report")
|
||||
lines.append("=" * 70)
|
||||
lines.append("")
|
||||
|
||||
# Metadata
|
||||
lines.append("## System Information")
|
||||
lines.append(f" Timestamp: {result.timestamp}")
|
||||
lines.append(f" Python Version: {result.metadata.get('python_version', 'N/A')}")
|
||||
lines.append(f" Platform: {result.metadata.get('platform', 'N/A')}")
|
||||
lines.append(f" PraisonAI: {result.metadata.get('praisonai_version', 'N/A')}")
|
||||
lines.append("")
|
||||
|
||||
# Startup times
|
||||
lines.append("## Startup Times")
|
||||
lines.append(f" Cold Start: {result.startup_cold_ms:>10.2f} ms")
|
||||
lines.append(f" Warm Start: {result.startup_warm_ms:>10.2f} ms")
|
||||
lines.append("")
|
||||
|
||||
# Import analysis
|
||||
if result.import_analysis:
|
||||
lines.append("## Top Import Times")
|
||||
lines.append("-" * 70)
|
||||
lines.append(f"{'Module':<45} {'Self (ms)':>10} {'Cumul (ms)':>12}")
|
||||
lines.append("-" * 70)
|
||||
for imp in result.import_analysis[:15]:
|
||||
module = imp["module"]
|
||||
if len(module) > 43:
|
||||
module = "..." + module[-40:]
|
||||
lines.append(f"{module:<45} {imp['self_ms']:>10.2f} {imp['cumulative_ms']:>12.2f}")
|
||||
lines.append("")
|
||||
|
||||
# Scenario results
|
||||
lines.append("## Scenario Results")
|
||||
lines.append("-" * 70)
|
||||
lines.append(f"{'Scenario':<25} {'Stream':>8} {'Iters':>6} {'Mean (ms)':>12} {'StdDev':>10}")
|
||||
lines.append("-" * 70)
|
||||
|
||||
for scenario in result.scenarios:
|
||||
stats = scenario.get_stats(scenario.total_times)
|
||||
stream_str = "Yes" if scenario.config.stream else "No"
|
||||
lines.append(
|
||||
f"{scenario.name:<25} {stream_str:>8} {len(scenario.results):>6} "
|
||||
f"{stats['mean']:>12.2f} {stats['stdev']:>10.2f}"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
# Summary statistics
|
||||
all_import_times = []
|
||||
all_total_times = []
|
||||
for scenario in result.scenarios:
|
||||
all_import_times.extend(scenario.import_times)
|
||||
all_total_times.extend(scenario.total_times)
|
||||
|
||||
if all_import_times:
|
||||
lines.append("## Overall Statistics")
|
||||
lines.append(f" Import Time (mean): {statistics.mean(all_import_times):>10.2f} ms")
|
||||
lines.append(f" Total Time (mean): {statistics.mean(all_total_times):>10.2f} ms")
|
||||
if len(all_total_times) > 1:
|
||||
lines.append(f" Total Time (stdev): {statistics.stdev(all_total_times):>10.2f} ms")
|
||||
lines.append("")
|
||||
|
||||
lines.append("=" * 70)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def run_profile_suite(
|
||||
output_dir: Optional[str] = None,
|
||||
scenarios: Optional[List[ScenarioConfig]] = None,
|
||||
verbose: bool = True,
|
||||
) -> SuiteResult:
|
||||
"""
|
||||
Run a profile suite and return results.
|
||||
|
||||
Args:
|
||||
output_dir: Directory to save results
|
||||
scenarios: Custom scenarios to run
|
||||
verbose: Print progress
|
||||
|
||||
Returns:
|
||||
SuiteResult with all profiling data
|
||||
"""
|
||||
runner = ProfileSuiteRunner(
|
||||
scenarios=scenarios,
|
||||
output_dir=output_dir,
|
||||
verbose=verbose,
|
||||
)
|
||||
return runner.run()
|
||||
@@ -0,0 +1,431 @@
|
||||
"""
|
||||
Queue Manager for PraisonAI.
|
||||
|
||||
High-level interface that coordinates scheduler, workers, and persistence.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Coroutine, Dict, List, Optional
|
||||
|
||||
from .models import QueuedRun, RunState, RunPriority, QueueConfig, QueueStats, QueueEvent
|
||||
from .scheduler import QueueScheduler, QueueFullError
|
||||
from .worker import WorkerPool
|
||||
from .persistence import QueuePersistence
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QueueManager:
|
||||
"""
|
||||
High-level queue manager that coordinates all queue components.
|
||||
|
||||
Provides a unified interface for:
|
||||
- Submitting and managing runs
|
||||
- Starting/stopping workers
|
||||
- Persistence and crash recovery
|
||||
- Event handling
|
||||
"""
|
||||
|
||||
# Runtime tools registry - stores tools that can't be serialized to JSON
|
||||
# Key: run_id or "default", Value: list of tool callables
|
||||
_tools_registry: Dict[str, list] = {}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[QueueConfig] = None,
|
||||
on_output: Optional[Callable[[str, str], Coroutine[Any, Any, None]]] = None,
|
||||
on_complete: Optional[Callable[[str, QueuedRun], Coroutine[Any, Any, None]]] = None,
|
||||
on_error: Optional[Callable[[str, Exception], Coroutine[Any, Any, None]]] = None,
|
||||
on_event: Optional[Callable[[QueueEvent], Coroutine[Any, Any, None]]] = None,
|
||||
default_tools: Optional[list] = None,
|
||||
):
|
||||
"""
|
||||
Initialize queue manager.
|
||||
|
||||
Args:
|
||||
config: Queue configuration.
|
||||
on_output: Callback for streaming output.
|
||||
on_complete: Callback when run completes.
|
||||
on_error: Callback when run fails.
|
||||
on_event: Callback for queue events.
|
||||
default_tools: Default tools to use for all runs (stored in runtime registry).
|
||||
"""
|
||||
self.config = config or QueueConfig()
|
||||
|
||||
# Initialize components
|
||||
self.scheduler = QueueScheduler(self.config)
|
||||
self.persistence = QueuePersistence(self.config.db_path) if self.config.enable_persistence else None
|
||||
|
||||
# Store callbacks
|
||||
self._on_output = on_output
|
||||
self._on_complete = on_complete
|
||||
self._on_error = on_error
|
||||
self._on_event = on_event
|
||||
|
||||
# Worker pool (created on start)
|
||||
self.workers: Optional[WorkerPool] = None
|
||||
|
||||
# Autosave task
|
||||
self._autosave_task: Optional[asyncio.Task] = None
|
||||
|
||||
# Running state
|
||||
self._running = False
|
||||
|
||||
# Session tracking
|
||||
self._current_session_id: Optional[str] = None
|
||||
|
||||
# Store default tools in registry
|
||||
if default_tools:
|
||||
QueueManager._tools_registry["default"] = default_tools
|
||||
|
||||
async def start(self, recover: bool = True) -> None:
|
||||
"""
|
||||
Start the queue manager.
|
||||
|
||||
Args:
|
||||
recover: If True, recover pending runs from persistence.
|
||||
"""
|
||||
if self._running:
|
||||
return
|
||||
|
||||
logger.info("Starting queue manager")
|
||||
|
||||
# Initialize persistence
|
||||
if self.persistence:
|
||||
self.persistence.initialize()
|
||||
|
||||
if recover:
|
||||
await self._recover_runs()
|
||||
|
||||
# Create and start worker pool
|
||||
self.workers = WorkerPool(
|
||||
scheduler=self.scheduler,
|
||||
on_output=self._on_output,
|
||||
on_complete=self._wrap_complete_callback(),
|
||||
on_error=self._wrap_error_callback(),
|
||||
on_event=self._on_event,
|
||||
max_workers=self.config.max_concurrent_global,
|
||||
poll_interval=self.config.worker_poll_interval,
|
||||
stream_buffer_size=self.config.stream_buffer_size,
|
||||
)
|
||||
|
||||
await self.workers.start()
|
||||
|
||||
# Start autosave
|
||||
if self.persistence and self.config.autosave_interval_seconds > 0:
|
||||
self._autosave_task = asyncio.create_task(
|
||||
self._autosave_loop(),
|
||||
name="queue-autosave"
|
||||
)
|
||||
|
||||
self._running = True
|
||||
logger.info("Queue manager started")
|
||||
|
||||
async def stop(self, timeout: float = 10.0) -> None:
|
||||
"""
|
||||
Stop the queue manager gracefully.
|
||||
|
||||
Args:
|
||||
timeout: Seconds to wait for workers to finish.
|
||||
"""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
logger.info("Stopping queue manager")
|
||||
|
||||
# Stop autosave
|
||||
if self._autosave_task:
|
||||
self._autosave_task.cancel()
|
||||
try:
|
||||
await self._autosave_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._autosave_task = None
|
||||
|
||||
# Stop workers
|
||||
if self.workers:
|
||||
await self.workers.stop(timeout=timeout)
|
||||
self.workers = None
|
||||
|
||||
# Final save
|
||||
if self.persistence:
|
||||
await self._save_all_runs()
|
||||
self.persistence.close()
|
||||
|
||||
self._running = False
|
||||
logger.info("Queue manager stopped")
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
input_content: str,
|
||||
agent_name: str = "Assistant",
|
||||
priority: RunPriority = RunPriority.NORMAL,
|
||||
session_id: Optional[str] = None,
|
||||
workspace: Optional[str] = None,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
run_id: Optional[str] = None,
|
||||
chat_history: Optional[List[Dict[str, str]]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Submit a new run to the queue.
|
||||
|
||||
Args:
|
||||
input_content: The input/prompt for the agent.
|
||||
agent_name: Name of the agent to use.
|
||||
priority: Run priority.
|
||||
session_id: Session ID for grouping.
|
||||
workspace: Workspace path.
|
||||
config: Additional configuration.
|
||||
run_id: Optional specific run ID.
|
||||
chat_history: Previous chat messages for context continuity.
|
||||
|
||||
Returns:
|
||||
The run ID.
|
||||
"""
|
||||
run = QueuedRun(
|
||||
agent_name=agent_name,
|
||||
input_content=input_content,
|
||||
priority=priority,
|
||||
session_id=session_id or self._current_session_id,
|
||||
workspace=workspace or os.getcwd(),
|
||||
config=config or {},
|
||||
max_retries=self.config.default_max_retries,
|
||||
chat_history=chat_history or [],
|
||||
)
|
||||
|
||||
if run_id:
|
||||
run.run_id = run_id
|
||||
|
||||
# Submit to scheduler
|
||||
result_id = await self.scheduler.submit(run)
|
||||
|
||||
# Persist
|
||||
if self.persistence:
|
||||
self.persistence.save_run(run)
|
||||
|
||||
return result_id
|
||||
|
||||
async def cancel(self, run_id: str) -> bool:
|
||||
"""Cancel a run."""
|
||||
result = await self.scheduler.cancel(run_id)
|
||||
|
||||
if result and self.persistence:
|
||||
run = self.scheduler.get_run(run_id)
|
||||
if run:
|
||||
self.persistence.save_run(run)
|
||||
|
||||
return result
|
||||
|
||||
async def retry(self, run_id: str) -> Optional[str]:
|
||||
"""Retry a failed run."""
|
||||
new_id = await self.scheduler.retry(run_id)
|
||||
|
||||
if new_id and self.persistence:
|
||||
run = self.scheduler.get_run(new_id)
|
||||
if run:
|
||||
self.persistence.save_run(run)
|
||||
|
||||
return new_id
|
||||
|
||||
async def pause(self, run_id: str) -> bool:
|
||||
"""Pause a running run."""
|
||||
return await self.scheduler.pause(run_id)
|
||||
|
||||
async def resume(self, run_id: str) -> bool:
|
||||
"""Resume a paused run."""
|
||||
return await self.scheduler.resume(run_id)
|
||||
|
||||
async def update_input(self, run_id: str, new_input: str) -> bool:
|
||||
"""Update the input content of a queued run."""
|
||||
return await self.scheduler.update_input(run_id, new_input)
|
||||
|
||||
def get_run(self, run_id: str) -> Optional[QueuedRun]:
|
||||
"""Get a run by ID."""
|
||||
return self.scheduler.get_run(run_id)
|
||||
|
||||
def get_queued(self) -> List[QueuedRun]:
|
||||
"""Get all queued runs."""
|
||||
return self.scheduler.get_queued()
|
||||
|
||||
def get_running(self) -> List[QueuedRun]:
|
||||
"""Get all running runs."""
|
||||
return self.scheduler.get_running()
|
||||
|
||||
def list_runs(
|
||||
self,
|
||||
state: Optional[RunState] = None,
|
||||
session_id: Optional[str] = None,
|
||||
limit: int = 100,
|
||||
) -> List[QueuedRun]:
|
||||
"""
|
||||
List runs with optional filters.
|
||||
|
||||
For active runs, uses scheduler. For historical, uses persistence.
|
||||
"""
|
||||
if state and state.is_terminal() and self.persistence:
|
||||
return self.persistence.list_runs(
|
||||
state=state,
|
||||
session_id=session_id,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
# Get from scheduler
|
||||
runs = self.scheduler.get_all()
|
||||
|
||||
if state:
|
||||
runs = [r for r in runs if r.state == state]
|
||||
if session_id:
|
||||
runs = [r for r in runs if r.session_id == session_id]
|
||||
|
||||
return runs[:limit]
|
||||
|
||||
async def clear_queue(self) -> int:
|
||||
"""Clear all queued runs."""
|
||||
count = await self.scheduler.clear_queue()
|
||||
return count
|
||||
|
||||
def get_stats(self) -> QueueStats:
|
||||
"""Get queue statistics."""
|
||||
if self.persistence:
|
||||
return self.persistence.get_stats(self._current_session_id)
|
||||
|
||||
# Calculate from scheduler
|
||||
queued = self.scheduler.get_queued()
|
||||
running = self.scheduler.get_running()
|
||||
|
||||
return QueueStats(
|
||||
queued_count=len(queued),
|
||||
running_count=len(running),
|
||||
)
|
||||
|
||||
# Session management
|
||||
|
||||
def set_session(self, session_id: str) -> None:
|
||||
"""Set the current session ID."""
|
||||
self._current_session_id = session_id
|
||||
|
||||
def save_session_state(self, state: Dict[str, Any]) -> None:
|
||||
"""Save session state."""
|
||||
if self.persistence and self._current_session_id:
|
||||
self.persistence.save_session(
|
||||
self._current_session_id,
|
||||
state=state,
|
||||
)
|
||||
|
||||
def load_session_state(self, session_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Load session state."""
|
||||
if self.persistence:
|
||||
return self.persistence.load_session(session_id)
|
||||
return None
|
||||
|
||||
def list_sessions(self, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
"""List recent sessions."""
|
||||
if self.persistence:
|
||||
return self.persistence.list_sessions(limit)
|
||||
return []
|
||||
|
||||
# Internal methods
|
||||
|
||||
async def _recover_runs(self) -> None:
|
||||
"""Recover pending runs from persistence."""
|
||||
if not self.persistence:
|
||||
return
|
||||
|
||||
# Mark interrupted runs as failed
|
||||
interrupted = self.persistence.mark_interrupted_as_failed()
|
||||
if interrupted:
|
||||
logger.warning(f"Marked {interrupted} interrupted runs as failed")
|
||||
|
||||
# Load pending runs
|
||||
pending = self.persistence.load_pending_runs()
|
||||
if pending:
|
||||
logger.info(f"Recovering {len(pending)} pending runs")
|
||||
self.scheduler.load_runs(pending)
|
||||
|
||||
for run in pending:
|
||||
self.persistence.mark_recovered(run.run_id)
|
||||
|
||||
async def _save_all_runs(self) -> None:
|
||||
"""Save all active runs to persistence."""
|
||||
if not self.persistence:
|
||||
return
|
||||
|
||||
for run in self.scheduler.get_all():
|
||||
self.persistence.save_run(run)
|
||||
|
||||
async def _autosave_loop(self) -> None:
|
||||
"""Periodically save runs to persistence."""
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(self.config.autosave_interval_seconds)
|
||||
await self._save_all_runs()
|
||||
logger.debug("Autosaved queue state")
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Autosave error: {e}")
|
||||
|
||||
def _wrap_complete_callback(self):
|
||||
"""Wrap complete callback to include persistence."""
|
||||
async def wrapper(run_id: str, run: QueuedRun):
|
||||
# Persist
|
||||
if self.persistence:
|
||||
self.persistence.save_run(run)
|
||||
|
||||
# Call user callback
|
||||
if self._on_complete:
|
||||
await self._on_complete(run_id, run)
|
||||
|
||||
return wrapper
|
||||
|
||||
def _wrap_error_callback(self):
|
||||
"""Wrap error callback to include persistence."""
|
||||
async def wrapper(run_id: str, error: Exception):
|
||||
# Persist
|
||||
if self.persistence:
|
||||
run = self.scheduler.get_run(run_id)
|
||||
if run:
|
||||
self.persistence.save_run(run)
|
||||
|
||||
# Call user callback
|
||||
if self._on_error:
|
||||
await self._on_error(run_id, error)
|
||||
|
||||
return wrapper
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Whether the queue manager is running."""
|
||||
return self._running
|
||||
|
||||
@property
|
||||
def queued_count(self) -> int:
|
||||
"""Number of queued runs."""
|
||||
return self.scheduler.queued_count
|
||||
|
||||
@property
|
||||
def running_count(self) -> int:
|
||||
"""Number of running runs."""
|
||||
return self.scheduler.running_count
|
||||
|
||||
@classmethod
|
||||
def get_tools_for_run(cls, run_id: str) -> list:
|
||||
"""
|
||||
Get tools for a specific run from the registry.
|
||||
|
||||
Falls back to default tools if no run-specific tools are registered.
|
||||
"""
|
||||
return cls._tools_registry.get(run_id) or cls._tools_registry.get("default", [])
|
||||
|
||||
@classmethod
|
||||
def register_tools(cls, run_id: str, tools: list) -> None:
|
||||
"""Register tools for a specific run."""
|
||||
cls._tools_registry[run_id] = tools
|
||||
|
||||
@classmethod
|
||||
def unregister_tools(cls, run_id: str) -> None:
|
||||
"""Unregister tools for a specific run."""
|
||||
cls._tools_registry.pop(run_id, None)
|
||||
@@ -85,6 +85,9 @@ class QueuedRun:
|
||||
# Configuration
|
||||
config: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Chat history for session continuity
|
||||
chat_history: List[Dict[str, str]] = field(default_factory=list)
|
||||
|
||||
# Streaming state
|
||||
output_chunks: List[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@@ -0,0 +1,564 @@
|
||||
"""
|
||||
SQLite persistence layer for the PraisonAI Queue System.
|
||||
|
||||
Provides crash recovery, session persistence, and run history.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .models import QueuedRun, RunState, RunPriority, QueueStats
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SCHEMA_VERSION = "1.0"
|
||||
|
||||
SCHEMA_SQL = """
|
||||
-- Schema version tracking
|
||||
CREATE TABLE IF NOT EXISTS schema_version (
|
||||
version TEXT PRIMARY KEY,
|
||||
applied_at REAL NOT NULL
|
||||
);
|
||||
|
||||
-- Queue runs
|
||||
CREATE TABLE IF NOT EXISTS runs (
|
||||
run_id TEXT PRIMARY KEY,
|
||||
agent_name TEXT NOT NULL,
|
||||
input_content TEXT,
|
||||
output_content TEXT,
|
||||
state TEXT NOT NULL DEFAULT 'queued',
|
||||
priority INTEGER NOT NULL DEFAULT 1,
|
||||
session_id TEXT,
|
||||
trace_id TEXT,
|
||||
workspace TEXT,
|
||||
user_id TEXT,
|
||||
created_at REAL NOT NULL,
|
||||
started_at REAL,
|
||||
ended_at REAL,
|
||||
error TEXT,
|
||||
retry_count INTEGER DEFAULT 0,
|
||||
max_retries INTEGER DEFAULT 3,
|
||||
parent_run_id TEXT,
|
||||
config TEXT,
|
||||
metrics TEXT,
|
||||
recovered INTEGER DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_runs_state ON runs(state);
|
||||
CREATE INDEX IF NOT EXISTS idx_runs_session ON runs(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_runs_priority_created ON runs(priority DESC, created_at ASC);
|
||||
CREATE INDEX IF NOT EXISTS idx_runs_workspace ON runs(workspace);
|
||||
|
||||
-- Messages (chat history per run)
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
message_id TEXT PRIMARY KEY,
|
||||
run_id TEXT REFERENCES runs(run_id),
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
timestamp REAL NOT NULL,
|
||||
metadata TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_run ON messages(run_id);
|
||||
|
||||
-- Tool calls
|
||||
CREATE TABLE IF NOT EXISTS tool_calls (
|
||||
call_id TEXT PRIMARY KEY,
|
||||
run_id TEXT REFERENCES runs(run_id),
|
||||
tool_name TEXT NOT NULL,
|
||||
args TEXT,
|
||||
result TEXT,
|
||||
started_at REAL NOT NULL,
|
||||
ended_at REAL,
|
||||
error TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_calls_run ON tool_calls(run_id);
|
||||
|
||||
-- Sessions
|
||||
CREATE TABLE IF NOT EXISTS sessions (
|
||||
session_id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
created_at REAL NOT NULL,
|
||||
updated_at REAL NOT NULL,
|
||||
state TEXT,
|
||||
config TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_sessions_user ON sessions(user_id);
|
||||
"""
|
||||
|
||||
|
||||
class QueuePersistence:
|
||||
"""SQLite-backed persistence for the queue system."""
|
||||
|
||||
def __init__(self, db_path: str = ".praison/queue.db"):
|
||||
"""
|
||||
Initialize persistence layer.
|
||||
|
||||
Args:
|
||||
db_path: Path to SQLite database file.
|
||||
"""
|
||||
self.db_path = db_path
|
||||
self._conn: Optional[sqlite3.Connection] = None
|
||||
self._lock = threading.Lock()
|
||||
self._initialized = False
|
||||
|
||||
def _ensure_dir(self) -> None:
|
||||
"""Ensure database directory exists."""
|
||||
db_dir = os.path.dirname(self.db_path)
|
||||
if db_dir:
|
||||
os.makedirs(db_dir, exist_ok=True)
|
||||
|
||||
def _get_connection(self) -> sqlite3.Connection:
|
||||
"""Get or create database connection."""
|
||||
if self._conn is None:
|
||||
self._ensure_dir()
|
||||
self._conn = sqlite3.connect(
|
||||
self.db_path,
|
||||
check_same_thread=False,
|
||||
timeout=30.0
|
||||
)
|
||||
self._conn.row_factory = sqlite3.Row
|
||||
# Enable WAL mode for better concurrency
|
||||
self._conn.execute("PRAGMA journal_mode=WAL")
|
||||
self._conn.execute("PRAGMA synchronous=NORMAL")
|
||||
return self._conn
|
||||
|
||||
@contextmanager
|
||||
def _transaction(self):
|
||||
"""Context manager for database transactions."""
|
||||
conn = self._get_connection()
|
||||
with self._lock:
|
||||
try:
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
def initialize(self) -> None:
|
||||
"""Initialize database schema."""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
with self._transaction() as conn:
|
||||
conn.executescript(SCHEMA_SQL)
|
||||
|
||||
# Check/set schema version
|
||||
cursor = conn.execute(
|
||||
"SELECT version FROM schema_version ORDER BY applied_at DESC LIMIT 1"
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
conn.execute(
|
||||
"INSERT INTO schema_version (version, applied_at) VALUES (?, ?)",
|
||||
(SCHEMA_VERSION, time.time())
|
||||
)
|
||||
elif row["version"] != SCHEMA_VERSION:
|
||||
# Future: handle migrations
|
||||
logger.warning(
|
||||
f"Schema version mismatch: {row['version']} vs {SCHEMA_VERSION}"
|
||||
)
|
||||
|
||||
self._initialized = True
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close database connection."""
|
||||
if self._conn is not None:
|
||||
self._conn.close()
|
||||
self._conn = None
|
||||
self._initialized = False
|
||||
|
||||
# Run operations
|
||||
|
||||
def save_run(self, run: QueuedRun) -> None:
|
||||
"""Save or update a run."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
conn.execute("""
|
||||
INSERT OR REPLACE INTO runs (
|
||||
run_id, agent_name, input_content, output_content,
|
||||
state, priority, session_id, trace_id, workspace, user_id,
|
||||
created_at, started_at, ended_at, error,
|
||||
retry_count, max_retries, parent_run_id,
|
||||
config, metrics
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
run.run_id,
|
||||
run.agent_name,
|
||||
run.input_content,
|
||||
run.output_content,
|
||||
run.state.value,
|
||||
int(run.priority),
|
||||
run.session_id,
|
||||
run.trace_id,
|
||||
run.workspace,
|
||||
run.user_id,
|
||||
run.created_at,
|
||||
run.started_at,
|
||||
run.ended_at,
|
||||
run.error,
|
||||
run.retry_count,
|
||||
run.max_retries,
|
||||
run.parent_run_id,
|
||||
json.dumps(run.config) if run.config else None,
|
||||
json.dumps(run.metrics) if run.metrics else None,
|
||||
))
|
||||
|
||||
def load_run(self, run_id: str) -> Optional[QueuedRun]:
|
||||
"""Load a run by ID."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM runs WHERE run_id = ?",
|
||||
(run_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
return self._row_to_run(row)
|
||||
|
||||
def list_runs(
|
||||
self,
|
||||
state: Optional[RunState] = None,
|
||||
session_id: Optional[str] = None,
|
||||
workspace: Optional[str] = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> List[QueuedRun]:
|
||||
"""List runs with optional filters."""
|
||||
self.initialize()
|
||||
|
||||
query = "SELECT * FROM runs WHERE 1=1"
|
||||
params: List[Any] = []
|
||||
|
||||
if state is not None:
|
||||
query += " AND state = ?"
|
||||
params.append(state.value)
|
||||
|
||||
if session_id is not None:
|
||||
query += " AND session_id = ?"
|
||||
params.append(session_id)
|
||||
|
||||
if workspace is not None:
|
||||
query += " AND workspace = ?"
|
||||
params.append(workspace)
|
||||
|
||||
query += " ORDER BY priority DESC, created_at ASC LIMIT ? OFFSET ?"
|
||||
params.extend([limit, offset])
|
||||
|
||||
with self._transaction() as conn:
|
||||
cursor = conn.execute(query, params)
|
||||
return [self._row_to_run(row) for row in cursor.fetchall()]
|
||||
|
||||
def delete_run(self, run_id: str) -> bool:
|
||||
"""Delete a run and its related data."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
# Delete related data first
|
||||
conn.execute("DELETE FROM messages WHERE run_id = ?", (run_id,))
|
||||
conn.execute("DELETE FROM tool_calls WHERE run_id = ?", (run_id,))
|
||||
|
||||
cursor = conn.execute("DELETE FROM runs WHERE run_id = ?", (run_id,))
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def update_run_state(
|
||||
self,
|
||||
run_id: str,
|
||||
state: RunState,
|
||||
error: Optional[str] = None,
|
||||
output: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Update run state."""
|
||||
self.initialize()
|
||||
|
||||
updates = ["state = ?"]
|
||||
params: List[Any] = [state.value]
|
||||
|
||||
if state == RunState.RUNNING:
|
||||
updates.append("started_at = ?")
|
||||
params.append(time.time())
|
||||
elif state.is_terminal():
|
||||
updates.append("ended_at = ?")
|
||||
params.append(time.time())
|
||||
|
||||
if error is not None:
|
||||
updates.append("error = ?")
|
||||
params.append(error)
|
||||
|
||||
if output is not None:
|
||||
updates.append("output_content = ?")
|
||||
params.append(output)
|
||||
|
||||
params.append(run_id)
|
||||
|
||||
with self._transaction() as conn:
|
||||
cursor = conn.execute(
|
||||
f"UPDATE runs SET {', '.join(updates)} WHERE run_id = ?",
|
||||
params
|
||||
)
|
||||
return cursor.rowcount > 0
|
||||
|
||||
# Crash recovery
|
||||
|
||||
def load_pending_runs(self) -> List[QueuedRun]:
|
||||
"""Load runs that were QUEUED or RUNNING at crash."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
cursor = conn.execute("""
|
||||
SELECT * FROM runs
|
||||
WHERE state IN ('queued', 'running')
|
||||
ORDER BY priority DESC, created_at ASC
|
||||
""")
|
||||
return [self._row_to_run(row) for row in cursor.fetchall()]
|
||||
|
||||
def mark_recovered(self, run_id: str) -> None:
|
||||
"""Mark a run as recovered after restart."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
conn.execute(
|
||||
"UPDATE runs SET recovered = 1 WHERE run_id = ?",
|
||||
(run_id,)
|
||||
)
|
||||
|
||||
def mark_interrupted_as_failed(self) -> int:
|
||||
"""Mark all RUNNING runs as FAILED (for crash recovery)."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
cursor = conn.execute("""
|
||||
UPDATE runs
|
||||
SET state = 'failed',
|
||||
error = 'Interrupted by crash/restart',
|
||||
ended_at = ?
|
||||
WHERE state = 'running'
|
||||
""", (time.time(),))
|
||||
return cursor.rowcount
|
||||
|
||||
# Statistics
|
||||
|
||||
def get_stats(self, session_id: Optional[str] = None) -> QueueStats:
|
||||
"""Get queue statistics."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
where = ""
|
||||
params: List[Any] = []
|
||||
if session_id:
|
||||
where = "WHERE session_id = ?"
|
||||
params = [session_id]
|
||||
|
||||
# Count by state
|
||||
cursor = conn.execute(f"""
|
||||
SELECT state, COUNT(*) as count
|
||||
FROM runs {where}
|
||||
GROUP BY state
|
||||
""", params)
|
||||
|
||||
counts = {row["state"]: row["count"] for row in cursor.fetchall()}
|
||||
|
||||
# Average wait time
|
||||
cursor = conn.execute(f"""
|
||||
SELECT AVG(started_at - created_at) as avg_wait
|
||||
FROM runs
|
||||
{where + ' AND' if where else 'WHERE'} started_at IS NOT NULL
|
||||
""", params)
|
||||
row = cursor.fetchone()
|
||||
avg_wait = row["avg_wait"] or 0.0
|
||||
|
||||
# Average duration
|
||||
cursor = conn.execute(f"""
|
||||
SELECT AVG(ended_at - started_at) as avg_duration
|
||||
FROM runs
|
||||
{where + ' AND' if where else 'WHERE'} ended_at IS NOT NULL AND started_at IS NOT NULL
|
||||
""", params)
|
||||
row = cursor.fetchone()
|
||||
avg_duration = row["avg_duration"] or 0.0
|
||||
|
||||
return QueueStats(
|
||||
queued_count=counts.get("queued", 0),
|
||||
running_count=counts.get("running", 0),
|
||||
succeeded_count=counts.get("succeeded", 0),
|
||||
failed_count=counts.get("failed", 0),
|
||||
cancelled_count=counts.get("cancelled", 0),
|
||||
total_runs=sum(counts.values()),
|
||||
avg_wait_seconds=avg_wait,
|
||||
avg_duration_seconds=avg_duration,
|
||||
)
|
||||
|
||||
# Session operations
|
||||
|
||||
def save_session(
|
||||
self,
|
||||
session_id: str,
|
||||
user_id: Optional[str] = None,
|
||||
state: Optional[Dict[str, Any]] = None,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Save or update a session."""
|
||||
self.initialize()
|
||||
|
||||
now = time.time()
|
||||
|
||||
with self._transaction() as conn:
|
||||
# Check if exists
|
||||
cursor = conn.execute(
|
||||
"SELECT session_id FROM sessions WHERE session_id = ?",
|
||||
(session_id,)
|
||||
)
|
||||
exists = cursor.fetchone() is not None
|
||||
|
||||
if exists:
|
||||
conn.execute("""
|
||||
UPDATE sessions SET
|
||||
user_id = COALESCE(?, user_id),
|
||||
updated_at = ?,
|
||||
state = COALESCE(?, state),
|
||||
config = COALESCE(?, config)
|
||||
WHERE session_id = ?
|
||||
""", (
|
||||
user_id,
|
||||
now,
|
||||
json.dumps(state) if state else None,
|
||||
json.dumps(config) if config else None,
|
||||
session_id,
|
||||
))
|
||||
else:
|
||||
conn.execute("""
|
||||
INSERT INTO sessions (session_id, user_id, created_at, updated_at, state, config)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
session_id,
|
||||
user_id,
|
||||
now,
|
||||
now,
|
||||
json.dumps(state) if state else None,
|
||||
json.dumps(config) if config else None,
|
||||
))
|
||||
|
||||
def load_session(self, session_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Load a session."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM sessions WHERE session_id = ?",
|
||||
(session_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
return {
|
||||
"session_id": row["session_id"],
|
||||
"user_id": row["user_id"],
|
||||
"created_at": row["created_at"],
|
||||
"updated_at": row["updated_at"],
|
||||
"state": json.loads(row["state"]) if row["state"] else None,
|
||||
"config": json.loads(row["config"]) if row["config"] else None,
|
||||
}
|
||||
|
||||
def list_sessions(self, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
"""List recent sessions."""
|
||||
self.initialize()
|
||||
|
||||
with self._transaction() as conn:
|
||||
cursor = conn.execute("""
|
||||
SELECT * FROM sessions
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT ?
|
||||
""", (limit,))
|
||||
|
||||
return [
|
||||
{
|
||||
"session_id": row["session_id"],
|
||||
"user_id": row["user_id"],
|
||||
"created_at": row["created_at"],
|
||||
"updated_at": row["updated_at"],
|
||||
"state": json.loads(row["state"]) if row["state"] else None,
|
||||
"config": json.loads(row["config"]) if row["config"] else None,
|
||||
}
|
||||
for row in cursor.fetchall()
|
||||
]
|
||||
|
||||
# Helper methods
|
||||
|
||||
def _row_to_run(self, row: sqlite3.Row) -> QueuedRun:
|
||||
"""Convert database row to QueuedRun."""
|
||||
return QueuedRun(
|
||||
run_id=row["run_id"],
|
||||
agent_name=row["agent_name"],
|
||||
input_content=row["input_content"] or "",
|
||||
output_content=row["output_content"],
|
||||
state=RunState(row["state"]),
|
||||
priority=RunPriority(row["priority"]),
|
||||
session_id=row["session_id"],
|
||||
trace_id=row["trace_id"],
|
||||
workspace=row["workspace"],
|
||||
user_id=row["user_id"],
|
||||
created_at=row["created_at"],
|
||||
started_at=row["started_at"],
|
||||
ended_at=row["ended_at"],
|
||||
error=row["error"],
|
||||
retry_count=row["retry_count"] or 0,
|
||||
max_retries=row["max_retries"] or 3,
|
||||
parent_run_id=row["parent_run_id"],
|
||||
config=json.loads(row["config"]) if row["config"] else {},
|
||||
metrics=json.loads(row["metrics"]) if row["metrics"] else {},
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
|
||||
def cleanup_old_runs(self, days: int = 30) -> int:
|
||||
"""Delete runs older than specified days."""
|
||||
self.initialize()
|
||||
|
||||
cutoff = time.time() - (days * 24 * 60 * 60)
|
||||
|
||||
with self._transaction() as conn:
|
||||
# Get run IDs to delete
|
||||
cursor = conn.execute("""
|
||||
SELECT run_id FROM runs
|
||||
WHERE created_at < ? AND state IN ('succeeded', 'failed', 'cancelled')
|
||||
""", (cutoff,))
|
||||
run_ids = [row["run_id"] for row in cursor.fetchall()]
|
||||
|
||||
if not run_ids:
|
||||
return 0
|
||||
|
||||
placeholders = ",".join("?" * len(run_ids))
|
||||
|
||||
# Delete related data
|
||||
conn.execute(
|
||||
f"DELETE FROM messages WHERE run_id IN ({placeholders})",
|
||||
run_ids
|
||||
)
|
||||
conn.execute(
|
||||
f"DELETE FROM tool_calls WHERE run_id IN ({placeholders})",
|
||||
run_ids
|
||||
)
|
||||
|
||||
# Delete runs
|
||||
cursor = conn.execute(
|
||||
f"DELETE FROM runs WHERE run_id IN ({placeholders})",
|
||||
run_ids
|
||||
)
|
||||
|
||||
return cursor.rowcount
|
||||
@@ -0,0 +1,518 @@
|
||||
"""
|
||||
Queue Scheduler for PraisonAI.
|
||||
|
||||
Priority-based FIFO scheduler with concurrency limits.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from typing import Any, Callable, Dict, List, Optional, Set
|
||||
|
||||
from .models import QueuedRun, RunState, RunPriority, QueueConfig, QueueEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QueueFullError(Exception):
|
||||
"""Raised when queue is at capacity."""
|
||||
pass
|
||||
|
||||
|
||||
class RunNotFoundError(Exception):
|
||||
"""Raised when a run is not found."""
|
||||
pass
|
||||
|
||||
|
||||
class QueueScheduler:
|
||||
"""
|
||||
Priority-based FIFO scheduler with concurrency limits.
|
||||
|
||||
Runs are organized by priority (URGENT > HIGH > NORMAL > LOW),
|
||||
with FIFO ordering within each priority level.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[QueueConfig] = None):
|
||||
"""
|
||||
Initialize scheduler.
|
||||
|
||||
Args:
|
||||
config: Queue configuration. Uses defaults if not provided.
|
||||
"""
|
||||
self.config = config or QueueConfig()
|
||||
|
||||
# Priority queues: priority -> list of runs (FIFO order)
|
||||
self._queues: Dict[RunPriority, List[QueuedRun]] = {
|
||||
p: [] for p in RunPriority
|
||||
}
|
||||
|
||||
# Running runs: run_id -> run
|
||||
self._running: Dict[str, QueuedRun] = {}
|
||||
|
||||
# All runs by ID for quick lookup
|
||||
self._all_runs: Dict[str, QueuedRun] = {}
|
||||
|
||||
# Lock for thread safety
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# Event callbacks
|
||||
self._event_callbacks: List[Callable[[QueueEvent], None]] = []
|
||||
|
||||
# Cancellation tokens
|
||||
self._cancel_tokens: Set[str] = set()
|
||||
|
||||
def add_event_callback(self, callback: Callable[[QueueEvent], None]) -> None:
|
||||
"""Add a callback for queue events."""
|
||||
self._event_callbacks.append(callback)
|
||||
|
||||
def remove_event_callback(self, callback: Callable[[QueueEvent], None]) -> None:
|
||||
"""Remove an event callback."""
|
||||
if callback in self._event_callbacks:
|
||||
self._event_callbacks.remove(callback)
|
||||
|
||||
def _emit_event(self, event: QueueEvent) -> None:
|
||||
"""Emit an event to all callbacks."""
|
||||
for callback in self._event_callbacks:
|
||||
try:
|
||||
callback(event)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in event callback: {e}")
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
run: QueuedRun,
|
||||
check_duplicate: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Submit a run to the queue.
|
||||
|
||||
Args:
|
||||
run: The run to submit.
|
||||
check_duplicate: If True, reject duplicate run_ids.
|
||||
|
||||
Returns:
|
||||
The run_id.
|
||||
|
||||
Raises:
|
||||
QueueFullError: If queue is at capacity.
|
||||
ValueError: If run_id already exists (when check_duplicate=True).
|
||||
"""
|
||||
async with self._lock:
|
||||
# Check for duplicates
|
||||
if check_duplicate and run.run_id in self._all_runs:
|
||||
raise ValueError(f"Run {run.run_id} already exists")
|
||||
|
||||
# Check queue capacity
|
||||
total_queued = sum(len(q) for q in self._queues.values())
|
||||
if total_queued >= self.config.max_queue_size:
|
||||
raise QueueFullError(
|
||||
f"Queue is full ({total_queued}/{self.config.max_queue_size})"
|
||||
)
|
||||
|
||||
# Ensure state is QUEUED
|
||||
run.state = RunState.QUEUED
|
||||
|
||||
# Add to appropriate priority queue
|
||||
self._queues[run.priority].append(run)
|
||||
self._all_runs[run.run_id] = run
|
||||
|
||||
logger.debug(f"Submitted run {run.run_id} with priority {run.priority.name}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_submitted",
|
||||
run_id=run.run_id,
|
||||
data={"priority": run.priority.name, "agent": run.agent_name}
|
||||
))
|
||||
|
||||
return run.run_id
|
||||
|
||||
async def next(self) -> Optional[QueuedRun]:
|
||||
"""
|
||||
Get the next run to execute, respecting priority and concurrency limits.
|
||||
|
||||
Returns:
|
||||
The next run, or None if no run is available.
|
||||
"""
|
||||
async with self._lock:
|
||||
if not self._can_start_new():
|
||||
return None
|
||||
|
||||
# Check priorities in order (URGENT first = highest value)
|
||||
for priority in reversed(list(RunPriority)):
|
||||
queue = self._queues[priority]
|
||||
|
||||
for i, run in enumerate(queue):
|
||||
if self._can_run(run):
|
||||
# Remove from queue
|
||||
queue.pop(i)
|
||||
|
||||
# Update state
|
||||
run.state = RunState.RUNNING
|
||||
run.started_at = time.time()
|
||||
|
||||
# Add to running
|
||||
self._running[run.run_id] = run
|
||||
|
||||
logger.debug(f"Starting run {run.run_id}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_started",
|
||||
run_id=run.run_id,
|
||||
data={"agent": run.agent_name}
|
||||
))
|
||||
|
||||
return run
|
||||
|
||||
return None
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
run_id: str,
|
||||
output: Optional[str] = None,
|
||||
metrics: Optional[Dict[str, Any]] = None,
|
||||
) -> Optional[QueuedRun]:
|
||||
"""
|
||||
Mark a run as completed successfully.
|
||||
|
||||
Args:
|
||||
run_id: The run ID.
|
||||
output: The output content.
|
||||
metrics: Optional metrics.
|
||||
|
||||
Returns:
|
||||
The completed run, or None if not found.
|
||||
"""
|
||||
async with self._lock:
|
||||
run = self._running.pop(run_id, None)
|
||||
if run is None:
|
||||
logger.warning(f"Run {run_id} not found in running")
|
||||
return None
|
||||
|
||||
run.state = RunState.SUCCEEDED
|
||||
run.ended_at = time.time()
|
||||
run.output_content = output
|
||||
if metrics:
|
||||
run.metrics.update(metrics)
|
||||
|
||||
logger.debug(f"Completed run {run_id}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_completed",
|
||||
run_id=run_id,
|
||||
data={"duration": run.duration_seconds}
|
||||
))
|
||||
|
||||
return run
|
||||
|
||||
async def fail(
|
||||
self,
|
||||
run_id: str,
|
||||
error: str,
|
||||
metrics: Optional[Dict[str, Any]] = None,
|
||||
) -> Optional[QueuedRun]:
|
||||
"""
|
||||
Mark a run as failed.
|
||||
|
||||
Args:
|
||||
run_id: The run ID.
|
||||
error: The error message.
|
||||
metrics: Optional metrics.
|
||||
|
||||
Returns:
|
||||
The failed run, or None if not found.
|
||||
"""
|
||||
async with self._lock:
|
||||
run = self._running.pop(run_id, None)
|
||||
if run is None:
|
||||
logger.warning(f"Run {run_id} not found in running")
|
||||
return None
|
||||
|
||||
run.state = RunState.FAILED
|
||||
run.ended_at = time.time()
|
||||
run.error = error
|
||||
if metrics:
|
||||
run.metrics.update(metrics)
|
||||
|
||||
logger.debug(f"Failed run {run_id}: {error}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_failed",
|
||||
run_id=run_id,
|
||||
data={"error": error}
|
||||
))
|
||||
|
||||
return run
|
||||
|
||||
async def cancel(self, run_id: str) -> bool:
|
||||
"""
|
||||
Cancel a queued or running run.
|
||||
|
||||
Args:
|
||||
run_id: The run ID to cancel.
|
||||
|
||||
Returns:
|
||||
True if cancelled, False if not found.
|
||||
"""
|
||||
async with self._lock:
|
||||
# Check if running
|
||||
if run_id in self._running:
|
||||
run = self._running.pop(run_id)
|
||||
run.state = RunState.CANCELLED
|
||||
run.ended_at = time.time()
|
||||
self._cancel_tokens.add(run_id)
|
||||
|
||||
logger.debug(f"Cancelled running run {run_id}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_cancelled",
|
||||
run_id=run_id,
|
||||
data={"was_running": True}
|
||||
))
|
||||
|
||||
return True
|
||||
|
||||
# Check if queued
|
||||
for priority in RunPriority:
|
||||
queue = self._queues[priority]
|
||||
for i, run in enumerate(queue):
|
||||
if run.run_id == run_id:
|
||||
queue.pop(i)
|
||||
run.state = RunState.CANCELLED
|
||||
run.ended_at = time.time()
|
||||
|
||||
logger.debug(f"Cancelled queued run {run_id}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_cancelled",
|
||||
run_id=run_id,
|
||||
data={"was_running": False}
|
||||
))
|
||||
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def is_cancelled(self, run_id: str) -> bool:
|
||||
"""Check if a run has been cancelled."""
|
||||
return run_id in self._cancel_tokens
|
||||
|
||||
def clear_cancel_token(self, run_id: str) -> None:
|
||||
"""Clear a cancellation token."""
|
||||
self._cancel_tokens.discard(run_id)
|
||||
|
||||
async def update_input(self, run_id: str, new_input: str) -> bool:
|
||||
"""
|
||||
Update the input content of a queued run.
|
||||
|
||||
Only works for runs in QUEUED state.
|
||||
|
||||
Args:
|
||||
run_id: The run ID to update.
|
||||
new_input: The new input content.
|
||||
|
||||
Returns:
|
||||
True if updated, False if not found or not editable.
|
||||
"""
|
||||
async with self._lock:
|
||||
run = self._all_runs.get(run_id)
|
||||
if run is None:
|
||||
return False
|
||||
|
||||
if run.state != RunState.QUEUED:
|
||||
logger.warning(f"Cannot edit run {run_id} in state {run.state}")
|
||||
return False
|
||||
|
||||
run.input_content = new_input
|
||||
|
||||
logger.debug(f"Updated input for run {run_id}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_updated",
|
||||
run_id=run_id,
|
||||
data={"new_input": new_input[:50]}
|
||||
))
|
||||
|
||||
return True
|
||||
|
||||
async def retry(self, run_id: str) -> Optional[str]:
|
||||
"""
|
||||
Retry a failed run.
|
||||
|
||||
Creates a new run with incremented retry count and link to parent.
|
||||
|
||||
Args:
|
||||
run_id: The run ID to retry.
|
||||
|
||||
Returns:
|
||||
The new run_id, or None if retry not allowed.
|
||||
"""
|
||||
async with self._lock:
|
||||
# Find the original run
|
||||
original = self._all_runs.get(run_id)
|
||||
if original is None:
|
||||
logger.warning(f"Run {run_id} not found for retry")
|
||||
return None
|
||||
|
||||
if not original.can_retry():
|
||||
logger.warning(
|
||||
f"Run {run_id} cannot be retried "
|
||||
f"(state={original.state}, retries={original.retry_count}/{original.max_retries})"
|
||||
)
|
||||
return None
|
||||
|
||||
# Create new run
|
||||
new_run = QueuedRun(
|
||||
run_id=str(uuid.uuid4())[:8],
|
||||
agent_name=original.agent_name,
|
||||
input_content=original.input_content,
|
||||
state=RunState.QUEUED,
|
||||
priority=original.priority,
|
||||
session_id=original.session_id,
|
||||
trace_id=original.trace_id,
|
||||
workspace=original.workspace,
|
||||
user_id=original.user_id,
|
||||
retry_count=original.retry_count + 1,
|
||||
max_retries=original.max_retries,
|
||||
parent_run_id=original.run_id,
|
||||
config=original.config.copy(),
|
||||
)
|
||||
|
||||
# Submit outside lock to avoid deadlock
|
||||
await self.submit(new_run, check_duplicate=True)
|
||||
|
||||
logger.debug(f"Retrying run {run_id} as {new_run.run_id}")
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_retried",
|
||||
run_id=new_run.run_id,
|
||||
data={"parent_run_id": run_id, "retry_count": new_run.retry_count}
|
||||
))
|
||||
|
||||
return new_run.run_id
|
||||
|
||||
async def pause(self, run_id: str) -> bool:
|
||||
"""
|
||||
Pause a running run.
|
||||
|
||||
Note: The worker must check for pause state and handle accordingly.
|
||||
"""
|
||||
async with self._lock:
|
||||
if run_id not in self._running:
|
||||
return False
|
||||
|
||||
run = self._running[run_id]
|
||||
run.state = RunState.PAUSED
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_paused",
|
||||
run_id=run_id,
|
||||
))
|
||||
|
||||
return True
|
||||
|
||||
async def resume(self, run_id: str) -> bool:
|
||||
"""Resume a paused run."""
|
||||
async with self._lock:
|
||||
if run_id not in self._running:
|
||||
return False
|
||||
|
||||
run = self._running[run_id]
|
||||
if run.state != RunState.PAUSED:
|
||||
return False
|
||||
|
||||
run.state = RunState.RUNNING
|
||||
|
||||
self._emit_event(QueueEvent(
|
||||
event_type="run_resumed",
|
||||
run_id=run_id,
|
||||
))
|
||||
|
||||
return True
|
||||
|
||||
def get_run(self, run_id: str) -> Optional[QueuedRun]:
|
||||
"""Get a run by ID."""
|
||||
return self._all_runs.get(run_id)
|
||||
|
||||
def get_queued(self) -> List[QueuedRun]:
|
||||
"""Get all queued runs in priority order."""
|
||||
result = []
|
||||
for priority in reversed(list(RunPriority)):
|
||||
result.extend(self._queues[priority])
|
||||
return result
|
||||
|
||||
def get_running(self) -> List[QueuedRun]:
|
||||
"""Get all running runs."""
|
||||
return list(self._running.values())
|
||||
|
||||
def get_all(self) -> List[QueuedRun]:
|
||||
"""Get all runs."""
|
||||
return list(self._all_runs.values())
|
||||
|
||||
@property
|
||||
def queued_count(self) -> int:
|
||||
"""Number of queued runs."""
|
||||
return sum(len(q) for q in self._queues.values())
|
||||
|
||||
@property
|
||||
def running_count(self) -> int:
|
||||
"""Number of running runs."""
|
||||
return len(self._running)
|
||||
|
||||
def _can_start_new(self) -> bool:
|
||||
"""Check if we can start a new run (global limit)."""
|
||||
return len(self._running) < self.config.max_concurrent_global
|
||||
|
||||
def _can_run(self, run: QueuedRun) -> bool:
|
||||
"""Check if a specific run can start (per-agent and per-workspace limits)."""
|
||||
# Count running by agent
|
||||
agent_count = sum(
|
||||
1 for r in self._running.values()
|
||||
if r.agent_name == run.agent_name
|
||||
)
|
||||
if agent_count >= self.config.max_concurrent_per_agent:
|
||||
return False
|
||||
|
||||
# Count running by workspace
|
||||
if run.workspace:
|
||||
workspace_count = sum(
|
||||
1 for r in self._running.values()
|
||||
if r.workspace == run.workspace
|
||||
)
|
||||
if workspace_count >= self.config.max_concurrent_per_workspace:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def clear_queue(self) -> int:
|
||||
"""Clear all queued (not running) runs."""
|
||||
async with self._lock:
|
||||
count = 0
|
||||
for priority in RunPriority:
|
||||
count += len(self._queues[priority])
|
||||
for run in self._queues[priority]:
|
||||
run.state = RunState.CANCELLED
|
||||
del self._all_runs[run.run_id]
|
||||
self._queues[priority] = []
|
||||
|
||||
return count
|
||||
|
||||
def load_runs(self, runs: List[QueuedRun]) -> None:
|
||||
"""
|
||||
Load runs (e.g., from persistence after restart).
|
||||
|
||||
Queued runs are added to queues, running runs are re-queued.
|
||||
"""
|
||||
for run in runs:
|
||||
if run.state == RunState.RUNNING:
|
||||
# Re-queue running runs (they were interrupted)
|
||||
run.state = RunState.QUEUED
|
||||
run.started_at = None
|
||||
|
||||
if run.state == RunState.QUEUED:
|
||||
self._queues[run.priority].append(run)
|
||||
|
||||
self._all_runs[run.run_id] = run
|
||||
|
||||
logger.info(f"Loaded {len(runs)} runs from persistence")
|
||||
@@ -0,0 +1,383 @@
|
||||
"""
|
||||
Worker Pool for PraisonAI Queue System.
|
||||
|
||||
Manages async workers that execute queued runs.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections import deque
|
||||
from typing import Any, Callable, Coroutine, Dict, List, Optional, Set
|
||||
|
||||
from .models import QueuedRun, RunState, StreamChunk, QueueEvent
|
||||
from .scheduler import QueueScheduler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StreamBuffer:
|
||||
"""Bounded buffer with backpressure for streaming output."""
|
||||
|
||||
def __init__(self, max_size: int = 1000, drop_strategy: str = "oldest"):
|
||||
"""
|
||||
Initialize stream buffer.
|
||||
|
||||
Args:
|
||||
max_size: Maximum number of chunks to buffer.
|
||||
drop_strategy: "oldest" to drop oldest, "newest" to reject new.
|
||||
"""
|
||||
self._buffer: deque = deque(maxlen=max_size)
|
||||
self._lock = asyncio.Lock()
|
||||
self._drop_strategy = drop_strategy
|
||||
self._dropped_count = 0
|
||||
self._max_size = max_size
|
||||
|
||||
async def push(self, chunk: StreamChunk) -> bool:
|
||||
"""
|
||||
Push a chunk to the buffer.
|
||||
|
||||
Returns:
|
||||
True if added, False if dropped.
|
||||
"""
|
||||
async with self._lock:
|
||||
if len(self._buffer) >= self._max_size:
|
||||
if self._drop_strategy == "oldest":
|
||||
self._buffer.popleft()
|
||||
self._dropped_count += 1
|
||||
elif self._drop_strategy == "newest":
|
||||
self._dropped_count += 1
|
||||
return False
|
||||
self._buffer.append(chunk)
|
||||
return True
|
||||
|
||||
async def drain(self, batch_size: int = 50) -> List[StreamChunk]:
|
||||
"""Drain up to batch_size chunks."""
|
||||
async with self._lock:
|
||||
result = []
|
||||
for _ in range(min(batch_size, len(self._buffer))):
|
||||
result.append(self._buffer.popleft())
|
||||
return result
|
||||
|
||||
async def drain_all(self) -> List[StreamChunk]:
|
||||
"""Drain all chunks."""
|
||||
async with self._lock:
|
||||
result = list(self._buffer)
|
||||
self._buffer.clear()
|
||||
return result
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
"""Current buffer size."""
|
||||
return len(self._buffer)
|
||||
|
||||
@property
|
||||
def dropped_count(self) -> int:
|
||||
"""Number of dropped chunks."""
|
||||
return self._dropped_count
|
||||
|
||||
|
||||
class WorkerPool:
|
||||
"""
|
||||
Manages worker tasks for executing queued runs.
|
||||
|
||||
Workers poll the scheduler for runs and execute them with streaming output.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scheduler: QueueScheduler,
|
||||
on_output: Optional[Callable[[str, str], Coroutine[Any, Any, None]]] = None,
|
||||
on_complete: Optional[Callable[[str, QueuedRun], Coroutine[Any, Any, None]]] = None,
|
||||
on_error: Optional[Callable[[str, Exception], Coroutine[Any, Any, None]]] = None,
|
||||
on_event: Optional[Callable[[QueueEvent], Coroutine[Any, Any, None]]] = None,
|
||||
max_workers: int = 4,
|
||||
poll_interval: float = 0.1,
|
||||
stream_buffer_size: int = 1000,
|
||||
):
|
||||
"""
|
||||
Initialize worker pool.
|
||||
|
||||
Args:
|
||||
scheduler: The queue scheduler.
|
||||
on_output: Callback for streaming output (run_id, chunk).
|
||||
on_complete: Callback when run completes (run_id, run).
|
||||
on_error: Callback when run fails (run_id, exception).
|
||||
on_event: Callback for queue events.
|
||||
max_workers: Maximum number of concurrent workers.
|
||||
poll_interval: Seconds between scheduler polls.
|
||||
stream_buffer_size: Size of stream buffer per run.
|
||||
"""
|
||||
self.scheduler = scheduler
|
||||
self.on_output = on_output
|
||||
self.on_complete = on_complete
|
||||
self.on_error = on_error
|
||||
self.on_event = on_event
|
||||
self.max_workers = max_workers
|
||||
self.poll_interval = poll_interval
|
||||
self.stream_buffer_size = stream_buffer_size
|
||||
|
||||
self._workers: List[asyncio.Task] = []
|
||||
self._running = False
|
||||
self._stream_buffers: Dict[str, StreamBuffer] = {}
|
||||
self._active_runs: Set[str] = set()
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the worker pool."""
|
||||
if self._running:
|
||||
return
|
||||
|
||||
self._running = True
|
||||
|
||||
for i in range(self.max_workers):
|
||||
task = asyncio.create_task(
|
||||
self._worker_loop(i),
|
||||
name=f"queue-worker-{i}"
|
||||
)
|
||||
self._workers.append(task)
|
||||
|
||||
logger.info(f"Started {self.max_workers} queue workers")
|
||||
|
||||
async def stop(self, timeout: float = 5.0) -> None:
|
||||
"""
|
||||
Stop all workers gracefully.
|
||||
|
||||
Args:
|
||||
timeout: Seconds to wait for workers to finish.
|
||||
"""
|
||||
self._running = False
|
||||
|
||||
if not self._workers:
|
||||
return
|
||||
|
||||
# Wait for workers to finish current work
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(*self._workers, return_exceptions=True),
|
||||
timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
# Cancel remaining workers
|
||||
for task in self._workers:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
|
||||
await asyncio.gather(*self._workers, return_exceptions=True)
|
||||
|
||||
self._workers.clear()
|
||||
logger.info("Stopped queue workers")
|
||||
|
||||
async def _worker_loop(self, worker_id: int) -> None:
|
||||
"""Main worker loop."""
|
||||
logger.debug(f"Worker {worker_id} started")
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
# Get next run from scheduler
|
||||
run = await self.scheduler.next()
|
||||
|
||||
if run is None:
|
||||
# No work available, wait and retry
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
continue
|
||||
|
||||
# Execute the run
|
||||
self._active_runs.add(run.run_id)
|
||||
try:
|
||||
await self._execute_run(run, worker_id)
|
||||
finally:
|
||||
self._active_runs.discard(run.run_id)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.debug(f"Worker {worker_id} cancelled")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Worker {worker_id} error: {e}", exc_info=True)
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
logger.debug(f"Worker {worker_id} stopped")
|
||||
|
||||
async def _execute_run(self, run: QueuedRun, worker_id: int) -> None:
|
||||
"""Execute a single run with streaming output."""
|
||||
logger.debug(f"Worker {worker_id} executing run {run.run_id}")
|
||||
|
||||
# Create stream buffer for this run
|
||||
buffer = StreamBuffer(max_size=self.stream_buffer_size)
|
||||
self._stream_buffers[run.run_id] = buffer
|
||||
|
||||
output_chunks: List[str] = []
|
||||
|
||||
try:
|
||||
# Import agent lazily to avoid circular imports
|
||||
from praisonaiagents import Agent
|
||||
from .manager import QueueManager
|
||||
|
||||
# Get agent configuration
|
||||
agent_config = run.config.get("agent_config", {})
|
||||
agent_name = run.agent_name or agent_config.get("name", "Assistant")
|
||||
|
||||
# Get tools from runtime registry (not from config, as functions can't be serialized)
|
||||
tools = QueueManager.get_tools_for_run(run.run_id)
|
||||
|
||||
# Create agent with session_id for history persistence
|
||||
# Note: Agent uses 'llm' parameter, not 'model'
|
||||
agent = Agent(
|
||||
name=agent_name,
|
||||
instructions=agent_config.get("instructions", "You are a helpful assistant."),
|
||||
llm=agent_config.get("model"),
|
||||
tools=tools,
|
||||
verbose=agent_config.get("verbose", False),
|
||||
session_id=agent_config.get("session_id") or run.session_id, # Use session for history
|
||||
)
|
||||
|
||||
# Inject chat history from session store for context continuity
|
||||
if run.chat_history:
|
||||
agent.chat_history = list(run.chat_history)
|
||||
logger.debug(f"Injected {len(run.chat_history)} messages from session history")
|
||||
|
||||
# Check for cancellation before starting
|
||||
if self.scheduler.is_cancelled(run.run_id):
|
||||
logger.debug(f"Run {run.run_id} was cancelled before execution")
|
||||
return
|
||||
|
||||
# Execute with streaming if available
|
||||
chunk_index = 0
|
||||
|
||||
# Try streaming first
|
||||
try:
|
||||
async for chunk in self._stream_agent(agent, run.input_content):
|
||||
# Check for cancellation
|
||||
if self.scheduler.is_cancelled(run.run_id):
|
||||
logger.debug(f"Run {run.run_id} cancelled during execution")
|
||||
self.scheduler.clear_cancel_token(run.run_id)
|
||||
return
|
||||
|
||||
# Check for pause
|
||||
while run.state == RunState.PAUSED:
|
||||
await asyncio.sleep(0.1)
|
||||
if self.scheduler.is_cancelled(run.run_id):
|
||||
return
|
||||
|
||||
output_chunks.append(chunk)
|
||||
|
||||
# Buffer the chunk
|
||||
stream_chunk = StreamChunk(
|
||||
run_id=run.run_id,
|
||||
content=chunk,
|
||||
chunk_index=chunk_index,
|
||||
)
|
||||
await buffer.push(stream_chunk)
|
||||
chunk_index += 1
|
||||
|
||||
# Emit output callback
|
||||
if self.on_output:
|
||||
try:
|
||||
await self.on_output(run.run_id, chunk)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in output callback: {e}")
|
||||
|
||||
except Exception as stream_error:
|
||||
# Fallback to non-streaming
|
||||
logger.debug(f"Streaming failed, falling back to sync: {stream_error}")
|
||||
result = agent.chat(run.input_content)
|
||||
output_chunks = [result]
|
||||
|
||||
if self.on_output:
|
||||
try:
|
||||
await self.on_output(run.run_id, result)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in output callback: {e}")
|
||||
|
||||
# Combine output
|
||||
full_output = "".join(output_chunks)
|
||||
|
||||
# Mark final chunk
|
||||
final_chunk = StreamChunk(
|
||||
run_id=run.run_id,
|
||||
content="",
|
||||
chunk_index=chunk_index,
|
||||
is_final=True,
|
||||
)
|
||||
await buffer.push(final_chunk)
|
||||
|
||||
# Complete the run
|
||||
completed_run = await self.scheduler.complete(
|
||||
run.run_id,
|
||||
output=full_output,
|
||||
metrics={
|
||||
"chunks": chunk_index,
|
||||
"output_length": len(full_output),
|
||||
}
|
||||
)
|
||||
|
||||
if completed_run and self.on_complete:
|
||||
try:
|
||||
await self.on_complete(run.run_id, completed_run)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in complete callback: {e}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# Worker was cancelled
|
||||
await self.scheduler.cancel(run.run_id)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Run {run.run_id} failed: {e}", exc_info=True)
|
||||
|
||||
# Mark as failed
|
||||
failed_run = await self.scheduler.fail(run.run_id, str(e))
|
||||
|
||||
if self.on_error:
|
||||
try:
|
||||
await self.on_error(run.run_id, e)
|
||||
except Exception as callback_error:
|
||||
logger.error(f"Error in error callback: {callback_error}")
|
||||
|
||||
finally:
|
||||
# Cleanup buffer
|
||||
self._stream_buffers.pop(run.run_id, None)
|
||||
|
||||
async def _stream_agent(self, agent, input_content: str):
|
||||
"""
|
||||
Stream output from agent.
|
||||
|
||||
Yields chunks of output text.
|
||||
"""
|
||||
# Try to use agent's async streaming if available
|
||||
if hasattr(agent, 'astream'):
|
||||
async for chunk in agent.astream(input_content):
|
||||
yield chunk
|
||||
elif hasattr(agent, 'stream'):
|
||||
# Wrap sync stream in async
|
||||
for chunk in agent.stream(input_content):
|
||||
yield chunk
|
||||
await asyncio.sleep(0) # Yield control
|
||||
else:
|
||||
# Fallback: get full response and yield in chunks
|
||||
result = agent.chat(input_content)
|
||||
|
||||
# Simulate streaming by yielding words
|
||||
words = result.split()
|
||||
chunk_size = 3 # Words per chunk
|
||||
|
||||
for i in range(0, len(words), chunk_size):
|
||||
chunk_words = words[i:i + chunk_size]
|
||||
chunk = " ".join(chunk_words)
|
||||
if i + chunk_size < len(words):
|
||||
chunk += " "
|
||||
yield chunk
|
||||
await asyncio.sleep(0.05) # Small delay for streaming effect
|
||||
|
||||
def get_stream_buffer(self, run_id: str) -> Optional[StreamBuffer]:
|
||||
"""Get the stream buffer for a run."""
|
||||
return self._stream_buffers.get(run_id)
|
||||
|
||||
@property
|
||||
def active_run_count(self) -> int:
|
||||
"""Number of currently executing runs."""
|
||||
return len(self._active_runs)
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Whether the worker pool is running."""
|
||||
return self._running
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
TUI (Terminal User Interface) for PraisonAI.
|
||||
|
||||
Provides an app-like interactive experience with:
|
||||
- Event-loop driven async UI
|
||||
- Multi-pane layout (chat, tools, queue, status)
|
||||
- Streaming output
|
||||
- Queue management
|
||||
- Session persistence
|
||||
- Headless simulation mode for CI/testing
|
||||
|
||||
Requires the [tui] extra: pip install praisonai[tui]
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .app import TUIApp
|
||||
from .events import TUIEvent, TUIEventType
|
||||
from .orchestrator import TuiOrchestrator, UIStateModel, SimulationRunner
|
||||
from .mock_provider import MockProvider, MockProviderConfig
|
||||
|
||||
_lazy_cache = {}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load TUI components."""
|
||||
global _lazy_cache
|
||||
|
||||
if name in _lazy_cache:
|
||||
return _lazy_cache[name]
|
||||
|
||||
if name == "TUIApp":
|
||||
from .app import TUIApp
|
||||
_lazy_cache[name] = TUIApp
|
||||
return TUIApp
|
||||
elif name == "TUIEvent":
|
||||
from .events import TUIEvent
|
||||
_lazy_cache[name] = TUIEvent
|
||||
return TUIEvent
|
||||
elif name == "TUIEventType":
|
||||
from .events import TUIEventType
|
||||
_lazy_cache[name] = TUIEventType
|
||||
return TUIEventType
|
||||
elif name == "run_tui":
|
||||
from .app import run_tui
|
||||
_lazy_cache[name] = run_tui
|
||||
return run_tui
|
||||
elif name == "TuiOrchestrator":
|
||||
from .orchestrator import TuiOrchestrator
|
||||
_lazy_cache[name] = TuiOrchestrator
|
||||
return TuiOrchestrator
|
||||
elif name == "UIStateModel":
|
||||
from .orchestrator import UIStateModel
|
||||
_lazy_cache[name] = UIStateModel
|
||||
return UIStateModel
|
||||
elif name == "SimulationRunner":
|
||||
from .orchestrator import SimulationRunner
|
||||
_lazy_cache[name] = SimulationRunner
|
||||
return SimulationRunner
|
||||
elif name == "MockProvider":
|
||||
from .mock_provider import MockProvider
|
||||
_lazy_cache[name] = MockProvider
|
||||
return MockProvider
|
||||
elif name == "MockProviderConfig":
|
||||
from .mock_provider import MockProviderConfig
|
||||
_lazy_cache[name] = MockProviderConfig
|
||||
return MockProviderConfig
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"TUIApp",
|
||||
"TUIEvent",
|
||||
"TUIEventType",
|
||||
"run_tui",
|
||||
"TuiOrchestrator",
|
||||
"UIStateModel",
|
||||
"SimulationRunner",
|
||||
"MockProvider",
|
||||
"MockProviderConfig",
|
||||
]
|
||||
@@ -0,0 +1,812 @@
|
||||
"""
|
||||
Main TUI Application for PraisonAI.
|
||||
|
||||
Event-loop driven async TUI with multi-pane layout.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
try:
|
||||
from textual.app import App, ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual import work
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
App = object
|
||||
work = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
from .screens.main import MainScreen
|
||||
from .screens.queue import QueueScreen
|
||||
from .screens.settings import SettingsScreen
|
||||
from .screens.session import SessionScreen
|
||||
from .screens.help import HelpScreen
|
||||
from .widgets.chat import ChatMessage
|
||||
from .widgets.status import StatusInfo
|
||||
from .widgets.queue_panel import QueueItem
|
||||
from ..queue import QueueManager, QueueConfig, QueuedRun, RunPriority
|
||||
from .session_store import SessionStore, get_session_store
|
||||
|
||||
class TUIApp(App):
|
||||
"""
|
||||
PraisonAI Terminal User Interface Application.
|
||||
|
||||
Features:
|
||||
- Event-loop driven async UI
|
||||
- Multi-pane layout (chat, tools, queue, status)
|
||||
- Streaming output
|
||||
- Queue management
|
||||
- Session persistence
|
||||
"""
|
||||
|
||||
TITLE = "PraisonAI"
|
||||
SUB_TITLE = "AI Agent Terminal"
|
||||
|
||||
CSS = """
|
||||
Screen {
|
||||
background: $surface;
|
||||
}
|
||||
"""
|
||||
|
||||
# Safe default bindings - Ctrl keys only at app level for emergency exit
|
||||
# MainScreen handles all other bindings with safe defaults
|
||||
BINDINGS = [
|
||||
# Emergency quit - kept at app level for reliability
|
||||
Binding("ctrl+q", "quit", "Quit", show=False, priority=True),
|
||||
Binding("ctrl+c", "cancel", "Cancel", show=False),
|
||||
]
|
||||
|
||||
SCREENS = {
|
||||
"main": MainScreen,
|
||||
"queue": QueueScreen,
|
||||
"settings": SettingsScreen,
|
||||
"sessions": SessionScreen,
|
||||
"help": HelpScreen,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace: Optional[str] = None,
|
||||
session_id: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
agent_config: Optional[Dict[str, Any]] = None,
|
||||
queue_config: Optional[QueueConfig] = None,
|
||||
enable_acp: bool = True,
|
||||
enable_lsp: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.workspace = workspace or os.getcwd()
|
||||
self.session_id = session_id or str(uuid.uuid4())[:8]
|
||||
self.model = model or os.environ.get("OPENAI_MODEL", "gpt-4o-mini")
|
||||
self.agent_config = agent_config or {}
|
||||
self.enable_acp = enable_acp
|
||||
self.enable_lsp = enable_lsp
|
||||
|
||||
# Load default interactive tools (ACP + LSP) if not provided
|
||||
if "tools" not in self.agent_config or not self.agent_config["tools"]:
|
||||
self.agent_config["tools"] = self._load_default_tools()
|
||||
|
||||
# Queue manager
|
||||
self.queue_config = queue_config or QueueConfig()
|
||||
self.queue_manager: Optional[QueueManager] = None
|
||||
|
||||
# Session store for chat history persistence
|
||||
self.session_store = get_session_store()
|
||||
|
||||
# State
|
||||
self._current_run_id: Optional[str] = None
|
||||
self._streaming_content: str = ""
|
||||
self._total_tokens: int = 0
|
||||
self._total_cost: float = 0.0
|
||||
|
||||
# Slash command handlers
|
||||
self._command_handlers: Dict[str, callable] = {}
|
||||
self._setup_commands()
|
||||
|
||||
def _load_default_tools(self) -> list:
|
||||
"""
|
||||
Load default interactive tools (ACP + LSP) for TUI.
|
||||
|
||||
Uses the canonical interactive_tools provider.
|
||||
"""
|
||||
try:
|
||||
from ..interactive_tools import get_interactive_tools, ToolConfig
|
||||
|
||||
# Build disable list based on flags
|
||||
disable = []
|
||||
if not self.enable_acp:
|
||||
disable.append('acp')
|
||||
if not self.enable_lsp:
|
||||
disable.append('lsp')
|
||||
|
||||
config = ToolConfig(
|
||||
workspace=self.workspace,
|
||||
enable_acp=self.enable_acp,
|
||||
enable_lsp=self.enable_lsp,
|
||||
)
|
||||
|
||||
tools = get_interactive_tools(
|
||||
config=config,
|
||||
disable=disable if disable else None,
|
||||
)
|
||||
|
||||
logger.info(f"TUI loaded {len(tools)} default tools (ACP: {self.enable_acp}, LSP: {self.enable_lsp})")
|
||||
return tools
|
||||
|
||||
except ImportError as e:
|
||||
logger.warning(f"Could not load interactive tools: {e}")
|
||||
return []
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading default tools: {e}")
|
||||
return []
|
||||
|
||||
def _setup_commands(self) -> None:
|
||||
"""Setup slash command handlers."""
|
||||
self._command_handlers = {
|
||||
"help": self._cmd_help,
|
||||
"clear": self._cmd_clear,
|
||||
"model": self._cmd_model,
|
||||
"queue": self._cmd_queue,
|
||||
"cancel": self._cmd_cancel,
|
||||
"retry": self._cmd_retry,
|
||||
"settings": self._cmd_settings,
|
||||
"sessions": self._cmd_sessions,
|
||||
"cost": self._cmd_cost,
|
||||
"exit": self._cmd_exit,
|
||||
"quit": self._cmd_exit,
|
||||
# Additional commands
|
||||
"tokens": self._cmd_tokens,
|
||||
"plan": self._cmd_plan,
|
||||
"map": self._cmd_map,
|
||||
"undo": self._cmd_undo,
|
||||
"diff": self._cmd_diff,
|
||||
"commit": self._cmd_commit,
|
||||
"tools": self._cmd_tools,
|
||||
}
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
"""Handle app mount."""
|
||||
# Initialize queue manager with default tools (stored in runtime registry)
|
||||
self.queue_manager = QueueManager(
|
||||
config=self.queue_config,
|
||||
on_output=self._handle_output,
|
||||
on_complete=self._handle_complete,
|
||||
on_error=self._handle_error,
|
||||
default_tools=self.agent_config.get("tools", []),
|
||||
)
|
||||
|
||||
await self.queue_manager.start(recover=True)
|
||||
self.queue_manager.set_session(self.session_id)
|
||||
|
||||
# Push main screen
|
||||
await self.push_screen("main")
|
||||
|
||||
# Update tools panel with available tools
|
||||
self._update_tools_panel()
|
||||
|
||||
# Update status
|
||||
self._update_status()
|
||||
|
||||
async def on_unmount(self) -> None:
|
||||
"""Handle app unmount."""
|
||||
if self.queue_manager:
|
||||
await self.queue_manager.stop()
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
"""Compose the app."""
|
||||
yield from ()
|
||||
|
||||
# Event handlers
|
||||
|
||||
def on_main_screen_message_submitted(
|
||||
self, event: MainScreen.MessageSubmitted
|
||||
) -> None:
|
||||
"""Handle message submission - dispatch to worker for non-blocking execution."""
|
||||
# Use run_worker to avoid blocking the UI
|
||||
self.run_worker(
|
||||
self._process_message_submission(event.content),
|
||||
exclusive=False, # Allow multiple messages to queue
|
||||
exit_on_error=False,
|
||||
)
|
||||
|
||||
@work(exclusive=False, exit_on_error=False)
|
||||
async def _process_message_submission(self, content: str) -> None:
|
||||
"""Process message submission in background worker.
|
||||
|
||||
Using @work decorator ensures the UI remains responsive.
|
||||
User can type while agent is processing.
|
||||
"""
|
||||
# Add user message to chat display
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_user_message(content)
|
||||
main_screen.set_processing(True)
|
||||
|
||||
# Add user message to session store for history persistence
|
||||
self.session_store.add_user_message(self.session_id, content)
|
||||
|
||||
# Get chat history for context continuity
|
||||
chat_history = self.session_store.get_chat_history(self.session_id, max_messages=50)
|
||||
|
||||
# Submit to queue with chat history
|
||||
# Note: tools are NOT included in config (they can't be JSON serialized)
|
||||
# Instead, they are stored in QueueManager._tools_registry and retrieved by worker
|
||||
try:
|
||||
run_id = await self.queue_manager.submit(
|
||||
input_content=content,
|
||||
agent_name=self.agent_config.get("name", "Assistant"),
|
||||
session_id=self.session_id, # Pass session_id for history
|
||||
chat_history=chat_history, # Pass previous messages for context
|
||||
config={
|
||||
"agent_config": {
|
||||
"name": self.agent_config.get("name", "Assistant"),
|
||||
"instructions": self.agent_config.get(
|
||||
"instructions",
|
||||
"You are a helpful AI assistant."
|
||||
),
|
||||
"model": self.model,
|
||||
"verbose": False,
|
||||
"session_id": self.session_id, # Pass to agent for history
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
self._current_run_id = run_id
|
||||
self._streaming_content = ""
|
||||
|
||||
# IMMEDIATELY update queue panel - increment the moment message is queued
|
||||
self._update_queue_panel_live()
|
||||
|
||||
# Add streaming placeholder
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(
|
||||
content="",
|
||||
run_id=run_id,
|
||||
agent_name=self.agent_config.get("name", "Assistant"),
|
||||
streaming=True,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to submit: {e}")
|
||||
if isinstance(main_screen, MainScreen):
|
||||
main_screen.set_processing(False)
|
||||
await main_screen.add_assistant_message(
|
||||
content=f"Error: {e}",
|
||||
agent_name="System",
|
||||
)
|
||||
|
||||
self._update_status()
|
||||
self._update_queue_panel_live()
|
||||
|
||||
async def on_main_screen_command_executed(
|
||||
self, event: MainScreen.CommandExecuted
|
||||
) -> None:
|
||||
"""Handle command execution."""
|
||||
command = event.command.lower()
|
||||
args = event.args
|
||||
|
||||
handler = self._command_handlers.get(command)
|
||||
if handler:
|
||||
await handler(args)
|
||||
else:
|
||||
# Unknown command
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(
|
||||
content=f"Unknown command: /{command}\nType /help for available commands.",
|
||||
agent_name="System",
|
||||
)
|
||||
|
||||
async def on_main_screen_cancel_requested(
|
||||
self, event: MainScreen.CancelRequested
|
||||
) -> None:
|
||||
"""Handle cancel request."""
|
||||
if self._current_run_id:
|
||||
await self.queue_manager.cancel(self._current_run_id)
|
||||
self._current_run_id = None
|
||||
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
main_screen.set_processing(False)
|
||||
await main_screen.add_assistant_message(
|
||||
content="Operation cancelled.",
|
||||
agent_name="System",
|
||||
)
|
||||
|
||||
self._update_status()
|
||||
|
||||
async def on_queue_screen_run_cancelled(
|
||||
self, event: QueueScreen.RunCancelled
|
||||
) -> None:
|
||||
"""Handle run cancellation from queue screen."""
|
||||
await self.queue_manager.cancel(event.run_id)
|
||||
self._update_queue_screen()
|
||||
|
||||
async def on_queue_screen_run_retried(
|
||||
self, event: QueueScreen.RunRetried
|
||||
) -> None:
|
||||
"""Handle run retry from queue screen."""
|
||||
await self.queue_manager.retry(event.run_id)
|
||||
self._update_queue_screen()
|
||||
|
||||
async def on_queue_screen_queue_cleared(
|
||||
self, event: QueueScreen.QueueCleared
|
||||
) -> None:
|
||||
"""Handle queue clear."""
|
||||
await self.queue_manager.clear_queue()
|
||||
self._update_queue_screen()
|
||||
|
||||
async def on_settings_screen_settings_saved(
|
||||
self, event: SettingsScreen.SettingsSaved
|
||||
) -> None:
|
||||
"""Handle settings save."""
|
||||
settings = event.settings
|
||||
|
||||
if "model" in settings:
|
||||
self.model = settings["model"]
|
||||
|
||||
if "max_concurrent" in settings:
|
||||
self.queue_config.max_concurrent_global = settings["max_concurrent"]
|
||||
|
||||
if "autosave_interval" in settings:
|
||||
self.queue_config.autosave_interval_seconds = settings["autosave_interval"]
|
||||
|
||||
self._update_status()
|
||||
|
||||
async def on_session_screen_session_selected(
|
||||
self, event: SessionScreen.SessionSelected
|
||||
) -> None:
|
||||
"""Handle session selection."""
|
||||
self.session_id = event.session_id
|
||||
self.queue_manager.set_session(self.session_id)
|
||||
self._update_status()
|
||||
|
||||
async def on_queue_panel_widget_edit_requested(
|
||||
self, event
|
||||
) -> None:
|
||||
"""Handle queue item edit request."""
|
||||
from .widgets.queue_panel import QueuePanelWidget
|
||||
if isinstance(event, QueuePanelWidget.EditRequested):
|
||||
success = await self.queue_manager.update_input(event.run_id, event.new_content)
|
||||
if success:
|
||||
self._update_queue_panel_live()
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(
|
||||
content=f"Updated queue item {event.run_id[:8]}",
|
||||
agent_name="System",
|
||||
)
|
||||
|
||||
# Queue callbacks
|
||||
|
||||
async def _handle_output(self, run_id: str, chunk: str) -> None:
|
||||
"""Handle streaming output."""
|
||||
if run_id != self._current_run_id:
|
||||
return
|
||||
|
||||
self._streaming_content += chunk
|
||||
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.update_streaming(run_id, self._streaming_content)
|
||||
|
||||
async def _handle_complete(self, run_id: str, run: QueuedRun) -> None:
|
||||
"""Handle run completion."""
|
||||
output_content = run.output_content or self._streaming_content
|
||||
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.complete_streaming(run_id, output_content)
|
||||
main_screen.set_processing(False)
|
||||
|
||||
# Store assistant response in session store for history persistence
|
||||
if output_content and run.session_id:
|
||||
self.session_store.add_assistant_message(run.session_id, output_content, run_id)
|
||||
|
||||
# Update metrics
|
||||
if run.metrics:
|
||||
self._total_tokens += run.metrics.get("tokens", 0)
|
||||
self._total_cost += run.metrics.get("cost", 0.0)
|
||||
|
||||
if run_id == self._current_run_id:
|
||||
self._current_run_id = None
|
||||
self._streaming_content = ""
|
||||
|
||||
self._update_status()
|
||||
self._update_queue_panel_live()
|
||||
|
||||
async def _handle_error(self, run_id: str, error: Exception) -> None:
|
||||
"""Handle run error."""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.complete_streaming(run_id, f"Error: {error}")
|
||||
main_screen.set_processing(False)
|
||||
|
||||
if run_id == self._current_run_id:
|
||||
self._current_run_id = None
|
||||
self._streaming_content = ""
|
||||
|
||||
self._update_status()
|
||||
self._update_queue_panel_live()
|
||||
|
||||
# Command handlers
|
||||
|
||||
async def _cmd_help(self, args: str) -> None:
|
||||
"""Show help."""
|
||||
help_text = """
|
||||
**Available Commands:**
|
||||
|
||||
- `/help` - Show this help
|
||||
- `/clear` - Clear chat history
|
||||
- `/model [name]` - Show or set model
|
||||
- `/queue` - Show queue screen
|
||||
- `/cancel [run_id]` - Cancel a run
|
||||
- `/retry [run_id]` - Retry a failed run
|
||||
- `/settings` - Show settings
|
||||
- `/sessions` - Browse sessions
|
||||
- `/cost` - Show cost summary
|
||||
- `/exit` or `/quit` - Exit TUI
|
||||
|
||||
**Keyboard Shortcuts (Safe Defaults):**
|
||||
|
||||
- `Enter` - Send message
|
||||
- `Shift+Enter` - New line
|
||||
- `q` - Quit (when not typing)
|
||||
- `:` - Command mode
|
||||
- `?` - Help
|
||||
- `Esc` - Cancel/close
|
||||
|
||||
**Command Mode:**
|
||||
|
||||
Press `:` then type a command:
|
||||
- `:quit` or `:q` - Quit
|
||||
- `:clear` or `:cl` - Clear chat
|
||||
- `:help` or `:h` - Help
|
||||
- `:tools` or `:t` - Toggle tools
|
||||
- `:queue` or `:qu` - Toggle queue
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(
|
||||
content=help_text,
|
||||
agent_name="Help",
|
||||
)
|
||||
|
||||
async def _cmd_clear(self, args: str) -> None:
|
||||
"""Clear chat."""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.query_one("#chat-widget").clear()
|
||||
|
||||
async def _cmd_model(self, args: str) -> None:
|
||||
"""Show or set model."""
|
||||
main_screen = self.screen
|
||||
if args:
|
||||
self.model = args.strip()
|
||||
msg = f"Model set to: {self.model}"
|
||||
else:
|
||||
msg = f"Current model: {self.model}"
|
||||
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
self._update_status()
|
||||
|
||||
async def _cmd_queue(self, args: str) -> None:
|
||||
"""Show queue screen."""
|
||||
await self.push_screen("queue")
|
||||
self._update_queue_screen()
|
||||
|
||||
async def _cmd_cancel(self, args: str) -> None:
|
||||
"""Cancel a run."""
|
||||
run_id = args.strip() if args else self._current_run_id
|
||||
if run_id:
|
||||
await self.queue_manager.cancel(run_id)
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(
|
||||
content=f"Cancelled run: {run_id}",
|
||||
agent_name="System",
|
||||
)
|
||||
self._update_status()
|
||||
|
||||
async def _cmd_retry(self, args: str) -> None:
|
||||
"""Retry a failed run."""
|
||||
run_id = args.strip()
|
||||
if run_id:
|
||||
new_id = await self.queue_manager.retry(run_id)
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
if new_id:
|
||||
await main_screen.add_assistant_message(
|
||||
content=f"Retrying run {run_id} as {new_id}",
|
||||
agent_name="System",
|
||||
)
|
||||
else:
|
||||
await main_screen.add_assistant_message(
|
||||
content=f"Cannot retry run: {run_id}",
|
||||
agent_name="System",
|
||||
)
|
||||
self._update_status()
|
||||
|
||||
async def _cmd_settings(self, args: str) -> None:
|
||||
"""Show settings."""
|
||||
await self.push_screen(SettingsScreen(current_settings={
|
||||
"model": self.model,
|
||||
"max_concurrent": self.queue_config.max_concurrent_global,
|
||||
"autosave_interval": int(self.queue_config.autosave_interval_seconds),
|
||||
}))
|
||||
|
||||
async def _cmd_sessions(self, args: str) -> None:
|
||||
"""Show sessions."""
|
||||
await self.push_screen("sessions")
|
||||
# Load and display sessions
|
||||
await self._update_sessions_screen()
|
||||
|
||||
async def _cmd_cost(self, args: str) -> None:
|
||||
"""Show cost summary."""
|
||||
msg = f"""
|
||||
**Cost Summary:**
|
||||
|
||||
- Total Tokens: {self._total_tokens:,}
|
||||
- Total Cost: ${self._total_cost:.4f}
|
||||
- Session: {self.session_id}
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
async def _cmd_exit(self, args: str) -> None:
|
||||
"""Exit TUI."""
|
||||
self.exit()
|
||||
|
||||
async def _cmd_tokens(self, args: str) -> None:
|
||||
"""Show token usage breakdown."""
|
||||
msg = f"""
|
||||
**Token Usage:**
|
||||
|
||||
- Total Tokens: {self._total_tokens:,}
|
||||
- Estimated Cost: ${self._total_cost:.4f}
|
||||
- Session: {self.session_id}
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
async def _cmd_plan(self, args: str) -> None:
|
||||
"""Show or create a plan."""
|
||||
msg = """
|
||||
**Plan Command:**
|
||||
|
||||
The /plan command is used to create or view task plans.
|
||||
Currently, plans are managed through the agent's task execution.
|
||||
|
||||
Usage:
|
||||
- `/plan` - View current plan
|
||||
- `/plan <task>` - Create a plan for a task
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
async def _cmd_map(self, args: str) -> None:
|
||||
"""Show repository map."""
|
||||
msg = """
|
||||
**Repository Map:**
|
||||
|
||||
The /map command shows a map of the repository structure.
|
||||
Use the file tools to explore the codebase.
|
||||
|
||||
Available tools:
|
||||
- `acp_list_files` - List files in a directory
|
||||
- `acp_read_file` - Read file contents
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
async def _cmd_undo(self, args: str) -> None:
|
||||
"""Undo last change."""
|
||||
msg = """
|
||||
**Undo:**
|
||||
|
||||
The /undo command reverts the last file change.
|
||||
Use git commands or the agent to manage file changes.
|
||||
|
||||
Tip: Ask the agent to "undo the last change" or use git.
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
async def _cmd_diff(self, args: str) -> None:
|
||||
"""Show diff of changes."""
|
||||
msg = """
|
||||
**Diff:**
|
||||
|
||||
The /diff command shows changes made to files.
|
||||
Use git diff or ask the agent to show changes.
|
||||
|
||||
Tip: Ask the agent "show me the diff" or use `git diff`.
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
async def _cmd_commit(self, args: str) -> None:
|
||||
"""Commit changes."""
|
||||
msg = """
|
||||
**Commit:**
|
||||
|
||||
The /commit command creates a git commit.
|
||||
Use git commands or ask the agent to commit changes.
|
||||
|
||||
Tip: Ask the agent to "commit the changes with message X".
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
async def _cmd_tools(self, args: str) -> None:
|
||||
"""Show available tools."""
|
||||
tools = self.agent_config.get("tools", [])
|
||||
tool_names = []
|
||||
for tool in tools:
|
||||
if hasattr(tool, '__name__'):
|
||||
tool_names.append(tool.__name__)
|
||||
elif hasattr(tool, 'name'):
|
||||
tool_names.append(tool.name)
|
||||
else:
|
||||
tool_names.append(str(tool)[:30])
|
||||
|
||||
msg = f"""
|
||||
**Available Tools ({len(tools)}):**
|
||||
|
||||
{chr(10).join(f'- {name}' for name in tool_names[:20])}
|
||||
{"..." if len(tool_names) > 20 else ""}
|
||||
"""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
await main_screen.add_assistant_message(content=msg, agent_name="System")
|
||||
|
||||
# Helper methods
|
||||
|
||||
def _update_status(self) -> None:
|
||||
"""Update status bar."""
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
main_screen.update_status(StatusInfo(
|
||||
session_id=self.session_id,
|
||||
model=self.model,
|
||||
total_tokens=self._total_tokens,
|
||||
total_cost=self._total_cost,
|
||||
queued_count=self.queue_manager.queued_count if self.queue_manager else 0,
|
||||
running_count=self.queue_manager.running_count if self.queue_manager else 0,
|
||||
is_processing=self._current_run_id is not None,
|
||||
))
|
||||
|
||||
def _update_queue_panel_live(self) -> None:
|
||||
"""Update queue panel on MainScreen with live counts."""
|
||||
from .widgets.queue_panel import QueuePanelWidget
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
try:
|
||||
queue_panel = main_screen.query_one("#queue-panel", QueuePanelWidget)
|
||||
# Update counts immediately
|
||||
queue_panel.queued_count = self.queue_manager.queued_count if self.queue_manager else 0
|
||||
queue_panel.running_count = self.queue_manager.running_count if self.queue_manager else 0
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _update_queue_screen(self) -> None:
|
||||
"""Update queue screen if visible."""
|
||||
if isinstance(self.screen, QueueScreen):
|
||||
runs = self.queue_manager.list_runs(limit=100)
|
||||
self.screen.update_runs([
|
||||
{
|
||||
"run_id": r.run_id,
|
||||
"agent_name": r.agent_name,
|
||||
"input": r.input_content,
|
||||
"state": r.state.value,
|
||||
"priority": r.priority.name.lower(),
|
||||
"wait_time": f"{r.wait_seconds:.1f}s" if r.wait_seconds else "-",
|
||||
"duration": f"{r.duration_seconds:.1f}s" if r.duration_seconds else "-",
|
||||
}
|
||||
for r in runs
|
||||
])
|
||||
|
||||
async def _update_sessions_screen(self) -> None:
|
||||
"""Update sessions screen with session list."""
|
||||
if isinstance(self.screen, SessionScreen):
|
||||
# Get sessions from queue manager or session store
|
||||
sessions = []
|
||||
if self.queue_manager:
|
||||
try:
|
||||
# Try to get sessions from the persistence layer
|
||||
from ..queue.persistence import QueuePersistence
|
||||
persistence = QueuePersistence()
|
||||
sessions = persistence.list_sessions(limit=50)
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not load sessions: {e}")
|
||||
# Fallback: show current session only
|
||||
sessions = [{
|
||||
"session_id": self.session_id,
|
||||
"created_at": "-",
|
||||
"updated_at": "-",
|
||||
"run_count": self.queue_manager.queued_count + self.queue_manager.running_count,
|
||||
}]
|
||||
|
||||
self.screen.update_sessions(sessions)
|
||||
|
||||
def _update_tools_panel(self) -> None:
|
||||
"""Update tools panel with available tools."""
|
||||
from .widgets.tool_panel import ToolPanelWidget
|
||||
main_screen = self.screen
|
||||
if isinstance(main_screen, MainScreen):
|
||||
try:
|
||||
tool_panel = main_screen.query_one("#tool-panel", ToolPanelWidget)
|
||||
tools = self.agent_config.get("tools", [])
|
||||
tool_panel.set_available_tools(tools)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def action_quit(self) -> None:
|
||||
"""Quit the application."""
|
||||
self.exit()
|
||||
|
||||
def action_cancel(self) -> None:
|
||||
"""Cancel current operation."""
|
||||
if self._current_run_id and self.queue_manager:
|
||||
asyncio.create_task(self.queue_manager.cancel(self._current_run_id))
|
||||
|
||||
else:
|
||||
class TUIApp:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
|
||||
|
||||
def run_tui(
|
||||
workspace: Optional[str] = None,
|
||||
session_id: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
agent_config: Optional[Dict[str, Any]] = None,
|
||||
enable_acp: bool = True,
|
||||
enable_lsp: bool = True,
|
||||
) -> None:
|
||||
"""
|
||||
Run the PraisonAI TUI.
|
||||
|
||||
Args:
|
||||
workspace: Working directory.
|
||||
session_id: Session ID to resume.
|
||||
model: Default model.
|
||||
agent_config: Agent configuration.
|
||||
enable_acp: Enable ACP tools (default: True).
|
||||
enable_lsp: Enable LSP tools (default: True).
|
||||
"""
|
||||
if not TEXTUAL_AVAILABLE:
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
|
||||
app = TUIApp(
|
||||
workspace=workspace,
|
||||
session_id=session_id,
|
||||
model=model,
|
||||
agent_config=agent_config,
|
||||
enable_acp=enable_acp,
|
||||
enable_lsp=enable_lsp,
|
||||
)
|
||||
app.run()
|
||||
@@ -0,0 +1,566 @@
|
||||
"""
|
||||
CLI Commands for PraisonAI TUI and Queue.
|
||||
|
||||
Provides Typer-based CLI commands for:
|
||||
- praison tui - Launch interactive TUI
|
||||
- praison run - Execute single task
|
||||
- praison queue ls/cancel/retry/clear
|
||||
- praison session ls/resume
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
try:
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from rich.panel import Panel
|
||||
TYPER_AVAILABLE = True
|
||||
except ImportError:
|
||||
TYPER_AVAILABLE = False
|
||||
typer = None
|
||||
|
||||
console = Console() if TYPER_AVAILABLE else None
|
||||
|
||||
|
||||
def create_tui_app():
|
||||
"""Create the TUI CLI app."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return None
|
||||
|
||||
app = typer.Typer(
|
||||
name="tui",
|
||||
help="PraisonAI TUI and Queue commands",
|
||||
no_args_is_help=True,
|
||||
)
|
||||
|
||||
@app.command("launch")
|
||||
def tui_launch(
|
||||
workspace: Optional[str] = typer.Option(
|
||||
None, "--workspace", "-w",
|
||||
help="Workspace directory"
|
||||
),
|
||||
session: Optional[str] = typer.Option(
|
||||
None, "--session", "-s",
|
||||
help="Resume session ID"
|
||||
),
|
||||
model: Optional[str] = typer.Option(
|
||||
None, "--model", "-m",
|
||||
help="Default model"
|
||||
),
|
||||
agent: Optional[str] = typer.Option(
|
||||
None, "--agent", "-a",
|
||||
help="Agent config file (YAML)"
|
||||
),
|
||||
no_acp: bool = typer.Option(
|
||||
False, "--no-acp",
|
||||
help="Disable ACP tools (file operations with plan/approve/apply)"
|
||||
),
|
||||
no_lsp: bool = typer.Option(
|
||||
False, "--no-lsp",
|
||||
help="Disable LSP tools (code intelligence: symbols, definitions)"
|
||||
),
|
||||
):
|
||||
"""Launch the interactive TUI with ACP + LSP tools enabled by default."""
|
||||
try:
|
||||
from .app import run_tui
|
||||
except ImportError:
|
||||
console.print(
|
||||
"[red]Textual is required for TUI. Install with:[/red]\n"
|
||||
" pip install praisonai[tui]"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Load agent config if provided
|
||||
agent_config = {}
|
||||
if agent:
|
||||
import yaml
|
||||
with open(agent, "r") as f:
|
||||
agent_config = yaml.safe_load(f)
|
||||
|
||||
run_tui(
|
||||
workspace=workspace,
|
||||
session_id=session,
|
||||
model=model,
|
||||
agent_config=agent_config,
|
||||
enable_acp=not no_acp,
|
||||
enable_lsp=not no_lsp,
|
||||
)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def create_queue_app():
|
||||
"""Create the queue CLI app."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return None
|
||||
|
||||
app = typer.Typer(
|
||||
name="queue",
|
||||
help="Queue management commands",
|
||||
no_args_is_help=True,
|
||||
)
|
||||
|
||||
@app.command("ls")
|
||||
def queue_ls(
|
||||
state: Optional[str] = typer.Option(
|
||||
None, "--state", "-s",
|
||||
help="Filter by state (queued, running, succeeded, failed, cancelled)"
|
||||
),
|
||||
session: Optional[str] = typer.Option(
|
||||
None, "--session",
|
||||
help="Filter by session ID"
|
||||
),
|
||||
limit: int = typer.Option(
|
||||
20, "--limit", "-n",
|
||||
help="Maximum number of results"
|
||||
),
|
||||
json_output: bool = typer.Option(
|
||||
False, "--json", "-j",
|
||||
help="Output as JSON"
|
||||
),
|
||||
):
|
||||
"""List queued runs."""
|
||||
from ..queue import QueuePersistence, RunState
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
# Parse state filter
|
||||
state_filter = None
|
||||
if state:
|
||||
try:
|
||||
state_filter = RunState(state.lower())
|
||||
except ValueError:
|
||||
console.print(f"[red]Invalid state: {state}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
runs = persistence.list_runs(
|
||||
state=state_filter,
|
||||
session_id=session,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
if json_output:
|
||||
output = [r.to_dict() for r in runs]
|
||||
console.print(json.dumps(output, indent=2, default=str))
|
||||
return
|
||||
|
||||
if not runs:
|
||||
console.print("[dim]No runs found.[/dim]")
|
||||
return
|
||||
|
||||
# Create table
|
||||
table = Table(title="Queue Runs")
|
||||
table.add_column("ID", style="cyan")
|
||||
table.add_column("Agent", style="green")
|
||||
table.add_column("Input", max_width=30)
|
||||
table.add_column("State", style="yellow")
|
||||
table.add_column("Priority")
|
||||
table.add_column("Wait", style="dim")
|
||||
table.add_column("Duration", style="dim")
|
||||
|
||||
for run in runs:
|
||||
state_style = {
|
||||
"queued": "yellow",
|
||||
"running": "green",
|
||||
"succeeded": "cyan",
|
||||
"failed": "red",
|
||||
"cancelled": "dim",
|
||||
}.get(run.state.value, "")
|
||||
|
||||
input_preview = run.input_content[:27] + "..." if len(run.input_content) > 30 else run.input_content
|
||||
|
||||
table.add_row(
|
||||
run.run_id[:8],
|
||||
run.agent_name,
|
||||
input_preview,
|
||||
f"[{state_style}]{run.state.value}[/{state_style}]",
|
||||
run.priority.name.lower(),
|
||||
f"{run.wait_seconds:.1f}s" if run.wait_seconds else "-",
|
||||
f"{run.duration_seconds:.1f}s" if run.duration_seconds else "-",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
persistence.close()
|
||||
|
||||
@app.command("cancel")
|
||||
def queue_cancel(
|
||||
run_id: str = typer.Argument(..., help="Run ID to cancel"),
|
||||
):
|
||||
"""Cancel a queued or running run."""
|
||||
from ..queue import QueuePersistence, RunState
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
# Find the run
|
||||
run = persistence.load_run(run_id)
|
||||
if not run:
|
||||
# Try partial match
|
||||
runs = persistence.list_runs(limit=100)
|
||||
matches = [r for r in runs if r.run_id.startswith(run_id)]
|
||||
if len(matches) == 1:
|
||||
run = matches[0]
|
||||
elif len(matches) > 1:
|
||||
console.print(f"[yellow]Multiple matches for '{run_id}':[/yellow]")
|
||||
for r in matches:
|
||||
console.print(f" {r.run_id}")
|
||||
raise typer.Exit(1)
|
||||
else:
|
||||
console.print(f"[red]Run not found: {run_id}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if run.state.is_terminal():
|
||||
console.print(f"[yellow]Run {run.run_id[:8]} is already {run.state.value}[/yellow]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Update state
|
||||
persistence.update_run_state(run.run_id, RunState.CANCELLED)
|
||||
console.print(f"[green]Cancelled run: {run.run_id[:8]}[/green]")
|
||||
persistence.close()
|
||||
|
||||
@app.command("retry")
|
||||
def queue_retry(
|
||||
run_id: str = typer.Argument(..., help="Run ID to retry"),
|
||||
):
|
||||
"""Retry a failed run."""
|
||||
from ..queue import QueuePersistence, QueuedRun, RunState
|
||||
import uuid
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
# Find the run
|
||||
run = persistence.load_run(run_id)
|
||||
if not run:
|
||||
runs = persistence.list_runs(limit=100)
|
||||
matches = [r for r in runs if r.run_id.startswith(run_id)]
|
||||
if len(matches) == 1:
|
||||
run = matches[0]
|
||||
else:
|
||||
console.print(f"[red]Run not found: {run_id}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if run.state != RunState.FAILED:
|
||||
console.print(f"[yellow]Can only retry failed runs (current: {run.state.value})[/yellow]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if run.retry_count >= run.max_retries:
|
||||
console.print(f"[yellow]Max retries reached ({run.retry_count}/{run.max_retries})[/yellow]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Create new run
|
||||
new_run = QueuedRun(
|
||||
run_id=str(uuid.uuid4())[:8],
|
||||
agent_name=run.agent_name,
|
||||
input_content=run.input_content,
|
||||
state=RunState.QUEUED,
|
||||
priority=run.priority,
|
||||
session_id=run.session_id,
|
||||
workspace=run.workspace,
|
||||
retry_count=run.retry_count + 1,
|
||||
max_retries=run.max_retries,
|
||||
parent_run_id=run.run_id,
|
||||
config=run.config,
|
||||
)
|
||||
|
||||
persistence.save_run(new_run)
|
||||
console.print(f"[green]Created retry run: {new_run.run_id} (from {run.run_id[:8]})[/green]")
|
||||
persistence.close()
|
||||
|
||||
@app.command("clear")
|
||||
def queue_clear(
|
||||
force: bool = typer.Option(
|
||||
False, "--force", "-f",
|
||||
help="Skip confirmation"
|
||||
),
|
||||
):
|
||||
"""Clear all queued runs."""
|
||||
from ..queue import QueuePersistence, RunState
|
||||
|
||||
if not force:
|
||||
confirm = typer.confirm("Clear all queued runs?")
|
||||
if not confirm:
|
||||
raise typer.Abort()
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
# Get queued runs
|
||||
runs = persistence.list_runs(state=RunState.QUEUED, limit=1000)
|
||||
|
||||
count = 0
|
||||
for run in runs:
|
||||
persistence.update_run_state(run.run_id, RunState.CANCELLED)
|
||||
count += 1
|
||||
|
||||
console.print(f"[green]Cleared {count} queued runs[/green]")
|
||||
persistence.close()
|
||||
|
||||
@app.command("stats")
|
||||
def queue_stats(
|
||||
session: Optional[str] = typer.Option(
|
||||
None, "--session",
|
||||
help="Filter by session ID"
|
||||
),
|
||||
):
|
||||
"""Show queue statistics."""
|
||||
from ..queue import QueuePersistence
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
stats = persistence.get_stats(session)
|
||||
|
||||
console.print(Panel(
|
||||
f"""
|
||||
[bold]Queue Statistics[/bold]
|
||||
|
||||
Queued: {stats.queued_count}
|
||||
Running: {stats.running_count}
|
||||
Succeeded: {stats.succeeded_count}
|
||||
Failed: {stats.failed_count}
|
||||
Cancelled: {stats.cancelled_count}
|
||||
|
||||
Total Runs: {stats.total_runs}
|
||||
Avg Wait: {stats.avg_wait_seconds:.1f}s
|
||||
Avg Duration: {stats.avg_duration_seconds:.1f}s
|
||||
""",
|
||||
title="Queue Stats"
|
||||
))
|
||||
persistence.close()
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def create_session_app():
|
||||
"""Create the session CLI app."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return None
|
||||
|
||||
app = typer.Typer(
|
||||
name="session",
|
||||
help="Session management commands",
|
||||
no_args_is_help=True,
|
||||
)
|
||||
|
||||
@app.command("ls")
|
||||
def session_ls(
|
||||
limit: int = typer.Option(
|
||||
20, "--limit", "-n",
|
||||
help="Maximum number of results"
|
||||
),
|
||||
):
|
||||
"""List recent sessions."""
|
||||
from ..queue import QueuePersistence
|
||||
from datetime import datetime
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
sessions = persistence.list_sessions(limit=limit)
|
||||
|
||||
if not sessions:
|
||||
console.print("[dim]No sessions found.[/dim]")
|
||||
return
|
||||
|
||||
table = Table(title="Sessions")
|
||||
table.add_column("ID", style="cyan")
|
||||
table.add_column("User")
|
||||
table.add_column("Created", style="dim")
|
||||
table.add_column("Updated", style="dim")
|
||||
|
||||
for session in sessions:
|
||||
created = datetime.fromtimestamp(session["created_at"]).strftime("%Y-%m-%d %H:%M")
|
||||
updated = datetime.fromtimestamp(session["updated_at"]).strftime("%Y-%m-%d %H:%M")
|
||||
|
||||
table.add_row(
|
||||
session["session_id"][:8],
|
||||
session.get("user_id", "-"),
|
||||
created,
|
||||
updated,
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
persistence.close()
|
||||
|
||||
@app.command("resume")
|
||||
def session_resume(
|
||||
session_id: str = typer.Argument(..., help="Session ID to resume"),
|
||||
):
|
||||
"""Resume a session in TUI."""
|
||||
try:
|
||||
from .app import run_tui
|
||||
except ImportError:
|
||||
console.print(
|
||||
"[red]Textual is required for TUI. Install with:[/red]\n"
|
||||
" pip install praisonai[tui]"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
run_tui(session_id=session_id)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def create_run_command():
|
||||
"""Create the run command."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return None
|
||||
|
||||
def run_command(
|
||||
prompt: str = typer.Argument(..., help="Prompt to execute"),
|
||||
agent: Optional[str] = typer.Option(
|
||||
None, "--agent", "-a",
|
||||
help="Agent config file (YAML)"
|
||||
),
|
||||
model: Optional[str] = typer.Option(
|
||||
None, "--model", "-m",
|
||||
help="Model to use"
|
||||
),
|
||||
stream: bool = typer.Option(
|
||||
True, "--stream/--no-stream",
|
||||
help="Stream output"
|
||||
),
|
||||
priority: str = typer.Option(
|
||||
"normal", "--priority", "-p",
|
||||
help="Run priority (low, normal, high, urgent)"
|
||||
),
|
||||
):
|
||||
"""Run a single agent task."""
|
||||
from ..queue import QueueManager, QueueConfig, RunPriority
|
||||
|
||||
# Parse priority
|
||||
try:
|
||||
run_priority = RunPriority.from_string(priority)
|
||||
except (ValueError, KeyError):
|
||||
console.print(f"[red]Invalid priority: {priority}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Load agent config
|
||||
agent_config = {}
|
||||
if agent:
|
||||
import yaml
|
||||
with open(agent, "r") as f:
|
||||
agent_config = yaml.safe_load(f)
|
||||
|
||||
# Set model
|
||||
if model:
|
||||
agent_config["model"] = model
|
||||
|
||||
async def execute():
|
||||
output_chunks = []
|
||||
|
||||
async def on_output(run_id: str, chunk: str):
|
||||
output_chunks.append(chunk)
|
||||
if stream:
|
||||
console.print(chunk, end="")
|
||||
|
||||
async def on_complete(run_id: str, run):
|
||||
if not stream:
|
||||
console.print(run.output_content or "".join(output_chunks))
|
||||
|
||||
async def on_error(run_id: str, error: Exception):
|
||||
console.print(f"\n[red]Error: {error}[/red]")
|
||||
|
||||
config = QueueConfig(enable_persistence=False)
|
||||
manager = QueueManager(
|
||||
config=config,
|
||||
on_output=on_output,
|
||||
on_complete=on_complete,
|
||||
on_error=on_error,
|
||||
)
|
||||
|
||||
await manager.start(recover=False)
|
||||
|
||||
try:
|
||||
run_id = await manager.submit(
|
||||
input_content=prompt,
|
||||
agent_name=agent_config.get("name", "Assistant"),
|
||||
priority=run_priority,
|
||||
config={"agent_config": agent_config},
|
||||
)
|
||||
|
||||
# Wait for completion
|
||||
while True:
|
||||
run = manager.get_run(run_id)
|
||||
if run and run.state.is_terminal():
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
if stream:
|
||||
console.print() # Newline after streaming
|
||||
|
||||
finally:
|
||||
await manager.stop()
|
||||
|
||||
asyncio.run(execute())
|
||||
|
||||
return run_command
|
||||
|
||||
|
||||
# Main CLI integration
|
||||
def register_tui_commands(app):
|
||||
"""Register TUI commands with the main CLI app."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return
|
||||
|
||||
tui_app = create_tui_app()
|
||||
queue_app = create_queue_app()
|
||||
session_app = create_session_app()
|
||||
run_cmd = create_run_command()
|
||||
|
||||
if tui_app:
|
||||
app.add_typer(tui_app, name="tui")
|
||||
if queue_app:
|
||||
app.add_typer(queue_app, name="queue")
|
||||
if session_app:
|
||||
app.add_typer(session_app, name="session")
|
||||
if run_cmd:
|
||||
app.command("run")(run_cmd)
|
||||
|
||||
|
||||
# Standalone entry point
|
||||
def main():
|
||||
"""Main entry point for TUI CLI."""
|
||||
if not TYPER_AVAILABLE:
|
||||
print("Typer is required. Install with: pip install typer")
|
||||
sys.exit(1)
|
||||
|
||||
app = typer.Typer(
|
||||
name="praison-tui",
|
||||
help="PraisonAI TUI and Queue CLI",
|
||||
)
|
||||
|
||||
register_tui_commands(app)
|
||||
|
||||
# Add direct tui command
|
||||
@app.callback(invoke_without_command=True)
|
||||
def default(
|
||||
ctx: typer.Context,
|
||||
workspace: Optional[str] = typer.Option(None, "--workspace", "-w"),
|
||||
session: Optional[str] = typer.Option(None, "--session", "-s"),
|
||||
model: Optional[str] = typer.Option(None, "--model", "-m"),
|
||||
):
|
||||
"""Launch TUI if no subcommand given."""
|
||||
if ctx.invoked_subcommand is None:
|
||||
try:
|
||||
from .app import run_tui
|
||||
run_tui(workspace=workspace, session_id=session, model=model)
|
||||
except ImportError:
|
||||
console.print(
|
||||
"[red]Textual is required for TUI. Install with:[/red]\n"
|
||||
" pip install praisonai[tui]"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
app()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,150 @@
|
||||
"""
|
||||
TUI Configuration for PraisonAI.
|
||||
|
||||
Handles keyboard shortcuts and display settings with safe defaults.
|
||||
"""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class TUIConfig:
|
||||
"""
|
||||
Configuration for TUI behavior.
|
||||
|
||||
Safe defaults:
|
||||
- Ctrl keys disabled by default (conflict with IDEs/terminals)
|
||||
- Function keys disabled by default (conflict with IDEs/OS)
|
||||
- Single-character shortcuts enabled by default
|
||||
- Command mode (:) enabled by default
|
||||
"""
|
||||
|
||||
# Keyboard configuration
|
||||
enable_ctrl_keys: bool = False # Ctrl+Q, Ctrl+L, etc.
|
||||
enable_fn_keys: bool = False # F1-F12
|
||||
enable_single_char: bool = True # q, l, etc. (when not in input)
|
||||
enable_command_mode: bool = True # :quit, :clear, etc.
|
||||
leader_key: str = ":" # Command mode prefix
|
||||
|
||||
# Display configuration
|
||||
auto_scroll: bool = True
|
||||
show_tools_panel: bool = True
|
||||
show_queue_panel: bool = True
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "TUIConfig":
|
||||
"""Create config from environment variables."""
|
||||
config = cls()
|
||||
|
||||
# Check env vars for enabling Ctrl keys (opt-in)
|
||||
if os.environ.get("PRAISONAI_TUI_CTRL_KEYS", "").lower() in ("1", "true", "yes"):
|
||||
config.enable_ctrl_keys = True
|
||||
|
||||
# Check env vars for enabling Function keys (opt-in)
|
||||
if os.environ.get("PRAISONAI_TUI_FN_KEYS", "").lower() in ("1", "true", "yes"):
|
||||
config.enable_fn_keys = True
|
||||
|
||||
# Check for disabling single-char shortcuts
|
||||
if os.environ.get("PRAISONAI_TUI_NO_SINGLE_CHAR", "").lower() in ("1", "true", "yes"):
|
||||
config.enable_single_char = False
|
||||
|
||||
# Check for custom leader key
|
||||
leader = os.environ.get("PRAISONAI_TUI_LEADER", "")
|
||||
if leader and len(leader) == 1:
|
||||
config.leader_key = leader
|
||||
|
||||
return config
|
||||
|
||||
@classmethod
|
||||
def from_cli_args(
|
||||
cls,
|
||||
enable_ctrl_keys: bool = False,
|
||||
enable_fn_keys: bool = False,
|
||||
leader_key: Optional[str] = None,
|
||||
) -> "TUIConfig":
|
||||
"""Create config from CLI arguments, merged with env."""
|
||||
config = cls.from_env()
|
||||
|
||||
# CLI args override env
|
||||
if enable_ctrl_keys:
|
||||
config.enable_ctrl_keys = True
|
||||
if enable_fn_keys:
|
||||
config.enable_fn_keys = True
|
||||
if leader_key and len(leader_key) == 1:
|
||||
config.leader_key = leader_key
|
||||
|
||||
return config
|
||||
|
||||
|
||||
# Command mode commands
|
||||
COMMANDS = {
|
||||
"quit": {"aliases": ["q", "exit"], "description": "Quit the TUI"},
|
||||
"clear": {"aliases": ["cl"], "description": "Clear the chat history"},
|
||||
"help": {"aliases": ["h", "?"], "description": "Show help"},
|
||||
"tools": {"aliases": ["t"], "description": "Toggle tools panel"},
|
||||
"queue": {"aliases": ["qu"], "description": "Toggle queue panel"},
|
||||
"settings": {"aliases": ["set"], "description": "Open settings"},
|
||||
"cancel": {"aliases": ["c"], "description": "Cancel current operation"},
|
||||
"sessions": {"aliases": ["sess"], "description": "Manage sessions"},
|
||||
}
|
||||
|
||||
|
||||
def get_command(input_str: str) -> Optional[str]:
|
||||
"""
|
||||
Parse command input and return the canonical command name.
|
||||
|
||||
Args:
|
||||
input_str: The command string (without leader key)
|
||||
|
||||
Returns:
|
||||
Canonical command name or None if not found
|
||||
"""
|
||||
input_str = input_str.strip().lower()
|
||||
|
||||
# Check direct match
|
||||
if input_str in COMMANDS:
|
||||
return input_str
|
||||
|
||||
# Check aliases
|
||||
for cmd, info in COMMANDS.items():
|
||||
if input_str in info.get("aliases", []):
|
||||
return cmd
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_help_text() -> str:
|
||||
"""Get formatted help text for all commands."""
|
||||
lines = [
|
||||
"PraisonAI TUI Commands",
|
||||
"=" * 40,
|
||||
"",
|
||||
"Command Mode (press : to enter):",
|
||||
]
|
||||
|
||||
for cmd, info in COMMANDS.items():
|
||||
aliases = ", ".join(info.get("aliases", []))
|
||||
desc = info.get("description", "")
|
||||
if aliases:
|
||||
lines.append(f" :{cmd} (or :{aliases}) - {desc}")
|
||||
else:
|
||||
lines.append(f" :{cmd} - {desc}")
|
||||
|
||||
lines.extend([
|
||||
"",
|
||||
"Single-key shortcuts (when not typing):",
|
||||
" q - Quit",
|
||||
" l - Clear screen",
|
||||
" ? - Help",
|
||||
" / - Search/filter",
|
||||
"",
|
||||
"Input shortcuts:",
|
||||
" Enter - Send message",
|
||||
" Shift+Enter - New line",
|
||||
" Escape - Cancel",
|
||||
" Up/Down - History navigation",
|
||||
])
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,511 @@
|
||||
"""
|
||||
Debug and Simulation CLI Commands for PraisonAI TUI.
|
||||
|
||||
Provides headless TUI simulation, tracing, and debugging capabilities.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
try:
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from rich.panel import Panel
|
||||
from rich.live import Live
|
||||
from rich.text import Text
|
||||
TYPER_AVAILABLE = True
|
||||
except ImportError:
|
||||
TYPER_AVAILABLE = False
|
||||
typer = None
|
||||
|
||||
console = Console() if TYPER_AVAILABLE else None
|
||||
|
||||
|
||||
def create_debug_app():
|
||||
"""Create the debug CLI app with simulation commands."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return None
|
||||
|
||||
app = typer.Typer(
|
||||
name="tui",
|
||||
help="TUI and simulation commands",
|
||||
no_args_is_help=True,
|
||||
)
|
||||
|
||||
@app.command("launch")
|
||||
def tui_launch(
|
||||
workspace: Optional[str] = typer.Option(None, "--workspace", "-w"),
|
||||
session: Optional[str] = typer.Option(None, "--session", "-s"),
|
||||
model: Optional[str] = typer.Option(None, "--model", "-m"),
|
||||
debug: bool = typer.Option(False, "--debug", "-d", help="Enable debug overlays"),
|
||||
log_jsonl: Optional[str] = typer.Option(None, "--log-jsonl", help="Write events to JSONL file"),
|
||||
profile: bool = typer.Option(False, "--profile", help="Enable performance profiling"),
|
||||
):
|
||||
"""Launch the interactive TUI."""
|
||||
try:
|
||||
from .app import run_tui
|
||||
except ImportError:
|
||||
console.print("[red]Textual is required for TUI. Install with:[/red]\n pip install praisonai[tui]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Set debug environment
|
||||
if debug:
|
||||
import os
|
||||
os.environ["PRAISONAI_TUI_DEBUG"] = "1"
|
||||
|
||||
if log_jsonl:
|
||||
import os
|
||||
os.environ["PRAISONAI_TUI_JSONL"] = log_jsonl
|
||||
|
||||
run_tui(workspace=workspace, session_id=session, model=model)
|
||||
|
||||
@app.command("simulate")
|
||||
def tui_simulate(
|
||||
script: str = typer.Argument(..., help="Path to simulation script (YAML/JSON)"),
|
||||
mock: bool = typer.Option(True, "--mock/--real-llm", help="Use mock provider"),
|
||||
pretty: bool = typer.Option(True, "--pretty/--jsonl", help="Output format"),
|
||||
assert_mode: bool = typer.Option(False, "--assert", help="Validate expected outcomes"),
|
||||
timeout: float = typer.Option(60.0, "--timeout", help="Max execution time"),
|
||||
):
|
||||
"""Run a headless TUI simulation script."""
|
||||
import os
|
||||
|
||||
# Safety gate for real LLM
|
||||
if not mock:
|
||||
if not os.environ.get("PRAISONAI_REAL_LLM"):
|
||||
console.print("[red]Real LLM mode requires PRAISONAI_REAL_LLM=1 environment variable[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Load script
|
||||
script_path = Path(script)
|
||||
if not script_path.exists():
|
||||
console.print(f"[red]Script not found: {script}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
with open(script_path) as f:
|
||||
if script_path.suffix in (".yaml", ".yml"):
|
||||
import yaml
|
||||
script_data = yaml.safe_load(f)
|
||||
else:
|
||||
script_data = json.load(f)
|
||||
|
||||
# Run simulation
|
||||
from .orchestrator import TuiOrchestrator, SimulationRunner, OutputMode
|
||||
from ..queue import QueueConfig
|
||||
|
||||
output_mode = OutputMode.PRETTY if pretty else OutputMode.JSONL
|
||||
|
||||
config = QueueConfig(enable_persistence=not mock)
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=output_mode,
|
||||
debug=True,
|
||||
)
|
||||
|
||||
runner = SimulationRunner(orchestrator, assert_mode=assert_mode)
|
||||
|
||||
async def run():
|
||||
return await asyncio.wait_for(
|
||||
runner.run_script(script_data),
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
try:
|
||||
success = asyncio.run(run())
|
||||
except asyncio.TimeoutError:
|
||||
console.print(f"[red]Simulation timed out after {timeout}s[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Print summary
|
||||
summary = runner.get_summary()
|
||||
if assert_mode:
|
||||
console.print(Panel(
|
||||
f"Passed: {summary['assertions_passed']}\n"
|
||||
f"Failed: {summary['assertions_failed']}\n"
|
||||
f"Errors: {len(summary['errors'])}",
|
||||
title="Simulation Results"
|
||||
))
|
||||
|
||||
if summary['errors']:
|
||||
for error in summary['errors']:
|
||||
console.print(f"[red] ✗ {error}[/red]")
|
||||
|
||||
if not success:
|
||||
raise typer.Exit(1)
|
||||
|
||||
@app.command("snapshot")
|
||||
def tui_snapshot(
|
||||
session: Optional[str] = typer.Option(None, "--session", "-s"),
|
||||
run_id: Optional[str] = typer.Option(None, "--run", "-r"),
|
||||
json_output: bool = typer.Option(False, "--json", "-j"),
|
||||
):
|
||||
"""Print a TUI-like snapshot of current state."""
|
||||
from ..queue import QueuePersistence, RunState
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
# Build snapshot
|
||||
snapshot = {
|
||||
"timestamp": time.time(),
|
||||
"session_id": session,
|
||||
}
|
||||
|
||||
# Get runs
|
||||
runs = persistence.list_runs(session_id=session, limit=100)
|
||||
|
||||
queued = [r for r in runs if r.state == RunState.QUEUED]
|
||||
running = [r for r in runs if r.state == RunState.RUNNING]
|
||||
recent = [r for r in runs if r.state.is_terminal()][-5:]
|
||||
|
||||
snapshot["queued_count"] = len(queued)
|
||||
snapshot["running_count"] = len(running)
|
||||
snapshot["recent_runs"] = [
|
||||
{
|
||||
"run_id": r.run_id[:8],
|
||||
"agent": r.agent_name,
|
||||
"state": r.state.value,
|
||||
"input": r.input_content[:50],
|
||||
}
|
||||
for r in recent
|
||||
]
|
||||
|
||||
# Get specific run if requested
|
||||
if run_id:
|
||||
run = persistence.load_run(run_id)
|
||||
if run:
|
||||
snapshot["run"] = run.to_dict()
|
||||
|
||||
# Get session info
|
||||
if session:
|
||||
session_data = persistence.load_session(session)
|
||||
if session_data:
|
||||
snapshot["session"] = session_data
|
||||
|
||||
persistence.close()
|
||||
|
||||
if json_output:
|
||||
console.print(json.dumps(snapshot, indent=2, default=str))
|
||||
else:
|
||||
# Pretty print
|
||||
console.print("─" * 60)
|
||||
console.print(f"◉ PraisonAI Snapshot │ {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
console.print("─" * 60)
|
||||
|
||||
console.print(f"\n[bold]Queue Status[/bold]")
|
||||
console.print(f" Queued: {snapshot['queued_count']} │ Running: {snapshot['running_count']}")
|
||||
|
||||
if snapshot["recent_runs"]:
|
||||
console.print(f"\n[bold]Recent Runs[/bold]")
|
||||
for r in snapshot["recent_runs"]:
|
||||
state_color = {
|
||||
"succeeded": "green",
|
||||
"failed": "red",
|
||||
"cancelled": "dim",
|
||||
}.get(r["state"], "yellow")
|
||||
console.print(f" [{state_color}]{r['run_id']}[/{state_color}] {r['agent']}: {r['input'][:30]}...")
|
||||
|
||||
if "run" in snapshot:
|
||||
console.print(f"\n[bold]Run Details[/bold]")
|
||||
run = snapshot["run"]
|
||||
console.print(f" ID: {run.get('run_id', 'N/A')}")
|
||||
console.print(f" State: {run.get('state', 'N/A')}")
|
||||
console.print(f" Agent: {run.get('agent_name', 'N/A')}")
|
||||
|
||||
console.print("─" * 60)
|
||||
|
||||
@app.command("trace")
|
||||
def tui_trace(
|
||||
id: str = typer.Argument(..., help="Session or run ID to trace"),
|
||||
follow: bool = typer.Option(False, "--follow", "-f", help="Follow new events"),
|
||||
limit: int = typer.Option(50, "--limit", "-n", help="Max events to show"),
|
||||
):
|
||||
"""Replay events from persistence like a timeline."""
|
||||
from ..queue import QueuePersistence
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
# Try to find as session or run
|
||||
session = persistence.load_session(id)
|
||||
run = persistence.load_run(id)
|
||||
|
||||
if not session and not run:
|
||||
# Try partial match
|
||||
runs = persistence.list_runs(limit=100)
|
||||
matches = [r for r in runs if r.run_id.startswith(id) or r.session_id.startswith(id)]
|
||||
if matches:
|
||||
run = matches[0]
|
||||
|
||||
if not session and not run:
|
||||
console.print(f"[red]No session or run found for: {id}[/red]")
|
||||
persistence.close()
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Get runs for timeline
|
||||
if session:
|
||||
runs = persistence.list_runs(session_id=id, limit=limit)
|
||||
else:
|
||||
runs = [run] if run else []
|
||||
|
||||
# Build timeline
|
||||
events = []
|
||||
for r in runs:
|
||||
events.append({
|
||||
"timestamp": r.created_at,
|
||||
"type": "run_created",
|
||||
"run_id": r.run_id[:8],
|
||||
"agent": r.agent_name,
|
||||
"input": r.input_content[:50],
|
||||
})
|
||||
if r.started_at:
|
||||
events.append({
|
||||
"timestamp": r.started_at,
|
||||
"type": "run_started",
|
||||
"run_id": r.run_id[:8],
|
||||
})
|
||||
if r.ended_at:
|
||||
events.append({
|
||||
"timestamp": r.ended_at,
|
||||
"type": f"run_{r.state.value}",
|
||||
"run_id": r.run_id[:8],
|
||||
})
|
||||
|
||||
# Sort by timestamp
|
||||
events.sort(key=lambda e: e["timestamp"])
|
||||
|
||||
persistence.close()
|
||||
|
||||
# Print timeline
|
||||
console.print(f"\n[bold]Event Timeline for {id[:8]}[/bold]\n")
|
||||
|
||||
for event in events[-limit:]:
|
||||
ts = time.strftime("%H:%M:%S", time.localtime(event["timestamp"]))
|
||||
event_type = event["type"]
|
||||
|
||||
type_colors = {
|
||||
"run_created": "cyan",
|
||||
"run_started": "yellow",
|
||||
"run_succeeded": "green",
|
||||
"run_failed": "red",
|
||||
"run_cancelled": "dim",
|
||||
}
|
||||
color = type_colors.get(event_type, "white")
|
||||
|
||||
line = f"[dim]{ts}[/dim] [{color}]{event_type:15}[/{color}]"
|
||||
if "run_id" in event:
|
||||
line += f" run={event['run_id']}"
|
||||
if "agent" in event:
|
||||
line += f" agent={event['agent']}"
|
||||
if "input" in event:
|
||||
line += f" \"{event['input']}...\""
|
||||
|
||||
console.print(line)
|
||||
|
||||
if follow:
|
||||
console.print("\n[dim]Following... (Ctrl+C to stop)[/dim]")
|
||||
# In a real implementation, this would poll for new events
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def create_queue_watch_command():
|
||||
"""Create the queue watch command."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return None
|
||||
|
||||
def queue_watch(
|
||||
state: Optional[str] = typer.Option(None, "--state", "-s"),
|
||||
agent: Optional[str] = typer.Option(None, "--agent", "-a"),
|
||||
run_id: Optional[str] = typer.Option(None, "--run", "-r"),
|
||||
jsonl: bool = typer.Option(False, "--jsonl", "-j"),
|
||||
interval: float = typer.Option(1.0, "--interval", "-i"),
|
||||
):
|
||||
"""Watch queue events in real time."""
|
||||
from ..queue import QueuePersistence, RunState
|
||||
|
||||
persistence = QueuePersistence()
|
||||
persistence.initialize()
|
||||
|
||||
last_seen = {}
|
||||
|
||||
try:
|
||||
while True:
|
||||
runs = persistence.list_runs(limit=100)
|
||||
|
||||
# Filter
|
||||
if state:
|
||||
try:
|
||||
state_filter = RunState(state.lower())
|
||||
runs = [r for r in runs if r.state == state_filter]
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if agent:
|
||||
runs = [r for r in runs if agent.lower() in r.agent_name.lower()]
|
||||
|
||||
if run_id:
|
||||
runs = [r for r in runs if r.run_id.startswith(run_id)]
|
||||
|
||||
# Check for changes
|
||||
for run in runs:
|
||||
key = run.run_id
|
||||
current_state = run.state.value
|
||||
|
||||
if key not in last_seen or last_seen[key] != current_state:
|
||||
last_seen[key] = current_state
|
||||
|
||||
if jsonl:
|
||||
event = {
|
||||
"timestamp": time.time(),
|
||||
"run_id": run.run_id,
|
||||
"agent": run.agent_name,
|
||||
"state": current_state,
|
||||
}
|
||||
print(json.dumps(event))
|
||||
else:
|
||||
ts = time.strftime("%H:%M:%S")
|
||||
state_colors = {
|
||||
"queued": "yellow",
|
||||
"running": "cyan",
|
||||
"succeeded": "green",
|
||||
"failed": "red",
|
||||
"cancelled": "dim",
|
||||
}
|
||||
color = state_colors.get(current_state, "white")
|
||||
console.print(
|
||||
f"[dim]{ts}[/dim] [{color}]{current_state:10}[/{color}] "
|
||||
f"{run.run_id[:8]} {run.agent_name}"
|
||||
)
|
||||
|
||||
time.sleep(interval)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
console.print("\n[dim]Stopped watching[/dim]")
|
||||
finally:
|
||||
persistence.close()
|
||||
|
||||
return queue_watch
|
||||
|
||||
|
||||
def create_doctor_command():
|
||||
"""Create the doctor command for TUI diagnostics."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return None
|
||||
|
||||
def doctor_tui():
|
||||
"""Validate TUI dependencies, DB schema, and wiring."""
|
||||
checks = []
|
||||
|
||||
# Check Textual
|
||||
try:
|
||||
import textual
|
||||
checks.append(("Textual", "✓", f"v{textual.__version__}", "green"))
|
||||
except ImportError:
|
||||
checks.append(("Textual", "✗", "Not installed", "red"))
|
||||
|
||||
# Check Rich
|
||||
try:
|
||||
import rich
|
||||
checks.append(("Rich", "✓", f"v{rich.__version__}", "green"))
|
||||
except ImportError:
|
||||
checks.append(("Rich", "✗", "Not installed", "red"))
|
||||
|
||||
# Check Typer
|
||||
try:
|
||||
import typer
|
||||
checks.append(("Typer", "✓", f"v{typer.__version__}", "green"))
|
||||
except ImportError:
|
||||
checks.append(("Typer", "✗", "Not installed", "red"))
|
||||
|
||||
# Check queue modules
|
||||
try:
|
||||
from ..queue import QueueManager, QueueScheduler, QueuePersistence
|
||||
checks.append(("Queue System", "✓", "All modules loaded", "green"))
|
||||
except ImportError as e:
|
||||
checks.append(("Queue System", "✗", str(e), "red"))
|
||||
|
||||
# Check TUI modules
|
||||
try:
|
||||
from .orchestrator import TuiOrchestrator
|
||||
checks.append(("TUI Orchestrator", "✓", "Loaded", "green"))
|
||||
except ImportError as e:
|
||||
checks.append(("TUI Orchestrator", "✗", str(e), "red"))
|
||||
|
||||
# Check events
|
||||
try:
|
||||
from .events import TUIEvent, TUIEventType
|
||||
event_count = len(TUIEventType)
|
||||
checks.append(("TUI Events", "✓", f"{event_count} event types", "green"))
|
||||
except ImportError as e:
|
||||
checks.append(("TUI Events", "✗", str(e), "red"))
|
||||
|
||||
# Check DB
|
||||
try:
|
||||
from ..queue import QueuePersistence
|
||||
p = QueuePersistence()
|
||||
p.initialize()
|
||||
stats = p.get_stats()
|
||||
p.close()
|
||||
checks.append(("Database", "✓", f"{stats.total_runs} runs in DB", "green"))
|
||||
except Exception as e:
|
||||
checks.append(("Database", "✗", str(e), "red"))
|
||||
|
||||
# Check slash commands
|
||||
try:
|
||||
from ..slash_commands import SlashCommandRegistry
|
||||
registry = SlashCommandRegistry()
|
||||
cmd_count = len(registry.commands)
|
||||
checks.append(("Slash Commands", "✓", f"{cmd_count} commands", "green"))
|
||||
except Exception as e:
|
||||
checks.append(("Slash Commands", "⚠", str(e), "yellow"))
|
||||
|
||||
# Print results
|
||||
console.print("\n[bold]PraisonAI TUI Diagnostics[/bold]\n")
|
||||
|
||||
table = Table(show_header=True)
|
||||
table.add_column("Component")
|
||||
table.add_column("Status")
|
||||
table.add_column("Details")
|
||||
|
||||
all_ok = True
|
||||
for name, status, details, color in checks:
|
||||
table.add_row(name, f"[{color}]{status}[/{color}]", details)
|
||||
if status == "✗":
|
||||
all_ok = False
|
||||
|
||||
console.print(table)
|
||||
|
||||
if all_ok:
|
||||
console.print("\n[green]All checks passed![/green]")
|
||||
else:
|
||||
console.print("\n[red]Some checks failed. Run 'pip install praisonai[tui]' to install missing dependencies.[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
return doctor_tui
|
||||
|
||||
|
||||
def register_debug_commands(app):
|
||||
"""Register debug commands with the main CLI app."""
|
||||
if not TYPER_AVAILABLE:
|
||||
return
|
||||
|
||||
debug_app = create_debug_app()
|
||||
watch_cmd = create_queue_watch_command()
|
||||
doctor_cmd = create_doctor_command()
|
||||
|
||||
if debug_app:
|
||||
app.add_typer(debug_app, name="tui")
|
||||
|
||||
if watch_cmd:
|
||||
# Add to queue subcommand
|
||||
pass # Will be added via queue app
|
||||
|
||||
if doctor_cmd:
|
||||
# Add doctor --tui option
|
||||
pass # Will be added via main app
|
||||
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
Custom events for PraisonAI TUI.
|
||||
|
||||
Defines event types for communication between TUI components.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional
|
||||
import time
|
||||
|
||||
|
||||
class TUIEventType(str, Enum):
|
||||
"""Types of TUI events."""
|
||||
# User input events
|
||||
MESSAGE_SUBMITTED = "message_submitted"
|
||||
COMMAND_EXECUTED = "command_executed"
|
||||
|
||||
# Queue events
|
||||
RUN_QUEUED = "run_queued"
|
||||
RUN_STARTED = "run_started"
|
||||
RUN_COMPLETED = "run_completed"
|
||||
RUN_FAILED = "run_failed"
|
||||
RUN_CANCELLED = "run_cancelled"
|
||||
|
||||
# Output events
|
||||
OUTPUT_CHUNK = "output_chunk"
|
||||
OUTPUT_COMPLETE = "output_complete"
|
||||
|
||||
# Tool events
|
||||
TOOL_CALL_STARTED = "tool_call_started"
|
||||
TOOL_CALL_COMPLETED = "tool_call_completed"
|
||||
TOOL_APPROVAL_REQUIRED = "tool_approval_required"
|
||||
|
||||
# Session events
|
||||
SESSION_STARTED = "session_started"
|
||||
SESSION_SAVED = "session_saved"
|
||||
SESSION_LOADED = "session_loaded"
|
||||
|
||||
# UI events
|
||||
SCREEN_CHANGED = "screen_changed"
|
||||
FOCUS_CHANGED = "focus_changed"
|
||||
ERROR_OCCURRED = "error_occurred"
|
||||
STATUS_UPDATED = "status_updated"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TUIEvent:
|
||||
"""A TUI event."""
|
||||
event_type: TUIEventType
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
data: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Optional identifiers
|
||||
run_id: Optional[str] = None
|
||||
session_id: Optional[str] = None
|
||||
agent_name: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def message_submitted(cls, content: str, **kwargs) -> "TUIEvent":
|
||||
"""Create a message submitted event."""
|
||||
return cls(
|
||||
event_type=TUIEventType.MESSAGE_SUBMITTED,
|
||||
data={"content": content, **kwargs}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def output_chunk(cls, run_id: str, content: str, **kwargs) -> "TUIEvent":
|
||||
"""Create an output chunk event."""
|
||||
return cls(
|
||||
event_type=TUIEventType.OUTPUT_CHUNK,
|
||||
run_id=run_id,
|
||||
data={"content": content, **kwargs}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def run_completed(cls, run_id: str, output: str, **kwargs) -> "TUIEvent":
|
||||
"""Create a run completed event."""
|
||||
return cls(
|
||||
event_type=TUIEventType.RUN_COMPLETED,
|
||||
run_id=run_id,
|
||||
data={"output": output, **kwargs}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def error(cls, message: str, **kwargs) -> "TUIEvent":
|
||||
"""Create an error event."""
|
||||
return cls(
|
||||
event_type=TUIEventType.ERROR_OCCURRED,
|
||||
data={"message": message, **kwargs}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def status_update(cls, status: str, **kwargs) -> "TUIEvent":
|
||||
"""Create a status update event."""
|
||||
return cls(
|
||||
event_type=TUIEventType.STATUS_UPDATED,
|
||||
data={"status": status, **kwargs}
|
||||
)
|
||||
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
Mock Provider for PraisonAI TUI Testing.
|
||||
|
||||
Provides deterministic, replayable responses for CI and headless testing.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator, Callable, Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockResponse:
|
||||
"""A mock response configuration."""
|
||||
content: str
|
||||
chunks: List[str] = field(default_factory=list)
|
||||
delay_per_chunk: float = 0.05
|
||||
tool_calls: List[Dict[str, Any]] = field(default_factory=list)
|
||||
error: Optional[str] = None
|
||||
tokens: int = 0
|
||||
cost: float = 0.0
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.chunks and self.content:
|
||||
# Split content into word chunks
|
||||
words = self.content.split()
|
||||
self.chunks = []
|
||||
for i in range(0, len(words), 3):
|
||||
chunk = " ".join(words[i:i+3])
|
||||
if i > 0:
|
||||
chunk = " " + chunk
|
||||
self.chunks.append(chunk)
|
||||
|
||||
if not self.tokens:
|
||||
self.tokens = len(self.content.split()) * 2 # Rough estimate
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockProviderConfig:
|
||||
"""Configuration for mock provider."""
|
||||
seed: int = 42
|
||||
default_delay: float = 0.05
|
||||
simulate_errors: bool = False
|
||||
error_rate: float = 0.1
|
||||
record_mode: bool = False
|
||||
playback_file: Optional[str] = None
|
||||
|
||||
# Canned responses by input pattern
|
||||
responses: Dict[str, MockResponse] = field(default_factory=dict)
|
||||
|
||||
# Default response if no pattern matches
|
||||
default_response: Optional[MockResponse] = None
|
||||
|
||||
|
||||
class MockProvider:
|
||||
"""
|
||||
Mock LLM provider for testing.
|
||||
|
||||
Features:
|
||||
- Deterministic responses based on seed
|
||||
- Streaming simulation
|
||||
- Tool call simulation
|
||||
- Error injection
|
||||
- Record/playback mode
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[MockProviderConfig] = None):
|
||||
self.config = config or MockProviderConfig()
|
||||
self._rng = random.Random(self.config.seed)
|
||||
self._recordings: List[Dict[str, Any]] = []
|
||||
self._playback_index = 0
|
||||
|
||||
# Load playback file if specified
|
||||
if self.config.playback_file:
|
||||
self._load_playback()
|
||||
|
||||
# Setup default responses
|
||||
self._setup_defaults()
|
||||
|
||||
def _setup_defaults(self) -> None:
|
||||
"""Setup default canned responses."""
|
||||
defaults = {
|
||||
"hello": MockResponse(
|
||||
content="Hello! I'm a mock AI assistant. How can I help you today?",
|
||||
tokens=20,
|
||||
cost=0.0001,
|
||||
),
|
||||
"help": MockResponse(
|
||||
content="I can help you with various tasks. Just ask me anything!",
|
||||
tokens=15,
|
||||
cost=0.00008,
|
||||
),
|
||||
"test": MockResponse(
|
||||
content="This is a test response. Everything is working correctly.",
|
||||
tokens=12,
|
||||
cost=0.00006,
|
||||
),
|
||||
"error": MockResponse(
|
||||
content="",
|
||||
error="Simulated error for testing",
|
||||
),
|
||||
"tool": MockResponse(
|
||||
content="I'll use a tool to help with that.",
|
||||
tool_calls=[{
|
||||
"id": "call_mock_001",
|
||||
"name": "mock_tool",
|
||||
"arguments": {"query": "test"},
|
||||
}],
|
||||
tokens=25,
|
||||
cost=0.00012,
|
||||
),
|
||||
}
|
||||
|
||||
for pattern, response in defaults.items():
|
||||
if pattern not in self.config.responses:
|
||||
self.config.responses[pattern] = response
|
||||
|
||||
if not self.config.default_response:
|
||||
self.config.default_response = MockResponse(
|
||||
content="I understand your request. Here's a mock response for testing purposes. "
|
||||
"This response is generated deterministically based on the input.",
|
||||
tokens=30,
|
||||
cost=0.00015,
|
||||
)
|
||||
|
||||
def _load_playback(self) -> None:
|
||||
"""Load recordings for playback."""
|
||||
if not self.config.playback_file:
|
||||
return
|
||||
|
||||
path = Path(self.config.playback_file)
|
||||
if path.exists():
|
||||
with open(path) as f:
|
||||
for line in f:
|
||||
self._recordings.append(json.loads(line))
|
||||
|
||||
def _save_recording(self, input_content: str, response: MockResponse) -> None:
|
||||
"""Save a recording."""
|
||||
if not self.config.record_mode or not self.config.playback_file:
|
||||
return
|
||||
|
||||
record = {
|
||||
"timestamp": time.time(),
|
||||
"input": input_content,
|
||||
"response": {
|
||||
"content": response.content,
|
||||
"chunks": response.chunks,
|
||||
"tool_calls": response.tool_calls,
|
||||
"tokens": response.tokens,
|
||||
"cost": response.cost,
|
||||
}
|
||||
}
|
||||
|
||||
path = Path(self.config.playback_file)
|
||||
with open(path, "a") as f:
|
||||
f.write(json.dumps(record) + "\n")
|
||||
|
||||
def _get_response(self, input_content: str) -> MockResponse:
|
||||
"""Get response for input."""
|
||||
# Check playback mode
|
||||
if self._recordings and self._playback_index < len(self._recordings):
|
||||
record = self._recordings[self._playback_index]
|
||||
self._playback_index += 1
|
||||
return MockResponse(**record["response"])
|
||||
|
||||
# Check for pattern match
|
||||
input_lower = input_content.lower()
|
||||
for pattern, response in self.config.responses.items():
|
||||
if pattern in input_lower:
|
||||
return response
|
||||
|
||||
# Simulate random error
|
||||
if self.config.simulate_errors and self._rng.random() < self.config.error_rate:
|
||||
return MockResponse(content="", error="Random simulated error")
|
||||
|
||||
# Generate deterministic response based on input hash
|
||||
input_hash = hashlib.md5(input_content.encode()).hexdigest()
|
||||
seed = int(input_hash[:8], 16)
|
||||
rng = random.Random(seed)
|
||||
|
||||
# Generate varied response
|
||||
templates = [
|
||||
"I've processed your request about '{topic}'. Here's what I found...",
|
||||
"Regarding '{topic}', I can provide the following information...",
|
||||
"Thank you for asking about '{topic}'. Let me explain...",
|
||||
"Based on your query about '{topic}', here's my response...",
|
||||
]
|
||||
|
||||
topic = input_content[:30] if len(input_content) > 30 else input_content
|
||||
template = rng.choice(templates)
|
||||
content = template.format(topic=topic)
|
||||
|
||||
# Add some mock content
|
||||
content += " This is a deterministic mock response generated for testing. "
|
||||
content += f"Input hash: {input_hash[:8]}. "
|
||||
content += "The response is reproducible given the same input."
|
||||
|
||||
return MockResponse(
|
||||
content=content,
|
||||
tokens=len(content.split()) * 2,
|
||||
cost=len(content) * 0.000001,
|
||||
)
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
input_content: str,
|
||||
model: str = "mock-model",
|
||||
stream: bool = True,
|
||||
on_chunk: Optional[Callable[[str], None]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate a mock response.
|
||||
|
||||
Args:
|
||||
input_content: The input prompt
|
||||
model: Model name (ignored for mock)
|
||||
stream: Whether to stream the response
|
||||
on_chunk: Callback for streaming chunks
|
||||
|
||||
Returns:
|
||||
Response dict with content, tokens, cost
|
||||
"""
|
||||
response = self._get_response(input_content)
|
||||
|
||||
# Record if in record mode
|
||||
if self.config.record_mode:
|
||||
self._save_recording(input_content, response)
|
||||
|
||||
# Check for error
|
||||
if response.error:
|
||||
raise Exception(response.error)
|
||||
|
||||
# Stream chunks
|
||||
if stream and on_chunk:
|
||||
for chunk in response.chunks:
|
||||
await asyncio.sleep(response.delay_per_chunk)
|
||||
on_chunk(chunk)
|
||||
else:
|
||||
# Non-streaming delay
|
||||
await asyncio.sleep(len(response.chunks) * response.delay_per_chunk)
|
||||
|
||||
return {
|
||||
"content": response.content,
|
||||
"tokens": response.tokens,
|
||||
"cost": response.cost,
|
||||
"tool_calls": response.tool_calls,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
input_content: str,
|
||||
model: str = "mock-model",
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream a mock response.
|
||||
|
||||
Yields chunks of the response.
|
||||
"""
|
||||
response = self._get_response(input_content)
|
||||
|
||||
if response.error:
|
||||
raise Exception(response.error)
|
||||
|
||||
for chunk in response.chunks:
|
||||
await asyncio.sleep(response.delay_per_chunk)
|
||||
yield chunk
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset provider state."""
|
||||
self._rng = random.Random(self.config.seed)
|
||||
self._playback_index = 0
|
||||
|
||||
|
||||
class MockExecutor:
|
||||
"""
|
||||
Mock executor for queue worker testing.
|
||||
|
||||
Simulates agent execution with configurable behavior.
|
||||
"""
|
||||
|
||||
def __init__(self, provider: Optional[MockProvider] = None):
|
||||
self.provider = provider or MockProvider()
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
input_content: str,
|
||||
agent_config: Dict[str, Any],
|
||||
on_output: Optional[Callable[[str], None]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Execute a mock agent run."""
|
||||
model = agent_config.get("model", "mock-model")
|
||||
|
||||
result = await self.provider.generate(
|
||||
input_content=input_content,
|
||||
model=model,
|
||||
stream=True,
|
||||
on_chunk=on_output,
|
||||
)
|
||||
|
||||
return {
|
||||
"output": result["content"],
|
||||
"metrics": {
|
||||
"tokens": result["tokens"],
|
||||
"cost": result["cost"],
|
||||
},
|
||||
"tool_calls": result.get("tool_calls", []),
|
||||
}
|
||||
|
||||
|
||||
# Convenience function for tests
|
||||
def create_mock_provider(
|
||||
seed: int = 42,
|
||||
responses: Optional[Dict[str, str]] = None,
|
||||
) -> MockProvider:
|
||||
"""Create a mock provider with optional custom responses."""
|
||||
config = MockProviderConfig(seed=seed)
|
||||
|
||||
if responses:
|
||||
for pattern, content in responses.items():
|
||||
config.responses[pattern] = MockResponse(content=content)
|
||||
|
||||
return MockProvider(config)
|
||||
@@ -0,0 +1,652 @@
|
||||
"""
|
||||
TUI Orchestrator for PraisonAI.
|
||||
|
||||
Provides unified event handling for both interactive TUI and headless simulation modes.
|
||||
Inspired by gemini-cli's event-driven architecture and codex-cli's state management.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional, TextIO
|
||||
import sys
|
||||
|
||||
from .events import TUIEvent, TUIEventType
|
||||
from ..queue import QueueManager, QueueConfig, QueuedRun, RunState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OutputMode(str, Enum):
|
||||
"""Output mode for headless simulation."""
|
||||
PRETTY = "pretty"
|
||||
JSONL = "jsonl"
|
||||
SILENT = "silent"
|
||||
|
||||
|
||||
@dataclass
|
||||
class UIStateModel:
|
||||
"""
|
||||
In-memory UI state model.
|
||||
|
||||
Mirrors the state that would be displayed in the TUI,
|
||||
enabling headless simulation and snapshot generation.
|
||||
"""
|
||||
# Session info
|
||||
session_id: str = ""
|
||||
workspace: str = ""
|
||||
|
||||
# Model/config
|
||||
model: str = "gpt-4o-mini"
|
||||
|
||||
# Chat history
|
||||
messages: List[Dict[str, Any]] = field(default_factory=list)
|
||||
max_messages: int = 1000
|
||||
|
||||
# Current streaming state
|
||||
current_run_id: Optional[str] = None
|
||||
streaming_content: str = ""
|
||||
is_processing: bool = False
|
||||
|
||||
# Queue state
|
||||
queued_runs: List[Dict[str, Any]] = field(default_factory=list)
|
||||
running_runs: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
# Tool calls
|
||||
pending_tool_calls: List[Dict[str, Any]] = field(default_factory=list)
|
||||
recent_tool_calls: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
# Metrics
|
||||
total_tokens: int = 0
|
||||
total_cost: float = 0.0
|
||||
|
||||
# Focus/screen state (for simulation)
|
||||
current_screen: str = "main"
|
||||
focused_widget: str = "composer"
|
||||
|
||||
# Events log (for trace/replay)
|
||||
events: List[Dict[str, Any]] = field(default_factory=list)
|
||||
max_events: int = 10000
|
||||
|
||||
def add_message(self, role: str, content: str, **kwargs) -> None:
|
||||
"""Add a message to history."""
|
||||
msg = {
|
||||
"role": role,
|
||||
"content": content,
|
||||
"timestamp": time.time(),
|
||||
**kwargs
|
||||
}
|
||||
self.messages.append(msg)
|
||||
if len(self.messages) > self.max_messages:
|
||||
self.messages = self.messages[-self.max_messages:]
|
||||
|
||||
def add_event(self, event: TUIEvent) -> None:
|
||||
"""Add an event to the log."""
|
||||
evt = {
|
||||
"type": event.event_type.value,
|
||||
"timestamp": event.timestamp,
|
||||
"run_id": event.run_id,
|
||||
"session_id": event.session_id,
|
||||
"agent_name": event.agent_name,
|
||||
"data": event.data,
|
||||
}
|
||||
self.events.append(evt)
|
||||
if len(self.events) > self.max_events:
|
||||
self.events = self.events[-self.max_events:]
|
||||
|
||||
def to_snapshot(self) -> Dict[str, Any]:
|
||||
"""Generate a snapshot of current state."""
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"model": self.model,
|
||||
"current_screen": self.current_screen,
|
||||
"focused_widget": self.focused_widget,
|
||||
"is_processing": self.is_processing,
|
||||
"current_run_id": self.current_run_id,
|
||||
"streaming_content_length": len(self.streaming_content),
|
||||
"message_count": len(self.messages),
|
||||
"last_messages": self.messages[-5:] if self.messages else [],
|
||||
"queued_count": len(self.queued_runs),
|
||||
"running_count": len(self.running_runs),
|
||||
"pending_tool_calls": len(self.pending_tool_calls),
|
||||
"total_tokens": self.total_tokens,
|
||||
"total_cost": self.total_cost,
|
||||
}
|
||||
|
||||
def render_snapshot_pretty(self) -> str:
|
||||
"""Render a pretty text snapshot like TUI would show."""
|
||||
lines = []
|
||||
|
||||
# Status bar
|
||||
status = f"◉ PraisonAI │ Session: {self.session_id[:8] if self.session_id else 'new'}"
|
||||
status += f" │ Model: {self.model}"
|
||||
if self.total_tokens > 0:
|
||||
status += f" │ Tokens: {self.total_tokens:,}"
|
||||
if self.total_cost > 0:
|
||||
status += f" │ ${self.total_cost:.4f}"
|
||||
if self.is_processing:
|
||||
status += " │ ⟳ Processing..."
|
||||
lines.append("─" * 60)
|
||||
lines.append(status)
|
||||
lines.append("─" * 60)
|
||||
|
||||
# Chat messages (last 5)
|
||||
lines.append("\n[Chat History]")
|
||||
for msg in self.messages[-5:]:
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")[:100]
|
||||
if len(msg.get("content", "")) > 100:
|
||||
content += "..."
|
||||
lines.append(f" {role.upper()}: {content}")
|
||||
|
||||
# Streaming content
|
||||
if self.streaming_content:
|
||||
lines.append(f"\n[Streaming] ({len(self.streaming_content)} chars)")
|
||||
preview = self.streaming_content[-200:]
|
||||
if len(self.streaming_content) > 200:
|
||||
preview = "..." + preview
|
||||
lines.append(f" {preview}")
|
||||
|
||||
# Queue status
|
||||
lines.append(f"\n[Queue] Queued: {len(self.queued_runs)} │ Running: {len(self.running_runs)}")
|
||||
|
||||
# Tool calls
|
||||
if self.pending_tool_calls:
|
||||
lines.append(f"\n[Pending Approvals] {len(self.pending_tool_calls)}")
|
||||
for tc in self.pending_tool_calls[:3]:
|
||||
lines.append(f" ⚠ {tc.get('tool_name', 'unknown')}")
|
||||
|
||||
# Screen/focus
|
||||
lines.append(f"\n[UI] Screen: {self.current_screen} │ Focus: {self.focused_widget}")
|
||||
lines.append("─" * 60)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
class TuiOrchestrator:
|
||||
"""
|
||||
Unified orchestrator for TUI and headless modes.
|
||||
|
||||
Subscribes to the event bus and maintains an in-memory UI state model.
|
||||
Can drive Textual widgets (interactive) or output snapshots (headless).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
queue_manager: Optional[QueueManager] = None,
|
||||
queue_config: Optional[QueueConfig] = None,
|
||||
output_mode: OutputMode = OutputMode.PRETTY,
|
||||
output_stream: Optional[TextIO] = None,
|
||||
jsonl_path: Optional[str] = None,
|
||||
debug: bool = False,
|
||||
):
|
||||
self.queue_config = queue_config or QueueConfig()
|
||||
self.queue_manager = queue_manager
|
||||
self.output_mode = output_mode
|
||||
self.output_stream = output_stream or sys.stdout
|
||||
self.jsonl_path = jsonl_path
|
||||
self.debug = debug
|
||||
|
||||
# State
|
||||
self.state = UIStateModel()
|
||||
self._event_callbacks: List[Callable[[TUIEvent], None]] = []
|
||||
self._jsonl_file: Optional[TextIO] = None
|
||||
self._running = False
|
||||
|
||||
# Trace ID for this orchestrator session
|
||||
self.trace_id = str(uuid.uuid4())[:8]
|
||||
|
||||
async def start(self, session_id: Optional[str] = None, recover: bool = True) -> None:
|
||||
"""Start the orchestrator."""
|
||||
self.state.session_id = session_id or str(uuid.uuid4())[:8]
|
||||
self._running = True
|
||||
|
||||
# Open JSONL file if specified
|
||||
if self.jsonl_path:
|
||||
self._jsonl_file = open(self.jsonl_path, "a")
|
||||
|
||||
# Initialize queue manager if not provided
|
||||
if not self.queue_manager:
|
||||
self.queue_manager = QueueManager(
|
||||
config=self.queue_config,
|
||||
on_output=self._handle_output,
|
||||
on_complete=self._handle_complete,
|
||||
on_error=self._handle_error,
|
||||
)
|
||||
await self.queue_manager.start(recover=recover)
|
||||
self.queue_manager.set_session(self.state.session_id)
|
||||
|
||||
self._emit_event(TUIEvent(
|
||||
event_type=TUIEventType.SESSION_STARTED,
|
||||
session_id=self.state.session_id,
|
||||
data={"trace_id": self.trace_id}
|
||||
))
|
||||
|
||||
if self.debug:
|
||||
self._log_debug(f"Orchestrator started: session={self.state.session_id}, trace={self.trace_id}")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the orchestrator."""
|
||||
self._running = False
|
||||
|
||||
if self.queue_manager:
|
||||
await self.queue_manager.stop()
|
||||
|
||||
if self._jsonl_file:
|
||||
self._jsonl_file.close()
|
||||
self._jsonl_file = None
|
||||
|
||||
if self.debug:
|
||||
self._log_debug("Orchestrator stopped")
|
||||
|
||||
def add_event_callback(self, callback: Callable[[TUIEvent], None]) -> None:
|
||||
"""Add an event callback."""
|
||||
self._event_callbacks.append(callback)
|
||||
|
||||
def _emit_event(self, event: TUIEvent) -> None:
|
||||
"""Emit an event to all listeners."""
|
||||
# Add to state log
|
||||
self.state.add_event(event)
|
||||
|
||||
# Write to JSONL if enabled
|
||||
if self._jsonl_file:
|
||||
self._write_jsonl(event)
|
||||
|
||||
# Call callbacks
|
||||
for callback in self._event_callbacks:
|
||||
try:
|
||||
callback(event)
|
||||
except Exception as e:
|
||||
logger.error(f"Event callback error: {e}")
|
||||
|
||||
# Output based on mode
|
||||
if self.output_mode == OutputMode.JSONL:
|
||||
self._print_jsonl(event)
|
||||
elif self.output_mode == OutputMode.PRETTY and self.debug:
|
||||
self._print_event_pretty(event)
|
||||
|
||||
def _write_jsonl(self, event: TUIEvent) -> None:
|
||||
"""Write event to JSONL file."""
|
||||
if not self._jsonl_file:
|
||||
return
|
||||
|
||||
record = {
|
||||
"timestamp": event.timestamp,
|
||||
"trace_id": self.trace_id,
|
||||
"session_id": event.session_id or self.state.session_id,
|
||||
"event_type": event.event_type.value,
|
||||
"run_id": event.run_id,
|
||||
"agent_name": event.agent_name,
|
||||
"data": event.data,
|
||||
}
|
||||
self._jsonl_file.write(json.dumps(record, default=str) + "\n")
|
||||
self._jsonl_file.flush()
|
||||
|
||||
def _print_jsonl(self, event: TUIEvent) -> None:
|
||||
"""Print event as JSONL to output stream."""
|
||||
record = {
|
||||
"timestamp": event.timestamp,
|
||||
"event_type": event.event_type.value,
|
||||
"run_id": event.run_id,
|
||||
"data": event.data,
|
||||
}
|
||||
print(json.dumps(record, default=str), file=self.output_stream)
|
||||
|
||||
def _print_event_pretty(self, event: TUIEvent) -> None:
|
||||
"""Print event in pretty format."""
|
||||
ts = time.strftime("%H:%M:%S", time.localtime(event.timestamp))
|
||||
print(f"[{ts}] {event.event_type.value}", end="", file=self.output_stream)
|
||||
if event.run_id:
|
||||
print(f" run={event.run_id[:8]}", end="", file=self.output_stream)
|
||||
if event.data:
|
||||
data_preview = str(event.data)[:50]
|
||||
print(f" {data_preview}", end="", file=self.output_stream)
|
||||
print(file=self.output_stream)
|
||||
|
||||
def _log_debug(self, message: str) -> None:
|
||||
"""Log debug message."""
|
||||
if self.debug:
|
||||
ts = time.strftime("%H:%M:%S")
|
||||
print(f"[DEBUG {ts}] {message}", file=self.output_stream)
|
||||
|
||||
# Queue callbacks
|
||||
|
||||
async def _handle_output(self, run_id: str, chunk: str) -> None:
|
||||
"""Handle streaming output."""
|
||||
if run_id == self.state.current_run_id:
|
||||
self.state.streaming_content += chunk
|
||||
|
||||
self._emit_event(TUIEvent.output_chunk(run_id, chunk))
|
||||
|
||||
async def _handle_complete(self, run_id: str, run: QueuedRun) -> None:
|
||||
"""Handle run completion."""
|
||||
if run_id == self.state.current_run_id:
|
||||
# Add assistant message
|
||||
self.state.add_message(
|
||||
"assistant",
|
||||
run.output_content or self.state.streaming_content,
|
||||
run_id=run_id,
|
||||
agent_name=run.agent_name,
|
||||
)
|
||||
|
||||
# Update metrics
|
||||
if run.metrics:
|
||||
self.state.total_tokens += run.metrics.get("tokens", 0)
|
||||
self.state.total_cost += run.metrics.get("cost", 0.0)
|
||||
|
||||
# Clear streaming state
|
||||
self.state.current_run_id = None
|
||||
self.state.streaming_content = ""
|
||||
self.state.is_processing = False
|
||||
|
||||
# Update queue state
|
||||
self._update_queue_state()
|
||||
|
||||
self._emit_event(TUIEvent.run_completed(
|
||||
run_id,
|
||||
run.output_content or "",
|
||||
agent_name=run.agent_name,
|
||||
))
|
||||
|
||||
async def _handle_error(self, run_id: str, error: Exception) -> None:
|
||||
"""Handle run error."""
|
||||
if run_id == self.state.current_run_id:
|
||||
self.state.add_message(
|
||||
"system",
|
||||
f"Error: {error}",
|
||||
run_id=run_id,
|
||||
)
|
||||
self.state.current_run_id = None
|
||||
self.state.streaming_content = ""
|
||||
self.state.is_processing = False
|
||||
|
||||
self._update_queue_state()
|
||||
|
||||
self._emit_event(TUIEvent.error(str(error), run_id=run_id))
|
||||
|
||||
def _update_queue_state(self) -> None:
|
||||
"""Update queue state from manager."""
|
||||
if not self.queue_manager:
|
||||
return
|
||||
|
||||
runs = self.queue_manager.list_runs(limit=100)
|
||||
self.state.queued_runs = [
|
||||
r.to_dict() for r in runs if r.state == RunState.QUEUED
|
||||
]
|
||||
self.state.running_runs = [
|
||||
r.to_dict() for r in runs if r.state == RunState.RUNNING
|
||||
]
|
||||
|
||||
# Public API
|
||||
|
||||
async def submit_message(self, content: str, agent_name: str = "Assistant") -> str:
|
||||
"""Submit a message for processing."""
|
||||
# Add user message
|
||||
self.state.add_message("user", content)
|
||||
self.state.is_processing = True
|
||||
|
||||
self._emit_event(TUIEvent.message_submitted(content))
|
||||
|
||||
# Submit to queue
|
||||
run_id = await self.queue_manager.submit(
|
||||
input_content=content,
|
||||
agent_name=agent_name,
|
||||
config={"agent_config": {"name": agent_name, "model": self.state.model}}
|
||||
)
|
||||
|
||||
self.state.current_run_id = run_id
|
||||
self.state.streaming_content = ""
|
||||
|
||||
# Add placeholder for streaming
|
||||
self.state.add_message(
|
||||
"assistant",
|
||||
"",
|
||||
run_id=run_id,
|
||||
agent_name=agent_name,
|
||||
is_streaming=True,
|
||||
)
|
||||
|
||||
self._update_queue_state()
|
||||
|
||||
return run_id
|
||||
|
||||
async def cancel_run(self, run_id: Optional[str] = None) -> bool:
|
||||
"""Cancel a run."""
|
||||
target_id = run_id or self.state.current_run_id
|
||||
if not target_id:
|
||||
return False
|
||||
|
||||
result = await self.queue_manager.cancel(target_id)
|
||||
|
||||
if target_id == self.state.current_run_id:
|
||||
self.state.current_run_id = None
|
||||
self.state.streaming_content = ""
|
||||
self.state.is_processing = False
|
||||
|
||||
self._update_queue_state()
|
||||
|
||||
self._emit_event(TUIEvent(
|
||||
event_type=TUIEventType.RUN_CANCELLED,
|
||||
run_id=target_id,
|
||||
))
|
||||
|
||||
return result
|
||||
|
||||
async def retry_run(self, run_id: str) -> Optional[str]:
|
||||
"""Retry a failed run."""
|
||||
new_id = await self.queue_manager.retry(run_id)
|
||||
self._update_queue_state()
|
||||
return new_id
|
||||
|
||||
def set_model(self, model: str) -> None:
|
||||
"""Set the current model."""
|
||||
self.state.model = model
|
||||
self._emit_event(TUIEvent.status_update(f"Model set to {model}"))
|
||||
|
||||
def navigate_screen(self, screen: str) -> None:
|
||||
"""Navigate to a screen (for simulation)."""
|
||||
self.state.current_screen = screen
|
||||
self._emit_event(TUIEvent(
|
||||
event_type=TUIEventType.SCREEN_CHANGED,
|
||||
data={"screen": screen}
|
||||
))
|
||||
|
||||
def set_focus(self, widget: str) -> None:
|
||||
"""Set focus to a widget (for simulation)."""
|
||||
self.state.focused_widget = widget
|
||||
self._emit_event(TUIEvent(
|
||||
event_type=TUIEventType.FOCUS_CHANGED,
|
||||
data={"widget": widget}
|
||||
))
|
||||
|
||||
def get_snapshot(self) -> Dict[str, Any]:
|
||||
"""Get current state snapshot."""
|
||||
return self.state.to_snapshot()
|
||||
|
||||
def render_snapshot(self) -> str:
|
||||
"""Render pretty snapshot."""
|
||||
return self.state.render_snapshot_pretty()
|
||||
|
||||
async def wait_for_idle(self, timeout: float = 60.0) -> bool:
|
||||
"""Wait until no runs are processing."""
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
if not self.state.is_processing and not self.state.running_runs:
|
||||
return True
|
||||
await asyncio.sleep(0.1)
|
||||
return False
|
||||
|
||||
async def wait_for_run(self, run_id: str, timeout: float = 60.0) -> bool:
|
||||
"""Wait for a specific run to complete."""
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
run = self.queue_manager.get_run(run_id)
|
||||
if run and run.state.is_terminal():
|
||||
return True
|
||||
await asyncio.sleep(0.1)
|
||||
return False
|
||||
|
||||
|
||||
@dataclass
|
||||
class SimulationStep:
|
||||
"""A step in a simulation script."""
|
||||
action: str # "submit", "cancel", "retry", "navigate", "focus", "wait", "approve", "deny"
|
||||
args: Dict[str, Any] = field(default_factory=dict)
|
||||
expected: Optional[Dict[str, Any]] = None # For assertion mode
|
||||
|
||||
|
||||
class SimulationRunner:
|
||||
"""
|
||||
Runs simulation scripts for headless TUI testing.
|
||||
|
||||
Script format (YAML):
|
||||
```yaml
|
||||
session_id: test-session
|
||||
model: gpt-4o-mini
|
||||
steps:
|
||||
- action: submit
|
||||
args:
|
||||
content: "Hello, world!"
|
||||
- action: wait
|
||||
args:
|
||||
condition: idle
|
||||
timeout: 30
|
||||
- action: navigate
|
||||
args:
|
||||
screen: queue
|
||||
- action: cancel
|
||||
args:
|
||||
run_id: current
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
orchestrator: TuiOrchestrator,
|
||||
assert_mode: bool = False,
|
||||
):
|
||||
self.orchestrator = orchestrator
|
||||
self.assert_mode = assert_mode
|
||||
self.assertions_passed = 0
|
||||
self.assertions_failed = 0
|
||||
self.errors: List[str] = []
|
||||
|
||||
async def run_script(self, script: Dict[str, Any]) -> bool:
|
||||
"""Run a simulation script."""
|
||||
# Initialize session
|
||||
session_id = script.get("session_id")
|
||||
await self.orchestrator.start(session_id=session_id)
|
||||
|
||||
if "model" in script:
|
||||
self.orchestrator.set_model(script["model"])
|
||||
|
||||
# Run steps
|
||||
steps = script.get("steps", [])
|
||||
for i, step_data in enumerate(steps):
|
||||
step = SimulationStep(
|
||||
action=step_data.get("action", ""),
|
||||
args=step_data.get("args", {}),
|
||||
expected=step_data.get("expected"),
|
||||
)
|
||||
|
||||
try:
|
||||
await self._run_step(step, i)
|
||||
except Exception as e:
|
||||
self.errors.append(f"Step {i} ({step.action}): {e}")
|
||||
if self.assert_mode:
|
||||
break
|
||||
|
||||
await self.orchestrator.stop()
|
||||
|
||||
return len(self.errors) == 0
|
||||
|
||||
async def _run_step(self, step: SimulationStep, index: int) -> None:
|
||||
"""Run a single simulation step."""
|
||||
action = step.action.lower()
|
||||
args = step.args
|
||||
|
||||
if action == "submit":
|
||||
content = args.get("content", "")
|
||||
agent = args.get("agent", "Assistant")
|
||||
await self.orchestrator.submit_message(content, agent)
|
||||
|
||||
elif action == "cancel":
|
||||
run_id = args.get("run_id")
|
||||
if run_id == "current":
|
||||
run_id = None
|
||||
await self.orchestrator.cancel_run(run_id)
|
||||
|
||||
elif action == "retry":
|
||||
run_id = args.get("run_id", "")
|
||||
await self.orchestrator.retry_run(run_id)
|
||||
|
||||
elif action == "navigate":
|
||||
screen = args.get("screen", "main")
|
||||
self.orchestrator.navigate_screen(screen)
|
||||
|
||||
elif action == "focus":
|
||||
widget = args.get("widget", "composer")
|
||||
self.orchestrator.set_focus(widget)
|
||||
|
||||
elif action == "wait":
|
||||
condition = args.get("condition", "idle")
|
||||
timeout = args.get("timeout", 30.0)
|
||||
|
||||
if condition == "idle":
|
||||
success = await self.orchestrator.wait_for_idle(timeout)
|
||||
elif condition == "run":
|
||||
run_id = args.get("run_id", "")
|
||||
success = await self.orchestrator.wait_for_run(run_id, timeout)
|
||||
else:
|
||||
await asyncio.sleep(timeout)
|
||||
success = True
|
||||
|
||||
if not success and self.assert_mode:
|
||||
raise TimeoutError(f"Wait condition '{condition}' timed out")
|
||||
|
||||
elif action == "model":
|
||||
model = args.get("model", "gpt-4o-mini")
|
||||
self.orchestrator.set_model(model)
|
||||
|
||||
elif action == "snapshot":
|
||||
snapshot = self.orchestrator.render_snapshot()
|
||||
print(snapshot)
|
||||
|
||||
elif action == "sleep":
|
||||
duration = args.get("seconds", 1.0)
|
||||
await asyncio.sleep(duration)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown action: {action}")
|
||||
|
||||
# Check assertions
|
||||
if self.assert_mode and step.expected:
|
||||
self._check_assertions(step.expected, index)
|
||||
|
||||
def _check_assertions(self, expected: Dict[str, Any], step_index: int) -> None:
|
||||
"""Check assertions against current state."""
|
||||
snapshot = self.orchestrator.get_snapshot()
|
||||
|
||||
for key, expected_value in expected.items():
|
||||
actual_value = snapshot.get(key)
|
||||
|
||||
if actual_value != expected_value:
|
||||
self.assertions_failed += 1
|
||||
self.errors.append(
|
||||
f"Step {step_index}: Expected {key}={expected_value}, got {actual_value}"
|
||||
)
|
||||
else:
|
||||
self.assertions_passed += 1
|
||||
|
||||
def get_summary(self) -> Dict[str, Any]:
|
||||
"""Get simulation summary."""
|
||||
return {
|
||||
"assertions_passed": self.assertions_passed,
|
||||
"assertions_failed": self.assertions_failed,
|
||||
"errors": self.errors,
|
||||
"success": len(self.errors) == 0,
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
TUI Screens for PraisonAI.
|
||||
|
||||
Screen definitions for different views in the TUI application.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .main import MainScreen
|
||||
from .queue import QueueScreen
|
||||
from .settings import SettingsScreen
|
||||
from .session import SessionScreen
|
||||
|
||||
_lazy_cache = {}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load screens."""
|
||||
global _lazy_cache
|
||||
|
||||
if name in _lazy_cache:
|
||||
return _lazy_cache[name]
|
||||
|
||||
if name == "MainScreen":
|
||||
from .main import MainScreen
|
||||
_lazy_cache[name] = MainScreen
|
||||
return MainScreen
|
||||
elif name == "QueueScreen":
|
||||
from .queue import QueueScreen
|
||||
_lazy_cache[name] = QueueScreen
|
||||
return QueueScreen
|
||||
elif name == "SettingsScreen":
|
||||
from .settings import SettingsScreen
|
||||
_lazy_cache[name] = SettingsScreen
|
||||
return SettingsScreen
|
||||
elif name == "SessionScreen":
|
||||
from .session import SessionScreen
|
||||
_lazy_cache[name] = SessionScreen
|
||||
return SessionScreen
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MainScreen",
|
||||
"QueueScreen",
|
||||
"SettingsScreen",
|
||||
"SessionScreen",
|
||||
]
|
||||
@@ -0,0 +1,157 @@
|
||||
"""
|
||||
Help Screen for PraisonAI TUI.
|
||||
|
||||
Displays keyboard shortcuts and usage information.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
try:
|
||||
from textual.screen import Screen
|
||||
from textual.containers import Vertical, Container
|
||||
from textual.widgets import Static, Footer, Button
|
||||
from textual.binding import Binding
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Screen = object
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
from rich.text import Text
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
|
||||
class HelpScreen(Screen):
|
||||
"""
|
||||
Help screen showing keyboard shortcuts and usage.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "dismiss", "Close", show=True),
|
||||
Binding("q", "dismiss", "Close", show=True),
|
||||
]
|
||||
|
||||
DEFAULT_CSS = """
|
||||
HelpScreen {
|
||||
align: center middle;
|
||||
}
|
||||
|
||||
HelpScreen #help-container {
|
||||
width: 80;
|
||||
height: auto;
|
||||
max-height: 90%;
|
||||
background: $surface;
|
||||
border: solid $primary;
|
||||
padding: 1 2;
|
||||
}
|
||||
|
||||
HelpScreen #help-title {
|
||||
text-align: center;
|
||||
text-style: bold;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
HelpScreen #close-button {
|
||||
margin-top: 1;
|
||||
width: 100%;
|
||||
}
|
||||
"""
|
||||
|
||||
def compose(self):
|
||||
"""Compose the help screen."""
|
||||
with Container(id="help-container"):
|
||||
yield Static("PraisonAI TUI Help", id="help-title")
|
||||
yield Static(self._build_help_content())
|
||||
yield Button("Close (Esc)", id="close-button", variant="primary")
|
||||
|
||||
def _build_help_content(self) -> Text:
|
||||
"""Build the help content."""
|
||||
content = Text()
|
||||
|
||||
# Safe keyboard shortcuts section
|
||||
content.append("\n📋 Keyboard Shortcuts (Safe Defaults)\n", style="bold cyan")
|
||||
content.append("─" * 40 + "\n", style="dim")
|
||||
|
||||
shortcuts = [
|
||||
("\\", "Show command popup"),
|
||||
("@path", "Reference file (type freely)"),
|
||||
("/", "Show command discovery"),
|
||||
("Enter", "Send message"),
|
||||
("Shift+Enter", "New line in message"),
|
||||
("q", "Quit (when not typing)"),
|
||||
("?", "Show this help"),
|
||||
(":", "Enter command mode"),
|
||||
("Escape", "Cancel/Close/Dismiss"),
|
||||
("Tab", "Focus next element"),
|
||||
("Up/Down", "Navigate history"),
|
||||
]
|
||||
|
||||
for key, desc in shortcuts:
|
||||
content.append(f" {key:<15}", style="yellow")
|
||||
content.append(f" {desc}\n")
|
||||
|
||||
# Command mode section
|
||||
content.append("\n⌨️ Command Mode (: prefix)\n", style="bold cyan")
|
||||
content.append("─" * 40 + "\n", style="dim")
|
||||
|
||||
colon_cmds = [
|
||||
(":quit, :q", "Quit application"),
|
||||
(":clear, :cl", "Clear chat history"),
|
||||
(":help, :h", "Show help"),
|
||||
(":tools, :t", "Toggle tools panel"),
|
||||
(":queue, :qu", "Toggle queue panel"),
|
||||
(":settings", "Open settings"),
|
||||
(":cancel", "Cancel current operation"),
|
||||
]
|
||||
|
||||
for cmd, desc in colon_cmds:
|
||||
content.append(f" {cmd:<15}", style="green")
|
||||
content.append(f" {desc}\n")
|
||||
|
||||
# Slash commands section
|
||||
content.append("\n💬 Slash Commands (/ prefix)\n", style="bold cyan")
|
||||
content.append("─" * 40 + "\n", style="dim")
|
||||
|
||||
commands = [
|
||||
("/help", "Show help"),
|
||||
("/clear", "Clear chat"),
|
||||
("/model", "Show/change model"),
|
||||
("/queue", "Show queue status"),
|
||||
("/cancel", "Cancel current run"),
|
||||
("/cost", "Show cost summary"),
|
||||
("/settings", "Open settings"),
|
||||
("/sessions", "Manage sessions"),
|
||||
("/exit", "Exit TUI"),
|
||||
]
|
||||
|
||||
for cmd, desc in commands:
|
||||
content.append(f" {cmd:<15}", style="green")
|
||||
content.append(f" {desc}\n")
|
||||
|
||||
# Optional keys note
|
||||
content.append("\n⚙️ Optional Keys (disabled by default)\n", style="bold cyan")
|
||||
content.append("─" * 40 + "\n", style="dim")
|
||||
content.append(" Ctrl keys: ", style="dim")
|
||||
content.append("PRAISONAI_TUI_CTRL_KEYS=1\n", style="yellow")
|
||||
content.append(" Fn keys: ", style="dim")
|
||||
content.append("PRAISONAI_TUI_FN_KEYS=1\n", style="yellow")
|
||||
|
||||
return content
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button press."""
|
||||
if event.button.id == "close-button":
|
||||
self.dismiss()
|
||||
|
||||
def action_dismiss(self) -> None:
|
||||
"""Dismiss the help screen."""
|
||||
self.dismiss()
|
||||
|
||||
else:
|
||||
class HelpScreen:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,503 @@
|
||||
"""
|
||||
Main Screen for PraisonAI TUI.
|
||||
|
||||
The primary chat interface screen.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
try:
|
||||
from textual.screen import Screen
|
||||
from textual.containers import Horizontal, Vertical, Container
|
||||
from textual.widgets import Static, Footer, Header
|
||||
from textual.binding import Binding
|
||||
from textual.message import Message
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Screen = object
|
||||
Message = object
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
from ..widgets.chat import ChatWidget, ChatMessage
|
||||
from ..widgets.composer import ComposerWidget
|
||||
from ..widgets.status import StatusWidget, StatusInfo
|
||||
from ..widgets.queue_panel import QueuePanelWidget
|
||||
from ..widgets.tool_panel import ToolPanelWidget
|
||||
from ..widgets.command_popup import CommandPopupWidget, get_default_commands
|
||||
from ..widgets.file_popup import FilePopupWidget
|
||||
from ..config import TUIConfig, get_command
|
||||
|
||||
class MainScreen(Screen):
|
||||
"""
|
||||
Main chat screen.
|
||||
|
||||
Layout:
|
||||
- Status bar (top)
|
||||
- Chat history (main area)
|
||||
- Tools panel (right sidebar)
|
||||
- Queue panel (right sidebar)
|
||||
- Composer (bottom)
|
||||
- Footer with keybindings
|
||||
"""
|
||||
|
||||
# Safe default bindings - no Ctrl or Fn keys by default
|
||||
# These work in VS Code, Windsurf, tmux, and plain terminals
|
||||
BINDINGS = [
|
||||
# Single-character shortcuts (safe defaults)
|
||||
Binding("q", "quit", "Quit", show=True, priority=True),
|
||||
Binding("question_mark", "help", "Help", show=True),
|
||||
Binding("backslash", "show_commands", "\\cmds", show=True),
|
||||
Binding("colon", "command_mode", ":cmd", show=False),
|
||||
Binding("slash", "search", "/search", show=False),
|
||||
# Navigation
|
||||
Binding("tab", "focus_next", "Next", show=False),
|
||||
Binding("escape", "cancel_or_dismiss", "Cancel", show=True),
|
||||
]
|
||||
|
||||
DEFAULT_CSS = """
|
||||
MainScreen {
|
||||
layout: grid;
|
||||
grid-size: 2;
|
||||
grid-columns: 3fr 1fr;
|
||||
grid-rows: 1 1fr auto 1;
|
||||
layers: base popup;
|
||||
}
|
||||
|
||||
MainScreen #status-bar {
|
||||
column-span: 2;
|
||||
}
|
||||
|
||||
MainScreen #main-area {
|
||||
height: 100%;
|
||||
}
|
||||
|
||||
MainScreen #sidebar {
|
||||
height: 100%;
|
||||
}
|
||||
|
||||
MainScreen #composer-area {
|
||||
column-span: 2;
|
||||
height: auto;
|
||||
max-height: 10;
|
||||
}
|
||||
|
||||
MainScreen Footer {
|
||||
column-span: 2;
|
||||
}
|
||||
|
||||
MainScreen .sidebar-panel {
|
||||
height: 1fr;
|
||||
}
|
||||
|
||||
MainScreen #command-popup {
|
||||
layer: popup;
|
||||
dock: bottom;
|
||||
margin: 0 2 2 2;
|
||||
}
|
||||
|
||||
MainScreen #file-popup {
|
||||
layer: popup;
|
||||
dock: bottom;
|
||||
margin: 0 2 2 2;
|
||||
}
|
||||
"""
|
||||
|
||||
class MessageSubmitted(Message):
|
||||
"""Event when a message is submitted."""
|
||||
def __init__(self, content: str):
|
||||
self.content = content
|
||||
super().__init__()
|
||||
|
||||
class CommandExecuted(Message):
|
||||
"""Event when a command is executed."""
|
||||
def __init__(self, command: str, args: str):
|
||||
self.command = command
|
||||
self.args = args
|
||||
super().__init__()
|
||||
|
||||
class CancelRequested(Message):
|
||||
"""Event when cancel is requested."""
|
||||
pass
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
show_queue: bool = True,
|
||||
show_tools: bool = True,
|
||||
config: Optional[TUIConfig] = None,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._show_queue = show_queue
|
||||
self._show_tools = show_tools
|
||||
self._config = config or TUIConfig.from_env()
|
||||
self._command_mode = False
|
||||
self._command_buffer = ""
|
||||
self._command_popup_visible = False
|
||||
self._file_popup_visible = False
|
||||
self._workspace = None
|
||||
|
||||
# Dynamically add Ctrl/Fn bindings if enabled
|
||||
self._setup_optional_bindings()
|
||||
|
||||
def compose(self):
|
||||
"""Compose the screen."""
|
||||
yield StatusWidget(id="status-bar")
|
||||
|
||||
yield ChatWidget(id="chat-widget")
|
||||
|
||||
with Vertical(id="sidebar"):
|
||||
if self._show_tools:
|
||||
yield ToolPanelWidget(id="tool-panel", classes="sidebar-panel")
|
||||
if self._show_queue:
|
||||
yield QueuePanelWidget(id="queue-panel", classes="sidebar-panel")
|
||||
|
||||
yield ComposerWidget(id="composer")
|
||||
|
||||
yield Footer()
|
||||
|
||||
def _setup_optional_bindings(self) -> None:
|
||||
"""Setup optional Ctrl/Fn key bindings based on config."""
|
||||
if self._config.enable_ctrl_keys:
|
||||
# Add Ctrl key bindings (opt-in)
|
||||
self._bindings.bind("ctrl+q", "quit", "^Q Quit", show=True)
|
||||
self._bindings.bind("ctrl+c", "cancel", "^C Cancel", show=True)
|
||||
self._bindings.bind("ctrl+l", "clear_chat", "^L Clear", show=True)
|
||||
|
||||
if self._config.enable_fn_keys:
|
||||
# Add Function key bindings (opt-in)
|
||||
self._bindings.bind("f1", "help", "F1 Help", show=True)
|
||||
self._bindings.bind("f2", "toggle_queue", "F2 Queue", show=True)
|
||||
self._bindings.bind("f3", "settings", "F3 Settings", show=True)
|
||||
self._bindings.bind("f5", "clear_chat", "F5 Clear", show=True)
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
# Focus composer
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.focus_input()
|
||||
|
||||
def on_composer_widget_submitted(self, event: ComposerWidget.Submitted) -> None:
|
||||
"""Handle message submission."""
|
||||
if event.is_command:
|
||||
# Parse command (supports both / and : prefix)
|
||||
content = event.content
|
||||
if content.startswith("/") or content.startswith(":"):
|
||||
content = content[1:] # Remove prefix
|
||||
|
||||
parts = content.split(maxsplit=1)
|
||||
command = parts[0] if parts else ""
|
||||
args = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
# Handle local screen commands first
|
||||
if self._handle_local_command(command, args):
|
||||
return
|
||||
|
||||
# Forward to app for other commands
|
||||
self.post_message(self.CommandExecuted(command, args))
|
||||
else:
|
||||
self.post_message(self.MessageSubmitted(event.content))
|
||||
|
||||
def _handle_local_command(self, command: str, args: str) -> bool:
|
||||
"""Handle commands that can be processed locally without app involvement."""
|
||||
cmd = get_command(command)
|
||||
|
||||
if cmd == "quit":
|
||||
self.app.exit()
|
||||
return True
|
||||
elif cmd == "clear":
|
||||
self.action_clear_chat()
|
||||
return True
|
||||
elif cmd == "help":
|
||||
self.action_help()
|
||||
return True
|
||||
elif cmd == "tools":
|
||||
self.action_toggle_tools()
|
||||
return True
|
||||
elif cmd == "queue":
|
||||
self.action_toggle_queue()
|
||||
return True
|
||||
elif cmd == "settings":
|
||||
self.action_settings()
|
||||
return True
|
||||
elif cmd == "cancel":
|
||||
self.action_cancel()
|
||||
return True
|
||||
|
||||
return False # Not handled locally
|
||||
|
||||
def action_quit(self) -> None:
|
||||
"""Quit the application."""
|
||||
self.app.exit()
|
||||
|
||||
def action_cancel(self) -> None:
|
||||
"""Cancel current operation."""
|
||||
self.post_message(self.CancelRequested())
|
||||
|
||||
def action_clear(self) -> None:
|
||||
"""Clear screen."""
|
||||
pass
|
||||
|
||||
def action_help(self) -> None:
|
||||
"""Show help."""
|
||||
self.app.push_screen("help")
|
||||
|
||||
def action_toggle_queue(self) -> None:
|
||||
"""Toggle queue panel visibility."""
|
||||
try:
|
||||
queue_panel = self.query_one("#queue-panel", QueuePanelWidget)
|
||||
queue_panel.display = not queue_panel.display
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def action_settings(self) -> None:
|
||||
"""Show settings."""
|
||||
self.app.push_screen("settings")
|
||||
|
||||
def action_clear_chat(self) -> None:
|
||||
"""Clear chat history."""
|
||||
chat = self.query_one("#chat-widget", ChatWidget)
|
||||
self.app.call_later(chat.clear)
|
||||
|
||||
def action_cancel_input(self) -> None:
|
||||
"""Cancel current input."""
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.action_cancel()
|
||||
|
||||
def action_command_mode(self) -> None:
|
||||
"""Enter command mode - focus composer with : prefix."""
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.focus_input()
|
||||
composer.set_text(":")
|
||||
|
||||
def action_search(self) -> None:
|
||||
"""Enter search mode - focus composer with / prefix."""
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.focus_input()
|
||||
composer.set_text("/")
|
||||
|
||||
def action_toggle_tools(self) -> None:
|
||||
"""Toggle tools panel visibility."""
|
||||
try:
|
||||
tool_panel = self.query_one("#tool-panel", ToolPanelWidget)
|
||||
tool_panel.display = not tool_panel.display
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def action_show_commands(self) -> None:
|
||||
"""Show command popup (triggered by backslash)."""
|
||||
if self._command_popup_visible:
|
||||
return
|
||||
|
||||
self._command_popup_visible = True
|
||||
|
||||
# Create and mount command popup
|
||||
popup = CommandPopupWidget(
|
||||
commands=get_default_commands(),
|
||||
id="command-popup",
|
||||
)
|
||||
self.mount(popup)
|
||||
popup.focus_search()
|
||||
|
||||
def action_cancel_or_dismiss(self) -> None:
|
||||
"""Cancel input or dismiss popup."""
|
||||
if self._command_popup_visible:
|
||||
self._dismiss_command_popup()
|
||||
else:
|
||||
self.action_cancel_input()
|
||||
|
||||
def _dismiss_command_popup(self) -> None:
|
||||
"""Dismiss the command popup."""
|
||||
if not self._command_popup_visible:
|
||||
return
|
||||
|
||||
try:
|
||||
popup = self.query_one("#command-popup", CommandPopupWidget)
|
||||
popup.remove()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self._command_popup_visible = False
|
||||
|
||||
# Return focus to composer
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.focus_input()
|
||||
|
||||
def on_command_popup_widget_command_selected(
|
||||
self, event: CommandPopupWidget.CommandSelected
|
||||
) -> None:
|
||||
"""Handle command selection from popup."""
|
||||
self._dismiss_command_popup()
|
||||
|
||||
# Execute the selected command
|
||||
if self._handle_local_command(event.command, event.args):
|
||||
return
|
||||
|
||||
# Forward to app for other commands
|
||||
self.post_message(self.CommandExecuted(event.command, event.args))
|
||||
|
||||
def on_command_popup_widget_dismissed(
|
||||
self, event: CommandPopupWidget.Dismissed
|
||||
) -> None:
|
||||
"""Handle popup dismissal."""
|
||||
self._dismiss_command_popup()
|
||||
|
||||
# Auto-trigger event handlers from ComposerWidget
|
||||
|
||||
def on_composer_widget_backslash_typed(
|
||||
self, event: ComposerWidget.BackslashTyped
|
||||
) -> None:
|
||||
"""Handle backslash typed - auto-show command popup."""
|
||||
# Clear the backslash from composer and show popup
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.set_text("")
|
||||
self.action_show_commands()
|
||||
|
||||
def on_composer_widget_at_typed(
|
||||
self, event: ComposerWidget.AtTyped
|
||||
) -> None:
|
||||
"""Handle @ typed - DO NOT block typing.
|
||||
|
||||
MINIMUM VIABLE UX: User can freely type @path/to/file.
|
||||
The popup is optional and must NOT steal focus or block input.
|
||||
For now, we simply allow typing to continue uninterrupted.
|
||||
"""
|
||||
# DO NOT show popup that blocks typing
|
||||
# User can manually type: @README.md, @src/main.py, etc.
|
||||
pass
|
||||
|
||||
def on_composer_widget_slash_typed(
|
||||
self, event: ComposerWidget.SlashTyped
|
||||
) -> None:
|
||||
"""Handle / typed - auto-show command popup with filter."""
|
||||
if not self._command_popup_visible:
|
||||
self._command_popup_visible = True
|
||||
popup = CommandPopupWidget(
|
||||
commands=get_default_commands(),
|
||||
id="command-popup",
|
||||
)
|
||||
self.mount(popup)
|
||||
# Set filter to current query
|
||||
popup.filter_text = event.query
|
||||
|
||||
def _show_file_popup(self, query: str, at_pos: int) -> None:
|
||||
"""Show file autocomplete popup."""
|
||||
if self._file_popup_visible:
|
||||
# Update existing popup
|
||||
try:
|
||||
popup = self.query_one("#file-popup", FilePopupWidget)
|
||||
popup.set_query(query, at_pos)
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
self._file_popup_visible = True
|
||||
|
||||
# Create and mount file popup
|
||||
import os
|
||||
popup = FilePopupWidget(
|
||||
root_dir=self._workspace or os.getcwd(),
|
||||
at_pos=at_pos,
|
||||
id="file-popup",
|
||||
)
|
||||
self.mount(popup)
|
||||
popup.set_query(query, at_pos)
|
||||
|
||||
def _dismiss_file_popup(self) -> None:
|
||||
"""Dismiss the file popup."""
|
||||
if not self._file_popup_visible:
|
||||
return
|
||||
|
||||
try:
|
||||
popup = self.query_one("#file-popup", FilePopupWidget)
|
||||
popup.remove()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self._file_popup_visible = False
|
||||
|
||||
# Return focus to composer
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.focus_input()
|
||||
|
||||
def on_file_popup_widget_file_selected(
|
||||
self, event: FilePopupWidget.FileSelected
|
||||
) -> None:
|
||||
"""Handle file selection from popup."""
|
||||
self._dismiss_file_popup()
|
||||
|
||||
# Insert the file path into composer
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
current_text = composer.text
|
||||
|
||||
# Replace from @ position to end with selected path
|
||||
new_text = current_text[:event.at_pos] + "@" + event.path
|
||||
composer.set_text(new_text)
|
||||
composer.focus_input()
|
||||
|
||||
def on_file_popup_widget_dismissed(
|
||||
self, event: FilePopupWidget.Dismissed
|
||||
) -> None:
|
||||
"""Handle file popup dismissal."""
|
||||
self._dismiss_file_popup()
|
||||
|
||||
# Public methods for updating UI
|
||||
|
||||
async def add_user_message(self, content: str) -> None:
|
||||
"""Add a user message to chat."""
|
||||
chat = self.query_one("#chat-widget", ChatWidget)
|
||||
await chat.add_message(ChatMessage(
|
||||
role="user",
|
||||
content=content,
|
||||
))
|
||||
|
||||
async def add_assistant_message(
|
||||
self,
|
||||
content: str,
|
||||
run_id: Optional[str] = None,
|
||||
agent_name: Optional[str] = None,
|
||||
streaming: bool = False,
|
||||
) -> None:
|
||||
"""Add an assistant message to chat."""
|
||||
chat = self.query_one("#chat-widget", ChatWidget)
|
||||
await chat.add_message(ChatMessage(
|
||||
role="assistant",
|
||||
content=content,
|
||||
run_id=run_id,
|
||||
agent_name=agent_name,
|
||||
is_streaming=streaming,
|
||||
))
|
||||
|
||||
async def update_streaming(self, run_id: str, content: str) -> None:
|
||||
"""Update streaming message."""
|
||||
chat = self.query_one("#chat-widget", ChatWidget)
|
||||
await chat.update_streaming(run_id, content)
|
||||
|
||||
async def complete_streaming(self, run_id: str, content: str) -> None:
|
||||
"""Complete streaming message."""
|
||||
chat = self.query_one("#chat-widget", ChatWidget)
|
||||
await chat.complete_streaming(run_id, content)
|
||||
|
||||
def update_status(self, info: StatusInfo) -> None:
|
||||
"""Update status bar."""
|
||||
status = self.query_one("#status-bar", StatusWidget)
|
||||
status.update_info(info)
|
||||
|
||||
def set_processing(self, processing: bool) -> None:
|
||||
"""Set processing state."""
|
||||
composer = self.query_one("#composer", ComposerWidget)
|
||||
composer.set_processing(processing)
|
||||
|
||||
status = self.query_one("#status-bar", StatusWidget)
|
||||
status.is_processing = processing
|
||||
|
||||
else:
|
||||
class MainScreen:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
Queue Screen for PraisonAI TUI.
|
||||
|
||||
Queue management interface.
|
||||
"""
|
||||
|
||||
from typing import Optional, List
|
||||
|
||||
try:
|
||||
from textual.screen import Screen
|
||||
from textual.containers import Vertical, Horizontal
|
||||
from textual.widgets import Static, Footer, DataTable, Button
|
||||
from textual.binding import Binding
|
||||
from textual.message import Message
|
||||
from rich.text import Text
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Screen = object
|
||||
Message = object
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class QueueScreen(Screen):
|
||||
"""
|
||||
Queue management screen.
|
||||
|
||||
Displays full queue with management options.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "back", "Back", show=True),
|
||||
Binding("c", "cancel", "Cancel", show=True),
|
||||
Binding("r", "retry", "Retry", show=True),
|
||||
Binding("x", "clear", "Clear Queue", show=True),
|
||||
Binding("enter", "select", "Select", show=False),
|
||||
]
|
||||
|
||||
DEFAULT_CSS = """
|
||||
QueueScreen {
|
||||
background: $surface;
|
||||
}
|
||||
|
||||
QueueScreen #queue-header {
|
||||
height: 3;
|
||||
background: $primary;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
QueueScreen DataTable {
|
||||
height: 1fr;
|
||||
}
|
||||
|
||||
QueueScreen #queue-actions {
|
||||
height: 3;
|
||||
padding: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
class RunCancelled(Message):
|
||||
def __init__(self, run_id: str):
|
||||
self.run_id = run_id
|
||||
super().__init__()
|
||||
|
||||
class RunRetried(Message):
|
||||
def __init__(self, run_id: str):
|
||||
self.run_id = run_id
|
||||
super().__init__()
|
||||
|
||||
class QueueCleared(Message):
|
||||
pass
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._selected_run_id: Optional[str] = None
|
||||
|
||||
def compose(self):
|
||||
"""Compose the screen."""
|
||||
yield Static("Queue Management", id="queue-header")
|
||||
yield DataTable(id="queue-table")
|
||||
with Horizontal(id="queue-actions"):
|
||||
yield Button("Cancel [C]", id="btn-cancel", variant="error")
|
||||
yield Button("Retry [R]", id="btn-retry", variant="warning")
|
||||
yield Button("Clear [X]", id="btn-clear", variant="default")
|
||||
yield Button("Back [Esc]", id="btn-back", variant="primary")
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
table = self.query_one("#queue-table", DataTable)
|
||||
table.add_columns(
|
||||
"ID", "Agent", "Input", "State", "Priority",
|
||||
"Wait Time", "Duration"
|
||||
)
|
||||
table.cursor_type = "row"
|
||||
|
||||
def update_runs(self, runs: List[dict]) -> None:
|
||||
"""Update the runs display."""
|
||||
table = self.query_one("#queue-table", DataTable)
|
||||
table.clear()
|
||||
|
||||
for run in runs:
|
||||
state = run.get("state", "unknown")
|
||||
state_style = {
|
||||
"queued": "yellow",
|
||||
"running": "green",
|
||||
"succeeded": "cyan",
|
||||
"failed": "red",
|
||||
"cancelled": "dim",
|
||||
}.get(state, "")
|
||||
|
||||
state_text = Text(state, style=state_style)
|
||||
|
||||
input_preview = run.get("input", "")[:40]
|
||||
if len(run.get("input", "")) > 40:
|
||||
input_preview += "..."
|
||||
|
||||
table.add_row(
|
||||
run.get("run_id", "")[:8],
|
||||
run.get("agent_name", ""),
|
||||
input_preview,
|
||||
state_text,
|
||||
run.get("priority", "normal"),
|
||||
run.get("wait_time", "-"),
|
||||
run.get("duration", "-"),
|
||||
key=run.get("run_id"),
|
||||
)
|
||||
|
||||
def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
|
||||
"""Handle row selection."""
|
||||
if event.row_key:
|
||||
self._selected_run_id = str(event.row_key.value)
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button press."""
|
||||
if event.button.id == "btn-cancel":
|
||||
self.action_cancel()
|
||||
elif event.button.id == "btn-retry":
|
||||
self.action_retry()
|
||||
elif event.button.id == "btn-clear":
|
||||
self.action_clear()
|
||||
elif event.button.id == "btn-back":
|
||||
self.action_back()
|
||||
|
||||
def action_back(self) -> None:
|
||||
"""Go back to main screen."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_cancel(self) -> None:
|
||||
"""Cancel selected run."""
|
||||
if self._selected_run_id:
|
||||
self.post_message(self.RunCancelled(self._selected_run_id))
|
||||
|
||||
def action_retry(self) -> None:
|
||||
"""Retry selected run."""
|
||||
if self._selected_run_id:
|
||||
self.post_message(self.RunRetried(self._selected_run_id))
|
||||
|
||||
def action_clear(self) -> None:
|
||||
"""Clear the queue."""
|
||||
self.post_message(self.QueueCleared())
|
||||
|
||||
else:
|
||||
class QueueScreen:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
"""
|
||||
Session Screen for PraisonAI TUI.
|
||||
|
||||
Session browser and management.
|
||||
"""
|
||||
|
||||
from typing import Optional, List
|
||||
|
||||
try:
|
||||
from textual.screen import Screen
|
||||
from textual.containers import Vertical
|
||||
from textual.widgets import Static, Footer, DataTable, Button
|
||||
from textual.binding import Binding
|
||||
from textual.message import Message
|
||||
from rich.text import Text
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Screen = object
|
||||
Message = object
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class SessionScreen(Screen):
|
||||
"""
|
||||
Session browser screen.
|
||||
|
||||
Displays recent sessions and allows resuming.
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "back", "Back", show=True),
|
||||
Binding("enter", "resume", "Resume", show=True),
|
||||
Binding("d", "delete", "Delete", show=True),
|
||||
]
|
||||
|
||||
DEFAULT_CSS = """
|
||||
SessionScreen {
|
||||
background: $surface;
|
||||
}
|
||||
|
||||
SessionScreen #session-header {
|
||||
height: 3;
|
||||
background: $primary;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
SessionScreen DataTable {
|
||||
height: 1fr;
|
||||
}
|
||||
"""
|
||||
|
||||
class SessionSelected(Message):
|
||||
def __init__(self, session_id: str):
|
||||
self.session_id = session_id
|
||||
super().__init__()
|
||||
|
||||
class SessionDeleted(Message):
|
||||
def __init__(self, session_id: str):
|
||||
self.session_id = session_id
|
||||
super().__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._selected_session_id: Optional[str] = None
|
||||
|
||||
def compose(self):
|
||||
"""Compose the screen."""
|
||||
yield Static("Sessions", id="session-header")
|
||||
yield DataTable(id="session-table")
|
||||
yield Footer()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
table = self.query_one("#session-table", DataTable)
|
||||
table.add_columns("ID", "Created", "Updated", "Runs")
|
||||
table.cursor_type = "row"
|
||||
|
||||
def update_sessions(self, sessions: List[dict]) -> None:
|
||||
"""Update the sessions display."""
|
||||
table = self.query_one("#session-table", DataTable)
|
||||
table.clear()
|
||||
|
||||
for session in sessions:
|
||||
table.add_row(
|
||||
session.get("session_id", "")[:8],
|
||||
session.get("created_at", "-"),
|
||||
session.get("updated_at", "-"),
|
||||
str(session.get("run_count", 0)),
|
||||
key=session.get("session_id"),
|
||||
)
|
||||
|
||||
def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
|
||||
"""Handle row selection."""
|
||||
if event.row_key:
|
||||
self._selected_session_id = str(event.row_key.value)
|
||||
|
||||
def action_back(self) -> None:
|
||||
"""Go back."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_resume(self) -> None:
|
||||
"""Resume selected session."""
|
||||
if self._selected_session_id:
|
||||
self.post_message(self.SessionSelected(self._selected_session_id))
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_delete(self) -> None:
|
||||
"""Delete selected session."""
|
||||
if self._selected_session_id:
|
||||
self.post_message(self.SessionDeleted(self._selected_session_id))
|
||||
|
||||
else:
|
||||
class SessionScreen:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,148 @@
|
||||
"""
|
||||
Settings Screen for PraisonAI TUI.
|
||||
|
||||
Configuration interface.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
try:
|
||||
from textual.screen import Screen
|
||||
from textual.containers import Vertical
|
||||
from textual.widgets import Static, Footer, Input, Select, Switch, Button
|
||||
from textual.binding import Binding
|
||||
from textual.message import Message
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Screen = object
|
||||
Message = object
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class SettingsScreen(Screen):
|
||||
"""
|
||||
Settings screen.
|
||||
|
||||
Allows configuration of:
|
||||
- Model selection
|
||||
- Queue settings
|
||||
- Display preferences
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "back", "Back", show=True),
|
||||
Binding("s", "save", "Save", show=True),
|
||||
]
|
||||
|
||||
DEFAULT_CSS = """
|
||||
SettingsScreen {
|
||||
background: $surface;
|
||||
}
|
||||
|
||||
SettingsScreen #settings-header {
|
||||
height: 3;
|
||||
background: $primary;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
SettingsScreen .setting-row {
|
||||
height: 3;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
SettingsScreen .setting-label {
|
||||
width: 20;
|
||||
}
|
||||
|
||||
SettingsScreen Input {
|
||||
width: 40;
|
||||
}
|
||||
"""
|
||||
|
||||
class SettingsSaved(Message):
|
||||
def __init__(self, settings: dict):
|
||||
self.settings = settings
|
||||
super().__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
current_settings: Optional[dict] = None,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._settings = current_settings or {}
|
||||
|
||||
def compose(self):
|
||||
"""Compose the screen."""
|
||||
yield Static("Settings", id="settings-header")
|
||||
|
||||
with Vertical(id="settings-form"):
|
||||
# Model setting
|
||||
yield Static("Model:", classes="setting-label")
|
||||
yield Input(
|
||||
value=self._settings.get("model", "gpt-4o-mini"),
|
||||
id="setting-model",
|
||||
placeholder="Model name"
|
||||
)
|
||||
|
||||
# Max concurrent runs
|
||||
yield Static("Max Concurrent:", classes="setting-label")
|
||||
yield Input(
|
||||
value=str(self._settings.get("max_concurrent", 4)),
|
||||
id="setting-max-concurrent",
|
||||
placeholder="4"
|
||||
)
|
||||
|
||||
# Auto-save interval
|
||||
yield Static("Autosave (sec):", classes="setting-label")
|
||||
yield Input(
|
||||
value=str(self._settings.get("autosave_interval", 30)),
|
||||
id="setting-autosave",
|
||||
placeholder="30"
|
||||
)
|
||||
|
||||
# Buttons
|
||||
yield Button("Save", id="btn-save", variant="primary")
|
||||
yield Button("Cancel", id="btn-cancel", variant="default")
|
||||
|
||||
yield Footer()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
"""Handle button press."""
|
||||
if event.button.id == "btn-save":
|
||||
self.action_save()
|
||||
elif event.button.id == "btn-cancel":
|
||||
self.action_back()
|
||||
|
||||
def action_back(self) -> None:
|
||||
"""Go back."""
|
||||
self.app.pop_screen()
|
||||
|
||||
def action_save(self) -> None:
|
||||
"""Save settings."""
|
||||
try:
|
||||
model_input = self.query_one("#setting-model", Input)
|
||||
max_concurrent_input = self.query_one("#setting-max-concurrent", Input)
|
||||
autosave_input = self.query_one("#setting-autosave", Input)
|
||||
|
||||
settings = {
|
||||
"model": model_input.value,
|
||||
"max_concurrent": int(max_concurrent_input.value or "4"),
|
||||
"autosave_interval": int(autosave_input.value or "30"),
|
||||
}
|
||||
|
||||
self.post_message(self.SettingsSaved(settings))
|
||||
self.app.pop_screen()
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
else:
|
||||
class SettingsScreen:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,181 @@
|
||||
"""
|
||||
Session Store for PraisonAI TUI.
|
||||
|
||||
Maintains chat history in memory for session persistence.
|
||||
This allows the AI to remember previous messages within a session.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Optional
|
||||
from datetime import datetime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatMessage:
|
||||
"""A single chat message."""
|
||||
role: str # "user" or "assistant"
|
||||
content: str
|
||||
timestamp: datetime = field(default_factory=datetime.now)
|
||||
run_id: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
"""A chat session with history."""
|
||||
session_id: str
|
||||
messages: List[ChatMessage] = field(default_factory=list)
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
|
||||
def add_message(self, role: str, content: str, run_id: Optional[str] = None) -> None:
|
||||
"""Add a message to the session."""
|
||||
self.messages.append(ChatMessage(
|
||||
role=role,
|
||||
content=content,
|
||||
run_id=run_id,
|
||||
))
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def get_history_for_llm(self, max_messages: int = 50) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Get chat history in LLM-compatible format.
|
||||
|
||||
Returns list of {"role": "user/assistant", "content": "..."} dicts.
|
||||
"""
|
||||
# Get last N messages
|
||||
recent = self.messages[-max_messages:] if len(self.messages) > max_messages else self.messages
|
||||
return [{"role": msg.role, "content": msg.content} for msg in recent]
|
||||
|
||||
@property
|
||||
def run_count(self) -> int:
|
||||
"""Number of user messages (runs) in this session."""
|
||||
return sum(1 for msg in self.messages if msg.role == "user")
|
||||
|
||||
|
||||
class SessionStore:
|
||||
"""
|
||||
In-memory session store for TUI.
|
||||
|
||||
Maintains chat history across agent runs within a session.
|
||||
"""
|
||||
|
||||
def __init__(self, max_sessions: int = 100, max_messages_per_session: int = 1000):
|
||||
"""
|
||||
Initialize session store.
|
||||
|
||||
Args:
|
||||
max_sessions: Maximum number of sessions to keep in memory.
|
||||
max_messages_per_session: Maximum messages per session.
|
||||
"""
|
||||
self._sessions: Dict[str, Session] = {}
|
||||
self._max_sessions = max_sessions
|
||||
self._max_messages = max_messages_per_session
|
||||
|
||||
def get_or_create_session(self, session_id: str) -> Session:
|
||||
"""Get existing session or create new one."""
|
||||
if session_id not in self._sessions:
|
||||
self._sessions[session_id] = Session(session_id=session_id)
|
||||
self._cleanup_old_sessions()
|
||||
logger.debug(f"Created new session: {session_id}")
|
||||
return self._sessions[session_id]
|
||||
|
||||
def add_user_message(self, session_id: str, content: str, run_id: Optional[str] = None) -> None:
|
||||
"""Add a user message to the session."""
|
||||
session = self.get_or_create_session(session_id)
|
||||
session.add_message("user", content, run_id)
|
||||
self._trim_session(session)
|
||||
logger.debug(f"Added user message to session {session_id}: {content[:50]}...")
|
||||
|
||||
def add_assistant_message(self, session_id: str, content: str, run_id: Optional[str] = None) -> None:
|
||||
"""Add an assistant message to the session."""
|
||||
session = self.get_or_create_session(session_id)
|
||||
session.add_message("assistant", content, run_id)
|
||||
self._trim_session(session)
|
||||
logger.debug(f"Added assistant message to session {session_id}: {content[:50]}...")
|
||||
|
||||
def get_chat_history(self, session_id: str, max_messages: int = 50) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Get chat history for a session in LLM-compatible format.
|
||||
|
||||
Args:
|
||||
session_id: The session ID.
|
||||
max_messages: Maximum number of messages to return.
|
||||
|
||||
Returns:
|
||||
List of {"role": "user/assistant", "content": "..."} dicts.
|
||||
"""
|
||||
if session_id not in self._sessions:
|
||||
return []
|
||||
return self._sessions[session_id].get_history_for_llm(max_messages)
|
||||
|
||||
def get_session(self, session_id: str) -> Optional[Session]:
|
||||
"""Get a session by ID."""
|
||||
return self._sessions.get(session_id)
|
||||
|
||||
def list_sessions(self, limit: int = 50) -> List[Dict]:
|
||||
"""List all sessions with metadata."""
|
||||
sessions = sorted(
|
||||
self._sessions.values(),
|
||||
key=lambda s: s.updated_at,
|
||||
reverse=True
|
||||
)[:limit]
|
||||
|
||||
return [
|
||||
{
|
||||
"session_id": s.session_id,
|
||||
"created_at": s.created_at.strftime("%Y-%m-%d %H:%M"),
|
||||
"updated_at": s.updated_at.strftime("%Y-%m-%d %H:%M"),
|
||||
"run_count": s.run_count,
|
||||
}
|
||||
for s in sessions
|
||||
]
|
||||
|
||||
def delete_session(self, session_id: str) -> bool:
|
||||
"""Delete a session."""
|
||||
if session_id in self._sessions:
|
||||
del self._sessions[session_id]
|
||||
logger.debug(f"Deleted session: {session_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
def clear_session(self, session_id: str) -> None:
|
||||
"""Clear all messages from a session."""
|
||||
if session_id in self._sessions:
|
||||
self._sessions[session_id].messages.clear()
|
||||
self._sessions[session_id].updated_at = datetime.now()
|
||||
logger.debug(f"Cleared session: {session_id}")
|
||||
|
||||
def _trim_session(self, session: Session) -> None:
|
||||
"""Trim session to max messages."""
|
||||
if len(session.messages) > self._max_messages:
|
||||
# Keep the most recent messages
|
||||
session.messages = session.messages[-self._max_messages:]
|
||||
|
||||
def _cleanup_old_sessions(self) -> None:
|
||||
"""Remove oldest sessions if over limit."""
|
||||
if len(self._sessions) > self._max_sessions:
|
||||
# Sort by updated_at and remove oldest
|
||||
sorted_sessions = sorted(
|
||||
self._sessions.items(),
|
||||
key=lambda x: x[1].updated_at
|
||||
)
|
||||
# Remove oldest sessions
|
||||
to_remove = len(self._sessions) - self._max_sessions
|
||||
for session_id, _ in sorted_sessions[:to_remove]:
|
||||
del self._sessions[session_id]
|
||||
logger.debug(f"Removed old session: {session_id}")
|
||||
|
||||
|
||||
# Global session store instance for TUI
|
||||
_session_store: Optional[SessionStore] = None
|
||||
|
||||
|
||||
def get_session_store() -> SessionStore:
|
||||
"""Get the global session store instance."""
|
||||
global _session_store
|
||||
if _session_store is None:
|
||||
_session_store = SessionStore()
|
||||
return _session_store
|
||||
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
TUI Widgets for PraisonAI.
|
||||
|
||||
Reusable UI components for the TUI application.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .chat import ChatWidget
|
||||
from .composer import ComposerWidget
|
||||
from .status import StatusWidget
|
||||
from .queue_panel import QueuePanelWidget
|
||||
from .tool_panel import ToolPanelWidget
|
||||
|
||||
_lazy_cache = {}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load widgets."""
|
||||
global _lazy_cache
|
||||
|
||||
if name in _lazy_cache:
|
||||
return _lazy_cache[name]
|
||||
|
||||
if name == "ChatWidget":
|
||||
from .chat import ChatWidget
|
||||
_lazy_cache[name] = ChatWidget
|
||||
return ChatWidget
|
||||
elif name == "ComposerWidget":
|
||||
from .composer import ComposerWidget
|
||||
_lazy_cache[name] = ComposerWidget
|
||||
return ComposerWidget
|
||||
elif name == "StatusWidget":
|
||||
from .status import StatusWidget
|
||||
_lazy_cache[name] = StatusWidget
|
||||
return StatusWidget
|
||||
elif name == "QueuePanelWidget":
|
||||
from .queue_panel import QueuePanelWidget
|
||||
_lazy_cache[name] = QueuePanelWidget
|
||||
return QueuePanelWidget
|
||||
elif name == "ToolPanelWidget":
|
||||
from .tool_panel import ToolPanelWidget
|
||||
_lazy_cache[name] = ToolPanelWidget
|
||||
return ToolPanelWidget
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ChatWidget",
|
||||
"ComposerWidget",
|
||||
"StatusWidget",
|
||||
"QueuePanelWidget",
|
||||
"ToolPanelWidget",
|
||||
]
|
||||
@@ -0,0 +1,263 @@
|
||||
"""
|
||||
Chat Widget for PraisonAI TUI.
|
||||
|
||||
Displays chat history with streaming support.
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
from dataclasses import dataclass, field
|
||||
import time
|
||||
|
||||
try:
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static
|
||||
from textual.containers import VerticalScroll
|
||||
from textual.message import Message
|
||||
from rich.markdown import Markdown
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Widget = object
|
||||
Message = object
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatMessage:
|
||||
"""A single chat message."""
|
||||
role: str # "user", "assistant", "system", "tool"
|
||||
content: str
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
run_id: Optional[str] = None
|
||||
agent_name: Optional[str] = None
|
||||
is_streaming: bool = False
|
||||
|
||||
@property
|
||||
def display_role(self) -> str:
|
||||
"""Get display name for role."""
|
||||
if self.role == "user":
|
||||
return "You"
|
||||
elif self.role == "assistant":
|
||||
return self.agent_name or "Assistant"
|
||||
elif self.role == "system":
|
||||
return "System"
|
||||
elif self.role == "tool":
|
||||
return "Tool"
|
||||
return self.role.title()
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class ChatWidget(VerticalScroll):
|
||||
"""
|
||||
Widget for displaying chat history with scrollbar.
|
||||
|
||||
Uses VerticalScroll for proper scrollbar support.
|
||||
Messages are mounted directly to this container.
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ChatWidget {
|
||||
height: 1fr;
|
||||
border: solid $primary;
|
||||
background: $surface;
|
||||
padding: 0 1;
|
||||
overflow-y: auto;
|
||||
}
|
||||
|
||||
ChatWidget .message-user {
|
||||
background: $primary-darken-2;
|
||||
margin: 1 0;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
ChatWidget .message-assistant {
|
||||
background: $surface-darken-1;
|
||||
margin: 1 0;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
ChatWidget .message-system {
|
||||
background: $warning-darken-3;
|
||||
margin: 1 0;
|
||||
padding: 1;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
ChatWidget .message-streaming {
|
||||
border: dashed $accent;
|
||||
}
|
||||
"""
|
||||
|
||||
class MessageAdded(Message):
|
||||
"""Event when a message is added."""
|
||||
def __init__(self, message: ChatMessage):
|
||||
self.message = message
|
||||
super().__init__()
|
||||
|
||||
class StreamingUpdate(Message):
|
||||
"""Event when streaming content updates."""
|
||||
def __init__(self, run_id: str, content: str):
|
||||
self.run_id = run_id
|
||||
self.content = content
|
||||
super().__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_messages: int = 1000,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._messages: List[ChatMessage] = []
|
||||
self._max_messages = max_messages
|
||||
self._streaming_widgets: dict = {}
|
||||
|
||||
def compose(self):
|
||||
"""Compose the widget - no inner container needed."""
|
||||
# Messages are mounted directly to this VerticalScroll
|
||||
# Must yield from empty iterable (not return None)
|
||||
yield from ()
|
||||
|
||||
async def add_message(self, message: ChatMessage) -> None:
|
||||
"""Add a message to the chat.
|
||||
|
||||
NEW BEHAVIOR: Messages render NEWEST at TOP.
|
||||
This guarantees visibility without relying on scroll.
|
||||
"""
|
||||
self._messages.append(message)
|
||||
|
||||
# Trim old messages if needed
|
||||
if len(self._messages) > self._max_messages:
|
||||
self._messages = self._messages[-self._max_messages:]
|
||||
|
||||
# Render message and scroll to show it
|
||||
await self._render_message(message)
|
||||
|
||||
self.post_message(self.MessageAdded(message))
|
||||
|
||||
async def _render_message(self, message: ChatMessage) -> None:
|
||||
"""Render a message and scroll to show it."""
|
||||
# Create role label
|
||||
role_style = {
|
||||
"user": "bold cyan",
|
||||
"assistant": "bold green",
|
||||
"system": "bold yellow",
|
||||
"tool": "bold magenta",
|
||||
}.get(message.role, "bold")
|
||||
|
||||
role_text = Text(f"{message.display_role}:", style=role_style)
|
||||
|
||||
# Create content
|
||||
try:
|
||||
content = Markdown(message.content) if message.content else Text("")
|
||||
except Exception:
|
||||
content = Text(message.content)
|
||||
|
||||
# Create panel
|
||||
css_class = f"message-{message.role}"
|
||||
if message.is_streaming:
|
||||
css_class += " message-streaming"
|
||||
|
||||
widget_id = f"msg-{message.run_id or id(message)}"
|
||||
|
||||
panel = Static(
|
||||
Panel(content, title=str(role_text), border_style=role_style),
|
||||
id=widget_id,
|
||||
classes=css_class,
|
||||
)
|
||||
|
||||
# Mount directly to this VerticalScroll container
|
||||
await self.mount(panel)
|
||||
|
||||
# Scroll to end to show new message
|
||||
self.scroll_end(animate=False)
|
||||
|
||||
if message.is_streaming:
|
||||
self._streaming_widgets[message.run_id] = widget_id
|
||||
|
||||
async def update_streaming(self, run_id: str, content: str) -> None:
|
||||
"""Update a streaming message."""
|
||||
if run_id not in self._streaming_widgets:
|
||||
return
|
||||
|
||||
widget_id = self._streaming_widgets[run_id]
|
||||
|
||||
try:
|
||||
widget = self.query_one(f"#{widget_id}", Static)
|
||||
|
||||
# Find the message
|
||||
for msg in self._messages:
|
||||
if msg.run_id == run_id:
|
||||
msg.content = content
|
||||
break
|
||||
|
||||
# Update content
|
||||
try:
|
||||
rendered = Markdown(content + " ▌")
|
||||
except Exception:
|
||||
rendered = Text(content + " ▌")
|
||||
|
||||
widget.update(Panel(rendered, title="Assistant", border_style="bold green"))
|
||||
|
||||
# Scroll to end to keep streaming content visible
|
||||
self.scroll_end(animate=False)
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def complete_streaming(self, run_id: str, final_content: str) -> None:
|
||||
"""Complete a streaming message."""
|
||||
if run_id not in self._streaming_widgets:
|
||||
return
|
||||
|
||||
widget_id = self._streaming_widgets.pop(run_id)
|
||||
|
||||
try:
|
||||
widget = self.query_one(f"#{widget_id}", Static)
|
||||
|
||||
# Update message
|
||||
for msg in self._messages:
|
||||
if msg.run_id == run_id:
|
||||
msg.content = final_content
|
||||
msg.is_streaming = False
|
||||
break
|
||||
|
||||
# Update content without cursor
|
||||
try:
|
||||
rendered = Markdown(final_content)
|
||||
except Exception:
|
||||
rendered = Text(final_content)
|
||||
|
||||
widget.update(Panel(rendered, title="Assistant", border_style="bold green"))
|
||||
widget.remove_class("message-streaming")
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""Clear all messages."""
|
||||
self._messages.clear()
|
||||
self._streaming_widgets.clear()
|
||||
|
||||
# Remove all children from this container
|
||||
await self.remove_children()
|
||||
|
||||
@property
|
||||
def messages(self) -> List[ChatMessage]:
|
||||
"""Get all messages."""
|
||||
return self._messages.copy()
|
||||
|
||||
@property
|
||||
def message_count(self) -> int:
|
||||
"""Get message count."""
|
||||
return len(self._messages)
|
||||
|
||||
else:
|
||||
class ChatWidget:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,263 @@
|
||||
"""
|
||||
Command Popup Widget for PraisonAI TUI.
|
||||
|
||||
Shows a searchable list of available commands when user types backslash.
|
||||
Inspired by Claude Code's command discovery UX.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static, Input, OptionList
|
||||
from textual.widgets.option_list import Option
|
||||
from textual.containers import Vertical, Container
|
||||
from textual.reactive import reactive
|
||||
from textual.message import Message
|
||||
from textual import events
|
||||
from rich.text import Text
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Widget = object
|
||||
Message = object
|
||||
|
||||
|
||||
@dataclass
|
||||
class CommandInfo:
|
||||
"""Information about a command for display."""
|
||||
name: str
|
||||
description: str
|
||||
aliases: List[str]
|
||||
category: str = "general"
|
||||
|
||||
@property
|
||||
def display_name(self) -> str:
|
||||
"""Get display name with aliases."""
|
||||
if self.aliases:
|
||||
return f"{self.name} ({', '.join(self.aliases)})"
|
||||
return self.name
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class CommandPopupWidget(Container):
|
||||
"""
|
||||
Popup widget for command discovery and selection.
|
||||
|
||||
Features:
|
||||
- Shows all available commands
|
||||
- Searchable/filterable
|
||||
- Keyboard navigation
|
||||
- Category grouping
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
CommandPopupWidget {
|
||||
layer: popup;
|
||||
width: 60;
|
||||
height: auto;
|
||||
max-height: 20;
|
||||
background: $surface;
|
||||
border: solid $primary;
|
||||
padding: 1;
|
||||
margin: 0 0 0 2;
|
||||
}
|
||||
|
||||
CommandPopupWidget #popup-title {
|
||||
height: 1;
|
||||
background: $primary;
|
||||
color: $text;
|
||||
text-align: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
CommandPopupWidget #popup-search {
|
||||
height: 3;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
CommandPopupWidget #popup-list {
|
||||
height: auto;
|
||||
max-height: 12;
|
||||
background: $surface-darken-1;
|
||||
}
|
||||
|
||||
CommandPopupWidget .command-item {
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
CommandPopupWidget .command-item:hover {
|
||||
background: $primary-darken-1;
|
||||
}
|
||||
|
||||
CommandPopupWidget #popup-hint {
|
||||
height: 1;
|
||||
color: $text-muted;
|
||||
text-align: center;
|
||||
margin-top: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
class CommandSelected(Message):
|
||||
"""Event when a command is selected."""
|
||||
def __init__(self, command: str, args: str = ""):
|
||||
self.command = command
|
||||
self.args = args
|
||||
super().__init__()
|
||||
|
||||
class Dismissed(Message):
|
||||
"""Event when popup is dismissed."""
|
||||
pass
|
||||
|
||||
# Reactive properties
|
||||
filter_text: reactive[str] = reactive("")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
commands: Optional[List[CommandInfo]] = None,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._commands = commands or []
|
||||
self._filtered_commands: List[CommandInfo] = []
|
||||
|
||||
def compose(self):
|
||||
"""Compose the widget."""
|
||||
yield Static("Commands", id="popup-title")
|
||||
yield Input(placeholder="Type to filter...", id="popup-search")
|
||||
yield OptionList(id="popup-list")
|
||||
yield Static("↑↓ Navigate • Enter Select • Esc Cancel", id="popup-hint")
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
self._update_list()
|
||||
# Focus the search input
|
||||
search = self.query_one("#popup-search", Input)
|
||||
search.focus()
|
||||
|
||||
def set_commands(self, commands: List[CommandInfo]) -> None:
|
||||
"""Set the available commands."""
|
||||
self._commands = commands
|
||||
self._update_list()
|
||||
|
||||
def watch_filter_text(self, value: str) -> None:
|
||||
"""React to filter text changes."""
|
||||
self._update_list()
|
||||
|
||||
def _update_list(self) -> None:
|
||||
"""Update the command list based on filter."""
|
||||
filter_lower = self.filter_text.lower()
|
||||
|
||||
if filter_lower:
|
||||
self._filtered_commands = [
|
||||
cmd for cmd in self._commands
|
||||
if (filter_lower in cmd.name.lower() or
|
||||
filter_lower in cmd.description.lower() or
|
||||
any(filter_lower in alias.lower() for alias in cmd.aliases))
|
||||
]
|
||||
else:
|
||||
self._filtered_commands = self._commands.copy()
|
||||
|
||||
# Update the option list
|
||||
try:
|
||||
option_list = self.query_one("#popup-list", OptionList)
|
||||
option_list.clear_options()
|
||||
|
||||
for cmd in self._filtered_commands:
|
||||
# Create rich text for the option
|
||||
text = Text()
|
||||
text.append(f"/{cmd.name}", style="bold cyan")
|
||||
if cmd.aliases:
|
||||
text.append(f" ({', '.join(cmd.aliases)})", style="dim")
|
||||
text.append(f"\n {cmd.description}", style="")
|
||||
|
||||
option_list.add_option(Option(text, id=cmd.name))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def on_input_changed(self, event: Input.Changed) -> None:
|
||||
"""Handle search input changes."""
|
||||
if event.input.id == "popup-search":
|
||||
self.filter_text = event.value
|
||||
|
||||
def on_input_submitted(self, event: Input.Submitted) -> None:
|
||||
"""Handle Enter in search input."""
|
||||
if event.input.id == "popup-search":
|
||||
# Select first matching command
|
||||
if self._filtered_commands:
|
||||
self._select_command(self._filtered_commands[0].name)
|
||||
|
||||
def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None:
|
||||
"""Handle command selection from list."""
|
||||
if event.option.id:
|
||||
self._select_command(str(event.option.id))
|
||||
|
||||
def on_key(self, event: events.Key) -> None:
|
||||
"""Handle key events."""
|
||||
if event.key == "escape":
|
||||
self.post_message(self.Dismissed())
|
||||
event.stop()
|
||||
elif event.key == "down":
|
||||
# Move focus to list if in search
|
||||
try:
|
||||
option_list = self.query_one("#popup-list", OptionList)
|
||||
option_list.focus()
|
||||
except Exception:
|
||||
pass
|
||||
elif event.key == "up":
|
||||
# Move focus to search if at top of list
|
||||
try:
|
||||
search = self.query_one("#popup-search", Input)
|
||||
search.focus()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _select_command(self, command_name: str) -> None:
|
||||
"""Select a command and emit event."""
|
||||
self.post_message(self.CommandSelected(command_name))
|
||||
|
||||
def focus_search(self) -> None:
|
||||
"""Focus the search input."""
|
||||
try:
|
||||
search = self.query_one("#popup-search", Input)
|
||||
search.focus()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
else:
|
||||
class CommandPopupWidget:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
|
||||
|
||||
# Default commands list (matches slash_commands.py registry)
|
||||
DEFAULT_COMMANDS = [
|
||||
CommandInfo("help", "Show help for commands", ["h", "?"]),
|
||||
CommandInfo("clear", "Clear conversation history", ["reset"]),
|
||||
CommandInfo("model", "Show or change the current model", ["m"]),
|
||||
CommandInfo("cost", "Show session cost and token usage", ["usage", "stats"]),
|
||||
CommandInfo("tokens", "Show token usage breakdown", []),
|
||||
CommandInfo("queue", "Show queue status", ["q"]),
|
||||
CommandInfo("cancel", "Cancel current operation", ["c"]),
|
||||
CommandInfo("settings", "Show current settings", ["set"]),
|
||||
CommandInfo("sessions", "Browse saved sessions", ["sess"]),
|
||||
CommandInfo("tools", "Toggle tools panel", ["t"]),
|
||||
CommandInfo("exit", "Exit the TUI", ["quit", "q"]),
|
||||
CommandInfo("plan", "Create an execution plan", []),
|
||||
CommandInfo("undo", "Undo the last change", []),
|
||||
CommandInfo("diff", "Show git diff of changes", []),
|
||||
CommandInfo("commit", "Commit changes with AI message", []),
|
||||
CommandInfo("map", "Show repository map", ["repo"]),
|
||||
]
|
||||
|
||||
|
||||
def get_default_commands() -> List[CommandInfo]:
|
||||
"""Get the default command list."""
|
||||
return DEFAULT_COMMANDS.copy()
|
||||
@@ -0,0 +1,292 @@
|
||||
"""
|
||||
Composer Widget for PraisonAI TUI.
|
||||
|
||||
Input area for composing messages with slash command support.
|
||||
"""
|
||||
|
||||
from typing import Callable, List, Optional
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
try:
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import TextArea, Input
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.message import Message
|
||||
from textual.binding import Binding
|
||||
from textual import events
|
||||
from rich.text import Text
|
||||
TEXTUAL_AVAILABLE = True
|
||||
|
||||
class SubmitTextArea(TextArea):
|
||||
"""Custom TextArea that submits on Enter and inserts newline on Shift+Enter."""
|
||||
|
||||
class SubmitRequested(Message):
|
||||
"""Message sent when Enter is pressed."""
|
||||
pass
|
||||
|
||||
async def _on_key(self, event: events.Key) -> None:
|
||||
"""Handle key events before TextArea processes them."""
|
||||
if event.key == "enter":
|
||||
# Don't let TextArea handle Enter - we want to submit
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
self.post_message(self.SubmitRequested())
|
||||
return
|
||||
elif event.key == "shift+enter":
|
||||
# Insert newline on Shift+Enter
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
self.insert("\n")
|
||||
return
|
||||
# Let TextArea handle other keys
|
||||
await super()._on_key(event)
|
||||
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Widget = object
|
||||
Message = object
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class ComposerWidget(Vertical):
|
||||
"""
|
||||
Input composer for chat messages.
|
||||
|
||||
Features:
|
||||
- Multi-line input
|
||||
- Slash command detection
|
||||
- @ mention support
|
||||
- History navigation
|
||||
- Send on Enter or Ctrl+Enter
|
||||
"""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("enter", "submit", "Send", show=True),
|
||||
Binding("escape", "cancel", "Cancel", show=False),
|
||||
Binding("up", "history_prev", "Previous", show=False),
|
||||
Binding("down", "history_next", "Next", show=False),
|
||||
]
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ComposerWidget {
|
||||
height: auto;
|
||||
max-height: 10;
|
||||
border: solid $primary;
|
||||
background: $surface;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
ComposerWidget TextArea {
|
||||
height: auto;
|
||||
min-height: 3;
|
||||
max-height: 8;
|
||||
}
|
||||
|
||||
ComposerWidget .hint {
|
||||
color: $text-muted;
|
||||
text-style: italic;
|
||||
}
|
||||
"""
|
||||
|
||||
class Submitted(Message):
|
||||
"""Event when message is submitted."""
|
||||
def __init__(self, content: str, is_command: bool = False):
|
||||
self.content = content
|
||||
self.is_command = is_command
|
||||
super().__init__()
|
||||
|
||||
class CommandDetected(Message):
|
||||
"""Event when a slash command is detected."""
|
||||
def __init__(self, command: str, args: str):
|
||||
self.command = command
|
||||
self.args = args
|
||||
super().__init__()
|
||||
|
||||
class BackslashTyped(Message):
|
||||
"""Event when backslash is typed - triggers command popup."""
|
||||
pass
|
||||
|
||||
class AtTyped(Message):
|
||||
"""Event when @ is typed - triggers file autocomplete."""
|
||||
def __init__(self, query: str, cursor_pos: int):
|
||||
self.query = query
|
||||
self.cursor_pos = cursor_pos
|
||||
super().__init__()
|
||||
|
||||
class SlashTyped(Message):
|
||||
"""Event when / is typed - triggers command discovery."""
|
||||
def __init__(self, query: str):
|
||||
self.query = query
|
||||
super().__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
placeholder: str = "Type your message... (Enter to send, Shift+Enter for newline, /help for commands)",
|
||||
multiline: bool = True,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._placeholder = placeholder
|
||||
self._multiline = multiline
|
||||
self._history: List[str] = []
|
||||
self._history_index = -1
|
||||
self._current_input = ""
|
||||
self._is_processing = False
|
||||
|
||||
def compose(self):
|
||||
"""Compose the widget."""
|
||||
yield SubmitTextArea(
|
||||
"",
|
||||
id="composer-input",
|
||||
language=None,
|
||||
)
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount event."""
|
||||
text_area = self.query_one("#composer-input", SubmitTextArea)
|
||||
text_area.focus()
|
||||
|
||||
def on_submit_text_area_submit_requested(self, event: SubmitTextArea.SubmitRequested) -> None:
|
||||
"""Handle submit request from TextArea."""
|
||||
self.action_submit()
|
||||
|
||||
async def on_text_area_changed(self, event: TextArea.Changed) -> None:
|
||||
"""Handle text changes."""
|
||||
content = event.text_area.text
|
||||
|
||||
# Detect backslash - auto-trigger command popup
|
||||
if content == "\\":
|
||||
self.post_message(self.BackslashTyped())
|
||||
return
|
||||
|
||||
# Detect @ for file autocomplete
|
||||
if "@" in content:
|
||||
# Find the @ and extract query
|
||||
at_pos = content.rfind("@")
|
||||
if at_pos >= 0:
|
||||
query = content[at_pos + 1:]
|
||||
# Only trigger if no space after @
|
||||
if " " not in query or query.endswith(" ") is False:
|
||||
self.post_message(self.AtTyped(query, at_pos))
|
||||
|
||||
# Detect slash commands - auto-trigger command discovery
|
||||
if content.startswith("/"):
|
||||
parts = content[1:].split(maxsplit=1)
|
||||
if parts:
|
||||
command = parts[0]
|
||||
args = parts[1] if len(parts) > 1 else ""
|
||||
self.post_message(self.CommandDetected(command, args))
|
||||
# Also trigger slash popup for discovery
|
||||
self.post_message(self.SlashTyped(content[1:]))
|
||||
|
||||
def on_key(self, event: events.Key) -> None:
|
||||
"""Handle key events at widget level (backup handler)."""
|
||||
# These are handled by SubmitTextArea now, but keep as backup
|
||||
pass
|
||||
|
||||
def action_submit(self) -> None:
|
||||
"""Submit the current input.
|
||||
|
||||
Note: We allow submitting even while processing to enable
|
||||
queueing multiple messages.
|
||||
"""
|
||||
text_area = self.query_one("#composer-input", TextArea)
|
||||
content = text_area.text.strip()
|
||||
|
||||
if not content:
|
||||
return
|
||||
|
||||
# Add to history
|
||||
if content and (not self._history or self._history[-1] != content):
|
||||
self._history.append(content)
|
||||
self._history_index = -1
|
||||
|
||||
# Check if it's a command (/ or : prefix)
|
||||
is_command = content.startswith("/") or content.startswith(":")
|
||||
|
||||
# Clear input
|
||||
text_area.clear()
|
||||
|
||||
# Emit event
|
||||
self.post_message(self.Submitted(content, is_command))
|
||||
|
||||
def action_cancel(self) -> None:
|
||||
"""Cancel current input."""
|
||||
text_area = self.query_one("#composer-input", TextArea)
|
||||
text_area.clear()
|
||||
self._history_index = -1
|
||||
|
||||
def action_history_prev(self) -> None:
|
||||
"""Navigate to previous history item."""
|
||||
if not self._history:
|
||||
return
|
||||
|
||||
text_area = self.query_one("#composer-input", TextArea)
|
||||
|
||||
if self._history_index == -1:
|
||||
self._current_input = text_area.text
|
||||
self._history_index = len(self._history) - 1
|
||||
elif self._history_index > 0:
|
||||
self._history_index -= 1
|
||||
|
||||
text_area.clear()
|
||||
text_area.insert(self._history[self._history_index])
|
||||
|
||||
def action_history_next(self) -> None:
|
||||
"""Navigate to next history item."""
|
||||
if self._history_index == -1:
|
||||
return
|
||||
|
||||
text_area = self.query_one("#composer-input", TextArea)
|
||||
|
||||
if self._history_index < len(self._history) - 1:
|
||||
self._history_index += 1
|
||||
text_area.clear()
|
||||
text_area.insert(self._history[self._history_index])
|
||||
else:
|
||||
self._history_index = -1
|
||||
text_area.clear()
|
||||
text_area.insert(self._current_input)
|
||||
|
||||
def set_processing(self, processing: bool) -> None:
|
||||
"""Set processing state.
|
||||
|
||||
Note: We no longer disable the text area during processing
|
||||
to allow users to queue additional messages while waiting.
|
||||
"""
|
||||
self._is_processing = processing
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
"""Get current text."""
|
||||
text_area = self.query_one("#composer-input", TextArea)
|
||||
return text_area.text
|
||||
|
||||
@text.setter
|
||||
def text(self, value: str) -> None:
|
||||
"""Set current text."""
|
||||
text_area = self.query_one("#composer-input", TextArea)
|
||||
text_area.clear()
|
||||
text_area.insert(value)
|
||||
|
||||
def focus_input(self) -> None:
|
||||
"""Focus the input area."""
|
||||
text_area = self.query_one("#composer-input", TextArea)
|
||||
text_area.focus()
|
||||
|
||||
def set_text(self, value: str) -> None:
|
||||
"""Set the text content and move cursor to end."""
|
||||
text_area = self.query_one("#composer-input", SubmitTextArea)
|
||||
text_area.clear()
|
||||
text_area.insert(value)
|
||||
|
||||
else:
|
||||
class ComposerWidget:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
File Autocomplete Popup Widget for PraisonAI TUI.
|
||||
|
||||
Shows file/directory suggestions when user types @.
|
||||
Reuses FileSearchService from at_mentions.py.
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
from dataclasses import dataclass
|
||||
import os
|
||||
|
||||
try:
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static, Input, OptionList
|
||||
from textual.widgets.option_list import Option
|
||||
from textual.containers import Container
|
||||
from textual.reactive import reactive
|
||||
from textual.message import Message
|
||||
from textual import events
|
||||
from rich.text import Text
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Widget = object
|
||||
Message = object
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileInfo:
|
||||
"""Information about a file for display."""
|
||||
path: str
|
||||
file_type: str # "file" or "directory"
|
||||
score: int = 100
|
||||
|
||||
@property
|
||||
def icon(self) -> str:
|
||||
"""Get icon for file type."""
|
||||
return "📁 " if self.file_type == "directory" else "📄 "
|
||||
|
||||
@property
|
||||
def display(self) -> str:
|
||||
"""Get display string."""
|
||||
return f"{self.icon}{self.path}"
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
# Import FileSearchService from at_mentions
|
||||
try:
|
||||
from ..at_mentions import FileSearchService, detect_at_mention
|
||||
FILE_SEARCH_AVAILABLE = True
|
||||
except ImportError:
|
||||
FILE_SEARCH_AVAILABLE = False
|
||||
FileSearchService = None
|
||||
|
||||
class FilePopupWidget(Container):
|
||||
"""
|
||||
Popup widget for file autocomplete.
|
||||
|
||||
Features:
|
||||
- Shows files/directories matching query
|
||||
- Fuzzy search
|
||||
- Keyboard navigation
|
||||
- Supports ~/ and absolute paths
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
FilePopupWidget {
|
||||
layer: popup;
|
||||
width: 60;
|
||||
height: auto;
|
||||
max-height: 15;
|
||||
background: $surface;
|
||||
border: solid $accent;
|
||||
padding: 1;
|
||||
margin: 0 0 0 2;
|
||||
}
|
||||
|
||||
FilePopupWidget #file-popup-title {
|
||||
height: 1;
|
||||
background: $accent;
|
||||
color: $text;
|
||||
text-align: center;
|
||||
margin-bottom: 1;
|
||||
}
|
||||
|
||||
FilePopupWidget #file-popup-list {
|
||||
height: auto;
|
||||
max-height: 10;
|
||||
background: $surface-darken-1;
|
||||
}
|
||||
|
||||
FilePopupWidget #file-popup-hint {
|
||||
height: 1;
|
||||
color: $text-muted;
|
||||
text-align: center;
|
||||
margin-top: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
class FileSelected(Message):
|
||||
"""Event when a file is selected."""
|
||||
def __init__(self, path: str, at_pos: int):
|
||||
self.path = path
|
||||
self.at_pos = at_pos
|
||||
super().__init__()
|
||||
|
||||
class Dismissed(Message):
|
||||
"""Event when popup is dismissed."""
|
||||
pass
|
||||
|
||||
# Reactive properties
|
||||
filter_text: reactive[str] = reactive("")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root_dir: Optional[str] = None,
|
||||
at_pos: int = 0,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._root_dir = root_dir or os.getcwd()
|
||||
self._at_pos = at_pos
|
||||
self._files: List[FileInfo] = []
|
||||
self._file_service = None
|
||||
if FILE_SEARCH_AVAILABLE and FileSearchService:
|
||||
self._file_service = FileSearchService(self._root_dir)
|
||||
|
||||
def compose(self):
|
||||
"""Compose the widget."""
|
||||
yield Static("📂 Files", id="file-popup-title")
|
||||
yield OptionList(id="file-popup-list")
|
||||
yield Static("↑↓ Navigate • Enter Select • Esc Cancel", id="file-popup-hint")
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
self._update_list()
|
||||
|
||||
def set_query(self, query: str, at_pos: int) -> None:
|
||||
"""Set the search query."""
|
||||
self.filter_text = query
|
||||
self._at_pos = at_pos
|
||||
|
||||
def watch_filter_text(self, value: str) -> None:
|
||||
"""React to filter text changes."""
|
||||
self._update_list()
|
||||
|
||||
def _update_list(self) -> None:
|
||||
"""Update the file list based on filter."""
|
||||
if not self._file_service:
|
||||
return
|
||||
|
||||
# Search files
|
||||
results = self._file_service.search(self.filter_text, max_results=15)
|
||||
|
||||
self._files = [
|
||||
FileInfo(
|
||||
path=r.path,
|
||||
file_type=r.file_type,
|
||||
score=r.score
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
|
||||
# Update the option list
|
||||
try:
|
||||
option_list = self.query_one("#file-popup-list", OptionList)
|
||||
option_list.clear_options()
|
||||
|
||||
for file_info in self._files:
|
||||
text = Text()
|
||||
text.append(file_info.icon, style="bold")
|
||||
text.append(file_info.path, style="cyan" if file_info.file_type == "directory" else "")
|
||||
|
||||
option_list.add_option(Option(text, id=file_info.path))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None:
|
||||
"""Handle file selection from list."""
|
||||
if event.option.id:
|
||||
self._select_file(str(event.option.id))
|
||||
|
||||
def on_key(self, event: events.Key) -> None:
|
||||
"""Handle key events."""
|
||||
if event.key == "escape":
|
||||
self.post_message(self.Dismissed())
|
||||
event.stop()
|
||||
elif event.key == "enter":
|
||||
# Select first item if available
|
||||
if self._files:
|
||||
self._select_file(self._files[0].path)
|
||||
event.stop()
|
||||
|
||||
def _select_file(self, path: str) -> None:
|
||||
"""Select a file and emit event."""
|
||||
self.post_message(self.FileSelected(path, self._at_pos))
|
||||
|
||||
|
||||
else:
|
||||
class FilePopupWidget:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
Queue Panel Widget for PraisonAI TUI.
|
||||
|
||||
Displays queue status and allows queue management.
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static, DataTable, Button
|
||||
from textual.containers import Vertical
|
||||
from textual.reactive import reactive
|
||||
from textual.message import Message
|
||||
from rich.text import Text
|
||||
from rich.panel import Panel
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Widget = object
|
||||
Message = object
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueueItem:
|
||||
"""A queue item for display."""
|
||||
run_id: str
|
||||
agent_name: str
|
||||
input_preview: str
|
||||
state: str
|
||||
priority: str
|
||||
wait_time: str
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class QueuePanelWidget(Vertical):
|
||||
"""
|
||||
Queue panel widget.
|
||||
|
||||
Displays:
|
||||
- Queue summary (queued/running counts)
|
||||
- List of queued runs
|
||||
- Cancel/retry buttons
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
QueuePanelWidget {
|
||||
height: 100%;
|
||||
border: solid $secondary;
|
||||
background: $surface;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
QueuePanelWidget .queue-header {
|
||||
height: 1;
|
||||
background: $secondary;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
QueuePanelWidget DataTable {
|
||||
height: 1fr;
|
||||
}
|
||||
|
||||
QueuePanelWidget .queue-item-queued {
|
||||
color: $warning;
|
||||
}
|
||||
|
||||
QueuePanelWidget .queue-item-running {
|
||||
color: $success;
|
||||
}
|
||||
|
||||
QueuePanelWidget .queue-item-failed {
|
||||
color: $error;
|
||||
}
|
||||
"""
|
||||
|
||||
class RunSelected(Message):
|
||||
"""Event when a run is selected."""
|
||||
def __init__(self, run_id: str):
|
||||
self.run_id = run_id
|
||||
super().__init__()
|
||||
|
||||
class CancelRequested(Message):
|
||||
"""Event when cancel is requested."""
|
||||
def __init__(self, run_id: str):
|
||||
self.run_id = run_id
|
||||
super().__init__()
|
||||
|
||||
class RetryRequested(Message):
|
||||
"""Event when retry is requested."""
|
||||
def __init__(self, run_id: str):
|
||||
self.run_id = run_id
|
||||
super().__init__()
|
||||
|
||||
class DeleteRequested(Message):
|
||||
"""Event when delete is requested."""
|
||||
def __init__(self, run_id: str):
|
||||
self.run_id = run_id
|
||||
super().__init__()
|
||||
|
||||
class EditRequested(Message):
|
||||
"""Event when edit is requested."""
|
||||
def __init__(self, run_id: str, new_content: str):
|
||||
self.run_id = run_id
|
||||
self.new_content = new_content
|
||||
super().__init__()
|
||||
|
||||
queued_count: reactive[int] = reactive(0)
|
||||
running_count: reactive[int] = reactive(0)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._items: List[QueueItem] = []
|
||||
self._selected_run_id: Optional[str] = None
|
||||
|
||||
def compose(self):
|
||||
"""Compose the widget."""
|
||||
yield Static("Queue", id="queue-title", classes="queue-header")
|
||||
yield DataTable(id="queue-table")
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
table = self.query_one("#queue-table", DataTable)
|
||||
table.add_columns("ID", "Agent", "Input", "State", "Priority")
|
||||
table.cursor_type = "row"
|
||||
self._update_header()
|
||||
|
||||
def watch_queued_count(self, value: int) -> None:
|
||||
self._update_header()
|
||||
|
||||
def watch_running_count(self, value: int) -> None:
|
||||
self._update_header()
|
||||
|
||||
def _update_header(self) -> None:
|
||||
"""Update the header text."""
|
||||
try:
|
||||
header = self.query_one("#queue-title", Static)
|
||||
text = Text()
|
||||
text.append("Queue ", style="bold")
|
||||
text.append(f"⏳ {self.queued_count} queued", style="yellow")
|
||||
text.append(" │ ", style="dim")
|
||||
text.append(f"▶ {self.running_count} running", style="green")
|
||||
header.update(text)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def update_items(self, items: List[QueueItem]) -> None:
|
||||
"""Update the queue items."""
|
||||
self._items = items
|
||||
|
||||
table = self.query_one("#queue-table", DataTable)
|
||||
table.clear()
|
||||
|
||||
for item in items:
|
||||
# Style based on state
|
||||
state_style = {
|
||||
"queued": "yellow",
|
||||
"running": "green",
|
||||
"failed": "red",
|
||||
"cancelled": "dim",
|
||||
"succeeded": "cyan",
|
||||
}.get(item.state.lower(), "")
|
||||
|
||||
state_text = Text(item.state, style=state_style)
|
||||
|
||||
# Truncate input preview
|
||||
input_preview = item.input_preview[:30]
|
||||
if len(item.input_preview) > 30:
|
||||
input_preview += "..."
|
||||
|
||||
table.add_row(
|
||||
item.run_id[:8],
|
||||
item.agent_name,
|
||||
input_preview,
|
||||
state_text,
|
||||
item.priority,
|
||||
key=item.run_id,
|
||||
)
|
||||
|
||||
def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
|
||||
"""Handle row selection."""
|
||||
if event.row_key:
|
||||
self._selected_run_id = str(event.row_key.value)
|
||||
self.post_message(self.RunSelected(self._selected_run_id))
|
||||
|
||||
def cancel_selected(self) -> None:
|
||||
"""Cancel the selected run."""
|
||||
if self._selected_run_id:
|
||||
self.post_message(self.CancelRequested(self._selected_run_id))
|
||||
|
||||
def retry_selected(self) -> None:
|
||||
"""Retry the selected run."""
|
||||
if self._selected_run_id:
|
||||
self.post_message(self.RetryRequested(self._selected_run_id))
|
||||
|
||||
def delete_selected(self) -> None:
|
||||
"""Delete the selected run."""
|
||||
if self._selected_run_id:
|
||||
self.post_message(self.DeleteRequested(self._selected_run_id))
|
||||
|
||||
def edit_selected(self, new_content: str) -> None:
|
||||
"""Edit the selected run's input content."""
|
||||
if self._selected_run_id:
|
||||
self.post_message(self.EditRequested(self._selected_run_id, new_content))
|
||||
|
||||
@property
|
||||
def selected_run_id(self) -> Optional[str]:
|
||||
"""Get the selected run ID."""
|
||||
return self._selected_run_id
|
||||
|
||||
else:
|
||||
class QueuePanelWidget:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,167 @@
|
||||
"""
|
||||
Status Widget for PraisonAI TUI.
|
||||
|
||||
Displays status bar with session info, model, tokens, and cost.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static
|
||||
from textual.containers import Horizontal
|
||||
from textual.reactive import reactive
|
||||
from rich.text import Text
|
||||
from rich.table import Table
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Widget = object
|
||||
|
||||
|
||||
@dataclass
|
||||
class StatusInfo:
|
||||
"""Status information."""
|
||||
session_id: str = ""
|
||||
model: str = "gpt-4o-mini"
|
||||
total_tokens: int = 0
|
||||
total_cost: float = 0.0
|
||||
queued_count: int = 0
|
||||
running_count: int = 0
|
||||
is_processing: bool = False
|
||||
status_message: str = ""
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class StatusWidget(Static):
|
||||
"""
|
||||
Status bar widget.
|
||||
|
||||
Displays:
|
||||
- Session ID
|
||||
- Current model
|
||||
- Token usage
|
||||
- Cost
|
||||
- Queue status
|
||||
- Processing indicator
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
StatusWidget {
|
||||
height: 1;
|
||||
background: $primary;
|
||||
color: $text;
|
||||
padding: 0 1;
|
||||
}
|
||||
"""
|
||||
|
||||
session_id: reactive[str] = reactive("")
|
||||
model: reactive[str] = reactive("gpt-4o-mini")
|
||||
total_tokens: reactive[int] = reactive(0)
|
||||
total_cost: reactive[float] = reactive(0.0)
|
||||
queued_count: reactive[int] = reactive(0)
|
||||
running_count: reactive[int] = reactive(0)
|
||||
is_processing: reactive[bool] = reactive(False)
|
||||
status_message: reactive[str] = reactive("")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__("", name=name, id=id, classes=classes)
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
self._update_display()
|
||||
|
||||
def watch_session_id(self, value: str) -> None:
|
||||
self._update_display()
|
||||
|
||||
def watch_model(self, value: str) -> None:
|
||||
self._update_display()
|
||||
|
||||
def watch_total_tokens(self, value: int) -> None:
|
||||
self._update_display()
|
||||
|
||||
def watch_total_cost(self, value: float) -> None:
|
||||
self._update_display()
|
||||
|
||||
def watch_queued_count(self, value: int) -> None:
|
||||
self._update_display()
|
||||
|
||||
def watch_running_count(self, value: int) -> None:
|
||||
self._update_display()
|
||||
|
||||
def watch_is_processing(self, value: bool) -> None:
|
||||
self._update_display()
|
||||
|
||||
def watch_status_message(self, value: str) -> None:
|
||||
self._update_display()
|
||||
|
||||
def _update_display(self) -> None:
|
||||
"""Update the status display."""
|
||||
parts = []
|
||||
|
||||
# App name
|
||||
parts.append(("◉ PraisonAI", "bold cyan"))
|
||||
|
||||
# Session
|
||||
if self.session_id:
|
||||
parts.append((f"Session: {self.session_id[:8]}", ""))
|
||||
|
||||
# Model
|
||||
parts.append((f"Model: {self.model}", ""))
|
||||
|
||||
# Tokens
|
||||
if self.total_tokens > 0:
|
||||
parts.append((f"Tokens: {self.total_tokens:,}", ""))
|
||||
|
||||
# Cost
|
||||
if self.total_cost > 0:
|
||||
parts.append((f"${self.total_cost:.4f}", "green"))
|
||||
|
||||
# Queue status
|
||||
if self.queued_count > 0 or self.running_count > 0:
|
||||
queue_text = f"⏳ {self.queued_count} queued"
|
||||
if self.running_count > 0:
|
||||
queue_text += f" │ ▶ {self.running_count} running"
|
||||
parts.append((queue_text, "yellow"))
|
||||
|
||||
# Processing indicator
|
||||
if self.is_processing:
|
||||
parts.append(("⟳ Processing...", "bold magenta"))
|
||||
|
||||
# Status message
|
||||
if self.status_message:
|
||||
parts.append((self.status_message, "italic"))
|
||||
|
||||
# Build text
|
||||
text = Text()
|
||||
for i, (content, style) in enumerate(parts):
|
||||
if i > 0:
|
||||
text.append(" │ ", style="dim")
|
||||
text.append(content, style=style)
|
||||
|
||||
self.update(text)
|
||||
|
||||
def update_info(self, info: StatusInfo) -> None:
|
||||
"""Update all status info at once."""
|
||||
self.session_id = info.session_id
|
||||
self.model = info.model
|
||||
self.total_tokens = info.total_tokens
|
||||
self.total_cost = info.total_cost
|
||||
self.queued_count = info.queued_count
|
||||
self.running_count = info.running_count
|
||||
self.is_processing = info.is_processing
|
||||
self.status_message = info.status_message
|
||||
|
||||
else:
|
||||
class StatusWidget:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
Tool Panel Widget for PraisonAI TUI.
|
||||
|
||||
Displays tool execution status and approval dialogs.
|
||||
"""
|
||||
|
||||
from typing import Any, List, Optional
|
||||
from dataclasses import dataclass, field
|
||||
import time
|
||||
|
||||
try:
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static, ListView, ListItem
|
||||
from textual.containers import Vertical, ScrollableContainer
|
||||
from textual.reactive import reactive
|
||||
from textual.message import Message
|
||||
from rich.text import Text
|
||||
from rich.panel import Panel
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
Widget = object
|
||||
Message = object
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""A tool call for display."""
|
||||
call_id: str
|
||||
tool_name: str
|
||||
args_preview: str
|
||||
status: str # "pending", "running", "completed", "failed", "approved", "rejected"
|
||||
result_preview: Optional[str] = None
|
||||
started_at: float = field(default_factory=time.time)
|
||||
ended_at: Optional[float] = None
|
||||
requires_approval: bool = False
|
||||
|
||||
|
||||
if TEXTUAL_AVAILABLE:
|
||||
class ToolPanelWidget(Vertical):
|
||||
"""
|
||||
Tool panel widget.
|
||||
|
||||
Displays:
|
||||
- Recent tool calls
|
||||
- Tool execution status
|
||||
- Approval requests
|
||||
"""
|
||||
|
||||
DEFAULT_CSS = """
|
||||
ToolPanelWidget {
|
||||
height: 100%;
|
||||
border: solid $accent;
|
||||
background: $surface;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
ToolPanelWidget .tool-header {
|
||||
height: 1;
|
||||
background: $accent;
|
||||
padding: 0 1;
|
||||
}
|
||||
|
||||
ToolPanelWidget #available-tools-scroll {
|
||||
height: 1fr;
|
||||
min-height: 5;
|
||||
scrollbar-gutter: stable;
|
||||
}
|
||||
|
||||
ToolPanelWidget #tool-list-scroll {
|
||||
height: auto;
|
||||
max-height: 10;
|
||||
}
|
||||
|
||||
ToolPanelWidget .tool-pending {
|
||||
color: $warning;
|
||||
}
|
||||
|
||||
ToolPanelWidget .tool-running {
|
||||
color: $primary;
|
||||
}
|
||||
|
||||
ToolPanelWidget .tool-completed {
|
||||
color: $success;
|
||||
}
|
||||
|
||||
ToolPanelWidget .tool-failed {
|
||||
color: $error;
|
||||
}
|
||||
|
||||
ToolPanelWidget .tool-approval {
|
||||
background: $warning-darken-2;
|
||||
border: solid $warning;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
ToolPanelWidget .tools-count {
|
||||
color: $text-muted;
|
||||
text-style: italic;
|
||||
}
|
||||
"""
|
||||
|
||||
class ApprovalResponse(Message):
|
||||
"""Event when user responds to approval request."""
|
||||
def __init__(self, call_id: str, approved: bool):
|
||||
self.call_id = call_id
|
||||
self.approved = approved
|
||||
super().__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_items: int = 50,
|
||||
name: Optional[str] = None,
|
||||
id: Optional[str] = None,
|
||||
classes: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, id=id, classes=classes)
|
||||
self._calls: List[ToolCall] = []
|
||||
self._max_items = max_items
|
||||
self._pending_approval: Optional[ToolCall] = None
|
||||
self._available_tools: List[str] = []
|
||||
|
||||
def compose(self):
|
||||
"""Compose the widget."""
|
||||
yield Static("Tools", id="tool-title", classes="tool-header")
|
||||
yield ScrollableContainer(Vertical(id="tool-list"), id="tool-list-scroll")
|
||||
yield Static("", id="tools-count", classes="tools-count")
|
||||
yield ScrollableContainer(Vertical(id="available-tools"), id="available-tools-scroll")
|
||||
|
||||
def on_mount(self) -> None:
|
||||
"""Handle mount."""
|
||||
self._update_display()
|
||||
self._update_available_tools_display()
|
||||
|
||||
def set_available_tools(self, tools: List[Any]) -> None:
|
||||
"""Set the list of available tools to display."""
|
||||
self._available_tools = []
|
||||
for tool in tools:
|
||||
if hasattr(tool, '__name__'):
|
||||
self._available_tools.append(tool.__name__)
|
||||
elif hasattr(tool, 'name'):
|
||||
self._available_tools.append(tool.name)
|
||||
elif callable(tool):
|
||||
self._available_tools.append(str(tool))
|
||||
else:
|
||||
self._available_tools.append(str(type(tool).__name__))
|
||||
self._update_available_tools_display()
|
||||
|
||||
def _update_available_tools_display(self) -> None:
|
||||
"""Update the available tools display."""
|
||||
try:
|
||||
container = self.query_one("#available-tools", Vertical)
|
||||
count_widget = self.query_one("#tools-count", Static)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
container.remove_children()
|
||||
|
||||
if not self._available_tools:
|
||||
container.mount(Static(Text("No tools loaded", style="dim italic")))
|
||||
count_widget.update("")
|
||||
return
|
||||
|
||||
# Update count display
|
||||
count_widget.update(f"{len(self._available_tools)} tools available")
|
||||
|
||||
# Show ALL tools with scrolling support
|
||||
for tool_name in sorted(self._available_tools):
|
||||
text = Text()
|
||||
text.append("• ", style="cyan")
|
||||
text.append(tool_name, style="white")
|
||||
container.mount(Static(text))
|
||||
|
||||
def add_call(self, call: ToolCall) -> None:
|
||||
"""Add a tool call."""
|
||||
self._calls.append(call)
|
||||
|
||||
# Trim old calls
|
||||
if len(self._calls) > self._max_items:
|
||||
self._calls = self._calls[-self._max_items:]
|
||||
|
||||
# Check for approval
|
||||
if call.requires_approval and call.status == "pending":
|
||||
self._pending_approval = call
|
||||
|
||||
self._update_display()
|
||||
|
||||
def update_call(self, call_id: str, status: str, result: Optional[str] = None) -> None:
|
||||
"""Update a tool call status."""
|
||||
for call in self._calls:
|
||||
if call.call_id == call_id:
|
||||
call.status = status
|
||||
if result:
|
||||
call.result_preview = result[:100]
|
||||
if status in ("completed", "failed"):
|
||||
call.ended_at = time.time()
|
||||
break
|
||||
|
||||
# Clear pending approval if this was it
|
||||
if self._pending_approval and self._pending_approval.call_id == call_id:
|
||||
self._pending_approval = None
|
||||
|
||||
self._update_display()
|
||||
|
||||
def _update_display(self) -> None:
|
||||
"""Update the display."""
|
||||
try:
|
||||
container = self.query_one("#tool-list", Vertical)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
# Clear existing
|
||||
container.remove_children()
|
||||
|
||||
# Show pending approval first
|
||||
if self._pending_approval:
|
||||
approval_text = Text()
|
||||
approval_text.append("⚠ Approval Required\n", style="bold yellow")
|
||||
approval_text.append(f"Tool: {self._pending_approval.tool_name}\n")
|
||||
approval_text.append(f"Args: {self._pending_approval.args_preview}\n")
|
||||
approval_text.append("[Y] Approve [N] Reject", style="dim")
|
||||
|
||||
container.mount(Static(
|
||||
Panel(approval_text, border_style="yellow"),
|
||||
classes="tool-approval"
|
||||
))
|
||||
|
||||
# Show recent calls
|
||||
for call in reversed(self._calls[-10:]):
|
||||
if call == self._pending_approval:
|
||||
continue
|
||||
|
||||
# Status icon
|
||||
status_icons = {
|
||||
"pending": "⏳",
|
||||
"running": "⟳",
|
||||
"completed": "✓",
|
||||
"failed": "✗",
|
||||
"approved": "✓",
|
||||
"rejected": "✗",
|
||||
}
|
||||
icon = status_icons.get(call.status, "?")
|
||||
|
||||
# Status style
|
||||
status_styles = {
|
||||
"pending": "yellow",
|
||||
"running": "cyan",
|
||||
"completed": "green",
|
||||
"failed": "red",
|
||||
"approved": "green",
|
||||
"rejected": "red",
|
||||
}
|
||||
style = status_styles.get(call.status, "")
|
||||
|
||||
text = Text()
|
||||
text.append(f"{icon} ", style=style)
|
||||
text.append(call.tool_name, style="bold")
|
||||
|
||||
if call.result_preview:
|
||||
text.append(f" → {call.result_preview[:30]}...", style="dim")
|
||||
|
||||
container.mount(Static(text, classes=f"tool-{call.status}"))
|
||||
|
||||
def handle_approval_key(self, key: str) -> bool:
|
||||
"""
|
||||
Handle approval key press.
|
||||
|
||||
Returns True if handled.
|
||||
"""
|
||||
if not self._pending_approval:
|
||||
return False
|
||||
|
||||
if key.lower() == "y":
|
||||
self.post_message(self.ApprovalResponse(
|
||||
self._pending_approval.call_id,
|
||||
approved=True
|
||||
))
|
||||
self.update_call(self._pending_approval.call_id, "approved")
|
||||
return True
|
||||
elif key.lower() == "n":
|
||||
self.post_message(self.ApprovalResponse(
|
||||
self._pending_approval.call_id,
|
||||
approved=False
|
||||
))
|
||||
self.update_call(self._pending_approval.call_id, "rejected")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@property
|
||||
def has_pending_approval(self) -> bool:
|
||||
"""Check if there's a pending approval."""
|
||||
return self._pending_approval is not None
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all tool calls."""
|
||||
self._calls.clear()
|
||||
self._pending_approval = None
|
||||
self._update_display()
|
||||
|
||||
else:
|
||||
class ToolPanelWidget:
|
||||
"""Placeholder when Textual is not available."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"Textual is required for TUI. Install with: pip install praisonai[tui]"
|
||||
)
|
||||
@@ -27,6 +27,7 @@ LEGACY_COMMANDS = {
|
||||
TYPER_COMMANDS = {
|
||||
'config', 'traces', 'env', 'session', 'completion', 'version',
|
||||
'debug', 'lsp', 'diag', 'doctor', 'acp', 'mcp', 'serve', 'schedule', 'run',
|
||||
'tui', 'queue', 'profile', # TUI, queue management, and profiling commands
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -329,8 +329,7 @@ class PraisonAI:
|
||||
if combined_inputs:
|
||||
combined_prompt = f"{args.command} {' '.join(combined_inputs)}"
|
||||
result = self.handle_direct_prompt(combined_prompt)
|
||||
if result is not None:
|
||||
print(result)
|
||||
# Result already printed by handle_direct_prompt, don't print again
|
||||
return result
|
||||
else:
|
||||
self.agent_file = args.command
|
||||
@@ -357,8 +356,7 @@ class PraisonAI:
|
||||
prompt_parts.append(file_input)
|
||||
prompt = ' '.join(prompt_parts)
|
||||
result = self.handle_direct_prompt(prompt)
|
||||
if result is not None:
|
||||
print(result)
|
||||
# Result already printed by handle_direct_prompt, don't print again
|
||||
return result
|
||||
else:
|
||||
# Agent file was explicitly set, ignore direct prompt and use the file
|
||||
@@ -374,8 +372,7 @@ class PraisonAI:
|
||||
inputs.append(file_input)
|
||||
combined_input = ' '.join(inputs)
|
||||
result = self.handle_direct_prompt(combined_input)
|
||||
if result is not None:
|
||||
print(result)
|
||||
# Result already printed by handle_direct_prompt, don't print again
|
||||
return result
|
||||
# If no command or direct_prompt, preserve agent_file from constructor (don't overwrite)
|
||||
|
||||
@@ -737,6 +734,8 @@ class PraisonAI:
|
||||
parser.add_argument("--expand-tools", type=str, help="Tools for prompt expander (e.g., 'internet_search' or path to tools.py)")
|
||||
parser.add_argument("--tools", "-t", type=str, help="Path to tools.py file for research agent")
|
||||
parser.add_argument("--no-tools", action="store_true", help="Disable default built-in tools (for models that don't support tool calling)")
|
||||
parser.add_argument("--no-acp", action="store_true", help="Disable ACP tools (agentic file operations with plan/approve/apply)")
|
||||
parser.add_argument("--no-lsp", action="store_true", help="Disable LSP tools (code intelligence: symbols, definitions, references)")
|
||||
parser.add_argument("--save", "-s", action="store_true", help="Save research output to file (output/research/)")
|
||||
parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose output for research")
|
||||
parser.add_argument("--web-search", action="store_true", help="Enable native web search (OpenAI, Gemini, Anthropic, xAI, Perplexity)")
|
||||
@@ -786,6 +785,11 @@ class PraisonAI:
|
||||
# Telemetry - usage monitoring
|
||||
parser.add_argument("--telemetry", action="store_true", help="Enable usage monitoring and analytics")
|
||||
|
||||
# Profiling - unified execution profiling
|
||||
parser.add_argument("--profile", action="store_true", help="Enable profiling with timing breakdown")
|
||||
parser.add_argument("--profile-deep", action="store_true", dest="profile_deep", help="Enable deep profiling with call graph (higher overhead)")
|
||||
parser.add_argument("--profile-format", type=str, choices=["text", "json"], default="text", dest="profile_format", help="Profile output format")
|
||||
|
||||
# MCP - Model Context Protocol
|
||||
parser.add_argument("--mcp", type=str, help="MCP server command (e.g., 'npx -y @modelcontextprotocol/server-filesystem .')")
|
||||
parser.add_argument("--mcp-env", type=str, help="MCP environment variables (KEY=value,KEY2=value2)")
|
||||
@@ -1201,10 +1205,13 @@ class PraisonAI:
|
||||
sys.exit(exit_code)
|
||||
|
||||
elif args.command == 'profile':
|
||||
# Profile command - performance profiling and benchmarking
|
||||
from .features.profile import handle_profile_command
|
||||
exit_code = handle_profile_command(unknown_args)
|
||||
sys.exit(exit_code)
|
||||
# Profile command - delegate to Typer CLI for new profiler
|
||||
# This routes to commands/profile.py which has query, imports, startup subcommands
|
||||
from .app import app as typer_app
|
||||
import sys as _sys
|
||||
_sys.argv = ['praisonai', 'profile'] + unknown_args
|
||||
typer_app()
|
||||
sys.exit(0)
|
||||
|
||||
elif args.command == 'eval':
|
||||
# Eval command - evaluate model responses against expected outputs
|
||||
@@ -3457,6 +3464,10 @@ Provide ONLY the commit message, no explanations."""
|
||||
- @rule:name - Include specific rule
|
||||
- @url:https://... - Fetch URL content
|
||||
"""
|
||||
# Check for profiling mode - use unified profiler
|
||||
if hasattr(self, 'args') and getattr(self.args, 'profile', False):
|
||||
return self._handle_profiled_prompt(prompt)
|
||||
|
||||
# Check for inline workflow mode
|
||||
if hasattr(self, 'args') and getattr(self.args, 'workflow', None):
|
||||
return self._run_inline_workflow(prompt)
|
||||
@@ -3932,6 +3943,55 @@ Now, {final_instruction.lower()}:"""
|
||||
print("pip install praisonaiagents # For PraisonAIAgents\n")
|
||||
sys.exit(1)
|
||||
|
||||
def _handle_profiled_prompt(self, prompt):
|
||||
"""
|
||||
Handle direct prompt with unified profiling enabled.
|
||||
|
||||
Uses the unified execution module for consistent profiling across
|
||||
CLI direct invocation and profile command.
|
||||
"""
|
||||
from .execution import ExecutionRequest, Profiler, ProfilerConfig
|
||||
|
||||
# Build execution request
|
||||
model = None
|
||||
if hasattr(self, 'args') and self.args.llm:
|
||||
model = self.args.llm
|
||||
|
||||
request = ExecutionRequest(
|
||||
prompt=prompt,
|
||||
agent_name="DirectAgent",
|
||||
model=model,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Create profiler config from flags
|
||||
deep = getattr(self.args, 'profile_deep', False)
|
||||
output_format = getattr(self.args, 'profile_format', 'text')
|
||||
|
||||
config = ProfilerConfig(
|
||||
layer=2 if deep else 1,
|
||||
show_callers=deep,
|
||||
show_callees=deep,
|
||||
output_format=output_format,
|
||||
)
|
||||
|
||||
# Run with profiling
|
||||
print("[bold cyan]🔬 Starting profiled execution...[/bold cyan]")
|
||||
if deep:
|
||||
print("[yellow]⚠️ Deep profiling enabled - higher overhead[/yellow]")
|
||||
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request, invocation_method="cli_direct")
|
||||
|
||||
# Output profile report
|
||||
if output_format == "json":
|
||||
print(report.to_json())
|
||||
else:
|
||||
print(report.to_text())
|
||||
|
||||
# Return the actual result for any downstream processing
|
||||
return result.output
|
||||
|
||||
def _run_with_status_display(self, agent, prompt):
|
||||
"""
|
||||
Run agent with minimal status display (spinner + tool/handoff updates).
|
||||
@@ -4883,7 +4943,58 @@ Now, {final_instruction.lower()}:"""
|
||||
sys.exit(1)
|
||||
|
||||
def _load_interactive_tools(self):
|
||||
"""Load tools for interactive mode."""
|
||||
"""
|
||||
Load tools for interactive mode using the canonical provider.
|
||||
|
||||
This method uses the centralized interactive_tools module which provides:
|
||||
- ACP tools (acp_create_file, acp_edit_file, etc.) for safe file operations
|
||||
- LSP tools (lsp_list_symbols, lsp_find_definition, etc.) for code intelligence
|
||||
- Basic tools (read_file, write_file, etc.) for standard operations
|
||||
|
||||
Tool groups can be disabled via:
|
||||
- CLI flags: --no-acp, --no-lsp
|
||||
- Env vars: PRAISON_TOOLS_DISABLE=acp,lsp
|
||||
"""
|
||||
# Determine which groups to disable based on CLI args
|
||||
disable_groups = []
|
||||
if hasattr(self, 'args'):
|
||||
if getattr(self.args, 'no_acp', False):
|
||||
disable_groups.append('acp')
|
||||
if getattr(self.args, 'no_lsp', False):
|
||||
disable_groups.append('lsp')
|
||||
|
||||
# Get workspace
|
||||
workspace = os.getcwd()
|
||||
|
||||
try:
|
||||
from .features.interactive_tools import get_interactive_tools, ToolConfig
|
||||
|
||||
# Create config
|
||||
config = ToolConfig.from_env()
|
||||
config.workspace = workspace
|
||||
|
||||
# Apply CLI overrides
|
||||
if 'acp' in disable_groups:
|
||||
config.enable_acp = False
|
||||
if 'lsp' in disable_groups:
|
||||
config.enable_lsp = False
|
||||
|
||||
# Get tools from canonical provider
|
||||
tools_list = get_interactive_tools(
|
||||
config=config,
|
||||
disable=disable_groups if disable_groups else None,
|
||||
)
|
||||
|
||||
logging.debug(f"Loaded {len(tools_list)} interactive tools (ACP: {config.enable_acp}, LSP: {config.enable_lsp})")
|
||||
return tools_list
|
||||
|
||||
except ImportError as e:
|
||||
logging.debug(f"Interactive tools provider not available: {e}")
|
||||
# Fallback to basic tools only
|
||||
return self._load_basic_tools_fallback()
|
||||
|
||||
def _load_basic_tools_fallback(self):
|
||||
"""Fallback to load basic tools when interactive_tools module unavailable."""
|
||||
tools_list = []
|
||||
try:
|
||||
from praisonaiagents.tools import (
|
||||
|
||||
@@ -57,7 +57,7 @@ class CloudDeployer:
|
||||
file.write("FROM python:3.11-slim\n")
|
||||
file.write("WORKDIR /app\n")
|
||||
file.write("COPY . .\n")
|
||||
file.write("RUN pip install flask praisonai==2.9.1 gunicorn markdown\n")
|
||||
file.write("RUN pip install flask praisonai==3.0.2 gunicorn markdown\n")
|
||||
file.write("EXPOSE 8080\n")
|
||||
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "2.9.1"
|
||||
__version__ = "3.0.2"
|
||||
@@ -12,12 +12,13 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=0.5.3",
|
||||
"praisonaiagents>=0.6.0",
|
||||
"python-dotenv>=0.19.0",
|
||||
"instructor>=1.3.3",
|
||||
"PyYAML>=6.0",
|
||||
"mcp>=1.6.0",
|
||||
"typer>=0.9.0",
|
||||
"textual>=0.47.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
@@ -104,6 +105,15 @@ autogen-v4 = [
|
||||
acp = [
|
||||
"agent-client-protocol>=0.7.0",
|
||||
]
|
||||
lite = [
|
||||
# Minimal installation without TUI, litellm - just core agent functionality
|
||||
# Use: pip install praisonai[lite] for lightweight deployments
|
||||
]
|
||||
dev = [
|
||||
"textual-dev>=1.0.0",
|
||||
"pytest>=7.0.0",
|
||||
"pytest-asyncio>=0.21.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://docs.praison.ai"
|
||||
|
||||
+187
@@ -0,0 +1,187 @@
|
||||
#!/bin/bash
|
||||
# PraisonAI TUI Real-Key Verification Script
|
||||
#
|
||||
# This script verifies the TUI and Queue system works with real API keys.
|
||||
# It includes safety gates and secret masking.
|
||||
#
|
||||
# Usage:
|
||||
# ./scripts/verify_tui_real_keys.sh
|
||||
#
|
||||
# Required environment variables:
|
||||
# OPENAI_API_KEY - OpenAI API key
|
||||
#
|
||||
# Optional:
|
||||
# PRAISONAI_REAL_LLM=1 - Required to run real LLM tests
|
||||
|
||||
set -e
|
||||
|
||||
# Colors
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
echo " PraisonAI TUI + Queue System Verification"
|
||||
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
|
||||
# Function to mask secrets in output
|
||||
mask_secrets() {
|
||||
sed -E 's/(sk-[a-zA-Z0-9]{20})[a-zA-Z0-9]+/\1***/g' | \
|
||||
sed -E 's/(key_[a-zA-Z0-9]{20})[a-zA-Z0-9]+/\1***/g' | \
|
||||
sed -E 's/(tvly-[a-zA-Z0-9]{20})[a-zA-Z0-9]+/\1***/g' | \
|
||||
sed -E 's/(AIza[a-zA-Z0-9]{20})[a-zA-Z0-9]+/\1***/g'
|
||||
}
|
||||
|
||||
# Check environment variables
|
||||
echo -e "\n${YELLOW}[1/7] Checking environment variables...${NC}"
|
||||
|
||||
check_env() {
|
||||
local var_name=$1
|
||||
local value="${!var_name}"
|
||||
if [ -n "$value" ]; then
|
||||
local masked=$(echo "$value" | head -c 10)
|
||||
echo -e " ${GREEN}✓${NC} $var_name: ${masked}***"
|
||||
return 0
|
||||
else
|
||||
echo -e " ${YELLOW}⚠${NC} $var_name: not set"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
OPENAI_SET=false
|
||||
check_env "OPENAI_API_KEY" && OPENAI_SET=true
|
||||
check_env "ANTHROPIC_API_KEY" || true
|
||||
check_env "GOOGLE_API_KEY" || true
|
||||
check_env "TAVILY_API_KEY" || true
|
||||
|
||||
if [ "$OPENAI_SET" = false ]; then
|
||||
echo -e "\n${RED}Error: OPENAI_API_KEY is required${NC}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check Python imports
|
||||
echo -e "\n${YELLOW}[2/7] Checking Python imports...${NC}"
|
||||
|
||||
python3 -c "
|
||||
from praisonai.cli.features.queue import QueuedRun, RunState, RunPriority, QueueConfig
|
||||
from praisonai.cli.features.queue import QueueScheduler, QueuePersistence, QueueManager
|
||||
from praisonai.cli.features.tui import TUIEvent, TUIEventType
|
||||
from praisonai.cli.features.tui import TuiOrchestrator, UIStateModel, SimulationRunner
|
||||
from praisonai.cli.features.tui import MockProvider, MockProviderConfig
|
||||
print(' ✓ All imports successful')
|
||||
" 2>&1 | mask_secrets
|
||||
|
||||
# Run unit tests
|
||||
echo -e "\n${YELLOW}[3/7] Running unit tests...${NC}"
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
python3 -m pytest tests/unit/cli/queue/ tests/unit/cli/tui/ -v --tb=short --import-mode=importlib 2>&1 | \
|
||||
tail -20 | mask_secrets
|
||||
|
||||
# Test queue CLI commands
|
||||
echo -e "\n${YELLOW}[4/7] Testing queue CLI commands...${NC}"
|
||||
|
||||
python3 -c "
|
||||
from praisonai.cli.features.tui.cli import create_queue_app
|
||||
app = create_queue_app()
|
||||
print(' ✓ Queue CLI app created')
|
||||
print(' Commands:', [cmd.name for cmd in app.registered_commands])
|
||||
" 2>&1 | mask_secrets
|
||||
|
||||
# Test mock provider
|
||||
echo -e "\n${YELLOW}[5/7] Testing mock provider...${NC}"
|
||||
|
||||
python3 -c "
|
||||
import asyncio
|
||||
from praisonai.cli.features.tui.mock_provider import MockProvider, create_mock_provider
|
||||
|
||||
async def test():
|
||||
provider = create_mock_provider(seed=42)
|
||||
|
||||
# Test streaming
|
||||
chunks = []
|
||||
result = await provider.generate(
|
||||
'hello world',
|
||||
stream=True,
|
||||
on_chunk=lambda c: chunks.append(c)
|
||||
)
|
||||
|
||||
print(' ✓ Mock provider works')
|
||||
print(f' Content length: {len(result[\"content\"])} chars')
|
||||
print(f' Chunks: {len(chunks)}')
|
||||
print(f' Tokens: {result[\"tokens\"]}')
|
||||
print(f' Cost: \${result[\"cost\"]:.6f}')
|
||||
|
||||
asyncio.run(test())
|
||||
" 2>&1 | mask_secrets
|
||||
|
||||
# Test orchestrator
|
||||
echo -e "\n${YELLOW}[6/7] Testing TUI orchestrator...${NC}"
|
||||
|
||||
python3 -c "
|
||||
import asyncio
|
||||
from praisonai.cli.features.tui.orchestrator import TuiOrchestrator, OutputMode
|
||||
from praisonai.cli.features.queue import QueueConfig
|
||||
|
||||
async def test():
|
||||
config = QueueConfig(enable_persistence=False)
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
debug=False,
|
||||
)
|
||||
|
||||
await orchestrator.start(session_id='test-verify')
|
||||
|
||||
# Get snapshot
|
||||
snapshot = orchestrator.get_snapshot()
|
||||
print(' ✓ Orchestrator started')
|
||||
print(f' Session: {snapshot[\"session_id\"]}')
|
||||
print(f' Model: {snapshot[\"model\"]}')
|
||||
|
||||
await orchestrator.stop()
|
||||
print(' ✓ Orchestrator stopped cleanly')
|
||||
|
||||
asyncio.run(test())
|
||||
" 2>&1 | mask_secrets
|
||||
|
||||
# Real LLM test (gated)
|
||||
echo -e "\n${YELLOW}[7/7] Real LLM test...${NC}"
|
||||
|
||||
if [ "$PRAISONAI_REAL_LLM" = "1" ]; then
|
||||
echo " Running real LLM test (PRAISONAI_REAL_LLM=1)"
|
||||
|
||||
python3 -c "
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
# Safety: Use minimal tokens
|
||||
async def test():
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions='Reply with exactly one word.',
|
||||
model='gpt-4o-mini',
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
result = agent.start('Say hello')
|
||||
print(f' ✓ Real LLM response: {result[:50]}...')
|
||||
|
||||
asyncio.run(test())
|
||||
" 2>&1 | mask_secrets
|
||||
else
|
||||
echo -e " ${YELLOW}⚠${NC} Skipped (set PRAISONAI_REAL_LLM=1 to enable)"
|
||||
fi
|
||||
|
||||
# Summary
|
||||
echo -e "\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
echo -e "${GREEN} Verification Complete!${NC}"
|
||||
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
echo ""
|
||||
echo "All checks passed. The TUI + Queue system is working correctly."
|
||||
echo ""
|
||||
echo "To run with real LLM:"
|
||||
echo " PRAISONAI_REAL_LLM=1 ./scripts/verify_tui_real_keys.sh"
|
||||
echo ""
|
||||
@@ -0,0 +1,287 @@
|
||||
"""
|
||||
Unit tests for the interactive tools provider.
|
||||
|
||||
Tests the canonical source of truth for default tools in interactive modes.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
|
||||
class TestToolGroups:
|
||||
"""Tests for tool group definitions."""
|
||||
|
||||
def test_tool_groups_defined(self):
|
||||
"""Test that all expected tool groups are defined."""
|
||||
from praisonai.cli.features.interactive_tools import TOOL_GROUPS
|
||||
|
||||
assert "acp" in TOOL_GROUPS
|
||||
assert "lsp" in TOOL_GROUPS
|
||||
assert "basic" in TOOL_GROUPS
|
||||
assert "interactive" in TOOL_GROUPS
|
||||
|
||||
def test_acp_tools_list(self):
|
||||
"""Test ACP tool group contains expected tools."""
|
||||
from praisonai.cli.features.interactive_tools import TOOL_GROUPS
|
||||
|
||||
expected = {
|
||||
"acp_create_file",
|
||||
"acp_edit_file",
|
||||
"acp_delete_file",
|
||||
"acp_execute_command",
|
||||
}
|
||||
assert set(TOOL_GROUPS["acp"]) == expected
|
||||
|
||||
def test_lsp_tools_list(self):
|
||||
"""Test LSP tool group contains expected tools."""
|
||||
from praisonai.cli.features.interactive_tools import TOOL_GROUPS
|
||||
|
||||
expected = {
|
||||
"lsp_list_symbols",
|
||||
"lsp_find_definition",
|
||||
"lsp_find_references",
|
||||
"lsp_get_diagnostics",
|
||||
}
|
||||
assert set(TOOL_GROUPS["lsp"]) == expected
|
||||
|
||||
def test_basic_tools_list(self):
|
||||
"""Test basic tool group contains expected tools."""
|
||||
from praisonai.cli.features.interactive_tools import TOOL_GROUPS
|
||||
|
||||
expected = {
|
||||
"read_file",
|
||||
"write_file",
|
||||
"list_files",
|
||||
"execute_command",
|
||||
"internet_search",
|
||||
}
|
||||
assert set(TOOL_GROUPS["basic"]) == expected
|
||||
|
||||
def test_interactive_group_is_union(self):
|
||||
"""Test interactive group is union of all groups."""
|
||||
from praisonai.cli.features.interactive_tools import TOOL_GROUPS
|
||||
|
||||
expected = set(TOOL_GROUPS["acp"]) | set(TOOL_GROUPS["lsp"]) | set(TOOL_GROUPS["basic"])
|
||||
assert set(TOOL_GROUPS["interactive"]) == expected
|
||||
|
||||
|
||||
class TestToolConfig:
|
||||
"""Tests for ToolConfig dataclass."""
|
||||
|
||||
def test_default_config(self):
|
||||
"""Test default configuration values."""
|
||||
from praisonai.cli.features.interactive_tools import ToolConfig
|
||||
|
||||
config = ToolConfig()
|
||||
assert config.enable_acp is True
|
||||
assert config.enable_lsp is True
|
||||
assert config.enable_basic is True
|
||||
assert config.approval_mode == "auto"
|
||||
|
||||
def test_config_from_env_disable_acp(self):
|
||||
"""Test config from env with ACP disabled."""
|
||||
from praisonai.cli.features.interactive_tools import ToolConfig
|
||||
|
||||
with patch.dict(os.environ, {"PRAISON_TOOLS_DISABLE": "acp"}):
|
||||
config = ToolConfig.from_env()
|
||||
assert config.enable_acp is False
|
||||
assert config.enable_lsp is True
|
||||
|
||||
def test_config_from_env_disable_lsp(self):
|
||||
"""Test config from env with LSP disabled."""
|
||||
from praisonai.cli.features.interactive_tools import ToolConfig
|
||||
|
||||
with patch.dict(os.environ, {"PRAISON_TOOLS_DISABLE": "lsp"}):
|
||||
config = ToolConfig.from_env()
|
||||
assert config.enable_acp is True
|
||||
assert config.enable_lsp is False
|
||||
|
||||
def test_config_from_env_disable_multiple(self):
|
||||
"""Test config from env with multiple groups disabled."""
|
||||
from praisonai.cli.features.interactive_tools import ToolConfig
|
||||
|
||||
with patch.dict(os.environ, {"PRAISON_TOOLS_DISABLE": "acp,lsp"}):
|
||||
config = ToolConfig.from_env()
|
||||
assert config.enable_acp is False
|
||||
assert config.enable_lsp is False
|
||||
assert config.enable_basic is True
|
||||
|
||||
def test_config_from_env_workspace(self):
|
||||
"""Test config from env with workspace override."""
|
||||
from praisonai.cli.features.interactive_tools import ToolConfig
|
||||
|
||||
with patch.dict(os.environ, {"PRAISON_WORKSPACE": "/custom/path"}):
|
||||
config = ToolConfig.from_env()
|
||||
assert config.workspace == "/custom/path"
|
||||
|
||||
|
||||
class TestResolveToolGroups:
|
||||
"""Tests for resolve_tool_groups function."""
|
||||
|
||||
def test_resolve_all_enabled(self):
|
||||
"""Test resolving with all groups enabled."""
|
||||
from praisonai.cli.features.interactive_tools import resolve_tool_groups, ToolConfig, TOOL_GROUPS
|
||||
|
||||
config = ToolConfig()
|
||||
tool_names = resolve_tool_groups(config=config)
|
||||
|
||||
expected = set(TOOL_GROUPS["interactive"])
|
||||
assert tool_names == expected
|
||||
|
||||
def test_resolve_disable_acp(self):
|
||||
"""Test resolving with ACP disabled."""
|
||||
from praisonai.cli.features.interactive_tools import resolve_tool_groups, ToolConfig, TOOL_GROUPS
|
||||
|
||||
config = ToolConfig(enable_acp=False)
|
||||
tool_names = resolve_tool_groups(disable=["acp"], config=config)
|
||||
|
||||
# Should not contain ACP tools
|
||||
for acp_tool in TOOL_GROUPS["acp"]:
|
||||
assert acp_tool not in tool_names
|
||||
|
||||
# Should contain LSP and basic tools
|
||||
for lsp_tool in TOOL_GROUPS["lsp"]:
|
||||
assert lsp_tool in tool_names
|
||||
|
||||
def test_resolve_disable_lsp(self):
|
||||
"""Test resolving with LSP disabled."""
|
||||
from praisonai.cli.features.interactive_tools import resolve_tool_groups, ToolConfig, TOOL_GROUPS
|
||||
|
||||
config = ToolConfig(enable_lsp=False)
|
||||
tool_names = resolve_tool_groups(disable=["lsp"], config=config)
|
||||
|
||||
# Should not contain LSP tools
|
||||
for lsp_tool in TOOL_GROUPS["lsp"]:
|
||||
assert lsp_tool not in tool_names
|
||||
|
||||
# Should contain ACP and basic tools
|
||||
for acp_tool in TOOL_GROUPS["acp"]:
|
||||
assert acp_tool in tool_names
|
||||
|
||||
def test_resolve_specific_groups(self):
|
||||
"""Test resolving specific groups only."""
|
||||
from praisonai.cli.features.interactive_tools import resolve_tool_groups, TOOL_GROUPS
|
||||
|
||||
tool_names = resolve_tool_groups(groups=["basic"])
|
||||
|
||||
assert tool_names == set(TOOL_GROUPS["basic"])
|
||||
|
||||
|
||||
class TestGetInteractiveTools:
|
||||
"""Tests for get_interactive_tools function."""
|
||||
|
||||
def test_get_tools_returns_list(self):
|
||||
"""Test that get_interactive_tools returns a list."""
|
||||
from praisonai.cli.features.interactive_tools import get_interactive_tools, ToolConfig
|
||||
|
||||
config = ToolConfig(workspace="/tmp")
|
||||
tools = get_interactive_tools(config=config)
|
||||
|
||||
assert isinstance(tools, list)
|
||||
|
||||
def test_get_tools_deterministic_order(self):
|
||||
"""Test that tool order is deterministic."""
|
||||
from praisonai.cli.features.interactive_tools import get_interactive_tools, ToolConfig
|
||||
|
||||
config = ToolConfig(workspace="/tmp")
|
||||
tools1 = get_interactive_tools(config=config)
|
||||
tools2 = get_interactive_tools(config=config)
|
||||
|
||||
names1 = [t.__name__ for t in tools1]
|
||||
names2 = [t.__name__ for t in tools2]
|
||||
|
||||
assert names1 == names2
|
||||
|
||||
def test_get_tools_with_disable_acp(self):
|
||||
"""Test get_interactive_tools with ACP disabled."""
|
||||
from praisonai.cli.features.interactive_tools import get_interactive_tools, ToolConfig
|
||||
|
||||
config = ToolConfig(workspace="/tmp", enable_acp=False)
|
||||
tools = get_interactive_tools(config=config, disable=["acp"])
|
||||
|
||||
tool_names = [t.__name__ for t in tools]
|
||||
acp_tools = [n for n in tool_names if n.startswith("acp_")]
|
||||
|
||||
assert len(acp_tools) == 0
|
||||
|
||||
def test_get_tools_with_disable_lsp(self):
|
||||
"""Test get_interactive_tools with LSP disabled."""
|
||||
from praisonai.cli.features.interactive_tools import get_interactive_tools, ToolConfig
|
||||
|
||||
config = ToolConfig(workspace="/tmp", enable_lsp=False)
|
||||
tools = get_interactive_tools(config=config, disable=["lsp"])
|
||||
|
||||
tool_names = [t.__name__ for t in tools]
|
||||
lsp_tools = [n for n in tool_names if n.startswith("lsp_")]
|
||||
|
||||
assert len(lsp_tools) == 0
|
||||
|
||||
|
||||
class TestMergeToolOverrides:
|
||||
"""Tests for merge_tool_overrides function."""
|
||||
|
||||
def test_merge_add_tools(self):
|
||||
"""Test adding tools to default list."""
|
||||
from praisonai.cli.features.interactive_tools import merge_tool_overrides
|
||||
|
||||
def tool_a():
|
||||
pass
|
||||
|
||||
def tool_b():
|
||||
pass
|
||||
|
||||
def tool_c():
|
||||
pass
|
||||
|
||||
default = [tool_a, tool_b]
|
||||
result = merge_tool_overrides(default, add_tools=[tool_c])
|
||||
|
||||
names = [t.__name__ for t in result]
|
||||
assert "tool_c" in names
|
||||
|
||||
def test_merge_remove_tools(self):
|
||||
"""Test removing tools from default list."""
|
||||
from praisonai.cli.features.interactive_tools import merge_tool_overrides
|
||||
|
||||
def tool_a():
|
||||
pass
|
||||
|
||||
def tool_b():
|
||||
pass
|
||||
|
||||
default = [tool_a, tool_b]
|
||||
result = merge_tool_overrides(default, remove_names=["tool_a"])
|
||||
|
||||
names = [t.__name__ for t in result]
|
||||
assert "tool_a" not in names
|
||||
assert "tool_b" in names
|
||||
|
||||
|
||||
class TestLazyImports:
|
||||
"""Tests for lazy import behavior."""
|
||||
|
||||
def test_import_time_fast(self):
|
||||
"""Test that importing the module is fast (no heavy deps)."""
|
||||
import time
|
||||
|
||||
start = time.time()
|
||||
# Force reimport
|
||||
import importlib
|
||||
import praisonai.cli.features.interactive_tools as module
|
||||
importlib.reload(module)
|
||||
duration = time.time() - start
|
||||
|
||||
# Should be under 100ms
|
||||
assert duration < 0.1, f"Import took {duration:.3f}s, expected < 0.1s"
|
||||
|
||||
def test_tool_groups_available_without_runtime(self):
|
||||
"""Test that TOOL_GROUPS is available without runtime initialization."""
|
||||
from praisonai.cli.features.interactive_tools import TOOL_GROUPS
|
||||
|
||||
# Should be accessible immediately
|
||||
assert len(TOOL_GROUPS) > 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,366 @@
|
||||
"""
|
||||
Unit tests for CLI profiler module.
|
||||
"""
|
||||
|
||||
|
||||
class TestProfilerConfig:
|
||||
"""Tests for ProfilerConfig dataclass."""
|
||||
|
||||
def test_default_config(self):
|
||||
"""Test default configuration values."""
|
||||
from praisonai.cli.features.profiler import ProfilerConfig
|
||||
|
||||
config = ProfilerConfig()
|
||||
assert config.deep is False
|
||||
assert config.limit == 30
|
||||
assert config.sort_by == "cumulative"
|
||||
assert config.show_files is False
|
||||
assert config.show_callers is False
|
||||
assert config.show_callees is False
|
||||
assert config.importtime is False
|
||||
assert config.first_token is False
|
||||
assert config.save_path is None
|
||||
assert config.output_format == "text"
|
||||
assert config.stream is False
|
||||
|
||||
def test_custom_config(self):
|
||||
"""Test custom configuration."""
|
||||
from praisonai.cli.features.profiler import ProfilerConfig
|
||||
|
||||
config = ProfilerConfig(
|
||||
deep=True,
|
||||
limit=50,
|
||||
sort_by="tottime",
|
||||
show_files=True,
|
||||
)
|
||||
assert config.deep is True
|
||||
assert config.limit == 50
|
||||
assert config.sort_by == "tottime"
|
||||
assert config.show_files is True
|
||||
|
||||
|
||||
class TestTimingBreakdown:
|
||||
"""Tests for TimingBreakdown dataclass."""
|
||||
|
||||
def test_default_timing(self):
|
||||
"""Test default timing values."""
|
||||
from praisonai.cli.features.profiler.core import TimingBreakdown
|
||||
|
||||
timing = TimingBreakdown()
|
||||
assert timing.cli_parse_ms == 0.0
|
||||
assert timing.imports_ms == 0.0
|
||||
assert timing.agent_construction_ms == 0.0
|
||||
assert timing.model_init_ms == 0.0
|
||||
assert timing.first_token_ms == 0.0
|
||||
assert timing.total_run_ms == 0.0
|
||||
|
||||
|
||||
class TestFunctionStats:
|
||||
"""Tests for FunctionStats dataclass."""
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test conversion to dictionary."""
|
||||
from praisonai.cli.features.profiler.core import FunctionStats
|
||||
|
||||
stats = FunctionStats(
|
||||
name="test_func",
|
||||
filename="/path/to/file.py",
|
||||
lineno=42,
|
||||
calls=10,
|
||||
tottime=0.5,
|
||||
cumtime=1.0,
|
||||
)
|
||||
|
||||
result = stats.to_dict()
|
||||
assert result["name"] == "test_func"
|
||||
assert result["filename"] == "/path/to/file.py"
|
||||
assert result["lineno"] == 42
|
||||
assert result["calls"] == 10
|
||||
assert result["tottime_ms"] == 500.0
|
||||
assert result["cumtime_ms"] == 1000.0
|
||||
|
||||
|
||||
class TestFileStats:
|
||||
"""Tests for FileStats dataclass."""
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test conversion to dictionary."""
|
||||
from praisonai.cli.features.profiler.core import FileStats, FunctionStats
|
||||
|
||||
func_stats = FunctionStats(
|
||||
name="func1",
|
||||
filename="/path/to/file.py",
|
||||
lineno=10,
|
||||
calls=5,
|
||||
tottime=0.1,
|
||||
cumtime=0.2,
|
||||
)
|
||||
|
||||
file_stats = FileStats(
|
||||
filename="/path/to/file.py",
|
||||
total_time=0.5,
|
||||
functions=[func_stats],
|
||||
)
|
||||
|
||||
result = file_stats.to_dict()
|
||||
assert result["filename"] == "/path/to/file.py"
|
||||
assert result["total_time_ms"] == 500.0
|
||||
assert result["function_count"] == 1
|
||||
|
||||
|
||||
class TestProfilerResult:
|
||||
"""Tests for ProfilerResult dataclass."""
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test conversion to dictionary."""
|
||||
from praisonai.cli.features.profiler.core import (
|
||||
ProfilerResult,
|
||||
TimingBreakdown,
|
||||
FunctionStats,
|
||||
FileStats,
|
||||
)
|
||||
|
||||
timing = TimingBreakdown(
|
||||
cli_parse_ms=1.0,
|
||||
imports_ms=100.0,
|
||||
total_run_ms=500.0,
|
||||
)
|
||||
|
||||
func_stats = [FunctionStats(
|
||||
name="test",
|
||||
filename="test.py",
|
||||
lineno=1,
|
||||
calls=1,
|
||||
tottime=0.1,
|
||||
cumtime=0.2,
|
||||
)]
|
||||
|
||||
file_stats = [FileStats(
|
||||
filename="test.py",
|
||||
total_time=0.2,
|
||||
)]
|
||||
|
||||
result = ProfilerResult(
|
||||
prompt="test prompt",
|
||||
response="test response",
|
||||
timing=timing,
|
||||
function_stats=func_stats,
|
||||
file_stats=file_stats,
|
||||
metadata={"model": "test"},
|
||||
)
|
||||
|
||||
data = result.to_dict()
|
||||
assert "timestamp" in data
|
||||
assert data["prompt"] == "test prompt"
|
||||
assert data["response_preview"] == "test response"
|
||||
assert data["timing"]["cli_parse_ms"] == 1.0
|
||||
assert data["timing"]["imports_ms"] == 100.0
|
||||
assert len(data["top_functions"]) == 1
|
||||
assert len(data["top_files"]) == 1
|
||||
|
||||
|
||||
class TestQueryProfiler:
|
||||
"""Tests for QueryProfiler class."""
|
||||
|
||||
def test_init_default_config(self):
|
||||
"""Test initialization with default config."""
|
||||
from praisonai.cli.features.profiler import QueryProfiler, ProfilerConfig
|
||||
|
||||
profiler = QueryProfiler()
|
||||
assert profiler.config is not None
|
||||
assert isinstance(profiler.config, ProfilerConfig)
|
||||
|
||||
def test_init_custom_config(self):
|
||||
"""Test initialization with custom config."""
|
||||
from praisonai.cli.features.profiler import QueryProfiler, ProfilerConfig
|
||||
|
||||
config = ProfilerConfig(deep=True, limit=50)
|
||||
profiler = QueryProfiler(config)
|
||||
assert profiler.config.deep is True
|
||||
assert profiler.config.limit == 50
|
||||
|
||||
def test_collect_metadata(self):
|
||||
"""Test metadata collection."""
|
||||
from praisonai.cli.features.profiler import QueryProfiler
|
||||
|
||||
profiler = QueryProfiler()
|
||||
metadata = profiler._collect_metadata("gpt-4")
|
||||
|
||||
assert "python_version" in metadata
|
||||
assert "platform" in metadata
|
||||
assert "praisonai_version" in metadata
|
||||
assert metadata["model"] == "gpt-4"
|
||||
assert "timestamp" in metadata
|
||||
|
||||
|
||||
class TestFormatProfileReport:
|
||||
"""Tests for format_profile_report function."""
|
||||
|
||||
def test_basic_report(self):
|
||||
"""Test basic report formatting."""
|
||||
from praisonai.cli.features.profiler import (
|
||||
format_profile_report,
|
||||
ProfilerConfig,
|
||||
)
|
||||
from praisonai.cli.features.profiler.core import (
|
||||
ProfilerResult,
|
||||
TimingBreakdown,
|
||||
FunctionStats,
|
||||
FileStats,
|
||||
)
|
||||
|
||||
timing = TimingBreakdown(
|
||||
cli_parse_ms=1.0,
|
||||
imports_ms=100.0,
|
||||
total_run_ms=500.0,
|
||||
)
|
||||
|
||||
func_stats = [FunctionStats(
|
||||
name="test_func",
|
||||
filename="test.py",
|
||||
lineno=1,
|
||||
calls=10,
|
||||
tottime=0.1,
|
||||
cumtime=0.2,
|
||||
)]
|
||||
|
||||
file_stats = [FileStats(
|
||||
filename="test.py",
|
||||
total_time=0.2,
|
||||
)]
|
||||
|
||||
result = ProfilerResult(
|
||||
prompt="test",
|
||||
response="response",
|
||||
timing=timing,
|
||||
function_stats=func_stats,
|
||||
file_stats=file_stats,
|
||||
metadata={"python_version": "3.12", "platform": "test", "praisonai_version": "1.0", "model": "test"},
|
||||
)
|
||||
|
||||
report = format_profile_report(result)
|
||||
|
||||
assert "PraisonAI Profile Report" in report
|
||||
assert "System Information" in report
|
||||
assert "Timing Breakdown" in report
|
||||
assert "Per-Function Timing" in report
|
||||
assert "test_func" in report
|
||||
|
||||
def test_report_with_files(self):
|
||||
"""Test report with file grouping."""
|
||||
from praisonai.cli.features.profiler import (
|
||||
format_profile_report,
|
||||
ProfilerConfig,
|
||||
)
|
||||
from praisonai.cli.features.profiler.core import (
|
||||
ProfilerResult,
|
||||
TimingBreakdown,
|
||||
FunctionStats,
|
||||
FileStats,
|
||||
)
|
||||
|
||||
config = ProfilerConfig(show_files=True)
|
||||
|
||||
timing = TimingBreakdown()
|
||||
func_stats = []
|
||||
file_stats = [FileStats(filename="test.py", total_time=0.5)]
|
||||
|
||||
result = ProfilerResult(
|
||||
prompt="test",
|
||||
response="response",
|
||||
timing=timing,
|
||||
function_stats=func_stats,
|
||||
file_stats=file_stats,
|
||||
metadata={},
|
||||
)
|
||||
|
||||
report = format_profile_report(result, config)
|
||||
assert "Per-File Timing" in report
|
||||
|
||||
|
||||
class TestRedactSecrets:
|
||||
"""Tests for secret redaction."""
|
||||
|
||||
def test_redact_openai_key(self):
|
||||
"""Test redacting OpenAI API keys."""
|
||||
from praisonai.cli.features.profiler.core import redact_secrets
|
||||
|
||||
text = "API key: sk-1234567890abcdefghijklmnopqrstuvwxyz"
|
||||
result = redact_secrets(text)
|
||||
assert "sk-1234567890" not in result
|
||||
assert "[REDACTED]" in result
|
||||
|
||||
def test_redact_anthropic_key(self):
|
||||
"""Test redacting Anthropic API keys."""
|
||||
from praisonai.cli.features.profiler.core import redact_secrets
|
||||
|
||||
text = "Key: sk-ant-api03-abcdefghijklmnopqrstuvwxyz"
|
||||
result = redact_secrets(text)
|
||||
assert "sk-ant-api03" not in result
|
||||
assert "[REDACTED]" in result
|
||||
|
||||
def test_preserve_normal_text(self):
|
||||
"""Test that normal text is preserved."""
|
||||
from praisonai.cli.features.profiler.core import redact_secrets
|
||||
|
||||
text = "This is normal text without secrets"
|
||||
result = redact_secrets(text)
|
||||
assert result == text
|
||||
|
||||
|
||||
class TestCLIImports:
|
||||
"""Tests for CLI import functionality."""
|
||||
|
||||
def test_cli_app_imports(self):
|
||||
"""Test that CLI app imports without errors."""
|
||||
# This tests the fix for ModuleNotFoundError
|
||||
from praisonai.cli.app import app
|
||||
assert app is not None
|
||||
|
||||
def test_output_module_imports(self):
|
||||
"""Test that output module imports correctly."""
|
||||
from praisonai.cli.output import OutputController, OutputMode
|
||||
assert OutputController is not None
|
||||
assert OutputMode is not None
|
||||
|
||||
def test_profiler_module_imports(self):
|
||||
"""Test that profiler module imports correctly."""
|
||||
from praisonai.cli.features.profiler import (
|
||||
ProfilerConfig,
|
||||
ProfilerResult,
|
||||
QueryProfiler,
|
||||
run_profiled_query,
|
||||
format_profile_report,
|
||||
)
|
||||
assert ProfilerConfig is not None
|
||||
assert ProfilerResult is not None
|
||||
assert QueryProfiler is not None
|
||||
assert run_profiled_query is not None
|
||||
assert format_profile_report is not None
|
||||
|
||||
|
||||
class TestOutputController:
|
||||
"""Tests for OutputController."""
|
||||
|
||||
def test_default_mode(self):
|
||||
"""Test default output mode."""
|
||||
from praisonai.cli.output import OutputController, OutputMode
|
||||
|
||||
controller = OutputController()
|
||||
assert controller.mode == OutputMode.TEXT
|
||||
assert controller.is_json_mode is False
|
||||
assert controller.is_quiet is False
|
||||
|
||||
def test_json_mode(self):
|
||||
"""Test JSON output mode."""
|
||||
from praisonai.cli.output import OutputController, OutputMode
|
||||
|
||||
controller = OutputController(mode=OutputMode.JSON)
|
||||
assert controller.is_json_mode is True
|
||||
|
||||
def test_quiet_mode(self):
|
||||
"""Test quiet output mode."""
|
||||
from praisonai.cli.output import OutputController, OutputMode
|
||||
|
||||
controller = OutputController(mode=OutputMode.QUIET)
|
||||
assert controller.is_quiet is True
|
||||
@@ -0,0 +1,584 @@
|
||||
"""
|
||||
Unit tests for profiler optimizations module.
|
||||
|
||||
Tests Tier 0/1/2 optimization features:
|
||||
- Provider caching
|
||||
- Lazy imports
|
||||
- Client pooling
|
||||
- Prewarm manager
|
||||
- Lite mode
|
||||
- Performance snapshots
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class TestProviderCache:
|
||||
"""Tests for ProviderCache (Tier 0)."""
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Test that ProviderCache is a singleton."""
|
||||
from praisonai.cli.features.profiler import ProviderCache
|
||||
|
||||
cache1 = ProviderCache()
|
||||
cache2 = ProviderCache()
|
||||
assert cache1 is cache2
|
||||
|
||||
def test_get_set(self):
|
||||
"""Test basic get/set operations."""
|
||||
from praisonai.cli.features.profiler import get_provider_cache
|
||||
|
||||
cache = get_provider_cache()
|
||||
cache.invalidate() # Clear cache
|
||||
|
||||
# Miss
|
||||
result = cache.get('test_key')
|
||||
assert result is None
|
||||
|
||||
# Set
|
||||
cache.set('test_key', {'value': 'test'})
|
||||
|
||||
# Hit - cache stores as {'value': ..., 'expires': ...}
|
||||
result = cache.get('test_key')
|
||||
assert result is not None
|
||||
assert result['value']['value'] == 'test'
|
||||
|
||||
def test_invalidate(self):
|
||||
"""Test cache invalidation."""
|
||||
from praisonai.cli.features.profiler import get_provider_cache
|
||||
|
||||
cache = get_provider_cache()
|
||||
cache.set('key1', 'value1')
|
||||
cache.set('key2', 'value2')
|
||||
|
||||
# Invalidate single key
|
||||
cache.invalidate('key1')
|
||||
assert cache.get('key1') is None
|
||||
assert cache.get('key2') is not None
|
||||
|
||||
# Invalidate all
|
||||
cache.invalidate()
|
||||
assert cache.get('key2') is None
|
||||
|
||||
def test_stats(self):
|
||||
"""Test cache statistics."""
|
||||
from praisonai.cli.features.profiler import get_provider_cache
|
||||
|
||||
cache = get_provider_cache()
|
||||
cache.invalidate()
|
||||
|
||||
# Generate some hits and misses
|
||||
cache.get('miss1') # miss
|
||||
cache.set('hit1', 'value')
|
||||
cache.get('hit1') # hit
|
||||
cache.get('miss2') # miss
|
||||
|
||||
stats = cache.get_stats()
|
||||
assert 'hits' in stats
|
||||
assert 'misses' in stats
|
||||
assert 'size' in stats
|
||||
|
||||
def test_thread_safety(self):
|
||||
"""Test thread-safe operations."""
|
||||
from praisonai.cli.features.profiler import get_provider_cache
|
||||
|
||||
cache = get_provider_cache()
|
||||
cache.invalidate()
|
||||
|
||||
errors = []
|
||||
|
||||
def worker(thread_id):
|
||||
try:
|
||||
for i in range(100):
|
||||
key = f"thread_{thread_id}_key_{i}"
|
||||
cache.set(key, f"value_{i}")
|
||||
cache.get(key)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(i,)) for i in range(5)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert len(errors) == 0
|
||||
|
||||
|
||||
class TestResolveModelProvider:
|
||||
"""Tests for resolve_model_provider (Tier 0)."""
|
||||
|
||||
def test_openai_models(self):
|
||||
"""Test OpenAI model resolution."""
|
||||
from praisonai.cli.features.profiler import resolve_model_provider
|
||||
|
||||
assert resolve_model_provider('gpt-4o') == ('openai', 'gpt-4o')
|
||||
assert resolve_model_provider('gpt-4o-mini') == ('openai', 'gpt-4o-mini')
|
||||
assert resolve_model_provider('o1-preview') == ('openai', 'o1-preview')
|
||||
|
||||
def test_anthropic_models(self):
|
||||
"""Test Anthropic model resolution."""
|
||||
from praisonai.cli.features.profiler import resolve_model_provider
|
||||
|
||||
assert resolve_model_provider('claude-3-opus') == ('anthropic', 'claude-3-opus')
|
||||
assert resolve_model_provider('claude-3-sonnet') == ('anthropic', 'claude-3-sonnet')
|
||||
|
||||
def test_google_models(self):
|
||||
"""Test Google model resolution."""
|
||||
from praisonai.cli.features.profiler import resolve_model_provider
|
||||
|
||||
assert resolve_model_provider('gemini-pro') == ('google', 'gemini-pro')
|
||||
assert resolve_model_provider('gemini-1.5-flash') == ('google', 'gemini-1.5-flash')
|
||||
|
||||
def test_explicit_provider(self):
|
||||
"""Test explicit provider/model format."""
|
||||
from praisonai.cli.features.profiler import resolve_model_provider
|
||||
|
||||
assert resolve_model_provider('openai/gpt-4') == ('openai', 'gpt-4')
|
||||
assert resolve_model_provider('anthropic/claude-3') == ('anthropic', 'claude-3')
|
||||
|
||||
def test_caching(self):
|
||||
"""Test that results are cached."""
|
||||
from praisonai.cli.features.profiler import resolve_model_provider
|
||||
|
||||
# Clear cache
|
||||
resolve_model_provider.cache_clear()
|
||||
|
||||
# First call
|
||||
result1 = resolve_model_provider('gpt-4o')
|
||||
# Second call (should be cached)
|
||||
result2 = resolve_model_provider('gpt-4o')
|
||||
|
||||
assert result1 == result2
|
||||
|
||||
# Check cache info
|
||||
info = resolve_model_provider.cache_info()
|
||||
assert info.hits >= 1
|
||||
|
||||
|
||||
class TestLazyImporter:
|
||||
"""Tests for LazyImporter (Tier 0)."""
|
||||
|
||||
def test_import_module(self):
|
||||
"""Test lazy module import."""
|
||||
from praisonai.cli.features.profiler import LazyImporter
|
||||
|
||||
LazyImporter.clear()
|
||||
|
||||
# Import a standard library module
|
||||
json_module = LazyImporter.get('json')
|
||||
assert json_module is not None
|
||||
assert hasattr(json_module, 'dumps')
|
||||
|
||||
def test_import_attribute(self):
|
||||
"""Test lazy attribute import."""
|
||||
from praisonai.cli.features.profiler import LazyImporter
|
||||
|
||||
LazyImporter.clear()
|
||||
|
||||
# Import a specific attribute
|
||||
Path = LazyImporter.get('pathlib', 'Path')
|
||||
assert Path is not None
|
||||
assert callable(Path)
|
||||
|
||||
def test_caching(self):
|
||||
"""Test that imports are cached."""
|
||||
from praisonai.cli.features.profiler import LazyImporter
|
||||
|
||||
LazyImporter.clear()
|
||||
|
||||
# First import
|
||||
json1 = LazyImporter.get('json')
|
||||
# Second import (should be cached)
|
||||
json2 = LazyImporter.get('json')
|
||||
|
||||
assert json1 is json2
|
||||
|
||||
|
||||
class TestClientPool:
|
||||
"""Tests for ClientPool (Tier 1)."""
|
||||
|
||||
def test_get_or_create(self):
|
||||
"""Test client creation and reuse."""
|
||||
from praisonai.cli.features.profiler import ClientPool
|
||||
|
||||
pool = ClientPool(max_size=5)
|
||||
|
||||
# Create client
|
||||
client1 = pool.get_or_create('key1', lambda: {'id': 1})
|
||||
assert client1 == {'id': 1}
|
||||
|
||||
# Reuse client
|
||||
client2 = pool.get_or_create('key1', lambda: {'id': 2})
|
||||
assert client2 == {'id': 1} # Same as first
|
||||
|
||||
def test_eviction(self):
|
||||
"""Test LRU eviction when pool is full."""
|
||||
from praisonai.cli.features.profiler import ClientPool
|
||||
|
||||
pool = ClientPool(max_size=2)
|
||||
|
||||
pool.get_or_create('key1', lambda: 'client1')
|
||||
pool.get_or_create('key2', lambda: 'client2')
|
||||
|
||||
assert pool.size() == 2
|
||||
|
||||
# This should evict key1 (oldest)
|
||||
pool.get_or_create('key3', lambda: 'client3')
|
||||
|
||||
assert pool.size() == 2
|
||||
|
||||
def test_remove(self):
|
||||
"""Test explicit client removal."""
|
||||
from praisonai.cli.features.profiler import ClientPool
|
||||
|
||||
pool = ClientPool()
|
||||
pool.get_or_create('key1', lambda: 'client1')
|
||||
|
||||
assert pool.size() == 1
|
||||
|
||||
pool.remove('key1')
|
||||
|
||||
assert pool.size() == 0
|
||||
|
||||
def test_clear(self):
|
||||
"""Test clearing all clients."""
|
||||
from praisonai.cli.features.profiler import ClientPool
|
||||
|
||||
pool = ClientPool()
|
||||
pool.get_or_create('key1', lambda: 'client1')
|
||||
pool.get_or_create('key2', lambda: 'client2')
|
||||
|
||||
pool.clear()
|
||||
|
||||
assert pool.size() == 0
|
||||
|
||||
|
||||
class TestPrewarmManager:
|
||||
"""Tests for PrewarmManager (Tier 1)."""
|
||||
|
||||
def test_enable_disable(self):
|
||||
"""Test enabling and disabling prewarm."""
|
||||
from praisonai.cli.features.profiler import PrewarmManager
|
||||
|
||||
PrewarmManager.reset()
|
||||
|
||||
assert not PrewarmManager.is_enabled()
|
||||
|
||||
PrewarmManager.enable()
|
||||
assert PrewarmManager.is_enabled()
|
||||
|
||||
PrewarmManager.disable()
|
||||
assert not PrewarmManager.is_enabled()
|
||||
|
||||
def test_prewarm_disabled(self):
|
||||
"""Test that prewarm does nothing when disabled."""
|
||||
from praisonai.cli.features.profiler import PrewarmManager
|
||||
|
||||
PrewarmManager.reset()
|
||||
|
||||
# Should not raise even without API key
|
||||
PrewarmManager.prewarm_provider('openai')
|
||||
|
||||
assert not PrewarmManager.is_prewarmed('openai')
|
||||
|
||||
def test_reset(self):
|
||||
"""Test reset clears state."""
|
||||
from praisonai.cli.features.profiler import PrewarmManager
|
||||
|
||||
PrewarmManager.enable()
|
||||
PrewarmManager.reset()
|
||||
|
||||
assert not PrewarmManager.is_enabled()
|
||||
|
||||
|
||||
class TestLiteModeConfig:
|
||||
"""Tests for LiteModeConfig (Tier 2)."""
|
||||
|
||||
def test_default_disabled(self):
|
||||
"""Test that lite mode is disabled by default."""
|
||||
from praisonai.cli.features.profiler import LiteModeConfig
|
||||
|
||||
config = LiteModeConfig()
|
||||
|
||||
assert not config.enabled
|
||||
assert not config.skip_type_validation
|
||||
assert not config.skip_model_validation
|
||||
assert not config.minimal_imports
|
||||
|
||||
def test_from_env(self):
|
||||
"""Test loading config from environment."""
|
||||
from praisonai.cli.features.profiler import LiteModeConfig
|
||||
|
||||
with patch.dict(os.environ, {
|
||||
'PRAISONAI_LITE_MODE': '1',
|
||||
'PRAISONAI_SKIP_TYPE_VALIDATION': 'true',
|
||||
}):
|
||||
config = LiteModeConfig.from_env()
|
||||
|
||||
assert config.enabled
|
||||
assert config.skip_type_validation
|
||||
assert not config.skip_model_validation
|
||||
|
||||
def test_is_lite_mode(self):
|
||||
"""Test is_lite_mode helper."""
|
||||
from praisonai.cli.features.profiler import is_lite_mode
|
||||
|
||||
# Reset global config
|
||||
import praisonai.cli.features.profiler.optimizations as opt
|
||||
opt._lite_mode_config = None
|
||||
|
||||
with patch.dict(os.environ, {'PRAISONAI_LITE_MODE': ''}, clear=False):
|
||||
# Force reload
|
||||
opt._lite_mode_config = None
|
||||
assert not is_lite_mode()
|
||||
|
||||
|
||||
class TestPerfSnapshot:
|
||||
"""Tests for PerfSnapshot (Tier 2)."""
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test snapshot serialization."""
|
||||
from praisonai.cli.features.profiler import PerfSnapshot
|
||||
|
||||
snapshot = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='test',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
|
||||
data = snapshot.to_dict()
|
||||
|
||||
assert data['name'] == 'test'
|
||||
assert data['startup_cold_ms'] == 100.0
|
||||
assert data['query_time_ms'] == 2000.0
|
||||
|
||||
def test_from_dict(self):
|
||||
"""Test snapshot deserialization."""
|
||||
from praisonai.cli.features.profiler import PerfSnapshot
|
||||
|
||||
data = {
|
||||
'timestamp': '2025-01-01T00:00:00Z',
|
||||
'name': 'test',
|
||||
'startup_cold_ms': 100.0,
|
||||
'startup_warm_ms': 80.0,
|
||||
'import_time_ms': 500.0,
|
||||
'query_time_ms': 2000.0,
|
||||
'first_token_ms': 300.0,
|
||||
}
|
||||
|
||||
snapshot = PerfSnapshot.from_dict(data)
|
||||
|
||||
assert snapshot.name == 'test'
|
||||
assert snapshot.startup_cold_ms == 100.0
|
||||
|
||||
|
||||
class TestPerfComparison:
|
||||
"""Tests for PerfComparison (Tier 2)."""
|
||||
|
||||
def test_diff_calculation(self):
|
||||
"""Test difference calculations."""
|
||||
from praisonai.cli.features.profiler import PerfSnapshot, PerfComparison
|
||||
|
||||
baseline = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='baseline',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
|
||||
current = PerfSnapshot(
|
||||
timestamp='2025-01-02T00:00:00Z',
|
||||
name='current',
|
||||
startup_cold_ms=110.0, # 10% slower
|
||||
startup_warm_ms=85.0,
|
||||
import_time_ms=550.0, # 10% slower
|
||||
query_time_ms=2200.0, # 10% slower
|
||||
first_token_ms=330.0,
|
||||
)
|
||||
|
||||
comparison = PerfComparison(baseline=baseline, current=current)
|
||||
|
||||
assert comparison.startup_cold_diff_ms == 10.0
|
||||
assert comparison.startup_cold_diff_pct == 10.0
|
||||
assert comparison.import_time_diff_ms == 50.0
|
||||
assert comparison.query_time_diff_ms == 200.0
|
||||
|
||||
def test_regression_detection(self):
|
||||
"""Test regression detection."""
|
||||
from praisonai.cli.features.profiler import PerfSnapshot, PerfComparison
|
||||
|
||||
baseline = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='baseline',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
|
||||
# No regression (within 10%)
|
||||
current_ok = PerfSnapshot(
|
||||
timestamp='2025-01-02T00:00:00Z',
|
||||
name='current',
|
||||
startup_cold_ms=105.0,
|
||||
startup_warm_ms=82.0,
|
||||
import_time_ms=520.0,
|
||||
query_time_ms=2100.0,
|
||||
first_token_ms=310.0,
|
||||
)
|
||||
|
||||
comparison_ok = PerfComparison(baseline=baseline, current=current_ok)
|
||||
assert not comparison_ok.is_regression()
|
||||
|
||||
# Regression (>10%)
|
||||
current_bad = PerfSnapshot(
|
||||
timestamp='2025-01-02T00:00:00Z',
|
||||
name='current',
|
||||
startup_cold_ms=150.0, # 50% slower
|
||||
startup_warm_ms=82.0,
|
||||
import_time_ms=520.0,
|
||||
query_time_ms=2100.0,
|
||||
first_token_ms=310.0,
|
||||
)
|
||||
|
||||
comparison_bad = PerfComparison(baseline=baseline, current=current_bad)
|
||||
assert comparison_bad.is_regression()
|
||||
|
||||
|
||||
class TestPerfSnapshotManager:
|
||||
"""Tests for PerfSnapshotManager (Tier 2)."""
|
||||
|
||||
def test_save_and_load(self):
|
||||
"""Test saving and loading snapshots."""
|
||||
from praisonai.cli.features.profiler import PerfSnapshot, PerfSnapshotManager
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
manager = PerfSnapshotManager(snapshot_dir=Path(tmpdir))
|
||||
|
||||
snapshot = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='test',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
|
||||
path = manager.save(snapshot)
|
||||
assert path.exists()
|
||||
|
||||
loaded = manager.load('test')
|
||||
assert loaded is not None
|
||||
assert loaded.name == 'test'
|
||||
assert loaded.startup_cold_ms == 100.0
|
||||
|
||||
def test_baseline(self):
|
||||
"""Test baseline save and load."""
|
||||
from praisonai.cli.features.profiler import PerfSnapshot, PerfSnapshotManager
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
manager = PerfSnapshotManager(snapshot_dir=Path(tmpdir))
|
||||
|
||||
snapshot = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='original',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
|
||||
manager.save_baseline(snapshot)
|
||||
|
||||
baseline = manager.load_baseline()
|
||||
assert baseline is not None
|
||||
assert baseline.name == 'baseline'
|
||||
|
||||
def test_list_snapshots(self):
|
||||
"""Test listing snapshots."""
|
||||
from praisonai.cli.features.profiler import PerfSnapshot, PerfSnapshotManager
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
manager = PerfSnapshotManager(snapshot_dir=Path(tmpdir))
|
||||
|
||||
for name in ['snap1', 'snap2', 'snap3']:
|
||||
snapshot = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name=name,
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
manager.save(snapshot)
|
||||
|
||||
names = manager.list_snapshots()
|
||||
assert 'snap1' in names
|
||||
assert 'snap2' in names
|
||||
assert 'snap3' in names
|
||||
|
||||
|
||||
class TestFormatComparisonReport:
|
||||
"""Tests for format_comparison_report."""
|
||||
|
||||
def test_format_report(self):
|
||||
"""Test report formatting."""
|
||||
from praisonai.cli.features.profiler import (
|
||||
PerfSnapshot, PerfComparison, format_comparison_report
|
||||
)
|
||||
|
||||
baseline = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='baseline',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
|
||||
current = PerfSnapshot(
|
||||
timestamp='2025-01-02T00:00:00Z',
|
||||
name='current',
|
||||
startup_cold_ms=105.0,
|
||||
startup_warm_ms=82.0,
|
||||
import_time_ms=520.0,
|
||||
query_time_ms=2100.0,
|
||||
first_token_ms=310.0,
|
||||
)
|
||||
|
||||
comparison = PerfComparison(baseline=baseline, current=current)
|
||||
report = format_comparison_report(comparison)
|
||||
|
||||
assert 'Performance Comparison Report' in report
|
||||
assert 'baseline' in report
|
||||
assert 'current' in report
|
||||
assert 'No significant regression' in report
|
||||
|
||||
|
||||
class TestModuleImports:
|
||||
"""Test that all exports are available."""
|
||||
|
||||
def test_all_exports_available(self):
|
||||
"""Test that all __all__ exports are importable."""
|
||||
from praisonai.cli.features.profiler import __all__
|
||||
import praisonai.cli.features.profiler as profiler
|
||||
|
||||
for name in __all__:
|
||||
assert hasattr(profiler, name), f"Missing export: {name}"
|
||||
@@ -0,0 +1,228 @@
|
||||
"""
|
||||
Unit tests for CLI profiler suite module.
|
||||
"""
|
||||
|
||||
|
||||
class TestScenarioConfig:
|
||||
"""Tests for ScenarioConfig dataclass."""
|
||||
|
||||
def test_default_config(self):
|
||||
"""Test default configuration values."""
|
||||
from praisonai.cli.features.profiler import ScenarioConfig
|
||||
|
||||
config = ScenarioConfig(name="test", prompt="hi")
|
||||
assert config.name == "test"
|
||||
assert config.prompt == "hi"
|
||||
assert config.model is None
|
||||
assert config.stream is False
|
||||
assert config.iterations == 3
|
||||
assert config.warmup == 1
|
||||
assert config.show_files is True
|
||||
assert config.limit == 20
|
||||
|
||||
def test_custom_config(self):
|
||||
"""Test custom configuration."""
|
||||
from praisonai.cli.features.profiler import ScenarioConfig
|
||||
|
||||
config = ScenarioConfig(
|
||||
name="custom",
|
||||
prompt="test prompt",
|
||||
model="gpt-4o",
|
||||
stream=True,
|
||||
iterations=5,
|
||||
)
|
||||
assert config.name == "custom"
|
||||
assert config.model == "gpt-4o"
|
||||
assert config.stream is True
|
||||
assert config.iterations == 5
|
||||
|
||||
|
||||
class TestScenarioResult:
|
||||
"""Tests for ScenarioResult dataclass."""
|
||||
|
||||
def test_empty_result(self):
|
||||
"""Test empty result."""
|
||||
from praisonai.cli.features.profiler import ScenarioConfig, ScenarioResult
|
||||
|
||||
config = ScenarioConfig(name="test", prompt="hi")
|
||||
result = ScenarioResult(name="test", config=config)
|
||||
|
||||
assert result.name == "test"
|
||||
assert len(result.results) == 0
|
||||
assert result.import_times == []
|
||||
assert result.total_times == []
|
||||
|
||||
def test_get_stats_empty(self):
|
||||
"""Test stats with empty values."""
|
||||
from praisonai.cli.features.profiler import ScenarioConfig, ScenarioResult
|
||||
|
||||
config = ScenarioConfig(name="test", prompt="hi")
|
||||
result = ScenarioResult(name="test", config=config)
|
||||
|
||||
stats = result.get_stats([])
|
||||
assert stats["mean"] == 0
|
||||
assert stats["min"] == 0
|
||||
assert stats["max"] == 0
|
||||
|
||||
def test_get_stats_single_value(self):
|
||||
"""Test stats with single value."""
|
||||
from praisonai.cli.features.profiler import ScenarioConfig, ScenarioResult
|
||||
|
||||
config = ScenarioConfig(name="test", prompt="hi")
|
||||
result = ScenarioResult(name="test", config=config)
|
||||
|
||||
stats = result.get_stats([100.0])
|
||||
assert stats["mean"] == 100.0
|
||||
assert stats["min"] == 100.0
|
||||
assert stats["max"] == 100.0
|
||||
assert stats["stdev"] == 0
|
||||
|
||||
def test_get_stats_multiple_values(self):
|
||||
"""Test stats with multiple values."""
|
||||
from praisonai.cli.features.profiler import ScenarioConfig, ScenarioResult
|
||||
|
||||
config = ScenarioConfig(name="test", prompt="hi")
|
||||
result = ScenarioResult(name="test", config=config)
|
||||
|
||||
stats = result.get_stats([100.0, 200.0, 300.0])
|
||||
assert stats["mean"] == 200.0
|
||||
assert stats["min"] == 100.0
|
||||
assert stats["max"] == 300.0
|
||||
assert stats["median"] == 200.0
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test conversion to dictionary."""
|
||||
from praisonai.cli.features.profiler import ScenarioConfig, ScenarioResult
|
||||
|
||||
config = ScenarioConfig(name="test", prompt="hi", iterations=2)
|
||||
result = ScenarioResult(name="test", config=config)
|
||||
|
||||
data = result.to_dict()
|
||||
assert data["name"] == "test"
|
||||
assert "config" in data
|
||||
assert data["config"]["prompt"] == "hi"
|
||||
assert "import_time_stats" in data
|
||||
assert "total_time_stats" in data
|
||||
|
||||
|
||||
class TestSuiteResult:
|
||||
"""Tests for SuiteResult dataclass."""
|
||||
|
||||
def test_default_result(self):
|
||||
"""Test default result."""
|
||||
from praisonai.cli.features.profiler import SuiteResult
|
||||
|
||||
result = SuiteResult()
|
||||
assert len(result.scenarios) == 0
|
||||
assert result.startup_cold_ms == 0.0
|
||||
assert result.startup_warm_ms == 0.0
|
||||
assert len(result.import_analysis) == 0
|
||||
assert result.timestamp.endswith("Z") # ISO format with Z suffix
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test conversion to dictionary."""
|
||||
from praisonai.cli.features.profiler import SuiteResult
|
||||
|
||||
result = SuiteResult(
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=50.0,
|
||||
)
|
||||
|
||||
data = result.to_dict()
|
||||
assert "timestamp" in data
|
||||
assert "metadata" in data
|
||||
assert data["startup"]["cold_ms"] == 100.0
|
||||
assert data["startup"]["warm_ms"] == 50.0
|
||||
|
||||
|
||||
class TestProfileSuiteRunner:
|
||||
"""Tests for ProfileSuiteRunner class."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test initialization with defaults."""
|
||||
from praisonai.cli.features.profiler import ProfileSuiteRunner
|
||||
|
||||
runner = ProfileSuiteRunner()
|
||||
assert len(runner.scenarios) == 4 # DEFAULT_SCENARIOS
|
||||
assert runner.verbose is False
|
||||
|
||||
def test_init_custom_scenarios(self):
|
||||
"""Test initialization with custom scenarios."""
|
||||
from praisonai.cli.features.profiler import ProfileSuiteRunner, ScenarioConfig
|
||||
|
||||
scenarios = [
|
||||
ScenarioConfig(name="custom1", prompt="test1"),
|
||||
ScenarioConfig(name="custom2", prompt="test2"),
|
||||
]
|
||||
|
||||
runner = ProfileSuiteRunner(scenarios=scenarios)
|
||||
assert len(runner.scenarios) == 2
|
||||
assert runner.scenarios[0].name == "custom1"
|
||||
|
||||
def test_collect_metadata(self):
|
||||
"""Test metadata collection."""
|
||||
from praisonai.cli.features.profiler import ProfileSuiteRunner
|
||||
|
||||
runner = ProfileSuiteRunner()
|
||||
metadata = runner._collect_metadata()
|
||||
|
||||
assert "python_version" in metadata
|
||||
assert "platform" in metadata
|
||||
assert "praisonai_version" in metadata
|
||||
assert "timestamp" in metadata
|
||||
|
||||
def test_measure_startup(self):
|
||||
"""Test startup measurement."""
|
||||
from praisonai.cli.features.profiler import ProfileSuiteRunner
|
||||
|
||||
runner = ProfileSuiteRunner()
|
||||
cold, warm = runner._measure_startup()
|
||||
|
||||
# Both should be positive numbers
|
||||
assert cold > 0
|
||||
assert warm > 0
|
||||
|
||||
def test_analyze_imports(self):
|
||||
"""Test import analysis."""
|
||||
from praisonai.cli.features.profiler import ProfileSuiteRunner
|
||||
|
||||
runner = ProfileSuiteRunner()
|
||||
imports = runner._analyze_imports()
|
||||
|
||||
# Should return a list (may be empty if subprocess fails)
|
||||
assert isinstance(imports, list)
|
||||
|
||||
# If we got results, check structure
|
||||
if imports:
|
||||
assert "module" in imports[0]
|
||||
assert "cumulative_ms" in imports[0]
|
||||
|
||||
|
||||
class TestProfilerModuleImports:
|
||||
"""Tests for profiler module imports."""
|
||||
|
||||
def test_all_exports_available(self):
|
||||
"""Test that all exports are available."""
|
||||
from praisonai.cli.features.profiler import (
|
||||
ProfilerConfig,
|
||||
ProfilerResult,
|
||||
QueryProfiler,
|
||||
run_profiled_query,
|
||||
format_profile_report,
|
||||
ScenarioConfig,
|
||||
ScenarioResult,
|
||||
SuiteResult,
|
||||
ProfileSuiteRunner,
|
||||
run_profile_suite,
|
||||
)
|
||||
|
||||
assert ProfilerConfig is not None
|
||||
assert ProfilerResult is not None
|
||||
assert QueryProfiler is not None
|
||||
assert run_profiled_query is not None
|
||||
assert format_profile_report is not None
|
||||
assert ScenarioConfig is not None
|
||||
assert ScenarioResult is not None
|
||||
assert SuiteResult is not None
|
||||
assert ProfileSuiteRunner is not None
|
||||
assert run_profile_suite is not None
|
||||
@@ -0,0 +1,630 @@
|
||||
"""
|
||||
Unit tests for the Unified Profiling Architecture.
|
||||
|
||||
Tests the core invariants:
|
||||
1. Profiled execution produces identical output to non-profiled execution
|
||||
2. Schema equivalence across invocation methods
|
||||
3. Overhead bounds
|
||||
4. Profile data bounds
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class TestExecutionRequest:
|
||||
"""Tests for ExecutionRequest dataclass."""
|
||||
|
||||
def test_create_basic_request(self):
|
||||
"""Test creating a basic execution request."""
|
||||
from praisonai.cli.execution import ExecutionRequest
|
||||
|
||||
req = ExecutionRequest(prompt="What is 2+2?")
|
||||
|
||||
assert req.prompt == "What is 2+2?"
|
||||
assert req.agent_name == "Assistant"
|
||||
assert req.model is None
|
||||
assert req.stream is False
|
||||
assert req.tools == ()
|
||||
|
||||
def test_request_immutability(self):
|
||||
"""Test that ExecutionRequest is immutable (frozen)."""
|
||||
from praisonai.cli.execution import ExecutionRequest
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
|
||||
with pytest.raises(AttributeError):
|
||||
req.prompt = "modified"
|
||||
|
||||
def test_request_with_model(self):
|
||||
"""Test creating request with model."""
|
||||
from praisonai.cli.execution import ExecutionRequest
|
||||
|
||||
req = ExecutionRequest(prompt="test", model="gpt-4")
|
||||
new_req = req.with_model("gpt-3.5-turbo")
|
||||
|
||||
assert req.model == "gpt-4"
|
||||
assert new_req.model == "gpt-3.5-turbo"
|
||||
assert req.prompt == new_req.prompt
|
||||
|
||||
def test_request_validation(self):
|
||||
"""Test that empty prompt raises error."""
|
||||
from praisonai.cli.execution import ExecutionRequest
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ExecutionRequest(prompt="")
|
||||
|
||||
def test_request_to_dict(self):
|
||||
"""Test serialization to dict."""
|
||||
from praisonai.cli.execution import ExecutionRequest
|
||||
|
||||
req = ExecutionRequest(
|
||||
prompt="test",
|
||||
agent_name="TestAgent",
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
d = req.to_dict()
|
||||
|
||||
assert d["prompt"] == "test"
|
||||
assert d["agent_name"] == "TestAgent"
|
||||
assert d["model"] == "gpt-4"
|
||||
assert d["tools_count"] == 0
|
||||
|
||||
|
||||
class TestExecutionResult:
|
||||
"""Tests for ExecutionResult dataclass."""
|
||||
|
||||
def test_create_basic_result(self):
|
||||
"""Test creating a basic execution result."""
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
result = ExecutionResult(output="four")
|
||||
|
||||
assert result.output == "four"
|
||||
assert result.success is True
|
||||
assert result.error is None
|
||||
|
||||
def test_result_from_error(self):
|
||||
"""Test creating result from error."""
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
result = ExecutionResult.from_error("Something went wrong")
|
||||
|
||||
assert result.output == ""
|
||||
assert result.success is False
|
||||
assert result.error == "Something went wrong"
|
||||
|
||||
def test_result_duration(self):
|
||||
"""Test duration calculation."""
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
result = ExecutionResult(
|
||||
output="test",
|
||||
start_time=1000.0,
|
||||
end_time=1002.5,
|
||||
)
|
||||
|
||||
assert result.duration_ms == 2500.0
|
||||
|
||||
|
||||
class TestProfilerConfig:
|
||||
"""Tests for ProfilerConfig."""
|
||||
|
||||
def test_default_config(self):
|
||||
"""Test default profiler configuration."""
|
||||
from praisonai.cli.execution.profiler import ProfilerConfig
|
||||
|
||||
config = ProfilerConfig()
|
||||
|
||||
assert config.layer == 1
|
||||
assert config.limit == 30
|
||||
assert config.track_network is False
|
||||
|
||||
def test_from_flags_basic(self):
|
||||
"""Test creating config from basic --profile flag."""
|
||||
from praisonai.cli.execution.profiler import ProfilerConfig
|
||||
|
||||
config = ProfilerConfig.from_flags(profile=True)
|
||||
|
||||
assert config.layer == 1
|
||||
|
||||
def test_from_flags_deep(self):
|
||||
"""Test creating config from --profile --deep flags."""
|
||||
from praisonai.cli.execution.profiler import ProfilerConfig
|
||||
|
||||
config = ProfilerConfig.from_flags(profile=True, deep=True)
|
||||
|
||||
assert config.layer == 2
|
||||
assert config.show_callers is True
|
||||
assert config.show_callees is True
|
||||
|
||||
def test_from_flags_network(self):
|
||||
"""Test creating config with network tracking."""
|
||||
from praisonai.cli.execution.profiler import ProfilerConfig
|
||||
|
||||
config = ProfilerConfig.from_flags(profile=True, network=True)
|
||||
|
||||
assert config.track_network is True
|
||||
|
||||
|
||||
class TestTimingBreakdown:
|
||||
"""Tests for TimingBreakdown."""
|
||||
|
||||
def test_timing_to_dict(self):
|
||||
"""Test timing serialization."""
|
||||
from praisonai.cli.execution.profiler import TimingBreakdown
|
||||
|
||||
timing = TimingBreakdown(
|
||||
total_ms=1000.0,
|
||||
imports_ms=500.0,
|
||||
agent_init_ms=10.0,
|
||||
execution_ms=490.0,
|
||||
)
|
||||
|
||||
d = timing.to_dict()
|
||||
|
||||
assert d["total_ms"] == 1000.0
|
||||
assert d["imports_ms"] == 500.0
|
||||
assert d["agent_init_ms"] == 10.0
|
||||
assert d["execution_ms"] == 490.0
|
||||
|
||||
|
||||
class TestProfileReport:
|
||||
"""Tests for ProfileReport."""
|
||||
|
||||
def test_report_schema_version(self):
|
||||
"""Test that schema version is set."""
|
||||
from praisonai.cli.execution.profiler import ProfileReport, SCHEMA_VERSION
|
||||
|
||||
report = ProfileReport(run_id="test123")
|
||||
|
||||
assert report.schema_version == SCHEMA_VERSION
|
||||
assert report.schema_version == "1.0"
|
||||
|
||||
def test_report_timestamp(self):
|
||||
"""Test that timestamp is auto-generated."""
|
||||
from praisonai.cli.execution.profiler import ProfileReport
|
||||
|
||||
report = ProfileReport(run_id="test123")
|
||||
|
||||
assert report.timestamp is not None
|
||||
assert "T" in report.timestamp # ISO 8601 format
|
||||
|
||||
def test_report_to_dict(self):
|
||||
"""Test report serialization."""
|
||||
from praisonai.cli.execution.profiler import (
|
||||
ProfileReport, TimingBreakdown, InvocationInfo
|
||||
)
|
||||
|
||||
report = ProfileReport(
|
||||
run_id="test123",
|
||||
timing=TimingBreakdown(total_ms=100),
|
||||
invocation=InvocationInfo(method="test"),
|
||||
)
|
||||
|
||||
d = report.to_dict()
|
||||
|
||||
assert d["schema_version"] == "1.0"
|
||||
assert d["run_id"] == "test123"
|
||||
assert "timing" in d
|
||||
assert "invocation" in d
|
||||
|
||||
def test_report_to_json(self):
|
||||
"""Test JSON serialization."""
|
||||
from praisonai.cli.execution.profiler import ProfileReport
|
||||
import json
|
||||
|
||||
report = ProfileReport(run_id="test123")
|
||||
|
||||
json_str = report.to_json()
|
||||
parsed = json.loads(json_str)
|
||||
|
||||
assert parsed["run_id"] == "test123"
|
||||
|
||||
def test_report_to_text(self):
|
||||
"""Test text formatting."""
|
||||
from praisonai.cli.execution.profiler import ProfileReport
|
||||
|
||||
report = ProfileReport(run_id="test123")
|
||||
|
||||
text = report.to_text()
|
||||
|
||||
assert "test123" in text
|
||||
assert "Timing Breakdown" in text
|
||||
|
||||
|
||||
class TestProfilerBounds:
|
||||
"""Tests for profile data bounds."""
|
||||
|
||||
def test_max_function_stats(self):
|
||||
"""Test that function stats are bounded."""
|
||||
from praisonai.cli.execution.profiler import MAX_FUNCTION_STATS
|
||||
|
||||
assert MAX_FUNCTION_STATS == 1000
|
||||
|
||||
def test_max_call_graph_edges(self):
|
||||
"""Test that call graph edges are bounded."""
|
||||
from praisonai.cli.execution.profiler import MAX_CALL_GRAPH_EDGES
|
||||
|
||||
assert MAX_CALL_GRAPH_EDGES == 5000
|
||||
|
||||
def test_max_network_requests(self):
|
||||
"""Test that network requests are bounded."""
|
||||
from praisonai.cli.execution.profiler import MAX_NETWORK_REQUESTS
|
||||
|
||||
assert MAX_NETWORK_REQUESTS == 100
|
||||
|
||||
|
||||
class TestSchemaStability:
|
||||
"""Tests for schema stability guarantees."""
|
||||
|
||||
def test_mandatory_fields_present(self):
|
||||
"""Test that all mandatory fields are present in report."""
|
||||
from praisonai.cli.execution.profiler import ProfileReport
|
||||
|
||||
report = ProfileReport(run_id="test")
|
||||
d = report.to_dict()
|
||||
|
||||
# Mandatory fields
|
||||
assert "schema_version" in d
|
||||
assert "run_id" in d
|
||||
assert "timestamp" in d
|
||||
assert "invocation" in d
|
||||
assert "timing" in d
|
||||
assert "response_preview" in d
|
||||
|
||||
def test_optional_fields_nullable(self):
|
||||
"""Test that optional fields can be None."""
|
||||
from praisonai.cli.execution.profiler import ProfileReport
|
||||
|
||||
report = ProfileReport(run_id="test")
|
||||
d = report.to_dict()
|
||||
|
||||
# Optional fields should not be present if None
|
||||
assert "functions" not in d or d["functions"] is None
|
||||
assert "call_graph" not in d or d["call_graph"] is None
|
||||
assert "network" not in d or d["network"] is None
|
||||
|
||||
|
||||
class TestInvocationInfo:
|
||||
"""Tests for InvocationInfo."""
|
||||
|
||||
def test_version_auto_populated(self):
|
||||
"""Test that versions are auto-populated."""
|
||||
from praisonai.cli.execution.profiler import InvocationInfo
|
||||
|
||||
info = InvocationInfo(method="test")
|
||||
|
||||
assert info.python_version is not None
|
||||
assert len(info.python_version) > 0
|
||||
|
||||
|
||||
class TestCallGraph:
|
||||
"""Tests for CallGraph."""
|
||||
|
||||
def test_edges_property(self):
|
||||
"""Test edges property."""
|
||||
from praisonai.cli.execution.profiler import CallGraph
|
||||
|
||||
graph = CallGraph(
|
||||
callers={"func_b": ["func_a"]},
|
||||
callees={"func_a": ["func_b"]},
|
||||
)
|
||||
|
||||
edges = graph.edges
|
||||
|
||||
assert len(edges) == 1
|
||||
assert ("func_a", "func_b") in edges
|
||||
|
||||
|
||||
class TestRequestTiming:
|
||||
"""Tests for RequestTiming."""
|
||||
|
||||
def test_ttfb_calculation(self):
|
||||
"""Test time to first byte calculation."""
|
||||
from praisonai.cli.execution.profiler import RequestTiming
|
||||
|
||||
timing = RequestTiming(
|
||||
url="https://api.openai.com",
|
||||
method="POST",
|
||||
start_time=1000.0,
|
||||
first_byte_time=1001.5,
|
||||
end_time=1002.0,
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
assert timing.ttfb_ms == 1500.0
|
||||
assert timing.total_ms == 2000.0
|
||||
|
||||
|
||||
# Integration tests (require mocking)
|
||||
class TestProfilerIntegration:
|
||||
"""Integration tests for Profiler class."""
|
||||
|
||||
def test_profiler_creates_report(self):
|
||||
"""Test that profiler creates a valid report."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
# Mock at the profiler level by patching the imported _execute_core
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="four",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
config = ProfilerConfig(layer=0) # Minimal profiling
|
||||
profiler = Profiler(config)
|
||||
|
||||
req = ExecutionRequest(prompt="What is 2+2?")
|
||||
result, report = profiler.profile_sync(req)
|
||||
|
||||
assert result.output == "four"
|
||||
assert report.schema_version == "1.0"
|
||||
assert report.timing.total_ms > 0
|
||||
|
||||
def test_output_equivalence_profiled_vs_non_profiled(self):
|
||||
"""INVARIANT 1: Profiled execution produces identical output."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
expected_output = "test_output_value"
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output=expected_output,
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
|
||||
# Layer 0 (minimal)
|
||||
config0 = ProfilerConfig(layer=0)
|
||||
result0, _ = Profiler(config0).profile_sync(req)
|
||||
|
||||
# Layer 1 (basic)
|
||||
config1 = ProfilerConfig(layer=1)
|
||||
result1, _ = Profiler(config1).profile_sync(req)
|
||||
|
||||
# Layer 2 (deep)
|
||||
config2 = ProfilerConfig(layer=2)
|
||||
result2, _ = Profiler(config2).profile_sync(req)
|
||||
|
||||
# All outputs must be identical
|
||||
assert result0.output == expected_output
|
||||
assert result1.output == expected_output
|
||||
assert result2.output == expected_output
|
||||
|
||||
def test_schema_equivalence_across_invocation_methods(self):
|
||||
"""INVARIANT 2: Schema equivalence across invocation methods."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
# CLI direct
|
||||
_, report1 = Profiler(config).profile_sync(req, invocation_method="cli_direct")
|
||||
|
||||
# Profile command
|
||||
_, report2 = Profiler(config).profile_sync(req, invocation_method="profile_command")
|
||||
|
||||
# Schema must be identical
|
||||
dict1 = report1.to_dict()
|
||||
dict2 = report2.to_dict()
|
||||
|
||||
assert dict1["schema_version"] == dict2["schema_version"]
|
||||
assert set(dict1["timing"].keys()) == set(dict2["timing"].keys())
|
||||
assert set(dict1["invocation"].keys()) == set(dict2["invocation"].keys())
|
||||
|
||||
def test_data_bounds_function_stats(self):
|
||||
"""INVARIANT 4: Function stats are bounded."""
|
||||
from praisonai.cli.execution.profiler import (
|
||||
Profiler, ProfilerConfig, MAX_FUNCTION_STATS
|
||||
)
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=1, limit=50)
|
||||
|
||||
_, report = Profiler(config).profile_sync(req)
|
||||
|
||||
# Functions should be bounded by limit
|
||||
if report.functions:
|
||||
assert len(report.functions) <= config.limit
|
||||
assert len(report.functions) <= MAX_FUNCTION_STATS
|
||||
|
||||
def test_layer_0_minimal_overhead(self):
|
||||
"""INVARIANT 3: Layer 0 has minimal overhead (<1ms)."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
import time
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
# Fast mock execution
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=0)
|
||||
|
||||
start = time.perf_counter()
|
||||
Profiler(config).profile_sync(req)
|
||||
elapsed = (time.perf_counter() - start) * 1000
|
||||
|
||||
# Layer 0 overhead should be minimal (allow some margin for test environment)
|
||||
# The actual profiling overhead (excluding execution) should be < 1ms
|
||||
# We can't measure this precisely with mocks, but we verify it runs fast
|
||||
assert elapsed < 100 # Very generous bound for test environment
|
||||
|
||||
def test_ttfr_is_populated(self):
|
||||
"""Test that Time to First Response (TTFR) is populated and > 0."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
_, report = Profiler(config).profile_sync(req)
|
||||
|
||||
# TTFR must be populated
|
||||
assert report.timing.time_to_first_response_ms > 0
|
||||
assert report.timing.first_output_ms >= 0
|
||||
|
||||
def test_first_output_less_than_total(self):
|
||||
"""Test that first_output_ms <= total_ms."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
_, report = Profiler(config).profile_sync(req)
|
||||
|
||||
# first_output must be <= total
|
||||
assert report.timing.first_output_ms <= report.timing.total_ms
|
||||
|
||||
def test_timeline_phases_ordered(self):
|
||||
"""Test that timeline phases are correctly ordered."""
|
||||
from praisonai.cli.execution.profiler import TimingBreakdown
|
||||
|
||||
timing = TimingBreakdown(
|
||||
total_ms=1000.0,
|
||||
imports_ms=200.0,
|
||||
agent_init_ms=50.0,
|
||||
execution_ms=700.0,
|
||||
)
|
||||
|
||||
timeline = timing.to_timeline()
|
||||
|
||||
# Timeline should have entries
|
||||
assert len(timeline) > 0
|
||||
|
||||
# Each entry should have (name, start, duration)
|
||||
for entry in timeline:
|
||||
assert len(entry) == 3
|
||||
name, start, duration = entry
|
||||
assert isinstance(name, str)
|
||||
assert isinstance(start, float)
|
||||
assert isinstance(duration, float)
|
||||
|
||||
def test_decision_trace_in_deep_profile(self):
|
||||
"""Test that decision trace is present in deep profile (layer 2)."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test", model="gpt-4")
|
||||
config = ProfilerConfig(layer=2) # Deep profile
|
||||
|
||||
_, report = Profiler(config).profile_sync(req)
|
||||
|
||||
# Decision trace must be present in deep profile
|
||||
assert report.decision_trace is not None
|
||||
assert report.decision_trace.model_selected == "gpt-4"
|
||||
assert report.decision_trace.profile_layer == 2
|
||||
|
||||
def test_module_breakdown_in_deep_profile(self):
|
||||
"""Test that module breakdown is present in deep profile (layer 2)."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=2) # Deep profile
|
||||
|
||||
_, report = Profiler(config).profile_sync(req)
|
||||
|
||||
# Module breakdown must be present in deep profile
|
||||
assert report.module_breakdown is not None
|
||||
|
||||
def test_text_output_includes_timeline(self):
|
||||
"""Test that text output includes timeline section."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
_, report = Profiler(config).profile_sync(req)
|
||||
text = report.to_text()
|
||||
|
||||
# Text output must include timeline and TTFR
|
||||
assert "Execution Timeline" in text
|
||||
assert "Time to First Response" in text
|
||||
|
||||
def test_json_output_includes_ttfr(self):
|
||||
"""Test that JSON output includes time_to_first_response_ms."""
|
||||
from praisonai.cli.execution.profiler import Profiler, ProfilerConfig
|
||||
from praisonai.cli.execution.request import ExecutionRequest
|
||||
from praisonai.cli.execution.result import ExecutionResult
|
||||
import json
|
||||
|
||||
with patch('praisonai.cli.execution.profiler._execute_core') as mock_exec:
|
||||
mock_exec.return_value = ExecutionResult(
|
||||
output="test",
|
||||
run_id="test123",
|
||||
)
|
||||
|
||||
req = ExecutionRequest(prompt="test")
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
_, report = Profiler(config).profile_sync(req)
|
||||
data = json.loads(report.to_json())
|
||||
|
||||
# JSON must include TTFR
|
||||
assert "time_to_first_response_ms" in data["timing"]
|
||||
assert data["timing"]["time_to_first_response_ms"] > 0
|
||||
Whitespace-only changes.
@@ -0,0 +1,330 @@
|
||||
"""
|
||||
TDD Test Suite for PraisonAI TUI Widgets.
|
||||
|
||||
Tests for:
|
||||
- ComposerWidget (Enter to send, Shift+Enter for newline)
|
||||
- HelpScreen (F1 key)
|
||||
- Keyboard shortcuts (Ctrl+Q, Ctrl+C, Ctrl+L)
|
||||
- Queue panel (edit/delete)
|
||||
- Tool panel (display tools)
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
|
||||
# Check if textual is available
|
||||
try:
|
||||
from textual.pilot import Pilot
|
||||
from textual.app import App
|
||||
TEXTUAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
TEXTUAL_AVAILABLE = False
|
||||
pytest.skip("Textual not available", allow_module_level=True)
|
||||
|
||||
|
||||
class TestComposerWidget:
|
||||
"""Tests for ComposerWidget."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_text_area_enter_triggers_submit(self):
|
||||
"""Test that Enter key triggers submit."""
|
||||
from praisonai.cli.features.tui.widgets.composer import SubmitTextArea
|
||||
|
||||
# Create a mock app context
|
||||
text_area = SubmitTextArea("")
|
||||
|
||||
# Verify the class exists and has SubmitRequested message
|
||||
assert hasattr(SubmitTextArea, 'SubmitRequested')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composer_widget_has_correct_bindings(self):
|
||||
"""Test ComposerWidget has correct key bindings."""
|
||||
from praisonai.cli.features.tui.widgets.composer import ComposerWidget
|
||||
|
||||
# Check bindings include enter for submit
|
||||
binding_keys = [b.key for b in ComposerWidget.BINDINGS]
|
||||
assert "enter" in binding_keys
|
||||
assert "escape" in binding_keys
|
||||
|
||||
|
||||
class TestHelpScreen:
|
||||
"""Tests for HelpScreen."""
|
||||
|
||||
def test_help_screen_exists(self):
|
||||
"""Test HelpScreen can be imported."""
|
||||
from praisonai.cli.features.tui.screens.help import HelpScreen
|
||||
assert HelpScreen is not None
|
||||
|
||||
def test_help_screen_has_dismiss_binding(self):
|
||||
"""Test HelpScreen has escape binding to dismiss."""
|
||||
from praisonai.cli.features.tui.screens.help import HelpScreen
|
||||
|
||||
binding_keys = [b.key for b in HelpScreen.BINDINGS]
|
||||
assert "escape" in binding_keys
|
||||
assert "q" in binding_keys
|
||||
|
||||
|
||||
class TestTUIApp:
|
||||
"""Tests for TUIApp."""
|
||||
|
||||
def test_tui_app_has_help_screen_registered(self):
|
||||
"""Test TUIApp has help screen in SCREENS."""
|
||||
from praisonai.cli.features.tui.app import TUIApp
|
||||
|
||||
assert "help" in TUIApp.SCREENS
|
||||
assert "main" in TUIApp.SCREENS
|
||||
assert "settings" in TUIApp.SCREENS
|
||||
assert "queue" in TUIApp.SCREENS
|
||||
|
||||
def test_tui_app_has_ctrl_bindings(self):
|
||||
"""Test TUIApp has Ctrl key bindings."""
|
||||
from praisonai.cli.features.tui.app import TUIApp
|
||||
|
||||
binding_keys = [b.key for b in TUIApp.BINDINGS]
|
||||
assert "ctrl+q" in binding_keys
|
||||
assert "ctrl+c" in binding_keys
|
||||
|
||||
|
||||
class TestMainScreen:
|
||||
"""Tests for MainScreen."""
|
||||
|
||||
def test_main_screen_has_safe_default_bindings(self):
|
||||
"""Test MainScreen has safe default bindings (no Ctrl/Fn by default)."""
|
||||
from praisonai.cli.features.tui.screens.main import MainScreen
|
||||
|
||||
binding_keys = [b.key for b in MainScreen.BINDINGS]
|
||||
# Safe defaults - single char shortcuts
|
||||
assert "q" in binding_keys # quit
|
||||
assert "colon" in binding_keys # command mode
|
||||
assert "escape" in binding_keys # cancel
|
||||
# Ctrl keys should NOT be in default bindings
|
||||
assert "ctrl+q" not in binding_keys
|
||||
assert "ctrl+c" not in binding_keys
|
||||
# Function keys should NOT be in default bindings
|
||||
assert "f1" not in binding_keys
|
||||
|
||||
|
||||
class TestTUIConfig:
|
||||
"""Tests for TUIConfig."""
|
||||
|
||||
def test_default_config_disables_ctrl_keys(self):
|
||||
"""Test default config has Ctrl keys disabled."""
|
||||
from praisonai.cli.features.tui.config import TUIConfig
|
||||
|
||||
config = TUIConfig()
|
||||
assert config.enable_ctrl_keys is False
|
||||
assert config.enable_fn_keys is False
|
||||
assert config.enable_single_char is True
|
||||
assert config.leader_key == ":"
|
||||
|
||||
def test_command_parsing(self):
|
||||
"""Test command mode parsing."""
|
||||
from praisonai.cli.features.tui.config import get_command
|
||||
|
||||
assert get_command("quit") == "quit"
|
||||
assert get_command("q") == "quit"
|
||||
assert get_command("exit") == "quit"
|
||||
assert get_command("clear") == "clear"
|
||||
assert get_command("cl") == "clear"
|
||||
assert get_command("invalid") is None
|
||||
|
||||
|
||||
class TestQueuePanel:
|
||||
"""Tests for QueuePanelWidget."""
|
||||
|
||||
def test_queue_panel_has_delete_message(self):
|
||||
"""Test QueuePanelWidget has DeleteRequested message."""
|
||||
from praisonai.cli.features.tui.widgets.queue_panel import QueuePanelWidget
|
||||
|
||||
assert hasattr(QueuePanelWidget, 'DeleteRequested')
|
||||
assert hasattr(QueuePanelWidget, 'EditRequested')
|
||||
|
||||
def test_queue_panel_has_delete_method(self):
|
||||
"""Test QueuePanelWidget has delete_selected method."""
|
||||
from praisonai.cli.features.tui.widgets.queue_panel import QueuePanelWidget
|
||||
|
||||
assert hasattr(QueuePanelWidget, 'delete_selected')
|
||||
assert hasattr(QueuePanelWidget, 'edit_selected')
|
||||
|
||||
|
||||
class TestToolPanel:
|
||||
"""Tests for ToolPanelWidget."""
|
||||
|
||||
def test_tool_panel_has_set_available_tools(self):
|
||||
"""Test ToolPanelWidget has set_available_tools method."""
|
||||
from praisonai.cli.features.tui.widgets.tool_panel import ToolPanelWidget
|
||||
|
||||
assert hasattr(ToolPanelWidget, 'set_available_tools')
|
||||
|
||||
|
||||
class TestInteractiveTools:
|
||||
"""Tests for interactive tools loading."""
|
||||
|
||||
def test_get_interactive_tools_exists(self):
|
||||
"""Test get_interactive_tools function exists."""
|
||||
from praisonai.cli.features.interactive_tools import get_interactive_tools
|
||||
|
||||
assert callable(get_interactive_tools)
|
||||
|
||||
def test_tool_config_exists(self):
|
||||
"""Test ToolConfig class exists."""
|
||||
from praisonai.cli.features.interactive_tools import ToolConfig
|
||||
|
||||
assert ToolConfig is not None
|
||||
|
||||
|
||||
class TestCommandPopup:
|
||||
"""Tests for CommandPopupWidget."""
|
||||
|
||||
def test_command_popup_exists(self):
|
||||
"""Test CommandPopupWidget can be imported."""
|
||||
from praisonai.cli.features.tui.widgets.command_popup import CommandPopupWidget
|
||||
assert CommandPopupWidget is not None
|
||||
|
||||
def test_command_info_dataclass(self):
|
||||
"""Test CommandInfo dataclass works correctly."""
|
||||
from praisonai.cli.features.tui.widgets.command_popup import CommandInfo
|
||||
|
||||
cmd = CommandInfo(
|
||||
name="test",
|
||||
description="Test command",
|
||||
aliases=["t", "tst"]
|
||||
)
|
||||
assert cmd.name == "test"
|
||||
assert cmd.description == "Test command"
|
||||
assert cmd.aliases == ["t", "tst"]
|
||||
assert cmd.display_name == "test (t, tst)"
|
||||
|
||||
def test_get_default_commands(self):
|
||||
"""Test get_default_commands returns command list."""
|
||||
from praisonai.cli.features.tui.widgets.command_popup import get_default_commands
|
||||
|
||||
commands = get_default_commands()
|
||||
assert len(commands) > 0
|
||||
|
||||
# Check for essential commands
|
||||
command_names = [c.name for c in commands]
|
||||
assert "help" in command_names
|
||||
assert "exit" in command_names
|
||||
assert "clear" in command_names
|
||||
assert "model" in command_names
|
||||
|
||||
def test_command_popup_has_messages(self):
|
||||
"""Test CommandPopupWidget has required message classes."""
|
||||
from praisonai.cli.features.tui.widgets.command_popup import CommandPopupWidget
|
||||
|
||||
assert hasattr(CommandPopupWidget, 'CommandSelected')
|
||||
assert hasattr(CommandPopupWidget, 'Dismissed')
|
||||
|
||||
|
||||
class TestMainScreenBackslash:
|
||||
"""Tests for MainScreen backslash command support."""
|
||||
|
||||
def test_main_screen_has_backslash_binding(self):
|
||||
"""Test MainScreen has backslash binding for commands."""
|
||||
from praisonai.cli.features.tui.screens.main import MainScreen
|
||||
|
||||
binding_keys = [b.key for b in MainScreen.BINDINGS]
|
||||
assert "backslash" in binding_keys
|
||||
|
||||
def test_main_screen_has_show_commands_action(self):
|
||||
"""Test MainScreen has action_show_commands method."""
|
||||
from praisonai.cli.features.tui.screens.main import MainScreen
|
||||
|
||||
assert hasattr(MainScreen, 'action_show_commands')
|
||||
|
||||
def test_main_screen_has_command_popup_handlers(self):
|
||||
"""Test MainScreen has command popup event handlers."""
|
||||
from praisonai.cli.features.tui.screens.main import MainScreen
|
||||
|
||||
assert hasattr(MainScreen, 'on_command_popup_widget_command_selected')
|
||||
assert hasattr(MainScreen, 'on_command_popup_widget_dismissed')
|
||||
|
||||
|
||||
class TestAutoTriggerEvents:
|
||||
"""Tests for auto-trigger events (backslash, @, /)."""
|
||||
|
||||
def test_composer_has_backslash_typed_message(self):
|
||||
"""Test ComposerWidget has BackslashTyped message."""
|
||||
from praisonai.cli.features.tui.widgets.composer import ComposerWidget
|
||||
|
||||
assert hasattr(ComposerWidget, 'BackslashTyped')
|
||||
|
||||
def test_composer_has_at_typed_message(self):
|
||||
"""Test ComposerWidget has AtTyped message for @ file autocomplete."""
|
||||
from praisonai.cli.features.tui.widgets.composer import ComposerWidget
|
||||
|
||||
assert hasattr(ComposerWidget, 'AtTyped')
|
||||
|
||||
def test_composer_has_slash_typed_message(self):
|
||||
"""Test ComposerWidget has SlashTyped message for / commands."""
|
||||
from praisonai.cli.features.tui.widgets.composer import ComposerWidget
|
||||
|
||||
assert hasattr(ComposerWidget, 'SlashTyped')
|
||||
|
||||
def test_main_screen_has_auto_trigger_handlers(self):
|
||||
"""Test MainScreen has auto-trigger event handlers."""
|
||||
from praisonai.cli.features.tui.screens.main import MainScreen
|
||||
|
||||
assert hasattr(MainScreen, 'on_composer_widget_backslash_typed')
|
||||
assert hasattr(MainScreen, 'on_composer_widget_at_typed')
|
||||
assert hasattr(MainScreen, 'on_composer_widget_slash_typed')
|
||||
|
||||
|
||||
class TestFilePopup:
|
||||
"""Tests for FilePopupWidget."""
|
||||
|
||||
def test_file_popup_exists(self):
|
||||
"""Test FilePopupWidget can be imported."""
|
||||
from praisonai.cli.features.tui.widgets.file_popup import FilePopupWidget
|
||||
assert FilePopupWidget is not None
|
||||
|
||||
def test_file_popup_has_messages(self):
|
||||
"""Test FilePopupWidget has required message classes."""
|
||||
from praisonai.cli.features.tui.widgets.file_popup import FilePopupWidget
|
||||
|
||||
assert hasattr(FilePopupWidget, 'FileSelected')
|
||||
assert hasattr(FilePopupWidget, 'Dismissed')
|
||||
|
||||
def test_file_info_dataclass(self):
|
||||
"""Test FileInfo dataclass works correctly."""
|
||||
from praisonai.cli.features.tui.widgets.file_popup import FileInfo
|
||||
|
||||
file_info = FileInfo(path="test.py", file_type="file")
|
||||
assert file_info.path == "test.py"
|
||||
assert file_info.file_type == "file"
|
||||
assert file_info.icon == "📄 "
|
||||
|
||||
dir_info = FileInfo(path="src/", file_type="directory")
|
||||
assert dir_info.icon == "📁 "
|
||||
|
||||
def test_main_screen_has_file_popup_handlers(self):
|
||||
"""Test MainScreen has file popup event handlers."""
|
||||
from praisonai.cli.features.tui.screens.main import MainScreen
|
||||
|
||||
assert hasattr(MainScreen, 'on_file_popup_widget_file_selected')
|
||||
assert hasattr(MainScreen, 'on_file_popup_widget_dismissed')
|
||||
assert hasattr(MainScreen, '_show_file_popup')
|
||||
assert hasattr(MainScreen, '_dismiss_file_popup')
|
||||
|
||||
|
||||
class TestQueueLiveUpdates:
|
||||
"""Tests for queue panel live updates."""
|
||||
|
||||
def test_app_has_update_queue_panel_live(self):
|
||||
"""Test TUIApp has _update_queue_panel_live method."""
|
||||
from praisonai.cli.features.tui.app import TUIApp
|
||||
|
||||
assert hasattr(TUIApp, '_update_queue_panel_live')
|
||||
|
||||
def test_queue_panel_has_reactive_counts(self):
|
||||
"""Test QueuePanelWidget has reactive count properties."""
|
||||
from praisonai.cli.features.tui.widgets.queue_panel import QueuePanelWidget
|
||||
|
||||
# Check reactive properties exist
|
||||
assert hasattr(QueuePanelWidget, 'queued_count')
|
||||
assert hasattr(QueuePanelWidget, 'running_count')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Generated
+241
-8
@@ -133,6 +133,19 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/b8/62/c9fa5bafe03186a0e4699150a7fed9b1e73240996d0d2f0e5f70f3fdf471/aiohttp-3.11.11-cp313-cp313-win_amd64.whl", hash = "sha256:c7a06301c2fb096bdb0bd25fe2011531c1453b9f2c163c8031600ec73af1cc99", size = 436081, upload-time = "2024-12-18T21:20:04.557Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aiohttp-jinja2"
|
||||
version = "1.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
{ name = "jinja2" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e6/39/da5a94dd89b1af7241fb7fc99ae4e73505b5f898b540b6aba6dc7afe600e/aiohttp-jinja2-1.6.tar.gz", hash = "sha256:a3a7ff5264e5bca52e8ae547bbfd0761b72495230d438d05b6c0915be619b0e2", size = 53057, upload-time = "2023-11-18T15:30:52.559Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/eb/90/65238d4246307195411b87a07d03539049819b022c01bcc773826f600138/aiohttp_jinja2-1.6-py3-none-any.whl", hash = "sha256:0df405ee6ad1b58e5a068a105407dc7dcc1704544c559f1938babde954f945c7", size = 11736, upload-time = "2023-11-18T15:30:50.743Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aiohttp-retry"
|
||||
version = "2.8.3"
|
||||
@@ -398,6 +411,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/df/73/b6e24bd22e6720ca8ee9a85a0c4a2971af8497d8f3193fa05390cbd46e09/backoff-2.2.1-py3-none-any.whl", hash = "sha256:63579f9a0628e06278f7e47b7d7d5b6ce20dc65c5e96a6f3ca99a6adca0396e8", size = 15148, upload-time = "2022-10-05T19:19:30.546Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "backports-asyncio-runner"
|
||||
version = "1.2.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/8e/ff/70dca7d7cb1cbc0edb2c6cc0c38b65cba36cccc491eca64cabd5fe7f8670/backports_asyncio_runner-1.2.0.tar.gz", hash = "sha256:a5aa7b2b7d8f8bfcaa2b57313f70792df84e32a2a746f585213373f900b42162", size = 69893, upload-time = "2025-07-02T02:27:15.685Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/59/76ab57e3fe74484f48a53f8e337171b4a2349e506eabe136d7e01d059086/backports_asyncio_runner-1.2.0-py3-none-any.whl", hash = "sha256:0da0a936a8aeb554eccb426dc55af3ba63bcdc69fa1a600b5bb305413a4477b5", size = 12313, upload-time = "2025-07-02T02:27:14.263Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bcrypt"
|
||||
version = "4.2.1"
|
||||
@@ -1817,6 +1839,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/59/91/aa6bde563e0085a02a435aa99b49ef75b0a4b062635e606dab23ce18d720/inflection-0.5.1-py2.py3-none-any.whl", hash = "sha256:f38b2b640938a4f35ade69ac3d053042959b62a0f1076a5bbaa1b9526605a8a2", size = 9454, upload-time = "2020-08-22T08:16:27.816Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iniconfig"
|
||||
version = "2.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "instructor"
|
||||
version = "1.3.3"
|
||||
@@ -2283,6 +2314,18 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d9/28/3be76b591a2e14a031b681b8283acf1dec2ad521f6f1701b7957df68c466/lightgbm-4.5.0-py3-none-win_amd64.whl", hash = "sha256:7ccb73ee9fb74fbbf89ad24c57a6edad505aa8f2165d02b999a082dbbbb0ee57", size = 1444319, upload-time = "2024-07-26T14:39:10.001Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "linkify-it-py"
|
||||
version = "2.0.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "uc-micro-py" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/2a/ae/bb56c6828e4797ba5a4821eec7c43b8bf40f69cda4d4f5f8c8a2810ec96a/linkify-it-py-2.0.3.tar.gz", hash = "sha256:68cda27e162e9215c17d786649d1da0021a451bdc436ef9e0fa0ba5234b9b048", size = 27946, upload-time = "2024-02-04T14:48:04.179Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/04/1e/b832de447dee8b582cac175871d2f6c3d5077cc56d5575cadba1fd1cccfa/linkify_it_py-2.0.3-py3-none-any.whl", hash = "sha256:6bcbc417b0ac14323382aef5c5192c0075bf8a9d6b41820a2b66371eac6b6d79", size = 19820, upload-time = "2024-02-04T14:48:02.496Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lit"
|
||||
version = "18.1.8"
|
||||
@@ -2430,6 +2473,14 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/42/d7/1ec15b46af6af88f19b8e5ffea08fa375d433c998b8a7639e76935c14f1f/markdown_it_py-3.0.0-py3-none-any.whl", hash = "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1", size = 87528, upload-time = "2023-06-03T06:41:11.019Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
linkify = [
|
||||
{ name = "linkify-it-py" },
|
||||
]
|
||||
plugins = [
|
||||
{ name = "mdit-py-plugins" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "markupsafe"
|
||||
version = "2.1.5"
|
||||
@@ -2514,6 +2565,18 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ad/68/316cbc54b7163fa22571dcf42c9cc46562aae0a021b974e0a8141e897200/mcp-1.12.4-py3-none-any.whl", hash = "sha256:7aa884648969fab8e78b89399d59a683202972e12e6bc9a1c88ce7eda7743789", size = 160145, upload-time = "2025-08-07T20:31:15.69Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mdit-py-plugins"
|
||||
version = "0.5.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "markdown-it-py" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/fb/86/dd6e5db36df29e76c7a7699123569a4a18c1623ce68d826ed96c62643cae/mdit_py_plugins-0.5.0-py3-none-any.whl", hash = "sha256:07a08422fc1936a5d26d146759e9155ea466e842f5ab2f7d2266dd084c8dab1f", size = 57205, upload-time = "2025-08-11T07:25:47.597Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mdurl"
|
||||
version = "0.1.2"
|
||||
@@ -2613,6 +2676,67 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "msgpack"
|
||||
version = "1.1.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f5/a2/3b68a9e769db68668b25c6108444a35f9bd163bb848c0650d516761a59c0/msgpack-1.1.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0051fffef5a37ca2cd16978ae4f0aef92f164df86823871b5162812bebecd8e2", size = 81318, upload-time = "2025-10-08T09:14:38.722Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5b/e1/2b720cc341325c00be44e1ed59e7cfeae2678329fbf5aa68f5bda57fe728/msgpack-1.1.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a605409040f2da88676e9c9e5853b3449ba8011973616189ea5ee55ddbc5bc87", size = 83786, upload-time = "2025-10-08T09:14:40.082Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/71/e5/c2241de64bfceac456b140737812a2ab310b10538a7b34a1d393b748e095/msgpack-1.1.2-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8b696e83c9f1532b4af884045ba7f3aa741a63b2bc22617293a2c6a7c645f251", size = 398240, upload-time = "2025-10-08T09:14:41.151Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b7/09/2a06956383c0fdebaef5aa9246e2356776f12ea6f2a44bd1368abf0e46c4/msgpack-1.1.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:365c0bbe981a27d8932da71af63ef86acc59ed5c01ad929e09a0b88c6294e28a", size = 406070, upload-time = "2025-10-08T09:14:42.821Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/74/2957703f0e1ef20637d6aead4fbb314330c26f39aa046b348c7edcf6ca6b/msgpack-1.1.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:41d1a5d875680166d3ac5c38573896453bbbea7092936d2e107214daf43b1d4f", size = 393403, upload-time = "2025-10-08T09:14:44.38Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a5/09/3bfc12aa90f77b37322fc33e7a8a7c29ba7c8edeadfa27664451801b9860/msgpack-1.1.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:354e81bcdebaab427c3df4281187edc765d5d76bfb3a7c125af9da7a27e8458f", size = 398947, upload-time = "2025-10-08T09:14:45.56Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4b/4f/05fcebd3b4977cb3d840f7ef6b77c51f8582086de5e642f3fefee35c86fc/msgpack-1.1.2-cp310-cp310-win32.whl", hash = "sha256:e64c8d2f5e5d5fda7b842f55dec6133260ea8f53c4257d64494c534f306bf7a9", size = 64769, upload-time = "2025-10-08T09:14:47.334Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d0/3e/b4547e3a34210956382eed1c85935fff7e0f9b98be3106b3745d7dec9c5e/msgpack-1.1.2-cp310-cp310-win_amd64.whl", hash = "sha256:db6192777d943bdaaafb6ba66d44bf65aa0e9c5616fa1d2da9bb08828c6b39aa", size = 71293, upload-time = "2025-10-08T09:14:48.665Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2c/97/560d11202bcd537abca693fd85d81cebe2107ba17301de42b01ac1677b69/msgpack-1.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2e86a607e558d22985d856948c12a3fa7b42efad264dca8a3ebbcfa2735d786c", size = 82271, upload-time = "2025-10-08T09:14:49.967Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/83/04/28a41024ccbd67467380b6fb440ae916c1e4f25e2cd4c63abe6835ac566e/msgpack-1.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:283ae72fc89da59aa004ba147e8fc2f766647b1251500182fac0350d8af299c0", size = 84914, upload-time = "2025-10-08T09:14:50.958Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/71/46/b817349db6886d79e57a966346cf0902a426375aadc1e8e7a86a75e22f19/msgpack-1.1.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:61c8aa3bd513d87c72ed0b37b53dd5c5a0f58f2ff9f26e1555d3bd7948fb7296", size = 416962, upload-time = "2025-10-08T09:14:51.997Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/da/e0/6cc2e852837cd6086fe7d8406af4294e66827a60a4cf60b86575a4a65ca8/msgpack-1.1.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:454e29e186285d2ebe65be34629fa0e8605202c60fbc7c4c650ccd41870896ef", size = 426183, upload-time = "2025-10-08T09:14:53.477Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/25/98/6a19f030b3d2ea906696cedd1eb251708e50a5891d0978b012cb6107234c/msgpack-1.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:7bc8813f88417599564fafa59fd6f95be417179f76b40325b500b3c98409757c", size = 411454, upload-time = "2025-10-08T09:14:54.648Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b7/cd/9098fcb6adb32187a70b7ecaabf6339da50553351558f37600e53a4a2a23/msgpack-1.1.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:bafca952dc13907bdfdedfc6a5f579bf4f292bdd506fadb38389afa3ac5b208e", size = 422341, upload-time = "2025-10-08T09:14:56.328Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e6/ae/270cecbcf36c1dc85ec086b33a51a4d7d08fc4f404bdbc15b582255d05ff/msgpack-1.1.2-cp311-cp311-win32.whl", hash = "sha256:602b6740e95ffc55bfb078172d279de3773d7b7db1f703b2f1323566b878b90e", size = 64747, upload-time = "2025-10-08T09:14:57.882Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2a/79/309d0e637f6f37e83c711f547308b91af02b72d2326ddd860b966080ef29/msgpack-1.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:d198d275222dc54244bf3327eb8cbe00307d220241d9cec4d306d49a44e85f68", size = 71633, upload-time = "2025-10-08T09:14:59.177Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/73/4d/7c4e2b3d9b1106cd0aa6cb56cc57c6267f59fa8bfab7d91df5adc802c847/msgpack-1.1.2-cp311-cp311-win_arm64.whl", hash = "sha256:86f8136dfa5c116365a8a651a7d7484b65b13339731dd6faebb9a0242151c406", size = 64755, upload-time = "2025-10-08T09:15:00.48Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ad/bd/8b0d01c756203fbab65d265859749860682ccd2a59594609aeec3a144efa/msgpack-1.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:70a0dff9d1f8da25179ffcf880e10cf1aad55fdb63cd59c9a49a1b82290062aa", size = 81939, upload-time = "2025-10-08T09:15:01.472Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/34/68/ba4f155f793a74c1483d4bdef136e1023f7bcba557f0db4ef3db3c665cf1/msgpack-1.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:446abdd8b94b55c800ac34b102dffd2f6aa0ce643c55dfc017ad89347db3dbdb", size = 85064, upload-time = "2025-10-08T09:15:03.764Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f2/60/a064b0345fc36c4c3d2c743c82d9100c40388d77f0b48b2f04d6041dbec1/msgpack-1.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c63eea553c69ab05b6747901b97d620bb2a690633c77f23feb0c6a947a8a7b8f", size = 417131, upload-time = "2025-10-08T09:15:05.136Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/65/92/a5100f7185a800a5d29f8d14041f61475b9de465ffcc0f3b9fba606e4505/msgpack-1.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:372839311ccf6bdaf39b00b61288e0557916c3729529b301c52c2d88842add42", size = 427556, upload-time = "2025-10-08T09:15:06.837Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f5/87/ffe21d1bf7d9991354ad93949286f643b2bb6ddbeab66373922b44c3b8cc/msgpack-1.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2929af52106ca73fcb28576218476ffbb531a036c2adbcf54a3664de124303e9", size = 404920, upload-time = "2025-10-08T09:15:08.179Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ff/41/8543ed2b8604f7c0d89ce066f42007faac1eaa7d79a81555f206a5cdb889/msgpack-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:be52a8fc79e45b0364210eef5234a7cf8d330836d0a64dfbb878efa903d84620", size = 415013, upload-time = "2025-10-08T09:15:09.83Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/41/0d/2ddfaa8b7e1cee6c490d46cb0a39742b19e2481600a7a0e96537e9c22f43/msgpack-1.1.2-cp312-cp312-win32.whl", hash = "sha256:1fff3d825d7859ac888b0fbda39a42d59193543920eda9d9bea44d958a878029", size = 65096, upload-time = "2025-10-08T09:15:11.11Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8c/ec/d431eb7941fb55a31dd6ca3404d41fbb52d99172df2e7707754488390910/msgpack-1.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:1de460f0403172cff81169a30b9a92b260cb809c4cb7e2fc79ae8d0510c78b6b", size = 72708, upload-time = "2025-10-08T09:15:12.554Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c5/31/5b1a1f70eb0e87d1678e9624908f86317787b536060641d6798e3cf70ace/msgpack-1.1.2-cp312-cp312-win_arm64.whl", hash = "sha256:be5980f3ee0e6bd44f3a9e9dea01054f175b50c3e6cdb692bc9424c0bbb8bf69", size = 64119, upload-time = "2025-10-08T09:15:13.589Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6b/31/b46518ecc604d7edf3a4f94cb3bf021fc62aa301f0cb849936968164ef23/msgpack-1.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4efd7b5979ccb539c221a4c4e16aac1a533efc97f3b759bb5a5ac9f6d10383bf", size = 81212, upload-time = "2025-10-08T09:15:14.552Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/92/dc/c385f38f2c2433333345a82926c6bfa5ecfff3ef787201614317b58dd8be/msgpack-1.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42eefe2c3e2af97ed470eec850facbe1b5ad1d6eacdbadc42ec98e7dcf68b4b7", size = 84315, upload-time = "2025-10-08T09:15:15.543Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d3/68/93180dce57f684a61a88a45ed13047558ded2be46f03acb8dec6d7c513af/msgpack-1.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fdf7d83102bf09e7ce3357de96c59b627395352a4024f6e2458501f158bf999", size = 412721, upload-time = "2025-10-08T09:15:16.567Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5d/ba/459f18c16f2b3fc1a1ca871f72f07d70c07bf768ad0a507a698b8052ac58/msgpack-1.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fac4be746328f90caa3cd4bc67e6fe36ca2bf61d5c6eb6d895b6527e3f05071e", size = 424657, upload-time = "2025-10-08T09:15:17.825Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/38/f8/4398c46863b093252fe67368b44edc6c13b17f4e6b0e4929dbf0bdb13f23/msgpack-1.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fffee09044073e69f2bad787071aeec727183e7580443dfeb8556cbf1978d162", size = 402668, upload-time = "2025-10-08T09:15:19.003Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/28/ce/698c1eff75626e4124b4d78e21cca0b4cc90043afb80a507626ea354ab52/msgpack-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5928604de9b032bc17f5099496417f113c45bc6bc21b5c6920caf34b3c428794", size = 419040, upload-time = "2025-10-08T09:15:20.183Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/67/32/f3cd1667028424fa7001d82e10ee35386eea1408b93d399b09fb0aa7875f/msgpack-1.1.2-cp313-cp313-win32.whl", hash = "sha256:a7787d353595c7c7e145e2331abf8b7ff1e6673a6b974ded96e6d4ec09f00c8c", size = 65037, upload-time = "2025-10-08T09:15:21.416Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/74/07/1ed8277f8653c40ebc65985180b007879f6a836c525b3885dcc6448ae6cb/msgpack-1.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:a465f0dceb8e13a487e54c07d04ae3ba131c7c5b95e2612596eafde1dccf64a9", size = 72631, upload-time = "2025-10-08T09:15:22.431Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e5/db/0314e4e2db56ebcf450f277904ffd84a7988b9e5da8d0d61ab2d057df2b6/msgpack-1.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:e69b39f8c0aa5ec24b57737ebee40be647035158f14ed4b40e6f150077e21a84", size = 64118, upload-time = "2025-10-08T09:15:23.402Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/22/71/201105712d0a2ff07b7873ed3c220292fb2ea5120603c00c4b634bcdafb3/msgpack-1.1.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:e23ce8d5f7aa6ea6d2a2b326b4ba46c985dbb204523759984430db7114f8aa00", size = 81127, upload-time = "2025-10-08T09:15:24.408Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1b/9f/38ff9e57a2eade7bf9dfee5eae17f39fc0e998658050279cbb14d97d36d9/msgpack-1.1.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:6c15b7d74c939ebe620dd8e559384be806204d73b4f9356320632d783d1f7939", size = 84981, upload-time = "2025-10-08T09:15:25.812Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8e/a9/3536e385167b88c2cc8f4424c49e28d49a6fc35206d4a8060f136e71f94c/msgpack-1.1.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99e2cb7b9031568a2a5c73aa077180f93dd2e95b4f8d3b8e14a73ae94a9e667e", size = 411885, upload-time = "2025-10-08T09:15:27.22Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2f/40/dc34d1a8d5f1e51fc64640b62b191684da52ca469da9cd74e84936ffa4a6/msgpack-1.1.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:180759d89a057eab503cf62eeec0aa61c4ea1200dee709f3a8e9397dbb3b6931", size = 419658, upload-time = "2025-10-08T09:15:28.4Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/3b/ef/2b92e286366500a09a67e03496ee8b8ba00562797a52f3c117aa2b29514b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:04fb995247a6e83830b62f0b07bf36540c213f6eac8e851166d8d86d83cbd014", size = 403290, upload-time = "2025-10-08T09:15:29.764Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/78/90/e0ea7990abea5764e4655b8177aa7c63cdfa89945b6e7641055800f6c16b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:8e22ab046fa7ede9e36eeb4cfad44d46450f37bb05d5ec482b02868f451c95e2", size = 415234, upload-time = "2025-10-08T09:15:31.022Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/72/4e/9390aed5db983a2310818cd7d3ec0aecad45e1f7007e0cda79c79507bb0d/msgpack-1.1.2-cp314-cp314-win32.whl", hash = "sha256:80a0ff7d4abf5fecb995fcf235d4064b9a9a8a40a3ab80999e6ac1e30b702717", size = 66391, upload-time = "2025-10-08T09:15:32.265Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6e/f1/abd09c2ae91228c5f3998dbd7f41353def9eac64253de3c8105efa2082f7/msgpack-1.1.2-cp314-cp314-win_amd64.whl", hash = "sha256:9ade919fac6a3e7260b7f64cea89df6bec59104987cbea34d34a2fa15d74310b", size = 73787, upload-time = "2025-10-08T09:15:33.219Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6a/b0/9d9f667ab48b16ad4115c1935d94023b82b3198064cb84a123e97f7466c1/msgpack-1.1.2-cp314-cp314-win_arm64.whl", hash = "sha256:59415c6076b1e30e563eb732e23b994a61c159cec44deaf584e5cc1dd662f2af", size = 66453, upload-time = "2025-10-08T09:15:34.225Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/16/67/93f80545eb1792b61a217fa7f06d5e5cb9e0055bed867f43e2b8e012e137/msgpack-1.1.2-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:897c478140877e5307760b0ea66e0932738879e7aa68144d9b78ea4c8302a84a", size = 85264, upload-time = "2025-10-08T09:15:35.61Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/87/1c/33c8a24959cf193966ef11a6f6a2995a65eb066bd681fd085afd519a57ce/msgpack-1.1.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a668204fa43e6d02f89dbe79a30b0d67238d9ec4c5bd8a940fc3a004a47b721b", size = 89076, upload-time = "2025-10-08T09:15:36.619Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/fc/6b/62e85ff7193663fbea5c0254ef32f0c77134b4059f8da89b958beb7696f3/msgpack-1.1.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5559d03930d3aa0f3aacb4c42c776af1a2ace2611871c84a75afe436695e6245", size = 435242, upload-time = "2025-10-08T09:15:37.647Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c1/47/5c74ecb4cc277cf09f64e913947871682ffa82b3b93c8dad68083112f412/msgpack-1.1.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:70c5a7a9fea7f036b716191c29047374c10721c389c21e9ffafad04df8c52c90", size = 432509, upload-time = "2025-10-08T09:15:38.794Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/24/a4/e98ccdb56dc4e98c929a3f150de1799831c0a800583cde9fa022fa90602d/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f2cb069d8b981abc72b41aea1c580ce92d57c673ec61af4c500153a626cb9e20", size = 415957, upload-time = "2025-10-08T09:15:40.238Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/da/28/6951f7fb67bc0a4e184a6b38ab71a92d9ba58080b27a77d3e2fb0be5998f/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d62ce1f483f355f61adb5433ebfd8868c5f078d1a52d042b0a998682b4fa8c27", size = 422910, upload-time = "2025-10-08T09:15:41.505Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f0/03/42106dcded51f0a0b5284d3ce30a671e7bd3f7318d122b2ead66ad289fed/msgpack-1.1.2-cp314-cp314t-win32.whl", hash = "sha256:1d1418482b1ee984625d88aa9585db570180c286d942da463533b238b98b812b", size = 75197, upload-time = "2025-10-08T09:15:42.954Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/15/86/d0071e94987f8db59d4eeb386ddc64d0bb9b10820a8d82bcd3e53eeb2da6/msgpack-1.1.2-cp314-cp314t-win_amd64.whl", hash = "sha256:5a46bf7e831d09470ad92dff02b8b1ac92175ca36b087f904a0519857c6be3ff", size = 85772, upload-time = "2025-10-08T09:15:43.954Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/81/f2/08ace4142eb281c12701fc3b93a10795e4d4dc7f753911d836675050f886/msgpack-1.1.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d99ef64f349d5ec3293688e91486c5fdb925ed03807f64d98d205d2713c60b46", size = 70868, upload-time = "2025-10-08T09:15:44.959Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "multidict"
|
||||
version = "6.1.0"
|
||||
@@ -3948,6 +4072,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e5/ae/580600f441f6fc05218bd6c9d5794f4aef072a7d9093b291f1c50a9db8bc/plotly-5.24.1-py3-none-any.whl", hash = "sha256:f67073a1e637eb0dc3e46324d9d51e2fe76e9727c892dde64ddf1e1b51f29089", size = 19054220, upload-time = "2024-09-12T15:36:24.08Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.6.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "portalocker"
|
||||
version = "2.7.0"
|
||||
@@ -3988,6 +4121,7 @@ dependencies = [
|
||||
{ name = "python-dotenv" },
|
||||
{ name = "pyyaml" },
|
||||
{ name = "rich" },
|
||||
{ name = "textual" },
|
||||
{ name = "typer" },
|
||||
]
|
||||
|
||||
@@ -4059,6 +4193,11 @@ crewai = [
|
||||
{ name = "crewai" },
|
||||
{ name = "praisonai-tools" },
|
||||
]
|
||||
dev = [
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "textual-dev" },
|
||||
]
|
||||
google = [
|
||||
{ name = "langchain-google-genai" },
|
||||
]
|
||||
@@ -4137,13 +4276,15 @@ requires-dist = [
|
||||
{ name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.1.0" },
|
||||
{ name = "praisonai-tools", marker = "extra == 'autogen-v4'", specifier = ">=0.1.0" },
|
||||
{ name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.1.0" },
|
||||
{ name = "praisonaiagents", specifier = ">=0.5.3" },
|
||||
{ name = "praisonaiagents", specifier = ">=0.6.0" },
|
||||
{ name = "pyautogen", marker = "extra == 'autogen'", specifier = "==0.2.29" },
|
||||
{ name = "pydantic", marker = "extra == 'chat'", specifier = "<=2.10.1" },
|
||||
{ name = "pydantic", marker = "extra == 'code'", specifier = "<=2.10.1" },
|
||||
{ name = "pydantic", marker = "extra == 'ui'", specifier = "<=2.10.1" },
|
||||
{ name = "pyngrok", marker = "extra == 'call'", specifier = ">=1.4.0" },
|
||||
{ name = "pyparsing", specifier = ">=3.0.0" },
|
||||
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" },
|
||||
{ name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.21.0" },
|
||||
{ name = "python-dotenv", specifier = ">=0.19.0" },
|
||||
{ name = "pyyaml", specifier = ">=6.0" },
|
||||
{ name = "rich", specifier = ">=13.7" },
|
||||
@@ -4157,6 +4298,8 @@ requires-dist = [
|
||||
{ name = "tavily-python", marker = "extra == 'chat'", specifier = "==0.5.0" },
|
||||
{ name = "tavily-python", marker = "extra == 'code'", specifier = "==0.5.0" },
|
||||
{ name = "tavily-python", marker = "extra == 'realtime'", specifier = "==0.5.0" },
|
||||
{ name = "textual", specifier = ">=0.47.0" },
|
||||
{ name = "textual-dev", marker = "extra == 'dev'", specifier = ">=1.0.0" },
|
||||
{ name = "twilio", marker = "extra == 'call'", specifier = ">=7.0.0" },
|
||||
{ name = "typer", specifier = ">=0.9.0" },
|
||||
{ name = "uvicorn", marker = "extra == 'api'", specifier = ">=0.34.0" },
|
||||
@@ -4165,7 +4308,7 @@ requires-dist = [
|
||||
{ name = "websockets", marker = "extra == 'realtime'", specifier = ">=12.0" },
|
||||
{ name = "yfinance", marker = "extra == 'realtime'", specifier = ">=0.2.44" },
|
||||
]
|
||||
provides-extras = ["ui", "gradio", "api", "agentops", "google", "openai", "anthropic", "cohere", "chat", "code", "realtime", "call", "train", "crewai", "autogen", "autogen-v4", "acp"]
|
||||
provides-extras = ["ui", "gradio", "api", "agentops", "google", "openai", "anthropic", "cohere", "chat", "code", "realtime", "call", "train", "crewai", "autogen", "autogen-v4", "acp", "lite", "dev"]
|
||||
|
||||
[[package]]
|
||||
name = "praisonai-tools"
|
||||
@@ -4181,7 +4324,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.5.3"
|
||||
version = "0.6.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -4190,9 +4333,9 @@ dependencies = [
|
||||
{ name = "pydantic" },
|
||||
{ name = "rich" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/75/33/6091503673fb3c143293e9639bafe3acb6e7406e0ace1f893c34cbe22912/praisonaiagents-0.5.3.tar.gz", hash = "sha256:4163b00c28ef5e5d1a75411eb4b72fc7749790e52cdaa4bfc98c9a747fe42140", size = 780945, upload-time = "2025-12-31T16:36:45.335Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/4b/3c/a25740d7e3da4ebea0d83319d412650a016d91b6eb2f2653df1b28a3882f/praisonaiagents-0.6.0.tar.gz", hash = "sha256:f83528f6f1c0adea17cea25a1ea93f660e46ba43002e354717501e82112d9a96", size = 780850, upload-time = "2025-12-31T18:50:35.35Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/32/a2/459f111873f1751a34bc780e1b3a2cd264cf580f35acff868fa35aa9d917/praisonaiagents-0.5.3-py3-none-any.whl", hash = "sha256:35bd6c26dfff6c9e77575f72a2f077cdbcc0544d0d45b0ed29fb00ca96c3719c", size = 724879, upload-time = "2025-12-31T16:36:43.383Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8e/18/f4f8360024ef26806831dde4ba3dad23950a4d94dde86085dde510ef1833/praisonaiagents-0.6.0-py3-none-any.whl", hash = "sha256:416c4550b6460b0866c1d59b284aacb86f23d849b38e75542f7056e530e84d80", size = 724880, upload-time = "2025-12-31T18:50:33.21Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4532,11 +4675,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "pygments"
|
||||
version = "2.18.0"
|
||||
version = "2.19.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/8e/62/8336eff65bcbc8e4cb5d05b55faf041285951b6e80f33e2bff2024788f31/pygments-2.18.0.tar.gz", hash = "sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199", size = 4891905, upload-time = "2024-05-04T13:42:02.013Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631, upload-time = "2025-06-21T13:39:12.283Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f7/3f/01c8b82017c199075f8f788d0d906b9ffbbc5a47dc9918a945e13d5a2bda/pygments-2.18.0-py3-none-any.whl", hash = "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a", size = 1205513, upload-time = "2024-05-04T13:41:57.345Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4632,6 +4775,38 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/5a/dc/491b7661614ab97483abf2056be1deee4dc2490ecbf7bff9ab5cdbac86e1/pyreadline3-3.5.4-py3-none-any.whl", hash = "sha256:eaf8e6cc3c49bcccf145fc6067ba8643d1df34d604a1ec0eccbf7a18e6d3fae6", size = 83178, upload-time = "2024-09-19T02:40:08.598Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest"
|
||||
version = "9.0.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||
{ name = "exceptiongroup", marker = "python_full_version < '3.11'" },
|
||||
{ name = "iniconfig" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pluggy" },
|
||||
{ name = "pygments" },
|
||||
{ name = "tomli", marker = "python_full_version < '3.11'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-asyncio"
|
||||
version = "1.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "backports-asyncio-runner", marker = "python_full_version < '3.11'" },
|
||||
{ name = "pytest" },
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.13'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/90/2c/8af215c0f776415f3590cac4f9086ccefd6fd463befeae41cd4d3f193e5a/pytest_asyncio-1.3.0.tar.gz", hash = "sha256:d7f52f36d231b80ee124cd216ffb19369aa168fc10095013c6b014a34d3ee9e5", size = 50087, upload-time = "2025-11-10T16:07:47.256Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "python-dateutil"
|
||||
version = "2.9.0.post0"
|
||||
@@ -5525,6 +5700,55 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d9/5f/8c716e47b3a50cbd7c146f45881e11d9414def768b7cd9c5e6650ec2a80a/termcolor-2.4.0-py3-none-any.whl", hash = "sha256:9297c0df9c99445c2412e832e882a7884038a25617c60cea2ad69488d4040d63", size = 7719, upload-time = "2023-12-01T11:04:50.019Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "textual"
|
||||
version = "6.2.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "markdown-it-py", extra = ["linkify", "plugins"] },
|
||||
{ name = "platformdirs" },
|
||||
{ name = "pygments" },
|
||||
{ name = "rich" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a2/30/38b615f7d4b16f6fdd73e4dcd8913e2d880bbb655e68a076e3d91181a7ee/textual-6.2.1.tar.gz", hash = "sha256:4699d8dfae43503b9c417bd2a6fb0da1c89e323fe91c4baa012f9298acaa83e1", size = 1570645, upload-time = "2025-10-01T16:11:24.467Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c5/93/02c7adec57a594af28388d85da9972703a4af94ae1399542555cd9581952/textual-6.2.1-py3-none-any.whl", hash = "sha256:3c7190633cd4d8bfe6049ae66808b98da91ded2edb85cef54e82bf77b03d2a54", size = 710702, upload-time = "2025-10-01T16:11:22.161Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "textual-dev"
|
||||
version = "1.8.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
{ name = "click" },
|
||||
{ name = "msgpack" },
|
||||
{ name = "textual" },
|
||||
{ name = "textual-serve" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/38/fd/fd5ad9527b536c306a5a860a33a45e13983a59804b48b91ea52b42ba030a/textual_dev-1.8.0.tar.gz", hash = "sha256:7e56867b0341405a95e938cac0647e6d2763d38d0df08710469ad6b6a8db76df", size = 25026, upload-time = "2025-10-11T09:47:01.824Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/45/c6/1dc08ceee6d0bdf219c500cb2fbd605b681b63507c4ff27a6477f12f8245/textual_dev-1.8.0-py3-none-any.whl", hash = "sha256:227b6d24a485fbbc77e302aa21f4fdf3083beb57eb45cd95bae082c81cbeddeb", size = 27541, upload-time = "2025-10-11T09:47:00.531Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "textual-serve"
|
||||
version = "1.1.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
{ name = "aiohttp-jinja2" },
|
||||
{ name = "jinja2" },
|
||||
{ name = "rich" },
|
||||
{ name = "textual" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c2/7e/62fecc552853ec6a178cb1faa2d6f73b34d5512924770e7b08b58ff14148/textual_serve-1.1.3.tar.gz", hash = "sha256:f8f636ae2f5fd651b79d965473c3e9383d3521cdf896f9bc289709185da3f683", size = 448340, upload-time = "2025-11-01T16:22:36.723Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/b5/fe/108e7773349d500cf363328c3d0b7123e03feda51e310a3a5b136ac8ca71/textual_serve-1.1.3-py3-none-any.whl", hash = "sha256:207a472bc6604e725b1adab4ab8bf12f4c4dc25b04eea31e4d04731d8bf30f18", size = 447339, upload-time = "2025-11-01T16:22:35.209Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tf-playwright-stealth"
|
||||
version = "1.1.2"
|
||||
@@ -5951,6 +6175,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a6/ab/7e5f53c3b9d14972843a647d8d7a853969a58aecc7559cb3267302c94774/tzdata-2024.2-py2.py3-none-any.whl", hash = "sha256:a48093786cdcde33cad18c2555e8532f34422074448fbc874186f0abd79565cd", size = 346586, upload-time = "2024-09-23T18:56:45.478Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "uc-micro-py"
|
||||
version = "1.0.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/91/7a/146a99696aee0609e3712f2b44c6274566bc368dfe8375191278045186b8/uc-micro-py-1.0.3.tar.gz", hash = "sha256:d321b92cff673ec58027c04015fcaa8bb1e005478643ff4a500882eaab88c48a", size = 6043, upload-time = "2024-02-09T16:52:01.654Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/37/87/1f677586e8ac487e29672e4b17455758fce261de06a0d086167bb760361a/uc_micro_py-1.0.3-py3-none-any.whl", hash = "sha256:db1dffff340817673d7b466ec86114a9dc0e9d4d9b5ba229d9d60e5c12600cd5", size = 6229, upload-time = "2024-02-09T16:52:00.371Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "urllib3"
|
||||
version = "2.3.0"
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for PraisonAI Queue System."""
|
||||
@@ -0,0 +1,228 @@
|
||||
"""Tests for Queue System Data Models."""
|
||||
|
||||
import pytest
|
||||
import time
|
||||
from praisonai.cli.features.queue.models import (
|
||||
QueuedRun,
|
||||
RunState,
|
||||
RunPriority,
|
||||
QueueConfig,
|
||||
QueueStats,
|
||||
StreamChunk,
|
||||
QueueEvent,
|
||||
)
|
||||
|
||||
|
||||
class TestRunState:
|
||||
"""Tests for RunState enum."""
|
||||
|
||||
def test_terminal_states(self):
|
||||
"""Test terminal state detection."""
|
||||
assert RunState.SUCCEEDED.is_terminal()
|
||||
assert RunState.FAILED.is_terminal()
|
||||
assert RunState.CANCELLED.is_terminal()
|
||||
assert not RunState.QUEUED.is_terminal()
|
||||
assert not RunState.RUNNING.is_terminal()
|
||||
assert not RunState.PAUSED.is_terminal()
|
||||
|
||||
def test_active_states(self):
|
||||
"""Test active state detection."""
|
||||
assert RunState.QUEUED.is_active()
|
||||
assert RunState.RUNNING.is_active()
|
||||
assert RunState.PAUSED.is_active()
|
||||
assert not RunState.SUCCEEDED.is_active()
|
||||
assert not RunState.FAILED.is_active()
|
||||
|
||||
|
||||
class TestRunPriority:
|
||||
"""Tests for RunPriority enum."""
|
||||
|
||||
def test_priority_ordering(self):
|
||||
"""Test priority ordering."""
|
||||
assert RunPriority.URGENT > RunPriority.HIGH
|
||||
assert RunPriority.HIGH > RunPriority.NORMAL
|
||||
assert RunPriority.NORMAL > RunPriority.LOW
|
||||
|
||||
def test_from_string(self):
|
||||
"""Test parsing from string."""
|
||||
assert RunPriority.from_string("low") == RunPriority.LOW
|
||||
assert RunPriority.from_string("NORMAL") == RunPriority.NORMAL
|
||||
assert RunPriority.from_string("High") == RunPriority.HIGH
|
||||
assert RunPriority.from_string("urgent") == RunPriority.URGENT
|
||||
assert RunPriority.from_string("invalid") == RunPriority.NORMAL
|
||||
|
||||
|
||||
class TestQueuedRun:
|
||||
"""Tests for QueuedRun dataclass."""
|
||||
|
||||
def test_default_values(self):
|
||||
"""Test default values."""
|
||||
run = QueuedRun()
|
||||
assert run.run_id # Auto-generated
|
||||
assert run.agent_name == ""
|
||||
assert run.state == RunState.QUEUED
|
||||
assert run.priority == RunPriority.NORMAL
|
||||
assert run.retry_count == 0
|
||||
assert run.max_retries == 3
|
||||
|
||||
def test_custom_values(self):
|
||||
"""Test custom values."""
|
||||
run = QueuedRun(
|
||||
run_id="test123",
|
||||
agent_name="TestAgent",
|
||||
input_content="Hello",
|
||||
state=RunState.RUNNING,
|
||||
priority=RunPriority.HIGH,
|
||||
)
|
||||
assert run.run_id == "test123"
|
||||
assert run.agent_name == "TestAgent"
|
||||
assert run.input_content == "Hello"
|
||||
assert run.state == RunState.RUNNING
|
||||
assert run.priority == RunPriority.HIGH
|
||||
|
||||
def test_duration_calculation(self):
|
||||
"""Test duration calculation."""
|
||||
run = QueuedRun()
|
||||
assert run.duration_seconds is None
|
||||
|
||||
run.started_at = time.time() - 5
|
||||
assert run.duration_seconds is not None
|
||||
assert run.duration_seconds >= 5
|
||||
|
||||
run.ended_at = run.started_at + 10
|
||||
assert abs(run.duration_seconds - 10) < 0.1
|
||||
|
||||
def test_wait_time_calculation(self):
|
||||
"""Test wait time calculation."""
|
||||
run = QueuedRun()
|
||||
run.created_at = time.time() - 3
|
||||
assert run.wait_seconds >= 3
|
||||
|
||||
run.started_at = run.created_at + 2
|
||||
assert abs(run.wait_seconds - 2) < 0.1
|
||||
|
||||
def test_can_retry(self):
|
||||
"""Test retry eligibility."""
|
||||
run = QueuedRun(state=RunState.FAILED, retry_count=0, max_retries=3)
|
||||
assert run.can_retry()
|
||||
|
||||
run.retry_count = 3
|
||||
assert not run.can_retry()
|
||||
|
||||
run.retry_count = 0
|
||||
run.state = RunState.SUCCEEDED
|
||||
assert not run.can_retry()
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test serialization to dict."""
|
||||
run = QueuedRun(
|
||||
run_id="test123",
|
||||
agent_name="TestAgent",
|
||||
input_content="Hello",
|
||||
)
|
||||
d = run.to_dict()
|
||||
|
||||
assert d["run_id"] == "test123"
|
||||
assert d["agent_name"] == "TestAgent"
|
||||
assert d["input_content"] == "Hello"
|
||||
assert d["state"] == "queued"
|
||||
assert d["priority"] == 1
|
||||
|
||||
def test_from_dict(self):
|
||||
"""Test deserialization from dict."""
|
||||
d = {
|
||||
"run_id": "test123",
|
||||
"agent_name": "TestAgent",
|
||||
"input_content": "Hello",
|
||||
"state": "running",
|
||||
"priority": 2,
|
||||
}
|
||||
run = QueuedRun.from_dict(d)
|
||||
|
||||
assert run.run_id == "test123"
|
||||
assert run.agent_name == "TestAgent"
|
||||
assert run.state == RunState.RUNNING
|
||||
assert run.priority == RunPriority.HIGH
|
||||
|
||||
def test_post_init_conversion(self):
|
||||
"""Test automatic type conversion in __post_init__."""
|
||||
run = QueuedRun(state="running", priority=2)
|
||||
assert run.state == RunState.RUNNING
|
||||
assert run.priority == RunPriority.HIGH
|
||||
|
||||
|
||||
class TestQueueConfig:
|
||||
"""Tests for QueueConfig dataclass."""
|
||||
|
||||
def test_default_values(self):
|
||||
"""Test default values."""
|
||||
config = QueueConfig()
|
||||
assert config.max_concurrent_global == 4
|
||||
assert config.max_concurrent_per_agent == 2
|
||||
assert config.max_queue_size == 100
|
||||
assert config.enable_persistence
|
||||
assert config.db_path == ".praison/queue.db"
|
||||
|
||||
def test_custom_values(self):
|
||||
"""Test custom values."""
|
||||
config = QueueConfig(
|
||||
max_concurrent_global=8,
|
||||
max_queue_size=200,
|
||||
enable_persistence=False,
|
||||
)
|
||||
assert config.max_concurrent_global == 8
|
||||
assert config.max_queue_size == 200
|
||||
assert not config.enable_persistence
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test serialization."""
|
||||
config = QueueConfig()
|
||||
d = config.to_dict()
|
||||
assert "max_concurrent_global" in d
|
||||
assert "db_path" in d
|
||||
|
||||
def test_from_dict(self):
|
||||
"""Test deserialization."""
|
||||
d = {"max_concurrent_global": 10, "max_queue_size": 50}
|
||||
config = QueueConfig.from_dict(d)
|
||||
assert config.max_concurrent_global == 10
|
||||
assert config.max_queue_size == 50
|
||||
|
||||
|
||||
class TestQueueStats:
|
||||
"""Tests for QueueStats dataclass."""
|
||||
|
||||
def test_active_count(self):
|
||||
"""Test active count calculation."""
|
||||
stats = QueueStats(queued_count=5, running_count=3)
|
||||
assert stats.active_count == 8
|
||||
|
||||
|
||||
class TestStreamChunk:
|
||||
"""Tests for StreamChunk dataclass."""
|
||||
|
||||
def test_creation(self):
|
||||
"""Test chunk creation."""
|
||||
chunk = StreamChunk(
|
||||
run_id="test123",
|
||||
content="Hello",
|
||||
chunk_index=0,
|
||||
)
|
||||
assert chunk.run_id == "test123"
|
||||
assert chunk.content == "Hello"
|
||||
assert chunk.chunk_index == 0
|
||||
assert not chunk.is_final
|
||||
|
||||
|
||||
class TestQueueEvent:
|
||||
"""Tests for QueueEvent dataclass."""
|
||||
|
||||
def test_creation(self):
|
||||
"""Test event creation."""
|
||||
event = QueueEvent(
|
||||
event_type="run_started",
|
||||
run_id="test123",
|
||||
)
|
||||
assert event.event_type == "run_started"
|
||||
assert event.run_id == "test123"
|
||||
assert event.timestamp > 0
|
||||
@@ -0,0 +1,251 @@
|
||||
"""Tests for Queue Persistence Layer."""
|
||||
|
||||
import pytest
|
||||
import tempfile
|
||||
import os
|
||||
from praisonai.cli.features.queue.models import QueuedRun, RunState, RunPriority
|
||||
from praisonai.cli.features.queue.persistence import QueuePersistence
|
||||
|
||||
|
||||
class TestQueuePersistence:
|
||||
"""Tests for QueuePersistence."""
|
||||
|
||||
@pytest.fixture
|
||||
def db_path(self):
|
||||
"""Create temporary database path."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
yield os.path.join(tmpdir, "test_queue.db")
|
||||
|
||||
@pytest.fixture
|
||||
def persistence(self, db_path):
|
||||
"""Create test persistence instance."""
|
||||
p = QueuePersistence(db_path)
|
||||
p.initialize()
|
||||
yield p
|
||||
p.close()
|
||||
|
||||
def test_initialize_creates_tables(self, persistence):
|
||||
"""Test initialization creates tables."""
|
||||
# Should not raise
|
||||
persistence.initialize()
|
||||
|
||||
def test_save_and_load_run(self, persistence):
|
||||
"""Test saving and loading a run."""
|
||||
run = QueuedRun(
|
||||
run_id="test123",
|
||||
agent_name="TestAgent",
|
||||
input_content="Hello",
|
||||
state=RunState.QUEUED,
|
||||
priority=RunPriority.HIGH,
|
||||
)
|
||||
|
||||
persistence.save_run(run)
|
||||
loaded = persistence.load_run("test123")
|
||||
|
||||
assert loaded is not None
|
||||
assert loaded.run_id == "test123"
|
||||
assert loaded.agent_name == "TestAgent"
|
||||
assert loaded.input_content == "Hello"
|
||||
assert loaded.state == RunState.QUEUED
|
||||
assert loaded.priority == RunPriority.HIGH
|
||||
|
||||
def test_load_nonexistent_run(self, persistence):
|
||||
"""Test loading a nonexistent run."""
|
||||
loaded = persistence.load_run("nonexistent")
|
||||
assert loaded is None
|
||||
|
||||
def test_update_run(self, persistence):
|
||||
"""Test updating a run."""
|
||||
run = QueuedRun(run_id="test123", agent_name="TestAgent")
|
||||
persistence.save_run(run)
|
||||
|
||||
run.state = RunState.RUNNING
|
||||
run.output_content = "Result"
|
||||
persistence.save_run(run)
|
||||
|
||||
loaded = persistence.load_run("test123")
|
||||
assert loaded.state == RunState.RUNNING
|
||||
assert loaded.output_content == "Result"
|
||||
|
||||
def test_list_runs(self, persistence):
|
||||
"""Test listing runs."""
|
||||
for i in range(5):
|
||||
run = QueuedRun(
|
||||
run_id=f"run{i}",
|
||||
agent_name="TestAgent",
|
||||
priority=RunPriority.NORMAL if i % 2 == 0 else RunPriority.HIGH,
|
||||
)
|
||||
persistence.save_run(run)
|
||||
|
||||
runs = persistence.list_runs(limit=10)
|
||||
assert len(runs) == 5
|
||||
|
||||
def test_list_runs_filter_by_state(self, persistence):
|
||||
"""Test filtering runs by state."""
|
||||
run1 = QueuedRun(run_id="run1", agent_name="Agent", state=RunState.QUEUED)
|
||||
run2 = QueuedRun(run_id="run2", agent_name="Agent", state=RunState.RUNNING)
|
||||
run3 = QueuedRun(run_id="run3", agent_name="Agent", state=RunState.QUEUED)
|
||||
|
||||
persistence.save_run(run1)
|
||||
persistence.save_run(run2)
|
||||
persistence.save_run(run3)
|
||||
|
||||
queued = persistence.list_runs(state=RunState.QUEUED)
|
||||
assert len(queued) == 2
|
||||
|
||||
running = persistence.list_runs(state=RunState.RUNNING)
|
||||
assert len(running) == 1
|
||||
|
||||
def test_list_runs_filter_by_session(self, persistence):
|
||||
"""Test filtering runs by session."""
|
||||
run1 = QueuedRun(run_id="run1", agent_name="Agent", session_id="session1")
|
||||
run2 = QueuedRun(run_id="run2", agent_name="Agent", session_id="session2")
|
||||
run3 = QueuedRun(run_id="run3", agent_name="Agent", session_id="session1")
|
||||
|
||||
persistence.save_run(run1)
|
||||
persistence.save_run(run2)
|
||||
persistence.save_run(run3)
|
||||
|
||||
session1_runs = persistence.list_runs(session_id="session1")
|
||||
assert len(session1_runs) == 2
|
||||
|
||||
def test_delete_run(self, persistence):
|
||||
"""Test deleting a run."""
|
||||
run = QueuedRun(run_id="test123", agent_name="TestAgent")
|
||||
persistence.save_run(run)
|
||||
|
||||
result = persistence.delete_run("test123")
|
||||
assert result is True
|
||||
|
||||
loaded = persistence.load_run("test123")
|
||||
assert loaded is None
|
||||
|
||||
def test_delete_nonexistent_run(self, persistence):
|
||||
"""Test deleting a nonexistent run."""
|
||||
result = persistence.delete_run("nonexistent")
|
||||
assert result is False
|
||||
|
||||
def test_update_run_state(self, persistence):
|
||||
"""Test updating run state."""
|
||||
run = QueuedRun(run_id="test123", agent_name="TestAgent")
|
||||
persistence.save_run(run)
|
||||
|
||||
result = persistence.update_run_state(
|
||||
"test123",
|
||||
RunState.FAILED,
|
||||
error="Test error",
|
||||
)
|
||||
|
||||
assert result is True
|
||||
|
||||
loaded = persistence.load_run("test123")
|
||||
assert loaded.state == RunState.FAILED
|
||||
assert loaded.error == "Test error"
|
||||
assert loaded.ended_at is not None
|
||||
|
||||
def test_load_pending_runs(self, persistence):
|
||||
"""Test loading pending runs for crash recovery."""
|
||||
run1 = QueuedRun(run_id="run1", agent_name="Agent", state=RunState.QUEUED)
|
||||
run2 = QueuedRun(run_id="run2", agent_name="Agent", state=RunState.RUNNING)
|
||||
run3 = QueuedRun(run_id="run3", agent_name="Agent", state=RunState.SUCCEEDED)
|
||||
|
||||
persistence.save_run(run1)
|
||||
persistence.save_run(run2)
|
||||
persistence.save_run(run3)
|
||||
|
||||
pending = persistence.load_pending_runs()
|
||||
assert len(pending) == 2
|
||||
|
||||
run_ids = [r.run_id for r in pending]
|
||||
assert "run1" in run_ids
|
||||
assert "run2" in run_ids
|
||||
assert "run3" not in run_ids
|
||||
|
||||
def test_mark_interrupted_as_failed(self, persistence):
|
||||
"""Test marking interrupted runs as failed."""
|
||||
run1 = QueuedRun(run_id="run1", agent_name="Agent", state=RunState.RUNNING)
|
||||
run2 = QueuedRun(run_id="run2", agent_name="Agent", state=RunState.RUNNING)
|
||||
run3 = QueuedRun(run_id="run3", agent_name="Agent", state=RunState.QUEUED)
|
||||
|
||||
persistence.save_run(run1)
|
||||
persistence.save_run(run2)
|
||||
persistence.save_run(run3)
|
||||
|
||||
count = persistence.mark_interrupted_as_failed()
|
||||
assert count == 2
|
||||
|
||||
loaded1 = persistence.load_run("run1")
|
||||
assert loaded1.state == RunState.FAILED
|
||||
assert "Interrupted" in loaded1.error
|
||||
|
||||
def test_get_stats(self, persistence):
|
||||
"""Test getting queue statistics."""
|
||||
runs = [
|
||||
QueuedRun(run_id="run1", agent_name="Agent", state=RunState.QUEUED),
|
||||
QueuedRun(run_id="run2", agent_name="Agent", state=RunState.RUNNING),
|
||||
QueuedRun(run_id="run3", agent_name="Agent", state=RunState.SUCCEEDED),
|
||||
QueuedRun(run_id="run4", agent_name="Agent", state=RunState.FAILED),
|
||||
]
|
||||
|
||||
for run in runs:
|
||||
persistence.save_run(run)
|
||||
|
||||
stats = persistence.get_stats()
|
||||
|
||||
assert stats.queued_count == 1
|
||||
assert stats.running_count == 1
|
||||
assert stats.succeeded_count == 1
|
||||
assert stats.failed_count == 1
|
||||
assert stats.total_runs == 4
|
||||
|
||||
def test_save_and_load_session(self, persistence):
|
||||
"""Test saving and loading session."""
|
||||
persistence.save_session(
|
||||
session_id="session123",
|
||||
user_id="user456",
|
||||
state={"key": "value"},
|
||||
config={"model": "gpt-4"},
|
||||
)
|
||||
|
||||
loaded = persistence.load_session("session123")
|
||||
|
||||
assert loaded is not None
|
||||
assert loaded["session_id"] == "session123"
|
||||
assert loaded["user_id"] == "user456"
|
||||
assert loaded["state"] == {"key": "value"}
|
||||
assert loaded["config"] == {"model": "gpt-4"}
|
||||
|
||||
def test_list_sessions(self, persistence):
|
||||
"""Test listing sessions."""
|
||||
for i in range(3):
|
||||
persistence.save_session(f"session{i}")
|
||||
|
||||
sessions = persistence.list_sessions()
|
||||
assert len(sessions) == 3
|
||||
|
||||
def test_cleanup_old_runs(self, persistence):
|
||||
"""Test cleaning up old runs."""
|
||||
import time
|
||||
|
||||
# Create old run
|
||||
old_run = QueuedRun(
|
||||
run_id="old_run",
|
||||
agent_name="Agent",
|
||||
state=RunState.SUCCEEDED,
|
||||
)
|
||||
old_run.created_at = time.time() - (40 * 24 * 60 * 60) # 40 days ago
|
||||
persistence.save_run(old_run)
|
||||
|
||||
# Create recent run
|
||||
recent_run = QueuedRun(
|
||||
run_id="recent_run",
|
||||
agent_name="Agent",
|
||||
state=RunState.SUCCEEDED,
|
||||
)
|
||||
persistence.save_run(recent_run)
|
||||
|
||||
count = persistence.cleanup_old_runs(days=30)
|
||||
assert count == 1
|
||||
|
||||
assert persistence.load_run("old_run") is None
|
||||
assert persistence.load_run("recent_run") is not None
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Tests for Queue Scheduler."""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from praisonai.cli.features.queue.models import QueuedRun, RunState, RunPriority, QueueConfig
|
||||
from praisonai.cli.features.queue.scheduler import QueueScheduler, QueueFullError
|
||||
|
||||
|
||||
class TestQueueScheduler:
|
||||
"""Tests for QueueScheduler."""
|
||||
|
||||
@pytest.fixture
|
||||
def config(self):
|
||||
"""Create test config."""
|
||||
return QueueConfig(
|
||||
max_concurrent_global=2,
|
||||
max_concurrent_per_agent=1,
|
||||
max_queue_size=5,
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def scheduler(self, config):
|
||||
"""Create test scheduler."""
|
||||
return QueueScheduler(config)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_run(self, scheduler):
|
||||
"""Test submitting a run."""
|
||||
run = QueuedRun(
|
||||
agent_name="TestAgent",
|
||||
input_content="Hello",
|
||||
)
|
||||
|
||||
run_id = await scheduler.submit(run)
|
||||
|
||||
assert run_id == run.run_id
|
||||
assert scheduler.queued_count == 1
|
||||
assert run.state == RunState.QUEUED
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_duplicate_rejected(self, scheduler):
|
||||
"""Test duplicate run_id is rejected."""
|
||||
run = QueuedRun(run_id="test123", agent_name="TestAgent")
|
||||
await scheduler.submit(run)
|
||||
|
||||
run2 = QueuedRun(run_id="test123", agent_name="TestAgent")
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
await scheduler.submit(run2)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_full_error(self, scheduler):
|
||||
"""Test queue full error."""
|
||||
for i in range(5):
|
||||
run = QueuedRun(agent_name="TestAgent", input_content=f"Test {i}")
|
||||
await scheduler.submit(run)
|
||||
|
||||
run = QueuedRun(agent_name="TestAgent", input_content="Overflow")
|
||||
with pytest.raises(QueueFullError):
|
||||
await scheduler.submit(run)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_next_returns_highest_priority(self, scheduler):
|
||||
"""Test next() returns highest priority run."""
|
||||
low = QueuedRun(agent_name="Agent1", priority=RunPriority.LOW)
|
||||
normal = QueuedRun(agent_name="Agent2", priority=RunPriority.NORMAL)
|
||||
high = QueuedRun(agent_name="Agent3", priority=RunPriority.HIGH)
|
||||
|
||||
await scheduler.submit(low)
|
||||
await scheduler.submit(normal)
|
||||
await scheduler.submit(high)
|
||||
|
||||
next_run = await scheduler.next()
|
||||
assert next_run.priority == RunPriority.HIGH
|
||||
assert next_run.state == RunState.RUNNING
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fifo_within_priority(self, scheduler):
|
||||
"""Test FIFO ordering within same priority."""
|
||||
run1 = QueuedRun(run_id="first", agent_name="Agent1")
|
||||
run2 = QueuedRun(run_id="second", agent_name="Agent2")
|
||||
|
||||
await scheduler.submit(run1)
|
||||
await scheduler.submit(run2)
|
||||
|
||||
next_run = await scheduler.next()
|
||||
assert next_run.run_id == "first"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrency_limit_global(self, scheduler):
|
||||
"""Test global concurrency limit."""
|
||||
run1 = QueuedRun(agent_name="Agent1")
|
||||
run2 = QueuedRun(agent_name="Agent2")
|
||||
run3 = QueuedRun(agent_name="Agent3")
|
||||
|
||||
await scheduler.submit(run1)
|
||||
await scheduler.submit(run2)
|
||||
await scheduler.submit(run3)
|
||||
|
||||
# Start first two
|
||||
await scheduler.next()
|
||||
await scheduler.next()
|
||||
|
||||
# Third should not start (limit is 2)
|
||||
next_run = await scheduler.next()
|
||||
assert next_run is None
|
||||
assert scheduler.running_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrency_limit_per_agent(self, scheduler):
|
||||
"""Test per-agent concurrency limit."""
|
||||
run1 = QueuedRun(agent_name="SameAgent")
|
||||
run2 = QueuedRun(agent_name="SameAgent")
|
||||
run3 = QueuedRun(agent_name="DifferentAgent")
|
||||
|
||||
await scheduler.submit(run1)
|
||||
await scheduler.submit(run2)
|
||||
await scheduler.submit(run3)
|
||||
|
||||
# Start first
|
||||
await scheduler.next()
|
||||
|
||||
# Second same agent should not start
|
||||
next_run = await scheduler.next()
|
||||
assert next_run.agent_name == "DifferentAgent"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_run(self, scheduler):
|
||||
"""Test completing a run."""
|
||||
run = QueuedRun(agent_name="TestAgent")
|
||||
await scheduler.submit(run)
|
||||
await scheduler.next()
|
||||
|
||||
completed = await scheduler.complete(run.run_id, output="Result")
|
||||
|
||||
assert completed.state == RunState.SUCCEEDED
|
||||
assert completed.output_content == "Result"
|
||||
assert completed.ended_at is not None
|
||||
assert scheduler.running_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_run(self, scheduler):
|
||||
"""Test failing a run."""
|
||||
run = QueuedRun(agent_name="TestAgent")
|
||||
await scheduler.submit(run)
|
||||
await scheduler.next()
|
||||
|
||||
failed = await scheduler.fail(run.run_id, error="Test error")
|
||||
|
||||
assert failed.state == RunState.FAILED
|
||||
assert failed.error == "Test error"
|
||||
assert failed.ended_at is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_queued_run(self, scheduler):
|
||||
"""Test cancelling a queued run."""
|
||||
run = QueuedRun(agent_name="TestAgent")
|
||||
await scheduler.submit(run)
|
||||
|
||||
result = await scheduler.cancel(run.run_id)
|
||||
|
||||
assert result is True
|
||||
assert scheduler.queued_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_running_run(self, scheduler):
|
||||
"""Test cancelling a running run."""
|
||||
run = QueuedRun(agent_name="TestAgent")
|
||||
await scheduler.submit(run)
|
||||
await scheduler.next()
|
||||
|
||||
result = await scheduler.cancel(run.run_id)
|
||||
|
||||
assert result is True
|
||||
assert scheduler.running_count == 0
|
||||
assert scheduler.is_cancelled(run.run_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_failed_run(self, scheduler):
|
||||
"""Test retrying a failed run."""
|
||||
run = QueuedRun(agent_name="TestAgent", input_content="Hello")
|
||||
await scheduler.submit(run)
|
||||
await scheduler.next()
|
||||
await scheduler.fail(run.run_id, error="Test error")
|
||||
|
||||
new_id = await scheduler.retry(run.run_id)
|
||||
|
||||
assert new_id is not None
|
||||
assert new_id != run.run_id
|
||||
|
||||
new_run = scheduler.get_run(new_id)
|
||||
assert new_run.retry_count == 1
|
||||
assert new_run.parent_run_id == run.run_id
|
||||
assert new_run.input_content == "Hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_max_retries_exceeded(self, scheduler):
|
||||
"""Test retry fails when max retries exceeded."""
|
||||
run = QueuedRun(
|
||||
agent_name="TestAgent",
|
||||
retry_count=3,
|
||||
max_retries=3,
|
||||
state=RunState.FAILED,
|
||||
)
|
||||
await scheduler.submit(run, check_duplicate=False)
|
||||
|
||||
new_id = await scheduler.retry(run.run_id)
|
||||
assert new_id is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_queued(self, scheduler):
|
||||
"""Test getting queued runs."""
|
||||
run1 = QueuedRun(agent_name="Agent1", priority=RunPriority.LOW)
|
||||
run2 = QueuedRun(agent_name="Agent2", priority=RunPriority.HIGH)
|
||||
|
||||
await scheduler.submit(run1)
|
||||
await scheduler.submit(run2)
|
||||
|
||||
queued = scheduler.get_queued()
|
||||
|
||||
assert len(queued) == 2
|
||||
# Should be in priority order
|
||||
assert queued[0].priority == RunPriority.HIGH
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_queue(self, scheduler):
|
||||
"""Test clearing the queue."""
|
||||
for i in range(3):
|
||||
run = QueuedRun(agent_name=f"Agent{i}")
|
||||
await scheduler.submit(run)
|
||||
|
||||
count = await scheduler.clear_queue()
|
||||
|
||||
assert count == 3
|
||||
assert scheduler.queued_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_runs(self, scheduler):
|
||||
"""Test loading runs from persistence."""
|
||||
runs = [
|
||||
QueuedRun(run_id="run1", agent_name="Agent1", state=RunState.QUEUED),
|
||||
QueuedRun(run_id="run2", agent_name="Agent2", state=RunState.RUNNING),
|
||||
]
|
||||
|
||||
scheduler.load_runs(runs)
|
||||
|
||||
# Running run should be re-queued
|
||||
assert scheduler.queued_count == 2
|
||||
assert scheduler.running_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_callbacks(self, scheduler):
|
||||
"""Test event callbacks are called."""
|
||||
events = []
|
||||
|
||||
def callback(event):
|
||||
events.append(event)
|
||||
|
||||
scheduler.add_event_callback(callback)
|
||||
|
||||
run = QueuedRun(agent_name="TestAgent")
|
||||
await scheduler.submit(run)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].event_type == "run_submitted"
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for PraisonAI TUI."""
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Tests for TUI Events."""
|
||||
|
||||
import time
|
||||
from praisonai.cli.features.tui.events import TUIEvent, TUIEventType
|
||||
|
||||
|
||||
class TestTUIEventType:
|
||||
"""Tests for TUIEventType enum."""
|
||||
|
||||
def test_event_types_exist(self):
|
||||
"""Test all expected event types exist."""
|
||||
assert TUIEventType.MESSAGE_SUBMITTED
|
||||
assert TUIEventType.RUN_QUEUED
|
||||
assert TUIEventType.RUN_STARTED
|
||||
assert TUIEventType.RUN_COMPLETED
|
||||
assert TUIEventType.OUTPUT_CHUNK
|
||||
assert TUIEventType.TOOL_CALL_STARTED
|
||||
assert TUIEventType.ERROR_OCCURRED
|
||||
|
||||
|
||||
class TestTUIEvent:
|
||||
"""Tests for TUIEvent dataclass."""
|
||||
|
||||
def test_basic_creation(self):
|
||||
"""Test basic event creation."""
|
||||
event = TUIEvent(
|
||||
event_type=TUIEventType.MESSAGE_SUBMITTED,
|
||||
)
|
||||
assert event.event_type == TUIEventType.MESSAGE_SUBMITTED
|
||||
assert event.timestamp > 0
|
||||
assert event.data == {}
|
||||
|
||||
def test_with_data(self):
|
||||
"""Test event with data."""
|
||||
event = TUIEvent(
|
||||
event_type=TUIEventType.RUN_COMPLETED,
|
||||
run_id="test123",
|
||||
data={"output": "Hello"},
|
||||
)
|
||||
assert event.run_id == "test123"
|
||||
assert event.data["output"] == "Hello"
|
||||
|
||||
def test_message_submitted_factory(self):
|
||||
"""Test message_submitted factory method."""
|
||||
event = TUIEvent.message_submitted("Hello world")
|
||||
assert event.event_type == TUIEventType.MESSAGE_SUBMITTED
|
||||
assert event.data["content"] == "Hello world"
|
||||
|
||||
def test_output_chunk_factory(self):
|
||||
"""Test output_chunk factory method."""
|
||||
event = TUIEvent.output_chunk("run123", "chunk content")
|
||||
assert event.event_type == TUIEventType.OUTPUT_CHUNK
|
||||
assert event.run_id == "run123"
|
||||
assert event.data["content"] == "chunk content"
|
||||
|
||||
def test_run_completed_factory(self):
|
||||
"""Test run_completed factory method."""
|
||||
event = TUIEvent.run_completed("run123", "final output")
|
||||
assert event.event_type == TUIEventType.RUN_COMPLETED
|
||||
assert event.run_id == "run123"
|
||||
assert event.data["output"] == "final output"
|
||||
|
||||
def test_error_factory(self):
|
||||
"""Test error factory method."""
|
||||
event = TUIEvent.error("Something went wrong")
|
||||
assert event.event_type == TUIEventType.ERROR_OCCURRED
|
||||
assert event.data["message"] == "Something went wrong"
|
||||
|
||||
def test_status_update_factory(self):
|
||||
"""Test status_update factory method."""
|
||||
event = TUIEvent.status_update("Processing...")
|
||||
assert event.event_type == TUIEventType.STATUS_UPDATED
|
||||
assert event.data["status"] == "Processing..."
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Tests for Mock Provider."""
|
||||
|
||||
import pytest
|
||||
from praisonai.cli.features.tui.mock_provider import (
|
||||
MockProvider,
|
||||
MockProviderConfig,
|
||||
MockResponse,
|
||||
MockExecutor,
|
||||
create_mock_provider,
|
||||
)
|
||||
|
||||
|
||||
class TestMockResponse:
|
||||
"""Tests for MockResponse."""
|
||||
|
||||
def test_default_chunking(self):
|
||||
"""Test automatic chunking of content."""
|
||||
response = MockResponse(content="Hello world how are you today")
|
||||
|
||||
assert len(response.chunks) > 0
|
||||
assert "".join(response.chunks).strip() == response.content
|
||||
|
||||
def test_custom_chunks(self):
|
||||
"""Test custom chunks."""
|
||||
response = MockResponse(
|
||||
content="Hello world",
|
||||
chunks=["Hello", " world"],
|
||||
)
|
||||
|
||||
assert response.chunks == ["Hello", " world"]
|
||||
|
||||
def test_token_estimation(self):
|
||||
"""Test automatic token estimation."""
|
||||
response = MockResponse(content="one two three four five")
|
||||
|
||||
# Rough estimate: words * 2
|
||||
assert response.tokens == 10
|
||||
|
||||
|
||||
class TestMockProviderConfig:
|
||||
"""Tests for MockProviderConfig."""
|
||||
|
||||
def test_default_config(self):
|
||||
"""Test default configuration."""
|
||||
config = MockProviderConfig()
|
||||
|
||||
assert config.seed == 42
|
||||
assert config.default_delay == 0.05
|
||||
assert config.simulate_errors is False
|
||||
|
||||
def test_custom_responses(self):
|
||||
"""Test custom responses."""
|
||||
config = MockProviderConfig(
|
||||
responses={
|
||||
"custom": MockResponse(content="Custom response"),
|
||||
}
|
||||
)
|
||||
|
||||
assert "custom" in config.responses
|
||||
|
||||
|
||||
class TestMockProvider:
|
||||
"""Tests for MockProvider."""
|
||||
|
||||
def test_default_responses(self):
|
||||
"""Test default canned responses."""
|
||||
provider = MockProvider()
|
||||
|
||||
assert "hello" in provider.config.responses
|
||||
assert "help" in provider.config.responses
|
||||
assert "test" in provider.config.responses
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_hello(self):
|
||||
"""Test generating hello response."""
|
||||
provider = MockProvider()
|
||||
|
||||
result = await provider.generate("hello")
|
||||
|
||||
assert "content" in result
|
||||
assert "Hello" in result["content"]
|
||||
assert result["tokens"] > 0
|
||||
assert result["cost"] > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_streaming(self):
|
||||
"""Test streaming generation."""
|
||||
provider = MockProvider()
|
||||
|
||||
chunks = []
|
||||
result = await provider.generate(
|
||||
"hello",
|
||||
stream=True,
|
||||
on_chunk=lambda c: chunks.append(c),
|
||||
)
|
||||
|
||||
assert len(chunks) > 0
|
||||
assert "".join(chunks).strip() == result["content"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_non_streaming(self):
|
||||
"""Test non-streaming generation."""
|
||||
provider = MockProvider()
|
||||
|
||||
result = await provider.generate("hello", stream=False)
|
||||
|
||||
assert "content" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_error(self):
|
||||
"""Test error response."""
|
||||
provider = MockProvider()
|
||||
|
||||
with pytest.raises(Exception, match="Simulated error"):
|
||||
await provider.generate("error")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_tool_calls(self):
|
||||
"""Test tool call response."""
|
||||
provider = MockProvider()
|
||||
|
||||
result = await provider.generate("tool")
|
||||
|
||||
assert "tool_calls" in result
|
||||
assert len(result["tool_calls"]) > 0
|
||||
assert result["tool_calls"][0]["name"] == "mock_tool"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deterministic_responses(self):
|
||||
"""Test responses are deterministic with same seed."""
|
||||
provider1 = MockProvider(MockProviderConfig(seed=123))
|
||||
provider2 = MockProvider(MockProviderConfig(seed=123))
|
||||
|
||||
result1 = await provider1.generate("unique input xyz")
|
||||
result2 = await provider2.generate("unique input xyz")
|
||||
|
||||
assert result1["content"] == result2["content"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_seeds_different_responses(self):
|
||||
"""Test different seeds produce different responses."""
|
||||
provider1 = MockProvider(MockProviderConfig(seed=1))
|
||||
provider2 = MockProvider(MockProviderConfig(seed=2))
|
||||
|
||||
# Use input that doesn't match canned responses
|
||||
result1 = await provider1.generate("random query abc123")
|
||||
result2 = await provider2.generate("random query abc123")
|
||||
|
||||
# Content should be deterministic based on input hash, not seed
|
||||
# But the random error injection would differ
|
||||
assert result1["content"] == result2["content"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_iterator(self):
|
||||
"""Test stream iterator."""
|
||||
provider = MockProvider()
|
||||
|
||||
chunks = []
|
||||
async for chunk in provider.stream("hello"):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert len(chunks) > 0
|
||||
|
||||
def test_reset(self):
|
||||
"""Test reset functionality."""
|
||||
provider = MockProvider()
|
||||
provider._playback_index = 5
|
||||
|
||||
provider.reset()
|
||||
|
||||
assert provider._playback_index == 0
|
||||
|
||||
|
||||
class TestMockExecutor:
|
||||
"""Tests for MockExecutor."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute(self):
|
||||
"""Test executor execution."""
|
||||
executor = MockExecutor()
|
||||
|
||||
outputs = []
|
||||
result = await executor.execute(
|
||||
input_content="hello",
|
||||
agent_config={"model": "test-model"},
|
||||
on_output=lambda c: outputs.append(c),
|
||||
)
|
||||
|
||||
assert "output" in result
|
||||
assert "metrics" in result
|
||||
assert result["metrics"]["tokens"] > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_with_custom_provider(self):
|
||||
"""Test executor with custom provider."""
|
||||
provider = create_mock_provider(
|
||||
seed=42,
|
||||
responses={"custom": "Custom output"},
|
||||
)
|
||||
executor = MockExecutor(provider)
|
||||
|
||||
result = await executor.execute(
|
||||
input_content="custom query",
|
||||
agent_config={},
|
||||
)
|
||||
|
||||
assert "Custom output" in result["output"]
|
||||
|
||||
|
||||
class TestCreateMockProvider:
|
||||
"""Tests for create_mock_provider helper."""
|
||||
|
||||
def test_create_with_defaults(self):
|
||||
"""Test creating with defaults."""
|
||||
provider = create_mock_provider()
|
||||
|
||||
assert provider.config.seed == 42
|
||||
|
||||
def test_create_with_custom_seed(self):
|
||||
"""Test creating with custom seed."""
|
||||
provider = create_mock_provider(seed=123)
|
||||
|
||||
assert provider.config.seed == 123
|
||||
|
||||
def test_create_with_custom_responses(self):
|
||||
"""Test creating with custom responses."""
|
||||
provider = create_mock_provider(
|
||||
responses={"mypattern": "My response"}
|
||||
)
|
||||
|
||||
assert "mypattern" in provider.config.responses
|
||||
@@ -0,0 +1,412 @@
|
||||
"""Tests for TUI Orchestrator."""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from praisonai.cli.features.tui.orchestrator import (
|
||||
TuiOrchestrator,
|
||||
UIStateModel,
|
||||
SimulationRunner,
|
||||
OutputMode,
|
||||
)
|
||||
from praisonai.cli.features.tui.events import TUIEvent, TUIEventType
|
||||
from praisonai.cli.features.queue import QueueConfig
|
||||
|
||||
|
||||
class TestUIStateModel:
|
||||
"""Tests for UIStateModel."""
|
||||
|
||||
def test_default_state(self):
|
||||
"""Test default state values."""
|
||||
state = UIStateModel()
|
||||
assert state.session_id == ""
|
||||
assert state.model == "gpt-4o-mini"
|
||||
assert state.messages == []
|
||||
assert state.is_processing is False
|
||||
assert state.current_screen == "main"
|
||||
|
||||
def test_add_message(self):
|
||||
"""Test adding messages."""
|
||||
state = UIStateModel()
|
||||
state.add_message("user", "Hello")
|
||||
state.add_message("assistant", "Hi there!")
|
||||
|
||||
assert len(state.messages) == 2
|
||||
assert state.messages[0]["role"] == "user"
|
||||
assert state.messages[0]["content"] == "Hello"
|
||||
assert state.messages[1]["role"] == "assistant"
|
||||
|
||||
def test_add_message_with_extras(self):
|
||||
"""Test adding messages with extra fields."""
|
||||
state = UIStateModel()
|
||||
state.add_message("assistant", "Response", run_id="run123", agent_name="Agent1")
|
||||
|
||||
assert state.messages[0]["run_id"] == "run123"
|
||||
assert state.messages[0]["agent_name"] == "Agent1"
|
||||
|
||||
def test_message_limit(self):
|
||||
"""Test message limit enforcement."""
|
||||
state = UIStateModel(max_messages=5)
|
||||
|
||||
for i in range(10):
|
||||
state.add_message("user", f"Message {i}")
|
||||
|
||||
assert len(state.messages) == 5
|
||||
assert state.messages[0]["content"] == "Message 5"
|
||||
|
||||
def test_add_event(self):
|
||||
"""Test adding events."""
|
||||
state = UIStateModel()
|
||||
event = TUIEvent(event_type=TUIEventType.MESSAGE_SUBMITTED)
|
||||
state.add_event(event)
|
||||
|
||||
assert len(state.events) == 1
|
||||
assert state.events[0]["type"] == "message_submitted"
|
||||
|
||||
def test_to_snapshot(self):
|
||||
"""Test snapshot generation."""
|
||||
state = UIStateModel(session_id="test123", model="gpt-4")
|
||||
state.add_message("user", "Hello")
|
||||
state.is_processing = True
|
||||
|
||||
snapshot = state.to_snapshot()
|
||||
|
||||
assert snapshot["session_id"] == "test123"
|
||||
assert snapshot["model"] == "gpt-4"
|
||||
assert snapshot["message_count"] == 1
|
||||
assert snapshot["is_processing"] is True
|
||||
|
||||
def test_render_snapshot_pretty(self):
|
||||
"""Test pretty snapshot rendering."""
|
||||
state = UIStateModel(session_id="test123")
|
||||
state.add_message("user", "Hello")
|
||||
state.add_message("assistant", "Hi!")
|
||||
|
||||
output = state.render_snapshot_pretty()
|
||||
|
||||
assert "PraisonAI" in output
|
||||
assert "test123" in output
|
||||
assert "USER:" in output
|
||||
assert "ASSISTANT:" in output
|
||||
|
||||
|
||||
class TestTuiOrchestrator:
|
||||
"""Tests for TuiOrchestrator."""
|
||||
|
||||
@pytest.fixture
|
||||
def config(self):
|
||||
"""Create test config."""
|
||||
return QueueConfig(enable_persistence=False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_stop(self, config):
|
||||
"""Test starting and stopping orchestrator."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
await orchestrator.start(session_id="test-session")
|
||||
|
||||
assert orchestrator.state.session_id == "test-session"
|
||||
assert orchestrator._running is True
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
assert orchestrator._running is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_callbacks(self, config):
|
||||
"""Test event callbacks are called."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
events = []
|
||||
orchestrator.add_event_callback(lambda e: events.append(e))
|
||||
|
||||
await orchestrator.start()
|
||||
|
||||
# Should have session started event
|
||||
assert len(events) >= 1
|
||||
assert events[0].event_type == TUIEventType.SESSION_STARTED
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_model(self, config):
|
||||
"""Test setting model."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
await orchestrator.start()
|
||||
orchestrator.set_model("gpt-4")
|
||||
|
||||
assert orchestrator.state.model == "gpt-4"
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_navigate_screen(self, config):
|
||||
"""Test screen navigation."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
await orchestrator.start()
|
||||
orchestrator.navigate_screen("queue")
|
||||
|
||||
assert orchestrator.state.current_screen == "queue"
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_focus(self, config):
|
||||
"""Test focus setting."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
await orchestrator.start()
|
||||
orchestrator.set_focus("chat")
|
||||
|
||||
assert orchestrator.state.focused_widget == "chat"
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_snapshot(self, config):
|
||||
"""Test getting snapshot."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
await orchestrator.start(session_id="snap-test")
|
||||
|
||||
snapshot = orchestrator.get_snapshot()
|
||||
|
||||
assert snapshot["session_id"] == "snap-test"
|
||||
assert "model" in snapshot
|
||||
assert "is_processing" in snapshot
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_render_snapshot(self, config):
|
||||
"""Test rendering snapshot."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
await orchestrator.start()
|
||||
|
||||
output = orchestrator.render_snapshot()
|
||||
|
||||
assert "PraisonAI" in output
|
||||
assert "Queue" in output
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
|
||||
class TestSimulationRunner:
|
||||
"""Tests for SimulationRunner."""
|
||||
|
||||
@pytest.fixture
|
||||
def config(self):
|
||||
"""Create test config."""
|
||||
return QueueConfig(enable_persistence=False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_empty_script(self, config):
|
||||
"""Test running empty script."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator)
|
||||
|
||||
script = {"steps": []}
|
||||
success = await runner.run_script(script)
|
||||
|
||||
assert success is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_navigate_step(self, config):
|
||||
"""Test navigate step."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator)
|
||||
|
||||
script = {
|
||||
"steps": [
|
||||
{"action": "navigate", "args": {"screen": "queue"}},
|
||||
{"action": "navigate", "args": {"screen": "main"}},
|
||||
]
|
||||
}
|
||||
|
||||
success = await runner.run_script(script)
|
||||
|
||||
assert success is True
|
||||
assert orchestrator.state.current_screen == "main"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_focus_step(self, config):
|
||||
"""Test focus step."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator)
|
||||
|
||||
script = {
|
||||
"steps": [
|
||||
{"action": "focus", "args": {"widget": "chat"}},
|
||||
]
|
||||
}
|
||||
|
||||
success = await runner.run_script(script)
|
||||
|
||||
assert success is True
|
||||
assert orchestrator.state.focused_widget == "chat"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_model_step(self, config):
|
||||
"""Test model step."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator)
|
||||
|
||||
script = {
|
||||
"steps": [
|
||||
{"action": "model", "args": {"model": "gpt-4"}},
|
||||
]
|
||||
}
|
||||
|
||||
success = await runner.run_script(script)
|
||||
|
||||
assert success is True
|
||||
assert orchestrator.state.model == "gpt-4"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_sleep_step(self, config):
|
||||
"""Test sleep step."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator)
|
||||
|
||||
script = {
|
||||
"steps": [
|
||||
{"action": "sleep", "args": {"seconds": 0.1}},
|
||||
]
|
||||
}
|
||||
|
||||
import time
|
||||
start = time.time()
|
||||
success = await runner.run_script(script)
|
||||
elapsed = time.time() - start
|
||||
|
||||
assert success is True
|
||||
assert elapsed >= 0.1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assertion_mode_pass(self, config):
|
||||
"""Test assertion mode with passing assertions."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator, assert_mode=True)
|
||||
|
||||
script = {
|
||||
"session_id": "assert-test",
|
||||
"steps": [
|
||||
{
|
||||
"action": "model",
|
||||
"args": {"model": "gpt-4"},
|
||||
"expected": {"model": "gpt-4"}
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
success = await runner.run_script(script)
|
||||
summary = runner.get_summary()
|
||||
|
||||
assert success is True
|
||||
assert summary["assertions_passed"] == 1
|
||||
assert summary["assertions_failed"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assertion_mode_fail(self, config):
|
||||
"""Test assertion mode with failing assertions."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator, assert_mode=True)
|
||||
|
||||
script = {
|
||||
"steps": [
|
||||
{
|
||||
"action": "model",
|
||||
"args": {"model": "gpt-4"},
|
||||
"expected": {"model": "gpt-3.5"} # Wrong expectation
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
success = await runner.run_script(script)
|
||||
summary = runner.get_summary()
|
||||
|
||||
assert success is False
|
||||
assert summary["assertions_failed"] == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_action(self, config):
|
||||
"""Test unknown action raises error."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator)
|
||||
|
||||
script = {
|
||||
"steps": [
|
||||
{"action": "unknown_action", "args": {}},
|
||||
]
|
||||
}
|
||||
|
||||
success = await runner.run_script(script)
|
||||
|
||||
assert success is False
|
||||
assert len(runner.errors) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_summary(self, config):
|
||||
"""Test getting summary."""
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
runner = SimulationRunner(orchestrator)
|
||||
|
||||
script = {"steps": []}
|
||||
await runner.run_script(script)
|
||||
|
||||
summary = runner.get_summary()
|
||||
|
||||
assert "assertions_passed" in summary
|
||||
assert "assertions_failed" in summary
|
||||
assert "errors" in summary
|
||||
assert "success" in summary
|
||||
Reference in new issue
Block a user