mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 08:11:07 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ebb9b24bff | ||
|
|
759ef2c655 | ||
|
|
cb4318f6e8 | ||
|
|
d9bb869d5e | ||
|
|
34b6aeb76b | ||
|
|
078b7801e1 | ||
|
|
a970e849b4 | ||
|
|
8e68f4e71c | ||
|
|
308758cdc2 | ||
|
|
f8fa61b082 | ||
|
|
955d236924 | ||
|
|
d9d63c3a13 | ||
|
|
5bee6eb470 | ||
|
|
7cb1676b7c | ||
|
|
e2e9c3f9ee | ||
|
|
cbe8beeca6 | ||
|
|
f7b66a52b7 | ||
|
|
6dbf194312 | ||
|
|
335e8d7e47 | ||
|
|
d65840e092 | ||
|
|
5a13571fdd |
No files matched your search
@@ -261,6 +261,12 @@ npm install praisonai
|
||||
| ↳ 100+ LLM Support | [Example](examples/python/providers/openai/openai_gpt4_example.py) | [📖](https://docs.praison.ai/models) |
|
||||
| ↳ Callback Agents | [Example](examples/python/general/advanced-callback-systems.py) | [📖](https://docs.praison.ai/features/callbacks) |
|
||||
| ↳ Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
|
||||
| ↳ Shadow Git Checkpoints | [Example](#10-shadow-git-checkpoints) | [📖](https://docs.praison.ai/features/checkpoints) |
|
||||
| ↳ Background Tasks | [Example](examples/background/basic_background.py) | [📖](https://docs.praison.ai/features/background-tasks) |
|
||||
| ↳ Policy Engine | [Example](examples/policy/basic_policy.py) | [📖](https://docs.praison.ai/features/policy-engine) |
|
||||
| ↳ Thinking Budgets | [Example](examples/thinking/basic_thinking.py) | [📖](https://docs.praison.ai/features/thinking-budgets) |
|
||||
| ↳ Output Styles | [Example](examples/output/basic_output.py) | [📖](https://docs.praison.ai/features/output-styles) |
|
||||
| ↳ Context Compaction | [Example](examples/compaction/basic_compaction.py) | [📖](https://docs.praison.ai/features/context-compaction) |
|
||||
| **📊 Monitoring & Management** | | |
|
||||
| ↳ Sessions Management | [Example](examples/python/sessions/comprehensive-session-management.py) | [📖](https://docs.praison.ai/features/sessions) |
|
||||
| ↳ Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
@@ -278,6 +284,12 @@ npm install praisonai
|
||||
| ↳ Git Integration | [Example](examples/python/cli/git_integration_example.py) | [📖](https://docs.praison.ai/docs/cli/git-integration) |
|
||||
| ↳ Sandbox Execution | [Example](examples/python/cli/sandbox_execution_example.py) | [📖](https://docs.praison.ai/docs/cli/sandbox-execution) |
|
||||
| ↳ CLI Compare | [Example](examples/compare/cli_compare_basic.py) | [📖](https://docs.praison.ai/docs/cli/compare) |
|
||||
| ↳ Profile/Benchmark | [Example](#profile-benchmark) | [📖](https://docs.praison.ai/docs/cli/profile) |
|
||||
| ↳ Auto Mode | [Example](#auto-mode) | [📖](https://docs.praison.ai/docs/cli/auto) |
|
||||
| ↳ Init | [Example](#init) | [📖](https://docs.praison.ai/docs/cli/init) |
|
||||
| ↳ File Input | [Example](#file-input) | [📖](https://docs.praison.ai/docs/cli/file-input) |
|
||||
| ↳ Final Agent | [Example](#final-agent) | [📖](https://docs.praison.ai/docs/cli/final-agent) |
|
||||
| ↳ Max Tokens | [Example](#max-tokens) | [📖](https://docs.praison.ai/docs/cli/max-tokens) |
|
||||
| **🧪 Evaluation** | | |
|
||||
| ↳ Accuracy Evaluation | [Example](examples/eval/accuracy_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
| ↳ Performance Evaluation | [Example](examples/eval/performance_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
@@ -857,21 +869,258 @@ callbacks:
|
||||
|
||||
### 9. Hooks
|
||||
|
||||
Configure in `.praison/hooks.json`:
|
||||
Intercept and modify agent behavior at various lifecycle points:
|
||||
|
||||
```python
|
||||
from praisonaiagents.memory import HooksManager
|
||||
from praisonaiagents.hooks import (
|
||||
HookRegistry, HookRunner, HookEvent, HookResult,
|
||||
BeforeToolInput
|
||||
)
|
||||
|
||||
hooks = HooksManager()
|
||||
# Create a hook registry
|
||||
registry = HookRegistry()
|
||||
|
||||
# Register Python hooks
|
||||
hooks.register("pre_write_code", lambda ctx: print(f"Writing {ctx['file']}"))
|
||||
# Log all tool calls
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def log_tools(event_data: BeforeToolInput) -> HookResult:
|
||||
print(f"Tool: {event_data.tool_name}")
|
||||
return HookResult.allow()
|
||||
|
||||
# Block dangerous operations
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def security_check(event_data: BeforeToolInput) -> HookResult:
|
||||
if "delete" in event_data.tool_name.lower():
|
||||
return HookResult.deny("Delete operations blocked")
|
||||
return HookResult.allow()
|
||||
|
||||
# Execute hooks
|
||||
result = hooks.execute("pre_write_code", {"file": "main.py"})
|
||||
runner = HookRunner(registry)
|
||||
```
|
||||
|
||||
### 10. Field Names Reference (A-I-G-S)
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai hooks list # List registered hooks
|
||||
praisonai hooks test before_tool # Test hooks for an event
|
||||
praisonai hooks run "echo test" # Run a command hook
|
||||
praisonai hooks validate hooks.json # Validate configuration
|
||||
```
|
||||
|
||||
|
||||
### 10. Shadow Git Checkpoints
|
||||
|
||||
File-level undo/restore using shadow git:
|
||||
|
||||
```python
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
|
||||
service = CheckpointService(workspace_dir="./my_project")
|
||||
await service.initialize()
|
||||
|
||||
# Save checkpoint before changes
|
||||
result = await service.save("Before refactoring")
|
||||
|
||||
# Make changes...
|
||||
|
||||
# Restore if needed
|
||||
await service.restore(result.checkpoint.id)
|
||||
|
||||
# View diff
|
||||
diff = await service.diff()
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai checkpoint save "Before changes" # Save checkpoint
|
||||
praisonai checkpoint list # List checkpoints
|
||||
praisonai checkpoint diff # Show changes
|
||||
praisonai checkpoint restore abc123 # Restore to checkpoint
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/checkpoints)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/checkpoint)
|
||||
- [💻 Example](examples/checkpoints/basic_checkpoints.py)
|
||||
|
||||
---
|
||||
|
||||
### 11. Background Tasks
|
||||
|
||||
Run agent tasks asynchronously without blocking:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from praisonaiagents.background import BackgroundRunner, BackgroundConfig
|
||||
|
||||
async def main():
|
||||
config = BackgroundConfig(max_concurrent_tasks=3)
|
||||
runner = BackgroundRunner(config=config)
|
||||
|
||||
async def my_task(name: str) -> str:
|
||||
await asyncio.sleep(2)
|
||||
return f"Task {name} completed"
|
||||
|
||||
task = await runner.submit(my_task, args=("example",), name="my_task")
|
||||
await task.wait(timeout=10.0)
|
||||
print(task.result)
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai background list # List running tasks
|
||||
praisonai background status <id> # Check task status
|
||||
praisonai background cancel <id> # Cancel a task
|
||||
praisonai background clear # Clear completed tasks
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/background-tasks)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/background)
|
||||
- [💻 Example](examples/background/basic_background.py)
|
||||
|
||||
---
|
||||
|
||||
### 12. Policy Engine
|
||||
|
||||
Control what agents can and cannot do with policy-based execution:
|
||||
|
||||
```python
|
||||
from praisonaiagents.policy import (
|
||||
PolicyEngine, Policy, PolicyRule, PolicyAction
|
||||
)
|
||||
|
||||
engine = PolicyEngine()
|
||||
|
||||
policy = Policy(
|
||||
name="no_delete",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="Delete operations blocked"
|
||||
)
|
||||
]
|
||||
)
|
||||
engine.add_policy(policy)
|
||||
|
||||
result = engine.check("tool:delete_file", {})
|
||||
print(f"Allowed: {result.allowed}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai policy list # List policies
|
||||
praisonai policy check "tool:name" # Check if allowed
|
||||
praisonai policy init # Create template
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/policy-engine)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/policy)
|
||||
- [💻 Example](examples/policy/basic_policy.py)
|
||||
|
||||
---
|
||||
|
||||
### 13. Thinking Budgets
|
||||
|
||||
Configure token budgets for extended thinking:
|
||||
|
||||
```python
|
||||
from praisonaiagents.thinking import ThinkingBudget, ThinkingTracker
|
||||
|
||||
# Use predefined levels
|
||||
budget = ThinkingBudget.high() # 16,000 tokens
|
||||
|
||||
# Track usage
|
||||
tracker = ThinkingTracker()
|
||||
session = tracker.start_session(budget_tokens=16000)
|
||||
tracker.end_session(session, tokens_used=12000)
|
||||
|
||||
summary = tracker.get_summary()
|
||||
print(f"Utilization: {summary['average_utilization']:.1%}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai thinking status # Show current budget
|
||||
praisonai thinking set high # Set budget level
|
||||
praisonai thinking stats # Show usage statistics
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/thinking-budgets)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/thinking)
|
||||
- [💻 Example](examples/thinking/basic_thinking.py)
|
||||
|
||||
---
|
||||
|
||||
### 14. Output Styles
|
||||
|
||||
Configure how agents format their responses:
|
||||
|
||||
```python
|
||||
from praisonaiagents.output import OutputStyle, OutputFormatter
|
||||
|
||||
# Use preset styles
|
||||
style = OutputStyle.concise()
|
||||
formatter = OutputFormatter(style)
|
||||
|
||||
# Format output
|
||||
text = "# Hello\n\nThis is **bold** text."
|
||||
plain = formatter.format(text)
|
||||
print(plain)
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai output status # Show current style
|
||||
praisonai output set concise # Set output style
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/output-styles)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/output-style)
|
||||
- [💻 Example](examples/output/basic_output.py)
|
||||
|
||||
---
|
||||
|
||||
### 15. Context Compaction
|
||||
|
||||
Automatically manage context window size:
|
||||
|
||||
```python
|
||||
from praisonaiagents.compaction import (
|
||||
ContextCompactor, CompactionStrategy
|
||||
)
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=4000,
|
||||
strategy=CompactionStrategy.SLIDING,
|
||||
preserve_recent=3
|
||||
)
|
||||
|
||||
messages = [...] # Your conversation history
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f"Compression: {result.compression_ratio:.1%}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai compaction status # Show settings
|
||||
praisonai compaction set sliding # Set strategy
|
||||
praisonai compaction stats # Show statistics
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/context-compaction)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/compaction)
|
||||
- [💻 Example](examples/compaction/basic_compaction.py)
|
||||
|
||||
---
|
||||
|
||||
### 16. Field Names Reference (A-I-G-S)
|
||||
|
||||
PraisonAI accepts both old (agents.yaml) and new (workflow.yaml) field names. Use the **canonical names** for new projects:
|
||||
|
||||
@@ -905,7 +1154,7 @@ steps: # Define steps (not 'tasks')
|
||||
|
||||
> **Note:** The parser accepts both old and new names. Run `praisonai workflow validate <file.yaml>` to see suggestions for canonical names.
|
||||
|
||||
### 11. Extended agents.yaml with Workflow Patterns
|
||||
### 12. Extended agents.yaml with Workflow Patterns
|
||||
|
||||
**Feature Parity:** Both `agents.yaml` and `workflow.yaml` now support the same features:
|
||||
- All workflow patterns (route, parallel, loop, repeat)
|
||||
@@ -981,6 +1230,124 @@ Run with the same simple command:
|
||||
praisonai agents.yaml
|
||||
```
|
||||
|
||||
### 12. MCP (Model Context Protocol)
|
||||
|
||||
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
|
||||
|
||||
#### MCP Client (Consume MCP Servers)
|
||||
```python
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
# stdio - Local NPX/Python servers
|
||||
agent = Agent(tools=MCP("npx @modelcontextprotocol/server-memory"))
|
||||
|
||||
# Streamable HTTP - Production servers
|
||||
agent = Agent(tools=MCP("https://api.example.com/mcp"))
|
||||
|
||||
# WebSocket - Real-time bidirectional
|
||||
agent = Agent(tools=MCP("wss://api.example.com/mcp", auth_token="token"))
|
||||
|
||||
# SSE (Legacy) - Backward compatibility
|
||||
agent = Agent(tools=MCP("http://localhost:8080/sse"))
|
||||
|
||||
# With environment variables
|
||||
agent = Agent(
|
||||
tools=MCP(
|
||||
command="npx",
|
||||
args=["-y", "@modelcontextprotocol/server-brave-search"],
|
||||
env={"BRAVE_API_KEY": "your-key"}
|
||||
)
|
||||
)
|
||||
|
||||
# Multiple MCP servers + regular functions
|
||||
def my_custom_tool(query: str) -> str:
|
||||
"""Custom tool function."""
|
||||
return f"Result: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="MultiToolAgent",
|
||||
instructions="Agent with multiple MCP servers",
|
||||
tools=[
|
||||
MCP("uvx mcp-server-time"), # Time tools
|
||||
MCP("npx @modelcontextprotocol/server-memory"), # Memory tools
|
||||
my_custom_tool # Regular function
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
#### MCP Server (Expose Tools as MCP Server)
|
||||
|
||||
Expose your Python functions as MCP tools for Claude Desktop, Cursor, and other MCP clients:
|
||||
|
||||
```python
|
||||
from praisonaiagents.mcp import ToolsMCPServer
|
||||
|
||||
def search_web(query: str, max_results: int = 5) -> dict:
|
||||
"""Search the web for information."""
|
||||
return {"results": [f"Result for {query}"]}
|
||||
|
||||
def calculate(expression: str) -> dict:
|
||||
"""Evaluate a mathematical expression."""
|
||||
return {"result": eval(expression)}
|
||||
|
||||
# Create and run MCP server
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
server.register_tools([search_web, calculate])
|
||||
server.run() # stdio for Claude Desktop
|
||||
# server.run_sse(host="0.0.0.0", port=8080) # SSE for web clients
|
||||
```
|
||||
|
||||
#### MCP Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Session Management | Automatic Mcp-Session-Id handling |
|
||||
| Protocol Versioning | Mcp-Protocol-Version header |
|
||||
| Resumability | SSE stream recovery via Last-Event-ID |
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
### 13. A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
#### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
#### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
---
|
||||
|
||||
## 🎯 CLI / No-Code Interface
|
||||
@@ -1730,124 +2097,6 @@ agent = Agent(
|
||||
| Web Fetch | Anthropic |
|
||||
| Prompt Caching | OpenAI (auto), Anthropic, Bedrock, Deepseek |
|
||||
|
||||
## MCP (Model Context Protocol)
|
||||
|
||||
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
|
||||
|
||||
### MCP Client (Consume MCP Servers)
|
||||
```python
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
# stdio - Local NPX/Python servers
|
||||
agent = Agent(tools=MCP("npx @modelcontextprotocol/server-memory"))
|
||||
|
||||
# Streamable HTTP - Production servers
|
||||
agent = Agent(tools=MCP("https://api.example.com/mcp"))
|
||||
|
||||
# WebSocket - Real-time bidirectional
|
||||
agent = Agent(tools=MCP("wss://api.example.com/mcp", auth_token="token"))
|
||||
|
||||
# SSE (Legacy) - Backward compatibility
|
||||
agent = Agent(tools=MCP("http://localhost:8080/sse"))
|
||||
|
||||
# With environment variables
|
||||
agent = Agent(
|
||||
tools=MCP(
|
||||
command="npx",
|
||||
args=["-y", "@modelcontextprotocol/server-brave-search"],
|
||||
env={"BRAVE_API_KEY": "your-key"}
|
||||
)
|
||||
)
|
||||
|
||||
# Multiple MCP servers + regular functions
|
||||
def my_custom_tool(query: str) -> str:
|
||||
"""Custom tool function."""
|
||||
return f"Result: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="MultiToolAgent",
|
||||
instructions="Agent with multiple MCP servers",
|
||||
tools=[
|
||||
MCP("uvx mcp-server-time"), # Time tools
|
||||
MCP("npx @modelcontextprotocol/server-memory"), # Memory tools
|
||||
my_custom_tool # Regular function
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### MCP Server (Expose Tools as MCP Server)
|
||||
|
||||
Expose your Python functions as MCP tools for Claude Desktop, Cursor, and other MCP clients:
|
||||
|
||||
```python
|
||||
from praisonaiagents.mcp import ToolsMCPServer
|
||||
|
||||
def search_web(query: str, max_results: int = 5) -> dict:
|
||||
"""Search the web for information."""
|
||||
return {"results": [f"Result for {query}"]}
|
||||
|
||||
def calculate(expression: str) -> dict:
|
||||
"""Evaluate a mathematical expression."""
|
||||
return {"result": eval(expression)}
|
||||
|
||||
# Create and run MCP server
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
server.register_tools([search_web, calculate])
|
||||
server.run() # stdio for Claude Desktop
|
||||
# server.run_sse(host="0.0.0.0", port=8080) # SSE for web clients
|
||||
```
|
||||
|
||||
### MCP Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Session Management | Automatic Mcp-Session-Id handling |
|
||||
| Protocol Versioning | Mcp-Protocol-Version header |
|
||||
| Resumability | SSE stream recovery via Last-Event-ID |
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
## A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
## CLI Features
|
||||
|
||||
| Feature | Docs |
|
||||
@@ -2475,6 +2724,103 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
|
||||
---
|
||||
|
||||
## 💾 Persistence (Databases)
|
||||
|
||||
Enable automatic conversation persistence with 2 lines of code:
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
db=db(database_url="postgresql://localhost/mydb"), # db(...) shortcut
|
||||
session_id="my-session" # Optional: defaults to per-hour ID (YYYYMMDDHH)
|
||||
)
|
||||
agent.chat("Hello!") # Auto-persists messages, runs, traces
|
||||
```
|
||||
|
||||
### Persistence CLI Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai persistence doctor` | Validate DB connectivity |
|
||||
| `praisonai persistence run` | Run agent with persistence |
|
||||
| `praisonai persistence resume` | Resume existing session |
|
||||
| `praisonai persistence export` | Export session to JSONL |
|
||||
| `praisonai persistence import` | Import session from JSONL |
|
||||
| `praisonai persistence migrate` | Apply schema migrations |
|
||||
| `praisonai persistence status` | Show schema status |
|
||||
|
||||
### Databases Table
|
||||
|
||||
| Database | Store Type | Install | Example | Docs |
|
||||
|----------|------------|---------|---------|------|
|
||||
| PostgreSQL | Conversation | `pip install "praisonai[tools]"` | [simple_db_agent.py](examples/persistence/simple_db_agent.py) | [docs](https://docs.praison.ai/docs/databases/postgres) |
|
||||
| MySQL | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SQLite | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SingleStore | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Supabase | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SurrealDB | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Qdrant | Knowledge | `pip install "praisonai[tools]"` | [knowledge_qdrant.py](examples/persistence/knowledge_qdrant.py) | [docs](https://docs.praison.ai/docs/databases/qdrant) |
|
||||
| ChromaDB | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Pinecone | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Weaviate | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| LanceDB | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Milvus | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| PGVector | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Redis Vector | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Cassandra | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| ClickHouse | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Redis | State | `pip install "praisonai[tools]"` | [state_redis.py](examples/persistence/state_redis.py) | [docs](https://docs.praison.ai/docs/databases/redis) |
|
||||
| MongoDB | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| DynamoDB | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Firestore | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Upstash | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Memory | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Tools Table
|
||||
|
||||
Install all tools with: `pip install "praisonai[tools]"`
|
||||
|
||||
| Tool | Category | Import | Docs |
|
||||
|------|----------|--------|------|
|
||||
| Tavily | Web Search | `from praisonai_tools import TavilyTool` | [docs](https://docs.praison.ai/docs/tools/external/tavily) |
|
||||
| DuckDuckGo | Web Search | `from praisonai_tools import DuckDuckGoTool` | [docs](https://docs.praison.ai/docs/tools/external/duckduckgo) |
|
||||
| Exa | Web Search | `from praisonai_tools import ExaTool` | [docs](https://docs.praison.ai/docs/tools/external/exa) |
|
||||
| Serper | Web Search | `from praisonai_tools import SerperTool` | [docs](https://docs.praison.ai/docs/tools/external/serper) |
|
||||
| Jina | Web Reader | `from praisonai_tools import JinaTool` | [docs](https://docs.praison.ai/docs/tools/external/jina) |
|
||||
| Firecrawl | Web Scraping | `from praisonai_tools import FirecrawlTool` | [docs](https://docs.praison.ai/docs/tools/external/firecrawl) |
|
||||
| Crawl4AI | Web Scraping | `from praisonai_tools import Crawl4AITool` | [docs](https://docs.praison.ai/docs/tools/external/crawl4ai) |
|
||||
| Wikipedia | Knowledge | `from praisonai_tools import WikipediaTool` | [docs](https://docs.praison.ai/docs/tools/external/wikipedia) |
|
||||
| ArXiv | Research | `from praisonai_tools import ArxivTool` | [docs](https://docs.praison.ai/docs/tools/external/arxiv) |
|
||||
| HackerNews | News | `from praisonai_tools import HackerNewsTool` | [docs](https://docs.praison.ai/docs/tools/external/hackernews) |
|
||||
| YouTube | Media | `from praisonai_tools import YouTubeTool` | [docs](https://docs.praison.ai/docs/tools/external/youtube) |
|
||||
| Weather | Data | `from praisonai_tools import WeatherTool` | [docs](https://docs.praison.ai/docs/tools/external/weather) |
|
||||
| PostgreSQL | Database | `from praisonai_tools import PostgresTool` | [docs](https://docs.praison.ai/docs/tools/external/postgres) |
|
||||
| MySQL | Database | `from praisonai_tools import MySQLTool` | [docs](https://docs.praison.ai/docs/tools/external/mysql) |
|
||||
| SQLite | Database | `from praisonai_tools import SQLiteTool` | [docs](https://docs.praison.ai/docs/tools/external/sqlite) |
|
||||
| MongoDB | Database | `from praisonai_tools import MongoDBTool` | [docs](https://docs.praison.ai/docs/tools/external/mongodb) |
|
||||
| Redis | Database | `from praisonai_tools import RedisTool` | [docs](https://docs.praison.ai/docs/tools/external/redis) |
|
||||
| Qdrant | Vector DB | `from praisonai_tools import QdrantTool` | [docs](https://docs.praison.ai/docs/tools/external/qdrant) |
|
||||
| GitHub | DevOps | `from praisonai_tools import GitHubTool` | [docs](https://docs.praison.ai/docs/tools/external/github) |
|
||||
| Slack | Communication | `from praisonai_tools import SlackTool` | [docs](https://docs.praison.ai/docs/tools/external/slack) |
|
||||
| Discord | Communication | `from praisonai_tools import DiscordTool` | [docs](https://docs.praison.ai/docs/tools/external/discord) |
|
||||
| Telegram | Communication | `from praisonai_tools import TelegramTool` | [docs](https://docs.praison.ai/docs/tools/external/telegram) |
|
||||
| Email | Communication | `from praisonai_tools import EmailTool` | [docs](https://docs.praison.ai/docs/tools/external/email) |
|
||||
| Notion | Productivity | `from praisonai_tools import NotionTool` | [docs](https://docs.praison.ai/docs/tools/external/notion) |
|
||||
| File | File System | `from praisonai_tools import FileTool` | [docs](https://docs.praison.ai/docs/tools/external/file) |
|
||||
| Shell | System | `from praisonai_tools import ShellTool` | [docs](https://docs.praison.ai/docs/tools/external/shell) |
|
||||
| Python | Code | `from praisonai_tools import PythonTool` | [docs](https://docs.praison.ai/docs/tools/external/python) |
|
||||
| JSON | Data | `from praisonai_tools import JSONTool` | [docs](https://docs.praison.ai/docs/tools/external/json) |
|
||||
| CSV | Data | `from praisonai_tools import CSVTool` | [docs](https://docs.praison.ai/docs/tools/external/csv) |
|
||||
| Calculator | Math | `from praisonai_tools import CalculatorTool` | [docs](https://docs.praison.ai/docs/tools/external/calculator) |
|
||||
|
||||
> See [full tools documentation](https://docs.praison.ai/docs/tools/tools) for all 100+ available tools.
|
||||
|
||||
---
|
||||
|
||||
## 🎓 Video Tutorials
|
||||
|
||||
Learn PraisonAI through our comprehensive video series:
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.3.74" \
|
||||
"praisonai>=2.4.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.3.74" \
|
||||
"praisonai>=2.4.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.3.74" \
|
||||
"praisonai>=2.4.2" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Background Tasks Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the background task system to:
|
||||
1. Run tasks asynchronously
|
||||
2. Track task progress
|
||||
3. Cancel running tasks
|
||||
4. Handle task completion
|
||||
|
||||
Usage:
|
||||
python basic_background.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
from praisonaiagents.background import BackgroundRunner, BackgroundConfig, TaskStatus
|
||||
|
||||
|
||||
async def main():
|
||||
print("=" * 60)
|
||||
print("Background Tasks Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a background runner with config
|
||||
config = BackgroundConfig(max_concurrent_tasks=3)
|
||||
runner = BackgroundRunner(config=config)
|
||||
|
||||
# Define some example tasks
|
||||
def slow_task(duration: float, name: str) -> str:
|
||||
"""A slow task that simulates work."""
|
||||
print(f" 🔄 Task '{name}' starting (will take {duration}s)")
|
||||
time.sleep(duration)
|
||||
return f"Task '{name}' completed after {duration}s"
|
||||
|
||||
async def async_task(duration: float, name: str) -> str:
|
||||
"""An async task that simulates work."""
|
||||
print(f" 🔄 Async task '{name}' starting (will take {duration}s)")
|
||||
await asyncio.sleep(duration)
|
||||
return f"Async task '{name}' completed after {duration}s"
|
||||
|
||||
# Submit tasks
|
||||
print("\n--- Submitting Tasks ---")
|
||||
|
||||
task1 = await runner.submit(slow_task, args=(1.0, "quick-job"), name="task1")
|
||||
print(f"✅ Submitted task 1: {task1.id[:8]}")
|
||||
|
||||
task2 = await runner.submit(slow_task, args=(2.0, "medium-job"), name="task2")
|
||||
print(f"✅ Submitted task 2: {task2.id[:8]}")
|
||||
|
||||
task3 = await runner.submit(async_task, args=(1.5, "async-job"), name="task3")
|
||||
print(f"✅ Submitted async task 3: {task3.id[:8]}")
|
||||
|
||||
# List all tasks
|
||||
print("\n--- Task Status ---")
|
||||
for task in runner.tasks:
|
||||
print(f" [{task.id[:8]}] {task.status.value}")
|
||||
|
||||
# Wait for task 1 to complete
|
||||
print("\n--- Waiting for Task 1 ---")
|
||||
await task1.wait(timeout=5.0)
|
||||
print(f"✅ Task 1 result: {task1.result}")
|
||||
|
||||
# Check status of all tasks
|
||||
print("\n--- Updated Task Status ---")
|
||||
for task in [task1, task2, task3]:
|
||||
status = "✅" if task.status == TaskStatus.COMPLETED else "🔄"
|
||||
print(f" {status} [{task.id[:8]}] {task.status.value}")
|
||||
|
||||
# Wait for all remaining tasks
|
||||
print("\n--- Waiting for All Tasks ---")
|
||||
await task2.wait(timeout=10.0)
|
||||
await task3.wait(timeout=10.0)
|
||||
|
||||
# Final status
|
||||
print("\n--- Final Results ---")
|
||||
for task in [task1, task2, task3]:
|
||||
if task.result:
|
||||
print(f" ✅ {task.result}")
|
||||
|
||||
# Cleanup
|
||||
runner.clear_completed()
|
||||
print("\n✅ Cleared completed tasks")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,110 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Checkpoints Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the shadow git checkpointing system to:
|
||||
1. Save checkpoints before making changes
|
||||
2. List available checkpoints
|
||||
3. View diffs between checkpoints
|
||||
4. Restore to a previous checkpoint
|
||||
|
||||
Usage:
|
||||
python basic_checkpoints.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
|
||||
|
||||
async def main():
|
||||
# Create a temporary workspace for demonstration
|
||||
with tempfile.TemporaryDirectory() as workspace:
|
||||
print("=" * 60)
|
||||
print("Shadow Git Checkpointing Demo")
|
||||
print("=" * 60)
|
||||
print(f"\nWorkspace: {workspace}")
|
||||
|
||||
# Create some initial files
|
||||
test_file = Path(workspace) / "example.py"
|
||||
test_file.write_text("# Initial content\nprint('Hello, World!')\n")
|
||||
print(f"\n✅ Created initial file: {test_file.name}")
|
||||
|
||||
# Initialize the checkpoint service
|
||||
service = CheckpointService(workspace_dir=workspace)
|
||||
|
||||
try:
|
||||
# Initialize the shadow git repository
|
||||
print("\n--- Initializing Checkpoint Service ---")
|
||||
await service.initialize()
|
||||
print("✅ Shadow git repository initialized")
|
||||
|
||||
# Save first checkpoint
|
||||
print("\n--- Saving Checkpoint 1: Initial State ---")
|
||||
result1 = await service.save("Initial state")
|
||||
print(f"✅ Checkpoint saved: {result1.checkpoint.id[:8]}")
|
||||
print(f" Message: {result1.checkpoint.message}")
|
||||
|
||||
# Make some changes
|
||||
print("\n--- Making Changes ---")
|
||||
test_file.write_text("# Modified content\nprint('Hello, PraisonAI!')\nprint('Checkpoints are awesome!')\n")
|
||||
print("✅ Modified example.py")
|
||||
|
||||
# Create a new file
|
||||
new_file = Path(workspace) / "utils.py"
|
||||
new_file.write_text("def helper():\n return 'I help!'\n")
|
||||
print("✅ Created utils.py")
|
||||
|
||||
# Save second checkpoint
|
||||
print("\n--- Saving Checkpoint 2: After Changes ---")
|
||||
result2 = await service.save("Added features")
|
||||
print(f"✅ Checkpoint saved: {result2.checkpoint.id[:8]}")
|
||||
|
||||
# List all checkpoints
|
||||
print("\n--- Listing All Checkpoints ---")
|
||||
checkpoints = await service.list_checkpoints()
|
||||
for i, cp in enumerate(checkpoints, 1):
|
||||
print(f" {i}. [{cp.id[:8]}] {cp.message} ({cp.timestamp})")
|
||||
|
||||
# Show diff
|
||||
print("\n--- Current Diff (uncommitted changes) ---")
|
||||
diff = await service.diff()
|
||||
if diff.files:
|
||||
for f in diff.files:
|
||||
print(f" {f.status}: {f.path}")
|
||||
else:
|
||||
print(" No uncommitted changes")
|
||||
|
||||
# Make more changes to show diff
|
||||
test_file.write_text("# Even more changes\nprint('This will be reverted')\n")
|
||||
print("\n✅ Made more changes to example.py")
|
||||
|
||||
diff = await service.diff()
|
||||
print("\n--- New Diff ---")
|
||||
if diff.files:
|
||||
for f in diff.files:
|
||||
print(f" {f.status}: {f.path}")
|
||||
|
||||
# Restore to checkpoint 2
|
||||
print(f"\n--- Restoring to Checkpoint: {result2.checkpoint.id[:8]} ---")
|
||||
await service.restore(result2.checkpoint.id)
|
||||
print("✅ Restored successfully!")
|
||||
|
||||
# Verify restoration
|
||||
content = test_file.read_text()
|
||||
print(f"\n--- Verified Content After Restore ---")
|
||||
print(content)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n❌ Error: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,93 @@
|
||||
# Persistence CLI Commands
|
||||
|
||||
## Quick Reference
|
||||
|
||||
```bash
|
||||
# Check database connectivity
|
||||
praisonai persistence doctor \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--state-url redis://localhost:6379
|
||||
|
||||
# Run agent with persistence
|
||||
praisonai persistence run \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--session-id my-session
|
||||
|
||||
# Resume a session
|
||||
praisonai persistence resume \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--session-id my-session
|
||||
|
||||
# Export session to file
|
||||
praisonai persistence export \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--session-id my-session \
|
||||
--output session_backup.jsonl
|
||||
|
||||
# Import session from file
|
||||
praisonai persistence import \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--file session_backup.jsonl
|
||||
|
||||
# Check schema status
|
||||
praisonai persistence status \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--state-url redis://localhost:6379
|
||||
|
||||
# Apply migrations
|
||||
praisonai persistence migrate \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--state-url redis://localhost:6379
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
Instead of CLI flags, you can use environment variables:
|
||||
|
||||
```bash
|
||||
export PRAISON_CONVERSATION_URL=postgresql://localhost:5432/praisonai
|
||||
export PRAISON_STATE_URL=redis://localhost:6379
|
||||
export PRAISON_KNOWLEDGE_URL=http://localhost:6333
|
||||
|
||||
# Then run without URL flags
|
||||
praisonai persistence doctor
|
||||
praisonai persistence status
|
||||
```
|
||||
|
||||
## Docker Setup
|
||||
|
||||
```bash
|
||||
# Start all services
|
||||
docker run -d --name praison-postgres \
|
||||
-e POSTGRES_PASSWORD=praison123 \
|
||||
-e POSTGRES_DB=praisonai \
|
||||
-p 5432:5432 postgres:16
|
||||
|
||||
docker run -d --name praison-redis \
|
||||
-p 6379:6379 redis:7
|
||||
|
||||
docker run -d --name praison-qdrant \
|
||||
-p 6333:6333 -p 6334:6334 qdrant/qdrant
|
||||
|
||||
# Verify
|
||||
docker ps
|
||||
```
|
||||
|
||||
## Example Workflow
|
||||
|
||||
```bash
|
||||
# 1. Start services
|
||||
docker-compose up -d
|
||||
|
||||
# 2. Check connectivity
|
||||
praisonai persistence doctor
|
||||
|
||||
# 3. Run agent
|
||||
praisonai persistence run --session-id demo-session
|
||||
|
||||
# 4. Export for backup
|
||||
praisonai persistence export --session-id demo-session --output backup.jsonl
|
||||
|
||||
# 5. Import to another environment
|
||||
praisonai persistence import --file backup.jsonl
|
||||
```
|
||||
@@ -0,0 +1,124 @@
|
||||
# PraisonAI Persistence CLI - End-to-End Guide
|
||||
|
||||
This guide demonstrates the CLI commands for database persistence.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
# Start Docker containers
|
||||
docker run -d --name praison-postgres -p 5432:5432 \
|
||||
-e POSTGRES_PASSWORD=praison123 \
|
||||
-e POSTGRES_DB=praisonai \
|
||||
postgres:16
|
||||
|
||||
docker run -d --name praison-qdrant -p 6333:6333 qdrant/qdrant
|
||||
docker run -d --name praison-redis -p 6379:6379 redis:7
|
||||
```
|
||||
|
||||
## 1. Doctor Command - Validate Connectivity
|
||||
|
||||
```bash
|
||||
# Check all stores
|
||||
praisonai persistence doctor \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--knowledge-url "http://localhost:6333" \
|
||||
--state-url "redis://localhost:6379"
|
||||
```
|
||||
|
||||
**Expected Output:**
|
||||
```
|
||||
==================================================
|
||||
PraisonAI Persistence Doctor
|
||||
==================================================
|
||||
|
||||
[Conversation Store] Testing: postgresql://postgres:****@localhost:5432/praisonai
|
||||
✅ Connected (postgres)
|
||||
|
||||
[Knowledge Store] Testing: http://localhost:6333
|
||||
✅ Connected (qdrant)
|
||||
|
||||
[State Store] Testing: redis://localhost:6379
|
||||
✅ Connected (redis)
|
||||
|
||||
==================================================
|
||||
Results: 3/3 stores connected successfully
|
||||
==================================================
|
||||
```
|
||||
|
||||
## 2. Run Command - Execute Agent with Persistence
|
||||
|
||||
```bash
|
||||
# Run agent with persistence (dry-run first)
|
||||
praisonai persistence run \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--dry-run \
|
||||
"Hello, my name is Bob"
|
||||
|
||||
# Actually run it
|
||||
praisonai persistence run \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
"Hello, my name is Bob"
|
||||
```
|
||||
|
||||
**Expected Output:**
|
||||
```
|
||||
[Session: cli-demo-001]
|
||||
Agent: Hello Bob! Nice to meet you. How can I help you today?
|
||||
```
|
||||
|
||||
## 3. Resume Command - Continue a Session
|
||||
|
||||
```bash
|
||||
# Show history
|
||||
praisonai persistence resume \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--show-history
|
||||
|
||||
# Continue conversation
|
||||
praisonai persistence resume \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--continue "What's my name?"
|
||||
```
|
||||
|
||||
**Expected Output:**
|
||||
```
|
||||
[Session: cli-demo-001]
|
||||
Messages in history: 2
|
||||
|
||||
--- Conversation History ---
|
||||
[USER] Hello, my name is Bob
|
||||
[ASSISTANT] Hello Bob! Nice to meet you...
|
||||
--- End History ---
|
||||
|
||||
Agent: Your name is Bob!
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
Instead of passing URLs on command line, use environment variables:
|
||||
|
||||
```bash
|
||||
export PRAISON_CONVERSATION_URL="postgresql://postgres:praison123@localhost:5432/praisonai"
|
||||
export PRAISON_KNOWLEDGE_URL="http://localhost:6333"
|
||||
export PRAISON_STATE_URL="redis://localhost:6379"
|
||||
export OPENAI_API_KEY="your-key"
|
||||
|
||||
# Now commands are simpler
|
||||
praisonai persistence doctor --all
|
||||
praisonai persistence run --session-id "my-session" "Hello!"
|
||||
```
|
||||
|
||||
## Full CLI Reference
|
||||
|
||||
```
|
||||
praisonai persistence --help
|
||||
praisonai persistence doctor --help
|
||||
praisonai persistence run --help
|
||||
praisonai persistence resume --help
|
||||
```
|
||||
@@ -0,0 +1,146 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Context Compaction Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use context compaction to:
|
||||
1. Manage conversation context length
|
||||
2. Use different compaction strategies
|
||||
3. Track compaction results
|
||||
4. Preserve important messages
|
||||
|
||||
Usage:
|
||||
python basic_compaction.py
|
||||
"""
|
||||
|
||||
from praisonaiagents.compaction import (
|
||||
ContextCompactor, CompactionConfig, CompactionStrategy, CompactionResult
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("Context Compaction Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# ==========================================================================
|
||||
# Sample Messages
|
||||
# ==========================================================================
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful AI assistant."},
|
||||
{"role": "user", "content": "Hello! Can you help me with Python?"},
|
||||
{"role": "assistant", "content": "Of course! I'd be happy to help you with Python. What would you like to know?"},
|
||||
{"role": "user", "content": "How do I read a file in Python?"},
|
||||
{"role": "assistant", "content": "To read a file in Python, you can use the built-in open() function. Here's an example:\n\nwith open('filename.txt', 'r') as f:\n content = f.read()\n\nThis opens the file, reads its content, and automatically closes it."},
|
||||
{"role": "user", "content": "What about writing to a file?"},
|
||||
{"role": "assistant", "content": "To write to a file, use the 'w' mode:\n\nwith open('filename.txt', 'w') as f:\n f.write('Hello, World!')\n\nUse 'a' mode to append instead of overwriting."},
|
||||
{"role": "user", "content": "Can you explain list comprehensions?"},
|
||||
{"role": "assistant", "content": "List comprehensions are a concise way to create lists. Instead of:\n\nresult = []\nfor x in range(10):\n result.append(x * 2)\n\nYou can write:\n\nresult = [x * 2 for x in range(10)]"},
|
||||
]
|
||||
|
||||
print(f"\n--- Original Messages ---")
|
||||
print(f" Message count: {len(messages)}")
|
||||
|
||||
# ==========================================================================
|
||||
# Basic Compaction
|
||||
# ==========================================================================
|
||||
print("\n--- Basic Compaction (Truncate Strategy) ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=50, # Low limit to trigger compaction
|
||||
strategy=CompactionStrategy.TRUNCATE,
|
||||
preserve_system=True,
|
||||
preserve_recent=2
|
||||
)
|
||||
|
||||
# Check stats before
|
||||
stats = compactor.get_stats(messages)
|
||||
print(f" Total tokens: {stats['total_tokens']}")
|
||||
print(f" Max tokens: {stats['max_tokens']}")
|
||||
print(f" Needs compaction: {stats['needs_compaction']}")
|
||||
|
||||
# Compact
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f"\n After compaction:")
|
||||
print(f" Messages: {len(messages)} -> {len(compacted)}")
|
||||
print(f" Tokens: {result.original_tokens} -> {result.compacted_tokens}")
|
||||
print(f" Saved: {result.tokens_saved} tokens")
|
||||
print(f" Compression: {result.compression_ratio:.1%}")
|
||||
|
||||
# ==========================================================================
|
||||
# Sliding Window Strategy
|
||||
# ==========================================================================
|
||||
print("\n--- Sliding Window Strategy ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=50,
|
||||
strategy=CompactionStrategy.SLIDING,
|
||||
preserve_system=True,
|
||||
preserve_recent=3
|
||||
)
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f" Messages kept: {result.messages_kept}")
|
||||
print(f" Messages removed: {result.messages_removed}")
|
||||
print(f" Strategy: {result.strategy_used.value}")
|
||||
|
||||
# ==========================================================================
|
||||
# Summarize Strategy
|
||||
# ==========================================================================
|
||||
print("\n--- Summarize Strategy ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=50,
|
||||
strategy=CompactionStrategy.SUMMARIZE,
|
||||
preserve_system=True,
|
||||
preserve_recent=2
|
||||
)
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f" Messages kept: {result.messages_kept}")
|
||||
print(f" Strategy: {result.strategy_used.value}")
|
||||
|
||||
# Show what messages remain
|
||||
print(f"\n Remaining message roles:")
|
||||
for msg in compacted:
|
||||
role = msg.get("role", "unknown")
|
||||
content_preview = msg.get("content", "")[:50]
|
||||
print(f" [{role}] {content_preview}...")
|
||||
|
||||
# ==========================================================================
|
||||
# No Compaction Needed
|
||||
# ==========================================================================
|
||||
print("\n--- No Compaction Needed ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=10000, # High limit
|
||||
strategy=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f" Was compacted: {result.was_compacted}")
|
||||
print(f" Messages unchanged: {len(compacted) == len(messages)}")
|
||||
|
||||
# ==========================================================================
|
||||
# Result Serialization
|
||||
# ==========================================================================
|
||||
print("\n--- Result Serialization ---")
|
||||
|
||||
compactor = ContextCompactor(max_tokens=50)
|
||||
_, result = compactor.compact(messages)
|
||||
|
||||
data = result.to_dict()
|
||||
print(f" Original tokens: {data['original_tokens']}")
|
||||
print(f" Compacted tokens: {data['compacted_tokens']}")
|
||||
print(f" Compression ratio: {data['compression_ratio']:.2f}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Hooks Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the hooks system to:
|
||||
1. Log tool calls before execution
|
||||
2. Block dangerous operations
|
||||
3. Add context to agent responses
|
||||
4. Monitor session lifecycle
|
||||
|
||||
Usage:
|
||||
python basic_hooks.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.hooks import (
|
||||
HookRegistry, HookRunner, HookEvent, HookResult,
|
||||
BeforeToolInput, AfterToolInput
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
# Create a hook registry
|
||||
registry = HookRegistry()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 1: Log all tool calls
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def log_tool_calls(event_data: BeforeToolInput) -> HookResult:
|
||||
"""Log every tool call before execution."""
|
||||
print(f"\n📝 [LOG] Tool: {event_data.tool_name}")
|
||||
print(f" Input: {event_data.tool_input}")
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 2: Block dangerous file operations (using function logic)
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def block_delete_operations(event_data: BeforeToolInput) -> HookResult:
|
||||
"""Block any delete or remove operations."""
|
||||
dangerous_keywords = ['delete', 'remove', 'rm', 'unlink']
|
||||
tool_lower = event_data.tool_name.lower()
|
||||
|
||||
if any(kw in tool_lower for kw in dangerous_keywords):
|
||||
print(f"\n🚫 [BLOCKED] Dangerous operation: {event_data.tool_name}")
|
||||
return HookResult.deny(
|
||||
f"Operation '{event_data.tool_name}' is blocked by security policy"
|
||||
)
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 3: Require confirmation for write operations
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def confirm_write_operations(event_data: BeforeToolInput) -> HookResult:
|
||||
"""Request confirmation for write operations."""
|
||||
write_keywords = ['write', 'save', 'create', 'update']
|
||||
tool_lower = event_data.tool_name.lower()
|
||||
|
||||
if any(kw in tool_lower for kw in write_keywords):
|
||||
print(f"\n⚠️ [CONFIRM] Write operation: {event_data.tool_name}")
|
||||
# In a real scenario, you might prompt the user here
|
||||
# For this example, we'll auto-approve
|
||||
return HookResult.allow("Auto-approved for demo")
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 4: Add timing information after tool execution
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.AFTER_TOOL)
|
||||
def log_tool_completion(event_data: AfterToolInput) -> HookResult:
|
||||
"""Log tool completion with timing."""
|
||||
print(f"\n✅ [DONE] Tool: {event_data.tool_name}")
|
||||
print(f" Duration: {event_data.execution_time_ms:.2f}ms")
|
||||
if event_data.tool_error:
|
||||
print(f" Error: {event_data.tool_error}")
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Test the hooks
|
||||
# ==========================================================================
|
||||
print("=" * 60)
|
||||
print("Testing Hooks System")
|
||||
print("=" * 60)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
|
||||
# Test 1: Normal tool call (should be logged and allowed)
|
||||
print("\n--- Test 1: Normal read operation ---")
|
||||
input1 = BeforeToolInput(
|
||||
session_id="test",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="read_file",
|
||||
tool_input={"path": "/tmp/test.txt"}
|
||||
)
|
||||
results = asyncio.run(runner.execute(HookEvent.BEFORE_TOOL, input1))
|
||||
print(f"Result: {'ALLOWED' if not HookRunner.is_blocked(results) else 'BLOCKED'}")
|
||||
|
||||
# Test 2: Delete operation (should be blocked)
|
||||
print("\n--- Test 2: Delete operation ---")
|
||||
input2 = BeforeToolInput(
|
||||
session_id="test",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="delete_file",
|
||||
tool_input={"path": "/important/data.txt"}
|
||||
)
|
||||
results = asyncio.run(runner.execute(HookEvent.BEFORE_TOOL, input2))
|
||||
print(f"Result: {'ALLOWED' if not HookRunner.is_blocked(results) else 'BLOCKED'}")
|
||||
if HookRunner.is_blocked(results):
|
||||
print(f"Reason: {HookRunner.get_blocking_reason(results)}")
|
||||
|
||||
# Test 3: Write operation (should be confirmed)
|
||||
print("\n--- Test 3: Write operation ---")
|
||||
input3 = BeforeToolInput(
|
||||
session_id="test",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="write_file",
|
||||
tool_input={"path": "/tmp/output.txt", "content": "Hello"}
|
||||
)
|
||||
results = asyncio.run(runner.execute(HookEvent.BEFORE_TOOL, input3))
|
||||
print(f"Result: {'ALLOWED' if not HookRunner.is_blocked(results) else 'BLOCKED'}")
|
||||
|
||||
# Show registered hooks
|
||||
print("\n" + "=" * 60)
|
||||
print("Registered Hooks:")
|
||||
print("=" * 60)
|
||||
for event, hooks in registry.list_hooks().items():
|
||||
print(f"\n{event}:")
|
||||
for hook in hooks:
|
||||
print(f" - {hook['name']} (matcher: {hook['matcher'] or '*'})")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,79 @@
|
||||
"""
|
||||
Qdrant Knowledge Store Example for PraisonAI.
|
||||
|
||||
Demonstrates vector insert + search through the knowledge store.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name praison-qdrant -p 6333:6333 qdrant/qdrant
|
||||
|
||||
Run:
|
||||
python knowledge_qdrant.py
|
||||
|
||||
Expected Output:
|
||||
=== Qdrant Knowledge Store Demo ===
|
||||
Created collection: demo_knowledge
|
||||
Inserted 3 documents
|
||||
Search results: 2 documents found
|
||||
"""
|
||||
|
||||
from praisonai.persistence.factory import create_knowledge_store
|
||||
from praisonai.persistence.knowledge.base import KnowledgeDocument
|
||||
import random
|
||||
|
||||
# Create Qdrant knowledge store
|
||||
store = create_knowledge_store("qdrant", url="http://localhost:6333")
|
||||
|
||||
collection_name = "demo_knowledge"
|
||||
dimension = 384 # Embedding dimension
|
||||
|
||||
# Create collection
|
||||
print("=== Qdrant Knowledge Store Demo ===")
|
||||
if store.collection_exists(collection_name):
|
||||
store.delete_collection(collection_name)
|
||||
|
||||
store.create_collection(collection_name, dimension=dimension)
|
||||
print(f"Created collection: {collection_name}")
|
||||
|
||||
# Create sample documents with embeddings
|
||||
random.seed(42)
|
||||
documents = [
|
||||
KnowledgeDocument(
|
||||
id="doc-1",
|
||||
content="Python is a versatile programming language used for web development, data science, and AI.",
|
||||
embedding=[random.random() for _ in range(dimension)],
|
||||
metadata={"category": "programming", "language": "python"}
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc-2",
|
||||
content="Machine learning enables computers to learn patterns from data without explicit programming.",
|
||||
embedding=[random.random() for _ in range(dimension)],
|
||||
metadata={"category": "ai", "topic": "ml"}
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc-3",
|
||||
content="PostgreSQL is a powerful open-source relational database system.",
|
||||
embedding=[random.random() for _ in range(dimension)],
|
||||
metadata={"category": "database", "type": "sql"}
|
||||
),
|
||||
]
|
||||
|
||||
# Insert documents
|
||||
ids = store.insert(collection_name, documents)
|
||||
print(f"Inserted {len(ids)} documents")
|
||||
|
||||
# Search for similar documents
|
||||
query_embedding = [random.random() for _ in range(dimension)]
|
||||
results = store.search(collection_name, query_embedding, limit=2)
|
||||
print(f"Search results: {len(results)} documents found")
|
||||
|
||||
for doc in results:
|
||||
print(f" - {doc.id}: {doc.content[:50]}...")
|
||||
|
||||
# Cleanup
|
||||
store.delete_collection(collection_name)
|
||||
store.close()
|
||||
|
||||
print("\n=== Demo Complete ===")
|
||||
@@ -0,0 +1,43 @@
|
||||
"""
|
||||
Minimal Agent with Database Persistence
|
||||
|
||||
This is the simplest way to add persistence to your agent.
|
||||
Messages, runs, and traces are automatically saved.
|
||||
|
||||
Run:
|
||||
python minimal_agent_db.py
|
||||
|
||||
Expected output:
|
||||
- Agent responds to your question
|
||||
- Data is persisted to PostgreSQL
|
||||
- Session can be resumed later with same session_id
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create database connection (PostgreSQL + Redis)
|
||||
my_db = db(
|
||||
database_url="postgresql://postgres:praison123@localhost:5432/praisonai",
|
||||
state_url="redis://localhost:6379" # Optional: for state/tracing
|
||||
)
|
||||
|
||||
# Create agent with persistence
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant. Be concise.",
|
||||
db=my_db,
|
||||
# session_id is optional - defaults to per-hour ID (YYYYMMDDHH-hash)
|
||||
# Set explicitly for session continuity across runs:
|
||||
session_id="my-persistent-session"
|
||||
)
|
||||
|
||||
# Chat - messages are automatically persisted
|
||||
response = agent.chat("What is the capital of France?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Verify persistence
|
||||
data = my_db.export_session("my-persistent-session")
|
||||
print(f"\nMessages stored: {len(data.get('messages', []))}")
|
||||
|
||||
# Clean up
|
||||
my_db.close()
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
PostgreSQL with Runs and Traces
|
||||
|
||||
Demonstrates full persistence with run tracking and tracing.
|
||||
|
||||
Run:
|
||||
python postgres_runs_traces.py
|
||||
|
||||
Expected output:
|
||||
- Agent responds with tool usage
|
||||
- Runs and traces are persisted
|
||||
- Export shows full session data
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
def get_weather(city: str) -> str:
|
||||
"""Get weather for a city."""
|
||||
return f"Weather in {city}: Sunny, 72°F"
|
||||
|
||||
# Create database with PostgreSQL + Redis
|
||||
my_db = db(
|
||||
database_url="postgresql://postgres:praison123@localhost:5432/praisonai",
|
||||
state_url="redis://localhost:6379"
|
||||
)
|
||||
|
||||
# Create agent with tools
|
||||
agent = Agent(
|
||||
name="WeatherBot",
|
||||
instructions="You help with weather. Use the get_weather tool when asked.",
|
||||
tools=[get_weather],
|
||||
db=my_db,
|
||||
session_id="weather-session-001"
|
||||
)
|
||||
|
||||
# First chat - triggers run tracking
|
||||
print("=== First Question ===")
|
||||
response1 = agent.chat("What's the weather in Tokyo?")
|
||||
print(f"Response: {response1}\n")
|
||||
|
||||
# Second chat - same session, new run
|
||||
print("=== Second Question ===")
|
||||
response2 = agent.chat("How about New York?")
|
||||
print(f"Response: {response2}\n")
|
||||
|
||||
# Export session data
|
||||
print("=== Session Data ===")
|
||||
data = my_db.export_session("weather-session-001")
|
||||
print(f"Session ID: {data.get('session_id')}")
|
||||
print(f"Messages: {len(data.get('messages', []))}")
|
||||
print(f"Agent: {data.get('agent_name')}")
|
||||
|
||||
# Show messages
|
||||
print("\n=== Messages ===")
|
||||
for msg in data.get('messages', []):
|
||||
role = msg.get('role', 'unknown')
|
||||
content = msg.get('content', '')[:60]
|
||||
print(f" [{role}] {content}...")
|
||||
|
||||
my_db.close()
|
||||
print("\n✅ Done")
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
Redis State Store
|
||||
|
||||
Use Redis for state management, tracing, and caching.
|
||||
|
||||
Run:
|
||||
python redis_state.py
|
||||
|
||||
Expected output:
|
||||
- State operations (set/get/delete)
|
||||
- Tracing operations
|
||||
"""
|
||||
|
||||
from praisonaiagents import db
|
||||
|
||||
# Create Redis-only database (for state/tracing)
|
||||
redis_db = db.RedisDB(host="localhost", port=6379)
|
||||
|
||||
print("=== Redis State Store ===")
|
||||
|
||||
# Test state operations
|
||||
print("\n1. State Operations:")
|
||||
redis_db._init_stores() # Initialize stores
|
||||
|
||||
# Set a value
|
||||
redis_db._state_store.set("user:123:preferences", {
|
||||
"theme": "dark",
|
||||
"language": "en"
|
||||
})
|
||||
print(" Set user preferences")
|
||||
|
||||
# Get the value
|
||||
prefs = redis_db._state_store.get("user:123:preferences")
|
||||
print(f" Got preferences: {prefs}")
|
||||
|
||||
# Set with TTL (expires in 60 seconds)
|
||||
redis_db._state_store.set("temp:token", "abc123", ttl=60)
|
||||
print(" Set temporary token (60s TTL)")
|
||||
|
||||
# Delete
|
||||
redis_db._state_store.delete("user:123:preferences")
|
||||
print(" Deleted preferences")
|
||||
|
||||
# Test tracing
|
||||
print("\n2. Tracing Operations:")
|
||||
|
||||
# Start a trace
|
||||
trace_id = "trace-example-001"
|
||||
redis_db.on_trace_start(
|
||||
trace_id=trace_id,
|
||||
session_id="example-session",
|
||||
agent_name="ExampleAgent",
|
||||
metadata={"purpose": "demo"}
|
||||
)
|
||||
print(f" Started trace: {trace_id}")
|
||||
|
||||
# Add spans
|
||||
span_id = "span-llm-001"
|
||||
redis_db.on_span_start(
|
||||
span_id=span_id,
|
||||
trace_id=trace_id,
|
||||
name="llm_call",
|
||||
attributes={"model": "gpt-4", "temperature": 0.7}
|
||||
)
|
||||
print(f" Started span: {span_id}")
|
||||
|
||||
redis_db.on_span_end(
|
||||
span_id=span_id,
|
||||
status="ok",
|
||||
attributes={"tokens": 150, "latency_ms": 450}
|
||||
)
|
||||
print(f" Ended span: {span_id}")
|
||||
|
||||
# End trace
|
||||
redis_db.on_trace_end(
|
||||
trace_id=trace_id,
|
||||
status="ok",
|
||||
metadata={"total_spans": 1}
|
||||
)
|
||||
print(f" Ended trace: {trace_id}")
|
||||
|
||||
redis_db.close()
|
||||
print("\n✅ Done")
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
Session Resume Example
|
||||
|
||||
Demonstrates how to resume a conversation from a previous session.
|
||||
|
||||
Run:
|
||||
python session_resume.py
|
||||
|
||||
Expected output:
|
||||
- First run: Creates new session
|
||||
- Second run: Resumes with history
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create database
|
||||
my_db = db(
|
||||
database_url="postgresql://postgres:praison123@localhost:5432/praisonai",
|
||||
state_url="redis://localhost:6379"
|
||||
)
|
||||
|
||||
SESSION_ID = "resume-demo-session"
|
||||
|
||||
# Create agent with explicit session_id
|
||||
agent = Agent(
|
||||
name="MemoryBot",
|
||||
instructions="You remember our conversation. Be helpful and concise.",
|
||||
db=my_db,
|
||||
session_id=SESSION_ID
|
||||
)
|
||||
|
||||
# Check if this is a new or resumed session
|
||||
data = my_db.export_session(SESSION_ID)
|
||||
existing_messages = len(data.get('messages', []))
|
||||
|
||||
if existing_messages == 0:
|
||||
print("=== NEW SESSION ===")
|
||||
print("This is a new conversation.\n")
|
||||
|
||||
response = agent.chat("My name is Alice and I love Python programming.")
|
||||
print(f"Response: {response}\n")
|
||||
|
||||
response = agent.chat("What's my favorite programming language?")
|
||||
print(f"Response: {response}\n")
|
||||
else:
|
||||
print(f"=== RESUMED SESSION ({existing_messages} previous messages) ===")
|
||||
print("Continuing previous conversation.\n")
|
||||
|
||||
# Show previous messages
|
||||
print("Previous messages:")
|
||||
for msg in data.get('messages', [])[:4]:
|
||||
role = msg.get('role', 'unknown')
|
||||
content = msg.get('content', '')[:60]
|
||||
print(f" [{role}] {content}...")
|
||||
print()
|
||||
|
||||
# Continue conversation
|
||||
response = agent.chat("What's my name and what do I like?")
|
||||
print(f"Response: {response}\n")
|
||||
|
||||
# Show final state
|
||||
final_data = my_db.export_session(SESSION_ID)
|
||||
print(f"Total messages now: {len(final_data.get('messages', []))}")
|
||||
|
||||
my_db.close()
|
||||
print("✅ Done")
|
||||
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
Simple Database Persistence Example for PraisonAI Agents.
|
||||
|
||||
This example shows how to enable automatic conversation persistence
|
||||
with just 2 lines of code.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name praison-postgres -p 5432:5432 \
|
||||
-e POSTGRES_PASSWORD=praison123 \
|
||||
-e POSTGRES_DB=praisonai \
|
||||
postgres:16
|
||||
|
||||
Environment:
|
||||
export OPENAI_API_KEY=your-key
|
||||
|
||||
Run:
|
||||
python simple_agent_postgres.py
|
||||
|
||||
Expected Output:
|
||||
=== First Run ===
|
||||
Agent: Nice to meet you, Alice! Python is a great choice...
|
||||
Session ID: demo-session-001
|
||||
Messages persisted: 2
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonai.db import PostgresDB
|
||||
|
||||
# Create database adapter
|
||||
db = PostgresDB(
|
||||
host="localhost",
|
||||
port=5432,
|
||||
database="praisonai",
|
||||
user="postgres",
|
||||
password="praison123"
|
||||
)
|
||||
|
||||
# Create agent with persistence enabled
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant. Keep responses brief.",
|
||||
db=db,
|
||||
session_id="demo-session-001", # Same session_id = resume conversation
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Chat with persistence
|
||||
print("=== First Run ===")
|
||||
response = agent.chat("My name is Alice and I love Python programming.")
|
||||
print(f"Agent: {response}\n")
|
||||
|
||||
# Show session info
|
||||
print(f"Session ID: {agent.session_id}")
|
||||
print(f"Messages in history: {len(agent.chat_history)}")
|
||||
|
||||
# Close DB connection
|
||||
db.close()
|
||||
|
||||
print("\n=== Session persisted! ===")
|
||||
print("Run again with same session_id to resume.")
|
||||
@@ -0,0 +1,41 @@
|
||||
"""
|
||||
SQLite Local Database
|
||||
|
||||
Use SQLite for simple local persistence without external services.
|
||||
|
||||
Run:
|
||||
python sqlite_local.py
|
||||
|
||||
Expected output:
|
||||
- Agent responds
|
||||
- Data persisted to local SQLite file
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
# Create SQLite database (local file)
|
||||
db_path = os.path.join(tempfile.gettempdir(), "my_agent.db")
|
||||
my_db = db.SQLiteDB(path=db_path)
|
||||
|
||||
print(f"Database: {db_path}")
|
||||
|
||||
# Create agent
|
||||
agent = Agent(
|
||||
name="LocalBot",
|
||||
instructions="You are a helpful assistant.",
|
||||
db=my_db,
|
||||
session_id="local-session"
|
||||
)
|
||||
|
||||
# Chat
|
||||
response = agent.chat("Hello! What can you help me with?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Verify
|
||||
data = my_db.export_session("local-session")
|
||||
print(f"\nMessages stored: {len(data.get('messages', []))}")
|
||||
|
||||
my_db.close()
|
||||
print("✅ Done")
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
Redis State Store Example for PraisonAI.
|
||||
|
||||
Demonstrates set/get/delete + TTL operations.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name praison-redis -p 6379:6379 redis:7
|
||||
|
||||
Run:
|
||||
python state_redis.py
|
||||
|
||||
Expected Output:
|
||||
=== Redis State Store Demo ===
|
||||
Set key: session:user123:preferences
|
||||
Get value: {'theme': 'dark', 'language': 'en'}
|
||||
TTL remaining: 3600 seconds
|
||||
Hash operations: {'name': 'Alice', 'role': 'admin'}
|
||||
"""
|
||||
|
||||
from praisonai.persistence.factory import create_state_store
|
||||
|
||||
# Create Redis state store
|
||||
store = create_state_store("redis", url="redis://localhost:6379")
|
||||
|
||||
print("=== Redis State Store Demo ===")
|
||||
|
||||
# Basic set/get
|
||||
key = "session:user123:preferences"
|
||||
value = {"theme": "dark", "language": "en", "notifications": True}
|
||||
|
||||
store.set(key, value)
|
||||
print(f"Set key: {key}")
|
||||
|
||||
retrieved = store.get(key)
|
||||
print(f"Get value: {retrieved}")
|
||||
|
||||
# Set with TTL (expires in 1 hour)
|
||||
ttl_key = "session:user123:token"
|
||||
store.set(ttl_key, "abc123xyz", ttl=3600)
|
||||
remaining = store.ttl(ttl_key)
|
||||
print(f"TTL remaining: {remaining} seconds")
|
||||
|
||||
# Hash operations (for structured data)
|
||||
hash_key = "user:123:profile"
|
||||
store.hset(hash_key, "name", "Alice")
|
||||
store.hset(hash_key, "role", "admin")
|
||||
store.hset(hash_key, "last_login", "2024-12-24")
|
||||
|
||||
profile = store.hgetall(hash_key)
|
||||
print(f"Hash operations: {profile}")
|
||||
|
||||
# Check existence
|
||||
exists = store.exists(key)
|
||||
print(f"Key exists: {exists}")
|
||||
|
||||
# Delete keys
|
||||
store.delete(key)
|
||||
store.delete(ttl_key)
|
||||
store.delete(hash_key)
|
||||
print("Cleaned up test keys")
|
||||
|
||||
store.close()
|
||||
|
||||
print("\n=== Demo Complete ===")
|
||||
@@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Policy Engine Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the policy engine to:
|
||||
1. Create policies with rules
|
||||
2. Check if actions are allowed
|
||||
3. Block dangerous operations
|
||||
4. Use convenience policy creators
|
||||
|
||||
Usage:
|
||||
python basic_policy.py
|
||||
"""
|
||||
|
||||
from praisonaiagents.policy import (
|
||||
PolicyEngine, Policy, PolicyRule, PolicyAction, PolicyConfig,
|
||||
create_deny_tools_policy, create_read_only_policy
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("Policy Engine Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a policy engine
|
||||
engine = PolicyEngine()
|
||||
|
||||
# ==========================================================================
|
||||
# Policy 1: Block dangerous tools
|
||||
# ==========================================================================
|
||||
print("\n--- Creating Policies ---")
|
||||
|
||||
dangerous_tools_policy = Policy(
|
||||
name="block_dangerous",
|
||||
description="Block dangerous file operations",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="Delete operations are blocked",
|
||||
name="deny_delete"
|
||||
),
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:rm_*",
|
||||
reason="Remove operations are blocked",
|
||||
name="deny_rm"
|
||||
),
|
||||
],
|
||||
priority=100
|
||||
)
|
||||
engine.add_policy(dangerous_tools_policy)
|
||||
print(f"✅ Added policy: {dangerous_tools_policy.name}")
|
||||
|
||||
# ==========================================================================
|
||||
# Policy 2: Allow read operations
|
||||
# ==========================================================================
|
||||
read_policy = Policy(
|
||||
name="allow_read",
|
||||
description="Allow read operations",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.ALLOW,
|
||||
resource="tool:read_*",
|
||||
name="allow_read"
|
||||
),
|
||||
PolicyRule(
|
||||
action=PolicyAction.ALLOW,
|
||||
resource="tool:list_*",
|
||||
name="allow_list"
|
||||
),
|
||||
]
|
||||
)
|
||||
engine.add_policy(read_policy)
|
||||
print(f"✅ Added policy: {read_policy.name}")
|
||||
|
||||
# ==========================================================================
|
||||
# Test policy checks
|
||||
# ==========================================================================
|
||||
print("\n--- Testing Policy Checks ---")
|
||||
|
||||
test_cases = [
|
||||
("tool:read_file", "Reading a file"),
|
||||
("tool:list_directory", "Listing directory"),
|
||||
("tool:delete_file", "Deleting a file"),
|
||||
("tool:rm_directory", "Removing directory"),
|
||||
("tool:write_file", "Writing a file"),
|
||||
]
|
||||
|
||||
for resource, description in test_cases:
|
||||
result = engine.check(resource, {})
|
||||
status = "✅ ALLOWED" if result.allowed else "🚫 DENIED"
|
||||
reason = f" ({result.reason})" if result.reason else ""
|
||||
print(f" {status}: {description}{reason}")
|
||||
|
||||
# ==========================================================================
|
||||
# Using convenience functions
|
||||
# ==========================================================================
|
||||
print("\n--- Using Convenience Functions ---")
|
||||
|
||||
# Create a read-only policy
|
||||
read_only = create_read_only_policy()
|
||||
print(f"✅ Created read-only policy with {len(read_only.rules)} rules")
|
||||
|
||||
# Create a deny tools policy
|
||||
deny_tools = create_deny_tools_policy(
|
||||
["execute_*", "shell_*", "system_*"],
|
||||
reason="System commands are blocked"
|
||||
)
|
||||
print(f"✅ Created deny-tools policy with {len(deny_tools.rules)} rules")
|
||||
|
||||
# ==========================================================================
|
||||
# Strict mode example
|
||||
# ==========================================================================
|
||||
print("\n--- Strict Mode Example ---")
|
||||
|
||||
strict_engine = PolicyEngine(PolicyConfig(strict_mode=True))
|
||||
|
||||
# In strict mode, unknown resources are denied
|
||||
result = strict_engine.check("tool:unknown_operation", {})
|
||||
print(f" Unknown operation in strict mode: {'DENIED' if not result.allowed else 'ALLOWED'}")
|
||||
|
||||
# ==========================================================================
|
||||
# Policy serialization
|
||||
# ==========================================================================
|
||||
print("\n--- Policy Serialization ---")
|
||||
|
||||
policy_dict = dangerous_tools_policy.to_dict()
|
||||
print(f" Serialized policy: {policy_dict['name']} with {len(policy_dict['rules'])} rules")
|
||||
|
||||
# Recreate from dict
|
||||
restored = Policy.from_dict(policy_dict)
|
||||
print(f" Restored policy: {restored.name}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,94 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example: Creating Skills with PraisonAI CLI
|
||||
|
||||
This example demonstrates how to create Agent Skills using the CLI.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import tempfile
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def example_create_skill_with_template():
|
||||
"""Create a skill using template mode (no API key required)."""
|
||||
print("=" * 60)
|
||||
print("Example 1: Create skill with template")
|
||||
print("=" * 60)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
result = subprocess.run([
|
||||
"praisonai", "skills", "create",
|
||||
"--name", "my-template-skill",
|
||||
"--description", "A skill for processing text files",
|
||||
"--author", "example-team",
|
||||
"--license", "MIT",
|
||||
"--template",
|
||||
"--output-dir", tmpdir
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode == 0:
|
||||
skill_md = Path(tmpdir) / "my-template-skill" / "SKILL.md"
|
||||
print(f"\nGenerated SKILL.md:\n{skill_md.read_text()[:500]}...")
|
||||
else:
|
||||
print(f"Error: {result.stderr}")
|
||||
|
||||
|
||||
def example_create_skill_with_ai():
|
||||
"""Create a skill using AI generation (requires API key)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Create skill with AI generation")
|
||||
print("=" * 60)
|
||||
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Skipping: OPENAI_API_KEY not set")
|
||||
return
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
result = subprocess.run([
|
||||
"praisonai", "skills", "create",
|
||||
"--name", "csv-analyzer",
|
||||
"--description", "Analyze CSV files and generate statistical reports",
|
||||
"--output-dir", tmpdir
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode == 0:
|
||||
skill_md = Path(tmpdir) / "csv-analyzer" / "SKILL.md"
|
||||
print(f"\nGenerated SKILL.md:\n{skill_md.read_text()[:800]}...")
|
||||
|
||||
# Check if script was generated
|
||||
script_path = Path(tmpdir) / "csv-analyzer" / "scripts" / "skill.py"
|
||||
if script_path.exists():
|
||||
print(f"\nGenerated skill.py:\n{script_path.read_text()[:500]}...")
|
||||
|
||||
|
||||
def example_create_skill_with_script():
|
||||
"""Create a skill with auto-generated Python script."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 3: Create skill with script")
|
||||
print("=" * 60)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
result = subprocess.run([
|
||||
"praisonai", "skills", "create",
|
||||
"--name", "data-processor",
|
||||
"--description", "Process and transform data files",
|
||||
"--template",
|
||||
"--script",
|
||||
"--output-dir", tmpdir
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode == 0:
|
||||
script_path = Path(tmpdir) / "data-processor" / "scripts" / "skill.py"
|
||||
if script_path.exists():
|
||||
print(f"\nGenerated skill.py:\n{script_path.read_text()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
example_create_skill_with_template()
|
||||
example_create_skill_with_ai()
|
||||
example_create_skill_with_script()
|
||||
@@ -0,0 +1,137 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example: Multi-Agent CLI Usage
|
||||
|
||||
This example demonstrates how to use the praisonai agents command
|
||||
to define and run multiple agents from the command line.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
|
||||
|
||||
def example_single_agent():
|
||||
"""Run a single agent with a simple task."""
|
||||
print("=" * 60)
|
||||
print("Example 1: Single Agent")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "assistant:Helper",
|
||||
"--task", "What is the capital of France?"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode != 0:
|
||||
print(f"Error: {result.stderr}")
|
||||
|
||||
|
||||
def example_agent_with_tools():
|
||||
"""Run an agent with tools."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Agent with Tools")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "researcher:Research Analyst:internet_search",
|
||||
"--task", "Find the latest news about renewable energy"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_multiple_agents_sequential():
|
||||
"""Run multiple agents in sequential mode."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 3: Multiple Agents (Sequential)")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "researcher:Research Analyst:internet_search",
|
||||
"--agent", "writer:Content Writer:write_file",
|
||||
"--task", "Research AI trends and write a brief summary"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_multiple_agents_parallel():
|
||||
"""Run multiple agents in parallel mode."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 4: Multiple Agents (Parallel)")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "analyst1:Market Analyst",
|
||||
"--agent", "analyst2:Tech Analyst",
|
||||
"--process", "parallel",
|
||||
"--task", "Analyze the current state of AI industry"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_with_custom_llm():
|
||||
"""Run agents with a custom LLM model."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 5: Custom LLM Model")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "coder:Developer:execute_command",
|
||||
"--llm", "gpt-4o-mini",
|
||||
"--task", "Write a simple hello world Python script"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_with_instructions():
|
||||
"""Run agents with additional instructions."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 6: With Additional Instructions")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "analyst:Data Analyst",
|
||||
"--instructions", "Be concise and provide only key insights",
|
||||
"--task", "Explain the benefits of using AI in healthcare"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_list_tools():
|
||||
"""List available tools for agents."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 7: List Available Tools")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "list"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Check for API key
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Warning: OPENAI_API_KEY not set. Some examples may fail.")
|
||||
print()
|
||||
|
||||
# Run examples
|
||||
example_list_tools()
|
||||
example_single_agent()
|
||||
# Uncomment to run more examples:
|
||||
# example_agent_with_tools()
|
||||
# example_multiple_agents_sequential()
|
||||
# example_multiple_agents_parallel()
|
||||
# example_with_custom_llm()
|
||||
# example_with_instructions()
|
||||
@@ -0,0 +1,137 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Thinking Budgets Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use thinking budgets to:
|
||||
1. Configure token budgets for extended thinking
|
||||
2. Use predefined budget levels
|
||||
3. Track thinking usage
|
||||
4. Adapt budgets based on complexity
|
||||
|
||||
Usage:
|
||||
python basic_thinking.py
|
||||
"""
|
||||
|
||||
from praisonaiagents.thinking import (
|
||||
ThinkingBudget, ThinkingConfig, ThinkingUsage, ThinkingTracker
|
||||
)
|
||||
from praisonaiagents.thinking.budget import BudgetLevel
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("Thinking Budgets Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# ==========================================================================
|
||||
# Budget Levels
|
||||
# ==========================================================================
|
||||
print("\n--- Budget Levels ---")
|
||||
|
||||
levels = [
|
||||
("Minimal", ThinkingBudget.minimal()),
|
||||
("Low", ThinkingBudget.low()),
|
||||
("Medium", ThinkingBudget.medium()),
|
||||
("High", ThinkingBudget.high()),
|
||||
("Maximum", ThinkingBudget.maximum()),
|
||||
]
|
||||
|
||||
for name, budget in levels:
|
||||
print(f" {name}: {budget.max_tokens:,} tokens")
|
||||
|
||||
# ==========================================================================
|
||||
# Custom Budget
|
||||
# ==========================================================================
|
||||
print("\n--- Custom Budget ---")
|
||||
|
||||
custom_budget = ThinkingBudget(
|
||||
max_tokens=12000,
|
||||
min_tokens=2000,
|
||||
adaptive=True,
|
||||
max_time_seconds=120.0
|
||||
)
|
||||
print(f" Max tokens: {custom_budget.max_tokens:,}")
|
||||
print(f" Min tokens: {custom_budget.min_tokens:,}")
|
||||
print(f" Adaptive: {custom_budget.adaptive}")
|
||||
print(f" Time limit: {custom_budget.max_time_seconds}s")
|
||||
|
||||
# ==========================================================================
|
||||
# Complexity-Based Scaling
|
||||
# ==========================================================================
|
||||
print("\n--- Complexity-Based Token Allocation ---")
|
||||
|
||||
budget = ThinkingBudget(
|
||||
max_tokens=16000,
|
||||
min_tokens=2000,
|
||||
adaptive=True
|
||||
)
|
||||
|
||||
complexities = [0.0, 0.25, 0.5, 0.75, 1.0]
|
||||
for complexity in complexities:
|
||||
tokens = budget.get_tokens_for_complexity(complexity)
|
||||
bar = "█" * int(tokens / 1000)
|
||||
print(f" Complexity {complexity:.2f}: {tokens:,} tokens {bar}")
|
||||
|
||||
# ==========================================================================
|
||||
# Usage Tracking
|
||||
# ==========================================================================
|
||||
print("\n--- Usage Tracking ---")
|
||||
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=5000,
|
||||
time_seconds=30.0,
|
||||
budget_time=60.0
|
||||
)
|
||||
|
||||
print(f" Tokens used: {usage.tokens_used:,} / {usage.budget_tokens:,}")
|
||||
print(f" Tokens remaining: {usage.tokens_remaining:,}")
|
||||
print(f" Utilization: {usage.token_utilization:.1%}")
|
||||
print(f" Over budget: {usage.is_over_budget}")
|
||||
print(f" Time remaining: {usage.time_remaining}s")
|
||||
|
||||
# ==========================================================================
|
||||
# Session Tracking
|
||||
# ==========================================================================
|
||||
print("\n--- Session Tracking ---")
|
||||
|
||||
tracker = ThinkingTracker()
|
||||
|
||||
# Simulate multiple thinking sessions
|
||||
sessions = [
|
||||
(8000, 4000, 20.0),
|
||||
(8000, 6000, 35.0),
|
||||
(16000, 12000, 60.0),
|
||||
(8000, 9000, 45.0), # Over budget
|
||||
]
|
||||
|
||||
for budget_tokens, tokens_used, time_seconds in sessions:
|
||||
session = tracker.start_session(budget_tokens=budget_tokens)
|
||||
tracker.end_session(session, tokens_used=tokens_used, time_seconds=time_seconds)
|
||||
|
||||
summary = tracker.get_summary()
|
||||
print(f" Total sessions: {summary['session_count']}")
|
||||
print(f" Total tokens used: {summary['total_tokens_used']:,}")
|
||||
print(f" Average tokens/session: {summary['average_tokens_per_session']:,.0f}")
|
||||
print(f" Average utilization: {summary['average_utilization']:.1%}")
|
||||
print(f" Over-budget sessions: {summary['over_budget_count']}")
|
||||
|
||||
# ==========================================================================
|
||||
# Serialization
|
||||
# ==========================================================================
|
||||
print("\n--- Serialization ---")
|
||||
|
||||
budget = ThinkingBudget.from_level(BudgetLevel.HIGH)
|
||||
data = budget.to_dict()
|
||||
print(f" Serialized: level={data['level']}, max_tokens={data['max_tokens']}")
|
||||
|
||||
restored = ThinkingBudget.from_dict(data)
|
||||
print(f" Restored: {restored.max_tokens:,} tokens")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
"""ArXiv Tool Example.
|
||||
|
||||
This example demonstrates how to use the ArXiv tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import ArxivTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize ArXiv tool (no API key needed)
|
||||
arxiv = ArxivTool()
|
||||
|
||||
# Example 1: Search ArXiv
|
||||
print("=" * 60)
|
||||
print("Example 1: Search ArXiv")
|
||||
print("=" * 60)
|
||||
|
||||
results = arxiv.search("transformer neural networks", max_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} papers:")
|
||||
for paper in results:
|
||||
print(f"\n Title: {paper.get('title', 'N/A')[:60]}...")
|
||||
authors = paper.get('authors', [])
|
||||
if authors:
|
||||
print(f" Authors: {', '.join(authors[:2])}...")
|
||||
print(f" Published: {paper.get('published', 'N/A')}")
|
||||
print(f" PDF: {paper.get('pdf_url', 'N/A')}")
|
||||
|
||||
print("\n✅ ArXiv tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
"""Calculator Tool Example.
|
||||
|
||||
This example demonstrates how to use the Calculator tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import CalculatorTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Calculator tool
|
||||
calc = CalculatorTool()
|
||||
|
||||
# Example 1: Basic arithmetic
|
||||
print("=" * 60)
|
||||
print("Example 1: Basic Arithmetic")
|
||||
print("=" * 60)
|
||||
|
||||
expressions = [
|
||||
"2 + 2",
|
||||
"10 * 5",
|
||||
"100 / 4",
|
||||
"2 ** 10",
|
||||
"(10 + 5) * 2 - 3"
|
||||
]
|
||||
|
||||
for expr in expressions:
|
||||
result = calc.evaluate(expr)
|
||||
print(f" {expr} = {result.get('result', result)}")
|
||||
|
||||
# Example 2: Percentage calculation
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Percentage Calculation")
|
||||
print("=" * 60)
|
||||
|
||||
result = calc.evaluate("250 * 0.15")
|
||||
print(f" 15% of 250 = {result.get('result', result)}")
|
||||
|
||||
result = calc.evaluate("1000 * (1 + 0.05) ** 3")
|
||||
print(f" $1000 at 5% for 3 years = ${result.get('result', 0):.2f}")
|
||||
|
||||
print("\n✅ Calculator tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+75
@@ -0,0 +1,75 @@
|
||||
"""DuckDuckGo Search Tool Example.
|
||||
|
||||
This example demonstrates how to use the DuckDuckGo search tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import DuckDuckGoTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize DuckDuckGo tool (no API key needed)
|
||||
ddg = DuckDuckGoTool()
|
||||
|
||||
# Example 1: Web search
|
||||
print("=" * 60)
|
||||
print("Example 1: Web Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = ddg.search("Python programming best practices", max_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r['title']}")
|
||||
print(f" URL: {r['url']}")
|
||||
print(f" Snippet: {r['snippet'][:100]}...")
|
||||
print()
|
||||
|
||||
# Example 2: News search
|
||||
print("=" * 60)
|
||||
print("Example 2: News Search")
|
||||
print("=" * 60)
|
||||
|
||||
news = ddg.news("artificial intelligence", max_results=3)
|
||||
|
||||
if news and "error" in news[0]:
|
||||
print(f"Error: {news[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(news)} news articles:")
|
||||
for n in news:
|
||||
print(f" - {n['title']}")
|
||||
print(f" Source: {n.get('source', 'N/A')}")
|
||||
print(f" Date: {n.get('date', 'N/A')}")
|
||||
print()
|
||||
|
||||
# Example 3: Image search
|
||||
print("=" * 60)
|
||||
print("Example 3: Image Search")
|
||||
print("=" * 60)
|
||||
|
||||
images = ddg.images("sunset landscape", max_results=3)
|
||||
|
||||
if images and "error" in images[0]:
|
||||
print(f"Error: {images[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(images)} images:")
|
||||
for img in images:
|
||||
print(f" - {img['title']}")
|
||||
print(f" Image URL: {img['image_url'][:60]}...")
|
||||
print()
|
||||
|
||||
print("✅ DuckDuckGo tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Vendored
+48
@@ -0,0 +1,48 @@
|
||||
"""Exa Search Tool Example.
|
||||
|
||||
This example demonstrates how to use the Exa search tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export EXA_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import ExaTool
|
||||
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("EXA_API_KEY"):
|
||||
print("Error: EXA_API_KEY environment variable not set")
|
||||
print("Get your API key from https://exa.ai/")
|
||||
return
|
||||
|
||||
# Initialize Exa tool
|
||||
exa = ExaTool()
|
||||
|
||||
# Example 1: Neural search
|
||||
print("=" * 60)
|
||||
print("Example 1: Neural Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = exa.search("best practices for Python development", num_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r.get('title', 'N/A')[:50]}...")
|
||||
print(f" URL: {r.get('url', 'N/A')[:60]}...")
|
||||
print(f" Score: {r.get('score', 'N/A')}")
|
||||
print()
|
||||
|
||||
print("✅ Exa tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
"""Firecrawl Tool Example.
|
||||
|
||||
This example demonstrates how to use the Firecrawl tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export FIRECRAWL_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import FirecrawlTool
|
||||
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("FIRECRAWL_API_KEY"):
|
||||
print("Error: FIRECRAWL_API_KEY environment variable not set")
|
||||
print("Get your API key from https://firecrawl.dev/")
|
||||
return
|
||||
|
||||
# Initialize Firecrawl tool
|
||||
firecrawl = FirecrawlTool()
|
||||
|
||||
# Example 1: Scrape a page
|
||||
print("=" * 60)
|
||||
print("Example 1: Scrape a Page")
|
||||
print("=" * 60)
|
||||
|
||||
result = firecrawl.scrape("https://example.com")
|
||||
|
||||
if "error" in result:
|
||||
print(f"Error: {result['error']}")
|
||||
else:
|
||||
print(f"URL: {result.get('url', 'N/A')}")
|
||||
markdown = result.get('markdown', '')
|
||||
print(f"Content preview: {markdown[:300]}...")
|
||||
|
||||
print("\n✅ Firecrawl tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
"""Hacker News Tool Example.
|
||||
|
||||
This example demonstrates how to use the Hacker News tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import HackerNewsTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Hacker News tool
|
||||
hn = HackerNewsTool()
|
||||
|
||||
# Example 1: Get top stories
|
||||
print("=" * 60)
|
||||
print("Example 1: Top Stories")
|
||||
print("=" * 60)
|
||||
|
||||
stories = hn.get_top_stories(limit=5)
|
||||
|
||||
if stories and "error" in stories[0]:
|
||||
print(f"Error: {stories[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(stories)} top stories:")
|
||||
for story in stories:
|
||||
print(f" - {story.get('title', 'N/A')[:50]}...")
|
||||
print(f" Score: {story.get('score', 'N/A')} | Comments: {story.get('descendants', 'N/A')}")
|
||||
print(f" URL: {story.get('url', 'N/A')[:60]}...")
|
||||
print()
|
||||
|
||||
print("✅ Hacker News tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Vendored
+40
@@ -0,0 +1,40 @@
|
||||
"""Jina Reader Tool Example.
|
||||
|
||||
This example demonstrates how to use the Jina Reader tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: Works without API key (with rate limits)
|
||||
"""
|
||||
|
||||
from praisonai_tools import JinaTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Jina tool (works without API key)
|
||||
jina = JinaTool()
|
||||
|
||||
# Example 1: Read a URL
|
||||
print("=" * 60)
|
||||
print("Example 1: Read URL Content")
|
||||
print("=" * 60)
|
||||
|
||||
result = jina.read("https://example.com")
|
||||
|
||||
if "error" in result:
|
||||
print(f"Error: {result['error']}")
|
||||
else:
|
||||
print(f"URL: {result.get('url', 'N/A')}")
|
||||
print(f"Title: {result.get('title', 'N/A')}")
|
||||
content = result.get('content', '')
|
||||
print(f"Content preview: {content[:300]}...")
|
||||
|
||||
print("\n✅ Jina tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+56
@@ -0,0 +1,56 @@
|
||||
"""PostgreSQL Tool Example.
|
||||
|
||||
This example demonstrates how to use the PostgreSQL tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
Docker: docker run -d --name postgres -e POSTGRES_PASSWORD=praison123 -e POSTGRES_DB=praisonai -p 5432:5432 postgres:16
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
from praisonai_tools import PostgresTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize PostgreSQL tool
|
||||
pg = PostgresTool(
|
||||
host="localhost",
|
||||
port=5432,
|
||||
database="praisonai",
|
||||
user="postgres",
|
||||
password="praison123"
|
||||
)
|
||||
|
||||
# Example 1: List tables
|
||||
print("=" * 60)
|
||||
print("Example 1: List Tables")
|
||||
print("=" * 60)
|
||||
|
||||
tables = pg.list_tables()
|
||||
|
||||
if isinstance(tables, dict) and "error" in tables:
|
||||
print(f"Error: {tables['error']}")
|
||||
else:
|
||||
print(f"Found {len(tables)} tables:")
|
||||
for t in tables[:10]:
|
||||
print(f" - {t.get('table_name', t)}")
|
||||
|
||||
# Example 2: Query
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Execute Query")
|
||||
print("=" * 60)
|
||||
|
||||
result = pg.query("SELECT version()")
|
||||
|
||||
if isinstance(result, dict) and "error" in result:
|
||||
print(f"Error: {result['error']}")
|
||||
else:
|
||||
print(f"PostgreSQL version: {result}")
|
||||
|
||||
print("\n✅ PostgreSQL tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
"""Serper Tool Example.
|
||||
|
||||
This example demonstrates how to use the Serper (Google Search) tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export SERPER_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import SerperTool
|
||||
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("SERPER_API_KEY"):
|
||||
print("Error: SERPER_API_KEY environment variable not set")
|
||||
print("Get your API key from https://serper.dev/")
|
||||
return
|
||||
|
||||
# Initialize Serper tool
|
||||
serper = SerperTool()
|
||||
|
||||
# Example 1: Web search
|
||||
print("=" * 60)
|
||||
print("Example 1: Google Web Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = serper.search("Python programming tutorials", max_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r.get('title', 'N/A')[:50]}...")
|
||||
print(f" URL: {r.get('link', 'N/A')[:60]}...")
|
||||
print()
|
||||
|
||||
print("✅ Serper tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
"""Tavily Search Tool Example.
|
||||
|
||||
This example demonstrates how to use the Tavily search tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export TAVILY_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import TavilyTool
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("TAVILY_API_KEY"):
|
||||
print("Error: TAVILY_API_KEY environment variable not set")
|
||||
print("Get your API key from https://tavily.com/")
|
||||
return
|
||||
|
||||
# Initialize Tavily tool
|
||||
tavily = TavilyTool()
|
||||
|
||||
# Example 1: Basic search
|
||||
print("=" * 60)
|
||||
print("Example 1: Basic Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = tavily.search("What is quantum computing?", max_results=3)
|
||||
|
||||
if "error" in results:
|
||||
print(f"Error: {results['error']}")
|
||||
else:
|
||||
print(f"Query: {results.get('query')}")
|
||||
if "answer" in results:
|
||||
print(f"\nAI Answer: {results['answer'][:200]}...")
|
||||
print(f"\nResults ({len(results.get('results', []))}):")
|
||||
for r in results.get("results", []):
|
||||
print(f" - {r['title']}")
|
||||
print(f" URL: {r['url']}")
|
||||
print(f" Score: {r.get('score', 'N/A')}")
|
||||
|
||||
# Example 2: Search context for RAG
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Search Context (for RAG)")
|
||||
print("=" * 60)
|
||||
|
||||
context = tavily.search_context("Python best practices 2024")
|
||||
print(f"Context (first 300 chars): {context[:300]}...")
|
||||
|
||||
print("\n✅ Tavily tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
"""Wikipedia Tool Example.
|
||||
|
||||
This example demonstrates how to use the Wikipedia tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import WikipediaTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Wikipedia tool (no API key needed)
|
||||
wiki = WikipediaTool()
|
||||
|
||||
# Example 1: Search Wikipedia
|
||||
print("=" * 60)
|
||||
print("Example 1: Search Wikipedia")
|
||||
print("=" * 60)
|
||||
|
||||
results = wiki.search("Python programming", max_results=5)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r['title']}")
|
||||
|
||||
# Example 2: Get page summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Get Page Summary")
|
||||
print("=" * 60)
|
||||
|
||||
summary = wiki.summary("Artificial intelligence", sentences=3)
|
||||
|
||||
if "error" in summary:
|
||||
print(f"Error: {summary['error']}")
|
||||
else:
|
||||
print(f"Title: {summary['title']}")
|
||||
print(f"Summary: {summary['summary'][:300]}...")
|
||||
|
||||
# Example 3: Get full page
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 3: Get Full Page")
|
||||
print("=" * 60)
|
||||
|
||||
page = wiki.get_page("Machine learning")
|
||||
|
||||
if "error" in page:
|
||||
print(f"Error: {page['error']}")
|
||||
if "options" in page:
|
||||
print(f"Options: {page['options'][:5]}")
|
||||
else:
|
||||
print(f"Title: {page['title']}")
|
||||
print(f"URL: {page['url']}")
|
||||
print(f"Categories: {page['categories'][:3]}")
|
||||
print(f"Content preview: {page['content'][:200]}...")
|
||||
|
||||
print("\n✅ Wikipedia tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -54,6 +54,7 @@ except ImportError:
|
||||
MCP = None
|
||||
from .session import Session
|
||||
from .memory.memory import Memory
|
||||
from .db import db
|
||||
# Workflows - import from dedicated workflows module
|
||||
from .workflows import (
|
||||
Workflow, WorkflowStep, WorkflowContext, StepResult,
|
||||
@@ -257,6 +258,7 @@ __all__ = [
|
||||
'AutoAgents',
|
||||
'Session',
|
||||
'Memory',
|
||||
'db',
|
||||
'display_interaction',
|
||||
'display_self_reflection',
|
||||
'display_instruction',
|
||||
|
||||
@@ -11,7 +11,7 @@ import sys
|
||||
warnings.filterwarnings("ignore", message=".*There is no current event loop.*", category=DeprecationWarning)
|
||||
warnings.filterwarnings("ignore", message=".*Pydantic serializer warnings.*")
|
||||
warnings.filterwarnings("ignore", message=".*PydanticSerializationUnexpectedValue.*")
|
||||
warnings.filterwarnings("ignore", message=".*Expected 9 fields but got.*")
|
||||
warnings.filterwarnings("ignore", message=".*Expected \\d+ fields but got.*")
|
||||
warnings.filterwarnings("ignore", message=".*Expected `StreamingChoices` but got.*")
|
||||
warnings.filterwarnings("ignore", message=".*serialized value may not be as expected.*")
|
||||
warnings.filterwarnings("ignore", message=".*Use 'content=<...>' to upload raw bytes/text content.*")
|
||||
@@ -28,7 +28,7 @@ SUPPRESSED_PATTERNS = [
|
||||
"The `dict` method is deprecated; use `model_dump` instead",
|
||||
"Pydantic serializer warnings",
|
||||
"PydanticSerializationUnexpectedValue",
|
||||
"Expected 9 fields but got 5 for type `Message`",
|
||||
"Expected", # Catches all "Expected N fields but got M" patterns
|
||||
"Expected `StreamingChoices` but got `Choices`",
|
||||
"serialized value may not be as expected",
|
||||
"Mixing V1 models and V2 models",
|
||||
|
||||
@@ -252,7 +252,9 @@ class Agent:
|
||||
history_in_context: Optional[int] = None,
|
||||
auto_save: Optional[str] = None,
|
||||
skills: Optional[List[str]] = None,
|
||||
skills_dirs: Optional[List[str]] = None
|
||||
skills_dirs: Optional[List[str]] = None,
|
||||
db: Optional[Any] = None,
|
||||
session_id: Optional[str] = None
|
||||
):
|
||||
"""Initialize an Agent instance.
|
||||
|
||||
@@ -615,6 +617,11 @@ Your Goal: {self.goal}
|
||||
self._skill_manager = None # Lazy loaded
|
||||
self._skills_initialized = False
|
||||
|
||||
# Database persistence (lazy - no imports until used)
|
||||
self._db = db
|
||||
self._session_id = session_id
|
||||
self._db_initialized = False
|
||||
|
||||
@property
|
||||
def console(self):
|
||||
"""Lazily initialize Rich Console only when needed."""
|
||||
@@ -651,8 +658,45 @@ Your Goal: {self.goal}
|
||||
self._skill_manager.discover(self._skills_dirs, include_defaults=False)
|
||||
|
||||
self._skills_initialized = True
|
||||
|
||||
# Auto-add skill execution tools if not already present
|
||||
self._add_skill_tools()
|
||||
return self._skill_manager
|
||||
|
||||
def _add_skill_tools(self):
|
||||
"""Add tools required for skill execution (read_file, run_skill_script).
|
||||
|
||||
Uses lazy imports from praisonaiagents.tools to avoid performance impact
|
||||
when skills are not used.
|
||||
"""
|
||||
# Check if tools already include required capabilities
|
||||
tool_names = set()
|
||||
for tool in self.tools:
|
||||
if callable(tool) and hasattr(tool, '__name__'):
|
||||
tool_names.add(tool.__name__)
|
||||
elif hasattr(tool, 'name'):
|
||||
tool_names.add(tool.name)
|
||||
|
||||
# Add read_file if not present
|
||||
if 'read_file' not in tool_names:
|
||||
try:
|
||||
from ..tools import read_file
|
||||
self.tools.append(read_file)
|
||||
logging.debug("Added read_file tool for skill support")
|
||||
except ImportError:
|
||||
logging.warning("Could not import read_file tool for skills")
|
||||
|
||||
# Add run_skill_script from skill_tools module
|
||||
if 'run_skill_script' not in tool_names:
|
||||
try:
|
||||
from ..tools.skill_tools import create_skill_tools
|
||||
# Create skill tools with current working directory
|
||||
skill_tools = create_skill_tools()
|
||||
self.tools.append(skill_tools.run_skill_script)
|
||||
logging.debug("Added run_skill_script tool for skill support")
|
||||
except ImportError as e:
|
||||
logging.warning(f"Could not import skill_tools: {e}")
|
||||
|
||||
def get_skills_prompt(self) -> str:
|
||||
"""Get the XML prompt for available skills.
|
||||
|
||||
@@ -1321,6 +1365,13 @@ Your Goal: {self.goal}"""
|
||||
if self.verbose:
|
||||
self._display_memory_info()
|
||||
|
||||
# Add skills prompt if skills are configured
|
||||
if self._skills or self._skills_dirs:
|
||||
skills_prompt = self.get_skills_prompt()
|
||||
if skills_prompt:
|
||||
system_prompt += f"\n\n## Available Skills\n{skills_prompt}"
|
||||
system_prompt += "\n\nWhen a skill is relevant to the task, read its SKILL.md file to get detailed instructions. If the skill has scripts in its scripts/ directory, you can execute them using the execute_code or run_script tool."
|
||||
|
||||
# Add tool usage instructions if tools are available
|
||||
# Use provided tools or fall back to self.tools
|
||||
tools_to_use = tools if tools is not None else self.tools
|
||||
@@ -1913,7 +1964,108 @@ Your Goal: {self.goal}"""
|
||||
# expand=False
|
||||
# )
|
||||
|
||||
def _init_db_session(self):
|
||||
"""Initialize DB session if db adapter is provided (lazy, first chat only)."""
|
||||
if self._db is None or self._db_initialized:
|
||||
return
|
||||
|
||||
# Generate session_id if not provided: default to per-hour ID (YYYYMMDDHH-agentname)
|
||||
if self._session_id is None:
|
||||
import hashlib
|
||||
from datetime import datetime, timezone
|
||||
# Per-hour session ID: YYYYMMDDHH (UTC) + agent name hash for uniqueness
|
||||
hour_str = datetime.now(timezone.utc).strftime("%Y%m%d%H")
|
||||
agent_hash = hashlib.md5(self.name.encode()).hexdigest()[:6]
|
||||
self._session_id = f"{hour_str}-{agent_hash}"
|
||||
|
||||
# Call db adapter's on_agent_start to get previous messages
|
||||
try:
|
||||
history = self._db.on_agent_start(
|
||||
agent_name=self.name,
|
||||
session_id=self._session_id,
|
||||
user_id=self.user_id,
|
||||
metadata={"role": self.role, "goal": self.goal}
|
||||
)
|
||||
|
||||
# Restore chat history from previous session
|
||||
if history:
|
||||
for msg in history:
|
||||
self.chat_history.append({
|
||||
"role": msg.role,
|
||||
"content": msg.content
|
||||
})
|
||||
logging.info(f"Resumed session {self._session_id} with {len(history)} messages")
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to initialize DB session: {e}")
|
||||
|
||||
self._db_initialized = True
|
||||
self._current_run_id = None # Track current run
|
||||
|
||||
def _start_run(self, input_content: str):
|
||||
"""Start a new run (turn) for persistence tracking."""
|
||||
if self._db is None:
|
||||
return
|
||||
|
||||
import uuid
|
||||
self._current_run_id = f"run-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
try:
|
||||
if hasattr(self._db, 'on_run_start'):
|
||||
self._db.on_run_start(
|
||||
session_id=self._session_id,
|
||||
run_id=self._current_run_id,
|
||||
input_content=input_content,
|
||||
metadata={"agent_name": self.name}
|
||||
)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to start run: {e}")
|
||||
|
||||
def _end_run(self, output_content: str, status: str = "completed", metrics: dict = None):
|
||||
"""End the current run (turn)."""
|
||||
if self._db is None or self._current_run_id is None:
|
||||
return
|
||||
|
||||
try:
|
||||
if hasattr(self._db, 'on_run_end'):
|
||||
self._db.on_run_end(
|
||||
session_id=self._session_id,
|
||||
run_id=self._current_run_id,
|
||||
output_content=output_content,
|
||||
status=status,
|
||||
metrics=metrics or {},
|
||||
metadata={"agent_name": self.name}
|
||||
)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to end run: {e}")
|
||||
|
||||
self._current_run_id = None
|
||||
|
||||
def _persist_message(self, role: str, content: str):
|
||||
"""Persist a message to the DB if adapter is provided."""
|
||||
if self._db is None:
|
||||
return
|
||||
|
||||
try:
|
||||
if role == "user":
|
||||
self._db.on_user_message(self._session_id, content)
|
||||
elif role == "assistant":
|
||||
self._db.on_agent_message(self._session_id, content)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to persist message: {e}")
|
||||
|
||||
@property
|
||||
def session_id(self) -> Optional[str]:
|
||||
"""Get the current session ID."""
|
||||
return self._session_id
|
||||
|
||||
def chat(self, prompt, temperature=1.0, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=None, task_name=None, task_description=None, task_id=None):
|
||||
# Initialize DB session on first chat (lazy)
|
||||
self._init_db_session()
|
||||
|
||||
# Start a new run for this chat turn
|
||||
prompt_str = prompt if isinstance(prompt, str) else str(prompt)
|
||||
self._start_run(prompt_str)
|
||||
|
||||
# Reset the final display flag for each new conversation
|
||||
self._final_display_shown = False
|
||||
|
||||
@@ -2010,6 +2162,8 @@ Your Goal: {self.goal}"""
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
# Persist user message to DB
|
||||
self._persist_message("user", normalized_content)
|
||||
|
||||
try:
|
||||
# Pass everything to LLM class
|
||||
@@ -2039,6 +2193,8 @@ Your Goal: {self.goal}"""
|
||||
)
|
||||
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
|
||||
# Log completion time if in debug mode
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
@@ -2083,6 +2239,8 @@ Your Goal: {self.goal}"""
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
# Persist user message to DB (OpenAI path)
|
||||
self._persist_message("user", normalized_content)
|
||||
|
||||
reflection_count = 0
|
||||
start_time = time.time()
|
||||
@@ -2122,6 +2280,8 @@ Your Goal: {self.goal}"""
|
||||
# Add to chat history and return raw response
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
# Apply guardrail validation even for JSON output
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
@@ -2137,6 +2297,8 @@ Your Goal: {self.goal}"""
|
||||
if not self.self_reflect:
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB (non-reflect path)
|
||||
self._persist_message("assistant", response_text)
|
||||
if self.verbose:
|
||||
logging.debug(f"Agent {self.name} final response: {response_text}")
|
||||
# Return only reasoning content if reasoning_steps is True
|
||||
@@ -2225,11 +2387,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
self._end_run(validated_response, "completed", {"duration_ms": (time.time() - start_time) * 1000})
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
self._end_run(None, "error", {"error": str(e)})
|
||||
return None
|
||||
|
||||
# Check if we've hit max reflections
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Background Agents Module for PraisonAI Agents.
|
||||
|
||||
Provides the ability to run agents in the background, allowing:
|
||||
- Long-running tasks without blocking
|
||||
- Task queuing and management
|
||||
- Progress monitoring and notifications
|
||||
- Graceful cancellation
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Background processing only when explicitly started
|
||||
- No overhead when not in use
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.background import BackgroundRunner, BackgroundTask
|
||||
|
||||
# Create a background runner
|
||||
runner = BackgroundRunner()
|
||||
|
||||
# Submit a task
|
||||
task = runner.submit(
|
||||
agent=my_agent,
|
||||
prompt="Research AI trends",
|
||||
callback=on_complete
|
||||
)
|
||||
|
||||
# Check status
|
||||
print(task.status) # "running", "completed", "failed"
|
||||
|
||||
# Wait for completion
|
||||
result = await task.wait()
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core classes
|
||||
"BackgroundRunner",
|
||||
"BackgroundTask",
|
||||
# Status enum
|
||||
"TaskStatus",
|
||||
# Configuration
|
||||
"BackgroundConfig",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "BackgroundRunner":
|
||||
from .runner import BackgroundRunner
|
||||
return BackgroundRunner
|
||||
|
||||
if name == "BackgroundTask":
|
||||
from .task import BackgroundTask
|
||||
return BackgroundTask
|
||||
|
||||
if name == "TaskStatus":
|
||||
from .task import TaskStatus
|
||||
return TaskStatus
|
||||
|
||||
if name == "BackgroundConfig":
|
||||
from .config import BackgroundConfig
|
||||
return BackgroundConfig
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
Background Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class BackgroundConfig:
|
||||
"""Configuration for background task execution."""
|
||||
max_concurrent_tasks: int = 5
|
||||
default_timeout: Optional[float] = None # None = no timeout
|
||||
auto_cleanup: bool = True # Auto-remove completed tasks
|
||||
cleanup_delay: float = 300.0 # Seconds to keep completed tasks
|
||||
enable_notifications: bool = True
|
||||
log_progress: bool = True
|
||||
@@ -0,0 +1,314 @@
|
||||
"""
|
||||
Background Runner for PraisonAI Agents.
|
||||
|
||||
Manages background task execution.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any, Callable, Union
|
||||
from datetime import datetime
|
||||
|
||||
from .task import BackgroundTask, TaskStatus
|
||||
from .config import BackgroundConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BackgroundRunner:
|
||||
"""
|
||||
Manages background task execution.
|
||||
|
||||
Provides:
|
||||
- Task submission and queuing
|
||||
- Concurrent execution management
|
||||
- Progress monitoring
|
||||
- Task cancellation
|
||||
|
||||
Example:
|
||||
runner = BackgroundRunner()
|
||||
|
||||
# Submit a task
|
||||
task = await runner.submit(
|
||||
func=my_agent.start,
|
||||
args=("Research AI trends",),
|
||||
name="research_task"
|
||||
)
|
||||
|
||||
# Check status
|
||||
print(task.status)
|
||||
|
||||
# Wait for result
|
||||
result = await task.wait()
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[BackgroundConfig] = None):
|
||||
"""
|
||||
Initialize the background runner.
|
||||
|
||||
Args:
|
||||
config: Configuration for background execution
|
||||
"""
|
||||
self.config = config or BackgroundConfig()
|
||||
self._tasks: Dict[str, BackgroundTask] = {}
|
||||
self._semaphore = asyncio.Semaphore(self.config.max_concurrent_tasks)
|
||||
self._running = False
|
||||
self._cleanup_task: Optional[asyncio.Task] = None
|
||||
|
||||
@property
|
||||
def tasks(self) -> List[BackgroundTask]:
|
||||
"""Get all tasks."""
|
||||
return list(self._tasks.values())
|
||||
|
||||
@property
|
||||
def running_tasks(self) -> List[BackgroundTask]:
|
||||
"""Get currently running tasks."""
|
||||
return [t for t in self._tasks.values() if t.is_running]
|
||||
|
||||
@property
|
||||
def pending_tasks(self) -> List[BackgroundTask]:
|
||||
"""Get pending tasks."""
|
||||
return [t for t in self._tasks.values() if t.status == TaskStatus.PENDING]
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
func: Callable,
|
||||
args: tuple = (),
|
||||
kwargs: Optional[Dict[str, Any]] = None,
|
||||
name: Optional[str] = None,
|
||||
timeout: Optional[float] = None,
|
||||
on_complete: Optional[Callable[[BackgroundTask], None]] = None,
|
||||
on_progress: Optional[Callable[[BackgroundTask], None]] = None
|
||||
) -> BackgroundTask:
|
||||
"""
|
||||
Submit a task for background execution.
|
||||
|
||||
Args:
|
||||
func: Function to execute (can be sync or async)
|
||||
args: Positional arguments for the function
|
||||
kwargs: Keyword arguments for the function
|
||||
name: Optional task name
|
||||
timeout: Optional timeout in seconds
|
||||
on_complete: Callback when task completes
|
||||
on_progress: Callback for progress updates
|
||||
|
||||
Returns:
|
||||
BackgroundTask object for tracking
|
||||
"""
|
||||
kwargs = kwargs or {}
|
||||
|
||||
task = BackgroundTask(
|
||||
name=name or func.__name__,
|
||||
metadata={
|
||||
"func_name": func.__name__,
|
||||
"timeout": timeout or self.config.default_timeout
|
||||
}
|
||||
)
|
||||
|
||||
self._tasks[task.id] = task
|
||||
|
||||
# Create the execution coroutine
|
||||
async def execute():
|
||||
async with self._semaphore:
|
||||
task.start()
|
||||
logger.info(f"Background task started: {task.name} ({task.id})")
|
||||
|
||||
try:
|
||||
# Check if function is async
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
result = await asyncio.wait_for(
|
||||
func(*args, **kwargs),
|
||||
timeout=timeout or self.config.default_timeout
|
||||
)
|
||||
else:
|
||||
# Run sync function in thread pool
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await asyncio.wait_for(
|
||||
loop.run_in_executor(None, lambda: func(*args, **kwargs)),
|
||||
timeout=timeout or self.config.default_timeout
|
||||
)
|
||||
|
||||
task.complete(result)
|
||||
logger.info(f"Background task completed: {task.name} ({task.id})")
|
||||
|
||||
if on_complete:
|
||||
try:
|
||||
on_complete(task)
|
||||
except Exception as e:
|
||||
logger.warning(f"on_complete callback error: {e}")
|
||||
|
||||
return result
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
task.fail(f"Task timed out after {timeout}s")
|
||||
logger.warning(f"Background task timed out: {task.name}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
task.cancel()
|
||||
logger.info(f"Background task cancelled: {task.name}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
task.fail(str(e))
|
||||
logger.error(f"Background task failed: {task.name} - {e}")
|
||||
|
||||
if on_complete:
|
||||
try:
|
||||
on_complete(task)
|
||||
except Exception as cb_error:
|
||||
logger.warning(f"on_complete callback error: {cb_error}")
|
||||
|
||||
# Create and store the future
|
||||
task._future = asyncio.create_task(execute())
|
||||
|
||||
return task
|
||||
|
||||
async def submit_agent(
|
||||
self,
|
||||
agent: Any,
|
||||
prompt: str,
|
||||
name: Optional[str] = None,
|
||||
timeout: Optional[float] = None,
|
||||
on_complete: Optional[Callable[[BackgroundTask], None]] = None
|
||||
) -> BackgroundTask:
|
||||
"""
|
||||
Submit an agent task for background execution.
|
||||
|
||||
Args:
|
||||
agent: Agent instance to run
|
||||
prompt: Prompt to send to the agent
|
||||
name: Optional task name
|
||||
timeout: Optional timeout in seconds
|
||||
on_complete: Callback when task completes
|
||||
|
||||
Returns:
|
||||
BackgroundTask object for tracking
|
||||
"""
|
||||
task_name = name or f"agent_{getattr(agent, 'name', 'unknown')}"
|
||||
|
||||
# Determine the method to call
|
||||
if hasattr(agent, 'start'):
|
||||
func = agent.start
|
||||
elif hasattr(agent, 'chat'):
|
||||
func = agent.chat
|
||||
elif hasattr(agent, 'run'):
|
||||
func = agent.run
|
||||
else:
|
||||
raise ValueError("Agent must have start(), chat(), or run() method")
|
||||
|
||||
return await self.submit(
|
||||
func=func,
|
||||
args=(prompt,),
|
||||
name=task_name,
|
||||
timeout=timeout,
|
||||
on_complete=on_complete
|
||||
)
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[BackgroundTask]:
|
||||
"""Get a task by ID."""
|
||||
return self._tasks.get(task_id)
|
||||
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
"""
|
||||
Cancel a task.
|
||||
|
||||
Args:
|
||||
task_id: ID of task to cancel
|
||||
|
||||
Returns:
|
||||
True if cancelled, False if not found or already completed
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if task is None:
|
||||
return False
|
||||
|
||||
if task.is_completed:
|
||||
return False
|
||||
|
||||
task.cancel()
|
||||
|
||||
if task._future and not task._future.done():
|
||||
task._future.cancel()
|
||||
|
||||
return True
|
||||
|
||||
async def cancel_all(self):
|
||||
"""Cancel all running tasks."""
|
||||
for task in self.running_tasks:
|
||||
await self.cancel_task(task.id)
|
||||
|
||||
async def wait_all(self, timeout: Optional[float] = None) -> List[BackgroundTask]:
|
||||
"""
|
||||
Wait for all tasks to complete.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait
|
||||
|
||||
Returns:
|
||||
List of completed tasks
|
||||
"""
|
||||
futures = [t._future for t in self._tasks.values() if t._future]
|
||||
|
||||
if not futures:
|
||||
return list(self._tasks.values())
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(*futures, return_exceptions=True),
|
||||
timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
return list(self._tasks.values())
|
||||
|
||||
def clear_completed(self):
|
||||
"""Remove completed tasks from tracking."""
|
||||
completed_ids = [
|
||||
task_id for task_id, task in self._tasks.items()
|
||||
if task.is_completed
|
||||
]
|
||||
for task_id in completed_ids:
|
||||
del self._tasks[task_id]
|
||||
|
||||
def list_tasks(self, status: Optional[TaskStatus] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List tasks with optional status filter.
|
||||
|
||||
Args:
|
||||
status: Optional status to filter by
|
||||
|
||||
Returns:
|
||||
List of task dictionaries
|
||||
"""
|
||||
tasks = self._tasks.values()
|
||||
|
||||
if status:
|
||||
tasks = [t for t in tasks if t.status == status]
|
||||
|
||||
return [t.to_dict() for t in tasks]
|
||||
|
||||
async def start_cleanup_loop(self):
|
||||
"""Start the automatic cleanup loop."""
|
||||
if self._cleanup_task is not None:
|
||||
return
|
||||
|
||||
async def cleanup_loop():
|
||||
while True:
|
||||
await asyncio.sleep(self.config.cleanup_delay)
|
||||
if self.config.auto_cleanup:
|
||||
self.clear_completed()
|
||||
|
||||
self._cleanup_task = asyncio.create_task(cleanup_loop())
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the runner and cancel all tasks."""
|
||||
await self.cancel_all()
|
||||
|
||||
if self._cleanup_task:
|
||||
self._cleanup_task.cancel()
|
||||
try:
|
||||
await self._cleanup_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._cleanup_task = None
|
||||
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
Background Task for PraisonAI Agents.
|
||||
|
||||
Represents a task running in the background.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
import asyncio
|
||||
import logging
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Any, Callable, Dict
|
||||
from datetime import datetime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TaskStatus(str, Enum):
|
||||
"""Status of a background task."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BackgroundTask:
|
||||
"""
|
||||
Represents a background task.
|
||||
|
||||
Tracks the status, progress, and result of a task running
|
||||
in the background.
|
||||
"""
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
|
||||
name: str = ""
|
||||
status: TaskStatus = TaskStatus.PENDING
|
||||
progress: float = 0.0 # 0.0 to 1.0
|
||||
result: Any = None
|
||||
error: Optional[str] = None
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
started_at: Optional[datetime] = None
|
||||
completed_at: Optional[datetime] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Internal
|
||||
_future: Optional[asyncio.Future] = field(default=None, repr=False)
|
||||
_cancel_event: Optional[asyncio.Event] = field(default=None, repr=False)
|
||||
|
||||
def __post_init__(self):
|
||||
self._cancel_event = asyncio.Event()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if task is currently running."""
|
||||
return self.status == TaskStatus.RUNNING
|
||||
|
||||
@property
|
||||
def is_completed(self) -> bool:
|
||||
"""Check if task has completed (success or failure)."""
|
||||
return self.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED)
|
||||
|
||||
@property
|
||||
def is_successful(self) -> bool:
|
||||
"""Check if task completed successfully."""
|
||||
return self.status == TaskStatus.COMPLETED
|
||||
|
||||
@property
|
||||
def duration_seconds(self) -> Optional[float]:
|
||||
"""Get task duration in seconds."""
|
||||
if self.started_at is None:
|
||||
return None
|
||||
end_time = self.completed_at or datetime.now()
|
||||
return (end_time - self.started_at).total_seconds()
|
||||
|
||||
def start(self):
|
||||
"""Mark task as started."""
|
||||
self.status = TaskStatus.RUNNING
|
||||
self.started_at = datetime.now()
|
||||
|
||||
def complete(self, result: Any = None):
|
||||
"""Mark task as completed."""
|
||||
self.status = TaskStatus.COMPLETED
|
||||
self.result = result
|
||||
self.progress = 1.0
|
||||
self.completed_at = datetime.now()
|
||||
|
||||
def fail(self, error: str):
|
||||
"""Mark task as failed."""
|
||||
self.status = TaskStatus.FAILED
|
||||
self.error = error
|
||||
self.completed_at = datetime.now()
|
||||
|
||||
def cancel(self):
|
||||
"""Request task cancellation."""
|
||||
if self._cancel_event:
|
||||
self._cancel_event.set()
|
||||
self.status = TaskStatus.CANCELLED
|
||||
self.completed_at = datetime.now()
|
||||
|
||||
def update_progress(self, progress: float, message: Optional[str] = None):
|
||||
"""Update task progress."""
|
||||
self.progress = min(max(progress, 0.0), 1.0)
|
||||
if message:
|
||||
self.metadata["last_message"] = message
|
||||
|
||||
async def wait(self, timeout: Optional[float] = None) -> Any:
|
||||
"""
|
||||
Wait for task completion.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait in seconds
|
||||
|
||||
Returns:
|
||||
Task result
|
||||
|
||||
Raises:
|
||||
asyncio.TimeoutError: If timeout exceeded
|
||||
RuntimeError: If task failed
|
||||
"""
|
||||
if self._future is None:
|
||||
if self.is_completed:
|
||||
if self.status == TaskStatus.FAILED:
|
||||
raise RuntimeError(self.error)
|
||||
return self.result
|
||||
raise RuntimeError("Task has no associated future")
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(self._future, timeout=timeout)
|
||||
return result
|
||||
except asyncio.TimeoutError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(str(e))
|
||||
|
||||
def should_cancel(self) -> bool:
|
||||
"""Check if cancellation was requested."""
|
||||
return self._cancel_event is not None and self._cancel_event.is_set()
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"name": self.name,
|
||||
"status": self.status.value,
|
||||
"progress": self.progress,
|
||||
"result": str(self.result) if self.result else None,
|
||||
"error": self.error,
|
||||
"created_at": self.created_at.isoformat(),
|
||||
"started_at": self.started_at.isoformat() if self.started_at else None,
|
||||
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
|
||||
"duration_seconds": self.duration_seconds,
|
||||
"metadata": self.metadata
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
"""
|
||||
Checkpoints Module for PraisonAI Agents.
|
||||
|
||||
Provides file-level checkpointing using a shadow git repository to track
|
||||
and restore file changes made by agents. This enables:
|
||||
|
||||
- Automatic checkpointing before file modifications
|
||||
- Rewind to any previous checkpoint
|
||||
- Diff between checkpoints
|
||||
- Restore files and conversation state together
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Checkpoints only created when enabled
|
||||
- No overhead when checkpoints are disabled
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
|
||||
# Create checkpoint service
|
||||
service = CheckpointService(
|
||||
workspace_dir="/path/to/project",
|
||||
storage_dir="~/.praison/checkpoints"
|
||||
)
|
||||
|
||||
# Initialize shadow git
|
||||
await service.initialize()
|
||||
|
||||
# Save a checkpoint
|
||||
checkpoint_id = await service.save("Before refactoring")
|
||||
|
||||
# Restore to a checkpoint
|
||||
await service.restore(checkpoint_id)
|
||||
|
||||
# Get diff between checkpoints
|
||||
diff = await service.diff(from_id, to_id)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core service
|
||||
"CheckpointService",
|
||||
# Data types
|
||||
"Checkpoint",
|
||||
"CheckpointDiff",
|
||||
"CheckpointConfig",
|
||||
# Events
|
||||
"CheckpointEvent",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "CheckpointService":
|
||||
from .service import CheckpointService
|
||||
return CheckpointService
|
||||
|
||||
if name in ("Checkpoint", "CheckpointDiff", "CheckpointConfig"):
|
||||
from .types import Checkpoint, CheckpointDiff, CheckpointConfig
|
||||
return locals()[name]
|
||||
|
||||
if name == "CheckpointEvent":
|
||||
from .types import CheckpointEvent
|
||||
return CheckpointEvent
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,497 @@
|
||||
"""
|
||||
Checkpoint Service for PraisonAI Agents.
|
||||
|
||||
Implements shadow git repository for file-level checkpointing.
|
||||
"""
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
import logging
|
||||
import shutil
|
||||
from typing import Optional, List, Dict, Any, Callable
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from .types import (
|
||||
CheckpointConfig, Checkpoint, CheckpointDiff, FileDiff,
|
||||
CheckpointResult, CheckpointEvent
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Protected paths that should never be checkpointed
|
||||
PROTECTED_PATHS = [
|
||||
os.path.expanduser("~"),
|
||||
os.path.expanduser("~/Desktop"),
|
||||
os.path.expanduser("~/Documents"),
|
||||
os.path.expanduser("~/Downloads"),
|
||||
"/",
|
||||
"/tmp",
|
||||
"/var",
|
||||
"/etc",
|
||||
]
|
||||
|
||||
|
||||
class CheckpointService:
|
||||
"""
|
||||
Shadow Git Checkpoint Service.
|
||||
|
||||
Uses a separate git repository to track file changes in the workspace,
|
||||
enabling checkpoint creation, restoration, and diffing without
|
||||
interfering with the workspace's own git repository.
|
||||
|
||||
Example:
|
||||
service = CheckpointService(
|
||||
workspace_dir="/path/to/project",
|
||||
storage_dir="~/.praison/checkpoints"
|
||||
)
|
||||
|
||||
await service.initialize()
|
||||
|
||||
# Save checkpoint before making changes
|
||||
result = await service.save("Before refactoring")
|
||||
|
||||
# Make changes...
|
||||
|
||||
# Restore if needed
|
||||
await service.restore(result.checkpoint.id)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace_dir: str,
|
||||
storage_dir: Optional[str] = None,
|
||||
enabled: bool = True,
|
||||
auto_checkpoint: bool = True,
|
||||
max_checkpoints: int = 100
|
||||
):
|
||||
"""
|
||||
Initialize the checkpoint service.
|
||||
|
||||
Args:
|
||||
workspace_dir: Directory to track
|
||||
storage_dir: Where to store checkpoint data
|
||||
enabled: Whether checkpoints are enabled
|
||||
auto_checkpoint: Auto-checkpoint before file modifications
|
||||
max_checkpoints: Maximum checkpoints to keep
|
||||
"""
|
||||
self.config = CheckpointConfig(
|
||||
workspace_dir=workspace_dir,
|
||||
storage_dir=storage_dir,
|
||||
enabled=enabled,
|
||||
auto_checkpoint=auto_checkpoint,
|
||||
max_checkpoints=max_checkpoints
|
||||
)
|
||||
|
||||
self._initialized = False
|
||||
self._checkpoints: List[Checkpoint] = []
|
||||
self._event_handlers: Dict[CheckpointEvent, List[Callable]] = {
|
||||
event: [] for event in CheckpointEvent
|
||||
}
|
||||
|
||||
@property
|
||||
def workspace_dir(self) -> str:
|
||||
"""Get the workspace directory."""
|
||||
return self.config.workspace_dir
|
||||
|
||||
@property
|
||||
def checkpoint_dir(self) -> str:
|
||||
"""Get the checkpoint storage directory."""
|
||||
return self.config.get_checkpoint_dir()
|
||||
|
||||
@property
|
||||
def git_dir(self) -> str:
|
||||
"""Get the shadow git directory."""
|
||||
return os.path.join(self.checkpoint_dir, ".git")
|
||||
|
||||
@property
|
||||
def is_initialized(self) -> bool:
|
||||
"""Check if the service is initialized."""
|
||||
return self._initialized
|
||||
|
||||
def on(self, event: CheckpointEvent, handler: Callable):
|
||||
"""Register an event handler."""
|
||||
self._event_handlers[event].append(handler)
|
||||
|
||||
def _emit(self, event: CheckpointEvent, data: Any = None):
|
||||
"""Emit an event to all handlers."""
|
||||
for handler in self._event_handlers[event]:
|
||||
try:
|
||||
handler(data)
|
||||
except Exception as e:
|
||||
logger.warning(f"Event handler error: {e}")
|
||||
|
||||
async def initialize(self) -> bool:
|
||||
"""
|
||||
Initialize the shadow git repository.
|
||||
|
||||
Returns:
|
||||
True if initialization succeeded
|
||||
"""
|
||||
if not self.config.enabled:
|
||||
return False
|
||||
|
||||
# Check for protected paths
|
||||
workspace_real = os.path.realpath(self.config.workspace_dir)
|
||||
for protected in PROTECTED_PATHS:
|
||||
if workspace_real == os.path.realpath(protected):
|
||||
logger.warning(f"Cannot checkpoint protected path: {workspace_real}")
|
||||
return False
|
||||
|
||||
# Check if workspace exists
|
||||
if not os.path.isdir(self.config.workspace_dir):
|
||||
logger.error(f"Workspace does not exist: {self.config.workspace_dir}")
|
||||
return False
|
||||
|
||||
# Create checkpoint directory
|
||||
os.makedirs(self.checkpoint_dir, exist_ok=True)
|
||||
|
||||
# Initialize shadow git if needed
|
||||
if not os.path.exists(self.git_dir):
|
||||
success = await self._init_shadow_git()
|
||||
if not success:
|
||||
return False
|
||||
|
||||
# Load existing checkpoints
|
||||
await self._load_checkpoints()
|
||||
|
||||
self._initialized = True
|
||||
self._emit(CheckpointEvent.INITIALIZED, {"workspace": self.config.workspace_dir})
|
||||
|
||||
return True
|
||||
|
||||
async def _init_shadow_git(self) -> bool:
|
||||
"""Initialize the shadow git repository."""
|
||||
try:
|
||||
# Initialize git repo
|
||||
await self._run_git("init")
|
||||
|
||||
# Configure git
|
||||
await self._run_git("config", "user.name", "PraisonAI Checkpoints")
|
||||
await self._run_git("config", "user.email", "[email protected]")
|
||||
|
||||
# Set worktree to workspace
|
||||
await self._run_git("config", "core.worktree", self.config.workspace_dir)
|
||||
|
||||
# Create exclude file
|
||||
exclude_dir = os.path.join(self.git_dir, "info")
|
||||
os.makedirs(exclude_dir, exist_ok=True)
|
||||
|
||||
exclude_file = os.path.join(exclude_dir, "exclude")
|
||||
with open(exclude_file, "w") as f:
|
||||
for pattern in self.config.exclude_patterns:
|
||||
f.write(f"{pattern}\n")
|
||||
|
||||
# Also exclude the checkpoint directory itself
|
||||
f.write(f"{self.checkpoint_dir}\n")
|
||||
|
||||
# Create initial empty commit
|
||||
await self._run_git("commit", "--allow-empty", "-m", "Initial checkpoint")
|
||||
|
||||
logger.info(f"Initialized shadow git at {self.checkpoint_dir}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize shadow git: {e}")
|
||||
return False
|
||||
|
||||
async def _run_git(self, *args: str) -> str:
|
||||
"""Run a git command in the shadow repository."""
|
||||
env = self._get_sanitized_env()
|
||||
|
||||
cmd = ["git", f"--git-dir={self.git_dir}", f"--work-tree={self.config.workspace_dir}"]
|
||||
cmd.extend(args)
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
cwd=self.config.workspace_dir
|
||||
)
|
||||
|
||||
stdout, stderr = await process.communicate()
|
||||
|
||||
if process.returncode != 0:
|
||||
error_msg = stderr.decode().strip()
|
||||
raise RuntimeError(f"Git command failed: {' '.join(args)}\n{error_msg}")
|
||||
|
||||
return stdout.decode().strip()
|
||||
|
||||
def _get_sanitized_env(self) -> Dict[str, str]:
|
||||
"""Get environment with git variables sanitized."""
|
||||
env = os.environ.copy()
|
||||
|
||||
# Remove git environment variables that could interfere
|
||||
git_vars = [
|
||||
"GIT_DIR", "GIT_WORK_TREE", "GIT_INDEX_FILE",
|
||||
"GIT_OBJECT_DIRECTORY", "GIT_ALTERNATE_OBJECT_DIRECTORIES"
|
||||
]
|
||||
for var in git_vars:
|
||||
env.pop(var, None)
|
||||
|
||||
return env
|
||||
|
||||
async def save(self, message: str, allow_empty: bool = False) -> CheckpointResult:
|
||||
"""
|
||||
Save a checkpoint.
|
||||
|
||||
Args:
|
||||
message: Checkpoint message
|
||||
allow_empty: Allow checkpoint even if no changes
|
||||
|
||||
Returns:
|
||||
CheckpointResult with the created checkpoint
|
||||
"""
|
||||
if not self._initialized:
|
||||
return CheckpointResult.fail("Service not initialized")
|
||||
|
||||
try:
|
||||
# Stage all changes
|
||||
await self._run_git("add", "-A")
|
||||
|
||||
# Check if there are changes
|
||||
try:
|
||||
await self._run_git("diff", "--cached", "--quiet")
|
||||
has_changes = False
|
||||
except RuntimeError:
|
||||
has_changes = True
|
||||
|
||||
if not has_changes and not allow_empty:
|
||||
return CheckpointResult.fail("No changes to checkpoint")
|
||||
|
||||
# Create commit
|
||||
commit_args = ["commit", "-m", message]
|
||||
if allow_empty:
|
||||
commit_args.append("--allow-empty")
|
||||
|
||||
await self._run_git(*commit_args)
|
||||
|
||||
# Get commit info
|
||||
commit_hash = await self._run_git("rev-parse", "HEAD")
|
||||
timestamp = await self._run_git("log", "-1", "--format=%aI")
|
||||
|
||||
checkpoint = Checkpoint(
|
||||
id=commit_hash,
|
||||
short_id=commit_hash[:8],
|
||||
message=message,
|
||||
timestamp=datetime.fromisoformat(timestamp)
|
||||
)
|
||||
|
||||
self._checkpoints.append(checkpoint)
|
||||
|
||||
# Prune old checkpoints if needed
|
||||
await self._prune_checkpoints()
|
||||
|
||||
self._emit(CheckpointEvent.CHECKPOINT_CREATED, checkpoint)
|
||||
logger.info(f"Created checkpoint: {checkpoint.short_id} - {message}")
|
||||
|
||||
return CheckpointResult.ok(checkpoint)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._emit(CheckpointEvent.ERROR, {"error": error_msg})
|
||||
return CheckpointResult.fail(error_msg)
|
||||
|
||||
async def restore(self, checkpoint_id: str) -> CheckpointResult:
|
||||
"""
|
||||
Restore workspace to a checkpoint.
|
||||
|
||||
Args:
|
||||
checkpoint_id: Checkpoint ID (commit hash) to restore
|
||||
|
||||
Returns:
|
||||
CheckpointResult indicating success/failure
|
||||
"""
|
||||
if not self._initialized:
|
||||
return CheckpointResult.fail("Service not initialized")
|
||||
|
||||
try:
|
||||
# Clean untracked files
|
||||
await self._run_git("clean", "-f", "-d")
|
||||
|
||||
# Reset to checkpoint
|
||||
await self._run_git("reset", "--hard", checkpoint_id)
|
||||
|
||||
# Get checkpoint info
|
||||
message = await self._run_git("log", "-1", "--format=%s", checkpoint_id)
|
||||
timestamp = await self._run_git("log", "-1", "--format=%aI", checkpoint_id)
|
||||
|
||||
checkpoint = Checkpoint(
|
||||
id=checkpoint_id,
|
||||
short_id=checkpoint_id[:8],
|
||||
message=message,
|
||||
timestamp=datetime.fromisoformat(timestamp)
|
||||
)
|
||||
|
||||
self._emit(CheckpointEvent.CHECKPOINT_RESTORED, checkpoint)
|
||||
logger.info(f"Restored to checkpoint: {checkpoint.short_id}")
|
||||
|
||||
return CheckpointResult.ok(checkpoint)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._emit(CheckpointEvent.ERROR, {"error": error_msg})
|
||||
return CheckpointResult.fail(error_msg)
|
||||
|
||||
async def diff(
|
||||
self,
|
||||
from_id: Optional[str] = None,
|
||||
to_id: Optional[str] = None
|
||||
) -> CheckpointDiff:
|
||||
"""
|
||||
Get diff between checkpoints.
|
||||
|
||||
Args:
|
||||
from_id: Starting checkpoint (default: previous checkpoint)
|
||||
to_id: Ending checkpoint (default: current working directory)
|
||||
|
||||
Returns:
|
||||
CheckpointDiff with file changes
|
||||
"""
|
||||
if not self._initialized:
|
||||
return CheckpointDiff(from_checkpoint="", to_checkpoint="", files=[])
|
||||
|
||||
try:
|
||||
# Default from_id to HEAD~1
|
||||
if from_id is None:
|
||||
try:
|
||||
from_id = await self._run_git("rev-parse", "HEAD~1")
|
||||
except RuntimeError:
|
||||
from_id = await self._run_git("rev-parse", "HEAD")
|
||||
|
||||
# Build diff command
|
||||
if to_id:
|
||||
diff_output = await self._run_git(
|
||||
"diff", "--stat", "--numstat", from_id, to_id
|
||||
)
|
||||
else:
|
||||
# Diff against working directory
|
||||
await self._run_git("add", "-A") # Stage to include untracked
|
||||
diff_output = await self._run_git(
|
||||
"diff", "--stat", "--numstat", "--cached", from_id
|
||||
)
|
||||
|
||||
# Parse diff output
|
||||
files = self._parse_diff_output(diff_output)
|
||||
|
||||
return CheckpointDiff(
|
||||
from_checkpoint=from_id[:8] if from_id else "",
|
||||
to_checkpoint=to_id[:8] if to_id else None,
|
||||
files=files
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get diff: {e}")
|
||||
return CheckpointDiff(from_checkpoint="", to_checkpoint="", files=[])
|
||||
|
||||
def _parse_diff_output(self, output: str) -> List[FileDiff]:
|
||||
"""Parse git diff --numstat output."""
|
||||
files = []
|
||||
|
||||
for line in output.strip().split("\n"):
|
||||
if not line or "\t" not in line:
|
||||
continue
|
||||
|
||||
parts = line.split("\t")
|
||||
if len(parts) >= 3:
|
||||
try:
|
||||
additions = int(parts[0]) if parts[0] != "-" else 0
|
||||
deletions = int(parts[1]) if parts[1] != "-" else 0
|
||||
path = parts[2]
|
||||
|
||||
# Determine status
|
||||
if additions > 0 and deletions == 0:
|
||||
status = "added"
|
||||
elif additions == 0 and deletions > 0:
|
||||
status = "deleted"
|
||||
else:
|
||||
status = "modified"
|
||||
|
||||
files.append(FileDiff(
|
||||
path=path,
|
||||
absolute_path=os.path.join(self.config.workspace_dir, path),
|
||||
status=status,
|
||||
additions=additions,
|
||||
deletions=deletions
|
||||
))
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return files
|
||||
|
||||
async def list_checkpoints(self, limit: int = 50) -> List[Checkpoint]:
|
||||
"""
|
||||
List all checkpoints.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of checkpoints to return
|
||||
|
||||
Returns:
|
||||
List of checkpoints, newest first
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
try:
|
||||
# Get commit log
|
||||
log_output = await self._run_git(
|
||||
"log", f"-{limit}", "--format=%H|%s|%aI"
|
||||
)
|
||||
|
||||
checkpoints = []
|
||||
for line in log_output.strip().split("\n"):
|
||||
if not line or "|" not in line:
|
||||
continue
|
||||
|
||||
parts = line.split("|", 2)
|
||||
if len(parts) >= 3:
|
||||
checkpoints.append(Checkpoint(
|
||||
id=parts[0],
|
||||
short_id=parts[0][:8],
|
||||
message=parts[1],
|
||||
timestamp=datetime.fromisoformat(parts[2])
|
||||
))
|
||||
|
||||
return checkpoints
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list checkpoints: {e}")
|
||||
return []
|
||||
|
||||
async def _load_checkpoints(self):
|
||||
"""Load existing checkpoints from git log."""
|
||||
self._checkpoints = await self.list_checkpoints()
|
||||
|
||||
async def _prune_checkpoints(self):
|
||||
"""Prune old checkpoints if over limit."""
|
||||
if len(self._checkpoints) <= self.config.max_checkpoints:
|
||||
return
|
||||
|
||||
# Keep only the most recent checkpoints
|
||||
# Note: This doesn't actually delete git history, just our tracking
|
||||
self._checkpoints = self._checkpoints[:self.config.max_checkpoints]
|
||||
|
||||
async def get_checkpoint(self, checkpoint_id: str) -> Optional[Checkpoint]:
|
||||
"""Get a specific checkpoint by ID."""
|
||||
for cp in self._checkpoints:
|
||||
if cp.id.startswith(checkpoint_id) or cp.short_id == checkpoint_id:
|
||||
return cp
|
||||
return None
|
||||
|
||||
async def cleanup(self):
|
||||
"""Clean up the checkpoint service."""
|
||||
# Nothing to clean up currently
|
||||
self._initialized = False
|
||||
|
||||
async def delete_all(self) -> bool:
|
||||
"""Delete all checkpoint data for this workspace."""
|
||||
try:
|
||||
if os.path.exists(self.checkpoint_dir):
|
||||
shutil.rmtree(self.checkpoint_dir)
|
||||
self._checkpoints = []
|
||||
self._initialized = False
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete checkpoints: {e}")
|
||||
return False
|
||||
@@ -0,0 +1,158 @@
|
||||
"""
|
||||
Checkpoint Types for PraisonAI Agents.
|
||||
|
||||
Defines the core types, enums, and dataclasses for the checkpoint system.
|
||||
"""
|
||||
|
||||
import os
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class CheckpointEvent(str, Enum):
|
||||
"""Events emitted by the checkpoint service."""
|
||||
INITIALIZED = "initialized"
|
||||
CHECKPOINT_CREATED = "checkpoint_created"
|
||||
CHECKPOINT_RESTORED = "checkpoint_restored"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointConfig:
|
||||
"""Configuration for the checkpoint service."""
|
||||
workspace_dir: str
|
||||
storage_dir: Optional[str] = None
|
||||
enabled: bool = True
|
||||
auto_checkpoint: bool = True # Auto-checkpoint before file modifications
|
||||
max_checkpoints: int = 100 # Maximum checkpoints to keep
|
||||
exclude_patterns: List[str] = field(default_factory=lambda: [
|
||||
".git",
|
||||
".praison",
|
||||
"__pycache__",
|
||||
"*.pyc",
|
||||
".env",
|
||||
"node_modules",
|
||||
".venv",
|
||||
"venv",
|
||||
"*.log"
|
||||
])
|
||||
|
||||
def __post_init__(self):
|
||||
# Expand paths
|
||||
self.workspace_dir = os.path.expanduser(self.workspace_dir)
|
||||
if self.storage_dir:
|
||||
self.storage_dir = os.path.expanduser(self.storage_dir)
|
||||
else:
|
||||
# Default storage in user's home directory
|
||||
self.storage_dir = os.path.expanduser("~/.praison/checkpoints")
|
||||
|
||||
def get_checkpoint_dir(self) -> str:
|
||||
"""Get the checkpoint directory for this workspace."""
|
||||
# Hash the workspace path to create a unique directory
|
||||
import hashlib
|
||||
workspace_hash = hashlib.sha256(
|
||||
self.workspace_dir.encode()
|
||||
).hexdigest()[:12]
|
||||
return os.path.join(self.storage_dir, workspace_hash)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Checkpoint:
|
||||
"""Represents a single checkpoint."""
|
||||
id: str # Git commit hash
|
||||
short_id: str # Short hash (8 chars)
|
||||
message: str
|
||||
timestamp: datetime
|
||||
files_changed: int = 0
|
||||
insertions: int = 0
|
||||
deletions: int = 0
|
||||
|
||||
@classmethod
|
||||
def from_git_commit(cls, commit_hash: str, message: str, timestamp: str) -> "Checkpoint":
|
||||
"""Create a Checkpoint from git commit info."""
|
||||
return cls(
|
||||
id=commit_hash,
|
||||
short_id=commit_hash[:8],
|
||||
message=message,
|
||||
timestamp=datetime.fromisoformat(timestamp) if isinstance(timestamp, str) else timestamp
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"short_id": self.short_id,
|
||||
"message": self.message,
|
||||
"timestamp": self.timestamp.isoformat(),
|
||||
"files_changed": self.files_changed,
|
||||
"insertions": self.insertions,
|
||||
"deletions": self.deletions
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileDiff:
|
||||
"""Represents a diff for a single file."""
|
||||
path: str
|
||||
absolute_path: str
|
||||
status: str # added, modified, deleted
|
||||
additions: int = 0
|
||||
deletions: int = 0
|
||||
before_content: Optional[str] = None
|
||||
after_content: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"path": self.path,
|
||||
"absolute_path": self.absolute_path,
|
||||
"status": self.status,
|
||||
"additions": self.additions,
|
||||
"deletions": self.deletions
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointDiff:
|
||||
"""Represents a diff between two checkpoints."""
|
||||
from_checkpoint: str
|
||||
to_checkpoint: Optional[str] # None means current working directory
|
||||
files: List[FileDiff] = field(default_factory=list)
|
||||
total_additions: int = 0
|
||||
total_deletions: int = 0
|
||||
|
||||
def __post_init__(self):
|
||||
# Calculate totals
|
||||
self.total_additions = sum(f.additions for f in self.files)
|
||||
self.total_deletions = sum(f.deletions for f in self.files)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"from_checkpoint": self.from_checkpoint,
|
||||
"to_checkpoint": self.to_checkpoint,
|
||||
"files": [f.to_dict() for f in self.files],
|
||||
"total_additions": self.total_additions,
|
||||
"total_deletions": self.total_deletions
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointResult:
|
||||
"""Result of a checkpoint operation."""
|
||||
success: bool
|
||||
checkpoint: Optional[Checkpoint] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def ok(cls, checkpoint: Checkpoint) -> "CheckpointResult":
|
||||
"""Create a successful result."""
|
||||
return cls(success=True, checkpoint=checkpoint)
|
||||
|
||||
@classmethod
|
||||
def fail(cls, error: str) -> "CheckpointResult":
|
||||
"""Create a failed result."""
|
||||
return cls(success=False, error=error)
|
||||
@@ -0,0 +1,57 @@
|
||||
"""
|
||||
Auto Context Compaction Module for PraisonAI Agents.
|
||||
|
||||
Provides automatic context management for long conversations:
|
||||
- Message summarization
|
||||
- Context window management
|
||||
- Token counting and limits
|
||||
- Intelligent message pruning
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Compaction only runs when needed
|
||||
- No overhead when context is within limits
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.compaction import ContextCompactor, CompactionConfig
|
||||
|
||||
# Create a compactor
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=8000,
|
||||
strategy="summarize"
|
||||
)
|
||||
|
||||
# Compact messages when needed
|
||||
compacted = compactor.compact(messages)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core classes
|
||||
"ContextCompactor",
|
||||
"CompactionConfig",
|
||||
# Strategies
|
||||
"CompactionStrategy",
|
||||
# Results
|
||||
"CompactionResult",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "ContextCompactor":
|
||||
from .compactor import ContextCompactor
|
||||
return ContextCompactor
|
||||
|
||||
if name == "CompactionConfig":
|
||||
from .config import CompactionConfig
|
||||
return CompactionConfig
|
||||
|
||||
if name == "CompactionStrategy":
|
||||
from .strategy import CompactionStrategy
|
||||
return CompactionStrategy
|
||||
|
||||
if name == "CompactionResult":
|
||||
from .result import CompactionResult
|
||||
return CompactionResult
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
Context Compactor for PraisonAI Agents.
|
||||
|
||||
Manages context window by compacting messages when needed.
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
from .config import CompactionConfig
|
||||
from .strategy import CompactionStrategy
|
||||
from .result import CompactionResult
|
||||
|
||||
|
||||
class ContextCompactor:
|
||||
"""
|
||||
Compacts conversation context to fit within token limits.
|
||||
|
||||
Example:
|
||||
compactor = ContextCompactor(max_tokens=8000)
|
||||
|
||||
# Check if compaction needed
|
||||
if compactor.needs_compaction(messages):
|
||||
result = compactor.compact(messages)
|
||||
messages = result.messages
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_tokens: int = 8000,
|
||||
target_tokens: Optional[int] = None,
|
||||
strategy: CompactionStrategy = CompactionStrategy.TRUNCATE,
|
||||
preserve_system: bool = True,
|
||||
preserve_recent: int = 5
|
||||
):
|
||||
"""
|
||||
Initialize the compactor.
|
||||
|
||||
Args:
|
||||
max_tokens: Maximum tokens before compaction
|
||||
target_tokens: Target tokens after compaction
|
||||
strategy: Compaction strategy to use
|
||||
preserve_system: Keep system messages
|
||||
preserve_recent: Number of recent messages to preserve
|
||||
"""
|
||||
self.max_tokens = max_tokens
|
||||
self.target_tokens = target_tokens or int(max_tokens * 0.75)
|
||||
self.strategy = strategy
|
||||
self.preserve_system = preserve_system
|
||||
self.preserve_recent = preserve_recent
|
||||
|
||||
def estimate_tokens(self, text: str) -> int:
|
||||
"""Estimate token count for text."""
|
||||
# Rough estimate: ~4 chars per token
|
||||
return len(text) // 4
|
||||
|
||||
def count_message_tokens(self, message: Dict[str, Any]) -> int:
|
||||
"""Count tokens in a message."""
|
||||
content = message.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return self.estimate_tokens(content)
|
||||
elif isinstance(content, list):
|
||||
# Handle multi-part content
|
||||
total = 0
|
||||
for part in content:
|
||||
if isinstance(part, dict):
|
||||
total += self.estimate_tokens(str(part.get("text", "")))
|
||||
else:
|
||||
total += self.estimate_tokens(str(part))
|
||||
return total
|
||||
return 0
|
||||
|
||||
def count_total_tokens(self, messages: List[Dict[str, Any]]) -> int:
|
||||
"""Count total tokens in messages."""
|
||||
return sum(self.count_message_tokens(m) for m in messages)
|
||||
|
||||
def needs_compaction(self, messages: List[Dict[str, Any]]) -> bool:
|
||||
"""Check if messages need compaction."""
|
||||
return self.count_total_tokens(messages) > self.max_tokens
|
||||
|
||||
def compact(
|
||||
self,
|
||||
messages: List[Dict[str, Any]]
|
||||
) -> tuple[List[Dict[str, Any]], CompactionResult]:
|
||||
"""
|
||||
Compact messages to fit within token limit.
|
||||
|
||||
Args:
|
||||
messages: List of messages to compact
|
||||
|
||||
Returns:
|
||||
Tuple of (compacted messages, result)
|
||||
"""
|
||||
original_tokens = self.count_total_tokens(messages)
|
||||
|
||||
if original_tokens <= self.max_tokens:
|
||||
return messages, CompactionResult(
|
||||
original_tokens=original_tokens,
|
||||
compacted_tokens=original_tokens,
|
||||
messages_removed=0,
|
||||
messages_kept=len(messages),
|
||||
strategy_used=self.strategy
|
||||
)
|
||||
|
||||
if self.strategy == CompactionStrategy.TRUNCATE:
|
||||
compacted = self._truncate(messages)
|
||||
elif self.strategy == CompactionStrategy.SLIDING:
|
||||
compacted = self._sliding_window(messages)
|
||||
elif self.strategy == CompactionStrategy.SUMMARIZE:
|
||||
compacted = self._summarize(messages)
|
||||
elif self.strategy == CompactionStrategy.SMART:
|
||||
compacted = self._smart_compact(messages)
|
||||
else:
|
||||
compacted = self._truncate(messages)
|
||||
|
||||
compacted_tokens = self.count_total_tokens(compacted)
|
||||
|
||||
result = CompactionResult(
|
||||
original_tokens=original_tokens,
|
||||
compacted_tokens=compacted_tokens,
|
||||
messages_removed=len(messages) - len(compacted),
|
||||
messages_kept=len(compacted),
|
||||
strategy_used=self.strategy
|
||||
)
|
||||
|
||||
return compacted, result
|
||||
|
||||
def _truncate(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Truncate oldest messages."""
|
||||
result = []
|
||||
|
||||
# Separate system and non-system messages
|
||||
system_msgs = []
|
||||
other_msgs = []
|
||||
|
||||
for msg in messages:
|
||||
if self.preserve_system and msg.get("role") == "system":
|
||||
system_msgs.append(msg)
|
||||
else:
|
||||
other_msgs.append(msg)
|
||||
|
||||
# Always keep system messages
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:] if other_msgs else []
|
||||
|
||||
# Add recent messages
|
||||
result.extend(recent)
|
||||
|
||||
# If still over limit, truncate more
|
||||
while self.count_total_tokens(result) > self.target_tokens and len(result) > 1:
|
||||
# Remove oldest non-system message
|
||||
for i, msg in enumerate(result):
|
||||
if msg.get("role") != "system":
|
||||
result.pop(i)
|
||||
break
|
||||
|
||||
return result
|
||||
|
||||
def _sliding_window(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Keep a sliding window of recent messages."""
|
||||
result = []
|
||||
|
||||
# Keep system messages
|
||||
for msg in messages:
|
||||
if self.preserve_system and msg.get("role") == "system":
|
||||
result.append(msg)
|
||||
|
||||
# Add messages from end until we hit target
|
||||
non_system = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
for msg in reversed(non_system):
|
||||
if self.count_total_tokens(result + [msg]) <= self.target_tokens:
|
||||
result.insert(len([m for m in result if m.get("role") == "system"]), msg)
|
||||
else:
|
||||
break
|
||||
|
||||
# Ensure messages are in order
|
||||
system_msgs = [m for m in result if m.get("role") == "system"]
|
||||
other_msgs = [m for m in result if m.get("role") != "system"]
|
||||
|
||||
return system_msgs + other_msgs
|
||||
|
||||
def _summarize(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Summarize old messages (simplified version)."""
|
||||
result = []
|
||||
|
||||
# Keep system messages
|
||||
system_msgs = [m for m in messages if m.get("role") == "system"]
|
||||
other_msgs = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:]
|
||||
|
||||
# Summarize older messages
|
||||
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
|
||||
if older:
|
||||
# Create a simple summary
|
||||
summary_parts = []
|
||||
for msg in older:
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str) and content:
|
||||
# Take first 100 chars
|
||||
summary_parts.append(f"{role}: {content[:100]}...")
|
||||
|
||||
if summary_parts:
|
||||
summary = "[Previous conversation summary]\n" + "\n".join(summary_parts[:5])
|
||||
result.append({
|
||||
"role": "system",
|
||||
"content": summary
|
||||
})
|
||||
|
||||
result.extend(recent)
|
||||
|
||||
return result
|
||||
|
||||
def _smart_compact(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Smart compaction based on message importance."""
|
||||
# For now, use sliding window as base
|
||||
return self._sliding_window(messages)
|
||||
|
||||
def get_stats(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""Get statistics about messages."""
|
||||
total_tokens = self.count_total_tokens(messages)
|
||||
|
||||
return {
|
||||
"message_count": len(messages),
|
||||
"total_tokens": total_tokens,
|
||||
"max_tokens": self.max_tokens,
|
||||
"target_tokens": self.target_tokens,
|
||||
"needs_compaction": total_tokens > self.max_tokens,
|
||||
"utilization": total_tokens / self.max_tokens if self.max_tokens > 0 else 0
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
Compaction Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionConfig:
|
||||
"""Configuration for context compaction."""
|
||||
enabled: bool = True
|
||||
max_tokens: int = 8000
|
||||
target_tokens: int = 6000 # Target after compaction
|
||||
preserve_system: bool = True # Always keep system messages
|
||||
preserve_recent: int = 5 # Keep last N messages
|
||||
auto_compact: bool = True # Automatically compact when needed
|
||||
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
Compaction Result for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from .strategy import CompactionStrategy
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionResult:
|
||||
"""Result of a compaction operation."""
|
||||
original_tokens: int
|
||||
compacted_tokens: int
|
||||
messages_removed: int
|
||||
messages_kept: int
|
||||
strategy_used: CompactionStrategy
|
||||
summary: str = ""
|
||||
|
||||
@property
|
||||
def tokens_saved(self) -> int:
|
||||
"""Get number of tokens saved."""
|
||||
return self.original_tokens - self.compacted_tokens
|
||||
|
||||
@property
|
||||
def compression_ratio(self) -> float:
|
||||
"""Get compression ratio."""
|
||||
if self.original_tokens == 0:
|
||||
return 0.0
|
||||
return self.compacted_tokens / self.original_tokens
|
||||
|
||||
@property
|
||||
def was_compacted(self) -> bool:
|
||||
"""Check if compaction occurred."""
|
||||
return self.messages_removed > 0 or self.tokens_saved > 0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"original_tokens": self.original_tokens,
|
||||
"compacted_tokens": self.compacted_tokens,
|
||||
"tokens_saved": self.tokens_saved,
|
||||
"messages_removed": self.messages_removed,
|
||||
"messages_kept": self.messages_kept,
|
||||
"strategy_used": self.strategy_used.value,
|
||||
"compression_ratio": self.compression_ratio,
|
||||
"summary": self.summary
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
Compaction Strategies for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CompactionStrategy(str, Enum):
|
||||
"""Strategy for compacting context."""
|
||||
TRUNCATE = "truncate" # Remove oldest messages
|
||||
SUMMARIZE = "summarize" # Summarize old messages
|
||||
SLIDING = "sliding" # Sliding window of recent messages
|
||||
SMART = "smart" # Intelligent selection based on relevance
|
||||
@@ -0,0 +1,97 @@
|
||||
"""
|
||||
Hooks Module for PraisonAI Agents.
|
||||
|
||||
Provides a powerful hook system for intercepting and modifying agent behavior
|
||||
at various lifecycle points. Unlike callbacks (which are for UI events),
|
||||
hooks can intercept, modify, or block tool execution.
|
||||
|
||||
Features:
|
||||
- Event-based hook system (BeforeTool, AfterTool, BeforeAgent, etc.)
|
||||
- Shell command hooks for external integrations
|
||||
- Python function hooks for in-process customization
|
||||
- Matcher patterns for selective hook execution
|
||||
- Sequential and parallel hook execution
|
||||
- Decision outcomes (allow, deny, block, ask)
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Hooks only execute when registered
|
||||
- No overhead when hooks are disabled
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.hooks import HookRegistry, HookEvent
|
||||
|
||||
# Register a Python function hook
|
||||
registry = HookRegistry()
|
||||
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def my_hook(event_data):
|
||||
if event_data.tool_name == "dangerous_tool":
|
||||
return HookResult(decision="deny", reason="Tool blocked by policy")
|
||||
return HookResult(decision="allow")
|
||||
|
||||
# Register a shell command hook
|
||||
registry.register_command_hook(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="python /path/to/validator.py",
|
||||
matcher="write_*" # Only match tools starting with write_
|
||||
)
|
||||
|
||||
# Use with Agent
|
||||
agent = Agent(
|
||||
name="MyAgent",
|
||||
hooks=registry
|
||||
)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core types
|
||||
"HookEvent",
|
||||
"HookDecision",
|
||||
"HookResult",
|
||||
"HookInput",
|
||||
"HookOutput",
|
||||
# Hook definitions
|
||||
"HookDefinition",
|
||||
"CommandHook",
|
||||
"FunctionHook",
|
||||
# Registry and runner
|
||||
"HookRegistry",
|
||||
"HookRunner",
|
||||
# Event-specific inputs
|
||||
"BeforeToolInput",
|
||||
"AfterToolInput",
|
||||
"BeforeAgentInput",
|
||||
"AfterAgentInput",
|
||||
"SessionStartInput",
|
||||
"SessionEndInput",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name in ("HookEvent", "HookDecision", "HookResult", "HookInput", "HookOutput"):
|
||||
from .types import HookEvent, HookDecision, HookResult, HookInput, HookOutput
|
||||
return locals()[name]
|
||||
|
||||
if name in ("HookDefinition", "CommandHook", "FunctionHook"):
|
||||
from .types import HookDefinition, CommandHook, FunctionHook
|
||||
return locals()[name]
|
||||
|
||||
if name == "HookRegistry":
|
||||
from .registry import HookRegistry
|
||||
return HookRegistry
|
||||
|
||||
if name == "HookRunner":
|
||||
from .runner import HookRunner
|
||||
return HookRunner
|
||||
|
||||
if name in ("BeforeToolInput", "AfterToolInput", "BeforeAgentInput",
|
||||
"AfterAgentInput", "SessionStartInput", "SessionEndInput"):
|
||||
from .events import (
|
||||
BeforeToolInput, AfterToolInput, BeforeAgentInput,
|
||||
AfterAgentInput, SessionStartInput, SessionEndInput
|
||||
)
|
||||
return locals()[name]
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,195 @@
|
||||
"""
|
||||
Event-specific input types for the hook system.
|
||||
|
||||
Each event type has its own input class with relevant fields.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List
|
||||
from .types import HookInput
|
||||
|
||||
|
||||
@dataclass
|
||||
class BeforeToolInput(HookInput):
|
||||
"""Input for BeforeTool hooks."""
|
||||
tool_name: str = ""
|
||||
tool_input: Dict[str, Any] = field(default_factory=dict)
|
||||
tool_description: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"tool_name": self.tool_name,
|
||||
"tool_input": self.tool_input,
|
||||
"tool_description": self.tool_description
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class AfterToolInput(HookInput):
|
||||
"""Input for AfterTool hooks."""
|
||||
tool_name: str = ""
|
||||
tool_input: Dict[str, Any] = field(default_factory=dict)
|
||||
tool_output: Any = None
|
||||
tool_error: Optional[str] = None
|
||||
execution_time_ms: float = 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"tool_name": self.tool_name,
|
||||
"tool_input": self.tool_input,
|
||||
"tool_output": str(self.tool_output) if self.tool_output else None,
|
||||
"tool_error": self.tool_error,
|
||||
"execution_time_ms": self.execution_time_ms
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class BeforeAgentInput(HookInput):
|
||||
"""Input for BeforeAgent hooks."""
|
||||
prompt: str = ""
|
||||
conversation_history: List[Dict[str, Any]] = field(default_factory=list)
|
||||
tools_available: List[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"prompt": self.prompt,
|
||||
"conversation_history_length": len(self.conversation_history),
|
||||
"tools_available": self.tools_available
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class AfterAgentInput(HookInput):
|
||||
"""Input for AfterAgent hooks."""
|
||||
prompt: str = ""
|
||||
response: str = ""
|
||||
tools_used: List[str] = field(default_factory=list)
|
||||
total_tokens: int = 0
|
||||
execution_time_ms: float = 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"prompt": self.prompt,
|
||||
"response": self.response[:500] if self.response else "", # Truncate for JSON
|
||||
"tools_used": self.tools_used,
|
||||
"total_tokens": self.total_tokens,
|
||||
"execution_time_ms": self.execution_time_ms
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionStartInput(HookInput):
|
||||
"""Input for SessionStart hooks."""
|
||||
source: str = "startup" # startup, resume, clear
|
||||
session_name: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"source": self.source,
|
||||
"session_name": self.session_name
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionEndInput(HookInput):
|
||||
"""Input for SessionEnd hooks."""
|
||||
reason: str = "exit" # exit, clear, logout, error
|
||||
total_turns: int = 0
|
||||
total_tokens: int = 0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"reason": self.reason,
|
||||
"total_turns": self.total_turns,
|
||||
"total_tokens": self.total_tokens
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class BeforeLLMInput(HookInput):
|
||||
"""Input for BeforeLLM hooks."""
|
||||
messages: List[Dict[str, Any]] = field(default_factory=list)
|
||||
model: str = ""
|
||||
temperature: Optional[float] = None
|
||||
max_tokens: Optional[int] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"messages_count": len(self.messages),
|
||||
"model": self.model,
|
||||
"temperature": self.temperature,
|
||||
"max_tokens": self.max_tokens
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class AfterLLMInput(HookInput):
|
||||
"""Input for AfterLLM hooks."""
|
||||
messages: List[Dict[str, Any]] = field(default_factory=list)
|
||||
response: str = ""
|
||||
model: str = ""
|
||||
tokens_used: int = 0
|
||||
latency_ms: float = 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"messages_count": len(self.messages),
|
||||
"response_length": len(self.response),
|
||||
"model": self.model,
|
||||
"tokens_used": self.tokens_used,
|
||||
"latency_ms": self.latency_ms
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class OnErrorInput(HookInput):
|
||||
"""Input for OnError hooks."""
|
||||
error_type: str = ""
|
||||
error_message: str = ""
|
||||
stack_trace: Optional[str] = None
|
||||
context: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"error_type": self.error_type,
|
||||
"error_message": self.error_message,
|
||||
"stack_trace": self.stack_trace,
|
||||
"context": self.context
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class OnRetryInput(HookInput):
|
||||
"""Input for OnRetry hooks."""
|
||||
retry_count: int = 0
|
||||
max_retries: int = 3
|
||||
error_message: str = ""
|
||||
operation: str = "" # tool_call, llm_request, etc.
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"retry_count": self.retry_count,
|
||||
"max_retries": self.max_retries,
|
||||
"error_message": self.error_message,
|
||||
"operation": self.operation
|
||||
})
|
||||
return base
|
||||
@@ -0,0 +1,336 @@
|
||||
"""
|
||||
Hook Registry for PraisonAI Agents.
|
||||
|
||||
Manages registration and lookup of hooks for different events.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Callable, Union
|
||||
from functools import wraps
|
||||
|
||||
from .types import (
|
||||
HookEvent, HookDefinition, CommandHook, FunctionHook,
|
||||
HookInput, HookResult
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HookRegistry:
|
||||
"""
|
||||
Registry for managing hooks.
|
||||
|
||||
Provides methods to register, unregister, and lookup hooks
|
||||
for different events.
|
||||
|
||||
Example:
|
||||
registry = HookRegistry()
|
||||
|
||||
# Register a function hook using decorator
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def validate_tool(event_data):
|
||||
if event_data.tool_name == "dangerous":
|
||||
return HookResult.deny("Tool is dangerous")
|
||||
return HookResult.allow()
|
||||
|
||||
# Register a command hook
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="python /path/to/validator.py",
|
||||
matcher="write_*"
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(self, enabled: bool = True):
|
||||
"""
|
||||
Initialize the hook registry.
|
||||
|
||||
Args:
|
||||
enabled: Whether hooks are enabled globally
|
||||
"""
|
||||
self._hooks: Dict[HookEvent, List[HookDefinition]] = {
|
||||
event: [] for event in HookEvent
|
||||
}
|
||||
self._enabled = enabled
|
||||
self._global_timeout = 60.0
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Check if hooks are enabled."""
|
||||
return self._enabled
|
||||
|
||||
@enabled.setter
|
||||
def enabled(self, value: bool):
|
||||
"""Enable or disable hooks."""
|
||||
self._enabled = value
|
||||
|
||||
def register(self, hook: HookDefinition) -> str:
|
||||
"""
|
||||
Register a hook definition.
|
||||
|
||||
Args:
|
||||
hook: The hook definition to register
|
||||
|
||||
Returns:
|
||||
The hook ID
|
||||
"""
|
||||
self._hooks[hook.event].append(hook)
|
||||
logger.debug(f"Registered hook '{hook.name}' for event '{hook.event.value}'")
|
||||
return hook.id
|
||||
|
||||
def register_function(
|
||||
self,
|
||||
event: HookEvent,
|
||||
func: Callable[[HookInput], HookResult],
|
||||
matcher: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0,
|
||||
is_async: bool = False
|
||||
) -> str:
|
||||
"""
|
||||
Register a Python function as a hook.
|
||||
|
||||
Args:
|
||||
event: The event to hook
|
||||
func: The function to call
|
||||
matcher: Optional regex pattern to match targets
|
||||
name: Optional name for the hook
|
||||
description: Optional description
|
||||
sequential: Whether to run sequentially
|
||||
timeout: Timeout in seconds
|
||||
is_async: Whether the function is async
|
||||
|
||||
Returns:
|
||||
The hook ID
|
||||
"""
|
||||
hook = FunctionHook(
|
||||
event=event,
|
||||
func=func,
|
||||
matcher=matcher,
|
||||
name=name or func.__name__,
|
||||
description=description,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
is_async=is_async
|
||||
)
|
||||
return self.register(hook)
|
||||
|
||||
def register_command(
|
||||
self,
|
||||
event: HookEvent,
|
||||
command: str,
|
||||
matcher: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
shell: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Register a shell command as a hook.
|
||||
|
||||
Args:
|
||||
event: The event to hook
|
||||
command: The shell command to execute
|
||||
matcher: Optional regex pattern to match targets
|
||||
name: Optional name for the hook
|
||||
description: Optional description
|
||||
sequential: Whether to run sequentially
|
||||
timeout: Timeout in seconds
|
||||
env: Additional environment variables
|
||||
shell: Whether to run in shell mode
|
||||
|
||||
Returns:
|
||||
The hook ID
|
||||
"""
|
||||
hook = CommandHook(
|
||||
event=event,
|
||||
command=command,
|
||||
matcher=matcher,
|
||||
name=name,
|
||||
description=description,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
env=env or {},
|
||||
shell=shell
|
||||
)
|
||||
return self.register(hook)
|
||||
|
||||
def on(
|
||||
self,
|
||||
event: HookEvent,
|
||||
matcher: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0
|
||||
) -> Callable:
|
||||
"""
|
||||
Decorator to register a function as a hook.
|
||||
|
||||
Args:
|
||||
event: The event to hook
|
||||
matcher: Optional regex pattern to match targets
|
||||
sequential: Whether to run sequentially
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
Decorator function
|
||||
|
||||
Example:
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def my_hook(event_data):
|
||||
return HookResult.allow()
|
||||
"""
|
||||
def decorator(func: Callable[[HookInput], HookResult]) -> Callable:
|
||||
import asyncio
|
||||
is_async = asyncio.iscoroutinefunction(func)
|
||||
|
||||
self.register_function(
|
||||
event=event,
|
||||
func=func,
|
||||
matcher=matcher,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
is_async=is_async
|
||||
)
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
def unregister(self, hook_id: str) -> bool:
|
||||
"""
|
||||
Unregister a hook by ID.
|
||||
|
||||
Args:
|
||||
hook_id: The hook ID to unregister
|
||||
|
||||
Returns:
|
||||
True if found and removed, False otherwise
|
||||
"""
|
||||
for event in HookEvent:
|
||||
for i, hook in enumerate(self._hooks[event]):
|
||||
if hook.id == hook_id:
|
||||
self._hooks[event].pop(i)
|
||||
logger.debug(f"Unregistered hook '{hook_id}'")
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_hooks(
|
||||
self,
|
||||
event: HookEvent,
|
||||
target: Optional[str] = None
|
||||
) -> List[HookDefinition]:
|
||||
"""
|
||||
Get all hooks for an event, optionally filtered by target.
|
||||
|
||||
Args:
|
||||
event: The event to get hooks for
|
||||
target: Optional target to filter by (e.g., tool name)
|
||||
|
||||
Returns:
|
||||
List of matching hook definitions
|
||||
"""
|
||||
if not self._enabled:
|
||||
return []
|
||||
|
||||
hooks = self._hooks.get(event, [])
|
||||
|
||||
if target is None:
|
||||
return [h for h in hooks if h.enabled]
|
||||
|
||||
return [h for h in hooks if h.enabled and h.matches(target)]
|
||||
|
||||
def has_hooks(self, event: HookEvent) -> bool:
|
||||
"""Check if there are any hooks registered for an event."""
|
||||
return bool(self.get_hooks(event))
|
||||
|
||||
def clear(self, event: Optional[HookEvent] = None):
|
||||
"""
|
||||
Clear all hooks or hooks for a specific event.
|
||||
|
||||
Args:
|
||||
event: Optional event to clear hooks for
|
||||
"""
|
||||
if event is None:
|
||||
for e in HookEvent:
|
||||
self._hooks[e] = []
|
||||
else:
|
||||
self._hooks[event] = []
|
||||
|
||||
def list_hooks(self) -> Dict[str, List[Dict]]:
|
||||
"""
|
||||
List all registered hooks.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping event names to hook info
|
||||
"""
|
||||
result = {}
|
||||
for event in HookEvent:
|
||||
hooks = self._hooks[event]
|
||||
if hooks:
|
||||
result[event.value] = [
|
||||
{
|
||||
"id": h.id,
|
||||
"name": h.name,
|
||||
"type": "command" if isinstance(h, CommandHook) else "function",
|
||||
"matcher": h.matcher,
|
||||
"enabled": h.enabled,
|
||||
"sequential": h.sequential
|
||||
}
|
||||
for h in hooks
|
||||
]
|
||||
return result
|
||||
|
||||
def enable_hook(self, hook_id: str) -> bool:
|
||||
"""Enable a specific hook."""
|
||||
for event in HookEvent:
|
||||
for hook in self._hooks[event]:
|
||||
if hook.id == hook_id:
|
||||
hook.enabled = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable_hook(self, hook_id: str) -> bool:
|
||||
"""Disable a specific hook."""
|
||||
for event in HookEvent:
|
||||
for hook in self._hooks[event]:
|
||||
if hook.id == hook_id:
|
||||
hook.enabled = False
|
||||
return True
|
||||
return False
|
||||
|
||||
def set_global_timeout(self, timeout: float):
|
||||
"""Set the global timeout for all hooks."""
|
||||
self._global_timeout = timeout
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Get total number of registered hooks."""
|
||||
return sum(len(hooks) for hooks in self._hooks.values())
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"HookRegistry(enabled={self._enabled}, hooks={len(self)})"
|
||||
|
||||
|
||||
# Global default registry
|
||||
_default_registry: Optional[HookRegistry] = None
|
||||
|
||||
|
||||
def get_default_registry() -> HookRegistry:
|
||||
"""Get the default global hook registry."""
|
||||
global _default_registry
|
||||
if _default_registry is None:
|
||||
_default_registry = HookRegistry()
|
||||
return _default_registry
|
||||
|
||||
|
||||
def set_default_registry(registry: HookRegistry):
|
||||
"""Set the default global hook registry."""
|
||||
global _default_registry
|
||||
_default_registry = registry
|
||||
@@ -0,0 +1,450 @@
|
||||
"""
|
||||
Hook Runner for PraisonAI Agents.
|
||||
|
||||
Executes hooks registered in the registry, supporting both
|
||||
Python functions and shell commands.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import asyncio
|
||||
import logging
|
||||
import subprocess
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
from .types import (
|
||||
HookEvent, HookDefinition, CommandHook, FunctionHook,
|
||||
HookInput, HookResult, HookExecutionResult
|
||||
)
|
||||
from .registry import HookRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Exit codes for command hooks
|
||||
EXIT_CODE_SUCCESS = 0
|
||||
EXIT_CODE_BLOCKING_ERROR = 2
|
||||
EXIT_CODE_NON_BLOCKING_ERROR = 1
|
||||
|
||||
|
||||
class HookRunner:
|
||||
"""
|
||||
Executes hooks from a registry.
|
||||
|
||||
Supports:
|
||||
- Python function hooks (sync and async)
|
||||
- Shell command hooks
|
||||
- Sequential and parallel execution
|
||||
- Timeout handling
|
||||
- Result aggregation
|
||||
|
||||
Example:
|
||||
registry = HookRegistry()
|
||||
runner = HookRunner(registry)
|
||||
|
||||
# Execute hooks for an event
|
||||
results = await runner.execute(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
input_data=BeforeToolInput(
|
||||
tool_name="write_file",
|
||||
tool_input={"path": "/tmp/test.txt"}
|
||||
)
|
||||
)
|
||||
|
||||
# Check if any hook blocked
|
||||
if runner.is_blocked(results):
|
||||
print("Tool execution blocked by hook")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
registry: Optional[HookRegistry] = None,
|
||||
default_timeout: float = 60.0,
|
||||
cwd: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
Initialize the hook runner.
|
||||
|
||||
Args:
|
||||
registry: Hook registry to use
|
||||
default_timeout: Default timeout for hooks
|
||||
cwd: Working directory for command hooks
|
||||
"""
|
||||
self._registry = registry or HookRegistry()
|
||||
self._default_timeout = default_timeout
|
||||
self._cwd = cwd or os.getcwd()
|
||||
|
||||
@property
|
||||
def registry(self) -> HookRegistry:
|
||||
"""Get the hook registry."""
|
||||
return self._registry
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
target: Optional[str] = None
|
||||
) -> List[HookExecutionResult]:
|
||||
"""
|
||||
Execute all hooks for an event.
|
||||
|
||||
Args:
|
||||
event: The event to execute hooks for
|
||||
input_data: Input data for the hooks
|
||||
target: Optional target to filter hooks (e.g., tool name)
|
||||
|
||||
Returns:
|
||||
List of execution results
|
||||
"""
|
||||
hooks = self._registry.get_hooks(event, target)
|
||||
|
||||
if not hooks:
|
||||
return []
|
||||
|
||||
# Separate sequential and parallel hooks
|
||||
sequential_hooks = [h for h in hooks if h.sequential]
|
||||
parallel_hooks = [h for h in hooks if not h.sequential]
|
||||
|
||||
results = []
|
||||
|
||||
# Execute parallel hooks first
|
||||
if parallel_hooks:
|
||||
parallel_results = await self._execute_parallel(
|
||||
parallel_hooks, event, input_data
|
||||
)
|
||||
results.extend(parallel_results)
|
||||
|
||||
# Execute sequential hooks
|
||||
if sequential_hooks:
|
||||
sequential_results = await self._execute_sequential(
|
||||
sequential_hooks, event, input_data
|
||||
)
|
||||
results.extend(sequential_results)
|
||||
|
||||
return results
|
||||
|
||||
def execute_sync(
|
||||
self,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
target: Optional[str] = None
|
||||
) -> List[HookExecutionResult]:
|
||||
"""
|
||||
Synchronous version of execute.
|
||||
|
||||
Args:
|
||||
event: The event to execute hooks for
|
||||
input_data: Input data for the hooks
|
||||
target: Optional target to filter hooks
|
||||
|
||||
Returns:
|
||||
List of execution results
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# If we're in an async context, create a task
|
||||
future = asyncio.ensure_future(self.execute(event, input_data, target))
|
||||
return loop.run_until_complete(future)
|
||||
except RuntimeError:
|
||||
# No event loop running, create one
|
||||
return asyncio.run(self.execute(event, input_data, target))
|
||||
|
||||
async def _execute_parallel(
|
||||
self,
|
||||
hooks: List[HookDefinition],
|
||||
event: HookEvent,
|
||||
input_data: HookInput
|
||||
) -> List[HookExecutionResult]:
|
||||
"""Execute hooks in parallel."""
|
||||
tasks = [
|
||||
self._execute_single(hook, event, input_data)
|
||||
for hook in hooks
|
||||
]
|
||||
return await asyncio.gather(*tasks)
|
||||
|
||||
async def _execute_sequential(
|
||||
self,
|
||||
hooks: List[HookDefinition],
|
||||
event: HookEvent,
|
||||
input_data: HookInput
|
||||
) -> List[HookExecutionResult]:
|
||||
"""Execute hooks sequentially, passing modified input to next hook."""
|
||||
results = []
|
||||
current_input = input_data
|
||||
|
||||
for hook in hooks:
|
||||
result = await self._execute_single(hook, event, current_input)
|
||||
results.append(result)
|
||||
|
||||
# If hook modified the input, use it for next hook
|
||||
if result.success and result.output and result.output.modified_input:
|
||||
# Update input with modifications
|
||||
for key, value in result.output.modified_input.items():
|
||||
if hasattr(current_input, key):
|
||||
setattr(current_input, key, value)
|
||||
|
||||
# Stop if hook blocked
|
||||
if result.output and result.output.is_denied():
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
async def _execute_single(
|
||||
self,
|
||||
hook: HookDefinition,
|
||||
event: HookEvent,
|
||||
input_data: HookInput
|
||||
) -> HookExecutionResult:
|
||||
"""Execute a single hook."""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
if isinstance(hook, CommandHook):
|
||||
return await self._execute_command_hook(hook, event, input_data, start_time)
|
||||
elif isinstance(hook, FunctionHook):
|
||||
return await self._execute_function_hook(hook, event, input_data, start_time)
|
||||
else:
|
||||
raise ValueError(f"Unknown hook type: {type(hook)}")
|
||||
except Exception as e:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
logger.warning(f"Hook '{hook.name}' failed: {e}")
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=str(e),
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
async def _execute_command_hook(
|
||||
self,
|
||||
hook: CommandHook,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
start_time: float
|
||||
) -> HookExecutionResult:
|
||||
"""Execute a command hook."""
|
||||
timeout = hook.timeout or self._default_timeout
|
||||
|
||||
# Prepare environment
|
||||
env = os.environ.copy()
|
||||
env.update(hook.env)
|
||||
env["PRAISON_PROJECT_DIR"] = self._cwd
|
||||
env["PRAISON_EVENT"] = event.value
|
||||
|
||||
# Expand command variables
|
||||
command = self._expand_command(hook.command, input_data)
|
||||
|
||||
try:
|
||||
# Create subprocess
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
cwd=self._cwd
|
||||
)
|
||||
|
||||
# Send input as JSON
|
||||
input_json = json.dumps(input_data.to_dict()).encode()
|
||||
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
process.communicate(input=input_json),
|
||||
timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=f"Hook timed out after {timeout}s",
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
duration = (time.time() - start_time) * 1000
|
||||
exit_code = process.returncode
|
||||
stdout_str = stdout.decode() if stdout else ""
|
||||
stderr_str = stderr.decode() if stderr else ""
|
||||
|
||||
# Parse output
|
||||
output = self._parse_command_output(stdout_str, stderr_str, exit_code)
|
||||
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=exit_code == EXIT_CODE_SUCCESS,
|
||||
output=output,
|
||||
stdout=stdout_str,
|
||||
stderr=stderr_str,
|
||||
exit_code=exit_code,
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=str(e),
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
async def _execute_function_hook(
|
||||
self,
|
||||
hook: FunctionHook,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
start_time: float
|
||||
) -> HookExecutionResult:
|
||||
"""Execute a function hook."""
|
||||
if hook.func is None:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error="Hook function is None",
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
timeout = hook.timeout or self._default_timeout
|
||||
|
||||
try:
|
||||
if hook.is_async:
|
||||
# Execute async function with timeout
|
||||
result = await asyncio.wait_for(
|
||||
hook.func(input_data),
|
||||
timeout=timeout
|
||||
)
|
||||
else:
|
||||
# Execute sync function in thread pool with timeout
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await asyncio.wait_for(
|
||||
loop.run_in_executor(None, hook.func, input_data),
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
duration = (time.time() - start_time) * 1000
|
||||
|
||||
# Normalize result
|
||||
if result is None:
|
||||
result = HookResult.allow()
|
||||
elif not isinstance(result, HookResult):
|
||||
result = HookResult.allow()
|
||||
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=True,
|
||||
output=result,
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=f"Hook timed out after {timeout}s",
|
||||
duration_ms=duration
|
||||
)
|
||||
except Exception as e:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=str(e),
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
def _expand_command(self, command: str, input_data: HookInput) -> str:
|
||||
"""Expand variables in command string."""
|
||||
return command.replace(
|
||||
"$PRAISON_PROJECT_DIR", self._cwd
|
||||
).replace(
|
||||
"$PRAISON_CWD", input_data.cwd
|
||||
)
|
||||
|
||||
def _parse_command_output(
|
||||
self,
|
||||
stdout: str,
|
||||
stderr: str,
|
||||
exit_code: int
|
||||
) -> HookResult:
|
||||
"""Parse command output into HookResult."""
|
||||
# Try to parse JSON from stdout
|
||||
if stdout.strip():
|
||||
try:
|
||||
data = json.loads(stdout.strip())
|
||||
|
||||
# Handle double-encoded JSON
|
||||
if isinstance(data, str):
|
||||
data = json.loads(data)
|
||||
|
||||
return HookResult(
|
||||
decision=data.get("decision", "allow"),
|
||||
reason=data.get("reason"),
|
||||
modified_input=data.get("modified_input"),
|
||||
additional_context=data.get("additional_context"),
|
||||
suppress_output=data.get("suppress_output", False)
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Convert exit code to decision
|
||||
if exit_code == EXIT_CODE_SUCCESS:
|
||||
return HookResult(
|
||||
decision="allow",
|
||||
additional_context=stdout.strip() if stdout.strip() else None
|
||||
)
|
||||
elif exit_code == EXIT_CODE_BLOCKING_ERROR:
|
||||
return HookResult(
|
||||
decision="deny",
|
||||
reason=stderr.strip() or stdout.strip() or "Blocked by hook"
|
||||
)
|
||||
else:
|
||||
return HookResult(
|
||||
decision="allow",
|
||||
additional_context=f"Warning: {stderr.strip() or stdout.strip()}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_blocked(results: List[HookExecutionResult]) -> bool:
|
||||
"""Check if any hook blocked execution."""
|
||||
for result in results:
|
||||
if result.output and result.output.is_denied():
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def get_blocking_reason(results: List[HookExecutionResult]) -> Optional[str]:
|
||||
"""Get the reason for blocking from results."""
|
||||
for result in results:
|
||||
if result.output and result.output.is_denied():
|
||||
return result.output.reason
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def aggregate_context(results: List[HookExecutionResult]) -> Optional[str]:
|
||||
"""Aggregate additional context from all results."""
|
||||
contexts = []
|
||||
for result in results:
|
||||
if result.output and result.output.additional_context:
|
||||
contexts.append(result.output.additional_context)
|
||||
return "\n".join(contexts) if contexts else None
|
||||
@@ -0,0 +1,197 @@
|
||||
"""
|
||||
Hook Types for PraisonAI Agents.
|
||||
|
||||
Defines the core types, enums, and dataclasses for the hook system.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List, Callable, Union, Literal
|
||||
|
||||
|
||||
class HookEvent(str, Enum):
|
||||
"""Event names for the hook system."""
|
||||
BEFORE_TOOL = "before_tool"
|
||||
AFTER_TOOL = "after_tool"
|
||||
BEFORE_AGENT = "before_agent"
|
||||
AFTER_AGENT = "after_agent"
|
||||
SESSION_START = "session_start"
|
||||
SESSION_END = "session_end"
|
||||
BEFORE_LLM = "before_llm"
|
||||
AFTER_LLM = "after_llm"
|
||||
ON_ERROR = "on_error"
|
||||
ON_RETRY = "on_retry"
|
||||
|
||||
|
||||
# Decision types for hook outputs
|
||||
HookDecision = Literal["allow", "deny", "block", "ask", None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookInput:
|
||||
"""Base hook input - common fields for all events."""
|
||||
session_id: str
|
||||
cwd: str
|
||||
event_name: str
|
||||
timestamp: str
|
||||
agent_name: Optional[str] = None
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"cwd": self.cwd,
|
||||
"event_name": self.event_name,
|
||||
"timestamp": self.timestamp,
|
||||
"agent_name": self.agent_name,
|
||||
**self.extra
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookOutput:
|
||||
"""Base hook output - common fields for all events."""
|
||||
proceed: bool = True
|
||||
stop_reason: Optional[str] = None
|
||||
suppress_output: bool = False
|
||||
system_message: Optional[str] = None
|
||||
decision: HookDecision = None
|
||||
reason: Optional[str] = None
|
||||
modified_data: Optional[Dict[str, Any]] = None
|
||||
|
||||
def is_blocking(self) -> bool:
|
||||
"""Check if this output represents a blocking decision."""
|
||||
return self.decision in ("block", "deny")
|
||||
|
||||
def should_stop(self) -> bool:
|
||||
"""Check if execution should stop."""
|
||||
return not self.proceed
|
||||
|
||||
def get_reason(self) -> str:
|
||||
"""Get the effective reason for blocking or stopping."""
|
||||
return self.reason or self.stop_reason or "No reason provided"
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookResult:
|
||||
"""Result from a hook execution."""
|
||||
decision: HookDecision = "allow"
|
||||
reason: Optional[str] = None
|
||||
modified_input: Optional[Dict[str, Any]] = None
|
||||
additional_context: Optional[str] = None
|
||||
suppress_output: bool = False
|
||||
|
||||
@classmethod
|
||||
def allow(cls, reason: Optional[str] = None) -> "HookResult":
|
||||
"""Create an allow result."""
|
||||
return cls(decision="allow", reason=reason)
|
||||
|
||||
@classmethod
|
||||
def deny(cls, reason: str) -> "HookResult":
|
||||
"""Create a deny result."""
|
||||
return cls(decision="deny", reason=reason)
|
||||
|
||||
@classmethod
|
||||
def block(cls, reason: str) -> "HookResult":
|
||||
"""Create a block result."""
|
||||
return cls(decision="block", reason=reason)
|
||||
|
||||
@classmethod
|
||||
def ask(cls, reason: str) -> "HookResult":
|
||||
"""Create an ask result (requires user confirmation)."""
|
||||
return cls(decision="ask", reason=reason)
|
||||
|
||||
def is_allowed(self) -> bool:
|
||||
"""Check if the result allows execution."""
|
||||
return self.decision in ("allow", None)
|
||||
|
||||
def is_denied(self) -> bool:
|
||||
"""Check if the result denies execution."""
|
||||
return self.decision in ("deny", "block")
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookDefinition:
|
||||
"""Hook definition with matcher and configuration."""
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
|
||||
event: HookEvent = HookEvent.BEFORE_TOOL
|
||||
matcher: Optional[str] = None # Regex pattern to match tool names
|
||||
sequential: bool = False # Execute hooks sequentially vs parallel
|
||||
enabled: bool = True
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
timeout: float = 60.0 # Timeout in seconds
|
||||
|
||||
def matches(self, target: str) -> bool:
|
||||
"""Check if this hook matches the target (tool name, etc.)."""
|
||||
if self.matcher is None:
|
||||
return True
|
||||
|
||||
import re
|
||||
try:
|
||||
return bool(re.match(self.matcher, target))
|
||||
except re.error:
|
||||
# Invalid regex, fall back to simple string match
|
||||
return self.matcher in target
|
||||
|
||||
|
||||
@dataclass
|
||||
class CommandHook(HookDefinition):
|
||||
"""Hook that executes a shell command."""
|
||||
command: str = ""
|
||||
shell: bool = True
|
||||
env: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.name:
|
||||
self.name = f"command_hook_{self.id}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionHook(HookDefinition):
|
||||
"""Hook that executes a Python function."""
|
||||
func: Optional[Callable[[HookInput], HookResult]] = None
|
||||
is_async: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.name and self.func:
|
||||
self.name = self.func.__name__
|
||||
elif not self.name:
|
||||
self.name = f"function_hook_{self.id}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookExecutionResult:
|
||||
"""Result of executing a single hook."""
|
||||
hook_id: str
|
||||
hook_name: str
|
||||
event: HookEvent
|
||||
success: bool
|
||||
output: Optional[HookResult] = None
|
||||
stdout: Optional[str] = None
|
||||
stderr: Optional[str] = None
|
||||
exit_code: Optional[int] = None
|
||||
duration_ms: float = 0.0
|
||||
error: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"hook_id": self.hook_id,
|
||||
"hook_name": self.hook_name,
|
||||
"event": self.event.value,
|
||||
"success": self.success,
|
||||
"output": self.output.__dict__ if self.output else None,
|
||||
"stdout": self.stdout,
|
||||
"stderr": self.stderr,
|
||||
"exit_code": self.exit_code,
|
||||
"duration_ms": self.duration_ms,
|
||||
"error": self.error
|
||||
}
|
||||
|
||||
|
||||
# Type alias for hook functions
|
||||
HookFunction = Callable[[HookInput], Union[HookResult, None]]
|
||||
AsyncHookFunction = Callable[[HookInput], Any] # Returns Awaitable[HookResult]
|
||||
@@ -0,0 +1,65 @@
|
||||
"""
|
||||
LSP Integration Module for PraisonAI Agents.
|
||||
|
||||
Provides Language Server Protocol integration for code intelligence:
|
||||
- Diagnostics (errors, warnings)
|
||||
- Code completion suggestions
|
||||
- Go to definition
|
||||
- Find references
|
||||
- Hover information
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- LSP client only initialized when needed
|
||||
- No overhead when not in use
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.lsp import LSPClient
|
||||
|
||||
# Create client for a language
|
||||
client = LSPClient(language="python")
|
||||
|
||||
# Start the language server
|
||||
await client.start()
|
||||
|
||||
# Get diagnostics for a file
|
||||
diagnostics = await client.get_diagnostics("/path/to/file.py")
|
||||
|
||||
# Get completions at a position
|
||||
completions = await client.get_completions("/path/to/file.py", line=10, character=5)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core client
|
||||
"LSPClient",
|
||||
# Data types
|
||||
"Diagnostic",
|
||||
"DiagnosticSeverity",
|
||||
"CompletionItem",
|
||||
"Location",
|
||||
"Position",
|
||||
"Range",
|
||||
# Configuration
|
||||
"LSPConfig",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "LSPClient":
|
||||
from .client import LSPClient
|
||||
return LSPClient
|
||||
|
||||
if name in ("Diagnostic", "DiagnosticSeverity", "CompletionItem",
|
||||
"Location", "Position", "Range"):
|
||||
from .types import (
|
||||
Diagnostic, DiagnosticSeverity, CompletionItem,
|
||||
Location, Position, Range
|
||||
)
|
||||
return locals()[name]
|
||||
|
||||
if name == "LSPConfig":
|
||||
from .config import LSPConfig
|
||||
return LSPConfig
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,489 @@
|
||||
"""
|
||||
LSP Client for PraisonAI Agents.
|
||||
|
||||
Provides Language Server Protocol client implementation.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any
|
||||
from pathlib import Path
|
||||
|
||||
from .types import (
|
||||
Diagnostic, DiagnosticSeverity, CompletionItem, Location,
|
||||
Position, Range, TextDocumentItem, TextDocumentIdentifier,
|
||||
TextDocumentPositionParams
|
||||
)
|
||||
from .config import LSPConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LSPClient:
|
||||
"""
|
||||
Language Server Protocol client.
|
||||
|
||||
Communicates with language servers to provide code intelligence
|
||||
features like diagnostics, completions, and navigation.
|
||||
|
||||
Example:
|
||||
client = LSPClient(language="python")
|
||||
await client.start()
|
||||
|
||||
# Get diagnostics
|
||||
diagnostics = await client.get_diagnostics("/path/to/file.py")
|
||||
|
||||
# Get completions
|
||||
completions = await client.get_completions(
|
||||
"/path/to/file.py", line=10, character=5
|
||||
)
|
||||
|
||||
await client.stop()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
language: str,
|
||||
command: Optional[str] = None,
|
||||
args: Optional[List[str]] = None,
|
||||
root_uri: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
Initialize the LSP client.
|
||||
|
||||
Args:
|
||||
language: Programming language (python, javascript, etc.)
|
||||
command: Language server command (uses default if not specified)
|
||||
args: Command arguments
|
||||
root_uri: Workspace root URI
|
||||
"""
|
||||
self.config = LSPConfig(
|
||||
language=language,
|
||||
command=command,
|
||||
args=args or [],
|
||||
root_uri=root_uri
|
||||
)
|
||||
|
||||
self._process: Optional[asyncio.subprocess.Process] = None
|
||||
self._request_id = 0
|
||||
self._pending_requests: Dict[int, asyncio.Future] = {}
|
||||
self._diagnostics: Dict[str, List[Diagnostic]] = {}
|
||||
self._initialized = False
|
||||
self._reader_task: Optional[asyncio.Task] = None
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if the language server is running."""
|
||||
return self._process is not None and self._process.returncode is None
|
||||
|
||||
async def start(self) -> bool:
|
||||
"""
|
||||
Start the language server.
|
||||
|
||||
Returns:
|
||||
True if started successfully
|
||||
"""
|
||||
if self.is_running:
|
||||
return True
|
||||
|
||||
try:
|
||||
cmd = [self.config.command] + self.config.args
|
||||
|
||||
self._process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
# Start reader task
|
||||
self._reader_task = asyncio.create_task(self._read_responses())
|
||||
|
||||
# Initialize the server
|
||||
root_uri = self.config.root_uri or f"file://{os.getcwd()}"
|
||||
|
||||
result = await self._send_request("initialize", {
|
||||
"processId": os.getpid(),
|
||||
"rootUri": root_uri,
|
||||
"capabilities": {
|
||||
"textDocument": {
|
||||
"completion": {"completionItem": {"snippetSupport": True}},
|
||||
"hover": {},
|
||||
"definition": {},
|
||||
"references": {},
|
||||
"publishDiagnostics": {}
|
||||
}
|
||||
},
|
||||
"initializationOptions": self.config.initialization_options
|
||||
})
|
||||
|
||||
if result:
|
||||
await self._send_notification("initialized", {})
|
||||
self._initialized = True
|
||||
logger.info(f"LSP server started: {self.config.command}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except FileNotFoundError:
|
||||
logger.error(f"Language server not found: {self.config.command}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start LSP server: {e}")
|
||||
return False
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the language server."""
|
||||
if not self.is_running:
|
||||
return
|
||||
|
||||
try:
|
||||
await self._send_request("shutdown", None)
|
||||
await self._send_notification("exit", None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if self._reader_task:
|
||||
self._reader_task.cancel()
|
||||
try:
|
||||
await self._reader_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self._process:
|
||||
self._process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(self._process.wait(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
self._process.kill()
|
||||
|
||||
self._process = None
|
||||
self._initialized = False
|
||||
|
||||
async def open_document(self, file_path: str, text: Optional[str] = None) -> bool:
|
||||
"""
|
||||
Open a document in the language server.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
text: File content (reads from disk if not provided)
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if not self._initialized:
|
||||
return False
|
||||
|
||||
if text is None:
|
||||
try:
|
||||
with open(file_path, "r") as f:
|
||||
text = f.read()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to read file: {e}")
|
||||
return False
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
await self._send_notification("textDocument/didOpen", {
|
||||
"textDocument": TextDocumentItem(
|
||||
uri=uri,
|
||||
language_id=self.config.language,
|
||||
version=1,
|
||||
text=text
|
||||
).to_dict()
|
||||
})
|
||||
|
||||
return True
|
||||
|
||||
async def close_document(self, file_path: str):
|
||||
"""Close a document in the language server."""
|
||||
if not self._initialized:
|
||||
return
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
await self._send_notification("textDocument/didClose", {
|
||||
"textDocument": TextDocumentIdentifier(uri=uri).to_dict()
|
||||
})
|
||||
|
||||
async def get_diagnostics(self, file_path: str) -> List[Diagnostic]:
|
||||
"""
|
||||
Get diagnostics for a file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
List of diagnostics
|
||||
"""
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
return self._diagnostics.get(uri, [])
|
||||
|
||||
async def get_completions(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int
|
||||
) -> List[CompletionItem]:
|
||||
"""
|
||||
Get completions at a position.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
|
||||
Returns:
|
||||
List of completion items
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/completion", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
items = result.get("items", result) if isinstance(result, dict) else result
|
||||
|
||||
return [
|
||||
CompletionItem.from_dict(item)
|
||||
for item in (items if isinstance(items, list) else [])
|
||||
]
|
||||
|
||||
async def get_definition(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int
|
||||
) -> List[Location]:
|
||||
"""
|
||||
Get definition location(s) for a symbol.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
|
||||
Returns:
|
||||
List of locations
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/definition", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
if isinstance(result, dict):
|
||||
return [Location.from_dict(result)]
|
||||
elif isinstance(result, list):
|
||||
return [Location.from_dict(loc) for loc in result]
|
||||
|
||||
return []
|
||||
|
||||
async def get_references(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int,
|
||||
include_declaration: bool = True
|
||||
) -> List[Location]:
|
||||
"""
|
||||
Get references to a symbol.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
include_declaration: Include the declaration in results
|
||||
|
||||
Returns:
|
||||
List of locations
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/references", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character},
|
||||
"context": {"includeDeclaration": include_declaration}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
return [Location.from_dict(loc) for loc in result]
|
||||
|
||||
async def get_hover(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Get hover information for a position.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
|
||||
Returns:
|
||||
Hover content as string
|
||||
"""
|
||||
if not self._initialized:
|
||||
return None
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/hover", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return None
|
||||
|
||||
contents = result.get("contents")
|
||||
if isinstance(contents, str):
|
||||
return contents
|
||||
elif isinstance(contents, dict):
|
||||
return contents.get("value", str(contents))
|
||||
elif isinstance(contents, list):
|
||||
return "\n".join(
|
||||
c.get("value", str(c)) if isinstance(c, dict) else str(c)
|
||||
for c in contents
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
async def _send_request(self, method: str, params: Any) -> Any:
|
||||
"""Send a request and wait for response."""
|
||||
if not self.is_running:
|
||||
return None
|
||||
|
||||
self._request_id += 1
|
||||
request_id = self._request_id
|
||||
|
||||
message = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"method": method,
|
||||
"params": params
|
||||
}
|
||||
|
||||
future = asyncio.get_event_loop().create_future()
|
||||
self._pending_requests[request_id] = future
|
||||
|
||||
await self._send_message(message)
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(future, timeout=self.config.timeout)
|
||||
return result
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(f"Request timed out: {method}")
|
||||
self._pending_requests.pop(request_id, None)
|
||||
return None
|
||||
|
||||
async def _send_notification(self, method: str, params: Any):
|
||||
"""Send a notification (no response expected)."""
|
||||
if not self.is_running:
|
||||
return
|
||||
|
||||
message = {
|
||||
"jsonrpc": "2.0",
|
||||
"method": method,
|
||||
"params": params
|
||||
}
|
||||
|
||||
await self._send_message(message)
|
||||
|
||||
async def _send_message(self, message: Dict[str, Any]):
|
||||
"""Send a JSON-RPC message."""
|
||||
if self._process is None or self._process.stdin is None:
|
||||
return
|
||||
|
||||
content = json.dumps(message)
|
||||
header = f"Content-Length: {len(content)}\r\n\r\n"
|
||||
|
||||
self._process.stdin.write(header.encode() + content.encode())
|
||||
await self._process.stdin.drain()
|
||||
|
||||
async def _read_responses(self):
|
||||
"""Read responses from the language server."""
|
||||
if self._process is None or self._process.stdout is None:
|
||||
return
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Read header
|
||||
header = b""
|
||||
while b"\r\n\r\n" not in header:
|
||||
chunk = await self._process.stdout.read(1)
|
||||
if not chunk:
|
||||
return
|
||||
header += chunk
|
||||
|
||||
# Parse content length
|
||||
content_length = 0
|
||||
for line in header.decode().split("\r\n"):
|
||||
if line.startswith("Content-Length:"):
|
||||
content_length = int(line.split(":")[1].strip())
|
||||
break
|
||||
|
||||
if content_length == 0:
|
||||
continue
|
||||
|
||||
# Read content
|
||||
content = await self._process.stdout.read(content_length)
|
||||
message = json.loads(content.decode())
|
||||
|
||||
await self._handle_message(message)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error reading LSP response: {e}")
|
||||
|
||||
async def _handle_message(self, message: Dict[str, Any]):
|
||||
"""Handle an incoming message."""
|
||||
if "id" in message:
|
||||
# Response to a request
|
||||
request_id = message["id"]
|
||||
if request_id in self._pending_requests:
|
||||
future = self._pending_requests.pop(request_id)
|
||||
if "error" in message:
|
||||
future.set_exception(
|
||||
RuntimeError(message["error"].get("message", "Unknown error"))
|
||||
)
|
||||
else:
|
||||
future.set_result(message.get("result"))
|
||||
|
||||
elif "method" in message:
|
||||
# Notification from server
|
||||
method = message["method"]
|
||||
params = message.get("params", {})
|
||||
|
||||
if method == "textDocument/publishDiagnostics":
|
||||
uri = params.get("uri", "")
|
||||
diagnostics = [
|
||||
Diagnostic.from_dict(d)
|
||||
for d in params.get("diagnostics", [])
|
||||
]
|
||||
self._diagnostics[uri] = diagnostics
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
LSP Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, List
|
||||
|
||||
|
||||
# Default language server commands
|
||||
DEFAULT_SERVERS = {
|
||||
"python": {
|
||||
"command": "pylsp",
|
||||
"args": [],
|
||||
},
|
||||
"javascript": {
|
||||
"command": "typescript-language-server",
|
||||
"args": ["--stdio"],
|
||||
},
|
||||
"typescript": {
|
||||
"command": "typescript-language-server",
|
||||
"args": ["--stdio"],
|
||||
},
|
||||
"rust": {
|
||||
"command": "rust-analyzer",
|
||||
"args": [],
|
||||
},
|
||||
"go": {
|
||||
"command": "gopls",
|
||||
"args": [],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class LSPConfig:
|
||||
"""Configuration for LSP client."""
|
||||
language: str
|
||||
command: Optional[str] = None
|
||||
args: List[str] = field(default_factory=list)
|
||||
root_uri: Optional[str] = None
|
||||
initialization_options: Dict = field(default_factory=dict)
|
||||
timeout: float = 30.0
|
||||
|
||||
def __post_init__(self):
|
||||
# Use default server if not specified
|
||||
if self.command is None:
|
||||
if self.language in DEFAULT_SERVERS:
|
||||
server = DEFAULT_SERVERS[self.language]
|
||||
self.command = server["command"]
|
||||
self.args = server["args"]
|
||||
else:
|
||||
raise ValueError(f"No default server for language: {self.language}")
|
||||
@@ -0,0 +1,162 @@
|
||||
"""
|
||||
LSP Types for PraisonAI Agents.
|
||||
|
||||
Defines the core types for Language Server Protocol integration.
|
||||
"""
|
||||
|
||||
from enum import IntEnum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
|
||||
class DiagnosticSeverity(IntEnum):
|
||||
"""Severity of a diagnostic."""
|
||||
ERROR = 1
|
||||
WARNING = 2
|
||||
INFORMATION = 3
|
||||
HINT = 4
|
||||
|
||||
|
||||
@dataclass
|
||||
class Position:
|
||||
"""Position in a text document (0-indexed)."""
|
||||
line: int
|
||||
character: int
|
||||
|
||||
def to_dict(self) -> Dict[str, int]:
|
||||
return {"line": self.line, "character": self.character}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, int]) -> "Position":
|
||||
return cls(line=data["line"], character=data["character"])
|
||||
|
||||
|
||||
@dataclass
|
||||
class Range:
|
||||
"""Range in a text document."""
|
||||
start: Position
|
||||
end: Position
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {"start": self.start.to_dict(), "end": self.end.to_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Range":
|
||||
return cls(
|
||||
start=Position.from_dict(data["start"]),
|
||||
end=Position.from_dict(data["end"])
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Location:
|
||||
"""Location in a document."""
|
||||
uri: str
|
||||
range: Range
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {"uri": self.uri, "range": self.range.to_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Location":
|
||||
return cls(uri=data["uri"], range=Range.from_dict(data["range"]))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Diagnostic:
|
||||
"""Represents a diagnostic (error, warning, etc.)."""
|
||||
range: Range
|
||||
message: str
|
||||
severity: DiagnosticSeverity = DiagnosticSeverity.ERROR
|
||||
code: Optional[str] = None
|
||||
source: Optional[str] = None
|
||||
related_information: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"range": self.range.to_dict(),
|
||||
"message": self.message,
|
||||
"severity": self.severity.value,
|
||||
"code": self.code,
|
||||
"source": self.source
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Diagnostic":
|
||||
return cls(
|
||||
range=Range.from_dict(data["range"]),
|
||||
message=data["message"],
|
||||
severity=DiagnosticSeverity(data.get("severity", 1)),
|
||||
code=data.get("code"),
|
||||
source=data.get("source")
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompletionItem:
|
||||
"""Represents a completion item."""
|
||||
label: str
|
||||
kind: int = 1 # Text
|
||||
detail: Optional[str] = None
|
||||
documentation: Optional[str] = None
|
||||
insert_text: Optional[str] = None
|
||||
sort_text: Optional[str] = None
|
||||
filter_text: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"label": self.label,
|
||||
"kind": self.kind,
|
||||
"detail": self.detail,
|
||||
"documentation": self.documentation,
|
||||
"insertText": self.insert_text
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "CompletionItem":
|
||||
return cls(
|
||||
label=data["label"],
|
||||
kind=data.get("kind", 1),
|
||||
detail=data.get("detail"),
|
||||
documentation=data.get("documentation"),
|
||||
insert_text=data.get("insertText")
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextDocumentItem:
|
||||
"""Represents a text document."""
|
||||
uri: str
|
||||
language_id: str
|
||||
version: int
|
||||
text: str
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"uri": self.uri,
|
||||
"languageId": self.language_id,
|
||||
"version": self.version,
|
||||
"text": self.text
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextDocumentIdentifier:
|
||||
"""Identifies a text document."""
|
||||
uri: str
|
||||
|
||||
def to_dict(self) -> Dict[str, str]:
|
||||
return {"uri": self.uri}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextDocumentPositionParams:
|
||||
"""Parameters for text document position requests."""
|
||||
text_document: TextDocumentIdentifier
|
||||
position: Position
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"textDocument": self.text_document.to_dict(),
|
||||
"position": self.position.to_dict()
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
"""
|
||||
Exec Policy Engine for PraisonAI Agents.
|
||||
|
||||
Provides policy-based execution control:
|
||||
- Define rules for what agents can/cannot do
|
||||
- Tool execution policies
|
||||
- Resource access control
|
||||
- Rate limiting and quotas
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Policies only evaluated when enabled
|
||||
- No overhead when no policies defined
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.policy import PolicyEngine, Policy, PolicyRule
|
||||
|
||||
# Create a policy engine
|
||||
engine = PolicyEngine()
|
||||
|
||||
# Add a policy
|
||||
policy = Policy(
|
||||
name="no_delete",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action="deny",
|
||||
resource="tool:delete_*",
|
||||
reason="Delete operations are not allowed"
|
||||
)
|
||||
]
|
||||
)
|
||||
engine.add_policy(policy)
|
||||
|
||||
# Check if action is allowed
|
||||
result = engine.check("tool:delete_file", context={})
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core classes
|
||||
"PolicyEngine",
|
||||
"Policy",
|
||||
"PolicyRule",
|
||||
# Result types
|
||||
"PolicyResult",
|
||||
"PolicyAction",
|
||||
# Configuration
|
||||
"PolicyConfig",
|
||||
# Convenience functions
|
||||
"create_deny_tools_policy",
|
||||
"create_allow_tools_policy",
|
||||
"create_read_only_policy",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "PolicyEngine":
|
||||
from .engine import PolicyEngine
|
||||
return PolicyEngine
|
||||
|
||||
if name in ("Policy", "PolicyRule"):
|
||||
from .policy import Policy, PolicyRule
|
||||
return locals()[name]
|
||||
|
||||
if name in ("PolicyResult", "PolicyAction"):
|
||||
from .types import PolicyResult, PolicyAction
|
||||
return locals()[name]
|
||||
|
||||
if name == "PolicyConfig":
|
||||
from .config import PolicyConfig
|
||||
return PolicyConfig
|
||||
|
||||
if name in ("create_deny_tools_policy", "create_allow_tools_policy", "create_read_only_policy"):
|
||||
from .engine import create_deny_tools_policy, create_allow_tools_policy, create_read_only_policy
|
||||
return locals()[name]
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
Policy Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, List
|
||||
|
||||
|
||||
@dataclass
|
||||
class PolicyConfig:
|
||||
"""Configuration for the policy engine."""
|
||||
enabled: bool = True
|
||||
default_action: str = "allow" # Default action when no policy matches
|
||||
log_decisions: bool = True
|
||||
strict_mode: bool = False # If True, deny when no policy matches
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
Policy Engine for PraisonAI Agents.
|
||||
|
||||
Manages and evaluates policies for execution control.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
from .types import PolicyAction, PolicyResult
|
||||
from .policy import Policy, PolicyRule
|
||||
from .config import PolicyConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PolicyEngine:
|
||||
"""
|
||||
Policy Engine for execution control.
|
||||
|
||||
Manages policies and evaluates them against resources
|
||||
to determine what actions are allowed.
|
||||
|
||||
Example:
|
||||
engine = PolicyEngine()
|
||||
|
||||
# Add a policy to block delete operations
|
||||
engine.add_policy(Policy(
|
||||
name="no_delete",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="Delete operations are not allowed"
|
||||
)
|
||||
]
|
||||
))
|
||||
|
||||
# Check if action is allowed
|
||||
result = engine.check("tool:delete_file", {})
|
||||
if not result.allowed:
|
||||
print(f"Denied: {result.reason}")
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[PolicyConfig] = None):
|
||||
"""
|
||||
Initialize the policy engine.
|
||||
|
||||
Args:
|
||||
config: Engine configuration
|
||||
"""
|
||||
self.config = config or PolicyConfig()
|
||||
self._policies: Dict[str, Policy] = {}
|
||||
|
||||
@property
|
||||
def policies(self) -> List[Policy]:
|
||||
"""Get all policies sorted by priority."""
|
||||
return sorted(
|
||||
self._policies.values(),
|
||||
key=lambda p: p.priority,
|
||||
reverse=True
|
||||
)
|
||||
|
||||
def add_policy(self, policy: Policy):
|
||||
"""
|
||||
Add a policy to the engine.
|
||||
|
||||
Args:
|
||||
policy: Policy to add
|
||||
"""
|
||||
self._policies[policy.name] = policy
|
||||
logger.debug(f"Added policy: {policy.name}")
|
||||
|
||||
def remove_policy(self, name: str) -> bool:
|
||||
"""
|
||||
Remove a policy by name.
|
||||
|
||||
Args:
|
||||
name: Policy name
|
||||
|
||||
Returns:
|
||||
True if removed, False if not found
|
||||
"""
|
||||
if name in self._policies:
|
||||
del self._policies[name]
|
||||
logger.debug(f"Removed policy: {name}")
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_policy(self, name: str) -> Optional[Policy]:
|
||||
"""Get a policy by name."""
|
||||
return self._policies.get(name)
|
||||
|
||||
def enable_policy(self, name: str) -> bool:
|
||||
"""Enable a policy."""
|
||||
policy = self._policies.get(name)
|
||||
if policy:
|
||||
policy.enabled = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable_policy(self, name: str) -> bool:
|
||||
"""Disable a policy."""
|
||||
policy = self._policies.get(name)
|
||||
if policy:
|
||||
policy.enabled = False
|
||||
return True
|
||||
return False
|
||||
|
||||
def check(
|
||||
self,
|
||||
resource: str,
|
||||
context: Optional[Dict[str, Any]] = None
|
||||
) -> PolicyResult:
|
||||
"""
|
||||
Check if an action is allowed.
|
||||
|
||||
Args:
|
||||
resource: Resource identifier (e.g., "tool:read_file")
|
||||
context: Additional context for policy evaluation
|
||||
|
||||
Returns:
|
||||
PolicyResult indicating if action is allowed
|
||||
"""
|
||||
if not self.config.enabled:
|
||||
return PolicyResult.allow("Policy engine disabled")
|
||||
|
||||
context = context or {}
|
||||
|
||||
# Evaluate policies in priority order
|
||||
for policy in self.policies:
|
||||
result = policy.evaluate(resource, context)
|
||||
if result is not None:
|
||||
if self.config.log_decisions:
|
||||
logger.info(
|
||||
f"Policy decision: {result.action.value} for {resource} "
|
||||
f"(policy: {result.policy_name}, reason: {result.reason})"
|
||||
)
|
||||
return result
|
||||
|
||||
# No policy matched - use default action
|
||||
if self.config.strict_mode:
|
||||
return PolicyResult.deny(
|
||||
"No policy matched and strict mode is enabled"
|
||||
)
|
||||
|
||||
return PolicyResult.allow("No policy matched")
|
||||
|
||||
def check_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
tool_input: Optional[Dict[str, Any]] = None
|
||||
) -> PolicyResult:
|
||||
"""
|
||||
Check if a tool execution is allowed.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
tool_input: Tool input parameters
|
||||
|
||||
Returns:
|
||||
PolicyResult
|
||||
"""
|
||||
return self.check(
|
||||
f"tool:{tool_name}",
|
||||
{"tool_name": tool_name, "tool_input": tool_input or {}}
|
||||
)
|
||||
|
||||
def check_file(
|
||||
self,
|
||||
operation: str,
|
||||
file_path: str
|
||||
) -> PolicyResult:
|
||||
"""
|
||||
Check if a file operation is allowed.
|
||||
|
||||
Args:
|
||||
operation: Operation type (read, write, delete)
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
PolicyResult
|
||||
"""
|
||||
return self.check(
|
||||
f"file:{operation}",
|
||||
{"operation": operation, "file_path": file_path}
|
||||
)
|
||||
|
||||
def clear(self):
|
||||
"""Remove all policies."""
|
||||
self._policies.clear()
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert engine state to dictionary."""
|
||||
return {
|
||||
"config": {
|
||||
"enabled": self.config.enabled,
|
||||
"default_action": self.config.default_action,
|
||||
"strict_mode": self.config.strict_mode
|
||||
},
|
||||
"policies": [p.to_dict() for p in self.policies]
|
||||
}
|
||||
|
||||
def load_from_dict(self, data: Dict[str, Any]):
|
||||
"""Load policies from a dictionary."""
|
||||
for policy_data in data.get("policies", []):
|
||||
policy = Policy.from_dict(policy_data)
|
||||
self.add_policy(policy)
|
||||
|
||||
|
||||
# Convenience functions for common policies
|
||||
|
||||
def create_deny_tools_policy(
|
||||
tool_patterns: List[str],
|
||||
name: str = "deny_tools",
|
||||
reason: str = "Tool not allowed"
|
||||
) -> Policy:
|
||||
"""
|
||||
Create a policy that denies specific tools.
|
||||
|
||||
Args:
|
||||
tool_patterns: List of tool name patterns to deny
|
||||
name: Policy name
|
||||
reason: Reason for denial
|
||||
|
||||
Returns:
|
||||
Policy
|
||||
"""
|
||||
rules = [
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource=f"tool:{pattern}",
|
||||
reason=reason,
|
||||
name=f"deny_{pattern}"
|
||||
)
|
||||
for pattern in tool_patterns
|
||||
]
|
||||
|
||||
return Policy(name=name, rules=rules)
|
||||
|
||||
|
||||
def create_allow_tools_policy(
|
||||
tool_patterns: List[str],
|
||||
name: str = "allow_tools"
|
||||
) -> Policy:
|
||||
"""
|
||||
Create a policy that allows specific tools.
|
||||
|
||||
Args:
|
||||
tool_patterns: List of tool name patterns to allow
|
||||
name: Policy name
|
||||
|
||||
Returns:
|
||||
Policy
|
||||
"""
|
||||
rules = [
|
||||
PolicyRule(
|
||||
action=PolicyAction.ALLOW,
|
||||
resource=f"tool:{pattern}",
|
||||
name=f"allow_{pattern}"
|
||||
)
|
||||
for pattern in tool_patterns
|
||||
]
|
||||
|
||||
return Policy(name=name, rules=rules, priority=10)
|
||||
|
||||
|
||||
def create_read_only_policy(name: str = "read_only") -> Policy:
|
||||
"""
|
||||
Create a read-only policy that denies write operations.
|
||||
|
||||
Args:
|
||||
name: Policy name
|
||||
|
||||
Returns:
|
||||
Policy
|
||||
"""
|
||||
return Policy(
|
||||
name=name,
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="file:write",
|
||||
reason="Write operations not allowed in read-only mode",
|
||||
name="deny_write"
|
||||
),
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="file:delete",
|
||||
reason="Delete operations not allowed in read-only mode",
|
||||
name="deny_delete"
|
||||
),
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:*_write*",
|
||||
reason="Write tools not allowed in read-only mode",
|
||||
name="deny_write_tools"
|
||||
),
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:*_delete*",
|
||||
reason="Delete tools not allowed in read-only mode",
|
||||
name="deny_delete_tools"
|
||||
),
|
||||
],
|
||||
priority=100
|
||||
)
|
||||
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
Policy and PolicyRule for PraisonAI Agents.
|
||||
|
||||
Defines policies and rules for execution control.
|
||||
"""
|
||||
|
||||
import re
|
||||
import fnmatch
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List, Callable
|
||||
|
||||
from .types import PolicyAction, PolicyResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class PolicyRule:
|
||||
"""
|
||||
A single rule within a policy.
|
||||
|
||||
Rules match against resources and contexts to determine
|
||||
what action to take.
|
||||
"""
|
||||
action: PolicyAction
|
||||
resource: Optional[str] = None # Pattern to match (supports wildcards)
|
||||
condition: Optional[Callable[[Dict[str, Any]], bool]] = None
|
||||
reason: Optional[str] = None
|
||||
name: Optional[str] = None
|
||||
priority: int = 0 # Higher priority rules are evaluated first
|
||||
|
||||
def matches(self, resource: str, context: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Check if this rule matches the given resource and context.
|
||||
|
||||
Args:
|
||||
resource: Resource identifier (e.g., "tool:read_file")
|
||||
context: Additional context for the check
|
||||
|
||||
Returns:
|
||||
True if the rule matches
|
||||
"""
|
||||
# Check resource pattern
|
||||
if self.resource:
|
||||
if not self._match_pattern(self.resource, resource):
|
||||
return False
|
||||
|
||||
# Check condition
|
||||
if self.condition:
|
||||
try:
|
||||
if not self.condition(context):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _match_pattern(self, pattern: str, value: str) -> bool:
|
||||
"""Match a pattern against a value (supports wildcards)."""
|
||||
# Convert glob pattern to regex
|
||||
if "*" in pattern or "?" in pattern:
|
||||
return fnmatch.fnmatch(value, pattern)
|
||||
return pattern == value
|
||||
|
||||
def evaluate(self, resource: str, context: Dict[str, Any]) -> Optional[PolicyResult]:
|
||||
"""
|
||||
Evaluate this rule against a resource.
|
||||
|
||||
Args:
|
||||
resource: Resource identifier
|
||||
context: Additional context
|
||||
|
||||
Returns:
|
||||
PolicyResult if rule matches, None otherwise
|
||||
"""
|
||||
if not self.matches(resource, context):
|
||||
return None
|
||||
|
||||
return PolicyResult(
|
||||
allowed=self.action in (PolicyAction.ALLOW, PolicyAction.LOG),
|
||||
action=self.action,
|
||||
rule_name=self.name,
|
||||
reason=self.reason
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"action": self.action.value,
|
||||
"resource": self.resource,
|
||||
"reason": self.reason,
|
||||
"name": self.name,
|
||||
"priority": self.priority
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Policy:
|
||||
"""
|
||||
A collection of rules that define execution policy.
|
||||
|
||||
Policies are evaluated in order, with higher priority rules
|
||||
taking precedence.
|
||||
"""
|
||||
name: str
|
||||
rules: List[PolicyRule] = field(default_factory=list)
|
||||
description: Optional[str] = None
|
||||
enabled: bool = True
|
||||
priority: int = 0 # Higher priority policies are evaluated first
|
||||
|
||||
def add_rule(self, rule: PolicyRule):
|
||||
"""Add a rule to this policy."""
|
||||
self.rules.append(rule)
|
||||
# Sort by priority (descending)
|
||||
self.rules.sort(key=lambda r: r.priority, reverse=True)
|
||||
|
||||
def remove_rule(self, rule_name: str) -> bool:
|
||||
"""Remove a rule by name."""
|
||||
for i, rule in enumerate(self.rules):
|
||||
if rule.name == rule_name:
|
||||
self.rules.pop(i)
|
||||
return True
|
||||
return False
|
||||
|
||||
def evaluate(self, resource: str, context: Dict[str, Any]) -> Optional[PolicyResult]:
|
||||
"""
|
||||
Evaluate this policy against a resource.
|
||||
|
||||
Args:
|
||||
resource: Resource identifier
|
||||
context: Additional context
|
||||
|
||||
Returns:
|
||||
PolicyResult if any rule matches, None otherwise
|
||||
"""
|
||||
if not self.enabled:
|
||||
return None
|
||||
|
||||
for rule in self.rules:
|
||||
result = rule.evaluate(resource, context)
|
||||
if result is not None:
|
||||
result.policy_name = self.name
|
||||
return result
|
||||
|
||||
return None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"enabled": self.enabled,
|
||||
"priority": self.priority,
|
||||
"rules": [r.to_dict() for r in self.rules]
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Policy":
|
||||
"""Create a policy from a dictionary."""
|
||||
rules = []
|
||||
for rule_data in data.get("rules", []):
|
||||
rules.append(PolicyRule(
|
||||
action=PolicyAction(rule_data["action"]),
|
||||
resource=rule_data.get("resource"),
|
||||
reason=rule_data.get("reason"),
|
||||
name=rule_data.get("name"),
|
||||
priority=rule_data.get("priority", 0)
|
||||
))
|
||||
|
||||
return cls(
|
||||
name=data["name"],
|
||||
rules=rules,
|
||||
description=data.get("description"),
|
||||
enabled=data.get("enabled", True),
|
||||
priority=data.get("priority", 0)
|
||||
)
|
||||
@@ -0,0 +1,60 @@
|
||||
"""
|
||||
Policy Types for PraisonAI Agents.
|
||||
|
||||
Defines the core types for the policy engine.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
|
||||
class PolicyAction(str, Enum):
|
||||
"""Action to take when a policy matches."""
|
||||
ALLOW = "allow"
|
||||
DENY = "deny"
|
||||
ASK = "ask" # Require user confirmation
|
||||
LOG = "log" # Allow but log the action
|
||||
RATE_LIMIT = "rate_limit"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PolicyResult:
|
||||
"""Result of a policy check."""
|
||||
allowed: bool
|
||||
action: PolicyAction
|
||||
policy_name: Optional[str] = None
|
||||
rule_name: Optional[str] = None
|
||||
reason: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def allow(cls, reason: Optional[str] = None) -> "PolicyResult":
|
||||
"""Create an allow result."""
|
||||
return cls(allowed=True, action=PolicyAction.ALLOW, reason=reason)
|
||||
|
||||
@classmethod
|
||||
def deny(cls, reason: str, policy_name: Optional[str] = None) -> "PolicyResult":
|
||||
"""Create a deny result."""
|
||||
return cls(
|
||||
allowed=False,
|
||||
action=PolicyAction.DENY,
|
||||
policy_name=policy_name,
|
||||
reason=reason
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def ask(cls, reason: str) -> "PolicyResult":
|
||||
"""Create an ask result (requires confirmation)."""
|
||||
return cls(allowed=False, action=PolicyAction.ASK, reason=reason)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"allowed": self.allowed,
|
||||
"action": self.action.value,
|
||||
"policy_name": self.policy_name,
|
||||
"rule_name": self.rule_name,
|
||||
"reason": self.reason,
|
||||
"metadata": self.metadata
|
||||
}
|
||||
@@ -28,20 +28,31 @@ def format_skill_for_prompt(skill: SkillMetadata) -> str:
|
||||
</skill>"""
|
||||
|
||||
|
||||
def generate_skills_xml(skills: List[SkillMetadata]) -> str:
|
||||
def generate_skills_xml(skills: List[SkillMetadata], working_directory: str = None) -> str:
|
||||
"""Generate XML block for available skills.
|
||||
|
||||
Args:
|
||||
skills: List of SkillMetadata instances
|
||||
working_directory: Current working directory for path resolution
|
||||
|
||||
Returns:
|
||||
XML string with <available_skills> block
|
||||
"""
|
||||
import os
|
||||
cwd = working_directory or os.getcwd()
|
||||
|
||||
if not skills:
|
||||
return "<available_skills>\n</available_skills>"
|
||||
return f"""<available_skills>
|
||||
<working_directory>{html.escape(cwd)}</working_directory>
|
||||
<note>When using run_skill_script, pass file paths relative to the working directory or as absolute paths. The tool will resolve relative paths automatically.</note>
|
||||
</available_skills>"""
|
||||
|
||||
skill_entries = "\n".join(format_skill_for_prompt(s) for s in skills)
|
||||
return f"<available_skills>\n{skill_entries}\n</available_skills>"
|
||||
return f"""<available_skills>
|
||||
<working_directory>{html.escape(cwd)}</working_directory>
|
||||
<note>When using run_skill_script, pass file paths relative to the working directory or as absolute paths. The tool will resolve relative paths automatically.</note>
|
||||
{skill_entries}
|
||||
</available_skills>"""
|
||||
|
||||
|
||||
def to_prompt(skill_dirs: List[Path]) -> str:
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
"""
|
||||
Extended Thinking Budgets Module for PraisonAI Agents.
|
||||
|
||||
Provides configurable thinking budgets for LLM reasoning:
|
||||
- Token budgets for extended thinking
|
||||
- Time budgets for reasoning
|
||||
- Adaptive budget allocation
|
||||
- Budget tracking and reporting
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Budgets only applied when configured
|
||||
- No overhead when not in use
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.thinking import ThinkingBudget, ThinkingConfig
|
||||
|
||||
# Create a thinking budget
|
||||
budget = ThinkingBudget(
|
||||
max_tokens=16000,
|
||||
max_time_seconds=60,
|
||||
adaptive=True
|
||||
)
|
||||
|
||||
# Apply to agent
|
||||
agent = Agent(
|
||||
instructions="...",
|
||||
thinking_budget=budget
|
||||
)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core classes
|
||||
"ThinkingBudget",
|
||||
"ThinkingConfig",
|
||||
# Tracking
|
||||
"ThinkingUsage",
|
||||
"ThinkingTracker",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "ThinkingBudget":
|
||||
from .budget import ThinkingBudget
|
||||
return ThinkingBudget
|
||||
|
||||
if name == "ThinkingConfig":
|
||||
from .config import ThinkingConfig
|
||||
return ThinkingConfig
|
||||
|
||||
if name == "ThinkingUsage":
|
||||
from .tracker import ThinkingUsage
|
||||
return ThinkingUsage
|
||||
|
||||
if name == "ThinkingTracker":
|
||||
from .tracker import ThinkingTracker
|
||||
return ThinkingTracker
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,131 @@
|
||||
"""
|
||||
Thinking Budget for PraisonAI Agents.
|
||||
|
||||
Defines budget constraints for extended thinking.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class BudgetLevel(str, Enum):
|
||||
"""Predefined budget levels."""
|
||||
MINIMAL = "minimal" # 2000 tokens
|
||||
LOW = "low" # 4000 tokens
|
||||
MEDIUM = "medium" # 8000 tokens
|
||||
HIGH = "high" # 16000 tokens
|
||||
MAXIMUM = "maximum" # 32000 tokens
|
||||
|
||||
|
||||
# Token allocations for each level
|
||||
BUDGET_TOKENS = {
|
||||
BudgetLevel.MINIMAL: 2000,
|
||||
BudgetLevel.LOW: 4000,
|
||||
BudgetLevel.MEDIUM: 8000,
|
||||
BudgetLevel.HIGH: 16000,
|
||||
BudgetLevel.MAXIMUM: 32000,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ThinkingBudget:
|
||||
"""
|
||||
Budget constraints for extended thinking.
|
||||
|
||||
Controls how much thinking/reasoning the LLM can do
|
||||
before producing a response.
|
||||
"""
|
||||
max_tokens: int = 8000
|
||||
max_time_seconds: Optional[float] = None
|
||||
adaptive: bool = True
|
||||
level: Optional[BudgetLevel] = None
|
||||
|
||||
# Adaptive settings
|
||||
min_tokens: int = 1000
|
||||
complexity_multiplier: float = 1.0
|
||||
|
||||
def __post_init__(self):
|
||||
# If level is specified, use its token allocation
|
||||
if self.level:
|
||||
self.max_tokens = BUDGET_TOKENS.get(self.level, self.max_tokens)
|
||||
|
||||
@classmethod
|
||||
def from_level(cls, level: BudgetLevel, **kwargs) -> "ThinkingBudget":
|
||||
"""Create a budget from a predefined level."""
|
||||
return cls(
|
||||
max_tokens=BUDGET_TOKENS[level],
|
||||
level=level,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def minimal(cls) -> "ThinkingBudget":
|
||||
"""Create a minimal budget."""
|
||||
return cls.from_level(BudgetLevel.MINIMAL)
|
||||
|
||||
@classmethod
|
||||
def low(cls) -> "ThinkingBudget":
|
||||
"""Create a low budget."""
|
||||
return cls.from_level(BudgetLevel.LOW)
|
||||
|
||||
@classmethod
|
||||
def medium(cls) -> "ThinkingBudget":
|
||||
"""Create a medium budget."""
|
||||
return cls.from_level(BudgetLevel.MEDIUM)
|
||||
|
||||
@classmethod
|
||||
def high(cls) -> "ThinkingBudget":
|
||||
"""Create a high budget."""
|
||||
return cls.from_level(BudgetLevel.HIGH)
|
||||
|
||||
@classmethod
|
||||
def maximum(cls) -> "ThinkingBudget":
|
||||
"""Create a maximum budget."""
|
||||
return cls.from_level(BudgetLevel.MAXIMUM)
|
||||
|
||||
def get_tokens_for_complexity(self, complexity: float) -> int:
|
||||
"""
|
||||
Get token budget based on task complexity.
|
||||
|
||||
Args:
|
||||
complexity: Complexity score (0.0 to 1.0)
|
||||
|
||||
Returns:
|
||||
Adjusted token budget
|
||||
"""
|
||||
if not self.adaptive:
|
||||
return self.max_tokens
|
||||
|
||||
# Scale between min and max based on complexity
|
||||
range_tokens = self.max_tokens - self.min_tokens
|
||||
adjusted = self.min_tokens + int(range_tokens * complexity * self.complexity_multiplier)
|
||||
|
||||
return min(adjusted, self.max_tokens)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"max_tokens": self.max_tokens,
|
||||
"max_time_seconds": self.max_time_seconds,
|
||||
"adaptive": self.adaptive,
|
||||
"level": self.level.value if self.level else None,
|
||||
"min_tokens": self.min_tokens,
|
||||
"complexity_multiplier": self.complexity_multiplier
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "ThinkingBudget":
|
||||
"""Create from dictionary."""
|
||||
level = None
|
||||
if data.get("level"):
|
||||
level = BudgetLevel(data["level"])
|
||||
|
||||
return cls(
|
||||
max_tokens=data.get("max_tokens", 8000),
|
||||
max_time_seconds=data.get("max_time_seconds"),
|
||||
adaptive=data.get("adaptive", True),
|
||||
level=level,
|
||||
min_tokens=data.get("min_tokens", 1000),
|
||||
complexity_multiplier=data.get("complexity_multiplier", 1.0)
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
Thinking Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class ThinkingConfig:
|
||||
"""Configuration for extended thinking."""
|
||||
enabled: bool = True
|
||||
default_budget_tokens: int = 8000
|
||||
max_budget_tokens: int = 32000
|
||||
default_time_seconds: Optional[float] = None
|
||||
adaptive: bool = True # Adjust budget based on task complexity
|
||||
log_thinking: bool = False # Log thinking process
|
||||
@@ -0,0 +1,173 @@
|
||||
"""
|
||||
Thinking Tracker for PraisonAI Agents.
|
||||
|
||||
Tracks thinking usage and provides reporting.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
@dataclass
|
||||
class ThinkingUsage:
|
||||
"""Usage statistics for a single thinking session."""
|
||||
tokens_used: int = 0
|
||||
time_seconds: float = 0.0
|
||||
budget_tokens: int = 0
|
||||
budget_time: Optional[float] = None
|
||||
complexity: float = 0.5
|
||||
started_at: Optional[datetime] = None
|
||||
ended_at: Optional[datetime] = None
|
||||
|
||||
@property
|
||||
def tokens_remaining(self) -> int:
|
||||
"""Get remaining token budget."""
|
||||
return max(0, self.budget_tokens - self.tokens_used)
|
||||
|
||||
@property
|
||||
def time_remaining(self) -> Optional[float]:
|
||||
"""Get remaining time budget."""
|
||||
if self.budget_time is None:
|
||||
return None
|
||||
return max(0.0, self.budget_time - self.time_seconds)
|
||||
|
||||
@property
|
||||
def token_utilization(self) -> float:
|
||||
"""Get token utilization percentage."""
|
||||
if self.budget_tokens == 0:
|
||||
return 0.0
|
||||
return self.tokens_used / self.budget_tokens
|
||||
|
||||
@property
|
||||
def is_over_budget(self) -> bool:
|
||||
"""Check if over token budget."""
|
||||
return self.tokens_used > self.budget_tokens
|
||||
|
||||
@property
|
||||
def is_over_time(self) -> bool:
|
||||
"""Check if over time budget."""
|
||||
if self.budget_time is None:
|
||||
return False
|
||||
return self.time_seconds > self.budget_time
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"tokens_used": self.tokens_used,
|
||||
"time_seconds": self.time_seconds,
|
||||
"budget_tokens": self.budget_tokens,
|
||||
"budget_time": self.budget_time,
|
||||
"complexity": self.complexity,
|
||||
"tokens_remaining": self.tokens_remaining,
|
||||
"token_utilization": self.token_utilization,
|
||||
"is_over_budget": self.is_over_budget
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ThinkingTracker:
|
||||
"""
|
||||
Tracks thinking usage across multiple sessions.
|
||||
|
||||
Provides aggregate statistics and reporting.
|
||||
"""
|
||||
sessions: List[ThinkingUsage] = field(default_factory=list)
|
||||
total_tokens_used: int = 0
|
||||
total_time_seconds: float = 0.0
|
||||
|
||||
def start_session(
|
||||
self,
|
||||
budget_tokens: int,
|
||||
budget_time: Optional[float] = None,
|
||||
complexity: float = 0.5
|
||||
) -> ThinkingUsage:
|
||||
"""
|
||||
Start a new thinking session.
|
||||
|
||||
Args:
|
||||
budget_tokens: Token budget for this session
|
||||
budget_time: Optional time budget
|
||||
complexity: Task complexity (0.0 to 1.0)
|
||||
|
||||
Returns:
|
||||
ThinkingUsage for tracking
|
||||
"""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=budget_tokens,
|
||||
budget_time=budget_time,
|
||||
complexity=complexity,
|
||||
started_at=datetime.now()
|
||||
)
|
||||
self.sessions.append(usage)
|
||||
return usage
|
||||
|
||||
def end_session(
|
||||
self,
|
||||
usage: ThinkingUsage,
|
||||
tokens_used: int,
|
||||
time_seconds: float
|
||||
):
|
||||
"""
|
||||
End a thinking session.
|
||||
|
||||
Args:
|
||||
usage: The usage object from start_session
|
||||
tokens_used: Actual tokens used
|
||||
time_seconds: Actual time taken
|
||||
"""
|
||||
usage.tokens_used = tokens_used
|
||||
usage.time_seconds = time_seconds
|
||||
usage.ended_at = datetime.now()
|
||||
|
||||
self.total_tokens_used += tokens_used
|
||||
self.total_time_seconds += time_seconds
|
||||
|
||||
@property
|
||||
def session_count(self) -> int:
|
||||
"""Get number of sessions."""
|
||||
return len(self.sessions)
|
||||
|
||||
@property
|
||||
def average_tokens_per_session(self) -> float:
|
||||
"""Get average tokens per session."""
|
||||
if not self.sessions:
|
||||
return 0.0
|
||||
return self.total_tokens_used / len(self.sessions)
|
||||
|
||||
@property
|
||||
def average_time_per_session(self) -> float:
|
||||
"""Get average time per session."""
|
||||
if not self.sessions:
|
||||
return 0.0
|
||||
return self.total_time_seconds / len(self.sessions)
|
||||
|
||||
@property
|
||||
def average_utilization(self) -> float:
|
||||
"""Get average budget utilization."""
|
||||
if not self.sessions:
|
||||
return 0.0
|
||||
return sum(s.token_utilization for s in self.sessions) / len(self.sessions)
|
||||
|
||||
@property
|
||||
def over_budget_count(self) -> int:
|
||||
"""Get number of sessions that went over budget."""
|
||||
return sum(1 for s in self.sessions if s.is_over_budget)
|
||||
|
||||
def get_summary(self) -> Dict[str, Any]:
|
||||
"""Get summary statistics."""
|
||||
return {
|
||||
"session_count": self.session_count,
|
||||
"total_tokens_used": self.total_tokens_used,
|
||||
"total_time_seconds": self.total_time_seconds,
|
||||
"average_tokens_per_session": self.average_tokens_per_session,
|
||||
"average_time_per_session": self.average_time_per_session,
|
||||
"average_utilization": self.average_utilization,
|
||||
"over_budget_count": self.over_budget_count
|
||||
}
|
||||
|
||||
def clear(self):
|
||||
"""Clear all tracking data."""
|
||||
self.sessions.clear()
|
||||
self.total_tokens_used = 0
|
||||
self.total_time_seconds = 0.0
|
||||
@@ -222,6 +222,14 @@ TOOL_MAPPINGS = {
|
||||
# Unified Web Search (auto-fallback across providers)
|
||||
'search_web': ('.web_search', None),
|
||||
'get_available_providers': ('.web_search', None),
|
||||
|
||||
# Skill Tools (for Agent Skills script execution)
|
||||
'run_skill_script': ('.skill_tools', None),
|
||||
'read_skill_file': ('.skill_tools', None),
|
||||
'list_skill_scripts': ('.skill_tools', None),
|
||||
'create_skill_tools': ('.skill_tools', None),
|
||||
'SkillTools': ('.skill_tools', 'SkillTools'),
|
||||
'skill_tools': ('.skill_tools', None),
|
||||
}
|
||||
|
||||
_instances = {} # Cache for class instances
|
||||
@@ -259,11 +267,12 @@ def __getattr__(name: str) -> Any:
|
||||
'exa_search_async', 'exa_search_contents_async', 'exa_answer_async',
|
||||
'crawl4ai', 'crawl4ai_many', 'crawl4ai_extract', 'crawl4ai_llm_extract',
|
||||
'crawl4ai_sync', 'crawl4ai_extract_sync',
|
||||
'search_web', 'get_available_providers'
|
||||
'search_web', 'get_available_providers',
|
||||
'run_skill_script', 'read_skill_file', 'list_skill_scripts', 'create_skill_tools'
|
||||
]:
|
||||
return getattr(module, name)
|
||||
if name in ['file_tools', 'pandas_tools', 'wikipedia_tools',
|
||||
'newspaper_tools', 'arxiv_tools', 'spider_tools', 'duckdb_tools', 'mongodb_tools', 'csv_tools', 'json_tools', 'excel_tools', 'xml_tools', 'yaml_tools', 'calculator_tools', 'python_tools', 'shell_tools', 'cot_tools', 'tavily_tools', 'youdotcom_tools', 'exa_tools', 'crawl4ai_tools']:
|
||||
'newspaper_tools', 'arxiv_tools', 'spider_tools', 'duckdb_tools', 'mongodb_tools', 'csv_tools', 'json_tools', 'excel_tools', 'xml_tools', 'yaml_tools', 'calculator_tools', 'python_tools', 'shell_tools', 'cot_tools', 'tavily_tools', 'youdotcom_tools', 'exa_tools', 'crawl4ai_tools', 'skill_tools']:
|
||||
return module # Returns the callable module
|
||||
return getattr(module, name)
|
||||
else:
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Tools for executing Agent Skills scripts.
|
||||
|
||||
This module provides tools for running scripts bundled with Agent Skills.
|
||||
Skills are modular packages that extend agent capabilities through SKILL.md
|
||||
files and optional scripts.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.tools import run_skill_script
|
||||
result = run_skill_script("./my-skill/scripts/skill.py", "./data.csv")
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
import logging
|
||||
|
||||
|
||||
class SkillTools:
|
||||
"""Tools for executing Agent Skills scripts."""
|
||||
|
||||
def __init__(self, working_directory: str = None):
|
||||
"""Initialize SkillTools.
|
||||
|
||||
Args:
|
||||
working_directory: Base directory for resolving relative paths.
|
||||
Defaults to current working directory.
|
||||
"""
|
||||
self._working_directory = working_directory or os.getcwd()
|
||||
|
||||
@property
|
||||
def working_directory(self) -> str:
|
||||
"""Get the working directory for path resolution."""
|
||||
return self._working_directory
|
||||
|
||||
@working_directory.setter
|
||||
def working_directory(self, value: str):
|
||||
"""Set the working directory for path resolution."""
|
||||
self._working_directory = value
|
||||
|
||||
def run_skill_script(
|
||||
self,
|
||||
script_path: str,
|
||||
args: str = "",
|
||||
timeout: int = 60
|
||||
) -> str:
|
||||
"""
|
||||
Execute a skill script from a skill's scripts/ directory.
|
||||
|
||||
Args:
|
||||
script_path: Path to the script file (e.g., ./csv-analyzer/scripts/skill.py)
|
||||
args: Arguments to pass to the script (e.g., file path to analyze)
|
||||
timeout: Maximum execution time in seconds (default: 60)
|
||||
|
||||
Returns:
|
||||
The output from running the script, or error message if failed
|
||||
"""
|
||||
try:
|
||||
# Resolve the script path
|
||||
script_path = os.path.expanduser(script_path)
|
||||
if not os.path.isabs(script_path):
|
||||
script_path = os.path.join(self._working_directory, script_path)
|
||||
script_path = os.path.abspath(script_path)
|
||||
|
||||
if not os.path.exists(script_path):
|
||||
return f"Error: Script not found at {script_path}"
|
||||
|
||||
# Resolve args paths relative to working directory
|
||||
resolved_args = []
|
||||
if args:
|
||||
for arg in args.split():
|
||||
# Skip flags
|
||||
if arg.startswith('-'):
|
||||
resolved_args.append(arg)
|
||||
continue
|
||||
|
||||
# Try to resolve as a path relative to working directory
|
||||
potential_path = os.path.join(self._working_directory, arg)
|
||||
if os.path.exists(potential_path):
|
||||
resolved_args.append(os.path.abspath(potential_path))
|
||||
elif os.path.isabs(arg) and os.path.exists(arg):
|
||||
# Already absolute and exists
|
||||
resolved_args.append(arg)
|
||||
else:
|
||||
# Keep original arg if path doesn't exist
|
||||
resolved_args.append(arg)
|
||||
|
||||
# Determine how to run the script based on extension
|
||||
ext = os.path.splitext(script_path)[1].lower()
|
||||
|
||||
if ext == '.py':
|
||||
cmd = ['python', script_path] + resolved_args
|
||||
elif ext == '.sh':
|
||||
cmd = ['bash', script_path] + resolved_args
|
||||
elif ext == '.js':
|
||||
cmd = ['node', script_path] + resolved_args
|
||||
else:
|
||||
# Try to run directly
|
||||
cmd = [script_path] + resolved_args
|
||||
|
||||
logging.debug(f"Running skill script: {' '.join(cmd)}")
|
||||
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
cwd=self._working_directory
|
||||
)
|
||||
|
||||
output = result.stdout
|
||||
if result.stderr:
|
||||
output += f"\nStderr: {result.stderr}"
|
||||
if result.returncode != 0:
|
||||
output += f"\nExit code: {result.returncode}"
|
||||
|
||||
return output if output else "Script executed successfully (no output)"
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return f"Error: Script execution timed out after {timeout} seconds"
|
||||
except FileNotFoundError as e:
|
||||
return f"Error: Could not find interpreter or script: {str(e)}"
|
||||
except PermissionError as e:
|
||||
return f"Error: Permission denied: {str(e)}"
|
||||
except Exception as e:
|
||||
return f"Error executing script: {str(e)}"
|
||||
|
||||
def read_skill_file(
|
||||
self,
|
||||
skill_path: str,
|
||||
file_path: str,
|
||||
encoding: str = 'utf-8'
|
||||
) -> str:
|
||||
"""
|
||||
Read a file from a skill directory.
|
||||
|
||||
Args:
|
||||
skill_path: Path to the skill directory (e.g., ./csv-analyzer)
|
||||
file_path: Relative path within the skill (e.g., SKILL.md, scripts/skill.py)
|
||||
encoding: File encoding (default: utf-8)
|
||||
|
||||
Returns:
|
||||
File contents or error message
|
||||
"""
|
||||
try:
|
||||
# Resolve skill path
|
||||
skill_path = os.path.expanduser(skill_path)
|
||||
if not os.path.isabs(skill_path):
|
||||
skill_path = os.path.join(self._working_directory, skill_path)
|
||||
skill_path = os.path.abspath(skill_path)
|
||||
|
||||
if not os.path.exists(skill_path):
|
||||
return f"Error: Skill directory not found at {skill_path}"
|
||||
|
||||
if not os.path.isdir(skill_path):
|
||||
return f"Error: {skill_path} is not a directory"
|
||||
|
||||
# Resolve file path within skill
|
||||
full_path = os.path.join(skill_path, file_path)
|
||||
full_path = os.path.abspath(full_path)
|
||||
|
||||
# Security check: ensure file is within skill directory
|
||||
if not full_path.startswith(skill_path):
|
||||
return f"Error: Path traversal detected - {file_path} is outside skill directory"
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
return f"Error: File not found at {full_path}"
|
||||
|
||||
with open(full_path, 'r', encoding=encoding) as f:
|
||||
return f.read()
|
||||
|
||||
except UnicodeDecodeError as e:
|
||||
return f"Error: Could not decode file with {encoding} encoding: {str(e)}"
|
||||
except PermissionError as e:
|
||||
return f"Error: Permission denied reading file: {str(e)}"
|
||||
except Exception as e:
|
||||
return f"Error reading skill file: {str(e)}"
|
||||
|
||||
def list_skill_scripts(self, skill_path: str) -> str:
|
||||
"""
|
||||
List available scripts in a skill's scripts/ directory.
|
||||
|
||||
Args:
|
||||
skill_path: Path to the skill directory
|
||||
|
||||
Returns:
|
||||
JSON string with list of scripts or error message
|
||||
"""
|
||||
import json
|
||||
|
||||
try:
|
||||
# Resolve skill path
|
||||
skill_path = os.path.expanduser(skill_path)
|
||||
if not os.path.isabs(skill_path):
|
||||
skill_path = os.path.join(self._working_directory, skill_path)
|
||||
skill_path = os.path.abspath(skill_path)
|
||||
|
||||
if not os.path.exists(skill_path):
|
||||
return f"Error: Skill directory not found at {skill_path}"
|
||||
|
||||
scripts_dir = os.path.join(skill_path, "scripts")
|
||||
if not os.path.exists(scripts_dir):
|
||||
return json.dumps({"scripts": [], "message": "No scripts/ directory found"})
|
||||
|
||||
scripts = []
|
||||
for item in os.listdir(scripts_dir):
|
||||
item_path = os.path.join(scripts_dir, item)
|
||||
if os.path.isfile(item_path):
|
||||
ext = os.path.splitext(item)[1].lower()
|
||||
if ext in ['.py', '.sh', '.js', '.rb', '.pl']:
|
||||
scripts.append({
|
||||
"name": item,
|
||||
"path": os.path.join("scripts", item),
|
||||
"type": ext[1:] # Remove the dot
|
||||
})
|
||||
|
||||
return json.dumps({"scripts": scripts}, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error listing skill scripts: {str(e)}"
|
||||
|
||||
|
||||
# Create default instance for direct function access
|
||||
_skill_tools = SkillTools()
|
||||
|
||||
def run_skill_script(script_path: str, args: str = "", timeout: int = 60) -> str:
|
||||
"""
|
||||
Execute a skill script from a skill's scripts/ directory.
|
||||
|
||||
Args:
|
||||
script_path: Path to the script file (e.g., ./csv-analyzer/scripts/script.py)
|
||||
args: Arguments to pass to the script (e.g., file path to analyze)
|
||||
timeout: Maximum execution time in seconds (default: 60)
|
||||
|
||||
Returns:
|
||||
The output from running the script, or error message if failed
|
||||
"""
|
||||
return _skill_tools.run_skill_script(script_path, args, timeout)
|
||||
|
||||
|
||||
def read_skill_file(skill_path: str, file_path: str, encoding: str = 'utf-8') -> str:
|
||||
"""
|
||||
Read a file from a skill directory.
|
||||
|
||||
Args:
|
||||
skill_path: Path to the skill directory (e.g., ./csv-analyzer)
|
||||
file_path: Relative path within the skill (e.g., SKILL.md, scripts/script.py)
|
||||
encoding: File encoding (default: utf-8)
|
||||
|
||||
Returns:
|
||||
File contents or error message
|
||||
"""
|
||||
return _skill_tools.read_skill_file(skill_path, file_path, encoding)
|
||||
|
||||
|
||||
def list_skill_scripts(skill_path: str) -> str:
|
||||
"""
|
||||
List available scripts in a skill's scripts/ directory.
|
||||
|
||||
Args:
|
||||
skill_path: Path to the skill directory
|
||||
|
||||
Returns:
|
||||
JSON string with list of scripts or error message
|
||||
"""
|
||||
return _skill_tools.list_skill_scripts(skill_path)
|
||||
|
||||
|
||||
def create_skill_tools(working_directory: str = None) -> SkillTools:
|
||||
"""
|
||||
Create a SkillTools instance with a specific working directory.
|
||||
|
||||
This is useful when you need to resolve paths relative to a specific
|
||||
directory rather than the current working directory.
|
||||
|
||||
Args:
|
||||
working_directory: Base directory for resolving relative paths
|
||||
|
||||
Returns:
|
||||
SkillTools instance
|
||||
"""
|
||||
return SkillTools(working_directory)
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.1.23"
|
||||
version = "0.2.1"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for background module."""
|
||||
@@ -0,0 +1,401 @@
|
||||
"""
|
||||
Unit tests for the Background module.
|
||||
|
||||
Tests cover:
|
||||
- BackgroundTask creation and status management
|
||||
- BackgroundRunner task submission
|
||||
- Task cancellation and timeout
|
||||
- Concurrent execution
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
|
||||
from praisonaiagents.background.task import BackgroundTask, TaskStatus
|
||||
from praisonaiagents.background.config import BackgroundConfig
|
||||
from praisonaiagents.background.runner import BackgroundRunner
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# BackgroundTask Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestBackgroundTask:
|
||||
"""Tests for BackgroundTask class."""
|
||||
|
||||
def test_task_creation(self):
|
||||
"""Test creating a task."""
|
||||
task = BackgroundTask(name="test_task")
|
||||
assert task.name == "test_task"
|
||||
assert task.status == TaskStatus.PENDING
|
||||
assert task.progress == 0.0
|
||||
assert not task.is_running
|
||||
assert not task.is_completed
|
||||
|
||||
def test_task_start(self):
|
||||
"""Test starting a task."""
|
||||
task = BackgroundTask(name="test")
|
||||
task.start()
|
||||
|
||||
assert task.status == TaskStatus.RUNNING
|
||||
assert task.is_running
|
||||
assert task.started_at is not None
|
||||
|
||||
def test_task_complete(self):
|
||||
"""Test completing a task."""
|
||||
task = BackgroundTask(name="test")
|
||||
task.start()
|
||||
task.complete(result="success")
|
||||
|
||||
assert task.status == TaskStatus.COMPLETED
|
||||
assert task.is_completed
|
||||
assert task.is_successful
|
||||
assert task.result == "success"
|
||||
assert task.progress == 1.0
|
||||
assert task.completed_at is not None
|
||||
|
||||
def test_task_fail(self):
|
||||
"""Test failing a task."""
|
||||
task = BackgroundTask(name="test")
|
||||
task.start()
|
||||
task.fail("Something went wrong")
|
||||
|
||||
assert task.status == TaskStatus.FAILED
|
||||
assert task.is_completed
|
||||
assert not task.is_successful
|
||||
assert task.error == "Something went wrong"
|
||||
|
||||
def test_task_cancel(self):
|
||||
"""Test cancelling a task."""
|
||||
task = BackgroundTask(name="test")
|
||||
task.start()
|
||||
task.cancel()
|
||||
|
||||
assert task.status == TaskStatus.CANCELLED
|
||||
assert task.is_completed
|
||||
assert task.should_cancel()
|
||||
|
||||
def test_task_progress(self):
|
||||
"""Test updating task progress."""
|
||||
task = BackgroundTask(name="test")
|
||||
task.start()
|
||||
|
||||
task.update_progress(0.5, "Halfway done")
|
||||
|
||||
assert task.progress == 0.5
|
||||
assert task.metadata.get("last_message") == "Halfway done"
|
||||
|
||||
def test_task_progress_bounds(self):
|
||||
"""Test that progress is bounded."""
|
||||
task = BackgroundTask(name="test")
|
||||
|
||||
task.update_progress(-0.5)
|
||||
assert task.progress == 0.0
|
||||
|
||||
task.update_progress(1.5)
|
||||
assert task.progress == 1.0
|
||||
|
||||
def test_task_duration(self):
|
||||
"""Test task duration calculation."""
|
||||
task = BackgroundTask(name="test")
|
||||
|
||||
assert task.duration_seconds is None
|
||||
|
||||
task.start()
|
||||
assert task.duration_seconds is not None
|
||||
assert task.duration_seconds >= 0
|
||||
|
||||
def test_task_to_dict(self):
|
||||
"""Test task serialization."""
|
||||
task = BackgroundTask(name="test")
|
||||
task.start()
|
||||
task.complete("result")
|
||||
|
||||
data = task.to_dict()
|
||||
|
||||
assert data["name"] == "test"
|
||||
assert data["status"] == "completed"
|
||||
assert data["result"] == "result"
|
||||
assert "created_at" in data
|
||||
assert "started_at" in data
|
||||
assert "completed_at" in data
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# BackgroundConfig Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestBackgroundConfig:
|
||||
"""Tests for BackgroundConfig class."""
|
||||
|
||||
def test_config_defaults(self):
|
||||
"""Test default configuration."""
|
||||
config = BackgroundConfig()
|
||||
|
||||
assert config.max_concurrent_tasks == 5
|
||||
assert config.default_timeout is None
|
||||
assert config.auto_cleanup
|
||||
assert config.cleanup_delay == 300.0
|
||||
|
||||
def test_config_custom(self):
|
||||
"""Test custom configuration."""
|
||||
config = BackgroundConfig(
|
||||
max_concurrent_tasks=10,
|
||||
default_timeout=60.0,
|
||||
auto_cleanup=False
|
||||
)
|
||||
|
||||
assert config.max_concurrent_tasks == 10
|
||||
assert config.default_timeout == 60.0
|
||||
assert not config.auto_cleanup
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# BackgroundRunner Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestBackgroundRunner:
|
||||
"""Tests for BackgroundRunner class."""
|
||||
|
||||
@pytest.fixture
|
||||
def runner(self):
|
||||
"""Create a test runner."""
|
||||
return BackgroundRunner()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_sync_function(self, runner):
|
||||
"""Test submitting a sync function."""
|
||||
def sync_func(x):
|
||||
return x * 2
|
||||
|
||||
task = await runner.submit(func=sync_func, args=(5,), name="double")
|
||||
result = await task.wait(timeout=5.0)
|
||||
|
||||
assert result == 10
|
||||
assert task.is_successful
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_async_function(self, runner):
|
||||
"""Test submitting an async function."""
|
||||
async def async_func(x):
|
||||
await asyncio.sleep(0.1)
|
||||
return x * 3
|
||||
|
||||
task = await runner.submit(func=async_func, args=(5,), name="triple")
|
||||
result = await task.wait(timeout=5.0)
|
||||
|
||||
assert result == 15
|
||||
assert task.is_successful
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_with_kwargs(self, runner):
|
||||
"""Test submitting with keyword arguments."""
|
||||
def func_with_kwargs(a, b=10):
|
||||
return a + b
|
||||
|
||||
task = await runner.submit(
|
||||
func=func_with_kwargs,
|
||||
args=(5,),
|
||||
kwargs={"b": 20}
|
||||
)
|
||||
result = await task.wait(timeout=5.0)
|
||||
|
||||
assert result == 25
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_task_timeout(self, runner):
|
||||
"""Test task timeout."""
|
||||
async def slow_func():
|
||||
await asyncio.sleep(10)
|
||||
return "done"
|
||||
|
||||
task = await runner.submit(func=slow_func, timeout=0.1)
|
||||
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
assert task.status == TaskStatus.FAILED
|
||||
assert "timed out" in task.error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_task_failure(self, runner):
|
||||
"""Test task that raises exception."""
|
||||
def failing_func():
|
||||
raise ValueError("Test error")
|
||||
|
||||
task = await runner.submit(func=failing_func)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
assert task.status == TaskStatus.FAILED
|
||||
assert "Test error" in task.error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_task(self, runner):
|
||||
"""Test cancelling a task."""
|
||||
async def long_running():
|
||||
await asyncio.sleep(10)
|
||||
return "done"
|
||||
|
||||
task = await runner.submit(func=long_running)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
result = await runner.cancel_task(task.id)
|
||||
|
||||
assert result
|
||||
assert task.status == TaskStatus.CANCELLED
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_task(self, runner):
|
||||
"""Test getting a task by ID."""
|
||||
task = await runner.submit(func=lambda: "test")
|
||||
|
||||
retrieved = runner.get_task(task.id)
|
||||
|
||||
assert retrieved is not None
|
||||
assert retrieved.id == task.id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tasks(self, runner):
|
||||
"""Test listing tasks."""
|
||||
await runner.submit(func=lambda: "a", name="task_a")
|
||||
await runner.submit(func=lambda: "b", name="task_b")
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
tasks = runner.list_tasks()
|
||||
|
||||
assert len(tasks) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tasks_by_status(self, runner):
|
||||
"""Test listing tasks filtered by status."""
|
||||
await runner.submit(func=lambda: "a")
|
||||
await runner.submit(func=lambda: "b")
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
completed = runner.list_tasks(status=TaskStatus.COMPLETED)
|
||||
|
||||
assert len(completed) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_completed(self, runner):
|
||||
"""Test clearing completed tasks."""
|
||||
await runner.submit(func=lambda: "a")
|
||||
await runner.submit(func=lambda: "b")
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
assert len(runner.tasks) == 2
|
||||
|
||||
runner.clear_completed()
|
||||
|
||||
assert len(runner.tasks) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_complete_callback(self, runner):
|
||||
"""Test on_complete callback."""
|
||||
callback_called = []
|
||||
|
||||
def on_complete(task):
|
||||
callback_called.append(task.id)
|
||||
|
||||
task = await runner.submit(
|
||||
func=lambda: "result",
|
||||
on_complete=on_complete
|
||||
)
|
||||
|
||||
await task.wait(timeout=5.0)
|
||||
|
||||
assert task.id in callback_called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_execution(self, runner):
|
||||
"""Test concurrent task execution."""
|
||||
results = []
|
||||
|
||||
async def task_func(n):
|
||||
await asyncio.sleep(0.1)
|
||||
results.append(n)
|
||||
return n
|
||||
|
||||
tasks = []
|
||||
for i in range(3):
|
||||
task = await runner.submit(func=task_func, args=(i,))
|
||||
tasks.append(task)
|
||||
|
||||
await runner.wait_all(timeout=5.0)
|
||||
|
||||
assert len(results) == 3
|
||||
assert all(t.is_successful for t in tasks)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_concurrent_limit(self):
|
||||
"""Test max concurrent tasks limit."""
|
||||
config = BackgroundConfig(max_concurrent_tasks=2)
|
||||
runner = BackgroundRunner(config=config)
|
||||
|
||||
execution_times = []
|
||||
|
||||
async def tracked_func(n):
|
||||
execution_times.append((n, "start", asyncio.get_event_loop().time()))
|
||||
await asyncio.sleep(0.2)
|
||||
execution_times.append((n, "end", asyncio.get_event_loop().time()))
|
||||
return n
|
||||
|
||||
tasks = []
|
||||
for i in range(4):
|
||||
task = await runner.submit(func=tracked_func, args=(i,))
|
||||
tasks.append(task)
|
||||
|
||||
await runner.wait_all(timeout=5.0)
|
||||
|
||||
# All tasks should complete
|
||||
assert all(t.is_successful for t in tasks)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wait_all(self, runner):
|
||||
"""Test waiting for all tasks."""
|
||||
async def slow_func(n):
|
||||
await asyncio.sleep(0.1 * n)
|
||||
return n
|
||||
|
||||
for i in range(1, 4):
|
||||
await runner.submit(func=slow_func, args=(i,))
|
||||
|
||||
tasks = await runner.wait_all(timeout=5.0)
|
||||
|
||||
assert len(tasks) == 3
|
||||
assert all(t.is_completed for t in tasks)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_all(self, runner):
|
||||
"""Test cancelling all tasks."""
|
||||
async def long_func():
|
||||
await asyncio.sleep(10)
|
||||
|
||||
for _ in range(3):
|
||||
await runner.submit(func=long_func)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
await runner.cancel_all()
|
||||
|
||||
assert all(t.status == TaskStatus.CANCELLED for t in runner.tasks)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_runner(self, runner):
|
||||
"""Test stopping the runner."""
|
||||
async def long_func():
|
||||
await asyncio.sleep(10)
|
||||
|
||||
await runner.submit(func=long_func)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
await runner.stop()
|
||||
|
||||
assert all(t.status == TaskStatus.CANCELLED for t in runner.tasks)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for checkpoints module."""
|
||||
@@ -0,0 +1,417 @@
|
||||
"""
|
||||
Unit tests for the Checkpoints module.
|
||||
|
||||
Tests cover:
|
||||
- CheckpointConfig creation and validation
|
||||
- Checkpoint data types
|
||||
- CheckpointService initialization
|
||||
- Save, restore, and diff operations
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
|
||||
from praisonaiagents.checkpoints.types import (
|
||||
CheckpointConfig, Checkpoint, CheckpointDiff, FileDiff,
|
||||
CheckpointResult, CheckpointEvent
|
||||
)
|
||||
from praisonaiagents.checkpoints.service import CheckpointService
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CheckpointConfig Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCheckpointConfig:
|
||||
"""Tests for CheckpointConfig class."""
|
||||
|
||||
def test_config_creation(self):
|
||||
"""Test creating a config."""
|
||||
config = CheckpointConfig(workspace_dir="/tmp/test")
|
||||
assert config.workspace_dir == "/tmp/test"
|
||||
assert config.enabled
|
||||
assert config.auto_checkpoint
|
||||
assert config.max_checkpoints == 100
|
||||
|
||||
def test_config_expands_paths(self):
|
||||
"""Test that paths are expanded."""
|
||||
config = CheckpointConfig(
|
||||
workspace_dir="~/test",
|
||||
storage_dir="~/.praison/checkpoints"
|
||||
)
|
||||
assert not config.workspace_dir.startswith("~")
|
||||
assert not config.storage_dir.startswith("~")
|
||||
|
||||
def test_config_default_storage(self):
|
||||
"""Test default storage directory."""
|
||||
config = CheckpointConfig(workspace_dir="/tmp/test")
|
||||
assert ".praison/checkpoints" in config.storage_dir
|
||||
|
||||
def test_config_exclude_patterns(self):
|
||||
"""Test default exclude patterns."""
|
||||
config = CheckpointConfig(workspace_dir="/tmp/test")
|
||||
assert ".git" in config.exclude_patterns
|
||||
assert "__pycache__" in config.exclude_patterns
|
||||
assert "node_modules" in config.exclude_patterns
|
||||
|
||||
def test_get_checkpoint_dir(self):
|
||||
"""Test checkpoint directory generation."""
|
||||
config = CheckpointConfig(workspace_dir="/tmp/test")
|
||||
checkpoint_dir = config.get_checkpoint_dir()
|
||||
assert config.storage_dir in checkpoint_dir
|
||||
# Should be a hash-based subdirectory
|
||||
assert len(os.path.basename(checkpoint_dir)) == 12
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Checkpoint Data Type Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCheckpointTypes:
|
||||
"""Tests for checkpoint data types."""
|
||||
|
||||
def test_checkpoint_creation(self):
|
||||
"""Test creating a checkpoint."""
|
||||
checkpoint = Checkpoint(
|
||||
id="abc123def456",
|
||||
short_id="abc123de",
|
||||
message="Test checkpoint",
|
||||
timestamp=datetime.now()
|
||||
)
|
||||
assert checkpoint.id == "abc123def456"
|
||||
assert checkpoint.short_id == "abc123de"
|
||||
assert checkpoint.message == "Test checkpoint"
|
||||
|
||||
def test_checkpoint_from_git_commit(self):
|
||||
"""Test creating checkpoint from git commit."""
|
||||
checkpoint = Checkpoint.from_git_commit(
|
||||
commit_hash="abc123def456789",
|
||||
message="Test message",
|
||||
timestamp="2024-01-01T00:00:00"
|
||||
)
|
||||
assert checkpoint.id == "abc123def456789"
|
||||
assert checkpoint.short_id == "abc123de"
|
||||
assert checkpoint.message == "Test message"
|
||||
|
||||
def test_checkpoint_to_dict(self):
|
||||
"""Test checkpoint serialization."""
|
||||
checkpoint = Checkpoint(
|
||||
id="abc123",
|
||||
short_id="abc1",
|
||||
message="Test",
|
||||
timestamp=datetime(2024, 1, 1)
|
||||
)
|
||||
data = checkpoint.to_dict()
|
||||
assert data["id"] == "abc123"
|
||||
assert data["message"] == "Test"
|
||||
assert "timestamp" in data
|
||||
|
||||
def test_file_diff(self):
|
||||
"""Test FileDiff creation."""
|
||||
diff = FileDiff(
|
||||
path="test.py",
|
||||
absolute_path="/tmp/test.py",
|
||||
status="modified",
|
||||
additions=10,
|
||||
deletions=5
|
||||
)
|
||||
assert diff.path == "test.py"
|
||||
assert diff.status == "modified"
|
||||
assert diff.additions == 10
|
||||
|
||||
def test_checkpoint_diff(self):
|
||||
"""Test CheckpointDiff creation."""
|
||||
files = [
|
||||
FileDiff("a.py", "/tmp/a.py", "added", 10, 0),
|
||||
FileDiff("b.py", "/tmp/b.py", "modified", 5, 3)
|
||||
]
|
||||
diff = CheckpointDiff(
|
||||
from_checkpoint="abc123",
|
||||
to_checkpoint="def456",
|
||||
files=files
|
||||
)
|
||||
assert diff.total_additions == 15
|
||||
assert diff.total_deletions == 3
|
||||
|
||||
def test_checkpoint_result_ok(self):
|
||||
"""Test successful checkpoint result."""
|
||||
checkpoint = Checkpoint(
|
||||
id="abc123",
|
||||
short_id="abc1",
|
||||
message="Test",
|
||||
timestamp=datetime.now()
|
||||
)
|
||||
result = CheckpointResult.ok(checkpoint)
|
||||
assert result.success
|
||||
assert result.checkpoint == checkpoint
|
||||
assert result.error is None
|
||||
|
||||
def test_checkpoint_result_fail(self):
|
||||
"""Test failed checkpoint result."""
|
||||
result = CheckpointResult.fail("Something went wrong")
|
||||
assert not result.success
|
||||
assert result.checkpoint is None
|
||||
assert result.error == "Something went wrong"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CheckpointService Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCheckpointService:
|
||||
"""Tests for CheckpointService class."""
|
||||
|
||||
@pytest.fixture
|
||||
def temp_workspace(self):
|
||||
"""Create a temporary workspace directory."""
|
||||
workspace = tempfile.mkdtemp(prefix="praison_test_")
|
||||
yield workspace
|
||||
# Cleanup
|
||||
if os.path.exists(workspace):
|
||||
shutil.rmtree(workspace)
|
||||
|
||||
@pytest.fixture
|
||||
def temp_storage(self):
|
||||
"""Create a temporary storage directory."""
|
||||
storage = tempfile.mkdtemp(prefix="praison_checkpoints_")
|
||||
yield storage
|
||||
# Cleanup
|
||||
if os.path.exists(storage):
|
||||
shutil.rmtree(storage)
|
||||
|
||||
def test_service_creation(self, temp_workspace, temp_storage):
|
||||
"""Test creating a checkpoint service."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
assert service.workspace_dir == temp_workspace
|
||||
assert not service.is_initialized
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_initialize(self, temp_workspace, temp_storage):
|
||||
"""Test initializing the service."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
|
||||
result = await service.initialize()
|
||||
|
||||
assert result
|
||||
assert service.is_initialized
|
||||
assert os.path.exists(service.git_dir)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_save_checkpoint(self, temp_workspace, temp_storage):
|
||||
"""Test saving a checkpoint."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
# Create a file
|
||||
test_file = os.path.join(temp_workspace, "test.txt")
|
||||
with open(test_file, "w") as f:
|
||||
f.write("Hello, World!")
|
||||
|
||||
# Save checkpoint
|
||||
result = await service.save("First checkpoint")
|
||||
|
||||
assert result.success
|
||||
assert result.checkpoint is not None
|
||||
assert result.checkpoint.message == "First checkpoint"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_save_no_changes(self, temp_workspace, temp_storage):
|
||||
"""Test saving checkpoint with no changes."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
# Try to save without changes
|
||||
result = await service.save("Empty checkpoint")
|
||||
|
||||
assert not result.success
|
||||
assert "No changes" in result.error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_save_allow_empty(self, temp_workspace, temp_storage):
|
||||
"""Test saving empty checkpoint when allowed."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
# Save empty checkpoint
|
||||
result = await service.save("Empty checkpoint", allow_empty=True)
|
||||
|
||||
assert result.success
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_restore(self, temp_workspace, temp_storage):
|
||||
"""Test restoring to a checkpoint."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
# Create initial file
|
||||
test_file = os.path.join(temp_workspace, "test.txt")
|
||||
with open(test_file, "w") as f:
|
||||
f.write("Version 1")
|
||||
|
||||
# Save checkpoint
|
||||
result1 = await service.save("Version 1")
|
||||
checkpoint_id = result1.checkpoint.id
|
||||
|
||||
# Modify file
|
||||
with open(test_file, "w") as f:
|
||||
f.write("Version 2")
|
||||
|
||||
# Restore to checkpoint
|
||||
result = await service.restore(checkpoint_id)
|
||||
|
||||
assert result.success
|
||||
|
||||
# Verify file content
|
||||
with open(test_file, "r") as f:
|
||||
content = f.read()
|
||||
assert content == "Version 1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_list_checkpoints(self, temp_workspace, temp_storage):
|
||||
"""Test listing checkpoints."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
# Create some checkpoints
|
||||
test_file = os.path.join(temp_workspace, "test.txt")
|
||||
|
||||
with open(test_file, "w") as f:
|
||||
f.write("v1")
|
||||
await service.save("Checkpoint 1")
|
||||
|
||||
with open(test_file, "w") as f:
|
||||
f.write("v2")
|
||||
await service.save("Checkpoint 2")
|
||||
|
||||
# List checkpoints
|
||||
checkpoints = await service.list_checkpoints()
|
||||
|
||||
# Should have initial + 2 checkpoints
|
||||
assert len(checkpoints) >= 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_diff(self, temp_workspace, temp_storage):
|
||||
"""Test getting diff between checkpoints."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
# Create file and checkpoint
|
||||
test_file = os.path.join(temp_workspace, "test.txt")
|
||||
with open(test_file, "w") as f:
|
||||
f.write("Hello")
|
||||
await service.save("Initial")
|
||||
|
||||
# Modify file
|
||||
with open(test_file, "w") as f:
|
||||
f.write("Hello World")
|
||||
|
||||
# Get diff
|
||||
diff = await service.diff()
|
||||
|
||||
assert diff.from_checkpoint != ""
|
||||
assert len(diff.files) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_protected_path(self, temp_storage):
|
||||
"""Test that protected paths are rejected."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=os.path.expanduser("~"),
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
|
||||
result = await service.initialize()
|
||||
|
||||
assert not result
|
||||
assert not service.is_initialized
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_event_handlers(self, temp_workspace, temp_storage):
|
||||
"""Test event handlers."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
|
||||
events_received = []
|
||||
|
||||
def on_initialized(data):
|
||||
events_received.append(("initialized", data))
|
||||
|
||||
def on_checkpoint(data):
|
||||
events_received.append(("checkpoint", data))
|
||||
|
||||
service.on(CheckpointEvent.INITIALIZED, on_initialized)
|
||||
service.on(CheckpointEvent.CHECKPOINT_CREATED, on_checkpoint)
|
||||
|
||||
await service.initialize()
|
||||
|
||||
test_file = os.path.join(temp_workspace, "test.txt")
|
||||
with open(test_file, "w") as f:
|
||||
f.write("test")
|
||||
await service.save("Test")
|
||||
|
||||
assert len(events_received) == 2
|
||||
assert events_received[0][0] == "initialized"
|
||||
assert events_received[1][0] == "checkpoint"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_cleanup(self, temp_workspace, temp_storage):
|
||||
"""Test service cleanup."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
await service.cleanup()
|
||||
|
||||
assert not service.is_initialized
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_delete_all(self, temp_workspace, temp_storage):
|
||||
"""Test deleting all checkpoints."""
|
||||
service = CheckpointService(
|
||||
workspace_dir=temp_workspace,
|
||||
storage_dir=temp_storage
|
||||
)
|
||||
await service.initialize()
|
||||
|
||||
checkpoint_dir = service.checkpoint_dir
|
||||
assert os.path.exists(checkpoint_dir)
|
||||
|
||||
result = await service.delete_all()
|
||||
|
||||
assert result
|
||||
assert not os.path.exists(checkpoint_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for compaction module."""
|
||||
@@ -0,0 +1,314 @@
|
||||
"""
|
||||
Unit tests for the Compaction module.
|
||||
|
||||
Tests cover:
|
||||
- CompactionConfig settings
|
||||
- CompactionStrategy enum
|
||||
- CompactionResult properties
|
||||
- ContextCompactor operations
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.compaction.config import CompactionConfig
|
||||
from praisonaiagents.compaction.strategy import CompactionStrategy
|
||||
from praisonaiagents.compaction.result import CompactionResult
|
||||
from praisonaiagents.compaction.compactor import ContextCompactor
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CompactionConfig Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCompactionConfig:
|
||||
"""Tests for CompactionConfig class."""
|
||||
|
||||
def test_config_defaults(self):
|
||||
"""Test default configuration."""
|
||||
config = CompactionConfig()
|
||||
|
||||
assert config.enabled
|
||||
assert config.max_tokens == 8000
|
||||
assert config.target_tokens == 6000
|
||||
assert config.preserve_system
|
||||
assert config.preserve_recent == 5
|
||||
|
||||
def test_config_custom(self):
|
||||
"""Test custom configuration."""
|
||||
config = CompactionConfig(
|
||||
enabled=False,
|
||||
max_tokens=16000,
|
||||
preserve_recent=10
|
||||
)
|
||||
|
||||
assert not config.enabled
|
||||
assert config.max_tokens == 16000
|
||||
assert config.preserve_recent == 10
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CompactionStrategy Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCompactionStrategy:
|
||||
"""Tests for CompactionStrategy enum."""
|
||||
|
||||
def test_strategy_values(self):
|
||||
"""Test strategy enum values."""
|
||||
assert CompactionStrategy.TRUNCATE.value == "truncate"
|
||||
assert CompactionStrategy.SUMMARIZE.value == "summarize"
|
||||
assert CompactionStrategy.SLIDING.value == "sliding"
|
||||
assert CompactionStrategy.SMART.value == "smart"
|
||||
|
||||
def test_strategy_from_string(self):
|
||||
"""Test creating strategy from string."""
|
||||
strategy = CompactionStrategy("truncate")
|
||||
assert strategy == CompactionStrategy.TRUNCATE
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CompactionResult Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCompactionResult:
|
||||
"""Tests for CompactionResult class."""
|
||||
|
||||
def test_result_creation(self):
|
||||
"""Test creating a result."""
|
||||
result = CompactionResult(
|
||||
original_tokens=10000,
|
||||
compacted_tokens=6000,
|
||||
messages_removed=5,
|
||||
messages_kept=10,
|
||||
strategy_used=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
|
||||
assert result.original_tokens == 10000
|
||||
assert result.compacted_tokens == 6000
|
||||
|
||||
def test_result_tokens_saved(self):
|
||||
"""Test tokens saved calculation."""
|
||||
result = CompactionResult(
|
||||
original_tokens=10000,
|
||||
compacted_tokens=6000,
|
||||
messages_removed=5,
|
||||
messages_kept=10,
|
||||
strategy_used=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
|
||||
assert result.tokens_saved == 4000
|
||||
|
||||
def test_result_compression_ratio(self):
|
||||
"""Test compression ratio calculation."""
|
||||
result = CompactionResult(
|
||||
original_tokens=10000,
|
||||
compacted_tokens=5000,
|
||||
messages_removed=5,
|
||||
messages_kept=10,
|
||||
strategy_used=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
|
||||
assert result.compression_ratio == 0.5
|
||||
|
||||
def test_result_was_compacted(self):
|
||||
"""Test was_compacted property."""
|
||||
result = CompactionResult(
|
||||
original_tokens=10000,
|
||||
compacted_tokens=6000,
|
||||
messages_removed=5,
|
||||
messages_kept=10,
|
||||
strategy_used=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
|
||||
assert result.was_compacted
|
||||
|
||||
def test_result_not_compacted(self):
|
||||
"""Test when not compacted."""
|
||||
result = CompactionResult(
|
||||
original_tokens=5000,
|
||||
compacted_tokens=5000,
|
||||
messages_removed=0,
|
||||
messages_kept=10,
|
||||
strategy_used=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
|
||||
assert not result.was_compacted
|
||||
|
||||
def test_result_to_dict(self):
|
||||
"""Test result serialization."""
|
||||
result = CompactionResult(
|
||||
original_tokens=10000,
|
||||
compacted_tokens=6000,
|
||||
messages_removed=5,
|
||||
messages_kept=10,
|
||||
strategy_used=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
data = result.to_dict()
|
||||
|
||||
assert data["original_tokens"] == 10000
|
||||
assert data["tokens_saved"] == 4000
|
||||
assert data["strategy_used"] == "truncate"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ContextCompactor Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestContextCompactor:
|
||||
"""Tests for ContextCompactor class."""
|
||||
|
||||
@pytest.fixture
|
||||
def compactor(self):
|
||||
"""Create a test compactor."""
|
||||
return ContextCompactor(max_tokens=100, preserve_recent=2)
|
||||
|
||||
@pytest.fixture
|
||||
def messages(self):
|
||||
"""Create test messages."""
|
||||
return [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello, how are you?"},
|
||||
{"role": "assistant", "content": "I'm doing well, thank you!"},
|
||||
{"role": "user", "content": "What's the weather like?"},
|
||||
{"role": "assistant", "content": "I don't have access to weather data."},
|
||||
]
|
||||
|
||||
def test_compactor_creation(self, compactor):
|
||||
"""Test creating a compactor."""
|
||||
assert compactor.max_tokens == 100
|
||||
assert compactor.preserve_recent == 2
|
||||
|
||||
def test_compactor_estimate_tokens(self, compactor):
|
||||
"""Test token estimation."""
|
||||
text = "Hello world" # 11 chars
|
||||
tokens = compactor.estimate_tokens(text)
|
||||
|
||||
assert tokens == 2 # 11 // 4
|
||||
|
||||
def test_compactor_count_message_tokens(self, compactor):
|
||||
"""Test message token counting."""
|
||||
message = {"role": "user", "content": "Hello world"}
|
||||
tokens = compactor.count_message_tokens(message)
|
||||
|
||||
assert tokens == 2
|
||||
|
||||
def test_compactor_count_total_tokens(self, compactor, messages):
|
||||
"""Test total token counting."""
|
||||
tokens = compactor.count_total_tokens(messages)
|
||||
|
||||
assert tokens > 0
|
||||
|
||||
def test_compactor_needs_compaction_true(self):
|
||||
"""Test needs_compaction when over limit."""
|
||||
# Create compactor with very low limit
|
||||
compactor = ContextCompactor(max_tokens=10, preserve_recent=2)
|
||||
messages = [
|
||||
{"role": "user", "content": "This is a very long message that should exceed the token limit."},
|
||||
{"role": "assistant", "content": "This is another long response that adds more tokens."},
|
||||
]
|
||||
|
||||
needs = compactor.needs_compaction(messages)
|
||||
|
||||
assert needs
|
||||
|
||||
def test_compactor_needs_compaction_false(self):
|
||||
"""Test needs_compaction when under limit."""
|
||||
compactor = ContextCompactor(max_tokens=10000)
|
||||
messages = [{"role": "user", "content": "Hi"}]
|
||||
|
||||
needs = compactor.needs_compaction(messages)
|
||||
|
||||
assert not needs
|
||||
|
||||
def test_compactor_compact_not_needed(self):
|
||||
"""Test compact when not needed."""
|
||||
compactor = ContextCompactor(max_tokens=10000)
|
||||
messages = [{"role": "user", "content": "Hi"}]
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
assert len(compacted) == len(messages)
|
||||
assert not result.was_compacted
|
||||
|
||||
def test_compactor_compact_truncate(self):
|
||||
"""Test truncate strategy."""
|
||||
compactor = ContextCompactor(max_tokens=10, preserve_recent=1)
|
||||
compactor.strategy = CompactionStrategy.TRUNCATE
|
||||
messages = [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "This is a very long message that should exceed the token limit."},
|
||||
{"role": "assistant", "content": "This is another long response."},
|
||||
{"role": "user", "content": "More content here."},
|
||||
]
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
assert len(compacted) < len(messages)
|
||||
assert result.was_compacted
|
||||
assert result.strategy_used == CompactionStrategy.TRUNCATE
|
||||
|
||||
def test_compactor_compact_sliding(self, compactor, messages):
|
||||
"""Test sliding window strategy."""
|
||||
compactor.strategy = CompactionStrategy.SLIDING
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
assert result.strategy_used == CompactionStrategy.SLIDING
|
||||
|
||||
def test_compactor_compact_summarize(self, compactor, messages):
|
||||
"""Test summarize strategy."""
|
||||
compactor.strategy = CompactionStrategy.SUMMARIZE
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
assert result.strategy_used == CompactionStrategy.SUMMARIZE
|
||||
|
||||
def test_compactor_compact_smart(self, compactor, messages):
|
||||
"""Test smart strategy."""
|
||||
compactor.strategy = CompactionStrategy.SMART
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
assert result.strategy_used == CompactionStrategy.SMART
|
||||
|
||||
def test_compactor_preserves_system(self, compactor, messages):
|
||||
"""Test that system messages are preserved."""
|
||||
compacted, _ = compactor.compact(messages)
|
||||
|
||||
system_msgs = [m for m in compacted if m.get("role") == "system"]
|
||||
assert len(system_msgs) >= 1
|
||||
|
||||
def test_compactor_preserves_recent(self, compactor, messages):
|
||||
"""Test that recent messages are preserved."""
|
||||
compacted, _ = compactor.compact(messages)
|
||||
|
||||
# Should have at least preserve_recent messages
|
||||
non_system = [m for m in compacted if m.get("role") != "system"]
|
||||
assert len(non_system) >= min(compactor.preserve_recent, len(messages) - 1)
|
||||
|
||||
def test_compactor_get_stats(self, compactor, messages):
|
||||
"""Test get_stats method."""
|
||||
stats = compactor.get_stats(messages)
|
||||
|
||||
assert "message_count" in stats
|
||||
assert "total_tokens" in stats
|
||||
assert "needs_compaction" in stats
|
||||
assert stats["message_count"] == len(messages)
|
||||
|
||||
def test_compactor_multipart_content(self, compactor):
|
||||
"""Test handling multipart content."""
|
||||
message = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello"},
|
||||
{"type": "text", "text": "World"}
|
||||
]
|
||||
}
|
||||
|
||||
tokens = compactor.count_message_tokens(message)
|
||||
|
||||
assert tokens > 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for hooks module."""
|
||||
@@ -0,0 +1,717 @@
|
||||
"""
|
||||
Unit tests for the Hooks module.
|
||||
|
||||
Tests cover:
|
||||
- HookRegistry registration and lookup
|
||||
- HookRunner execution (function and command hooks)
|
||||
- Event types and inputs
|
||||
- Decision handling (allow, deny, block)
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
# Import hook types
|
||||
from praisonaiagents.hooks.types import (
|
||||
HookEvent, HookResult, HookInput, HookOutput,
|
||||
HookDefinition, CommandHook, FunctionHook, HookExecutionResult
|
||||
)
|
||||
from praisonaiagents.hooks.events import (
|
||||
BeforeToolInput, AfterToolInput, BeforeAgentInput,
|
||||
AfterAgentInput, SessionStartInput, SessionEndInput
|
||||
)
|
||||
from praisonaiagents.hooks.registry import HookRegistry, get_default_registry
|
||||
from praisonaiagents.hooks.runner import HookRunner
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# HookResult Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestHookResult:
|
||||
"""Tests for HookResult class."""
|
||||
|
||||
def test_allow_result(self):
|
||||
"""Test creating an allow result."""
|
||||
result = HookResult.allow("Allowed by policy")
|
||||
assert result.decision == "allow"
|
||||
assert result.reason == "Allowed by policy"
|
||||
assert result.is_allowed()
|
||||
assert not result.is_denied()
|
||||
|
||||
def test_deny_result(self):
|
||||
"""Test creating a deny result."""
|
||||
result = HookResult.deny("Blocked by policy")
|
||||
assert result.decision == "deny"
|
||||
assert result.reason == "Blocked by policy"
|
||||
assert not result.is_allowed()
|
||||
assert result.is_denied()
|
||||
|
||||
def test_block_result(self):
|
||||
"""Test creating a block result."""
|
||||
result = HookResult.block("Security violation")
|
||||
assert result.decision == "block"
|
||||
assert result.reason == "Security violation"
|
||||
assert result.is_denied()
|
||||
|
||||
def test_ask_result(self):
|
||||
"""Test creating an ask result."""
|
||||
result = HookResult.ask("Requires confirmation")
|
||||
assert result.decision == "ask"
|
||||
assert result.reason == "Requires confirmation"
|
||||
|
||||
def test_default_result(self):
|
||||
"""Test default result is allow."""
|
||||
result = HookResult()
|
||||
assert result.decision == "allow"
|
||||
assert result.is_allowed()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# HookDefinition Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestHookDefinition:
|
||||
"""Tests for HookDefinition class."""
|
||||
|
||||
def test_definition_creation(self):
|
||||
"""Test creating a hook definition."""
|
||||
hook = HookDefinition(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
name="test_hook",
|
||||
matcher="write_*"
|
||||
)
|
||||
assert hook.event == HookEvent.BEFORE_TOOL
|
||||
assert hook.name == "test_hook"
|
||||
assert hook.matcher == "write_*"
|
||||
assert hook.enabled
|
||||
|
||||
def test_matcher_regex(self):
|
||||
"""Test regex matching."""
|
||||
hook = HookDefinition(matcher="write_.*")
|
||||
assert hook.matches("write_file")
|
||||
assert hook.matches("write_data")
|
||||
assert not hook.matches("read_file")
|
||||
|
||||
def test_matcher_none(self):
|
||||
"""Test that None matcher matches everything."""
|
||||
hook = HookDefinition(matcher=None)
|
||||
assert hook.matches("anything")
|
||||
assert hook.matches("write_file")
|
||||
|
||||
def test_command_hook(self):
|
||||
"""Test CommandHook creation."""
|
||||
hook = CommandHook(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="python validator.py",
|
||||
env={"API_KEY": "test"}
|
||||
)
|
||||
assert hook.command == "python validator.py"
|
||||
assert hook.env == {"API_KEY": "test"}
|
||||
assert hook.shell
|
||||
|
||||
def test_function_hook(self):
|
||||
"""Test FunctionHook creation."""
|
||||
def my_hook(data):
|
||||
return HookResult.allow()
|
||||
|
||||
hook = FunctionHook(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=my_hook
|
||||
)
|
||||
assert hook.func == my_hook
|
||||
assert hook.name == "my_hook"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# HookRegistry Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestHookRegistry:
|
||||
"""Tests for HookRegistry class."""
|
||||
|
||||
def test_registry_creation(self):
|
||||
"""Test creating a registry."""
|
||||
registry = HookRegistry()
|
||||
assert registry.enabled
|
||||
assert len(registry) == 0
|
||||
|
||||
def test_register_function(self):
|
||||
"""Test registering a function hook."""
|
||||
registry = HookRegistry()
|
||||
|
||||
def my_hook(data):
|
||||
return HookResult.allow()
|
||||
|
||||
hook_id = registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=my_hook,
|
||||
name="my_hook"
|
||||
)
|
||||
|
||||
assert hook_id is not None
|
||||
assert len(registry) == 1
|
||||
|
||||
hooks = registry.get_hooks(HookEvent.BEFORE_TOOL)
|
||||
assert len(hooks) == 1
|
||||
assert hooks[0].name == "my_hook"
|
||||
|
||||
def test_register_command(self):
|
||||
"""Test registering a command hook."""
|
||||
registry = HookRegistry()
|
||||
|
||||
hook_id = registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="python validator.py",
|
||||
name="validator"
|
||||
)
|
||||
|
||||
assert hook_id is not None
|
||||
hooks = registry.get_hooks(HookEvent.BEFORE_TOOL)
|
||||
assert len(hooks) == 1
|
||||
assert isinstance(hooks[0], CommandHook)
|
||||
|
||||
def test_decorator_registration(self):
|
||||
"""Test decorator-based registration."""
|
||||
registry = HookRegistry()
|
||||
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def validate_tool(data):
|
||||
return HookResult.allow()
|
||||
|
||||
hooks = registry.get_hooks(HookEvent.BEFORE_TOOL)
|
||||
assert len(hooks) == 1
|
||||
assert hooks[0].name == "validate_tool"
|
||||
|
||||
def test_unregister(self):
|
||||
"""Test unregistering a hook."""
|
||||
registry = HookRegistry()
|
||||
|
||||
hook_id = registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="echo test"
|
||||
)
|
||||
|
||||
assert len(registry) == 1
|
||||
assert registry.unregister(hook_id)
|
||||
assert len(registry) == 0
|
||||
|
||||
def test_get_hooks_with_matcher(self):
|
||||
"""Test getting hooks filtered by target."""
|
||||
registry = HookRegistry()
|
||||
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=lambda d: HookResult.allow(),
|
||||
matcher="write_.*",
|
||||
name="write_hook"
|
||||
)
|
||||
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=lambda d: HookResult.allow(),
|
||||
matcher="read_.*",
|
||||
name="read_hook"
|
||||
)
|
||||
|
||||
write_hooks = registry.get_hooks(HookEvent.BEFORE_TOOL, "write_file")
|
||||
assert len(write_hooks) == 1
|
||||
assert write_hooks[0].name == "write_hook"
|
||||
|
||||
read_hooks = registry.get_hooks(HookEvent.BEFORE_TOOL, "read_file")
|
||||
assert len(read_hooks) == 1
|
||||
assert read_hooks[0].name == "read_hook"
|
||||
|
||||
def test_disable_enable_hook(self):
|
||||
"""Test disabling and enabling hooks."""
|
||||
registry = HookRegistry()
|
||||
|
||||
hook_id = registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="echo test"
|
||||
)
|
||||
|
||||
assert len(registry.get_hooks(HookEvent.BEFORE_TOOL)) == 1
|
||||
|
||||
registry.disable_hook(hook_id)
|
||||
assert len(registry.get_hooks(HookEvent.BEFORE_TOOL)) == 0
|
||||
|
||||
registry.enable_hook(hook_id)
|
||||
assert len(registry.get_hooks(HookEvent.BEFORE_TOOL)) == 1
|
||||
|
||||
def test_clear_hooks(self):
|
||||
"""Test clearing hooks."""
|
||||
registry = HookRegistry()
|
||||
|
||||
registry.register_command(HookEvent.BEFORE_TOOL, "echo 1")
|
||||
registry.register_command(HookEvent.AFTER_TOOL, "echo 2")
|
||||
|
||||
assert len(registry) == 2
|
||||
|
||||
registry.clear(HookEvent.BEFORE_TOOL)
|
||||
assert len(registry) == 1
|
||||
|
||||
registry.clear()
|
||||
assert len(registry) == 0
|
||||
|
||||
def test_list_hooks(self):
|
||||
"""Test listing all hooks."""
|
||||
registry = HookRegistry()
|
||||
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="echo test",
|
||||
name="test_hook"
|
||||
)
|
||||
|
||||
hooks_list = registry.list_hooks()
|
||||
assert "before_tool" in hooks_list
|
||||
assert len(hooks_list["before_tool"]) == 1
|
||||
assert hooks_list["before_tool"][0]["name"] == "test_hook"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Event Input Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestEventInputs:
|
||||
"""Tests for event-specific input classes."""
|
||||
|
||||
def test_before_tool_input(self):
|
||||
"""Test BeforeToolInput."""
|
||||
input_data = BeforeToolInput(
|
||||
session_id="test-session",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="write_file",
|
||||
tool_input={"path": "/tmp/test.txt", "content": "hello"}
|
||||
)
|
||||
|
||||
assert input_data.tool_name == "write_file"
|
||||
assert input_data.tool_input["path"] == "/tmp/test.txt"
|
||||
|
||||
data_dict = input_data.to_dict()
|
||||
assert data_dict["tool_name"] == "write_file"
|
||||
|
||||
def test_after_tool_input(self):
|
||||
"""Test AfterToolInput."""
|
||||
input_data = AfterToolInput(
|
||||
session_id="test-session",
|
||||
cwd="/tmp",
|
||||
event_name="after_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="write_file",
|
||||
tool_output="File written successfully",
|
||||
execution_time_ms=100.5
|
||||
)
|
||||
|
||||
assert input_data.tool_output == "File written successfully"
|
||||
assert input_data.execution_time_ms == 100.5
|
||||
|
||||
def test_before_agent_input(self):
|
||||
"""Test BeforeAgentInput."""
|
||||
input_data = BeforeAgentInput(
|
||||
session_id="test-session",
|
||||
cwd="/tmp",
|
||||
event_name="before_agent",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
prompt="Write a file",
|
||||
tools_available=["write_file", "read_file"]
|
||||
)
|
||||
|
||||
assert input_data.prompt == "Write a file"
|
||||
assert "write_file" in input_data.tools_available
|
||||
|
||||
def test_session_start_input(self):
|
||||
"""Test SessionStartInput."""
|
||||
input_data = SessionStartInput(
|
||||
session_id="test-session",
|
||||
cwd="/tmp",
|
||||
event_name="session_start",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
source="startup"
|
||||
)
|
||||
|
||||
assert input_data.source == "startup"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# HookRunner Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestHookRunner:
|
||||
"""Tests for HookRunner class."""
|
||||
|
||||
def _create_sample_input(self):
|
||||
"""Create sample input data."""
|
||||
return BeforeToolInput(
|
||||
session_id="test-session",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="write_file",
|
||||
tool_input={"path": "/tmp/test.txt"}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_function_hook_allow(self):
|
||||
"""Test executing a function hook that allows."""
|
||||
registry = HookRegistry()
|
||||
|
||||
def allow_hook(data):
|
||||
return HookResult.allow("Allowed")
|
||||
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=allow_hook
|
||||
)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].success
|
||||
assert results[0].output.is_allowed()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_function_hook_deny(self):
|
||||
"""Test executing a function hook that denies."""
|
||||
registry = HookRegistry()
|
||||
|
||||
def deny_hook(data):
|
||||
return HookResult.deny("Not allowed")
|
||||
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=deny_hook
|
||||
)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].success
|
||||
assert results[0].output.is_denied()
|
||||
assert HookRunner.is_blocked(results)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_async_function_hook(self):
|
||||
"""Test executing an async function hook."""
|
||||
registry = HookRegistry()
|
||||
|
||||
async def async_hook(data):
|
||||
await asyncio.sleep(0.01)
|
||||
return HookResult.allow("Async allowed")
|
||||
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=async_hook,
|
||||
is_async=True
|
||||
)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].success
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_multiple_hooks_parallel(self):
|
||||
"""Test executing multiple hooks in parallel."""
|
||||
registry = HookRegistry()
|
||||
call_order = []
|
||||
|
||||
def hook1(data):
|
||||
call_order.append(1)
|
||||
return HookResult.allow()
|
||||
|
||||
def hook2(data):
|
||||
call_order.append(2)
|
||||
return HookResult.allow()
|
||||
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, hook1, sequential=False)
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, hook2, sequential=False)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 2
|
||||
assert all(r.success for r in results)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sequential_hooks(self):
|
||||
"""Test executing hooks sequentially."""
|
||||
registry = HookRegistry()
|
||||
call_order = []
|
||||
|
||||
def hook1(data):
|
||||
call_order.append(1)
|
||||
return HookResult.allow()
|
||||
|
||||
def hook2(data):
|
||||
call_order.append(2)
|
||||
return HookResult.allow()
|
||||
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, hook1, sequential=True)
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, hook2, sequential=True)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 2
|
||||
assert call_order == [1, 2]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sequential_stops_on_deny(self):
|
||||
"""Test that sequential execution stops on deny."""
|
||||
registry = HookRegistry()
|
||||
call_order = []
|
||||
|
||||
def hook1(data):
|
||||
call_order.append(1)
|
||||
return HookResult.deny("Blocked")
|
||||
|
||||
def hook2(data):
|
||||
call_order.append(2)
|
||||
return HookResult.allow()
|
||||
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, hook1, sequential=True)
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, hook2, sequential=True)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1 # Only first hook executed
|
||||
assert call_order == [1]
|
||||
assert HookRunner.is_blocked(results)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_timeout(self):
|
||||
"""Test hook timeout handling."""
|
||||
registry = HookRegistry()
|
||||
|
||||
async def slow_hook(data):
|
||||
await asyncio.sleep(10) # Very slow
|
||||
return HookResult.allow()
|
||||
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=slow_hook,
|
||||
is_async=True,
|
||||
timeout=0.1 # 100ms timeout
|
||||
)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert not results[0].success
|
||||
assert "timed out" in results[0].error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_exception_handling(self):
|
||||
"""Test hook exception handling."""
|
||||
registry = HookRegistry()
|
||||
|
||||
def failing_hook(data):
|
||||
raise ValueError("Hook failed")
|
||||
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=failing_hook
|
||||
)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert not results[0].success
|
||||
assert "Hook failed" in results[0].error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_hooks_returns_empty(self):
|
||||
"""Test that no hooks returns empty list."""
|
||||
registry = HookRegistry()
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
assert results == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_registry(self):
|
||||
"""Test that disabled registry returns no hooks."""
|
||||
registry = HookRegistry()
|
||||
registry.register_function(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
func=lambda d: HookResult.allow()
|
||||
)
|
||||
|
||||
registry.enabled = False
|
||||
runner = HookRunner(registry)
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_is_blocked_helper(self):
|
||||
"""Test is_blocked helper method."""
|
||||
allow_result = HookExecutionResult(
|
||||
hook_id="1",
|
||||
hook_name="test",
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
success=True,
|
||||
output=HookResult.allow()
|
||||
)
|
||||
|
||||
deny_result = HookExecutionResult(
|
||||
hook_id="2",
|
||||
hook_name="test",
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
success=True,
|
||||
output=HookResult.deny("Blocked")
|
||||
)
|
||||
|
||||
assert not HookRunner.is_blocked([allow_result])
|
||||
assert HookRunner.is_blocked([deny_result])
|
||||
assert HookRunner.is_blocked([allow_result, deny_result])
|
||||
|
||||
def test_get_blocking_reason(self):
|
||||
"""Test get_blocking_reason helper method."""
|
||||
deny_result = HookExecutionResult(
|
||||
hook_id="1",
|
||||
hook_name="test",
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
success=True,
|
||||
output=HookResult.deny("Security violation")
|
||||
)
|
||||
|
||||
reason = HookRunner.get_blocking_reason([deny_result])
|
||||
assert reason == "Security violation"
|
||||
|
||||
def test_aggregate_context(self):
|
||||
"""Test aggregate_context helper method."""
|
||||
result1 = HookExecutionResult(
|
||||
hook_id="1",
|
||||
hook_name="test1",
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
success=True,
|
||||
output=HookResult(additional_context="Context 1")
|
||||
)
|
||||
|
||||
result2 = HookExecutionResult(
|
||||
hook_id="2",
|
||||
hook_name="test2",
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
success=True,
|
||||
output=HookResult(additional_context="Context 2")
|
||||
)
|
||||
|
||||
context = HookRunner.aggregate_context([result1, result2])
|
||||
assert "Context 1" in context
|
||||
assert "Context 2" in context
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Command Hook Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCommandHooks:
|
||||
"""Tests for command hook execution."""
|
||||
|
||||
def _create_sample_input(self):
|
||||
return BeforeToolInput(
|
||||
session_id="test-session",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="write_file",
|
||||
tool_input={"path": "/tmp/test.txt"}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_command_hook_success(self):
|
||||
"""Test successful command hook execution."""
|
||||
registry = HookRegistry()
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command='echo \'{"decision": "allow"}\'',
|
||||
name="allow_hook"
|
||||
)
|
||||
|
||||
runner = HookRunner(registry, cwd="/tmp")
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].success
|
||||
assert results[0].output.decision == "allow"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_command_hook_deny(self):
|
||||
"""Test command hook that denies."""
|
||||
registry = HookRegistry()
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command='echo \'{"decision": "deny", "reason": "Blocked"}\''
|
||||
)
|
||||
|
||||
runner = HookRunner(registry, cwd="/tmp")
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].output.decision == "deny"
|
||||
assert results[0].output.reason == "Blocked"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_command_hook_exit_code_2(self):
|
||||
"""Test command hook with blocking exit code."""
|
||||
registry = HookRegistry()
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="exit 2"
|
||||
)
|
||||
|
||||
runner = HookRunner(registry, cwd="/tmp")
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].exit_code == 2
|
||||
assert results[0].output.decision == "deny"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_command_hook_timeout(self):
|
||||
"""Test command hook timeout."""
|
||||
registry = HookRegistry()
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="sleep 10",
|
||||
timeout=0.1
|
||||
)
|
||||
|
||||
runner = HookRunner(registry, cwd="/tmp")
|
||||
sample_input = self._create_sample_input()
|
||||
results = await runner.execute(HookEvent.BEFORE_TOOL, sample_input)
|
||||
|
||||
assert len(results) == 1
|
||||
assert not results[0].success
|
||||
assert "timed out" in results[0].error
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for LSP module."""
|
||||
@@ -0,0 +1,327 @@
|
||||
"""
|
||||
Unit tests for the LSP module.
|
||||
|
||||
Tests cover:
|
||||
- LSP types (Position, Range, Diagnostic, etc.)
|
||||
- LSPConfig creation and defaults
|
||||
- LSPClient initialization
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from praisonaiagents.lsp.types import (
|
||||
Position, Range, Location, Diagnostic, DiagnosticSeverity,
|
||||
CompletionItem, TextDocumentItem, TextDocumentIdentifier
|
||||
)
|
||||
from praisonaiagents.lsp.config import LSPConfig, DEFAULT_SERVERS
|
||||
from praisonaiagents.lsp.client import LSPClient
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Position Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestPosition:
|
||||
"""Tests for Position class."""
|
||||
|
||||
def test_position_creation(self):
|
||||
"""Test creating a position."""
|
||||
pos = Position(line=10, character=5)
|
||||
assert pos.line == 10
|
||||
assert pos.character == 5
|
||||
|
||||
def test_position_to_dict(self):
|
||||
"""Test position serialization."""
|
||||
pos = Position(line=10, character=5)
|
||||
data = pos.to_dict()
|
||||
assert data == {"line": 10, "character": 5}
|
||||
|
||||
def test_position_from_dict(self):
|
||||
"""Test position deserialization."""
|
||||
pos = Position.from_dict({"line": 10, "character": 5})
|
||||
assert pos.line == 10
|
||||
assert pos.character == 5
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Range Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestRange:
|
||||
"""Tests for Range class."""
|
||||
|
||||
def test_range_creation(self):
|
||||
"""Test creating a range."""
|
||||
start = Position(line=0, character=0)
|
||||
end = Position(line=10, character=20)
|
||||
r = Range(start=start, end=end)
|
||||
|
||||
assert r.start.line == 0
|
||||
assert r.end.line == 10
|
||||
|
||||
def test_range_to_dict(self):
|
||||
"""Test range serialization."""
|
||||
r = Range(
|
||||
start=Position(0, 0),
|
||||
end=Position(10, 20)
|
||||
)
|
||||
data = r.to_dict()
|
||||
|
||||
assert data["start"]["line"] == 0
|
||||
assert data["end"]["line"] == 10
|
||||
|
||||
def test_range_from_dict(self):
|
||||
"""Test range deserialization."""
|
||||
r = Range.from_dict({
|
||||
"start": {"line": 0, "character": 0},
|
||||
"end": {"line": 10, "character": 20}
|
||||
})
|
||||
|
||||
assert r.start.line == 0
|
||||
assert r.end.character == 20
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Location Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestLocation:
|
||||
"""Tests for Location class."""
|
||||
|
||||
def test_location_creation(self):
|
||||
"""Test creating a location."""
|
||||
loc = Location(
|
||||
uri="file:///test.py",
|
||||
range=Range(Position(0, 0), Position(10, 0))
|
||||
)
|
||||
|
||||
assert loc.uri == "file:///test.py"
|
||||
assert loc.range.start.line == 0
|
||||
|
||||
def test_location_to_dict(self):
|
||||
"""Test location serialization."""
|
||||
loc = Location(
|
||||
uri="file:///test.py",
|
||||
range=Range(Position(0, 0), Position(10, 0))
|
||||
)
|
||||
data = loc.to_dict()
|
||||
|
||||
assert data["uri"] == "file:///test.py"
|
||||
assert "range" in data
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Diagnostic Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestDiagnostic:
|
||||
"""Tests for Diagnostic class."""
|
||||
|
||||
def test_diagnostic_creation(self):
|
||||
"""Test creating a diagnostic."""
|
||||
diag = Diagnostic(
|
||||
range=Range(Position(5, 0), Position(5, 10)),
|
||||
message="Undefined variable 'x'",
|
||||
severity=DiagnosticSeverity.ERROR
|
||||
)
|
||||
|
||||
assert diag.message == "Undefined variable 'x'"
|
||||
assert diag.severity == DiagnosticSeverity.ERROR
|
||||
|
||||
def test_diagnostic_default_severity(self):
|
||||
"""Test default severity is ERROR."""
|
||||
diag = Diagnostic(
|
||||
range=Range(Position(0, 0), Position(0, 10)),
|
||||
message="Test"
|
||||
)
|
||||
|
||||
assert diag.severity == DiagnosticSeverity.ERROR
|
||||
|
||||
def test_diagnostic_to_dict(self):
|
||||
"""Test diagnostic serialization."""
|
||||
diag = Diagnostic(
|
||||
range=Range(Position(5, 0), Position(5, 10)),
|
||||
message="Test error",
|
||||
severity=DiagnosticSeverity.WARNING,
|
||||
code="W001",
|
||||
source="test"
|
||||
)
|
||||
data = diag.to_dict()
|
||||
|
||||
assert data["message"] == "Test error"
|
||||
assert data["severity"] == 2 # WARNING
|
||||
assert data["code"] == "W001"
|
||||
assert data["source"] == "test"
|
||||
|
||||
def test_diagnostic_from_dict(self):
|
||||
"""Test diagnostic deserialization."""
|
||||
diag = Diagnostic.from_dict({
|
||||
"range": {
|
||||
"start": {"line": 5, "character": 0},
|
||||
"end": {"line": 5, "character": 10}
|
||||
},
|
||||
"message": "Test error",
|
||||
"severity": 2,
|
||||
"code": "W001"
|
||||
})
|
||||
|
||||
assert diag.message == "Test error"
|
||||
assert diag.severity == DiagnosticSeverity.WARNING
|
||||
assert diag.code == "W001"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CompletionItem Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestCompletionItem:
|
||||
"""Tests for CompletionItem class."""
|
||||
|
||||
def test_completion_creation(self):
|
||||
"""Test creating a completion item."""
|
||||
item = CompletionItem(
|
||||
label="print",
|
||||
kind=3, # Function
|
||||
detail="Built-in function"
|
||||
)
|
||||
|
||||
assert item.label == "print"
|
||||
assert item.kind == 3
|
||||
assert item.detail == "Built-in function"
|
||||
|
||||
def test_completion_to_dict(self):
|
||||
"""Test completion serialization."""
|
||||
item = CompletionItem(
|
||||
label="print",
|
||||
kind=3,
|
||||
detail="Built-in function",
|
||||
insert_text="print($1)"
|
||||
)
|
||||
data = item.to_dict()
|
||||
|
||||
assert data["label"] == "print"
|
||||
assert data["insertText"] == "print($1)"
|
||||
|
||||
def test_completion_from_dict(self):
|
||||
"""Test completion deserialization."""
|
||||
item = CompletionItem.from_dict({
|
||||
"label": "print",
|
||||
"kind": 3,
|
||||
"detail": "Built-in function",
|
||||
"insertText": "print($1)"
|
||||
})
|
||||
|
||||
assert item.label == "print"
|
||||
assert item.insert_text == "print($1)"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TextDocument Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestTextDocumentTypes:
|
||||
"""Tests for text document types."""
|
||||
|
||||
def test_text_document_item(self):
|
||||
"""Test TextDocumentItem."""
|
||||
doc = TextDocumentItem(
|
||||
uri="file:///test.py",
|
||||
language_id="python",
|
||||
version=1,
|
||||
text="print('hello')"
|
||||
)
|
||||
data = doc.to_dict()
|
||||
|
||||
assert data["uri"] == "file:///test.py"
|
||||
assert data["languageId"] == "python"
|
||||
assert data["version"] == 1
|
||||
|
||||
def test_text_document_identifier(self):
|
||||
"""Test TextDocumentIdentifier."""
|
||||
doc = TextDocumentIdentifier(uri="file:///test.py")
|
||||
data = doc.to_dict()
|
||||
|
||||
assert data["uri"] == "file:///test.py"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# LSPConfig Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestLSPConfig:
|
||||
"""Tests for LSPConfig class."""
|
||||
|
||||
def test_config_with_defaults(self):
|
||||
"""Test config with default server."""
|
||||
config = LSPConfig(language="python")
|
||||
|
||||
assert config.language == "python"
|
||||
assert config.command == "pylsp"
|
||||
assert config.args == []
|
||||
|
||||
def test_config_custom_command(self):
|
||||
"""Test config with custom command."""
|
||||
config = LSPConfig(
|
||||
language="python",
|
||||
command="pyright-langserver",
|
||||
args=["--stdio"]
|
||||
)
|
||||
|
||||
assert config.command == "pyright-langserver"
|
||||
assert config.args == ["--stdio"]
|
||||
|
||||
def test_config_unknown_language(self):
|
||||
"""Test config with unknown language raises error."""
|
||||
with pytest.raises(ValueError):
|
||||
LSPConfig(language="unknown_language")
|
||||
|
||||
def test_default_servers_exist(self):
|
||||
"""Test that default servers are defined."""
|
||||
assert "python" in DEFAULT_SERVERS
|
||||
assert "javascript" in DEFAULT_SERVERS
|
||||
assert "typescript" in DEFAULT_SERVERS
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# LSPClient Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestLSPClient:
|
||||
"""Tests for LSPClient class."""
|
||||
|
||||
def test_client_creation(self):
|
||||
"""Test creating a client."""
|
||||
client = LSPClient(language="python")
|
||||
|
||||
assert client.config.language == "python"
|
||||
assert not client.is_running
|
||||
|
||||
def test_client_custom_command(self):
|
||||
"""Test client with custom command."""
|
||||
client = LSPClient(
|
||||
language="python",
|
||||
command="custom-server",
|
||||
args=["--mode", "stdio"]
|
||||
)
|
||||
|
||||
assert client.config.command == "custom-server"
|
||||
assert client.config.args == ["--mode", "stdio"]
|
||||
|
||||
def test_client_root_uri(self):
|
||||
"""Test client with root URI."""
|
||||
client = LSPClient(
|
||||
language="python",
|
||||
root_uri="file:///workspace"
|
||||
)
|
||||
|
||||
assert client.config.root_uri == "file:///workspace"
|
||||
|
||||
def test_client_not_running_initially(self):
|
||||
"""Test client is not running initially."""
|
||||
client = LSPClient(language="python")
|
||||
|
||||
assert not client.is_running
|
||||
assert not client._initialized
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for policy module."""
|
||||
@@ -0,0 +1,508 @@
|
||||
"""
|
||||
Unit tests for the Policy module.
|
||||
|
||||
Tests cover:
|
||||
- PolicyResult creation
|
||||
- PolicyRule matching and evaluation
|
||||
- Policy evaluation
|
||||
- PolicyEngine management and checks
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.policy.types import PolicyAction, PolicyResult
|
||||
from praisonaiagents.policy.policy import Policy, PolicyRule
|
||||
from praisonaiagents.policy.config import PolicyConfig
|
||||
from praisonaiagents.policy.engine import (
|
||||
PolicyEngine, create_deny_tools_policy,
|
||||
create_allow_tools_policy, create_read_only_policy
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PolicyResult Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestPolicyResult:
|
||||
"""Tests for PolicyResult class."""
|
||||
|
||||
def test_result_creation(self):
|
||||
"""Test creating a result."""
|
||||
result = PolicyResult(
|
||||
allowed=True,
|
||||
action=PolicyAction.ALLOW,
|
||||
reason="Test"
|
||||
)
|
||||
assert result.allowed
|
||||
assert result.action == PolicyAction.ALLOW
|
||||
|
||||
def test_result_allow(self):
|
||||
"""Test allow factory method."""
|
||||
result = PolicyResult.allow("Allowed")
|
||||
assert result.allowed
|
||||
assert result.action == PolicyAction.ALLOW
|
||||
assert result.reason == "Allowed"
|
||||
|
||||
def test_result_deny(self):
|
||||
"""Test deny factory method."""
|
||||
result = PolicyResult.deny("Not allowed", "test_policy")
|
||||
assert not result.allowed
|
||||
assert result.action == PolicyAction.DENY
|
||||
assert result.reason == "Not allowed"
|
||||
assert result.policy_name == "test_policy"
|
||||
|
||||
def test_result_ask(self):
|
||||
"""Test ask factory method."""
|
||||
result = PolicyResult.ask("Confirm action")
|
||||
assert not result.allowed
|
||||
assert result.action == PolicyAction.ASK
|
||||
|
||||
def test_result_to_dict(self):
|
||||
"""Test result serialization."""
|
||||
result = PolicyResult(
|
||||
allowed=False,
|
||||
action=PolicyAction.DENY,
|
||||
policy_name="test",
|
||||
reason="Test reason"
|
||||
)
|
||||
data = result.to_dict()
|
||||
|
||||
assert data["allowed"] is False
|
||||
assert data["action"] == "deny"
|
||||
assert data["policy_name"] == "test"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PolicyRule Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestPolicyRule:
|
||||
"""Tests for PolicyRule class."""
|
||||
|
||||
def test_rule_creation(self):
|
||||
"""Test creating a rule."""
|
||||
rule = PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="No delete"
|
||||
)
|
||||
assert rule.action == PolicyAction.DENY
|
||||
assert rule.resource == "tool:delete_*"
|
||||
|
||||
def test_rule_exact_match(self):
|
||||
"""Test exact resource matching."""
|
||||
rule = PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_file"
|
||||
)
|
||||
|
||||
assert rule.matches("tool:delete_file", {})
|
||||
assert not rule.matches("tool:read_file", {})
|
||||
|
||||
def test_rule_wildcard_match(self):
|
||||
"""Test wildcard resource matching."""
|
||||
rule = PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*"
|
||||
)
|
||||
|
||||
assert rule.matches("tool:delete_file", {})
|
||||
assert rule.matches("tool:delete_directory", {})
|
||||
assert not rule.matches("tool:read_file", {})
|
||||
|
||||
def test_rule_condition(self):
|
||||
"""Test rule with condition."""
|
||||
rule = PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:*",
|
||||
condition=lambda ctx: ctx.get("dangerous", False)
|
||||
)
|
||||
|
||||
assert rule.matches("tool:test", {"dangerous": True})
|
||||
assert not rule.matches("tool:test", {"dangerous": False})
|
||||
assert not rule.matches("tool:test", {})
|
||||
|
||||
def test_rule_evaluate_match(self):
|
||||
"""Test rule evaluation when matching."""
|
||||
rule = PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="No delete",
|
||||
name="deny_delete"
|
||||
)
|
||||
|
||||
result = rule.evaluate("tool:delete_file", {})
|
||||
|
||||
assert result is not None
|
||||
assert not result.allowed
|
||||
assert result.action == PolicyAction.DENY
|
||||
assert result.reason == "No delete"
|
||||
|
||||
def test_rule_evaluate_no_match(self):
|
||||
"""Test rule evaluation when not matching."""
|
||||
rule = PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*"
|
||||
)
|
||||
|
||||
result = rule.evaluate("tool:read_file", {})
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_rule_to_dict(self):
|
||||
"""Test rule serialization."""
|
||||
rule = PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:test",
|
||||
reason="Test",
|
||||
name="test_rule",
|
||||
priority=10
|
||||
)
|
||||
data = rule.to_dict()
|
||||
|
||||
assert data["action"] == "deny"
|
||||
assert data["resource"] == "tool:test"
|
||||
assert data["priority"] == 10
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Policy Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestPolicy:
|
||||
"""Tests for Policy class."""
|
||||
|
||||
def test_policy_creation(self):
|
||||
"""Test creating a policy."""
|
||||
policy = Policy(
|
||||
name="test_policy",
|
||||
description="Test policy"
|
||||
)
|
||||
assert policy.name == "test_policy"
|
||||
assert policy.enabled
|
||||
|
||||
def test_policy_add_rule(self):
|
||||
"""Test adding rules to a policy."""
|
||||
policy = Policy(name="test")
|
||||
|
||||
policy.add_rule(PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*"
|
||||
))
|
||||
|
||||
assert len(policy.rules) == 1
|
||||
|
||||
def test_policy_remove_rule(self):
|
||||
"""Test removing rules from a policy."""
|
||||
policy = Policy(name="test")
|
||||
policy.add_rule(PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
name="deny_delete"
|
||||
))
|
||||
|
||||
result = policy.remove_rule("deny_delete")
|
||||
|
||||
assert result
|
||||
assert len(policy.rules) == 0
|
||||
|
||||
def test_policy_evaluate_match(self):
|
||||
"""Test policy evaluation when rule matches."""
|
||||
policy = Policy(
|
||||
name="test",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="No delete"
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
result = policy.evaluate("tool:delete_file", {})
|
||||
|
||||
assert result is not None
|
||||
assert not result.allowed
|
||||
assert result.policy_name == "test"
|
||||
|
||||
def test_policy_evaluate_no_match(self):
|
||||
"""Test policy evaluation when no rule matches."""
|
||||
policy = Policy(
|
||||
name="test",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*"
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
result = policy.evaluate("tool:read_file", {})
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_policy_disabled(self):
|
||||
"""Test disabled policy."""
|
||||
policy = Policy(
|
||||
name="test",
|
||||
enabled=False,
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:*"
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
result = policy.evaluate("tool:delete_file", {})
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_policy_to_dict(self):
|
||||
"""Test policy serialization."""
|
||||
policy = Policy(
|
||||
name="test",
|
||||
description="Test policy",
|
||||
rules=[
|
||||
PolicyRule(action=PolicyAction.DENY, resource="tool:*")
|
||||
]
|
||||
)
|
||||
data = policy.to_dict()
|
||||
|
||||
assert data["name"] == "test"
|
||||
assert len(data["rules"]) == 1
|
||||
|
||||
def test_policy_from_dict(self):
|
||||
"""Test policy deserialization."""
|
||||
data = {
|
||||
"name": "test",
|
||||
"description": "Test",
|
||||
"rules": [
|
||||
{"action": "deny", "resource": "tool:delete_*"}
|
||||
]
|
||||
}
|
||||
|
||||
policy = Policy.from_dict(data)
|
||||
|
||||
assert policy.name == "test"
|
||||
assert len(policy.rules) == 1
|
||||
assert policy.rules[0].action == PolicyAction.DENY
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PolicyConfig Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestPolicyConfig:
|
||||
"""Tests for PolicyConfig class."""
|
||||
|
||||
def test_config_defaults(self):
|
||||
"""Test default configuration."""
|
||||
config = PolicyConfig()
|
||||
|
||||
assert config.enabled
|
||||
assert config.default_action == "allow"
|
||||
assert not config.strict_mode
|
||||
|
||||
def test_config_custom(self):
|
||||
"""Test custom configuration."""
|
||||
config = PolicyConfig(
|
||||
enabled=False,
|
||||
strict_mode=True
|
||||
)
|
||||
|
||||
assert not config.enabled
|
||||
assert config.strict_mode
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PolicyEngine Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestPolicyEngine:
|
||||
"""Tests for PolicyEngine class."""
|
||||
|
||||
@pytest.fixture
|
||||
def engine(self):
|
||||
"""Create a test engine."""
|
||||
return PolicyEngine()
|
||||
|
||||
def test_engine_creation(self, engine):
|
||||
"""Test creating an engine."""
|
||||
assert engine.config.enabled
|
||||
assert len(engine.policies) == 0
|
||||
|
||||
def test_engine_add_policy(self, engine):
|
||||
"""Test adding a policy."""
|
||||
policy = Policy(name="test")
|
||||
engine.add_policy(policy)
|
||||
|
||||
assert len(engine.policies) == 1
|
||||
assert engine.get_policy("test") is not None
|
||||
|
||||
def test_engine_remove_policy(self, engine):
|
||||
"""Test removing a policy."""
|
||||
engine.add_policy(Policy(name="test"))
|
||||
|
||||
result = engine.remove_policy("test")
|
||||
|
||||
assert result
|
||||
assert len(engine.policies) == 0
|
||||
|
||||
def test_engine_check_allow(self, engine):
|
||||
"""Test check that allows."""
|
||||
engine.add_policy(Policy(
|
||||
name="allow_read",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.ALLOW,
|
||||
resource="tool:read_*"
|
||||
)
|
||||
]
|
||||
))
|
||||
|
||||
result = engine.check("tool:read_file", {})
|
||||
|
||||
assert result.allowed
|
||||
|
||||
def test_engine_check_deny(self, engine):
|
||||
"""Test check that denies."""
|
||||
engine.add_policy(Policy(
|
||||
name="deny_delete",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="No delete"
|
||||
)
|
||||
]
|
||||
))
|
||||
|
||||
result = engine.check("tool:delete_file", {})
|
||||
|
||||
assert not result.allowed
|
||||
assert result.reason == "No delete"
|
||||
|
||||
def test_engine_check_no_match(self, engine):
|
||||
"""Test check with no matching policy."""
|
||||
result = engine.check("tool:anything", {})
|
||||
|
||||
assert result.allowed # Default is allow
|
||||
|
||||
def test_engine_strict_mode(self):
|
||||
"""Test strict mode denies when no match."""
|
||||
engine = PolicyEngine(PolicyConfig(strict_mode=True))
|
||||
|
||||
result = engine.check("tool:anything", {})
|
||||
|
||||
assert not result.allowed
|
||||
|
||||
def test_engine_disabled(self):
|
||||
"""Test disabled engine allows everything."""
|
||||
engine = PolicyEngine(PolicyConfig(enabled=False))
|
||||
engine.add_policy(Policy(
|
||||
name="deny_all",
|
||||
rules=[PolicyRule(action=PolicyAction.DENY, resource="*")]
|
||||
))
|
||||
|
||||
result = engine.check("tool:anything", {})
|
||||
|
||||
assert result.allowed
|
||||
|
||||
def test_engine_check_tool(self, engine):
|
||||
"""Test check_tool helper."""
|
||||
engine.add_policy(Policy(
|
||||
name="deny_delete",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_file"
|
||||
)
|
||||
]
|
||||
))
|
||||
|
||||
result = engine.check_tool("delete_file")
|
||||
|
||||
assert not result.allowed
|
||||
|
||||
def test_engine_check_file(self, engine):
|
||||
"""Test check_file helper."""
|
||||
engine.add_policy(Policy(
|
||||
name="deny_write",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="file:write"
|
||||
)
|
||||
]
|
||||
))
|
||||
|
||||
result = engine.check_file("write", "/path/to/file")
|
||||
|
||||
assert not result.allowed
|
||||
|
||||
def test_engine_enable_disable_policy(self, engine):
|
||||
"""Test enabling/disabling policies."""
|
||||
engine.add_policy(Policy(name="test"))
|
||||
|
||||
engine.disable_policy("test")
|
||||
assert not engine.get_policy("test").enabled
|
||||
|
||||
engine.enable_policy("test")
|
||||
assert engine.get_policy("test").enabled
|
||||
|
||||
def test_engine_clear(self, engine):
|
||||
"""Test clearing all policies."""
|
||||
engine.add_policy(Policy(name="test1"))
|
||||
engine.add_policy(Policy(name="test2"))
|
||||
|
||||
engine.clear()
|
||||
|
||||
assert len(engine.policies) == 0
|
||||
|
||||
def test_engine_to_dict(self, engine):
|
||||
"""Test engine serialization."""
|
||||
engine.add_policy(Policy(name="test"))
|
||||
|
||||
data = engine.to_dict()
|
||||
|
||||
assert "config" in data
|
||||
assert "policies" in data
|
||||
assert len(data["policies"]) == 1
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Convenience Function Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestConvenienceFunctions:
|
||||
"""Tests for convenience policy creation functions."""
|
||||
|
||||
def test_create_deny_tools_policy(self):
|
||||
"""Test creating deny tools policy."""
|
||||
policy = create_deny_tools_policy(
|
||||
["delete_*", "write_*"],
|
||||
reason="Not allowed"
|
||||
)
|
||||
|
||||
assert policy.name == "deny_tools"
|
||||
assert len(policy.rules) == 2
|
||||
|
||||
def test_create_allow_tools_policy(self):
|
||||
"""Test creating allow tools policy."""
|
||||
policy = create_allow_tools_policy(["read_*", "list_*"])
|
||||
|
||||
assert policy.name == "allow_tools"
|
||||
assert len(policy.rules) == 2
|
||||
|
||||
def test_create_read_only_policy(self):
|
||||
"""Test creating read-only policy."""
|
||||
policy = create_read_only_policy()
|
||||
|
||||
assert policy.name == "read_only"
|
||||
assert len(policy.rules) >= 2
|
||||
assert policy.priority == 100
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for thinking module."""
|
||||
@@ -0,0 +1,355 @@
|
||||
"""
|
||||
Unit tests for the Thinking module.
|
||||
|
||||
Tests cover:
|
||||
- ThinkingBudget creation and levels
|
||||
- ThinkingConfig settings
|
||||
- ThinkingUsage tracking
|
||||
- ThinkingTracker aggregation
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.thinking.budget import ThinkingBudget, BudgetLevel, BUDGET_TOKENS
|
||||
from praisonaiagents.thinking.config import ThinkingConfig
|
||||
from praisonaiagents.thinking.tracker import ThinkingUsage, ThinkingTracker
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ThinkingBudget Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestThinkingBudget:
|
||||
"""Tests for ThinkingBudget class."""
|
||||
|
||||
def test_budget_creation(self):
|
||||
"""Test creating a budget."""
|
||||
budget = ThinkingBudget(max_tokens=10000)
|
||||
assert budget.max_tokens == 10000
|
||||
assert budget.adaptive
|
||||
|
||||
def test_budget_from_level(self):
|
||||
"""Test creating budget from level."""
|
||||
budget = ThinkingBudget.from_level(BudgetLevel.HIGH)
|
||||
assert budget.max_tokens == BUDGET_TOKENS[BudgetLevel.HIGH]
|
||||
assert budget.level == BudgetLevel.HIGH
|
||||
|
||||
def test_budget_minimal(self):
|
||||
"""Test minimal budget."""
|
||||
budget = ThinkingBudget.minimal()
|
||||
assert budget.max_tokens == 2000
|
||||
|
||||
def test_budget_low(self):
|
||||
"""Test low budget."""
|
||||
budget = ThinkingBudget.low()
|
||||
assert budget.max_tokens == 4000
|
||||
|
||||
def test_budget_medium(self):
|
||||
"""Test medium budget."""
|
||||
budget = ThinkingBudget.medium()
|
||||
assert budget.max_tokens == 8000
|
||||
|
||||
def test_budget_high(self):
|
||||
"""Test high budget."""
|
||||
budget = ThinkingBudget.high()
|
||||
assert budget.max_tokens == 16000
|
||||
|
||||
def test_budget_maximum(self):
|
||||
"""Test maximum budget."""
|
||||
budget = ThinkingBudget.maximum()
|
||||
assert budget.max_tokens == 32000
|
||||
|
||||
def test_budget_complexity_scaling(self):
|
||||
"""Test complexity-based token scaling."""
|
||||
budget = ThinkingBudget(
|
||||
max_tokens=10000,
|
||||
min_tokens=2000,
|
||||
adaptive=True
|
||||
)
|
||||
|
||||
# Low complexity
|
||||
tokens_low = budget.get_tokens_for_complexity(0.0)
|
||||
assert tokens_low == 2000
|
||||
|
||||
# High complexity
|
||||
tokens_high = budget.get_tokens_for_complexity(1.0)
|
||||
assert tokens_high == 10000
|
||||
|
||||
# Medium complexity
|
||||
tokens_mid = budget.get_tokens_for_complexity(0.5)
|
||||
assert 2000 < tokens_mid < 10000
|
||||
|
||||
def test_budget_non_adaptive(self):
|
||||
"""Test non-adaptive budget."""
|
||||
budget = ThinkingBudget(
|
||||
max_tokens=10000,
|
||||
adaptive=False
|
||||
)
|
||||
|
||||
# Should always return max_tokens
|
||||
assert budget.get_tokens_for_complexity(0.0) == 10000
|
||||
assert budget.get_tokens_for_complexity(1.0) == 10000
|
||||
|
||||
def test_budget_to_dict(self):
|
||||
"""Test budget serialization."""
|
||||
budget = ThinkingBudget(
|
||||
max_tokens=8000,
|
||||
max_time_seconds=60.0,
|
||||
adaptive=True
|
||||
)
|
||||
data = budget.to_dict()
|
||||
|
||||
assert data["max_tokens"] == 8000
|
||||
assert data["max_time_seconds"] == 60.0
|
||||
assert data["adaptive"]
|
||||
|
||||
def test_budget_from_dict(self):
|
||||
"""Test budget deserialization."""
|
||||
data = {
|
||||
"max_tokens": 16000,
|
||||
"max_time_seconds": 120.0,
|
||||
"adaptive": False,
|
||||
"level": "high"
|
||||
}
|
||||
|
||||
budget = ThinkingBudget.from_dict(data)
|
||||
|
||||
assert budget.max_tokens == 16000
|
||||
assert budget.max_time_seconds == 120.0
|
||||
assert not budget.adaptive
|
||||
assert budget.level == BudgetLevel.HIGH
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ThinkingConfig Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestThinkingConfig:
|
||||
"""Tests for ThinkingConfig class."""
|
||||
|
||||
def test_config_defaults(self):
|
||||
"""Test default configuration."""
|
||||
config = ThinkingConfig()
|
||||
|
||||
assert config.enabled
|
||||
assert config.default_budget_tokens == 8000
|
||||
assert config.max_budget_tokens == 32000
|
||||
assert config.adaptive
|
||||
|
||||
def test_config_custom(self):
|
||||
"""Test custom configuration."""
|
||||
config = ThinkingConfig(
|
||||
enabled=False,
|
||||
default_budget_tokens=4000,
|
||||
adaptive=False
|
||||
)
|
||||
|
||||
assert not config.enabled
|
||||
assert config.default_budget_tokens == 4000
|
||||
assert not config.adaptive
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ThinkingUsage Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestThinkingUsage:
|
||||
"""Tests for ThinkingUsage class."""
|
||||
|
||||
def test_usage_creation(self):
|
||||
"""Test creating usage."""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=4000
|
||||
)
|
||||
|
||||
assert usage.budget_tokens == 8000
|
||||
assert usage.tokens_used == 4000
|
||||
|
||||
def test_usage_tokens_remaining(self):
|
||||
"""Test tokens remaining calculation."""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=3000
|
||||
)
|
||||
|
||||
assert usage.tokens_remaining == 5000
|
||||
|
||||
def test_usage_tokens_remaining_over_budget(self):
|
||||
"""Test tokens remaining when over budget."""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=10000
|
||||
)
|
||||
|
||||
assert usage.tokens_remaining == 0
|
||||
|
||||
def test_usage_time_remaining(self):
|
||||
"""Test time remaining calculation."""
|
||||
usage = ThinkingUsage(
|
||||
budget_time=60.0,
|
||||
time_seconds=20.0
|
||||
)
|
||||
|
||||
assert usage.time_remaining == 40.0
|
||||
|
||||
def test_usage_time_remaining_no_budget(self):
|
||||
"""Test time remaining with no budget."""
|
||||
usage = ThinkingUsage()
|
||||
|
||||
assert usage.time_remaining is None
|
||||
|
||||
def test_usage_token_utilization(self):
|
||||
"""Test token utilization calculation."""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=4000
|
||||
)
|
||||
|
||||
assert usage.token_utilization == 0.5
|
||||
|
||||
def test_usage_is_over_budget(self):
|
||||
"""Test over budget detection."""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=10000
|
||||
)
|
||||
|
||||
assert usage.is_over_budget
|
||||
|
||||
def test_usage_is_not_over_budget(self):
|
||||
"""Test not over budget."""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=4000
|
||||
)
|
||||
|
||||
assert not usage.is_over_budget
|
||||
|
||||
def test_usage_is_over_time(self):
|
||||
"""Test over time detection."""
|
||||
usage = ThinkingUsage(
|
||||
budget_time=60.0,
|
||||
time_seconds=90.0
|
||||
)
|
||||
|
||||
assert usage.is_over_time
|
||||
|
||||
def test_usage_to_dict(self):
|
||||
"""Test usage serialization."""
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=4000,
|
||||
time_seconds=30.0
|
||||
)
|
||||
data = usage.to_dict()
|
||||
|
||||
assert data["budget_tokens"] == 8000
|
||||
assert data["tokens_used"] == 4000
|
||||
assert data["token_utilization"] == 0.5
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ThinkingTracker Tests
|
||||
# =============================================================================
|
||||
|
||||
class TestThinkingTracker:
|
||||
"""Tests for ThinkingTracker class."""
|
||||
|
||||
@pytest.fixture
|
||||
def tracker(self):
|
||||
"""Create a test tracker."""
|
||||
return ThinkingTracker()
|
||||
|
||||
def test_tracker_creation(self, tracker):
|
||||
"""Test creating a tracker."""
|
||||
assert tracker.session_count == 0
|
||||
assert tracker.total_tokens_used == 0
|
||||
|
||||
def test_tracker_start_session(self, tracker):
|
||||
"""Test starting a session."""
|
||||
usage = tracker.start_session(
|
||||
budget_tokens=8000,
|
||||
complexity=0.7
|
||||
)
|
||||
|
||||
assert usage.budget_tokens == 8000
|
||||
assert usage.complexity == 0.7
|
||||
assert tracker.session_count == 1
|
||||
|
||||
def test_tracker_end_session(self, tracker):
|
||||
"""Test ending a session."""
|
||||
usage = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage, tokens_used=5000, time_seconds=30.0)
|
||||
|
||||
assert usage.tokens_used == 5000
|
||||
assert usage.time_seconds == 30.0
|
||||
assert tracker.total_tokens_used == 5000
|
||||
assert tracker.total_time_seconds == 30.0
|
||||
|
||||
def test_tracker_multiple_sessions(self, tracker):
|
||||
"""Test multiple sessions."""
|
||||
usage1 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage1, tokens_used=4000, time_seconds=20.0)
|
||||
|
||||
usage2 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage2, tokens_used=6000, time_seconds=40.0)
|
||||
|
||||
assert tracker.session_count == 2
|
||||
assert tracker.total_tokens_used == 10000
|
||||
assert tracker.total_time_seconds == 60.0
|
||||
|
||||
def test_tracker_average_tokens(self, tracker):
|
||||
"""Test average tokens calculation."""
|
||||
usage1 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage1, tokens_used=4000, time_seconds=20.0)
|
||||
|
||||
usage2 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage2, tokens_used=6000, time_seconds=40.0)
|
||||
|
||||
assert tracker.average_tokens_per_session == 5000
|
||||
|
||||
def test_tracker_average_time(self, tracker):
|
||||
"""Test average time calculation."""
|
||||
usage1 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage1, tokens_used=4000, time_seconds=20.0)
|
||||
|
||||
usage2 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage2, tokens_used=6000, time_seconds=40.0)
|
||||
|
||||
assert tracker.average_time_per_session == 30.0
|
||||
|
||||
def test_tracker_over_budget_count(self, tracker):
|
||||
"""Test over budget counting."""
|
||||
usage1 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage1, tokens_used=4000, time_seconds=20.0)
|
||||
|
||||
usage2 = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage2, tokens_used=10000, time_seconds=40.0)
|
||||
|
||||
assert tracker.over_budget_count == 1
|
||||
|
||||
def test_tracker_get_summary(self, tracker):
|
||||
"""Test summary generation."""
|
||||
usage = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage, tokens_used=4000, time_seconds=20.0)
|
||||
|
||||
summary = tracker.get_summary()
|
||||
|
||||
assert summary["session_count"] == 1
|
||||
assert summary["total_tokens_used"] == 4000
|
||||
assert "average_utilization" in summary
|
||||
|
||||
def test_tracker_clear(self, tracker):
|
||||
"""Test clearing tracker."""
|
||||
usage = tracker.start_session(budget_tokens=8000)
|
||||
tracker.end_session(usage, tokens_used=4000, time_seconds=20.0)
|
||||
|
||||
tracker.clear()
|
||||
|
||||
assert tracker.session_count == 0
|
||||
assert tracker.total_tokens_used == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Generated
+1
-1
@@ -2422,7 +2422,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.1.23"
|
||||
version = "0.2.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
+472
-126
@@ -261,6 +261,12 @@ npm install praisonai
|
||||
| ↳ 100+ LLM Support | [Example](examples/python/providers/openai/openai_gpt4_example.py) | [📖](https://docs.praison.ai/models) |
|
||||
| ↳ Callback Agents | [Example](examples/python/general/advanced-callback-systems.py) | [📖](https://docs.praison.ai/features/callbacks) |
|
||||
| ↳ Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
|
||||
| ↳ Shadow Git Checkpoints | [Example](#10-shadow-git-checkpoints) | [📖](https://docs.praison.ai/features/checkpoints) |
|
||||
| ↳ Background Tasks | [Example](examples/background/basic_background.py) | [📖](https://docs.praison.ai/features/background-tasks) |
|
||||
| ↳ Policy Engine | [Example](examples/policy/basic_policy.py) | [📖](https://docs.praison.ai/features/policy-engine) |
|
||||
| ↳ Thinking Budgets | [Example](examples/thinking/basic_thinking.py) | [📖](https://docs.praison.ai/features/thinking-budgets) |
|
||||
| ↳ Output Styles | [Example](examples/output/basic_output.py) | [📖](https://docs.praison.ai/features/output-styles) |
|
||||
| ↳ Context Compaction | [Example](examples/compaction/basic_compaction.py) | [📖](https://docs.praison.ai/features/context-compaction) |
|
||||
| **📊 Monitoring & Management** | | |
|
||||
| ↳ Sessions Management | [Example](examples/python/sessions/comprehensive-session-management.py) | [📖](https://docs.praison.ai/features/sessions) |
|
||||
| ↳ Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
@@ -278,6 +284,12 @@ npm install praisonai
|
||||
| ↳ Git Integration | [Example](examples/python/cli/git_integration_example.py) | [📖](https://docs.praison.ai/docs/cli/git-integration) |
|
||||
| ↳ Sandbox Execution | [Example](examples/python/cli/sandbox_execution_example.py) | [📖](https://docs.praison.ai/docs/cli/sandbox-execution) |
|
||||
| ↳ CLI Compare | [Example](examples/compare/cli_compare_basic.py) | [📖](https://docs.praison.ai/docs/cli/compare) |
|
||||
| ↳ Profile/Benchmark | [Example](#profile-benchmark) | [📖](https://docs.praison.ai/docs/cli/profile) |
|
||||
| ↳ Auto Mode | [Example](#auto-mode) | [📖](https://docs.praison.ai/docs/cli/auto) |
|
||||
| ↳ Init | [Example](#init) | [📖](https://docs.praison.ai/docs/cli/init) |
|
||||
| ↳ File Input | [Example](#file-input) | [📖](https://docs.praison.ai/docs/cli/file-input) |
|
||||
| ↳ Final Agent | [Example](#final-agent) | [📖](https://docs.praison.ai/docs/cli/final-agent) |
|
||||
| ↳ Max Tokens | [Example](#max-tokens) | [📖](https://docs.praison.ai/docs/cli/max-tokens) |
|
||||
| **🧪 Evaluation** | | |
|
||||
| ↳ Accuracy Evaluation | [Example](examples/eval/accuracy_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
| ↳ Performance Evaluation | [Example](examples/eval/performance_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
@@ -857,21 +869,258 @@ callbacks:
|
||||
|
||||
### 9. Hooks
|
||||
|
||||
Configure in `.praison/hooks.json`:
|
||||
Intercept and modify agent behavior at various lifecycle points:
|
||||
|
||||
```python
|
||||
from praisonaiagents.memory import HooksManager
|
||||
from praisonaiagents.hooks import (
|
||||
HookRegistry, HookRunner, HookEvent, HookResult,
|
||||
BeforeToolInput
|
||||
)
|
||||
|
||||
hooks = HooksManager()
|
||||
# Create a hook registry
|
||||
registry = HookRegistry()
|
||||
|
||||
# Register Python hooks
|
||||
hooks.register("pre_write_code", lambda ctx: print(f"Writing {ctx['file']}"))
|
||||
# Log all tool calls
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def log_tools(event_data: BeforeToolInput) -> HookResult:
|
||||
print(f"Tool: {event_data.tool_name}")
|
||||
return HookResult.allow()
|
||||
|
||||
# Block dangerous operations
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def security_check(event_data: BeforeToolInput) -> HookResult:
|
||||
if "delete" in event_data.tool_name.lower():
|
||||
return HookResult.deny("Delete operations blocked")
|
||||
return HookResult.allow()
|
||||
|
||||
# Execute hooks
|
||||
result = hooks.execute("pre_write_code", {"file": "main.py"})
|
||||
runner = HookRunner(registry)
|
||||
```
|
||||
|
||||
### 10. Field Names Reference (A-I-G-S)
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai hooks list # List registered hooks
|
||||
praisonai hooks test before_tool # Test hooks for an event
|
||||
praisonai hooks run "echo test" # Run a command hook
|
||||
praisonai hooks validate hooks.json # Validate configuration
|
||||
```
|
||||
|
||||
|
||||
### 10. Shadow Git Checkpoints
|
||||
|
||||
File-level undo/restore using shadow git:
|
||||
|
||||
```python
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
|
||||
service = CheckpointService(workspace_dir="./my_project")
|
||||
await service.initialize()
|
||||
|
||||
# Save checkpoint before changes
|
||||
result = await service.save("Before refactoring")
|
||||
|
||||
# Make changes...
|
||||
|
||||
# Restore if needed
|
||||
await service.restore(result.checkpoint.id)
|
||||
|
||||
# View diff
|
||||
diff = await service.diff()
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai checkpoint save "Before changes" # Save checkpoint
|
||||
praisonai checkpoint list # List checkpoints
|
||||
praisonai checkpoint diff # Show changes
|
||||
praisonai checkpoint restore abc123 # Restore to checkpoint
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/checkpoints)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/checkpoint)
|
||||
- [💻 Example](examples/checkpoints/basic_checkpoints.py)
|
||||
|
||||
---
|
||||
|
||||
### 11. Background Tasks
|
||||
|
||||
Run agent tasks asynchronously without blocking:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from praisonaiagents.background import BackgroundRunner, BackgroundConfig
|
||||
|
||||
async def main():
|
||||
config = BackgroundConfig(max_concurrent_tasks=3)
|
||||
runner = BackgroundRunner(config=config)
|
||||
|
||||
async def my_task(name: str) -> str:
|
||||
await asyncio.sleep(2)
|
||||
return f"Task {name} completed"
|
||||
|
||||
task = await runner.submit(my_task, args=("example",), name="my_task")
|
||||
await task.wait(timeout=10.0)
|
||||
print(task.result)
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai background list # List running tasks
|
||||
praisonai background status <id> # Check task status
|
||||
praisonai background cancel <id> # Cancel a task
|
||||
praisonai background clear # Clear completed tasks
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/background-tasks)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/background)
|
||||
- [💻 Example](examples/background/basic_background.py)
|
||||
|
||||
---
|
||||
|
||||
### 12. Policy Engine
|
||||
|
||||
Control what agents can and cannot do with policy-based execution:
|
||||
|
||||
```python
|
||||
from praisonaiagents.policy import (
|
||||
PolicyEngine, Policy, PolicyRule, PolicyAction
|
||||
)
|
||||
|
||||
engine = PolicyEngine()
|
||||
|
||||
policy = Policy(
|
||||
name="no_delete",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="Delete operations blocked"
|
||||
)
|
||||
]
|
||||
)
|
||||
engine.add_policy(policy)
|
||||
|
||||
result = engine.check("tool:delete_file", {})
|
||||
print(f"Allowed: {result.allowed}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai policy list # List policies
|
||||
praisonai policy check "tool:name" # Check if allowed
|
||||
praisonai policy init # Create template
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/policy-engine)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/policy)
|
||||
- [💻 Example](examples/policy/basic_policy.py)
|
||||
|
||||
---
|
||||
|
||||
### 13. Thinking Budgets
|
||||
|
||||
Configure token budgets for extended thinking:
|
||||
|
||||
```python
|
||||
from praisonaiagents.thinking import ThinkingBudget, ThinkingTracker
|
||||
|
||||
# Use predefined levels
|
||||
budget = ThinkingBudget.high() # 16,000 tokens
|
||||
|
||||
# Track usage
|
||||
tracker = ThinkingTracker()
|
||||
session = tracker.start_session(budget_tokens=16000)
|
||||
tracker.end_session(session, tokens_used=12000)
|
||||
|
||||
summary = tracker.get_summary()
|
||||
print(f"Utilization: {summary['average_utilization']:.1%}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai thinking status # Show current budget
|
||||
praisonai thinking set high # Set budget level
|
||||
praisonai thinking stats # Show usage statistics
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/thinking-budgets)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/thinking)
|
||||
- [💻 Example](examples/thinking/basic_thinking.py)
|
||||
|
||||
---
|
||||
|
||||
### 14. Output Styles
|
||||
|
||||
Configure how agents format their responses:
|
||||
|
||||
```python
|
||||
from praisonaiagents.output import OutputStyle, OutputFormatter
|
||||
|
||||
# Use preset styles
|
||||
style = OutputStyle.concise()
|
||||
formatter = OutputFormatter(style)
|
||||
|
||||
# Format output
|
||||
text = "# Hello\n\nThis is **bold** text."
|
||||
plain = formatter.format(text)
|
||||
print(plain)
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai output status # Show current style
|
||||
praisonai output set concise # Set output style
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/output-styles)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/output-style)
|
||||
- [💻 Example](examples/output/basic_output.py)
|
||||
|
||||
---
|
||||
|
||||
### 15. Context Compaction
|
||||
|
||||
Automatically manage context window size:
|
||||
|
||||
```python
|
||||
from praisonaiagents.compaction import (
|
||||
ContextCompactor, CompactionStrategy
|
||||
)
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=4000,
|
||||
strategy=CompactionStrategy.SLIDING,
|
||||
preserve_recent=3
|
||||
)
|
||||
|
||||
messages = [...] # Your conversation history
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f"Compression: {result.compression_ratio:.1%}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai compaction status # Show settings
|
||||
praisonai compaction set sliding # Set strategy
|
||||
praisonai compaction stats # Show statistics
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/context-compaction)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/compaction)
|
||||
- [💻 Example](examples/compaction/basic_compaction.py)
|
||||
|
||||
---
|
||||
|
||||
### 16. Field Names Reference (A-I-G-S)
|
||||
|
||||
PraisonAI accepts both old (agents.yaml) and new (workflow.yaml) field names. Use the **canonical names** for new projects:
|
||||
|
||||
@@ -905,7 +1154,7 @@ steps: # Define steps (not 'tasks')
|
||||
|
||||
> **Note:** The parser accepts both old and new names. Run `praisonai workflow validate <file.yaml>` to see suggestions for canonical names.
|
||||
|
||||
### 11. Extended agents.yaml with Workflow Patterns
|
||||
### 12. Extended agents.yaml with Workflow Patterns
|
||||
|
||||
**Feature Parity:** Both `agents.yaml` and `workflow.yaml` now support the same features:
|
||||
- All workflow patterns (route, parallel, loop, repeat)
|
||||
@@ -981,6 +1230,124 @@ Run with the same simple command:
|
||||
praisonai agents.yaml
|
||||
```
|
||||
|
||||
### 12. MCP (Model Context Protocol)
|
||||
|
||||
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
|
||||
|
||||
#### MCP Client (Consume MCP Servers)
|
||||
```python
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
# stdio - Local NPX/Python servers
|
||||
agent = Agent(tools=MCP("npx @modelcontextprotocol/server-memory"))
|
||||
|
||||
# Streamable HTTP - Production servers
|
||||
agent = Agent(tools=MCP("https://api.example.com/mcp"))
|
||||
|
||||
# WebSocket - Real-time bidirectional
|
||||
agent = Agent(tools=MCP("wss://api.example.com/mcp", auth_token="token"))
|
||||
|
||||
# SSE (Legacy) - Backward compatibility
|
||||
agent = Agent(tools=MCP("http://localhost:8080/sse"))
|
||||
|
||||
# With environment variables
|
||||
agent = Agent(
|
||||
tools=MCP(
|
||||
command="npx",
|
||||
args=["-y", "@modelcontextprotocol/server-brave-search"],
|
||||
env={"BRAVE_API_KEY": "your-key"}
|
||||
)
|
||||
)
|
||||
|
||||
# Multiple MCP servers + regular functions
|
||||
def my_custom_tool(query: str) -> str:
|
||||
"""Custom tool function."""
|
||||
return f"Result: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="MultiToolAgent",
|
||||
instructions="Agent with multiple MCP servers",
|
||||
tools=[
|
||||
MCP("uvx mcp-server-time"), # Time tools
|
||||
MCP("npx @modelcontextprotocol/server-memory"), # Memory tools
|
||||
my_custom_tool # Regular function
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
#### MCP Server (Expose Tools as MCP Server)
|
||||
|
||||
Expose your Python functions as MCP tools for Claude Desktop, Cursor, and other MCP clients:
|
||||
|
||||
```python
|
||||
from praisonaiagents.mcp import ToolsMCPServer
|
||||
|
||||
def search_web(query: str, max_results: int = 5) -> dict:
|
||||
"""Search the web for information."""
|
||||
return {"results": [f"Result for {query}"]}
|
||||
|
||||
def calculate(expression: str) -> dict:
|
||||
"""Evaluate a mathematical expression."""
|
||||
return {"result": eval(expression)}
|
||||
|
||||
# Create and run MCP server
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
server.register_tools([search_web, calculate])
|
||||
server.run() # stdio for Claude Desktop
|
||||
# server.run_sse(host="0.0.0.0", port=8080) # SSE for web clients
|
||||
```
|
||||
|
||||
#### MCP Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Session Management | Automatic Mcp-Session-Id handling |
|
||||
| Protocol Versioning | Mcp-Protocol-Version header |
|
||||
| Resumability | SSE stream recovery via Last-Event-ID |
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
### 13. A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
#### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
#### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
---
|
||||
|
||||
## 🎯 CLI / No-Code Interface
|
||||
@@ -1730,124 +2097,6 @@ agent = Agent(
|
||||
| Web Fetch | Anthropic |
|
||||
| Prompt Caching | OpenAI (auto), Anthropic, Bedrock, Deepseek |
|
||||
|
||||
## MCP (Model Context Protocol)
|
||||
|
||||
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
|
||||
|
||||
### MCP Client (Consume MCP Servers)
|
||||
```python
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
# stdio - Local NPX/Python servers
|
||||
agent = Agent(tools=MCP("npx @modelcontextprotocol/server-memory"))
|
||||
|
||||
# Streamable HTTP - Production servers
|
||||
agent = Agent(tools=MCP("https://api.example.com/mcp"))
|
||||
|
||||
# WebSocket - Real-time bidirectional
|
||||
agent = Agent(tools=MCP("wss://api.example.com/mcp", auth_token="token"))
|
||||
|
||||
# SSE (Legacy) - Backward compatibility
|
||||
agent = Agent(tools=MCP("http://localhost:8080/sse"))
|
||||
|
||||
# With environment variables
|
||||
agent = Agent(
|
||||
tools=MCP(
|
||||
command="npx",
|
||||
args=["-y", "@modelcontextprotocol/server-brave-search"],
|
||||
env={"BRAVE_API_KEY": "your-key"}
|
||||
)
|
||||
)
|
||||
|
||||
# Multiple MCP servers + regular functions
|
||||
def my_custom_tool(query: str) -> str:
|
||||
"""Custom tool function."""
|
||||
return f"Result: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="MultiToolAgent",
|
||||
instructions="Agent with multiple MCP servers",
|
||||
tools=[
|
||||
MCP("uvx mcp-server-time"), # Time tools
|
||||
MCP("npx @modelcontextprotocol/server-memory"), # Memory tools
|
||||
my_custom_tool # Regular function
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### MCP Server (Expose Tools as MCP Server)
|
||||
|
||||
Expose your Python functions as MCP tools for Claude Desktop, Cursor, and other MCP clients:
|
||||
|
||||
```python
|
||||
from praisonaiagents.mcp import ToolsMCPServer
|
||||
|
||||
def search_web(query: str, max_results: int = 5) -> dict:
|
||||
"""Search the web for information."""
|
||||
return {"results": [f"Result for {query}"]}
|
||||
|
||||
def calculate(expression: str) -> dict:
|
||||
"""Evaluate a mathematical expression."""
|
||||
return {"result": eval(expression)}
|
||||
|
||||
# Create and run MCP server
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
server.register_tools([search_web, calculate])
|
||||
server.run() # stdio for Claude Desktop
|
||||
# server.run_sse(host="0.0.0.0", port=8080) # SSE for web clients
|
||||
```
|
||||
|
||||
### MCP Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Session Management | Automatic Mcp-Session-Id handling |
|
||||
| Protocol Versioning | Mcp-Protocol-Version header |
|
||||
| Resumability | SSE stream recovery via Last-Event-ID |
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
## A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
## CLI Features
|
||||
|
||||
| Feature | Docs |
|
||||
@@ -2475,6 +2724,103 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
|
||||
---
|
||||
|
||||
## 💾 Persistence (Databases)
|
||||
|
||||
Enable automatic conversation persistence with 2 lines of code:
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
db=db(database_url="postgresql://localhost/mydb"), # db(...) shortcut
|
||||
session_id="my-session" # Optional: defaults to per-hour ID (YYYYMMDDHH)
|
||||
)
|
||||
agent.chat("Hello!") # Auto-persists messages, runs, traces
|
||||
```
|
||||
|
||||
### Persistence CLI Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai persistence doctor` | Validate DB connectivity |
|
||||
| `praisonai persistence run` | Run agent with persistence |
|
||||
| `praisonai persistence resume` | Resume existing session |
|
||||
| `praisonai persistence export` | Export session to JSONL |
|
||||
| `praisonai persistence import` | Import session from JSONL |
|
||||
| `praisonai persistence migrate` | Apply schema migrations |
|
||||
| `praisonai persistence status` | Show schema status |
|
||||
|
||||
### Databases Table
|
||||
|
||||
| Database | Store Type | Install | Example | Docs |
|
||||
|----------|------------|---------|---------|------|
|
||||
| PostgreSQL | Conversation | `pip install "praisonai[tools]"` | [simple_db_agent.py](examples/persistence/simple_db_agent.py) | [docs](https://docs.praison.ai/docs/databases/postgres) |
|
||||
| MySQL | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SQLite | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SingleStore | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Supabase | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SurrealDB | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Qdrant | Knowledge | `pip install "praisonai[tools]"` | [knowledge_qdrant.py](examples/persistence/knowledge_qdrant.py) | [docs](https://docs.praison.ai/docs/databases/qdrant) |
|
||||
| ChromaDB | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Pinecone | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Weaviate | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| LanceDB | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Milvus | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| PGVector | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Redis Vector | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Cassandra | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| ClickHouse | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Redis | State | `pip install "praisonai[tools]"` | [state_redis.py](examples/persistence/state_redis.py) | [docs](https://docs.praison.ai/docs/databases/redis) |
|
||||
| MongoDB | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| DynamoDB | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Firestore | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Upstash | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Memory | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Tools Table
|
||||
|
||||
Install all tools with: `pip install "praisonai[tools]"`
|
||||
|
||||
| Tool | Category | Import | Docs |
|
||||
|------|----------|--------|------|
|
||||
| Tavily | Web Search | `from praisonai_tools import TavilyTool` | [docs](https://docs.praison.ai/docs/tools/external/tavily) |
|
||||
| DuckDuckGo | Web Search | `from praisonai_tools import DuckDuckGoTool` | [docs](https://docs.praison.ai/docs/tools/external/duckduckgo) |
|
||||
| Exa | Web Search | `from praisonai_tools import ExaTool` | [docs](https://docs.praison.ai/docs/tools/external/exa) |
|
||||
| Serper | Web Search | `from praisonai_tools import SerperTool` | [docs](https://docs.praison.ai/docs/tools/external/serper) |
|
||||
| Jina | Web Reader | `from praisonai_tools import JinaTool` | [docs](https://docs.praison.ai/docs/tools/external/jina) |
|
||||
| Firecrawl | Web Scraping | `from praisonai_tools import FirecrawlTool` | [docs](https://docs.praison.ai/docs/tools/external/firecrawl) |
|
||||
| Crawl4AI | Web Scraping | `from praisonai_tools import Crawl4AITool` | [docs](https://docs.praison.ai/docs/tools/external/crawl4ai) |
|
||||
| Wikipedia | Knowledge | `from praisonai_tools import WikipediaTool` | [docs](https://docs.praison.ai/docs/tools/external/wikipedia) |
|
||||
| ArXiv | Research | `from praisonai_tools import ArxivTool` | [docs](https://docs.praison.ai/docs/tools/external/arxiv) |
|
||||
| HackerNews | News | `from praisonai_tools import HackerNewsTool` | [docs](https://docs.praison.ai/docs/tools/external/hackernews) |
|
||||
| YouTube | Media | `from praisonai_tools import YouTubeTool` | [docs](https://docs.praison.ai/docs/tools/external/youtube) |
|
||||
| Weather | Data | `from praisonai_tools import WeatherTool` | [docs](https://docs.praison.ai/docs/tools/external/weather) |
|
||||
| PostgreSQL | Database | `from praisonai_tools import PostgresTool` | [docs](https://docs.praison.ai/docs/tools/external/postgres) |
|
||||
| MySQL | Database | `from praisonai_tools import MySQLTool` | [docs](https://docs.praison.ai/docs/tools/external/mysql) |
|
||||
| SQLite | Database | `from praisonai_tools import SQLiteTool` | [docs](https://docs.praison.ai/docs/tools/external/sqlite) |
|
||||
| MongoDB | Database | `from praisonai_tools import MongoDBTool` | [docs](https://docs.praison.ai/docs/tools/external/mongodb) |
|
||||
| Redis | Database | `from praisonai_tools import RedisTool` | [docs](https://docs.praison.ai/docs/tools/external/redis) |
|
||||
| Qdrant | Vector DB | `from praisonai_tools import QdrantTool` | [docs](https://docs.praison.ai/docs/tools/external/qdrant) |
|
||||
| GitHub | DevOps | `from praisonai_tools import GitHubTool` | [docs](https://docs.praison.ai/docs/tools/external/github) |
|
||||
| Slack | Communication | `from praisonai_tools import SlackTool` | [docs](https://docs.praison.ai/docs/tools/external/slack) |
|
||||
| Discord | Communication | `from praisonai_tools import DiscordTool` | [docs](https://docs.praison.ai/docs/tools/external/discord) |
|
||||
| Telegram | Communication | `from praisonai_tools import TelegramTool` | [docs](https://docs.praison.ai/docs/tools/external/telegram) |
|
||||
| Email | Communication | `from praisonai_tools import EmailTool` | [docs](https://docs.praison.ai/docs/tools/external/email) |
|
||||
| Notion | Productivity | `from praisonai_tools import NotionTool` | [docs](https://docs.praison.ai/docs/tools/external/notion) |
|
||||
| File | File System | `from praisonai_tools import FileTool` | [docs](https://docs.praison.ai/docs/tools/external/file) |
|
||||
| Shell | System | `from praisonai_tools import ShellTool` | [docs](https://docs.praison.ai/docs/tools/external/shell) |
|
||||
| Python | Code | `from praisonai_tools import PythonTool` | [docs](https://docs.praison.ai/docs/tools/external/python) |
|
||||
| JSON | Data | `from praisonai_tools import JSONTool` | [docs](https://docs.praison.ai/docs/tools/external/json) |
|
||||
| CSV | Data | `from praisonai_tools import CSVTool` | [docs](https://docs.praison.ai/docs/tools/external/csv) |
|
||||
| Calculator | Math | `from praisonai_tools import CalculatorTool` | [docs](https://docs.praison.ai/docs/tools/external/calculator) |
|
||||
|
||||
> See [full tools documentation](https://docs.praison.ai/docs/tools/tools) for all 100+ available tools.
|
||||
|
||||
---
|
||||
|
||||
## 🎓 Video Tutorials
|
||||
|
||||
Learn PraisonAI through our comprehensive video series:
|
||||
|
||||
@@ -3,8 +3,8 @@ class Praisonai < Formula
|
||||
|
||||
desc "AI tools for various AI applications"
|
||||
homepage "https://github.com/MervinPraison/PraisonAI"
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.74.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.74.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.4.2.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.4.2.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
+255
-18
@@ -62,7 +62,7 @@ except ImportError:
|
||||
logging.basicConfig(level=os.environ.get('LOGLEVEL', 'INFO').upper(), format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
|
||||
# =============================================================================
|
||||
# Available Tools List (shared between generators)
|
||||
# Available Tools List (shared between generators) - Legacy for praisonai_tools
|
||||
# =============================================================================
|
||||
AVAILABLE_TOOLS = [
|
||||
"CodeDocsSearchTool", "CSVSearchTool", "DirectorySearchTool", "DOCXSearchTool",
|
||||
@@ -72,6 +72,206 @@ AVAILABLE_TOOLS = [
|
||||
"YoutubeChannelSearchTool", "YoutubeVideoSearchTool"
|
||||
]
|
||||
|
||||
# =============================================================================
|
||||
# Enhanced Tool Discovery from praisonaiagents.tools
|
||||
# =============================================================================
|
||||
|
||||
# Tool categories with their tools from praisonaiagents.tools
|
||||
TOOL_CATEGORIES = {
|
||||
'web_search': [
|
||||
'internet_search', 'duckduckgo', 'tavily_search', 'exa_search',
|
||||
'search_web', 'ydc_search', 'searxng_search'
|
||||
],
|
||||
'web_scraping': [
|
||||
'scrape_page', 'extract_links', 'crawl', 'extract_text',
|
||||
'crawl4ai', 'crawl4ai_extract', 'get_article'
|
||||
],
|
||||
'file_operations': [
|
||||
'read_file', 'write_file', 'list_files', 'get_file_info',
|
||||
'copy_file', 'move_file', 'delete_file'
|
||||
],
|
||||
'code_execution': [
|
||||
'execute_command', 'execute_code', 'analyze_code', 'format_code'
|
||||
],
|
||||
'data_processing': [
|
||||
'read_csv', 'write_csv', 'analyze_csv', 'read_json', 'write_json',
|
||||
'read_excel', 'write_excel', 'read_yaml', 'write_yaml', 'read_xml'
|
||||
],
|
||||
'research': [
|
||||
'search_arxiv', 'get_arxiv_paper', 'wiki_search', 'wiki_summary',
|
||||
'get_news_sources', 'get_trending_topics'
|
||||
],
|
||||
'finance': [
|
||||
'get_stock_price', 'get_stock_info', 'get_historical_data'
|
||||
],
|
||||
'math': [
|
||||
'evaluate', 'solve_equation', 'convert_units', 'calculate_statistics'
|
||||
],
|
||||
'database': [
|
||||
'query', 'create_table', 'load_data', 'find_documents', 'vector_search'
|
||||
]
|
||||
}
|
||||
|
||||
# Keywords that map to tool categories
|
||||
TASK_KEYWORD_TO_TOOLS = {
|
||||
# Web search keywords
|
||||
'search': 'web_search',
|
||||
'find': 'web_search',
|
||||
'look up': 'web_search',
|
||||
'google': 'web_search',
|
||||
'internet': 'web_search',
|
||||
'online': 'web_search',
|
||||
'web': 'web_search',
|
||||
|
||||
# Web scraping keywords
|
||||
'scrape': 'web_scraping',
|
||||
'crawl': 'web_scraping',
|
||||
'extract from website': 'web_scraping',
|
||||
'get from url': 'web_scraping',
|
||||
'fetch page': 'web_scraping',
|
||||
|
||||
# File operation keywords
|
||||
'read file': 'file_operations',
|
||||
'write file': 'file_operations',
|
||||
'save': 'file_operations',
|
||||
'load': 'file_operations',
|
||||
'open file': 'file_operations',
|
||||
'create file': 'file_operations',
|
||||
|
||||
# Code execution keywords
|
||||
'execute': 'code_execution',
|
||||
'run code': 'code_execution',
|
||||
'python': 'code_execution',
|
||||
'script': 'code_execution',
|
||||
'command': 'code_execution',
|
||||
'shell': 'code_execution',
|
||||
|
||||
# Data processing keywords
|
||||
'csv': 'data_processing',
|
||||
'excel': 'data_processing',
|
||||
'json': 'data_processing',
|
||||
'yaml': 'data_processing',
|
||||
'xml': 'data_processing',
|
||||
'data': 'data_processing',
|
||||
'spreadsheet': 'data_processing',
|
||||
|
||||
# Research keywords
|
||||
'research': 'research',
|
||||
'paper': 'research',
|
||||
'arxiv': 'research',
|
||||
'wikipedia': 'research',
|
||||
'academic': 'research',
|
||||
'news': 'research',
|
||||
|
||||
# Finance keywords
|
||||
'stock': 'finance',
|
||||
'price': 'finance',
|
||||
'market': 'finance',
|
||||
'financial': 'finance',
|
||||
'trading': 'finance',
|
||||
|
||||
# Math keywords
|
||||
'calculate': 'math',
|
||||
'math': 'math',
|
||||
'equation': 'math',
|
||||
'compute': 'math',
|
||||
'statistics': 'math',
|
||||
|
||||
# Database keywords
|
||||
'database': 'database',
|
||||
'sql': 'database',
|
||||
'query': 'database',
|
||||
'mongodb': 'database',
|
||||
'vector': 'database'
|
||||
}
|
||||
|
||||
|
||||
def get_all_available_tools() -> Dict[str, List[str]]:
|
||||
"""
|
||||
Get all available tools organized by category.
|
||||
|
||||
Returns:
|
||||
Dict mapping category names to lists of tool names
|
||||
"""
|
||||
return TOOL_CATEGORIES.copy()
|
||||
|
||||
|
||||
def get_tools_for_task(task_description: str) -> List[str]:
|
||||
"""
|
||||
Analyze a task description and return appropriate tools.
|
||||
|
||||
Args:
|
||||
task_description: The task to analyze
|
||||
|
||||
Returns:
|
||||
List of tool names appropriate for the task
|
||||
"""
|
||||
task_lower = task_description.lower()
|
||||
matched_categories = set()
|
||||
|
||||
# Match keywords to categories
|
||||
for keyword, category in TASK_KEYWORD_TO_TOOLS.items():
|
||||
if keyword in task_lower:
|
||||
matched_categories.add(category)
|
||||
|
||||
# Collect tools from matched categories
|
||||
tools = []
|
||||
for category in matched_categories:
|
||||
if category in TOOL_CATEGORIES:
|
||||
tools.extend(TOOL_CATEGORIES[category])
|
||||
|
||||
# Always include core tools for flexibility
|
||||
core_tools = ['read_file', 'write_file', 'execute_command']
|
||||
for tool in core_tools:
|
||||
if tool not in tools:
|
||||
tools.append(tool)
|
||||
|
||||
# Remove duplicates while preserving order
|
||||
seen = set()
|
||||
unique_tools = []
|
||||
for tool in tools:
|
||||
if tool not in seen:
|
||||
seen.add(tool)
|
||||
unique_tools.append(tool)
|
||||
|
||||
return unique_tools
|
||||
|
||||
|
||||
def recommend_agent_count(task_description: str) -> int:
|
||||
"""
|
||||
Recommend the optimal number of agents based on task complexity.
|
||||
|
||||
Args:
|
||||
task_description: The task to analyze
|
||||
|
||||
Returns:
|
||||
Recommended number of agents (1-4)
|
||||
"""
|
||||
complexity = BaseAutoGenerator.analyze_complexity(task_description)
|
||||
|
||||
if complexity == 'simple':
|
||||
return 1
|
||||
elif complexity == 'moderate':
|
||||
return 2
|
||||
else: # complex
|
||||
# Count distinct aspects of the task
|
||||
task_lower = task_description.lower()
|
||||
aspects = 0
|
||||
|
||||
aspect_keywords = [
|
||||
['research', 'search', 'find', 'gather'],
|
||||
['analyze', 'evaluate', 'assess', 'review'],
|
||||
['write', 'create', 'generate', 'produce'],
|
||||
['edit', 'refine', 'improve', 'polish'],
|
||||
['coordinate', 'manage', 'orchestrate', 'delegate']
|
||||
]
|
||||
|
||||
for keyword_group in aspect_keywords:
|
||||
if any(kw in task_lower for kw in keyword_group):
|
||||
aspects += 1
|
||||
|
||||
return min(max(aspects, 2), 4) # Between 2 and 4 agents
|
||||
|
||||
# =============================================================================
|
||||
# Base Generator Class (DRY - shared functionality)
|
||||
# =============================================================================
|
||||
@@ -458,9 +658,19 @@ Tools are not available for {framework}. To use tools, install:
|
||||
|
||||
return merged_data
|
||||
|
||||
def discover_tools_for_topic(self) -> List[str]:
|
||||
"""
|
||||
Discover appropriate tools for the topic using intelligent matching.
|
||||
|
||||
Returns:
|
||||
List of tool names appropriate for this topic
|
||||
"""
|
||||
return get_tools_for_task(self.topic)
|
||||
|
||||
def get_user_content(self):
|
||||
"""
|
||||
Generates a prompt for the OpenAI API to generate a team structure.
|
||||
Uses intelligent tool discovery based on task analysis.
|
||||
|
||||
Args:
|
||||
None
|
||||
@@ -484,42 +694,68 @@ Tools are not available for {framework}. To use tools, install:
|
||||
|
||||
workflow_guidance = pattern_guidance.get(self.pattern, pattern_guidance["sequential"])
|
||||
|
||||
# Get recommended tools based on task analysis
|
||||
recommended_tools = self.discover_tools_for_topic()
|
||||
recommended_agent_count = recommend_agent_count(self.topic)
|
||||
complexity = self.analyze_complexity(self.topic)
|
||||
|
||||
# Build comprehensive tool list with categories
|
||||
all_tools_by_category = []
|
||||
for category, tools in TOOL_CATEGORIES.items():
|
||||
all_tools_by_category.append(f" {category}: {', '.join(tools)}")
|
||||
tools_reference = "\n".join(all_tools_by_category)
|
||||
|
||||
# Also include legacy tools for backward compatibility
|
||||
legacy_tools = ", ".join(AVAILABLE_TOOLS)
|
||||
|
||||
user_content = f"""Analyze and generate a team structure for: "{self.topic}"
|
||||
|
||||
STEP 1: TASK COMPLEXITY ANALYSIS
|
||||
First, analyze the task complexity:
|
||||
- Is this a simple task that 1-2 agents could handle?
|
||||
- Does it require multiple specialized perspectives?
|
||||
- Are there clear sequential dependencies between steps?
|
||||
TASK COMPLEXITY ANALYSIS (Pre-computed):
|
||||
- Complexity: {complexity}
|
||||
- Recommended agents: {recommended_agent_count}
|
||||
- Recommended tools based on task keywords: {', '.join(recommended_tools)}
|
||||
|
||||
STEP 1: VALIDATE TASK ANALYSIS
|
||||
Review the pre-computed analysis above. Adjust if needed based on your understanding.
|
||||
|
||||
STEP 2: DETERMINE OPTIMAL TEAM SIZE
|
||||
Based on your analysis, create the minimum number of agents needed:
|
||||
- Simple tasks: 1-2 agents
|
||||
- Moderate tasks: 2-3 agents
|
||||
- Complex tasks: 3-4 agents (maximum)
|
||||
Based on complexity analysis:
|
||||
- Simple tasks: 1-2 agents (single focused agent or simple pair)
|
||||
- Moderate tasks: 2-3 agents (researcher + executor pattern)
|
||||
- Complex tasks: 3-4 agents (specialized team)
|
||||
|
||||
IMPORTANT: Only create agents that provide meaningful specialization. Avoid unnecessary complexity.
|
||||
Recommended for this task: {recommended_agent_count} agent(s)
|
||||
|
||||
IMPORTANT: Avoid unnecessary complexity. Only add agents if there is meaningful specialization.
|
||||
Each agent must have a distinct, non-overlapping responsibility.
|
||||
|
||||
STEP 3: DESIGN THE TEAM (Pattern: {self.pattern})
|
||||
{workflow_guidance}
|
||||
|
||||
Each agent should have:
|
||||
- A clear, distinct role
|
||||
- A clear, distinct role with meaningful specialization
|
||||
- A specific goal
|
||||
- Relevant backstory
|
||||
- 1 focused task with clear description and expected output
|
||||
- Appropriate tools (only if needed)
|
||||
- Appropriate tools from the recommended list
|
||||
|
||||
Available Tools: CodeDocsSearchTool, CSVSearchTool, DirectorySearchTool, DOCXSearchTool, DirectoryReadTool, FileReadTool, TXTSearchTool, JSONSearchTool, MDXSearchTool, PDFSearchTool, RagTool, ScrapeElementFromWebsiteTool, ScrapeWebsiteTool, WebsiteSearchTool, XMLSearchTool, YoutubeChannelSearchTool, YoutubeVideoSearchTool.
|
||||
Only use these tools if the task requires them. Use empty list [] if no tools needed.
|
||||
AVAILABLE TOOLS BY CATEGORY:
|
||||
{tools_reference}
|
||||
|
||||
Example structure (2 agents for a writing task):
|
||||
LEGACY TOOLS (for backward compatibility):
|
||||
{legacy_tools}
|
||||
|
||||
RECOMMENDED TOOLS FOR THIS TASK: {', '.join(recommended_tools)}
|
||||
Prioritize using the recommended tools. Only add others if specifically needed.
|
||||
|
||||
Example structure (2 agents for a research + writing task):
|
||||
{{
|
||||
"roles": {{
|
||||
"researcher": {{
|
||||
"role": "Research Analyst",
|
||||
"goal": "Gather comprehensive information on the topic",
|
||||
"backstory": "Expert researcher skilled at finding and synthesizing information.",
|
||||
"tools": ["WebsiteSearchTool"],
|
||||
"tools": ["internet_search", "read_file"],
|
||||
"tasks": {{
|
||||
"research_task": {{
|
||||
"description": "Research key information about the topic and compile findings.",
|
||||
@@ -531,7 +767,7 @@ Example structure (2 agents for a writing task):
|
||||
"role": "Content Writer",
|
||||
"goal": "Create polished final content",
|
||||
"backstory": "Skilled writer who transforms research into engaging content.",
|
||||
"tools": [],
|
||||
"tools": ["write_file"],
|
||||
"tasks": {{
|
||||
"writing_task": {{
|
||||
"description": "Write the final content based on research findings.",
|
||||
@@ -543,6 +779,7 @@ Example structure (2 agents for a writing task):
|
||||
}}
|
||||
|
||||
Now generate the optimal team structure for: {self.topic}
|
||||
Use the recommended tools: {', '.join(recommended_tools)}
|
||||
"""
|
||||
return user_content
|
||||
|
||||
|
||||
@@ -29,6 +29,9 @@ Features:
|
||||
- at_mentions: @ mention autocomplete for files and directories
|
||||
- compare: CLI command comparison feature
|
||||
- eval: Agent evaluation framework (accuracy, performance, reliability, criteria)
|
||||
- hooks: Hook system for intercepting and modifying agent behavior
|
||||
- checkpoints: Shadow git checkpointing for file-level undo/restore
|
||||
- background: Background agent task execution and management
|
||||
"""
|
||||
|
||||
# Type hints available for IDE support
|
||||
@@ -71,6 +74,12 @@ __all__ = [
|
||||
'CompareHandler',
|
||||
# Agent Skills
|
||||
'SkillsHandler',
|
||||
# Hooks
|
||||
'HooksHandler',
|
||||
# Checkpoints
|
||||
'CheckpointsHandler',
|
||||
# Background
|
||||
'BackgroundHandler',
|
||||
]
|
||||
|
||||
def __getattr__(name):
|
||||
@@ -170,4 +179,13 @@ def __getattr__(name):
|
||||
elif name == 'SkillsHandler':
|
||||
from .skills import SkillsHandler
|
||||
return SkillsHandler
|
||||
elif name == 'HooksHandler':
|
||||
from .hooks import HooksHandler
|
||||
return HooksHandler
|
||||
elif name == 'CheckpointsHandler':
|
||||
from .checkpoints import CheckpointsHandler
|
||||
return CheckpointsHandler
|
||||
elif name == 'BackgroundHandler':
|
||||
from .background import BackgroundHandler
|
||||
return BackgroundHandler
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,375 @@
|
||||
"""
|
||||
Multi-Agent CLI Feature Handler
|
||||
|
||||
Provides CLI commands for defining and running multiple agents:
|
||||
- praisonai agents run --agent "name:role:tools" --agent "..." --task "..."
|
||||
|
||||
Example:
|
||||
praisonai agents run \
|
||||
--agent "researcher:Research Analyst:internet_search" \
|
||||
--agent "writer:Content Writer:write_file" \
|
||||
--task "Research AI trends and write a report"
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
|
||||
def parse_agent_definition(definition: str) -> Dict[str, Any]:
|
||||
"""Parse an agent definition string.
|
||||
|
||||
Format: name:role:tools:goal (tools and goal are optional)
|
||||
|
||||
Args:
|
||||
definition: Agent definition string
|
||||
|
||||
Returns:
|
||||
Dictionary with name, role, tools, and optionally goal
|
||||
"""
|
||||
parts = definition.split(':')
|
||||
|
||||
result = {
|
||||
'name': parts[0].strip() if len(parts) > 0 else 'agent',
|
||||
'role': parts[1].strip() if len(parts) > 1 else 'Assistant',
|
||||
'tools': [],
|
||||
'goal': None
|
||||
}
|
||||
|
||||
# Parse tools (comma-separated)
|
||||
if len(parts) > 2 and parts[2].strip():
|
||||
result['tools'] = [t.strip() for t in parts[2].split(',') if t.strip()]
|
||||
|
||||
# Parse goal if provided
|
||||
if len(parts) > 3:
|
||||
result['goal'] = parts[3].strip()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class MultiAgentHandler:
|
||||
"""Handler for multi-agent CLI commands."""
|
||||
|
||||
def __init__(self, verbose: bool = True):
|
||||
"""Initialize the multi-agent handler.
|
||||
|
||||
Args:
|
||||
verbose: Whether to print verbose output
|
||||
"""
|
||||
self.verbose = verbose
|
||||
|
||||
def create_agents_from_definitions(
|
||||
self,
|
||||
definitions: List[str],
|
||||
llm: Optional[str] = None,
|
||||
instructions: Optional[str] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Create agent configurations from definition strings.
|
||||
|
||||
Args:
|
||||
definitions: List of agent definition strings
|
||||
llm: LLM model to use for all agents
|
||||
instructions: Additional instructions for all agents
|
||||
|
||||
Returns:
|
||||
List of agent configuration dictionaries
|
||||
"""
|
||||
agents = []
|
||||
|
||||
for definition in definitions:
|
||||
agent_config = parse_agent_definition(definition)
|
||||
|
||||
if llm:
|
||||
agent_config['llm'] = llm
|
||||
|
||||
if instructions:
|
||||
agent_config['instructions'] = instructions
|
||||
|
||||
agents.append(agent_config)
|
||||
|
||||
return agents
|
||||
|
||||
def prepare_execution_config(
|
||||
self,
|
||||
agent_definitions: List[str],
|
||||
task: str,
|
||||
process: str = "sequential",
|
||||
llm: Optional[str] = None,
|
||||
instructions: Optional[str] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Prepare configuration for agent execution.
|
||||
|
||||
Args:
|
||||
agent_definitions: List of agent definition strings
|
||||
task: Task description
|
||||
process: Execution process (sequential, parallel)
|
||||
llm: LLM model to use
|
||||
instructions: Additional instructions
|
||||
|
||||
Returns:
|
||||
Execution configuration dictionary
|
||||
"""
|
||||
agents = self.create_agents_from_definitions(
|
||||
agent_definitions, llm, instructions
|
||||
)
|
||||
|
||||
return {
|
||||
'agents': agents,
|
||||
'task': task,
|
||||
'process': process,
|
||||
'llm': llm
|
||||
}
|
||||
|
||||
def _load_tools(self, tool_names: List[str]) -> List:
|
||||
"""Load tool functions by name.
|
||||
|
||||
Args:
|
||||
tool_names: List of tool names to load
|
||||
|
||||
Returns:
|
||||
List of tool functions
|
||||
"""
|
||||
tools = []
|
||||
|
||||
try:
|
||||
from praisonaiagents.tools import (
|
||||
internet_search, read_file, write_file, list_files,
|
||||
execute_command, read_csv, write_csv, analyze_csv
|
||||
)
|
||||
|
||||
tool_map = {
|
||||
'internet_search': internet_search,
|
||||
'read_file': read_file,
|
||||
'write_file': write_file,
|
||||
'list_files': list_files,
|
||||
'execute_command': execute_command,
|
||||
'read_csv': read_csv,
|
||||
'write_csv': write_csv,
|
||||
'analyze_csv': analyze_csv,
|
||||
}
|
||||
|
||||
for name in tool_names:
|
||||
if name in tool_map:
|
||||
tools.append(tool_map[name])
|
||||
else:
|
||||
if self.verbose:
|
||||
print(f"⚠ Tool '{name}' not found, skipping")
|
||||
except ImportError as e:
|
||||
if self.verbose:
|
||||
print(f"⚠ Could not load tools: {e}")
|
||||
|
||||
return tools
|
||||
|
||||
def _execute_agents(
|
||||
self,
|
||||
agents_config: List[Dict[str, Any]],
|
||||
task: str,
|
||||
process: str = "sequential",
|
||||
llm: Optional[str] = None
|
||||
) -> str:
|
||||
"""Execute agents with the given configuration.
|
||||
|
||||
Args:
|
||||
agents_config: List of agent configurations
|
||||
task: Task to execute
|
||||
process: Execution process
|
||||
llm: LLM model to use
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
try:
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
# Create Agent objects
|
||||
agents = []
|
||||
tasks = []
|
||||
|
||||
for i, config in enumerate(agents_config):
|
||||
# Load tools for this agent
|
||||
tools = self._load_tools(config.get('tools', []))
|
||||
|
||||
agent = Agent(
|
||||
name=config['name'],
|
||||
role=config['role'],
|
||||
goal=config.get('goal') or f"Complete tasks as {config['role']}",
|
||||
backstory=f"You are a skilled {config['role']}.",
|
||||
instructions=config.get('instructions', ''),
|
||||
tools=tools if tools else None,
|
||||
llm=config.get('llm') or llm or os.environ.get('OPENAI_MODEL_NAME', 'gpt-4o-mini'),
|
||||
verbose=self.verbose
|
||||
)
|
||||
agents.append(agent)
|
||||
|
||||
# Create task for this agent
|
||||
agent_task = Task(
|
||||
name=f"task_{config['name']}",
|
||||
description=task if i == 0 else f"Continue the work: {task}",
|
||||
expected_output="Completed task output",
|
||||
agent=agent
|
||||
)
|
||||
tasks.append(agent_task)
|
||||
|
||||
# Run agents
|
||||
if self.verbose:
|
||||
print(f"\n🚀 Running {len(agents)} agents ({process} mode)...")
|
||||
for agent in agents:
|
||||
print(f" - {agent.name}: {agent.role}")
|
||||
print()
|
||||
|
||||
praison_agents = PraisonAIAgents(
|
||||
agents=agents,
|
||||
tasks=tasks,
|
||||
process=process,
|
||||
verbose=self.verbose
|
||||
)
|
||||
|
||||
result = praison_agents.start()
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
if self.verbose:
|
||||
print(f"❌ Execution failed: {e}")
|
||||
raise
|
||||
|
||||
def run(
|
||||
self,
|
||||
agent_definitions: List[str],
|
||||
task: str,
|
||||
process: str = "sequential",
|
||||
llm: Optional[str] = None,
|
||||
instructions: Optional[str] = None
|
||||
) -> str:
|
||||
"""Run multiple agents with the given task.
|
||||
|
||||
Args:
|
||||
agent_definitions: List of agent definition strings
|
||||
task: Task to execute
|
||||
process: Execution process (sequential, parallel)
|
||||
llm: LLM model to use
|
||||
instructions: Additional instructions
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
config = self.prepare_execution_config(
|
||||
agent_definitions, task, process, llm, instructions
|
||||
)
|
||||
|
||||
return self._execute_agents(
|
||||
config['agents'],
|
||||
config['task'],
|
||||
config['process'],
|
||||
config['llm']
|
||||
)
|
||||
|
||||
|
||||
def handle_agents_command(args) -> int:
|
||||
"""Handle agents subcommand from CLI.
|
||||
|
||||
Args:
|
||||
args: Parsed command line arguments
|
||||
|
||||
Returns:
|
||||
Exit code (0 for success, non-zero for error)
|
||||
"""
|
||||
handler = MultiAgentHandler(verbose=True)
|
||||
|
||||
try:
|
||||
if args.agents_command == "run":
|
||||
if not hasattr(args, 'agent') or not args.agent:
|
||||
print("Error: At least one --agent is required")
|
||||
return 1
|
||||
|
||||
if not hasattr(args, 'task') or not args.task:
|
||||
print("Error: --task is required")
|
||||
return 1
|
||||
|
||||
result = handler.run(
|
||||
agent_definitions=args.agent,
|
||||
task=args.task,
|
||||
process=getattr(args, 'process', 'sequential'),
|
||||
llm=getattr(args, 'llm', None),
|
||||
instructions=getattr(args, 'instructions', None)
|
||||
)
|
||||
|
||||
print("\n" + "="*50)
|
||||
print("RESULT:")
|
||||
print("="*50)
|
||||
print(result)
|
||||
|
||||
elif args.agents_command == "list":
|
||||
# List available tools
|
||||
print("\nAvailable tools for agents:")
|
||||
print("-" * 30)
|
||||
tools = [
|
||||
("internet_search", "Search the web"),
|
||||
("read_file", "Read file contents"),
|
||||
("write_file", "Write to a file"),
|
||||
("list_files", "List directory contents"),
|
||||
("execute_command", "Execute shell commands"),
|
||||
("read_csv", "Read CSV files"),
|
||||
("write_csv", "Write CSV files"),
|
||||
("analyze_csv", "Analyze CSV data"),
|
||||
]
|
||||
for name, desc in tools:
|
||||
print(f" {name}: {desc}")
|
||||
|
||||
else:
|
||||
print(f"Unknown agents command: {args.agents_command}")
|
||||
return 1
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def add_agents_parser(subparsers) -> None:
|
||||
"""Add agents subcommand to argument parser.
|
||||
|
||||
Args:
|
||||
subparsers: Subparsers object from argparse
|
||||
"""
|
||||
# run command
|
||||
run_parser = subparsers.add_parser(
|
||||
'run',
|
||||
help='Run multiple agents with a task'
|
||||
)
|
||||
run_parser.add_argument(
|
||||
'--agent', '-a',
|
||||
action='append',
|
||||
required=True,
|
||||
help='Agent definition (format: name:role:tools). Can be used multiple times.'
|
||||
)
|
||||
run_parser.add_argument(
|
||||
'--task', '-t',
|
||||
required=True,
|
||||
help='Task for the agents to complete'
|
||||
)
|
||||
run_parser.add_argument(
|
||||
'--process', '-p',
|
||||
choices=['sequential', 'parallel'],
|
||||
default='sequential',
|
||||
help='Execution process (default: sequential)'
|
||||
)
|
||||
run_parser.add_argument(
|
||||
'--llm', '-m',
|
||||
help='LLM model to use for all agents'
|
||||
)
|
||||
run_parser.add_argument(
|
||||
'--instructions', '-i',
|
||||
help='Additional instructions for all agents'
|
||||
)
|
||||
run_parser.add_argument(
|
||||
'--verbose', '-v',
|
||||
action='store_true',
|
||||
help='Enable verbose output'
|
||||
)
|
||||
|
||||
# list command - list available tools
|
||||
list_parser = subparsers.add_parser(
|
||||
'list',
|
||||
help='List available tools for agents'
|
||||
)
|
||||
@@ -105,6 +105,7 @@ class FileSearchService:
|
||||
- Fuzzy match file paths
|
||||
- Cache results with TTL
|
||||
- Respect .gitignore (basic)
|
||||
- Support absolute paths (@/Users/...) and home paths (@~/...)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -125,9 +126,10 @@ class FileSearchService:
|
||||
# Cache: query -> (results, timestamp)
|
||||
self._cache: Dict[str, Tuple[List[FileSuggestion], float]] = {}
|
||||
|
||||
# File list cache
|
||||
# File list cache per directory
|
||||
self._file_list: Optional[List[str]] = None
|
||||
self._file_list_time: float = 0
|
||||
self._file_list_dir: Optional[str] = None
|
||||
|
||||
def _should_ignore(self, name: str) -> bool:
|
||||
"""Check if file/dir should be ignored."""
|
||||
@@ -136,20 +138,23 @@ class FileSearchService:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _get_file_list(self) -> List[str]:
|
||||
def _get_file_list(self, search_dir: str = None) -> List[str]:
|
||||
"""Get list of all files (cached)."""
|
||||
now = time.time()
|
||||
target_dir = search_dir or self._root
|
||||
|
||||
# Return cached if fresh
|
||||
if self._file_list is not None and (now - self._file_list_time) < self._cache_ttl:
|
||||
# Return cached if fresh and same directory
|
||||
if (self._file_list is not None and
|
||||
self._file_list_dir == target_dir and
|
||||
(now - self._file_list_time) < self._cache_ttl):
|
||||
return self._file_list
|
||||
|
||||
files = []
|
||||
|
||||
try:
|
||||
for root, dirs, filenames in os.walk(self._root):
|
||||
for root, dirs, filenames in os.walk(target_dir):
|
||||
# Calculate depth
|
||||
rel_root = os.path.relpath(root, self._root)
|
||||
rel_root = os.path.relpath(root, target_dir)
|
||||
depth = 0 if rel_root == '.' else rel_root.count(os.sep) + 1
|
||||
|
||||
if depth > self._max_depth:
|
||||
@@ -161,19 +166,20 @@ class FileSearchService:
|
||||
|
||||
# Add directories
|
||||
for d in dirs:
|
||||
rel_path = os.path.relpath(os.path.join(root, d), self._root)
|
||||
rel_path = os.path.relpath(os.path.join(root, d), target_dir)
|
||||
files.append(rel_path + os.sep)
|
||||
|
||||
# Add files
|
||||
for f in filenames:
|
||||
if not self._should_ignore(f):
|
||||
rel_path = os.path.relpath(os.path.join(root, f), self._root)
|
||||
rel_path = os.path.relpath(os.path.join(root, f), target_dir)
|
||||
files.append(rel_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
self._file_list = files
|
||||
self._file_list_time = now
|
||||
self._file_list_dir = target_dir
|
||||
return files
|
||||
|
||||
def _fuzzy_match(self, text: str, query: str) -> int:
|
||||
@@ -221,6 +227,11 @@ class FileSearchService:
|
||||
"""
|
||||
Search for files matching query.
|
||||
|
||||
Supports:
|
||||
- Relative paths: @file.txt, @src/main.py
|
||||
- Home paths: @~/Documents, @~/.config
|
||||
- Absolute paths: @/Users/name/file.txt, @/etc/hosts
|
||||
|
||||
Args:
|
||||
query: Search query (fuzzy matched)
|
||||
max_results: Maximum results to return
|
||||
@@ -228,8 +239,51 @@ class FileSearchService:
|
||||
Returns:
|
||||
List of FileSuggestion sorted by score
|
||||
"""
|
||||
# Determine search directory based on query
|
||||
search_dir = self._root
|
||||
search_query = query
|
||||
prefix = ""
|
||||
|
||||
# Handle home directory paths (@~/)
|
||||
if query.startswith('~'):
|
||||
expanded = os.path.expanduser(query)
|
||||
# Find the directory part and filename part
|
||||
if os.path.isdir(expanded):
|
||||
search_dir = expanded
|
||||
search_query = ""
|
||||
prefix = query.rstrip('/') + '/'
|
||||
else:
|
||||
parent = os.path.dirname(expanded)
|
||||
if os.path.isdir(parent):
|
||||
search_dir = parent
|
||||
search_query = os.path.basename(expanded)
|
||||
prefix = os.path.dirname(query).rstrip('/') + '/'
|
||||
else:
|
||||
# Just expand ~ and search from home
|
||||
search_dir = os.path.expanduser('~')
|
||||
search_query = query[2:] if query.startswith('~/') else query[1:]
|
||||
prefix = "~/"
|
||||
|
||||
# Handle absolute paths (@/Users/..., @/etc/...)
|
||||
elif query.startswith('/'):
|
||||
if os.path.isdir(query):
|
||||
search_dir = query
|
||||
search_query = ""
|
||||
prefix = query.rstrip('/') + '/'
|
||||
else:
|
||||
parent = os.path.dirname(query)
|
||||
if os.path.isdir(parent):
|
||||
search_dir = parent
|
||||
search_query = os.path.basename(query)
|
||||
prefix = parent.rstrip('/') + '/'
|
||||
else:
|
||||
# Search from root with the query
|
||||
search_dir = '/'
|
||||
search_query = query[1:]
|
||||
prefix = "/"
|
||||
|
||||
# Check cache
|
||||
cache_key = f"{query}:{max_results}"
|
||||
cache_key = f"{search_dir}:{search_query}:{max_results}"
|
||||
now = time.time()
|
||||
|
||||
if cache_key in self._cache:
|
||||
@@ -237,25 +291,26 @@ class FileSearchService:
|
||||
if (now - cached_time) < self._cache_ttl:
|
||||
return results
|
||||
|
||||
# Get all files
|
||||
all_files = self._get_file_list()
|
||||
# Get all files from the appropriate directory
|
||||
all_files = self._get_file_list(search_dir)
|
||||
|
||||
# Score and filter
|
||||
scored: List[Tuple[int, str]] = []
|
||||
for path in all_files:
|
||||
score = self._fuzzy_match(path, query)
|
||||
score = self._fuzzy_match(path, search_query)
|
||||
if score > 0:
|
||||
scored.append((score, path))
|
||||
|
||||
# Sort by score descending
|
||||
scored.sort(key=lambda x: (-x[0], x[1]))
|
||||
|
||||
# Build results
|
||||
# Build results with proper prefix
|
||||
results = []
|
||||
for score, path in scored[:max_results]:
|
||||
file_type = "directory" if path.endswith(os.sep) else "file"
|
||||
display_path = prefix + path if prefix else path
|
||||
results.append(FileSuggestion(
|
||||
path=path,
|
||||
path=display_path,
|
||||
file_type=file_type,
|
||||
score=score
|
||||
))
|
||||
|
||||
@@ -45,8 +45,10 @@ class AutoMemoryHandler(FlagHandler):
|
||||
def _get_auto_memory(self, user_id: str = None):
|
||||
"""Get AutoMemory instance lazily."""
|
||||
try:
|
||||
from praisonaiagents.memory import AutoMemory
|
||||
return AutoMemory(user_id=user_id or "default")
|
||||
from praisonaiagents.memory import AutoMemory, FileMemory
|
||||
# AutoMemory requires a FileMemory instance
|
||||
memory = FileMemory(user_id=user_id or "default")
|
||||
return AutoMemory(memory=memory)
|
||||
except ImportError:
|
||||
self.print_status(
|
||||
"AutoMemory requires praisonaiagents. Install with: pip install praisonaiagents",
|
||||
@@ -71,33 +73,36 @@ class AutoMemoryHandler(FlagHandler):
|
||||
config['auto_memory_user_id'] = flag_value.get('user_id', 'default')
|
||||
return config
|
||||
|
||||
def extract_memories(self, text: str, user_id: str = None) -> Dict[str, Any]:
|
||||
def extract_memories(self, text: str, user_id: str = None, user_message: str = None) -> list:
|
||||
"""
|
||||
Extract memories from text.
|
||||
Extract and store memories from text.
|
||||
|
||||
Args:
|
||||
text: Text to extract memories from
|
||||
text: Assistant response text to extract memories from
|
||||
user_id: User ID for memory isolation
|
||||
user_message: Original user message (for context)
|
||||
|
||||
Returns:
|
||||
Dictionary of extracted memories
|
||||
List of extracted memories
|
||||
"""
|
||||
auto_memory = self._get_auto_memory(user_id)
|
||||
if not auto_memory:
|
||||
return {}
|
||||
return []
|
||||
|
||||
try:
|
||||
if hasattr(auto_memory, 'extract'):
|
||||
memories = auto_memory.extract(text)
|
||||
self.print_status(f"🧠 Extracted {len(memories)} memories", "success")
|
||||
return memories
|
||||
elif hasattr(auto_memory, 'process'):
|
||||
result = auto_memory.process(text)
|
||||
return result
|
||||
# Use process_interaction which extracts AND stores memories
|
||||
memories = auto_memory.process_interaction(
|
||||
user_message=user_message or text,
|
||||
assistant_response=text if user_message else None,
|
||||
store=True # Store the memories
|
||||
)
|
||||
if memories:
|
||||
self.print_status(f"🧠 Extracted and stored {len(memories)} memories", "success")
|
||||
return memories
|
||||
except Exception as e:
|
||||
self.log(f"Memory extraction failed: {e}", "error")
|
||||
|
||||
return {}
|
||||
return []
|
||||
|
||||
def store_memory(self, content: str, user_id: str = None, importance: float = 0.5) -> bool:
|
||||
"""
|
||||
@@ -127,13 +132,14 @@ class AutoMemoryHandler(FlagHandler):
|
||||
self.log(f"Memory storage failed: {e}", "error")
|
||||
return False
|
||||
|
||||
def post_process_result(self, result: Any, flag_value: Any) -> Any:
|
||||
def post_process_result(self, result: Any, flag_value: Any, user_message: str = None) -> Any:
|
||||
"""
|
||||
Post-process result to extract and store memories.
|
||||
|
||||
Args:
|
||||
result: Agent output
|
||||
flag_value: Boolean or dict with configuration
|
||||
user_message: Original user message for context
|
||||
|
||||
Returns:
|
||||
Original result (memories are stored)
|
||||
@@ -144,15 +150,18 @@ class AutoMemoryHandler(FlagHandler):
|
||||
user_id = None
|
||||
if isinstance(flag_value, dict):
|
||||
user_id = flag_value.get('user_id')
|
||||
user_message = flag_value.get('user_message', user_message)
|
||||
|
||||
# Extract memories from result
|
||||
# Extract and store memories from result
|
||||
text = str(result)
|
||||
memories = self.extract_memories(text, user_id)
|
||||
memories = self.extract_memories(text, user_id, user_message=user_message)
|
||||
|
||||
if memories:
|
||||
self.print_status("\n🧠 Auto-extracted Memories:", "info")
|
||||
for key, value in memories.items():
|
||||
self.print_status(f" • {key}: {value[:50]}...", "info")
|
||||
for mem in memories:
|
||||
mem_type = mem.get('type', 'unknown')
|
||||
content = mem.get('content', str(mem))[:50]
|
||||
self.print_status(f" • {mem_type}: {content}...", "info")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
"""
|
||||
Background Agents CLI Feature for PraisonAI.
|
||||
|
||||
Provides CLI commands for managing background agent tasks.
|
||||
|
||||
Commands:
|
||||
- praisonai background list # List background tasks
|
||||
- praisonai background status <id> # Get task status
|
||||
- praisonai background cancel <id> # Cancel a task
|
||||
- praisonai background clear # Clear completed tasks
|
||||
"""
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
from typing import Optional, List
|
||||
|
||||
|
||||
class BackgroundHandler:
|
||||
"""
|
||||
Handler for background CLI commands.
|
||||
|
||||
Provides functionality to:
|
||||
- List running and completed background tasks
|
||||
- Check task status
|
||||
- Cancel running tasks
|
||||
- Clear completed tasks
|
||||
"""
|
||||
|
||||
def __init__(self, verbose: bool = False):
|
||||
self.verbose = verbose
|
||||
self._runner = None
|
||||
|
||||
@property
|
||||
def feature_name(self) -> str:
|
||||
return "background"
|
||||
|
||||
async def _get_runner(self):
|
||||
"""Lazy load the background runner."""
|
||||
if self._runner is None:
|
||||
from praisonaiagents.background import BackgroundRunner
|
||||
self._runner = BackgroundRunner()
|
||||
return self._runner
|
||||
|
||||
async def list_tasks(self, status: Optional[str] = None) -> List[dict]:
|
||||
"""
|
||||
List all background tasks.
|
||||
|
||||
Args:
|
||||
status: Optional status filter (pending, running, completed, failed)
|
||||
|
||||
Returns:
|
||||
List of task dictionaries
|
||||
"""
|
||||
runner = await self._get_runner()
|
||||
|
||||
status_enum = None
|
||||
if status:
|
||||
from praisonaiagents.background import TaskStatus
|
||||
status_map = {
|
||||
"pending": TaskStatus.PENDING,
|
||||
"running": TaskStatus.RUNNING,
|
||||
"completed": TaskStatus.COMPLETED,
|
||||
"failed": TaskStatus.FAILED,
|
||||
"cancelled": TaskStatus.CANCELLED
|
||||
}
|
||||
status_enum = status_map.get(status.lower())
|
||||
|
||||
tasks = runner.list_tasks(status=status_enum)
|
||||
|
||||
self._print_tasks_table(tasks)
|
||||
return tasks
|
||||
|
||||
def _print_tasks_table(self, tasks: List[dict]):
|
||||
"""Print tasks in table format."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
if not tasks:
|
||||
console.print("[yellow]No background tasks[/yellow]")
|
||||
return
|
||||
|
||||
table = Table(title="Background Tasks")
|
||||
table.add_column("ID", style="cyan")
|
||||
table.add_column("Name", style="green")
|
||||
table.add_column("Status", style="yellow")
|
||||
table.add_column("Progress", style="blue")
|
||||
table.add_column("Duration", style="magenta")
|
||||
|
||||
for task in tasks:
|
||||
status = task.get("status", "unknown")
|
||||
status_color = {
|
||||
"pending": "dim",
|
||||
"running": "yellow",
|
||||
"completed": "green",
|
||||
"failed": "red",
|
||||
"cancelled": "dim"
|
||||
}.get(status, "white")
|
||||
|
||||
progress = f"{task.get('progress', 0) * 100:.0f}%"
|
||||
duration = task.get("duration_seconds")
|
||||
duration_str = f"{duration:.1f}s" if duration else "-"
|
||||
|
||||
table.add_row(
|
||||
task.get("id", "?"),
|
||||
task.get("name", "unnamed"),
|
||||
f"[{status_color}]{status}[/{status_color}]",
|
||||
progress,
|
||||
duration_str
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
except ImportError:
|
||||
print("Background Tasks:")
|
||||
for task in tasks:
|
||||
print(f" {task.get('id')} - {task.get('name')} ({task.get('status')})")
|
||||
|
||||
async def get_status(self, task_id: str) -> Optional[dict]:
|
||||
"""
|
||||
Get status of a specific task.
|
||||
|
||||
Args:
|
||||
task_id: Task ID
|
||||
|
||||
Returns:
|
||||
Task dictionary or None
|
||||
"""
|
||||
runner = await self._get_runner()
|
||||
task = runner.get_task(task_id)
|
||||
|
||||
if task is None:
|
||||
self._print_error(f"Task not found: {task_id}")
|
||||
return None
|
||||
|
||||
task_dict = task.to_dict()
|
||||
self._print_task_details(task_dict)
|
||||
return task_dict
|
||||
|
||||
def _print_task_details(self, task: dict):
|
||||
"""Print detailed task information."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
|
||||
console = Console()
|
||||
|
||||
status = task.get("status", "unknown")
|
||||
status_color = {
|
||||
"pending": "dim",
|
||||
"running": "yellow",
|
||||
"completed": "green",
|
||||
"failed": "red",
|
||||
"cancelled": "dim"
|
||||
}.get(status, "white")
|
||||
|
||||
content = f"""
|
||||
[bold]ID:[/bold] {task.get('id')}
|
||||
[bold]Name:[/bold] {task.get('name')}
|
||||
[bold]Status:[/bold] [{status_color}]{status}[/{status_color}]
|
||||
[bold]Progress:[/bold] {task.get('progress', 0) * 100:.0f}%
|
||||
[bold]Created:[/bold] {task.get('created_at')}
|
||||
[bold]Started:[/bold] {task.get('started_at') or '-'}
|
||||
[bold]Completed:[/bold] {task.get('completed_at') or '-'}
|
||||
[bold]Duration:[/bold] {task.get('duration_seconds', 0):.2f}s
|
||||
"""
|
||||
if task.get("error"):
|
||||
content += f"[bold]Error:[/bold] [red]{task.get('error')}[/red]\n"
|
||||
|
||||
if task.get("result"):
|
||||
result_str = str(task.get("result"))[:200]
|
||||
content += f"[bold]Result:[/bold] {result_str}\n"
|
||||
|
||||
console.print(Panel(content.strip(), title=f"Task: {task.get('name')}"))
|
||||
except ImportError:
|
||||
print(f"Task: {task.get('name')}")
|
||||
print(f" ID: {task.get('id')}")
|
||||
print(f" Status: {task.get('status')}")
|
||||
print(f" Progress: {task.get('progress', 0) * 100:.0f}%")
|
||||
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
"""
|
||||
Cancel a running task.
|
||||
|
||||
Args:
|
||||
task_id: Task ID to cancel
|
||||
|
||||
Returns:
|
||||
True if cancelled
|
||||
"""
|
||||
runner = await self._get_runner()
|
||||
result = await runner.cancel_task(task_id)
|
||||
|
||||
if result:
|
||||
self._print_success(f"Task cancelled: {task_id}")
|
||||
else:
|
||||
self._print_error(f"Could not cancel task: {task_id}")
|
||||
|
||||
return result
|
||||
|
||||
async def clear_completed(self) -> int:
|
||||
"""
|
||||
Clear all completed tasks.
|
||||
|
||||
Returns:
|
||||
Number of tasks cleared
|
||||
"""
|
||||
runner = await self._get_runner()
|
||||
count = len([t for t in runner.tasks if t.is_completed])
|
||||
runner.clear_completed()
|
||||
|
||||
self._print_success(f"Cleared {count} completed tasks")
|
||||
return count
|
||||
|
||||
def _print_success(self, message: str):
|
||||
"""Print success message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[green]✓[/green] {message}")
|
||||
except ImportError:
|
||||
print(f"✓ {message}")
|
||||
|
||||
def _print_error(self, message: str):
|
||||
"""Print error message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[red]✗[/red] {message}")
|
||||
except ImportError:
|
||||
print(f"✗ {message}")
|
||||
|
||||
|
||||
def handle_background_command(args: List[str], verbose: bool = False):
|
||||
"""
|
||||
Handle background CLI commands.
|
||||
|
||||
Usage:
|
||||
praisonai background list [--status <status>]
|
||||
praisonai background status <task_id>
|
||||
praisonai background cancel <task_id>
|
||||
praisonai background clear
|
||||
"""
|
||||
handler = BackgroundHandler(verbose=verbose)
|
||||
|
||||
if not args:
|
||||
print("Usage: praisonai background <command> [options]")
|
||||
print("\nCommands:")
|
||||
print(" list [--status <status>] List background tasks")
|
||||
print(" status <task_id> Get task status")
|
||||
print(" cancel <task_id> Cancel a task")
|
||||
print(" clear Clear completed tasks")
|
||||
return
|
||||
|
||||
command = args[0]
|
||||
|
||||
if command == "list":
|
||||
status = None
|
||||
if "--status" in args:
|
||||
idx = args.index("--status")
|
||||
if idx + 1 < len(args):
|
||||
status = args[idx + 1]
|
||||
|
||||
asyncio.run(handler.list_tasks(status=status))
|
||||
|
||||
elif command == "status":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai background status <task_id>")
|
||||
return
|
||||
|
||||
asyncio.run(handler.get_status(args[1]))
|
||||
|
||||
elif command == "cancel":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai background cancel <task_id>")
|
||||
return
|
||||
|
||||
asyncio.run(handler.cancel_task(args[1]))
|
||||
|
||||
elif command == "clear":
|
||||
asyncio.run(handler.clear_completed())
|
||||
|
||||
elif command == "help" or command == "--help":
|
||||
print("Background CLI Commands:")
|
||||
print("\n praisonai background list [--status <status>]")
|
||||
print(" List all background tasks")
|
||||
print(" Status: pending, running, completed, failed, cancelled")
|
||||
print("\n praisonai background status <task_id>")
|
||||
print(" Get detailed status of a specific task")
|
||||
print("\n praisonai background cancel <task_id>")
|
||||
print(" Cancel a running task")
|
||||
print("\n praisonai background clear")
|
||||
print(" Clear all completed tasks from the list")
|
||||
|
||||
else:
|
||||
print(f"Unknown command: {command}")
|
||||
print("Use 'praisonai background help' for available commands")
|
||||
@@ -0,0 +1,338 @@
|
||||
"""
|
||||
Checkpoints CLI Feature for PraisonAI.
|
||||
|
||||
Provides CLI commands for managing file-level checkpoints.
|
||||
|
||||
Commands:
|
||||
- praisonai checkpoint save <message> # Save a checkpoint
|
||||
- praisonai checkpoint restore <id> # Restore to a checkpoint
|
||||
- praisonai checkpoint list # List all checkpoints
|
||||
- praisonai checkpoint diff [from] [to] # Show diff between checkpoints
|
||||
- praisonai checkpoint delete # Delete all checkpoints
|
||||
"""
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
from typing import Optional, List
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
class CheckpointsHandler:
|
||||
"""
|
||||
Handler for checkpoints CLI commands.
|
||||
|
||||
Provides functionality to:
|
||||
- Save checkpoints before file modifications
|
||||
- Restore workspace to previous checkpoints
|
||||
- View diffs between checkpoints
|
||||
- List and manage checkpoints
|
||||
"""
|
||||
|
||||
def __init__(self, workspace_dir: Optional[str] = None, verbose: bool = False):
|
||||
self.workspace_dir = workspace_dir or os.getcwd()
|
||||
self.verbose = verbose
|
||||
self._service = None
|
||||
|
||||
@property
|
||||
def feature_name(self) -> str:
|
||||
return "checkpoints"
|
||||
|
||||
async def _get_service(self):
|
||||
"""Lazy load the checkpoint service."""
|
||||
if self._service is None:
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
self._service = CheckpointService(
|
||||
workspace_dir=self.workspace_dir
|
||||
)
|
||||
await self._service.initialize()
|
||||
return self._service
|
||||
|
||||
async def save(self, message: str, allow_empty: bool = False) -> bool:
|
||||
"""
|
||||
Save a checkpoint.
|
||||
|
||||
Args:
|
||||
message: Checkpoint message
|
||||
allow_empty: Allow checkpoint even if no changes
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
service = await self._get_service()
|
||||
result = await service.save(message, allow_empty=allow_empty)
|
||||
|
||||
if result.success:
|
||||
self._print_success(f"Checkpoint saved: {result.checkpoint.short_id}")
|
||||
self._print_info(f"Message: {result.checkpoint.message}")
|
||||
return True
|
||||
else:
|
||||
self._print_error(f"Failed to save checkpoint: {result.error}")
|
||||
return False
|
||||
|
||||
async def restore(self, checkpoint_id: str) -> bool:
|
||||
"""
|
||||
Restore to a checkpoint.
|
||||
|
||||
Args:
|
||||
checkpoint_id: Checkpoint ID to restore
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
service = await self._get_service()
|
||||
result = await service.restore(checkpoint_id)
|
||||
|
||||
if result.success:
|
||||
self._print_success(f"Restored to checkpoint: {result.checkpoint.short_id}")
|
||||
return True
|
||||
else:
|
||||
self._print_error(f"Failed to restore: {result.error}")
|
||||
return False
|
||||
|
||||
async def list_checkpoints(self, limit: int = 20) -> List[dict]:
|
||||
"""
|
||||
List all checkpoints.
|
||||
|
||||
Args:
|
||||
limit: Maximum number to show
|
||||
|
||||
Returns:
|
||||
List of checkpoint dictionaries
|
||||
"""
|
||||
service = await self._get_service()
|
||||
checkpoints = await service.list_checkpoints(limit=limit)
|
||||
|
||||
if not checkpoints:
|
||||
self._print_info("No checkpoints found")
|
||||
return []
|
||||
|
||||
self._print_checkpoints_table(checkpoints)
|
||||
return [cp.to_dict() for cp in checkpoints]
|
||||
|
||||
def _print_checkpoints_table(self, checkpoints):
|
||||
"""Print checkpoints in table format."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
table = Table(title="Checkpoints")
|
||||
table.add_column("ID", style="cyan")
|
||||
table.add_column("Message", style="green")
|
||||
table.add_column("Timestamp", style="yellow")
|
||||
|
||||
for cp in checkpoints:
|
||||
table.add_row(
|
||||
cp.short_id,
|
||||
cp.message[:50] + "..." if len(cp.message) > 50 else cp.message,
|
||||
cp.timestamp.strftime("%Y-%m-%d %H:%M:%S")
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
except ImportError:
|
||||
print("Checkpoints:")
|
||||
for cp in checkpoints:
|
||||
print(f" {cp.short_id} - {cp.message} ({cp.timestamp})")
|
||||
|
||||
async def diff(
|
||||
self,
|
||||
from_id: Optional[str] = None,
|
||||
to_id: Optional[str] = None
|
||||
) -> dict:
|
||||
"""
|
||||
Show diff between checkpoints.
|
||||
|
||||
Args:
|
||||
from_id: Starting checkpoint (default: previous)
|
||||
to_id: Ending checkpoint (default: current)
|
||||
|
||||
Returns:
|
||||
Diff dictionary
|
||||
"""
|
||||
service = await self._get_service()
|
||||
diff = await service.diff(from_id, to_id)
|
||||
|
||||
self._print_diff(diff)
|
||||
return diff.to_dict()
|
||||
|
||||
def _print_diff(self, diff):
|
||||
"""Print diff in a readable format."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
if not diff.files:
|
||||
console.print("[yellow]No changes[/yellow]")
|
||||
return
|
||||
|
||||
console.print(f"\n[bold]Changes from {diff.from_checkpoint} to {diff.to_checkpoint or 'working directory'}[/bold]\n")
|
||||
|
||||
table = Table()
|
||||
table.add_column("Status", style="cyan")
|
||||
table.add_column("File", style="white")
|
||||
table.add_column("+", style="green")
|
||||
table.add_column("-", style="red")
|
||||
|
||||
for f in diff.files:
|
||||
status_color = {
|
||||
"added": "green",
|
||||
"modified": "yellow",
|
||||
"deleted": "red"
|
||||
}.get(f.status, "white")
|
||||
|
||||
table.add_row(
|
||||
f"[{status_color}]{f.status}[/{status_color}]",
|
||||
f.path,
|
||||
str(f.additions),
|
||||
str(f.deletions)
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
console.print(f"\n[green]+{diff.total_additions}[/green] [red]-{diff.total_deletions}[/red]")
|
||||
except ImportError:
|
||||
print(f"Changes from {diff.from_checkpoint} to {diff.to_checkpoint or 'working directory'}:")
|
||||
for f in diff.files:
|
||||
print(f" {f.status}: {f.path} (+{f.additions}/-{f.deletions})")
|
||||
|
||||
async def delete_all(self) -> bool:
|
||||
"""
|
||||
Delete all checkpoints.
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
service = await self._get_service()
|
||||
result = await service.delete_all()
|
||||
|
||||
if result:
|
||||
self._print_success("All checkpoints deleted")
|
||||
return True
|
||||
else:
|
||||
self._print_error("Failed to delete checkpoints")
|
||||
return False
|
||||
|
||||
def _print_success(self, message: str):
|
||||
"""Print success message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[green]✓[/green] {message}")
|
||||
except ImportError:
|
||||
print(f"✓ {message}")
|
||||
|
||||
def _print_error(self, message: str):
|
||||
"""Print error message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[red]✗[/red] {message}")
|
||||
except ImportError:
|
||||
print(f"✗ {message}")
|
||||
|
||||
def _print_info(self, message: str):
|
||||
"""Print info message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[dim]{message}[/dim]")
|
||||
except ImportError:
|
||||
print(message)
|
||||
|
||||
|
||||
def handle_checkpoint_command(args: List[str], workspace_dir: Optional[str] = None, verbose: bool = False):
|
||||
"""
|
||||
Handle checkpoint CLI commands.
|
||||
|
||||
Usage:
|
||||
praisonai checkpoint save <message> # Save a checkpoint
|
||||
praisonai checkpoint restore <id> # Restore to a checkpoint
|
||||
praisonai checkpoint list [--limit N] # List checkpoints
|
||||
praisonai checkpoint diff [from] [to] # Show diff
|
||||
praisonai checkpoint delete # Delete all checkpoints
|
||||
"""
|
||||
handler = CheckpointsHandler(workspace_dir=workspace_dir, verbose=verbose)
|
||||
|
||||
if not args:
|
||||
print("Usage: praisonai checkpoint <command> [options]")
|
||||
print("\nCommands:")
|
||||
print(" save <message> Save a checkpoint with message")
|
||||
print(" restore <id> Restore to a checkpoint")
|
||||
print(" list [--limit N] List checkpoints (default: 20)")
|
||||
print(" diff [from] [to] Show diff between checkpoints")
|
||||
print(" delete Delete all checkpoints")
|
||||
return
|
||||
|
||||
command = args[0]
|
||||
|
||||
if command == "save":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai checkpoint save <message>")
|
||||
return
|
||||
|
||||
message = " ".join(args[1:])
|
||||
allow_empty = "--allow-empty" in args
|
||||
if allow_empty:
|
||||
message = message.replace("--allow-empty", "").strip()
|
||||
|
||||
asyncio.run(handler.save(message, allow_empty=allow_empty))
|
||||
|
||||
elif command == "restore":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai checkpoint restore <checkpoint_id>")
|
||||
return
|
||||
|
||||
asyncio.run(handler.restore(args[1]))
|
||||
|
||||
elif command == "list":
|
||||
limit = 20
|
||||
if "--limit" in args:
|
||||
idx = args.index("--limit")
|
||||
if idx + 1 < len(args):
|
||||
try:
|
||||
limit = int(args[idx + 1])
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
asyncio.run(handler.list_checkpoints(limit=limit))
|
||||
|
||||
elif command == "diff":
|
||||
from_id = args[1] if len(args) > 1 else None
|
||||
to_id = args[2] if len(args) > 2 else None
|
||||
|
||||
asyncio.run(handler.diff(from_id, to_id))
|
||||
|
||||
elif command == "delete":
|
||||
# Confirm deletion
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.prompt import Confirm
|
||||
|
||||
console = Console()
|
||||
if Confirm.ask("[yellow]Delete all checkpoints?[/yellow]"):
|
||||
asyncio.run(handler.delete_all())
|
||||
else:
|
||||
console.print("[dim]Cancelled[/dim]")
|
||||
except ImportError:
|
||||
response = input("Delete all checkpoints? [y/N] ")
|
||||
if response.lower() == 'y':
|
||||
asyncio.run(handler.delete_all())
|
||||
else:
|
||||
print("Cancelled")
|
||||
|
||||
elif command == "help" or command == "--help":
|
||||
print("Checkpoint CLI Commands:")
|
||||
print("\n praisonai checkpoint save <message>")
|
||||
print(" Save a checkpoint with the given message")
|
||||
print(" Options: --allow-empty Allow checkpoint with no changes")
|
||||
print("\n praisonai checkpoint restore <id>")
|
||||
print(" Restore workspace to a specific checkpoint")
|
||||
print("\n praisonai checkpoint list [--limit N]")
|
||||
print(" List all checkpoints (default limit: 20)")
|
||||
print("\n praisonai checkpoint diff [from_id] [to_id]")
|
||||
print(" Show diff between checkpoints")
|
||||
print(" If no IDs given, shows diff from last checkpoint to current")
|
||||
print("\n praisonai checkpoint delete")
|
||||
print(" Delete all checkpoints for this workspace")
|
||||
|
||||
else:
|
||||
print(f"Unknown command: {command}")
|
||||
print("Use 'praisonai checkpoint help' for available commands")
|
||||
@@ -16,7 +16,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
COMPARE_MODES = {
|
||||
"basic": {},
|
||||
"tools": {"tools": "internet_search"},
|
||||
"tools": {"tools": "calculator,internet_search"},
|
||||
"research": {"research": True},
|
||||
"planning": {"planning": True},
|
||||
"memory": {"memory": True},
|
||||
@@ -177,12 +177,13 @@ def save_compare_result(result: CompareResult, path: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def format_comparison_table(result: CompareResult) -> str:
|
||||
def format_comparison_table(result: CompareResult, show_responses: bool = True) -> str:
|
||||
"""
|
||||
Format comparison result as a table string.
|
||||
Format comparison result as a table string with optional response display.
|
||||
|
||||
Args:
|
||||
result: CompareResult to format
|
||||
show_responses: Whether to show actual responses from each mode
|
||||
|
||||
Returns:
|
||||
Formatted table string
|
||||
@@ -190,10 +191,12 @@ def format_comparison_table(result: CompareResult) -> str:
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from rich.panel import Panel
|
||||
from io import StringIO
|
||||
|
||||
console = Console(file=StringIO(), force_terminal=True)
|
||||
|
||||
# Main comparison table
|
||||
table = Table(title=f"Comparison: {result.query[:50]}...")
|
||||
table.add_column("Mode", style="cyan")
|
||||
table.add_column("Time (ms)", style="green")
|
||||
@@ -224,6 +227,26 @@ def format_comparison_table(result: CompareResult) -> str:
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
# Show responses if enabled
|
||||
if show_responses:
|
||||
console.print("\n[bold cyan]📝 Responses:[/bold cyan]")
|
||||
for c in result.comparisons:
|
||||
if c.error is None:
|
||||
# Truncate long responses
|
||||
response_preview = c.output[:500] + "..." if len(c.output) > 500 else c.output
|
||||
console.print(Panel(
|
||||
response_preview,
|
||||
title=f"[bold]{c.mode}[/bold] ({c.execution_time_ms:.0f}ms)",
|
||||
border_style="cyan"
|
||||
))
|
||||
else:
|
||||
console.print(Panel(
|
||||
f"[red]Error: {c.error}[/red]",
|
||||
title=f"[bold]{c.mode}[/bold]",
|
||||
border_style="red"
|
||||
))
|
||||
|
||||
return console.file.getvalue()
|
||||
except ImportError:
|
||||
lines = [f"Comparison: {result.query}"]
|
||||
@@ -231,6 +254,8 @@ def format_comparison_table(result: CompareResult) -> str:
|
||||
for c in result.comparisons:
|
||||
status = "OK" if c.error is None else f"ERROR: {c.error}"
|
||||
lines.append(f"{c.mode}: {c.execution_time_ms:.1f}ms - {status}")
|
||||
if show_responses and c.error is None:
|
||||
lines.append(f" Response: {c.output[:200]}...")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@@ -407,13 +432,7 @@ class CompareHandler:
|
||||
from praisonaiagents.tools import internet_search
|
||||
tools.append(internet_search)
|
||||
elif name == "calculator":
|
||||
def calculator(expression: str) -> str:
|
||||
"""Evaluate a mathematical expression."""
|
||||
try:
|
||||
return str(eval(expression))
|
||||
except Exception as e:
|
||||
return f"Error: {e}"
|
||||
tools.append(calculator)
|
||||
tools.append(_calculator_tool)
|
||||
except ImportError:
|
||||
logger.warning(f"Could not load tool: {name}")
|
||||
|
||||
@@ -466,3 +485,50 @@ class CompareHandler:
|
||||
print(f"[red]Failed to save results to: {output_path}[/red]")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _calculator_tool(expression: str) -> str:
|
||||
"""
|
||||
Evaluate a mathematical expression safely.
|
||||
|
||||
Args:
|
||||
expression: Mathematical expression to evaluate (e.g., "2+2", "15*3", "100/4")
|
||||
|
||||
Returns:
|
||||
The result of the calculation as a string
|
||||
"""
|
||||
import ast
|
||||
import operator
|
||||
|
||||
# Safe operators
|
||||
operators = {
|
||||
ast.Add: operator.add,
|
||||
ast.Sub: operator.sub,
|
||||
ast.Mult: operator.mul,
|
||||
ast.Div: operator.truediv,
|
||||
ast.Pow: operator.pow,
|
||||
ast.USub: operator.neg,
|
||||
ast.UAdd: operator.pos,
|
||||
ast.Mod: operator.mod,
|
||||
ast.FloorDiv: operator.floordiv,
|
||||
}
|
||||
|
||||
def _eval(node):
|
||||
if isinstance(node, ast.Constant):
|
||||
return node.value
|
||||
elif isinstance(node, ast.BinOp):
|
||||
left = _eval(node.left)
|
||||
right = _eval(node.right)
|
||||
return operators[type(node.op)](left, right)
|
||||
elif isinstance(node, ast.UnaryOp):
|
||||
operand = _eval(node.operand)
|
||||
return operators[type(node.op)](operand)
|
||||
else:
|
||||
raise ValueError(f"Unsupported operation: {type(node)}")
|
||||
|
||||
try:
|
||||
tree = ast.parse(expression, mode='eval')
|
||||
result = _eval(tree.body)
|
||||
return str(result)
|
||||
except Exception as e:
|
||||
return f"Error evaluating '{expression}': {e}"
|
||||
@@ -26,42 +26,68 @@ class EvalHandler:
|
||||
|
||||
def run_accuracy(
|
||||
self,
|
||||
agent_file: str,
|
||||
input_text: str,
|
||||
expected_output: str,
|
||||
agent_file: Optional[str] = None,
|
||||
input_text: str = "",
|
||||
expected_output: str = "",
|
||||
iterations: int = 1,
|
||||
model: Optional[str] = None,
|
||||
output_file: Optional[str] = None
|
||||
output_file: Optional[str] = None,
|
||||
prompt: Optional[str] = None,
|
||||
llm: Optional[str] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Run accuracy evaluation on an agent.
|
||||
|
||||
Args:
|
||||
agent_file: Path to agents.yaml file
|
||||
agent_file: Path to agents.yaml file (optional if prompt is provided)
|
||||
input_text: Input to provide to the agent
|
||||
expected_output: Expected output to compare against
|
||||
iterations: Number of evaluation iterations
|
||||
model: LLM model for judging
|
||||
output_file: Path to save results
|
||||
prompt: Direct prompt (alternative to agent_file)
|
||||
llm: LLM model for the agent (when using prompt)
|
||||
|
||||
Returns:
|
||||
Evaluation result dictionary
|
||||
"""
|
||||
try:
|
||||
from praisonaiagents.eval import AccuracyEvaluator
|
||||
from praisonai.agents_generator import AgentsGenerator
|
||||
from praisonaiagents import Agent
|
||||
except ImportError as e:
|
||||
logger.error(f"Failed to import evaluation modules: {e}")
|
||||
return {"error": str(e)}
|
||||
|
||||
try:
|
||||
generator = AgentsGenerator(agent_file)
|
||||
agents = generator.generate_agents()
|
||||
|
||||
if not agents:
|
||||
return {"error": "No agents found in configuration"}
|
||||
|
||||
agent = agents[0] if isinstance(agents, list) else agents
|
||||
# Create agent either from file or from prompt
|
||||
if prompt:
|
||||
# Direct prompt mode - create agent on the fly
|
||||
agent = Agent(
|
||||
name="EvalAgent",
|
||||
role="Assistant",
|
||||
goal="Complete the given task",
|
||||
backstory="You are a helpful assistant.",
|
||||
llm=llm or model or "gpt-4o-mini",
|
||||
verbose=False
|
||||
)
|
||||
# Use prompt as input if input_text not provided
|
||||
if not input_text:
|
||||
input_text = prompt
|
||||
elif agent_file:
|
||||
# Load from agents.yaml
|
||||
try:
|
||||
from praisonai.agents_generator import AgentsGenerator
|
||||
generator = AgentsGenerator(agent_file)
|
||||
agents = generator.generate_agents()
|
||||
|
||||
if not agents:
|
||||
return {"error": "No agents found in configuration"}
|
||||
|
||||
agent = agents[0] if isinstance(agents, list) else agents
|
||||
except Exception as e:
|
||||
return {"error": f"Failed to load agents from {agent_file}: {e}"}
|
||||
else:
|
||||
return {"error": "Either --agent or --prompt must be provided"}
|
||||
|
||||
evaluator = AccuracyEvaluator(
|
||||
agent=agent,
|
||||
@@ -321,18 +347,49 @@ class EvalHandler:
|
||||
return batch_result
|
||||
|
||||
|
||||
def handle_eval_command(args) -> None:
|
||||
def handle_eval_command(args) -> int:
|
||||
"""
|
||||
Handle the eval CLI command.
|
||||
|
||||
Args:
|
||||
args: Parsed command line arguments
|
||||
args: Command line arguments (list or parsed namespace)
|
||||
|
||||
Returns:
|
||||
Exit code
|
||||
"""
|
||||
import argparse
|
||||
|
||||
# If args is a list, parse it first
|
||||
if isinstance(args, list):
|
||||
parser = argparse.ArgumentParser(prog="praisonai eval")
|
||||
subparsers = parser.add_subparsers(dest='eval_type')
|
||||
add_eval_parser_subcommands(subparsers)
|
||||
|
||||
try:
|
||||
args = parser.parse_args(args)
|
||||
except SystemExit:
|
||||
return 1
|
||||
|
||||
if not args.eval_type:
|
||||
parser.print_help()
|
||||
print("\n[bold]Examples:[/bold]")
|
||||
print(" praisonai eval accuracy --prompt \"What is 2+2?\" --expected \"4\"")
|
||||
print(" praisonai eval performance --agent agents.yaml --input \"Hello\"")
|
||||
return 0
|
||||
|
||||
handler = EvalHandler(verbose=getattr(args, 'verbose', False))
|
||||
|
||||
eval_type = getattr(args, 'eval_type', 'accuracy')
|
||||
agent_file = getattr(args, 'agent', 'agents.yaml')
|
||||
agent_file = getattr(args, 'agent', None)
|
||||
output_file = getattr(args, 'output', None)
|
||||
prompt = getattr(args, 'prompt', None)
|
||||
llm = getattr(args, 'llm', None)
|
||||
|
||||
# If no agent file and no prompt, check if agents.yaml exists
|
||||
if not agent_file and not prompt:
|
||||
import os
|
||||
if os.path.exists('agents.yaml'):
|
||||
agent_file = 'agents.yaml'
|
||||
|
||||
if eval_type == 'accuracy':
|
||||
result = handler.run_accuracy(
|
||||
@@ -341,7 +398,9 @@ def handle_eval_command(args) -> None:
|
||||
expected_output=getattr(args, 'expected', ''),
|
||||
iterations=getattr(args, 'iterations', 1),
|
||||
model=getattr(args, 'model', None),
|
||||
output_file=output_file
|
||||
output_file=output_file,
|
||||
prompt=prompt,
|
||||
llm=llm
|
||||
)
|
||||
elif eval_type == 'performance':
|
||||
result = handler.run_performance(
|
||||
@@ -384,12 +443,40 @@ def handle_eval_command(args) -> None:
|
||||
)
|
||||
else:
|
||||
print(f"Unknown evaluation type: {eval_type}")
|
||||
return
|
||||
return 1
|
||||
|
||||
if 'error' in result:
|
||||
print(f"Error: {result['error']}")
|
||||
return 1
|
||||
elif not getattr(args, 'quiet', False):
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def add_eval_parser_subcommands(subparsers) -> None:
|
||||
"""Add eval subcommand parsers to an existing subparsers object."""
|
||||
accuracy_parser = subparsers.add_parser('accuracy', help='Run accuracy evaluation')
|
||||
accuracy_parser.add_argument('--agent', '-a', help='Agent config file (optional if --prompt used)')
|
||||
accuracy_parser.add_argument('--prompt', '-p', type=str, help='Direct prompt (alternative to --agent)')
|
||||
accuracy_parser.add_argument('--llm', help='LLM model for agent (when using --prompt)')
|
||||
accuracy_parser.add_argument('--input', '-i', help='Input text (defaults to --prompt if not provided)')
|
||||
accuracy_parser.add_argument('--expected', '-e', required=True, help='Expected output')
|
||||
accuracy_parser.add_argument('--iterations', '-n', type=int, default=1, help='Number of iterations')
|
||||
accuracy_parser.add_argument('--model', '-m', help='Judge model')
|
||||
accuracy_parser.add_argument('--output', '-o', help='Output file')
|
||||
accuracy_parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
|
||||
accuracy_parser.add_argument('--quiet', '-q', action='store_true', help='Suppress JSON output')
|
||||
|
||||
perf_parser = subparsers.add_parser('performance', help='Run performance evaluation')
|
||||
perf_parser.add_argument('--agent', '-a', default='agents.yaml', help='Agent config file')
|
||||
perf_parser.add_argument('--input', '-i', default='Hello', help='Input text')
|
||||
perf_parser.add_argument('--iterations', '-n', type=int, default=10, help='Number of iterations')
|
||||
perf_parser.add_argument('--warmup', '-w', type=int, default=2, help='Warmup runs')
|
||||
perf_parser.add_argument('--memory', action='store_true', default=True, help='Track memory')
|
||||
perf_parser.add_argument('--output', '-o', help='Output file')
|
||||
perf_parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
|
||||
perf_parser.add_argument('--quiet', '-q', action='store_true', help='Suppress JSON output')
|
||||
|
||||
|
||||
def add_eval_parser(subparsers) -> None:
|
||||
@@ -410,8 +497,10 @@ def add_eval_parser(subparsers) -> None:
|
||||
'accuracy',
|
||||
help='Run accuracy evaluation'
|
||||
)
|
||||
accuracy_parser.add_argument('--agent', '-a', default='agents.yaml', help='Agent config file')
|
||||
accuracy_parser.add_argument('--input', '-i', required=True, help='Input text')
|
||||
accuracy_parser.add_argument('--agent', '-a', help='Agent config file (optional if --prompt used)')
|
||||
accuracy_parser.add_argument('--prompt', '-p', help='Direct prompt (alternative to --agent)')
|
||||
accuracy_parser.add_argument('--llm', help='LLM model for agent (when using --prompt)')
|
||||
accuracy_parser.add_argument('--input', '-i', help='Input text (defaults to --prompt if not provided)')
|
||||
accuracy_parser.add_argument('--expected', '-e', required=True, help='Expected output')
|
||||
accuracy_parser.add_argument('--iterations', '-n', type=int, default=1, help='Number of iterations')
|
||||
accuracy_parser.add_argument('--model', '-m', help='Judge model')
|
||||
|
||||
@@ -3,12 +3,21 @@ Fast Context Handler for CLI.
|
||||
|
||||
Provides codebase search capability using FastContext.
|
||||
Usage: praisonai "Find authentication code" --fast-context ./src
|
||||
|
||||
Fast Context Strategy:
|
||||
1. Extract search keywords from natural language query using LLM
|
||||
2. Provide folder tree context to help understand codebase structure
|
||||
3. Use both file name search AND content search (ripgrep-style)
|
||||
4. Return relevant code snippets for context
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import Any, Dict, Tuple, List
|
||||
from .base import FlagHandler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FastContextHandler(FlagHandler):
|
||||
"""
|
||||
@@ -61,14 +70,208 @@ class FastContextHandler(FlagHandler):
|
||||
|
||||
return True, ""
|
||||
|
||||
def _count_files_quick(self, path: str, max_count: int = 1000) -> int:
|
||||
"""Quickly count files in directory, stopping at max_count."""
|
||||
count = 0
|
||||
ignore_dirs = {'.git', '__pycache__', 'node_modules', '.venv', 'venv', '.cache', 'dist', 'build'}
|
||||
|
||||
for root, dirs, files in os.walk(path):
|
||||
dirs[:] = [d for d in dirs if d not in ignore_dirs and not d.startswith('.')]
|
||||
count += len(files)
|
||||
if count >= max_count:
|
||||
return count
|
||||
return count
|
||||
|
||||
def _get_folder_tree(self, path: str, max_depth: int = 2) -> str:
|
||||
"""
|
||||
Get a compact folder tree structure for context.
|
||||
|
||||
Optimized for large directories:
|
||||
- Limits depth to max_depth (default 2)
|
||||
- Shows max 10 files per directory
|
||||
- Truncates total output to 50 lines
|
||||
- Skips common non-code directories
|
||||
"""
|
||||
# Quick check for very large directories
|
||||
file_count = self._count_files_quick(path, 500)
|
||||
if file_count >= 500:
|
||||
# For large directories, reduce depth and show summary
|
||||
max_depth = 1
|
||||
logger.debug(f"Large directory detected ({file_count}+ files), reducing tree depth")
|
||||
|
||||
tree_lines = []
|
||||
ignore_dirs = {'.git', '__pycache__', 'node_modules', '.venv', 'venv', '.cache', 'dist', 'build', 'coverage', '.pytest_cache', '.mypy_cache', 'eggs', '*.egg-info'}
|
||||
ignore_exts = {'.pyc', '.pyo', '.so', '.o', '.a', '.dylib', '.class', '.jar'}
|
||||
total_files = 0
|
||||
max_total_files = 100 # Stop after seeing this many files
|
||||
|
||||
def walk_tree(current_path, prefix="", depth=0):
|
||||
nonlocal total_files
|
||||
if depth > max_depth or total_files >= max_total_files:
|
||||
return
|
||||
|
||||
try:
|
||||
entries = sorted(os.listdir(current_path))
|
||||
except PermissionError:
|
||||
return
|
||||
|
||||
dirs = []
|
||||
files = []
|
||||
for entry in entries:
|
||||
full_path = os.path.join(current_path, entry)
|
||||
if os.path.isdir(full_path):
|
||||
if entry not in ignore_dirs and not entry.startswith('.'):
|
||||
dirs.append(entry)
|
||||
else:
|
||||
ext = os.path.splitext(entry)[1]
|
||||
if ext not in ignore_exts and not entry.startswith('.'):
|
||||
files.append(entry)
|
||||
|
||||
# Limit directories shown at each level
|
||||
dirs_to_show = dirs[:15]
|
||||
for i, d in enumerate(dirs_to_show):
|
||||
is_last = (i == len(dirs_to_show) - 1) and not files
|
||||
tree_lines.append(f"{prefix}{'└── ' if is_last else '├── '}{d}/")
|
||||
walk_tree(os.path.join(current_path, d),
|
||||
prefix + (' ' if is_last else '│ '), depth + 1)
|
||||
if len(dirs) > 15:
|
||||
tree_lines.append(f"{prefix} ... and {len(dirs) - 15} more directories")
|
||||
|
||||
# Limit files shown
|
||||
files_to_show = files[:10]
|
||||
total_files += len(files)
|
||||
for i, f in enumerate(files_to_show):
|
||||
is_last = i == len(files_to_show) - 1
|
||||
tree_lines.append(f"{prefix}{'└── ' if is_last else '├── '}{f}")
|
||||
if len(files) > 10:
|
||||
tree_lines.append(f"{prefix} ... and {len(files) - 10} more files")
|
||||
|
||||
tree_lines.append(os.path.basename(path) + "/")
|
||||
walk_tree(path)
|
||||
return "\n".join(tree_lines[:50]) # Limit total lines
|
||||
|
||||
def _extract_keywords(self, query: str, folder_tree: str) -> List[str]:
|
||||
"""Extract search keywords from natural language query using LLM."""
|
||||
try:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
extractor = Agent(
|
||||
name="KeywordExtractor",
|
||||
role="Search Query Analyzer",
|
||||
goal="Extract file search keywords from natural language",
|
||||
backstory="Expert at understanding code search queries",
|
||||
llm="gpt-4o-mini",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
prompt = f"""Given this folder structure:
|
||||
{folder_tree}
|
||||
|
||||
And this search query: "{query}"
|
||||
|
||||
Extract 1-5 specific keywords/patterns to search for in file names and content.
|
||||
Return ONLY the keywords, one per line, no explanations.
|
||||
Focus on: file names, function names, class names, variable names that might match.
|
||||
If the query mentions a folder name that exists, include files from that folder."""
|
||||
|
||||
response = extractor.chat(prompt, stream=False)
|
||||
keywords = [k.strip() for k in str(response).strip().split('\n') if k.strip()]
|
||||
logger.debug(f"Extracted keywords: {keywords}")
|
||||
return keywords[:5] # Limit to 5 keywords
|
||||
except Exception as e:
|
||||
logger.debug(f"Keyword extraction failed: {e}")
|
||||
# Fallback: extract words from query
|
||||
return [w for w in query.split() if len(w) > 2][:3]
|
||||
|
||||
def _search_files_by_name(self, path: str, keywords: List[str]) -> List[Dict[str, Any]]:
|
||||
"""Search for files matching keywords in their names."""
|
||||
matches = []
|
||||
path = os.path.abspath(path)
|
||||
|
||||
ignore_dirs = {'.git', '__pycache__', 'node_modules', '.venv', 'venv'}
|
||||
|
||||
for root, dirs, files in os.walk(path):
|
||||
dirs[:] = [d for d in dirs if d not in ignore_dirs]
|
||||
|
||||
for filename in files:
|
||||
filepath = os.path.join(root, filename)
|
||||
rel_path = os.path.relpath(filepath, path)
|
||||
|
||||
# Check if any keyword matches filename or path
|
||||
for kw in keywords:
|
||||
kw_lower = kw.lower()
|
||||
if kw_lower in filename.lower() or kw_lower in rel_path.lower():
|
||||
matches.append({
|
||||
'file': rel_path,
|
||||
'lines': [(1, 50)], # First 50 lines
|
||||
'relevance': 1.0,
|
||||
'match_type': 'filename'
|
||||
})
|
||||
break
|
||||
|
||||
return matches[:20] # Limit results
|
||||
|
||||
def _search_files_by_content(self, path: str, keywords: List[str]) -> List[Dict[str, Any]]:
|
||||
"""Search for files containing keywords in content."""
|
||||
matches = []
|
||||
path = os.path.abspath(path)
|
||||
|
||||
ignore_dirs = {'.git', '__pycache__', 'node_modules', '.venv', 'venv'}
|
||||
code_exts = {'.py', '.js', '.ts', '.jsx', '.tsx', '.java', '.go', '.rs', '.rb', '.sh', '.yaml', '.yml', '.json', '.md', '.txt'}
|
||||
|
||||
for root, dirs, files in os.walk(path):
|
||||
dirs[:] = [d for d in dirs if d not in ignore_dirs]
|
||||
|
||||
for filename in files:
|
||||
ext = os.path.splitext(filename)[1]
|
||||
if ext not in code_exts:
|
||||
continue
|
||||
|
||||
filepath = os.path.join(root, filename)
|
||||
rel_path = os.path.relpath(filepath, path)
|
||||
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
|
||||
content = f.read(50000) # Read first 50KB
|
||||
|
||||
content_lower = content.lower()
|
||||
for kw in keywords:
|
||||
if kw.lower() in content_lower:
|
||||
# Find line numbers with matches
|
||||
lines = content.split('\n')
|
||||
match_lines = []
|
||||
for i, line in enumerate(lines[:200], 1): # Check first 200 lines
|
||||
if kw.lower() in line.lower():
|
||||
match_lines.append((max(1, i-2), min(len(lines), i+2)))
|
||||
|
||||
if match_lines:
|
||||
matches.append({
|
||||
'file': rel_path,
|
||||
'lines': match_lines[:5], # Limit line ranges
|
||||
'relevance': 0.8,
|
||||
'match_type': 'content'
|
||||
})
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return matches[:15] # Limit results
|
||||
|
||||
def search_context(self, query: str, path: str, **kwargs) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search for relevant context in the codebase.
|
||||
|
||||
Uses a smart multi-step approach:
|
||||
1. Get folder tree structure
|
||||
2. Use LLM to extract search keywords from query
|
||||
3. Search by filename AND content
|
||||
4. Return combined results
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
query: Natural language search query
|
||||
path: Path to search
|
||||
**kwargs: Additional search options
|
||||
- use_llm_keywords: Use LLM to extract keywords (default: True)
|
||||
|
||||
Returns:
|
||||
List of matching results
|
||||
@@ -84,32 +287,50 @@ class FastContextHandler(FlagHandler):
|
||||
return []
|
||||
|
||||
try:
|
||||
from praisonaiagents import FastContext
|
||||
use_llm_keywords = kwargs.get('use_llm_keywords', True)
|
||||
|
||||
fc = FastContext(
|
||||
path=path,
|
||||
model=kwargs.get('model', 'gpt-4o-mini'),
|
||||
max_turns=kwargs.get('max_turns', 4),
|
||||
parallelism=kwargs.get('parallelism', 8),
|
||||
timeout=kwargs.get('timeout', 30.0)
|
||||
)
|
||||
# Step 1: Get folder tree for context
|
||||
logger.debug(f"Getting folder tree for: {path}")
|
||||
folder_tree = self._get_folder_tree(path)
|
||||
logger.debug(f"Folder tree:\n{folder_tree}")
|
||||
|
||||
result = fc.search(query)
|
||||
# Step 2: Extract keywords
|
||||
if use_llm_keywords:
|
||||
logger.debug(f"Extracting keywords from query: {query}")
|
||||
keywords = self._extract_keywords(query, folder_tree)
|
||||
else:
|
||||
keywords = [w for w in query.split() if len(w) > 2][:3]
|
||||
|
||||
# Convert to list of dicts
|
||||
logger.debug(f"Search keywords: {keywords}")
|
||||
|
||||
if not keywords:
|
||||
self.print_status("Could not extract search keywords", "warning")
|
||||
return []
|
||||
|
||||
# Step 3: Search by filename
|
||||
logger.debug("Searching by filename...")
|
||||
filename_matches = self._search_files_by_name(path, keywords)
|
||||
logger.debug(f"Found {len(filename_matches)} filename matches")
|
||||
|
||||
# Step 4: Search by content
|
||||
logger.debug("Searching by content...")
|
||||
content_matches = self._search_files_by_content(path, keywords)
|
||||
logger.debug(f"Found {len(content_matches)} content matches")
|
||||
|
||||
# Combine and deduplicate
|
||||
seen_files = set()
|
||||
matches = []
|
||||
if hasattr(result, 'files'):
|
||||
for file_match in result.files:
|
||||
matches.append({
|
||||
'file': file_match.path,
|
||||
'lines': [(lr.start, lr.end) for lr in file_match.line_ranges],
|
||||
'relevance': getattr(file_match, 'relevance', 1.0)
|
||||
})
|
||||
|
||||
for m in filename_matches + content_matches:
|
||||
if m['file'] not in seen_files:
|
||||
seen_files.add(m['file'])
|
||||
matches.append(m)
|
||||
|
||||
self.print_status(f"🔍 Found {len(matches)} relevant files", "success")
|
||||
return matches
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"FastContext search error: {e}")
|
||||
self.log(f"FastContext search error: {e}", "error")
|
||||
return []
|
||||
|
||||
|
||||
@@ -75,7 +75,7 @@ class GuardrailHandler(FlagHandler):
|
||||
|
||||
Args:
|
||||
description: Natural language description of validation criteria
|
||||
llm: Optional LLM instance for validation
|
||||
llm: Optional LLM instance for validation (defaults to gpt-4o-mini)
|
||||
|
||||
Returns:
|
||||
LLMGuardrail instance or None if unavailable
|
||||
@@ -87,6 +87,12 @@ class GuardrailHandler(FlagHandler):
|
||||
)
|
||||
return None
|
||||
|
||||
# Provide default LLM if none specified
|
||||
if llm is None:
|
||||
import os
|
||||
# Try to get from environment or use default
|
||||
llm = os.environ.get('OPENAI_MODEL_NAME', 'gpt-4o-mini')
|
||||
|
||||
return LLMGuardrail(description=description, llm=llm)
|
||||
|
||||
def validate_output(self, output: str, description: str, llm: Any = None) -> Tuple[bool, str]:
|
||||
|
||||
@@ -72,11 +72,11 @@ class HandoffHandler(FlagHandler):
|
||||
self.print_status(msg, "error")
|
||||
return []
|
||||
|
||||
from praisonaiagents import Agent, Handoff
|
||||
from praisonaiagents import Agent, handoff
|
||||
|
||||
agents = []
|
||||
|
||||
# Create agents for each role
|
||||
# First pass: Create all agents without handoffs
|
||||
for i, name in enumerate(agent_names):
|
||||
agent_config = {
|
||||
"name": name.title().replace("_", " "),
|
||||
@@ -91,26 +91,45 @@ class HandoffHandler(FlagHandler):
|
||||
agent = Agent(**agent_config)
|
||||
agents.append(agent)
|
||||
|
||||
# Set up handoffs between agents
|
||||
for i, agent in enumerate(agents[:-1]):
|
||||
next_agent = agents[i + 1]
|
||||
handoff = Handoff(
|
||||
target=next_agent,
|
||||
name=f"handoff_to_{agent_names[i + 1]}",
|
||||
description=f"Hand off task to {agent_names[i + 1]} for further processing"
|
||||
)
|
||||
# Second pass: Set up handoffs by recreating agents with handoff config
|
||||
final_agents = []
|
||||
for i, name in enumerate(agent_names):
|
||||
agent_config = {
|
||||
"name": name.title().replace("_", " "),
|
||||
"role": name.title().replace("_", " "),
|
||||
"goal": f"Complete tasks as {name}",
|
||||
"backstory": f"Expert {name} with specialized skills"
|
||||
}
|
||||
|
||||
# Add handoff to agent
|
||||
if not hasattr(agent, 'handoffs') or agent.handoffs is None:
|
||||
agent.handoffs = []
|
||||
agent.handoffs.append(handoff)
|
||||
if llm:
|
||||
agent_config["llm"] = llm
|
||||
|
||||
# Add handoff to next agent if not the last agent
|
||||
if i < len(agent_names) - 1:
|
||||
next_agent = agents[i + 1]
|
||||
next_agent_name = agent_names[i + 1]
|
||||
|
||||
# Pass handoffs via constructor - this properly registers them as tools
|
||||
agent_config["handoffs"] = [
|
||||
handoff(
|
||||
agent=next_agent,
|
||||
tool_name_override=f"handoff_to_{next_agent_name}",
|
||||
tool_description_override=f"Hand off the task to {next_agent_name}. Call this after completing your part."
|
||||
)
|
||||
]
|
||||
|
||||
# Add instructions to use handoff
|
||||
agent_config["instructions"] = f"After completing your task, you MUST use the handoff_to_{next_agent_name} tool to pass your work to the next agent."
|
||||
|
||||
final_agent = Agent(**agent_config)
|
||||
final_agents.append(final_agent)
|
||||
|
||||
self.print_status(f"🤝 Created {len(agents)} agents with handoff chain", "success")
|
||||
self.print_status(f"🤝 Created {len(final_agents)} agents with handoff chain", "success")
|
||||
for i, name in enumerate(agent_names):
|
||||
arrow = " → " if i < len(agent_names) - 1 else ""
|
||||
self.print_status(f" {name}{arrow}", "info")
|
||||
|
||||
return agents
|
||||
return final_agents
|
||||
|
||||
def apply_to_agent_config(self, config: Dict[str, Any], flag_value: Any) -> Dict[str, Any]:
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,583 @@
|
||||
"""
|
||||
Hooks CLI Feature for PraisonAI.
|
||||
|
||||
Provides CLI commands for managing and testing hooks.
|
||||
|
||||
Commands:
|
||||
- praisonai hooks list # List registered hooks
|
||||
- praisonai hooks test <event> <target> # Test hooks for an event
|
||||
- praisonai hooks validate <config> # Validate hooks configuration
|
||||
- praisonai hooks run <command> # Run a command hook manually
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
from typing import Optional, Dict, Any, List
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class HooksConfig:
|
||||
"""Configuration for hooks loaded from file."""
|
||||
hooks: List[Dict[str, Any]] = field(default_factory=list)
|
||||
enabled: bool = True
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, path: str) -> "HooksConfig":
|
||||
"""Load hooks configuration from a JSON/YAML file."""
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"Hooks config file not found: {path}")
|
||||
|
||||
with open(path, 'r') as f:
|
||||
if path.endswith('.yaml') or path.endswith('.yml'):
|
||||
try:
|
||||
import yaml
|
||||
data = yaml.safe_load(f)
|
||||
except ImportError:
|
||||
raise ImportError("PyYAML required for YAML config files")
|
||||
else:
|
||||
data = json.load(f)
|
||||
|
||||
return cls(
|
||||
hooks=data.get('hooks', []),
|
||||
enabled=data.get('enabled', True)
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
'hooks': self.hooks,
|
||||
'enabled': self.enabled
|
||||
}
|
||||
|
||||
|
||||
class HooksHandler:
|
||||
"""
|
||||
Handler for hooks CLI commands.
|
||||
|
||||
Provides functionality to:
|
||||
- List registered hooks
|
||||
- Test hooks with sample events
|
||||
- Validate hook configurations
|
||||
- Run command hooks manually
|
||||
"""
|
||||
|
||||
def __init__(self, verbose: bool = False):
|
||||
self.verbose = verbose
|
||||
self._registry = None
|
||||
self._runner = None
|
||||
|
||||
@property
|
||||
def feature_name(self) -> str:
|
||||
return "hooks"
|
||||
|
||||
def _get_registry(self):
|
||||
"""Lazy load the hook registry."""
|
||||
if self._registry is None:
|
||||
from praisonaiagents.hooks import HookRegistry
|
||||
self._registry = HookRegistry()
|
||||
return self._registry
|
||||
|
||||
def _get_runner(self):
|
||||
"""Lazy load the hook runner."""
|
||||
if self._runner is None:
|
||||
from praisonaiagents.hooks import HookRunner
|
||||
self._runner = HookRunner(self._get_registry())
|
||||
return self._runner
|
||||
|
||||
def list_hooks(self, format: str = "table") -> Dict[str, Any]:
|
||||
"""
|
||||
List all registered hooks.
|
||||
|
||||
Args:
|
||||
format: Output format (table, json)
|
||||
|
||||
Returns:
|
||||
Dictionary of hooks by event
|
||||
"""
|
||||
registry = self._get_registry()
|
||||
hooks = registry.list_hooks()
|
||||
|
||||
if format == "json":
|
||||
print(json.dumps(hooks, indent=2))
|
||||
else:
|
||||
self._print_hooks_table(hooks)
|
||||
|
||||
return hooks
|
||||
|
||||
def _print_hooks_table(self, hooks: Dict[str, List[Dict]]):
|
||||
"""Print hooks in table format."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
if not hooks:
|
||||
console.print("[yellow]No hooks registered[/yellow]")
|
||||
return
|
||||
|
||||
table = Table(title="Registered Hooks")
|
||||
table.add_column("Event", style="cyan")
|
||||
table.add_column("Name", style="green")
|
||||
table.add_column("Type", style="blue")
|
||||
table.add_column("Matcher", style="magenta")
|
||||
table.add_column("Enabled", style="yellow")
|
||||
|
||||
for event, event_hooks in hooks.items():
|
||||
for hook in event_hooks:
|
||||
table.add_row(
|
||||
event,
|
||||
hook.get('name', 'unknown'),
|
||||
hook.get('type', 'unknown'),
|
||||
hook.get('matcher') or '*',
|
||||
'✓' if hook.get('enabled', True) else '✗'
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
except ImportError:
|
||||
# Fallback to simple print
|
||||
print("Registered Hooks:")
|
||||
for event, event_hooks in hooks.items():
|
||||
print(f"\n {event}:")
|
||||
for hook in event_hooks:
|
||||
print(f" - {hook.get('name', 'unknown')} ({hook.get('type', 'unknown')})")
|
||||
|
||||
def test_hook(
|
||||
self,
|
||||
event: str,
|
||||
target: Optional[str] = None,
|
||||
input_data: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Test hooks for a specific event.
|
||||
|
||||
Args:
|
||||
event: Event name (before_tool, after_tool, etc.)
|
||||
target: Optional target to filter hooks (e.g., tool name)
|
||||
input_data: Optional custom input data
|
||||
|
||||
Returns:
|
||||
List of hook execution results
|
||||
"""
|
||||
from praisonaiagents.hooks import HookEvent
|
||||
from praisonaiagents.hooks.events import BeforeToolInput, BeforeAgentInput
|
||||
|
||||
# Map string to HookEvent
|
||||
event_map = {
|
||||
'before_tool': HookEvent.BEFORE_TOOL,
|
||||
'after_tool': HookEvent.AFTER_TOOL,
|
||||
'before_agent': HookEvent.BEFORE_AGENT,
|
||||
'after_agent': HookEvent.AFTER_AGENT,
|
||||
'session_start': HookEvent.SESSION_START,
|
||||
'session_end': HookEvent.SESSION_END,
|
||||
}
|
||||
|
||||
hook_event = event_map.get(event.lower())
|
||||
if hook_event is None:
|
||||
print(f"[red]Unknown event: {event}[/red]")
|
||||
print(f"Available events: {', '.join(event_map.keys())}")
|
||||
return []
|
||||
|
||||
# Create sample input based on event type
|
||||
from datetime import datetime
|
||||
timestamp = datetime.now().isoformat()
|
||||
|
||||
if hook_event in (HookEvent.BEFORE_TOOL, HookEvent.AFTER_TOOL):
|
||||
sample_input = BeforeToolInput(
|
||||
session_id="cli-test",
|
||||
cwd=os.getcwd(),
|
||||
event_name=event,
|
||||
timestamp=timestamp,
|
||||
tool_name=target or "test_tool",
|
||||
tool_input=input_data or {"test": "data"}
|
||||
)
|
||||
else:
|
||||
sample_input = BeforeAgentInput(
|
||||
session_id="cli-test",
|
||||
cwd=os.getcwd(),
|
||||
event_name=event,
|
||||
timestamp=timestamp,
|
||||
prompt=target or "Test prompt"
|
||||
)
|
||||
|
||||
runner = self._get_runner()
|
||||
results = asyncio.run(runner.execute(hook_event, sample_input, target))
|
||||
|
||||
# Print results
|
||||
self._print_test_results(results)
|
||||
|
||||
return [r.to_dict() for r in results]
|
||||
|
||||
def _print_test_results(self, results):
|
||||
"""Print test results."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
|
||||
console = Console()
|
||||
|
||||
if not results:
|
||||
console.print("[yellow]No hooks executed[/yellow]")
|
||||
return
|
||||
|
||||
console.print(f"\n[bold]Hook Execution Results ({len(results)} hooks)[/bold]\n")
|
||||
|
||||
for result in results:
|
||||
status = "[green]✓ PASSED[/green]" if result.success else "[red]✗ FAILED[/red]"
|
||||
decision = result.output.decision if result.output else "N/A"
|
||||
|
||||
content = f"""
|
||||
Hook: {result.hook_name}
|
||||
Status: {status}
|
||||
Decision: {decision}
|
||||
Duration: {result.duration_ms:.2f}ms
|
||||
"""
|
||||
if result.error:
|
||||
content += f"Error: {result.error}\n"
|
||||
if result.output and result.output.reason:
|
||||
content += f"Reason: {result.output.reason}\n"
|
||||
|
||||
panel_style = "green" if result.success else "red"
|
||||
console.print(Panel(content.strip(), title=result.hook_name, border_style=panel_style))
|
||||
except ImportError:
|
||||
print(f"\nHook Execution Results ({len(results)} hooks):")
|
||||
for result in results:
|
||||
status = "PASSED" if result.success else "FAILED"
|
||||
print(f" - {result.hook_name}: {status}")
|
||||
|
||||
def validate_config(self, config_path: str) -> bool:
|
||||
"""
|
||||
Validate a hooks configuration file.
|
||||
|
||||
Args:
|
||||
config_path: Path to the configuration file
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
"""
|
||||
try:
|
||||
config = HooksConfig.from_file(config_path)
|
||||
|
||||
errors = []
|
||||
warnings = []
|
||||
|
||||
for i, hook in enumerate(config.hooks):
|
||||
# Check required fields
|
||||
if 'event' not in hook:
|
||||
errors.append(f"Hook {i}: missing 'event' field")
|
||||
|
||||
if 'command' not in hook and 'function' not in hook:
|
||||
errors.append(f"Hook {i}: must have 'command' or 'function'")
|
||||
|
||||
# Check event name
|
||||
valid_events = ['before_tool', 'after_tool', 'before_agent',
|
||||
'after_agent', 'session_start', 'session_end']
|
||||
if hook.get('event') and hook['event'] not in valid_events:
|
||||
warnings.append(f"Hook {i}: unknown event '{hook['event']}'")
|
||||
|
||||
# Check timeout
|
||||
if 'timeout' in hook:
|
||||
try:
|
||||
timeout = float(hook['timeout'])
|
||||
if timeout <= 0:
|
||||
warnings.append(f"Hook {i}: timeout should be positive")
|
||||
except (ValueError, TypeError):
|
||||
errors.append(f"Hook {i}: invalid timeout value")
|
||||
|
||||
# Print results
|
||||
try:
|
||||
from rich.console import Console
|
||||
console = Console()
|
||||
|
||||
if errors:
|
||||
console.print("[red]Validation FAILED[/red]")
|
||||
for error in errors:
|
||||
console.print(f" [red]✗[/red] {error}")
|
||||
else:
|
||||
console.print("[green]Validation PASSED[/green]")
|
||||
|
||||
if warnings:
|
||||
console.print("\n[yellow]Warnings:[/yellow]")
|
||||
for warning in warnings:
|
||||
console.print(f" [yellow]![/yellow] {warning}")
|
||||
|
||||
console.print(f"\n[dim]Found {len(config.hooks)} hooks in config[/dim]")
|
||||
except ImportError:
|
||||
if errors:
|
||||
print("Validation FAILED")
|
||||
for error in errors:
|
||||
print(f" - {error}")
|
||||
else:
|
||||
print("Validation PASSED")
|
||||
|
||||
if warnings:
|
||||
print("\nWarnings:")
|
||||
for warning in warnings:
|
||||
print(f" - {warning}")
|
||||
|
||||
return len(errors) == 0
|
||||
|
||||
except FileNotFoundError as e:
|
||||
print(f"[red]Error: {e}[/red]")
|
||||
return False
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"[red]Invalid JSON: {e}[/red]")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"[red]Error validating config: {e}[/red]")
|
||||
return False
|
||||
|
||||
def run_command_hook(
|
||||
self,
|
||||
command: str,
|
||||
event: str = "before_tool",
|
||||
input_json: Optional[str] = None,
|
||||
timeout: float = 60.0
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Run a command hook manually.
|
||||
|
||||
Args:
|
||||
command: Shell command to execute
|
||||
event: Event type for the hook
|
||||
input_json: Optional JSON input to pass to the command
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
from praisonaiagents.hooks import HookEvent, HookRegistry, HookRunner
|
||||
from praisonaiagents.hooks.events import BeforeToolInput
|
||||
from datetime import datetime
|
||||
|
||||
registry = HookRegistry()
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command=command,
|
||||
name="cli_test_hook",
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
|
||||
# Create input
|
||||
if input_json:
|
||||
try:
|
||||
custom_input = json.loads(input_json)
|
||||
except json.JSONDecodeError:
|
||||
custom_input = {}
|
||||
else:
|
||||
custom_input = {}
|
||||
|
||||
sample_input = BeforeToolInput(
|
||||
session_id="cli-manual",
|
||||
cwd=os.getcwd(),
|
||||
event_name=event,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
tool_name=custom_input.get('tool_name', 'manual_test'),
|
||||
tool_input=custom_input.get('tool_input', {})
|
||||
)
|
||||
|
||||
results = asyncio.run(runner.execute(HookEvent.BEFORE_TOOL, sample_input))
|
||||
|
||||
if results:
|
||||
result = results[0]
|
||||
self._print_command_result(result)
|
||||
return result.to_dict()
|
||||
|
||||
return {}
|
||||
|
||||
def _print_command_result(self, result):
|
||||
"""Print command hook result."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.syntax import Syntax
|
||||
|
||||
console = Console()
|
||||
|
||||
console.print(f"\n[bold]Command Hook Result[/bold]")
|
||||
console.print(f"Exit Code: {result.exit_code}")
|
||||
console.print(f"Success: {'[green]Yes[/green]' if result.success else '[red]No[/red]'}")
|
||||
console.print(f"Duration: {result.duration_ms:.2f}ms")
|
||||
|
||||
if result.stdout:
|
||||
console.print("\n[bold]STDOUT:[/bold]")
|
||||
console.print(result.stdout)
|
||||
|
||||
if result.stderr:
|
||||
console.print("\n[bold]STDERR:[/bold]")
|
||||
console.print(f"[red]{result.stderr}[/red]")
|
||||
|
||||
if result.output:
|
||||
console.print(f"\n[bold]Decision:[/bold] {result.output.decision}")
|
||||
if result.output.reason:
|
||||
console.print(f"[bold]Reason:[/bold] {result.output.reason}")
|
||||
except ImportError:
|
||||
print(f"\nCommand Hook Result")
|
||||
print(f"Exit Code: {result.exit_code}")
|
||||
print(f"Success: {result.success}")
|
||||
if result.stdout:
|
||||
print(f"STDOUT: {result.stdout}")
|
||||
if result.stderr:
|
||||
print(f"STDERR: {result.stderr}")
|
||||
|
||||
def register_from_config(self, config_path: str) -> int:
|
||||
"""
|
||||
Register hooks from a configuration file.
|
||||
|
||||
Args:
|
||||
config_path: Path to the configuration file
|
||||
|
||||
Returns:
|
||||
Number of hooks registered
|
||||
"""
|
||||
from praisonaiagents.hooks import HookEvent
|
||||
|
||||
config = HooksConfig.from_file(config_path)
|
||||
registry = self._get_registry()
|
||||
|
||||
event_map = {
|
||||
'before_tool': HookEvent.BEFORE_TOOL,
|
||||
'after_tool': HookEvent.AFTER_TOOL,
|
||||
'before_agent': HookEvent.BEFORE_AGENT,
|
||||
'after_agent': HookEvent.AFTER_AGENT,
|
||||
'session_start': HookEvent.SESSION_START,
|
||||
'session_end': HookEvent.SESSION_END,
|
||||
}
|
||||
|
||||
count = 0
|
||||
for hook in config.hooks:
|
||||
event = event_map.get(hook.get('event', '').lower())
|
||||
if event is None:
|
||||
continue
|
||||
|
||||
if 'command' in hook:
|
||||
registry.register_command(
|
||||
event=event,
|
||||
command=hook['command'],
|
||||
name=hook.get('name'),
|
||||
matcher=hook.get('matcher'),
|
||||
timeout=hook.get('timeout', 60.0),
|
||||
env=hook.get('env', {})
|
||||
)
|
||||
count += 1
|
||||
|
||||
print(f"[green]Registered {count} hooks from {config_path}[/green]")
|
||||
return count
|
||||
|
||||
|
||||
def handle_hooks_command(args: List[str], verbose: bool = False):
|
||||
"""
|
||||
Handle hooks CLI commands.
|
||||
|
||||
Usage:
|
||||
praisonai hooks list [--format json|table]
|
||||
praisonai hooks test <event> [target] [--input '{"key": "value"}']
|
||||
praisonai hooks validate <config_path>
|
||||
praisonai hooks run <command> [--event <event>] [--input <json>] [--timeout <seconds>]
|
||||
praisonai hooks load <config_path>
|
||||
"""
|
||||
handler = HooksHandler(verbose=verbose)
|
||||
|
||||
if not args:
|
||||
print("Usage: praisonai hooks <command> [options]")
|
||||
print("\nCommands:")
|
||||
print(" list List registered hooks")
|
||||
print(" test <event> [target] Test hooks for an event")
|
||||
print(" validate <config> Validate hooks configuration")
|
||||
print(" run <command> Run a command hook manually")
|
||||
print(" load <config> Load hooks from configuration file")
|
||||
return
|
||||
|
||||
command = args[0]
|
||||
|
||||
if command == "list":
|
||||
format_type = "table"
|
||||
if "--format" in args:
|
||||
idx = args.index("--format")
|
||||
if idx + 1 < len(args):
|
||||
format_type = args[idx + 1]
|
||||
handler.list_hooks(format=format_type)
|
||||
|
||||
elif command == "test":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai hooks test <event> [target] [--input <json>]")
|
||||
return
|
||||
|
||||
event = args[1]
|
||||
target = args[2] if len(args) > 2 and not args[2].startswith("--") else None
|
||||
|
||||
input_data = None
|
||||
if "--input" in args:
|
||||
idx = args.index("--input")
|
||||
if idx + 1 < len(args):
|
||||
try:
|
||||
input_data = json.loads(args[idx + 1])
|
||||
except json.JSONDecodeError:
|
||||
print("[red]Invalid JSON input[/red]")
|
||||
return
|
||||
|
||||
handler.test_hook(event, target, input_data)
|
||||
|
||||
elif command == "validate":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai hooks validate <config_path>")
|
||||
return
|
||||
handler.validate_config(args[1])
|
||||
|
||||
elif command == "run":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai hooks run <command> [--event <event>] [--input <json>] [--timeout <seconds>]")
|
||||
return
|
||||
|
||||
cmd = args[1]
|
||||
event = "before_tool"
|
||||
input_json = None
|
||||
timeout = 60.0
|
||||
|
||||
if "--event" in args:
|
||||
idx = args.index("--event")
|
||||
if idx + 1 < len(args):
|
||||
event = args[idx + 1]
|
||||
|
||||
if "--input" in args:
|
||||
idx = args.index("--input")
|
||||
if idx + 1 < len(args):
|
||||
input_json = args[idx + 1]
|
||||
|
||||
if "--timeout" in args:
|
||||
idx = args.index("--timeout")
|
||||
if idx + 1 < len(args):
|
||||
try:
|
||||
timeout = float(args[idx + 1])
|
||||
except ValueError:
|
||||
print("[red]Invalid timeout value[/red]")
|
||||
return
|
||||
|
||||
handler.run_command_hook(cmd, event, input_json, timeout)
|
||||
|
||||
elif command == "load":
|
||||
if len(args) < 2:
|
||||
print("Usage: praisonai hooks load <config_path>")
|
||||
return
|
||||
handler.register_from_config(args[1])
|
||||
|
||||
elif command == "help" or command == "--help":
|
||||
print("Hooks CLI Commands:")
|
||||
print("\n praisonai hooks list [--format json|table]")
|
||||
print(" List all registered hooks")
|
||||
print("\n praisonai hooks test <event> [target] [--input <json>]")
|
||||
print(" Test hooks for a specific event")
|
||||
print(" Events: before_tool, after_tool, before_agent, after_agent, session_start, session_end")
|
||||
print("\n praisonai hooks validate <config_path>")
|
||||
print(" Validate a hooks configuration file (JSON or YAML)")
|
||||
print("\n praisonai hooks run <command> [--event <event>] [--input <json>] [--timeout <seconds>]")
|
||||
print(" Run a command hook manually for testing")
|
||||
print("\n praisonai hooks load <config_path>")
|
||||
print(" Load and register hooks from a configuration file")
|
||||
|
||||
else:
|
||||
print(f"[red]Unknown command: {command}[/red]")
|
||||
print("Use 'praisonai hooks help' for available commands")
|
||||
Loaded 100 of 157 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user