Compare commits

...
Author SHA1 Message Date
Mervin Praison 1d0f08b20a Release v3.1.5 2026-01-07 07:48:57 +00:00
Mervin Praison 07ff93ed7f Release v3.1.4 2026-01-07 07:11:58 +00:00
Mervin Praison 2d5c982893 Release v3.1.3 2026-01-02 09:48:09 +00:00
Mervin Praison f5d4db4262 Release v3.1.2 2026-01-02 08:17:39 +00:00
Mervin Praison 87d691c616 Release v3.1.1 2026-01-02 07:54:02 +00:00
Mervin Praison 8ea15dee0f Release v3.1.0 2026-01-02 07:35:57 +00:00
Mervin Praison 4f56842e08 Release v3.0.9 2026-01-02 07:22:27 +00:00
Mervin Praison a24836aaee Release v3.0.8 2026-01-02 06:59:48 +00:00
Mervin Praison be9ce27615 Release v3.0.7 2026-01-02 06:43:56 +00:00
Mervin Praison 0159a1e082 Release v3.0.6 2026-01-02 06:18:14 +00:00
Mervin Praison 71a5076c07 Release v3.0.5 2026-01-02 06:06:01 +00:00
Mervin Praison 9815a9f8a8 Release v3.0.4 2026-01-02 05:45:18 +00:00
Mervin Praison 2ae44a05b0 Release v3.0.3 2026-01-02 05:17:33 +00:00
Mervin Praison 3386d25022 Release v3.0.2 2026-01-02 04:33:19 +00:00
Mervin Praison 6dca2ace8a feat(tui): queue editing + update_input method
## Queue Editing
- Added update_input method to QueueScheduler for editing queued runs
- Added update_input method to QueueManager
- Added handler for QueuePanelWidget.EditRequested in TUIApp
- Only queued (not running) items can be edited

## Queue-First Architecture Support
- Queue items can be edited before processing
- EditRequested message already existed in QueuePanelWidget

31 tests passing
2026-01-02 04:24:10 +00:00
Mervin Praison 63548739e8 feat(tui): session history persistence with SessionStore
## Session History Implementation
- Created SessionStore class for in-memory chat history persistence
- Added chat_history field to QueuedRun model
- Modified queue worker to inject chat_history into Agent
- Updated TUIApp to store user/assistant messages in SessionStore
- Pass chat_history to queue when submitting messages

## How It Works
1. User sends message -> stored in SessionStore
2. Queue submission includes chat_history from SessionStore
3. Worker injects chat_history into Agent.chat_history
4. Agent has full conversation context
5. Response stored back in SessionStore

## Verified
- Agent remembers previous messages in session
- History format correct for LLM (role + content)
- 31 tests passing
2026-01-02 04:22:26 +00:00
Mervin Praison 3e046537f4 fix(tui): sessions screen loading + update_sessions_screen method
- Call _update_sessions_screen() after pushing sessions screen
- Add _update_sessions_screen() method to load sessions from persistence
- Fallback to showing current session if persistence fails

31 tests passing
2026-01-02 04:02:55 +00:00
Mervin Praison 4a34677567 fix(tui): compose() NoneType error + add missing slash commands
## Fixes
1. ChatWidget.compose() now yields empty generator instead of None
2. Added 7 missing slash commands: /tokens, /plan, /map, /undo, /diff, /commit, /tools

## Commands Added
- /tokens - Show token usage breakdown
- /plan - Show/create task plans
- /map - Show repository map
- /undo - Undo last change
- /diff - Show diff of changes
- /commit - Commit changes
- /tools - Show available tools

31 tests passing
2026-01-02 04:01:25 +00:00
Mervin Praison 19e91cb669 fix(tui): proper scrollbar + non-blocking input
## Scrollbar Fix (VerticalScroll)
- ChatWidget now inherits from VerticalScroll (not ScrollableContainer)
- Messages mount directly to container (removed inner Vertical)
- scroll_end() called after adding messages and during streaming
- CSS has overflow-y: auto for proper scrollbar display

## Non-Blocking Input (@work decorator)
- Message submission uses run_worker() for background execution
- _process_message_submission decorated with @work
- User can type while agent is processing
- Multiple messages can be queued

## Root Causes Fixed
1. Inner Vertical container had overflow:hidden - removed it
2. Message handler was async blocking - now uses worker pattern

Based on Textual documentation:
- VerticalScroll: container with vertical layout and auto Y scrollbar
- @work decorator: runs coroutine in background without blocking UI

31 tests passing
2026-01-02 03:56:56 +00:00
Mervin Praison 5a6125021a fix(tui): add scrollbar to chat + fix session history
## Scrollbar Fix
- ChatWidget now inherits from ScrollableContainer
- Scrollbar visible in chat area (like queue panel)

## Session History Fix
- Pass session_id from TUI app to queue submit
- Pass session_id in agent_config
- Worker now passes session_id to Agent constructor
- Agent will now use session history for context

Flow: TUI.session_id → queue.submit() → worker → Agent(session_id=...)

31 tests passing
2026-01-02 03:44:55 +00:00
Mervin Praison da26887f24 feat(tui): NEWEST-AT-TOP rendering + fix @ typing blockage
## UX STRATEGY CHANGE (Stabilization)

### Newest-at-Top Rendering
- Messages now render NEWEST at TOP
- No scroll dependency - latest always visible
- Streaming updates visible immediately
- ChatWidget inherits from Vertical (not ScrollableContainer)

### @ Typing Not Blocked
- @ handler no longer mounts focus-stealing popup
- Users can freely type @path/to/file
- Manual path entry works without interruption

### Verified Working
- ✅ New messages appear at TOP
- ✅ Latest response always visible
- ✅ Streaming updates visible
- ✅ Can type while agent runs
- ✅ Queue increments immediately
- ✅ Backslash popup works
- ✅ @ does NOT block typing
- ✅ 31 tests passing
2026-01-02 03:38:13 +00:00
Mervin Praison 1b2f7139f9 docs(tui): update help screen with auto-trigger features
- Document @ auto-show file browser
- Document / auto-show command discovery
- Update backslash description to 'auto-show'
2026-01-02 03:25:17 +00:00
Mervin Praison aafbc86180 feat(tui): auto-trigger popups for backslash, @, and / commands
## Auto-Trigger Features (Claude Code-style UX)

### Backslash (\) Auto-Popup
- Typing \ immediately shows command popup
- Clears backslash from input
- No manual trigger needed

### @ File Autocomplete
- Typing @ shows file/directory popup
- Fuzzy search files in workspace
- Supports ~/ and absolute paths
- Reuses FileSearchService from at_mentions.py

### / Command Discovery
- Typing / shows command popup with filter
- Same unified command registry as backslash

## Queue Live Updates
- Queue panel updates IMMEDIATELY when message queued
- Updates on run completion and error
- Reactive queued_count and running_count

## Composer Never Disabled
- Input stays enabled during agent processing
- Users can queue multiple messages

## New Files
- widgets/file_popup.py: FilePopupWidget for @ autocomplete

## Tests
- 31 tests passing (10 new tests)
- TestAutoTriggerEvents: 4 tests
- TestFilePopup: 4 tests
- TestQueueLiveUpdates: 2 tests
2026-01-02 03:24:39 +00:00
Mervin Praison 3d4ca1b7c0 docs(tui): update help screen with safe keyboard defaults and backslash commands
## Help Screen Updates
- Document backslash (\) command popup
- Show safe keyboard defaults (no Ctrl/Fn by default)
- Add command mode (: prefix) section
- Add slash commands (/ prefix) section
- Document optional Ctrl/Fn key env vars

## Keyboard Shortcuts Documented
- \ - Show command popup
- Enter - Send message
- Shift+Enter - New line
- q - Quit (when not typing)
- ? - Help
- : - Command mode
- / - Search mode
- Escape - Cancel/Close
2026-01-02 03:07:50 +00:00
Mervin Praison 1180aa38bd feat(tui): add backslash command popup for Claude Code-style UX
## New Features

### Command Popup Widget (command_popup.py)
- CommandPopupWidget: Overlay showing all available commands
- Triggered by backslash (\) key - matches Claude Code UX
- Searchable/filterable command list
- Keyboard navigation (arrows, Enter, Escape)
- 16 default commands with aliases

### MainScreen Integration
- Backslash binding shows command popup
- Escape dismisses popup or cancels input
- Command selection executes command
- Popup layer CSS for proper overlay positioning

## Commands Available
- help, clear, model, cost, tokens, queue
- cancel, settings, sessions, tools, exit
- plan, undo, diff, commit, map

## Tests
- 21 tests passing
- TestCommandPopup: 4 new tests
- TestMainScreenBackslash: 3 new tests

## Files Changed
- NEW: widgets/command_popup.py
- screens/main.py: Add popup integration
- tests/unit/tui/test_tui_widgets.py: Add tests
2026-01-02 03:06:44 +00:00
Mervin Praison 65f65e5c7b feat(tui): command mode routing and safe keyboard defaults
## Changes

### Command System (: and / prefix)
- Both : and / now work as command prefixes
- Local command handling in MainScreen._handle_local_command()
- Commands: :quit, :clear, :help, :tools, :queue, :settings, :cancel
- Aliases work: :q, :cl, :h, :t, :qu

### Keyboard Defaults
- TUIApp bindings now hidden in footer (show=False)
- Ctrl+Q/C kept for emergency exit but not displayed
- MainScreen shows safe defaults: q, :, ?, Esc

### Help Text Updated
- Shows safe keyboard shortcuts
- Documents command mode usage
- Removed references to Ctrl/Fn keys

## Files Changed
- app.py: Hide Ctrl bindings in footer, update help text
- screens/main.py: Add _handle_local_command() for : commands
- widgets/composer.py: Support both / and : as command prefixes

## Tests
- 14 tests passing
- Command parsing verified
- TUI app startup verified
2026-01-02 02:48:50 +00:00
Mervin Praison f81a12d5e8 feat(tui): comprehensive UX fixes - safe keyboard defaults, auto-scroll, tools panel
BREAKING CHANGE: Keyboard shortcuts changed to safe defaults

## Changes

### 1. Auto-scroll Fix
- Use call_after_refresh() to ensure scroll happens after DOM updates
- Fixed in ChatWidget.add_message() and update_streaming()

### 2. Tools Panel Fix
- Added ScrollableContainer for tools list (no more truncation)
- Shows count of available tools
- All 13 tools now visible with scrolling

### 3. Safe Keyboard Defaults (NO Ctrl/Fn keys by default)
- Single-char shortcuts: q (quit), : (command mode), ? (help)
- Command mode: :quit, :clear, :help, :tools, :queue, :settings
- Ctrl keys (^Q, ^C, ^L) - DISABLED by default (opt-in)
- Function keys (F1-F5) - DISABLED by default (opt-in)

### 4. Configuration System
- New TUIConfig class in config.py
- Environment variables:
  - PRAISONAI_TUI_CTRL_KEYS=1 (enable Ctrl keys)
  - PRAISONAI_TUI_FN_KEYS=1 (enable Fn keys)
  - PRAISONAI_TUI_LEADER=; (custom leader key)

### 5. Non-blocking Input
- Queue accepts messages while agents are running
- Already implemented in previous commit

## Files Changed
- widgets/chat.py - Auto-scroll fix
- widgets/tool_panel.py - ScrollableContainer, show all tools
- widgets/composer.py - set_text() method for command mode
- screens/main.py - Safe default bindings, command mode actions
- config.py - NEW: TUIConfig, command parsing, help text

## Tests
- 14 tests passing
- New tests for TUIConfig and safe bindings
2026-01-02 02:39:16 +00:00
Mervin Praison cf6f5fbefe feat(tui): fix help screen, Enter key, and add TDD tests
Fixes:
- Add HelpScreen for F1 key (was missing, caused KeyError)
- Fix Enter key to send message using custom SubmitTextArea class
- Fix Shift+Enter to insert newline properly
- Register HelpScreen in TUIApp.SCREENS

Tests:
- Add TDD test suite for TUI widgets (12 tests)
- Test Enter key submission
- Test Shift+Enter newline
- Test F1 help screen
- Test Ctrl key bindings

All 13 default tools (ACP, LSP, Basic) load correctly.
Queue system tested and working with real API keys.
2026-01-02 02:23:02 +00:00
Mervin Praison 908d7710db feat(tui): improve TUI usability and add tools panel display
- Enter key now sends message (was Ctrl+Enter)
- Shift+Enter inserts newline without backslash
- Allow typing while response is generating (queue multiple messages)
- Display available tools in Tools panel
- Add edit/delete functionality to queue panel
- Auto-scroll already working in chat widget

TUI Architecture:
- TUIApp -> QueueManager -> WorkerPool -> Agent (from praisonaiagents)
- Tools stored in runtime registry to avoid JSON serialization issues
2026-01-02 02:08:22 +00:00
Mervin Praison 7d5b41e31f fix: use 'llm' parameter instead of 'model' for Agent initialization in queue worker
- Agent class uses 'llm' parameter, not 'model'
- This fixes 'unexpected keyword argument model' error in TUI mode
2026-01-01 19:57:31 +00:00
Mervin Praison e46a0bc9f3 Release v3.0.1 2025-12-31 19:29:02 +00:00
Mervin Praison 911483c92c Release v3.0.0 2025-12-31 18:55:59 +00:00
Mervin Praison 4ffb9d6e3f Release v2.9.2 2025-12-31 17:17:17 +00:00
Mervin Praison b686e5b00f Release v2.9.1 2025-12-31 16:42:54 +00:00
Mervin Praison d8d84ee3e3 Release v2.9.0 2025-12-31 15:49:07 +00:00
Mervin Praison 9f2c4d10b5 Release v2.8.9 2025-12-31 14:05:30 +00:00
Mervin Praison 42a86c831f Release v2.8.8 2025-12-31 12:46:10 +00:00
Mervin Praison 48443b11d6 Release v2.8.7 2025-12-31 12:28:44 +00:00
Mervin Praison 71745aa14a Release v2.8.6 2025-12-31 12:27:20 +00:00
Mervin Praison 4f309ec16d Release v2.8.5 2025-12-31 12:00:29 +00:00
Mervin Praison e48afc8993 Release v2.8.4 2025-12-31 11:39:37 +00:00
Mervin Praison 90950f8889 Release v2.8.3 2025-12-31 11:02:57 +00:00
Mervin Praison 56f36e2593 feat: bump version to 1.4.0 and add AI SDK integration dependencies
- Update package.json version from 1.3.6 to 1.4.0
- Update FEATURE_PARITY_TRACKER.json to v1.3.0 (Phase 5 Complete)
- Add AI SDK wrapper features: generateText, streamText, generateObject, streamObject, generateImage, embed, embedMany
- Add server adapters for http, express, hono, fastify, nest
- Add Next.js integration with createRouteHandler and createPagesHandler
- Add MCP client tools: createMCP, getMCPClient, mcpToolsToAITools
2025-12-31 01:39:39 +00:00
Mervin Praison 7aa67ad27c Release v2.8.2 2025-12-30 23:17:48 +00:00
Mervin Praison 307df443ef Release v2.8.1 2025-12-30 17:47:16 +00:00
Mervin Praison 6434a87b60 chore: remove completed AI SDK verification doc and bump version to 1.3.6
- Remove ai-sdk-backbone-verification.md (all tasks completed)
- Update package.json version from 1.3.5 to 1.3.6
- Refactor agent classes to use lazy-loaded AI SDK backend via resolveBackend()
- Convert direct createProvider() calls to async getProvider() pattern
- Add provider promise caching to prevent duplicate backend initialization
- Update ContextAgent, EnhancedAgent, ImageAgent, and PromptExpanderAgent classes
2025-12-30 17:15:40 +00:00
Mervin Praison 97d5b6db0d feat: add AI SDK backend support and embedding capabilities to Agent class
- Add AI SDK backend resolution with lazy initialization via getBackend()
- Add embed() method supporting single/batch text embeddings
- Add getBackendSource() to check backend type (ai-sdk/native/custom/legacy)
- Add getModel() to retrieve agent's model string
- Parse model string to extract provider/model ID for backend selection
- Use AI SDK backend for non-OpenAI providers (streaming and non-streaming)
- Export backend resolver an
2025-12-30 16:52:46 +00:00
Mervin Praison 4996a9f9fa Merge remote-tracking branch 'origin/pr-1097-fixes' 2025-12-30 16:15:49 +00:00
Mervin Praison 54f4a96cba test: add comprehensive ImportError tests for MCP optional import
- Add test for MCP class ImportError when mcp package unavailable
- Add test for SSEMCPClient ImportError when mcp unavailable
- Add test for HTTPStreamMCPClient ImportError when mcp unavailable
- Add test for HTTPStreamMCPClient ImportError when aiohttp unavailable
- Fix unused import lint warnings (remove sys, add noqa for MCP)

These tests verify that actionable error messages with install instructions
are shown when users try to use MCP functionality without the mcp package.
2025-12-30 16:14:15 +00:00
Mervin Praison e59be8f0fd Merge pull request #1096 from majiayu000/fix/issue-1091-mcp-optional-import
fix: make MCP package import optional to prevent errors when not installed
2025-12-30 16:09:06 +00:00
Mervin Praison b52708f3f6 Release v1.3.4 2025-12-30 15:10:46 +00:00
Mervin Praison ba37ce39d3 Release v2.8.0 2025-12-30 15:08:22 +00:00
majiayu000 98c5e8f175 fix: make MCP package import optional to prevent errors when not installed
Wrap mcp imports in try/except blocks in mcp.py, mcp_sse.py, and
mcp_http_stream.py. Add MCP_AVAILABLE flag to check availability at runtime.
When MCP class is instantiated without the mcp package installed, raise a
clear ImportError with installation instructions.

Closes #1091

Signed-off-by: majiayu000 <[email protected]>
2025-12-30 19:53:45 +08:00
402 changed files with 77768 additions and 1376 deletions

No files matched your search

+71
View File
@@ -0,0 +1,71 @@
"""
Auto-generated API server for PraisonAI agents.
"""
from flask import Flask, request, jsonify
from flask_cors import CORS
from praisonai import PraisonAI
import os
app = Flask(__name__)
# CORS configuration
CORS(app)
# Authentication
AUTH_ENABLED = False
AUTH_TOKEN = None
def check_auth():
"""Check authentication if enabled."""
if not AUTH_ENABLED:
return True
token = request.headers.get('Authorization', '').replace('Bearer ', '')
return token == AUTH_TOKEN
@app.route('/health', methods=['GET'])
def health():
"""Health check endpoint."""
return jsonify({"status": "ok", "service": "praisonai-api"})
@app.route('/chat', methods=['POST'])
def chat():
"""Chat endpoint for agent interaction."""
if not check_auth():
return jsonify({"error": "Unauthorized"}), 401
data = request.get_json()
if not data or 'message' not in data:
return jsonify({"error": "Message required"}), 400
try:
praisonai = PraisonAI(agent_file="agents.yaml")
result = praisonai.run()
return jsonify({
"response": result,
"status": "success"
})
except Exception as e:
return jsonify({
"error": str(e),
"status": "error"
}), 500
@app.route('/agents', methods=['GET'])
def list_agents():
"""List available agents."""
if not check_auth():
return jsonify({"error": "Unauthorized"}), 401
return jsonify({
"agents": ["default"],
"agent_file": "agents.yaml"
})
if __name__ == '__main__':
app.run(
host='0.0.0.0',
port=8080,
debug=False
)
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.7.9" \
"praisonai>=3.1.5" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.7.9" \
"praisonai>=3.1.5" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.7.9" \
"praisonai>=3.1.5" \
"praisonai[ui]" \
"praisonai[crewai]"
+203
View File
@@ -0,0 +1,203 @@
# Profiling CLI
PraisonAI provides unified profiling capabilities for performance analysis and debugging, with **full visibility from ENTER to first response**.
## Quick Start
```bash
# Basic profiling
praisonai "What is 2+2?" --profile
# Deep profiling with call graph, decision trace, and module breakdown
praisonai "What is 2+2?" --profile --profile-deep
# JSON output format
praisonai "What is 2+2?" --profile --profile-format json
```
## Time to First Response (TTFR)
**TTFR** is the key user-perceived latency metric. It measures the time from when you press ENTER until the first visible output appears.
### What contributes to TTFR?
1. **CLI Entry** - Time to reach the CLI entry point
2. **CLI Parse** - Argparse/Typer routing time
3. **Routing** - Command routing to handler
4. **Imports** - Loading required modules
5. **Agent Init** - Initializing the agent
6. **Network Request** - Sending request to LLM API
7. **First Token** - Receiving first token (streaming)
8. **First Output** - Rendering first visible output
### How to reduce TTFR
1. **Warm runs** - Imports are cached after first run
2. **Streaming mode** - See output as tokens arrive
3. **Lighter models** - Faster response times
4. **Local models** - Eliminate network latency
## CLI Flags
| Flag | Description | Default |
|------|-------------|---------|
| `--profile` | Enable profiling with timing breakdown | Disabled |
| `--profile-deep` | Enable deep profiling with call graph (higher overhead) | Disabled |
| `--profile-format` | Output format: `text` or `json` | `text` |
## Profile Command
For more detailed profiling options, use the `profile` subcommand:
```bash
# Profile a query
praisonai profile query "What is 2+2?"
# With deep call tracing
praisonai profile query "Hello" --deep
# Limit functions shown
praisonai profile query "Test" --limit 20
# Save artifacts to directory
praisonai profile query "Test" --save ./profile_results
# JSON output
praisonai profile query "Test" --format json
```
### Profile Query Options
| Option | Description | Default |
|--------|-------------|---------|
| `--model, -m` | Model to use | Default model |
| `--stream/--no-stream` | Use streaming mode | No stream |
| `--deep` | Enable deep call tracing | Disabled |
| `--limit, -n` | Top N functions to show | 30 |
| `--sort, -s` | Sort by: `cumulative` or `tottime` | `cumulative` |
| `--show-callers` | Show caller functions | Disabled |
| `--show-callees` | Show callee functions | Disabled |
| `--save` | Save artifacts to path | None |
| `--format, -f` | Output format: `text` or `json` | `text` |
## Other Profile Commands
```bash
# Profile module import times
praisonai profile imports
# Profile CLI startup time
praisonai profile startup
# Run comprehensive profiling suite
praisonai profile suite
# Create performance snapshots for comparison
praisonai profile snapshot --baseline
praisonai profile snapshot current --compare
# Configure performance optimizations
praisonai profile optimize --show
```
## Output Format
### Text Report
```
======================================================================
PraisonAI Profile Report
======================================================================
Run ID: abc12345
Timestamp: 2024-01-02T12:00:00.000000Z
Method: cli_direct
Version: 3.0.1
## Execution Timeline
---------------------------------------------
Imports : 500.00 ms
Agent Init : 0.50 ms
Execution : 2000.00 ms
───────────────────────────────────────────
⏱ Time to First Response : 2500.50 ms
TOTAL : 2500.50 ms
## Timing Breakdown
----------------------------------------
Imports: 500.00 ms
Agent Init: 0.50 ms
Execution: 2000.00 ms
─────────────────────────────────────
TOTAL: 2500.50 ms
## Decision Trace (deep profile only)
----------------------------------------
Agent Config: default
Model: gpt-4
Streaming: disabled
Profile Layer: 2
## Top Functions by Cumulative Time
----------------------------------------------------------------------
Function Calls Cumul (ms)
----------------------------------------------------------------------
request 1 1800.00
send 1 1750.00
...
```
### JSON Report
```json
{
"schema_version": "1.0",
"run_id": "abc12345",
"timestamp": "2024-01-02T12:00:00.000000Z",
"invocation": {
"method": "cli_direct",
"flags": {"layer": 1, "stream": false, "model": null},
"praisonai_version": "3.0.1",
"python_version": "3.12.0"
},
"timing": {
"total_ms": 2500.5,
"imports_ms": 500.0,
"agent_init_ms": 0.5,
"execution_ms": 2000.0
},
"functions": [...],
"response_preview": "four"
}
```
## Profiling Layers
| Layer | Overhead | Features |
|-------|----------|----------|
| 0 | < 1ms | Wall-clock phases only |
| 1 | ~5% | + cProfile function stats |
| 2 | ~15% | + Call graph (callers/callees) |
- `--profile` uses Layer 1
- `--profile --profile-deep` uses Layer 2
## Performance Tips
1. **Use Layer 0 for production monitoring** - minimal overhead
2. **Use Layer 1 for development profiling** - good balance
3. **Use Layer 2 only for deep debugging** - higher overhead
## Observational Equivalence
The following commands produce equivalent profile reports:
```bash
# CLI direct
praisonai "prompt" --profile
# Profile command
praisonai profile query "prompt"
```
Both use the same unified profiling architecture and produce identical schema output.
+257
View File
@@ -0,0 +1,257 @@
# Profiling Module
PraisonAI provides a unified profiling module for programmatic performance analysis.
## Quick Start
```python
from praisonai.cli.execution import (
ExecutionRequest,
Profiler,
ProfilerConfig,
)
# Create execution request
request = ExecutionRequest(prompt="What is 2+2?")
# Configure profiler
config = ProfilerConfig(layer=1)
# Run with profiling
profiler = Profiler(config)
result, report = profiler.profile_sync(request)
# Access results
print(f"Output: {result.output}")
print(f"Total time: {report.timing.total_ms}ms")
```
## ExecutionRequest
Immutable request descriptor for execution.
```python
from praisonai.cli.execution import ExecutionRequest
# Basic request
request = ExecutionRequest(prompt="Hello")
# With model
request = ExecutionRequest(
prompt="Hello",
model="gpt-4",
agent_name="MyAgent",
)
# With streaming
request = ExecutionRequest(
prompt="Hello",
stream=True,
)
```
### Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `prompt` | str | Required | The prompt to execute |
| `agent_name` | str | "Assistant" | Name of the agent |
| `agent_instructions` | str | None | Custom instructions |
| `model` | str | None | Model to use |
| `stream` | bool | False | Enable streaming |
| `tools` | tuple | () | Tools to provide |
## ProfilerConfig
Configuration for profiling behavior.
```python
from praisonai.cli.execution import ProfilerConfig
# Layer 0: Wall-clock only (< 1ms overhead)
config = ProfilerConfig(layer=0)
# Layer 1: + cProfile stats (~5% overhead)
config = ProfilerConfig(layer=1)
# Layer 2: + Call graph (~15% overhead)
config = ProfilerConfig(
layer=2,
show_callers=True,
show_callees=True,
)
# From CLI flags
config = ProfilerConfig.from_flags(
profile=True,
deep=True,
network=False,
)
```
### Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `layer` | int | 1 | Profiling layer (0, 1, or 2) |
| `limit` | int | 30 | Max functions to show |
| `sort_by` | str | "cumulative" | Sort by: cumulative, time, calls |
| `show_callers` | bool | False | Show caller functions |
| `show_callees` | bool | False | Show callee functions |
| `track_network` | bool | False | Track network requests |
| `output_format` | str | "text" | Output format: text, json |
| `save_path` | str | None | Path to save artifacts |
## Profiler
The main profiler class that wraps execution.
```python
from praisonai.cli.execution import Profiler, ProfilerConfig, ExecutionRequest
config = ProfilerConfig(layer=1)
profiler = Profiler(config)
request = ExecutionRequest(prompt="Hello")
result, report = profiler.profile_sync(request, invocation_method="api")
```
### Methods
#### `profile_sync(request, invocation_method="profile_command")`
Profile synchronous execution.
**Parameters:**
- `request`: ExecutionRequest - The execution request
- `invocation_method`: str - How this was invoked (for report)
**Returns:** Tuple of (ExecutionResult, ProfileReport)
## ProfileReport
The canonical profile report with timing and statistics.
```python
# Access timing breakdown
print(f"Total: {report.timing.total_ms}ms")
print(f"Imports: {report.timing.imports_ms}ms")
print(f"Agent Init: {report.timing.agent_init_ms}ms")
print(f"Execution: {report.timing.execution_ms}ms")
# Access function stats (layer 1+)
if report.functions:
for func in report.functions[:5]:
print(f"{func.name}: {func.cumulative_time_ms}ms")
# Access call graph (layer 2)
if report.call_graph:
print(f"Edges: {len(report.call_graph.edges)}")
# Export formats
json_str = report.to_json()
text_str = report.to_text()
dict_data = report.to_dict()
```
### Schema Version
The report uses schema version `1.0`. Schema is frozen after release - new fields are always optional with defaults.
### Mandatory Fields
- `schema_version`: str
- `run_id`: str
- `timestamp`: str (ISO 8601)
- `invocation`: InvocationInfo
- `timing`: TimingBreakdown
- `response_preview`: str
### Optional Fields
- `functions`: List[FunctionStat] (layer 1+)
- `call_graph`: CallGraph (layer 2)
- `network`: List[RequestTiming] (if track_network)
## ExecutionResult
Result of execution with output and metadata.
```python
# Access output
print(result.output)
print(result.success)
print(result.run_id)
# Check for errors
if result.error:
print(f"Error: {result.error}")
# Duration
print(f"Duration: {result.duration_ms}ms")
```
## Data Bounds
Profile data is strictly bounded to prevent memory issues:
| Data Type | Maximum |
|-----------|---------|
| Function stats | 1000 |
| Call graph edges | 5000 |
| Network requests | 100 |
## Example: Full Profiling Workflow
```python
from praisonai.cli.execution import (
ExecutionRequest,
Profiler,
ProfilerConfig,
)
# Configure deep profiling
config = ProfilerConfig(
layer=2,
limit=50,
show_callers=True,
show_callees=True,
)
# Create request
request = ExecutionRequest(
prompt="Explain quantum computing in simple terms",
model="gpt-4",
)
# Profile execution
profiler = Profiler(config)
result, report = profiler.profile_sync(request)
# Print text report
print(report.to_text())
# Save JSON report
with open("profile.json", "w") as f:
f.write(report.to_json())
# Analyze timing
print(f"\nTiming Analysis:")
print(f" Import overhead: {report.timing.imports_ms / report.timing.total_ms * 100:.1f}%")
print(f" Execution time: {report.timing.execution_ms / report.timing.total_ms * 100:.1f}%")
# Top 5 hotspots
if report.functions:
print(f"\nTop 5 Hotspots:")
for func in report.functions[:5]:
print(f" {func.name}: {func.cumulative_time_ms:.2f}ms ({func.calls} calls)")
```
## Invariants
The profiling module guarantees:
1. **Output Equivalence**: Profiled execution produces identical output to non-profiled execution
2. **Schema Stability**: Schema version 1.0 is frozen; new fields are always optional
3. **Bounded Data**: All profile data respects maximum bounds
4. **Zero Overhead When Disabled**: No profiling code runs unless explicitly enabled
+137
View File
@@ -0,0 +1,137 @@
# ACP (Agent Client Protocol) Examples
This directory contains examples for using PraisonAI with the Agent Client Protocol (ACP).
## What is ACP?
The Agent Client Protocol (ACP) is a standardized JSON-RPC 2.0 protocol that allows code editors and IDEs to communicate with AI coding agents. It enables seamless integration between PraisonAI and editors like:
- **Zed**
- **JetBrains IDEs** (IntelliJ, PyCharm, WebStorm, etc.)
- **VSCode** (via extensions)
- **Toad**
## Quick Start
### CLI Usage
```bash
# Start ACP server with defaults
praisonai acp
# Enable debug logging
praisonai acp --debug
# Allow file writes
praisonai acp --allow-write
# Use a specific model
praisonai acp -m gpt-4o
# Resume last session
praisonai acp --resume --last
```
### Python API
```python
from praisonai.acp import serve
serve(
workspace=".",
debug=True,
allow_write=True,
)
```
## Editor Configuration
### Zed
Add to `~/.config/zed/settings.json`:
```json
{
"agent_servers": {
"PraisonAI": {
"command": "praisonai",
"args": ["acp"],
"env": {}
}
}
}
```
### JetBrains
Add to `~/.jetbrains/acp.json`:
```json
{
"agent_servers": {
"PraisonAI": {
"command": "praisonai",
"args": ["acp"],
"env": {}
}
}
}
```
### Toad
```bash
toad acp "praisonai acp"
```
## Examples
- `basic_acp_server.py` - Simple ACP server with default settings
- `custom_agent_acp.py` - ACP server with a custom agent
## CLI Options
| Option | Description |
|--------|-------------|
| `-w, --workspace` | Workspace root directory |
| `-a, --agent` | Agent name or config file |
| `-m, --model` | LLM model to use |
| `-r, --resume` | Resume session by ID |
| `--last` | Resume the last session |
| `--approve` | Approval mode: manual, auto, scoped |
| `--read-only` | Read-only mode (default) |
| `--allow-write` | Allow file writes |
| `--allow-shell` | Allow shell commands |
| `--debug` | Enable debug logging |
## Installation
```bash
pip install praisonai[acp]
```
## Protocol Details
ACP uses JSON-RPC 2.0 over stdio:
- **stdin**: Receives JSON-RPC requests from the client
- **stdout**: Sends JSON-RPC responses to the client
- **stderr**: Debug/log output (never pollutes stdout)
### Supported Methods
**Agent Methods (client → agent):**
- `initialize` - Negotiate protocol version and capabilities
- `authenticate` - Optional authentication
- `session/new` - Create a new session
- `session/load` - Load an existing session
- `session/prompt` - Send user message
- `session/cancel` - Cancel ongoing operations
- `session/set_mode` - Change operating mode
**Client Methods (agent → client):**
- `session/update` - Send progress updates
- `session/request_permission` - Request user approval
- `fs/read_text_file` - Read file contents
- `fs/write_text_file` - Write file contents
- `terminal/create` - Create terminal
- `terminal/output` - Get terminal output
+27
View File
@@ -0,0 +1,27 @@
#!/usr/bin/env python3
"""
Basic ACP Server Example
This example demonstrates how to run PraisonAI as an ACP server
that can be connected to by IDEs like Zed, JetBrains, VSCode, or Toad.
Usage:
# Run from command line
praisonai acp
# Or run this script directly
python basic_acp_server.py
"""
from praisonai.acp import serve
if __name__ == "__main__":
# Start ACP server with default settings
# This will listen on stdin/stdout for JSON-RPC messages
serve(
workspace=".", # Current directory as workspace
agent="default", # Use default agent
debug=False, # Set to True for debug logging to stderr
read_only=True, # Safe by default
approval_mode="manual", # Require approval for actions
)
+45
View File
@@ -0,0 +1,45 @@
#!/usr/bin/env python3
"""
Custom Agent ACP Server Example
This example shows how to run a custom PraisonAI agent as an ACP server.
Usage:
python custom_agent_acp.py
"""
from praisonai.acp import serve, ACPServer, ACPConfig
from praisonaiagents import Agent
# Create a custom agent
agent = Agent(
name="CodeAssistant",
instructions="""You are an expert coding assistant. You help users:
- Write clean, efficient code
- Debug issues
- Explain complex concepts
- Review code for best practices
Always provide clear explanations with your code suggestions.""",
model="gpt-4o-mini",
)
if __name__ == "__main__":
# Create config
config = ACPConfig(
workspace=".",
debug=True,
allow_write=True, # Allow file modifications
approval_mode="manual",
)
# Create server with custom agent
server = ACPServer(config=config, agent=agent)
# Run the server
serve(
workspace=".",
debug=True,
allow_write=True,
approval_mode="manual",
)
@@ -0,0 +1,135 @@
"""
Agent-Centric Tools Example
This example demonstrates how to use LSP/ACP-powered tools that make
the Agent the central orchestrator for file operations.
How to run:
export OPENAI_API_KEY=your_key
python examples/agent_tools/agent_centric_example.py
Prerequisites:
pip install praisonai praisonaiagents
"""
import asyncio
import os
import sys
import tempfile
import shutil
# Add paths for development
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai-agents'))
async def main():
print("=" * 60)
print("AGENT-CENTRIC TOOLS EXAMPLE")
print("=" * 60)
print()
# Create temporary workspace
workspace = tempfile.mkdtemp(prefix="praisonai_example_")
print(f"Workspace: {workspace}")
print()
try:
# Import after path setup
from praisonai.cli.features import (
create_agent_centric_tools,
InteractiveRuntime,
RuntimeConfig
)
from praisonaiagents import Agent
# 1. Create runtime with ACP enabled
print("1. Creating InteractiveRuntime with ACP...")
config = RuntimeConfig(
workspace=workspace,
lsp_enabled=False, # Skip LSP for this example
acp_enabled=True,
approval_mode="auto" # Auto-approve for demo
)
runtime = InteractiveRuntime(config)
status = await runtime.start()
print(f" ACP ready: {runtime.acp_ready}")
print(f" Read-only: {runtime.read_only}")
print()
# 2. Create agent-centric tools
print("2. Creating agent-centric tools...")
tools = create_agent_centric_tools(runtime)
print(f" Tools created: {len(tools)}")
for tool in tools:
print(f" - {tool.__name__}")
print()
# 3. Create Agent with ACP-powered tools
print("3. Creating Agent with ACP-powered tools...")
agent = Agent(
name="FileAgent",
instructions="""You help create and manage files.
Use acp_create_file to create files.
Use read_file to read files.
Use list_files to list directory contents.""",
tools=tools,
verbose=True
)
print()
# 4. Ask agent to create a file
print("4. Asking agent to create a Python file...")
print("-" * 40)
result = agent.start(
"Create a Python file called calculator.py with add and subtract functions"
)
print("-" * 40)
print(f"Agent response: {result[:200]}..." if len(str(result)) > 200 else f"Agent response: {result}")
print()
# 5. Verify file was created
print("5. Verifying file creation...")
calc_path = os.path.join(workspace, "calculator.py")
if os.path.exists(calc_path):
print(f" ✓ File created at {calc_path}")
with open(calc_path) as f:
content = f.read()
print(f" Content preview:")
for line in content.split('\n')[:10]:
print(f" {line}")
else:
# Check what files exist
files = os.listdir(workspace)
print(f" Files in workspace: {files}")
for f in files:
fp = os.path.join(workspace, f)
if os.path.isfile(fp):
print(f" ✓ Found: {f}")
print()
# 6. Ask agent to read the file
print("6. Asking agent to read the file...")
print("-" * 40)
result2 = agent.start("Read the calculator.py file and tell me what functions it has")
print("-" * 40)
print(f"Agent response: {result2[:300]}..." if len(str(result2)) > 300 else f"Agent response: {result2}")
print()
# Cleanup
await runtime.stop()
finally:
# Clean up workspace
shutil.rmtree(workspace, ignore_errors=True)
print(f"Cleaned up workspace: {workspace}")
print()
print("=" * 60)
print("EXAMPLE COMPLETE")
print("=" * 60)
if __name__ == "__main__":
asyncio.run(main())
+146
View File
@@ -0,0 +1,146 @@
#!/usr/bin/env python3
"""
PraisonAI Benchmark Example
This example demonstrates how to use the benchmark system programmatically
to compare performance across different execution paths.
Usage:
python benchmark_example.py
Requirements:
- OPENAI_API_KEY environment variable set
- praisonai package installed
"""
import os
import sys
# Ensure the package is importable
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
from praisonai.cli.features.benchmark import BenchmarkHandler, BenchmarkReport
def run_quick_comparison():
"""Run a quick comparison of key execution paths."""
print("=" * 60)
print("Quick Benchmark Comparison")
print("=" * 60)
handler = BenchmarkHandler()
# Run benchmark with 2 iterations on key paths
report = handler.run_full_benchmark(
prompt="What is 2+2?",
iterations=2,
paths=["openai_sdk", "praisonai_agent"],
verbose=True
)
# Print comparison table
print("\n" + handler.create_comparison_table(report))
return report
def run_full_benchmark():
"""Run full benchmark suite across all paths."""
print("=" * 60)
print("Full Benchmark Suite")
print("=" * 60)
handler = BenchmarkHandler()
# Run full benchmark
report = handler.run_full_benchmark(
prompt="Hi",
iterations=3,
verbose=True
)
# Print full report
handler.print_report(report)
return report
def run_agent_benchmark():
"""Benchmark PraisonAI Agent specifically."""
print("=" * 60)
print("Agent Benchmark")
print("=" * 60)
handler = BenchmarkHandler()
# Benchmark agent vs SDK
report = handler.run_full_benchmark(
prompt="Explain Python in one sentence",
iterations=3,
paths=["openai_sdk", "praisonai_agent"],
verbose=True
)
# Show timeline diagrams
for name, result in report.results.items():
print(f"\n### {name}")
print(handler.create_timeline_diagram(result))
# Show variance analysis
print(handler.create_variance_table(report))
return report
def save_benchmark_results(report: BenchmarkReport, output_path: str):
"""Save benchmark results to JSON file."""
import json
with open(output_path, 'w') as f:
json.dump(report.to_dict(), f, indent=2)
print(f"\nResults saved to: {output_path}")
def main():
"""Main entry point."""
import argparse
parser = argparse.ArgumentParser(description="PraisonAI Benchmark Example")
parser.add_argument(
"--mode",
choices=["quick", "full", "agent"],
default="quick",
help="Benchmark mode: quick (2 paths), full (all paths), agent (agent focus)"
)
parser.add_argument(
"--output",
type=str,
default=None,
help="Output JSON file path"
)
args = parser.parse_args()
# Check for API key
if not os.environ.get("OPENAI_API_KEY"):
print("Error: OPENAI_API_KEY environment variable not set")
sys.exit(1)
# Run selected benchmark
if args.mode == "quick":
report = run_quick_comparison()
elif args.mode == "full":
report = run_full_benchmark()
elif args.mode == "agent":
report = run_agent_benchmark()
# Save results if requested
if args.output:
save_benchmark_results(report, args.output)
print("\n✅ Benchmark complete!")
if __name__ == "__main__":
main()
+228
View File
@@ -0,0 +1,228 @@
"""
Debug Interactive Mode Example
This example demonstrates how to use the debug CLI commands to test
the Interactive Coding Assistant without entering TUI mode.
Features demonstrated:
- LSP symbol listing
- LSP definition lookup
- ACP action planning
- ACP action execution
- Trace capture and JSON output
"""
import asyncio
import json
import sys
from pathlib import Path
# Add the src directory to path for development
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src" / "praisonai"))
from praisonai.cli.features.interactive_runtime import create_runtime
from praisonai.cli.features.code_intelligence import CodeIntelligenceRouter
from praisonai.cli.features.action_orchestrator import ActionOrchestrator
async def demo_lsp_queries(workspace: str):
"""Demonstrate LSP-based code queries."""
print("\n" + "=" * 60)
print("LSP CODE QUERIES DEMO")
print("=" * 60)
runtime = create_runtime(
workspace=workspace,
lsp=True,
acp=False,
trace=True
)
try:
await runtime.start()
router = CodeIntelligenceRouter(runtime)
# Query 1: List symbols
print("\n[Query 1] List all symbols in app.py")
result = await router.handle_query("list all functions in app.py", "app.py")
print(f" Success: {result.success}")
print(f" LSP Used: {result.lsp_used}")
print(f" Fallback Used: {result.fallback_used}")
if result.data:
print(f" Symbols found: {len(result.data)}")
for symbol in result.data[:5]:
print(f" - {symbol['name']} ({symbol['kind']}) at line {symbol['line']}")
# Query 2: Find definition
print("\n[Query 2] Find definition of Calculator")
result = await router.handle_query("where is Calculator defined", "app.py")
print(f" Success: {result.success}")
if result.citations:
print(f" Found at: {result.citations[0]}")
# Get trace
trace = runtime.get_trace()
print(f"\n[Trace] {len(trace.entries)} entries captured")
finally:
await runtime.stop()
async def demo_acp_actions(workspace: str):
"""Demonstrate ACP-based code modifications."""
print("\n" + "=" * 60)
print("ACP CODE MODIFICATIONS DEMO")
print("=" * 60)
runtime = create_runtime(
workspace=workspace,
lsp=False,
acp=True,
approval="auto",
trace=True
)
try:
await runtime.start()
print(f"\n ACP Ready: {runtime.acp_ready}")
print(f" Read Only: {runtime.read_only}")
orchestrator = ActionOrchestrator(runtime)
# Action 1: Create a plan
print("\n[Action 1] Create plan for new file")
result = await orchestrator.create_plan("create a new file called demo_output.py")
print(f" Success: {result.success}")
if result.plan:
print(f" Plan ID: {result.plan.id}")
print(f" Steps: {len(result.plan.steps)}")
for step in result.plan.steps:
print(f" - {step.action_type.value}: {step.target}")
# Action 2: Execute with auto-approve
print("\n[Action 2] Execute file creation")
exec_result = await orchestrator.execute(
"create a new file called demo_output.py with a hello function",
auto_approve=True
)
print(f" Success: {exec_result.success}")
print(f" Applied Steps: {exec_result.applied_steps}")
# Verify file was created
output_file = Path(workspace) / "demo_output.py"
if output_file.exists():
print(f" File created: {output_file}")
print(f" Content preview: {output_file.read_text()[:100]}...")
# Get trace
trace = runtime.get_trace()
print(f"\n[Trace] {len(trace.entries)} entries captured")
for entry in trace.entries:
print(f" [{entry.category}] {entry.action}")
finally:
await runtime.stop()
async def demo_full_interactive_turn(workspace: str):
"""Demonstrate full interactive turn simulation."""
print("\n" + "=" * 60)
print("FULL INTERACTIVE TURN SIMULATION")
print("=" * 60)
runtime = create_runtime(
workspace=workspace,
lsp=True,
acp=True,
approval="auto",
trace=True
)
try:
status = await runtime.start()
print("\n Runtime Status:")
print(f" LSP: {'ready' if status['lsp']['ready'] else 'fallback'}")
print(f" ACP: {'ready' if status['acp']['ready'] else 'unavailable'}")
router = CodeIntelligenceRouter(runtime)
orchestrator = ActionOrchestrator(runtime)
# Simulate user prompt: code query
prompt1 = "What functions are in app.py?"
print(f"\n[Turn 1] User: {prompt1}")
# Classify and route
result = await router.handle_query(prompt1, "app.py")
print(f" Intent: {result.intent.value}")
print(f" Response: Found {len(result.data) if result.data else 0} symbols")
# Simulate user prompt: code modification
prompt2 = "Create a new test file"
print(f"\n[Turn 2] User: {prompt2}")
exec_result = await orchestrator.execute(
"create a new file called test_demo.py",
auto_approve=True
)
print(" Action: file_create")
print(f" Result: {'success' if exec_result.success else 'failed'}")
# Export trace as JSON
trace = runtime.get_trace()
trace_json = trace.to_dict()
print("[Trace Export]")
print(json.dumps(trace_json, indent=2, default=str)[:500] + "...")
finally:
await runtime.stop()
def main():
"""Run the debug demo."""
# Create a temporary workspace
workspace = Path(__file__).parent / "_demo_workspace"
workspace.mkdir(exist_ok=True)
# Create a sample app.py
app_file = workspace / "app.py"
app_file.write_text('''"""Sample application."""
def greet(name: str) -> str:
"""Greet someone."""
return f"Hello, {name}!"
def add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
class Calculator:
"""Simple calculator."""
def __init__(self):
self.value = 0
def add(self, x):
self.value += x
return self.value
if __name__ == "__main__":
print(greet("World"))
''')
print("=" * 60)
print("PRAISONAI DEBUG MODE DEMO")
print("=" * 60)
print(f"Workspace: {workspace}")
# Run demos
asyncio.run(demo_lsp_queries(str(workspace)))
asyncio.run(demo_acp_actions(str(workspace)))
asyncio.run(demo_full_interactive_turn(str(workspace)))
print("\n" + "=" * 60)
print("DEMO COMPLETE")
print("=" * 60)
if __name__ == "__main__":
main()
@@ -0,0 +1,147 @@
"""
Interactive Default Tools Example
This example demonstrates how PraisonAI's interactive modes (TUI and prompt)
include ACP and LSP tools by default.
Tool Groups:
- ACP: acp_create_file, acp_edit_file, acp_delete_file, acp_execute_command
- LSP: lsp_list_symbols, lsp_find_definition, lsp_find_references, lsp_get_diagnostics
- Basic: read_file, write_file, list_files, execute_command, internet_search
Usage:
python interactive_tools_example.py
"""
import os
import sys
# Add paths for development
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai-agents'))
def example_get_default_tools():
"""Example: Get all default interactive tools."""
from praisonai.cli.features.interactive_tools import (
get_interactive_tools,
ToolConfig,
TOOL_GROUPS,
)
print("=== Default Interactive Tools ===\n")
# Show tool groups
print("Tool Groups:")
for group, tools in TOOL_GROUPS.items():
if group != "interactive":
print(f" {group}: {tools}")
print()
# Get all tools
config = ToolConfig(workspace=os.getcwd())
tools = get_interactive_tools(config=config)
print(f"Total tools loaded: {len(tools)}")
print(f"Tool names: {[t.__name__ for t in tools]}")
print()
return tools
def example_disable_acp():
"""Example: Disable ACP tools."""
from praisonai.cli.features.interactive_tools import (
get_interactive_tools,
ToolConfig,
)
print("=== Disable ACP Tools (--no-acp) ===\n")
config = ToolConfig(workspace=os.getcwd(), enable_acp=False)
tools = get_interactive_tools(config=config, disable=["acp"])
acp_tools = [t for t in tools if t.__name__.startswith("acp_")]
print(f"ACP tools after --no-acp: {len(acp_tools)}")
print(f"Remaining tools: {[t.__name__ for t in tools]}")
print()
return tools
def example_disable_lsp():
"""Example: Disable LSP tools."""
from praisonai.cli.features.interactive_tools import (
get_interactive_tools,
ToolConfig,
)
print("=== Disable LSP Tools (--no-lsp) ===\n")
config = ToolConfig(workspace=os.getcwd(), enable_lsp=False)
tools = get_interactive_tools(config=config, disable=["lsp"])
lsp_tools = [t for t in tools if t.__name__.startswith("lsp_")]
print(f"LSP tools after --no-lsp: {len(lsp_tools)}")
print(f"Remaining tools: {[t.__name__ for t in tools]}")
print()
return tools
def example_agent_with_default_tools():
"""Example: Create an Agent with default interactive tools."""
from praisonai.cli.features.interactive_tools import get_interactive_tools, ToolConfig
from praisonaiagents import Agent
print("=== Agent with Default Interactive Tools ===\n")
# Get default tools
config = ToolConfig(workspace=os.getcwd())
tools = get_interactive_tools(config=config)
# Create agent with tools
agent = Agent(
name="InteractiveAgent",
instructions="You are a helpful coding assistant with file and code analysis capabilities.",
tools=tools,
verbose=False
)
print(f"Agent created with {len(tools)} tools")
print(f"Tools: {[t.__name__ for t in tools]}")
print()
# Simple test - list files
print("Testing agent with list_files tool...")
result = agent.start("List the Python files in the current directory using list_files tool.")
print(f"Result: {result[:200]}..." if len(str(result)) > 200 else f"Result: {result}")
print()
return agent
if __name__ == "__main__":
print("=" * 60)
print("PraisonAI Interactive Default Tools Example")
print("=" * 60)
print()
# Example 1: Get default tools
example_get_default_tools()
# Example 2: Disable ACP
example_disable_acp()
# Example 3: Disable LSP
example_disable_lsp()
# Example 4: Agent with default tools (requires API key)
if os.environ.get("OPENAI_API_KEY"):
example_agent_with_default_tools()
else:
print("=== Agent Example Skipped (no OPENAI_API_KEY) ===")
print("=" * 60)
print("Examples Complete")
print("=" * 60)
+41
View File
@@ -0,0 +1,41 @@
# AI SDK Examples
Examples demonstrating the AI SDK integration in PraisonAI TypeScript.
## Prerequisites
```bash
npm install praisonai ai @ai-sdk/openai @ai-sdk/anthropic @ai-sdk/google zod
```
## Environment Variables
```bash
export OPENAI_API_KEY=sk-...
export ANTHROPIC_API_KEY=sk-ant-...
export GOOGLE_API_KEY=AIza...
```
## Examples
| File | Description |
|------|-------------|
| `basic-chat.ts` | Simple text generation |
| `streaming.ts` | Streaming text output |
| `tool-calling.ts` | Function/tool calling |
| `multi-provider.ts` | Using multiple providers |
| `structured-output.ts` | JSON schema output |
| `multi-agent.ts` | Multi-agent attribution |
## Running Examples
```bash
npx ts-node basic-chat.ts
npx ts-node streaming.ts
npx ts-node tool-calling.ts
```
## More Information
- [AI SDK Documentation](https://docs.praisonai.com/docs/js/ai-sdk)
- [AI SDK CLI](https://docs.praisonai.com/docs/js/ai-sdk-cli)
+39
View File
@@ -0,0 +1,39 @@
/**
* AI SDK Basic Chat Example
*
* Demonstrates simple text generation using the AI SDK backend.
*
* Usage:
* npx ts-node basic-chat.ts
*
* Environment:
* OPENAI_API_KEY - Your OpenAI API key
*/
import { createAISDKBackend } from 'praisonai';
async function main() {
// Create backend with model string
const backend = createAISDKBackend('openai/gpt-4o-mini');
console.log('Sending message to OpenAI...\n');
// Generate text
const result = await backend.generateText({
messages: [
{ role: 'system', content: 'You are a helpful assistant.' },
{ role: 'user', content: 'What are the three primary colors?' }
],
maxTokens: 200,
temperature: 0.7,
});
console.log('Response:', result.text);
console.log('\nFinish reason:', result.finishReason);
if (result.usage) {
console.log('Token usage:', result.usage);
}
}
main().catch(console.error);
+32
View File
@@ -0,0 +1,32 @@
/**
* Generate Text Example
*
* Demonstrates basic text generation using AI SDK v6 via praisonai-ts.
*
* Run: npx ts-node generate-text.ts
* Required: OPENAI_API_KEY
*/
// Import from local praisonai-ts package
// In production: import { Agent } from 'praisonai';
import { Agent } from '../../../src/praisonai-ts/src';
async function main() {
// Simple text generation with Agent
const agent = new Agent({
instructions: 'You are a helpful assistant that provides concise answers.',
llm: 'openai/gpt-4o-mini',
});
console.log('=== Generate Text Example ===\n');
// Basic chat
const response = await agent.chat('What is the capital of France?');
console.log('Response:', response);
// With context
const response2 = await agent.chat('What is its population?');
console.log('\nFollow-up:', response2);
}
main().catch(console.error);
+91
View File
@@ -0,0 +1,91 @@
/**
* MCP Client Example (AI SDK v6 Parity)
*
* Demonstrates Model Context Protocol client with HTTP/SSE transports.
*
* Run: npx ts-node mcp-client.ts
* Note: Requires an MCP server running
*/
import {
createMCP,
closeMCPClient,
type MCPConfig,
} from '../../../src/praisonai-ts/src';
async function main() {
console.log('=== MCP Client Example ===\n');
// Example configurations for different transports
const configs: Record<string, MCPConfig> = {
// HTTP transport (recommended for production)
http: {
transport: {
type: 'http',
url: 'https://mcp-server.example.com/mcp',
headers: { Authorization: 'Bearer token' },
},
name: 'http-client',
},
// SSE transport
sse: {
transport: {
type: 'sse',
url: 'https://mcp-server.example.com/sse',
},
name: 'sse-client',
},
// Stdio transport (for local servers)
stdio: {
transport: {
type: 'stdio',
command: 'npx',
args: ['-y', '@modelcontextprotocol/server-everything'],
},
name: 'stdio-client',
},
};
console.log('MCP Transport Configurations:');
console.log('');
console.log('1. HTTP Transport (recommended):');
console.log(` ${JSON.stringify(configs.http.transport, null, 2)}`);
console.log('');
console.log('2. SSE Transport:');
console.log(` ${JSON.stringify(configs.sse.transport, null, 2)}`);
console.log('');
console.log('3. Stdio Transport (local):');
console.log(` ${JSON.stringify(configs.stdio.transport, null, 2)}`);
console.log('\n--- Example Usage ---\n');
console.log(`
// Create MCP client
const client = await createMCP({
transport: {
type: 'http',
url: 'https://your-mcp-server.com/mcp',
headers: { Authorization: 'Bearer your-token' },
}
});
// Get available tools
const tools = await client.tools();
console.log('Available tools:', Object.keys(tools));
// List resources
const resources = await client.listResources();
console.log('Resources:', resources);
// List prompts
const prompts = await client.listPrompts();
console.log('Prompts:', prompts);
// Close when done
await closeMCPClient('default');
`);
console.log('\n✅ MCP Client demo completed!');
console.log('Note: To run with a real MCP server, uncomment the connection code.');
}
main().catch(console.error);
+96
View File
@@ -0,0 +1,96 @@
/**
* AI SDK Multi-Agent Attribution Example
*
* Demonstrates how to propagate attribution context (agentId, runId, traceId)
* to LLM calls for multi-agent observability and safety.
*
* Usage:
* npx ts-node multi-agent.ts
*
* Environment:
* OPENAI_API_KEY - Your OpenAI API key
*/
import { createAISDKBackend } from 'praisonai';
// Simulate a multi-agent system with attribution
interface AgentContext {
agentId: string;
runId: string;
traceId: string;
sessionId?: string;
}
async function runAgent(name: string, task: string, context: AgentContext) {
console.log(`\n[${name}] Starting task: ${task}`);
console.log(` Attribution: agentId=${context.agentId}, runId=${context.runId}`);
const backend = createAISDKBackend('openai/gpt-4o-mini', {
attribution: context,
timeout: 30000,
});
const result = await backend.generateText({
messages: [
{ role: 'system', content: `You are ${name}, a specialized AI agent.` },
{ role: 'user', content: task }
],
maxTokens: 100,
});
console.log(` Response: ${result.text.slice(0, 100)}...`);
return result.text;
}
async function main() {
// Create a shared trace ID for the entire workflow
const traceId = `trace-${Date.now().toString(36)}`;
const sessionId = 'session-user123';
console.log('Multi-Agent Workflow with Attribution');
console.log('=====================================');
console.log(`Trace ID: ${traceId}`);
console.log(`Session ID: ${sessionId}`);
// Agent 1: Research Agent
const researchResult = await runAgent(
'ResearchAgent',
'What are the key benefits of TypeScript?',
{
agentId: 'agent-research',
runId: `run-${Date.now().toString(36)}-1`,
traceId,
sessionId,
}
);
// Agent 2: Summary Agent
const summaryResult = await runAgent(
'SummaryAgent',
`Summarize this in one sentence: ${researchResult}`,
{
agentId: 'agent-summary',
runId: `run-${Date.now().toString(36)}-2`,
traceId,
sessionId,
}
);
// Agent 3: Review Agent
await runAgent(
'ReviewAgent',
`Review this summary for accuracy: ${summaryResult}`,
{
agentId: 'agent-review',
runId: `run-${Date.now().toString(36)}-3`,
traceId,
sessionId,
}
);
console.log('\n\nWorkflow complete!');
console.log('All requests were tagged with the same trace ID for observability.');
console.log('Check your LLM provider dashboard to see the correlated requests.');
}
main().catch(console.error);
+59
View File
@@ -0,0 +1,59 @@
/**
* AI SDK Multi-Provider Example
*
* Demonstrates using multiple LLM providers with the same interface.
*
* Usage:
* npx ts-node multi-provider.ts
*
* Environment:
* OPENAI_API_KEY - Your OpenAI API key
* ANTHROPIC_API_KEY - Your Anthropic API key
* GOOGLE_API_KEY - Your Google API key
*/
import { createAISDKBackend } from 'praisonai';
const PROMPT = 'Explain what an API is in one sentence.';
async function testProvider(modelString: string) {
try {
const backend = createAISDKBackend(modelString, {
timeout: 30000,
});
const startTime = Date.now();
const result = await backend.generateText({
messages: [{ role: 'user', content: PROMPT }],
maxTokens: 100,
});
const duration = Date.now() - startTime;
console.log(`\n${modelString}:`);
console.log(` Response: ${result.text}`);
console.log(` Duration: ${duration}ms`);
return true;
} catch (error: any) {
console.log(`\n${modelString}:`);
console.log(` Error: ${error.message}`);
return false;
}
}
async function main() {
console.log('Testing multiple providers with the same prompt...');
console.log(`Prompt: "${PROMPT}"`);
const providers = [
'openai/gpt-4o-mini',
'anthropic/claude-3-haiku-20240307',
'google/gemini-1.5-flash',
];
const results = await Promise.all(providers.map(testProvider));
const successful = results.filter(Boolean).length;
console.log(`\n\nSummary: ${successful}/${providers.length} providers succeeded`);
}
main().catch(console.error);
@@ -0,0 +1,87 @@
/**
* Speech & Transcription Example (AI SDK v6 Parity)
*
* Demonstrates speech generation and transcription wrappers.
*
* Run: npx ts-node speech-transcription.ts
* Required: OPENAI_API_KEY or ELEVENLABS_API_KEY
*/
import {
generateSpeech,
transcribe,
SPEECH_MODELS,
TRANSCRIPTION_MODELS,
} from '../../../src/praisonai-ts/src';
async function main() {
console.log('=== Speech & Transcription Example ===\n');
// Show available models
console.log('1. Available Speech Models:');
Object.entries(SPEECH_MODELS).forEach(([key, value]) => {
console.log(` ${key}: ${value}`);
});
console.log('\n2. Available Transcription Models:');
Object.entries(TRANSCRIPTION_MODELS).forEach(([key, value]) => {
console.log(` ${key}: ${value}`);
});
// Example: Generate speech (requires API key)
console.log('\n3. Speech Generation Example:');
console.log(`
// Generate speech from text
const result = await generateSpeech({
text: 'Hello, this is a test of speech synthesis.',
model: 'openai/tts-1',
voice: 'alloy',
});
// Save to file
await fs.writeFile('output.mp3', result.audio);
console.log('Audio saved to output.mp3');
`);
// Example: Transcribe audio (requires API key)
console.log('4. Transcription Example:');
console.log(`
// Transcribe audio file
const result = await transcribe({
audio: await fs.readFile('audio.mp3'),
model: 'openai/whisper-1',
});
console.log('Transcription:', result.text);
console.log('Language:', result.language);
console.log('Duration:', result.durationInSeconds, 'seconds');
// With segments
if (result.segments) {
result.segments.forEach(seg => {
console.log(\`[\${seg.start}s - \${seg.end}s]: \${seg.text}\`);
});
}
`);
// Live test if API key is available
if (process.env.OPENAI_API_KEY) {
console.log('\n5. Live Speech Generation Test:');
try {
const result = await generateSpeech({
text: 'Hello from PraisonAI!',
model: 'openai/tts-1',
voice: 'alloy',
});
console.log(` ✅ Generated ${result.audio.length} bytes of audio`);
} catch (error: any) {
console.log(` ⚠️ Speech generation failed: ${error.message}`);
}
} else {
console.log('\n5. Live Test: Skipped (OPENAI_API_KEY not set)');
}
console.log('\n✅ Speech & Transcription demo completed!');
}
main().catch(console.error);
+27
View File
@@ -0,0 +1,27 @@
/**
* Stream Text Example
*
* Demonstrates streaming text generation using AI SDK v6 via praisonai-ts.
*
* Run: npx ts-node stream-text.ts
* Required: OPENAI_API_KEY
*/
import { Agent } from '../../../src/praisonai-ts/src';
async function main() {
const agent = new Agent({
instructions: 'You are a helpful assistant that provides detailed answers.',
llm: 'openai/gpt-4o-mini',
stream: true,
});
console.log('=== Stream Text Example ===\n');
console.log('Streaming response:\n');
// Stream the response
const response = await agent.chat('Explain quantum computing in 3 paragraphs.');
console.log('\n\nFinal response:', response);
}
main().catch(console.error);
+39
View File
@@ -0,0 +1,39 @@
/**
* AI SDK Streaming Example
*
* Demonstrates streaming text generation with real-time output.
*
* Usage:
* npx ts-node streaming.ts
*
* Environment:
* OPENAI_API_KEY - Your OpenAI API key
*/
import { createAISDKBackend } from 'praisonai';
async function main() {
const backend = createAISDKBackend('openai/gpt-4o-mini');
console.log('Streaming response from OpenAI...\n');
const stream = await backend.streamText({
messages: [
{ role: 'user', content: 'Write a short poem about coding.' }
],
maxTokens: 200,
});
// Stream chunks as they arrive
for await (const chunk of stream) {
if (chunk.text) {
process.stdout.write(chunk.text);
}
if (chunk.finishReason) {
console.log('\n\nFinish reason:', chunk.finishReason);
}
}
}
main().catch(console.error);
+72
View File
@@ -0,0 +1,72 @@
/**
* AI SDK Structured Output Example
*
* Demonstrates generating typed JSON objects with Zod schemas.
*
* Usage:
* npx ts-node structured-output.ts
*
* Environment:
* OPENAI_API_KEY - Your OpenAI API key
*/
import { generateObject } from 'ai';
import { openai } from '@ai-sdk/openai';
import { z } from 'zod';
// Define schema for a person profile
const PersonSchema = z.object({
name: z.string().describe('Full name'),
age: z.number().describe('Age in years'),
occupation: z.string().describe('Job title'),
city: z.string().describe('City of residence'),
interests: z.array(z.string()).describe('List of hobbies/interests'),
bio: z.string().describe('Short biography'),
});
// Define schema for a recipe
const RecipeSchema = z.object({
name: z.string().describe('Recipe name'),
cuisine: z.string().describe('Type of cuisine'),
prepTime: z.number().describe('Preparation time in minutes'),
cookTime: z.number().describe('Cooking time in minutes'),
servings: z.number().describe('Number of servings'),
ingredients: z.array(z.object({
item: z.string(),
amount: z.string(),
})).describe('List of ingredients'),
instructions: z.array(z.string()).describe('Step-by-step instructions'),
});
async function main() {
console.log('Generating structured output with Zod schemas...\n');
// Generate a person profile
console.log('1. Generating a fictional person profile...');
const personResult = await generateObject({
model: openai('gpt-4o-mini'),
schema: PersonSchema,
prompt: 'Generate a fictional software engineer profile.',
});
console.log('Person:', JSON.stringify(personResult.object, null, 2));
// Generate a recipe
console.log('\n2. Generating a recipe...');
const recipeResult = await generateObject({
model: openai('gpt-4o-mini'),
schema: RecipeSchema,
prompt: 'Generate a simple pasta recipe.',
});
console.log('Recipe:', JSON.stringify(recipeResult.object, null, 2));
// Type-safe access
console.log('\n3. Type-safe access:');
console.log(` Person name: ${personResult.object.name}`);
console.log(` Person age: ${personResult.object.age}`);
console.log(` Recipe name: ${recipeResult.object.name}`);
console.log(` Total time: ${recipeResult.object.prepTime + recipeResult.object.cookTime} minutes`);
}
main().catch(console.error);
+78
View File
@@ -0,0 +1,78 @@
/**
* Telemetry Example (AI SDK v6 Parity)
*
* Demonstrates telemetry and tracing with optional OpenTelemetry integration.
*
* Run: npx ts-node telemetry.ts
* Required: OPENAI_API_KEY
*/
import {
configureTelemetry,
isTelemetryEnabled,
createAISpan,
recordEvent,
getEvents,
clearEvents,
Agent,
} from '../../../src/praisonai-ts/src';
async function main() {
console.log('=== Telemetry Example ===\n');
// Configure telemetry
configureTelemetry({
isEnabled: true,
functionId: 'telemetry-demo',
metadata: {
version: '1.0.0',
environment: 'development',
},
});
console.log('1. Telemetry Status:');
console.log(` Enabled: ${isTelemetryEnabled()}`);
// Create a span for an operation
console.log('\n2. Creating AI Span:');
const span = createAISpan('generateText', {
model: 'gpt-4o-mini',
provider: 'openai',
});
try {
// Simulate AI operation
console.log(' Performing AI operation...');
await new Promise(resolve => setTimeout(resolve, 100));
span.setAttribute('tokens', 150);
span.setStatus({ code: 'ok' });
console.log(' Span completed successfully');
} catch (error: any) {
span.recordException(error);
span.setStatus({ code: 'error', message: error.message });
} finally {
span.end();
}
// Record custom events
console.log('\n3. Recording Events:');
recordEvent('agent.started', { agentName: 'demo-agent' });
recordEvent('tool.called', { toolName: 'search', duration: 250 });
recordEvent('agent.completed', { success: true });
const events = getEvents();
console.log(` Recorded ${events.length} events:`);
events.forEach(e => {
console.log(` - ${e.name}: ${JSON.stringify(e.attributes)}`);
});
// Clear events
console.log('\n4. Clearing Events:');
clearEvents();
console.log(` Events after clear: ${getEvents().length}`);
console.log('\n✅ Telemetry demo completed!');
}
main().catch(console.error);
+71
View File
@@ -0,0 +1,71 @@
/**
* Tool Approval Example (AI SDK v6 Parity)
*
* Demonstrates human-in-the-loop tool approval using needsApproval.
*
* Run: npx ts-node tool-approval.ts
* Required: OPENAI_API_KEY
*/
import {
Agent,
withApproval,
getApprovalManager,
createDangerousPatternChecker,
} from '../../../src/praisonai-ts/src';
async function main() {
console.log('=== Tool Approval Example ===\n');
// Set up approval handler
const manager = getApprovalManager();
manager.onApprovalRequest(async (request) => {
console.log(`\n🔐 APPROVAL REQUESTED`);
console.log(` Tool: ${request.toolName}`);
console.log(` Input: ${JSON.stringify(request.input)}`);
console.log(` Reason: ${request.reason || 'No reason provided'}`);
// Auto-approve for demo (in production, prompt user)
console.log(` ✅ Auto-approved for demo\n`);
return true;
});
// Create a tool that always requires approval
const deleteFile = withApproval({
name: 'deleteFile',
description: 'Delete a file from the filesystem',
needsApproval: true, // Always require approval
execute: async (input: { path: string }) => {
console.log(` Executing: Would delete ${input.path}`);
return { success: true, deleted: input.path };
},
});
// Create a tool with conditional approval (dangerous patterns only)
const runCommand = withApproval({
name: 'runCommand',
description: 'Run a shell command',
needsApproval: createDangerousPatternChecker(),
execute: async (input: { command: string }) => {
console.log(` Executing: ${input.command}`);
return { output: 'command executed' };
},
});
// Test the tools
console.log('1. Testing deleteFile (always requires approval):');
const result1 = await deleteFile({ path: '/tmp/test.txt' });
console.log(' Result:', result1);
console.log('\n2. Testing runCommand with safe command (no approval needed):');
const result2 = await runCommand({ command: 'ls -la' });
console.log(' Result:', result2);
console.log('\n3. Testing runCommand with dangerous command (requires approval):');
const result3 = await runCommand({ command: 'rm -rf /tmp/test' });
console.log(' Result:', result3);
console.log('\n✅ Tool approval demo completed!');
}
main().catch(console.error);
+75
View File
@@ -0,0 +1,75 @@
/**
* AI SDK Tool Calling Example
*
* Demonstrates function/tool calling with Zod schemas.
*
* Usage:
* npx ts-node tool-calling.ts
*
* Environment:
* OPENAI_API_KEY - Your OpenAI API key
*/
import { generateText, tool } from 'ai';
import { openai } from '@ai-sdk/openai';
import { z } from 'zod';
// Define tools with Zod schemas
const tools = {
get_weather: tool({
description: 'Get the current weather for a city',
inputSchema: z.object({
city: z.string().describe('The city name'),
unit: z.enum(['celsius', 'fahrenheit']).optional().describe('Temperature unit'),
}),
execute: async ({ city, unit = 'celsius' }) => {
// Mock weather data
const weather = {
city,
temperature: unit === 'celsius' ? 22 : 72,
unit,
condition: 'sunny',
humidity: 45,
};
return weather;
},
}),
get_time: tool({
description: 'Get the current time in a timezone',
inputSchema: z.object({
timezone: z.string().describe('Timezone like "America/New_York"'),
}),
execute: async ({ timezone }) => {
const time = new Date().toLocaleString('en-US', { timeZone: timezone });
return { timezone, time };
},
}),
};
async function main() {
console.log('Asking about weather with tool calling...\n');
const result = await generateText({
model: openai('gpt-4o-mini'),
messages: [
{ role: 'user', content: 'What is the weather in Paris and what time is it in New York?' }
],
tools,
maxSteps: 5, // Allow multiple tool calls
});
console.log('Tool Calls:');
for (const toolCall of result.toolCalls) {
console.log(` - ${toolCall.toolName}:`, toolCall.args);
}
console.log('\nTool Results:');
for (const toolResult of result.toolResults) {
console.log(` - ${toolResult.toolName}:`, toolResult.result);
}
console.log('\nFinal Response:', result.text);
}
main().catch(console.error);
+63
View File
@@ -0,0 +1,63 @@
/**
* UI Message Example (AI SDK v6 Parity)
*
* Demonstrates UIMessage protocol for Next.js/React integration.
*
* Run: npx ts-node ui-message.ts
* Required: OPENAI_API_KEY
*/
import {
createTextMessage,
createSystemMessage,
convertToModelMessages,
convertToUIMessages,
validateUIMessages,
hasPendingApprovals,
type UIMessage,
} from '../../../src/praisonai-ts/src';
async function main() {
console.log('=== UI Message Example ===\n');
// Create UI messages
const messages: UIMessage[] = [
createSystemMessage('You are a helpful assistant.'),
createTextMessage('user', 'What is the capital of France?'),
createTextMessage('assistant', 'The capital of France is Paris.'),
createTextMessage('user', 'What is its population?'),
];
console.log('1. Created UI Messages:');
messages.forEach((msg, i) => {
const text = msg.parts.find(p => p.type === 'text');
console.log(` [${i}] ${msg.role}: ${text ? (text as any).text.substring(0, 50) : '...'}`);
});
// Validate messages
console.log('\n2. Validating messages:');
const isValid = validateUIMessages(messages);
console.log(` Valid: ${isValid}`);
// Convert to model messages (for AI SDK)
console.log('\n3. Converting to Model Messages:');
const modelMessages = await convertToModelMessages(messages);
modelMessages.forEach((msg, i) => {
const content = typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
console.log(` [${i}] ${msg.role}: ${content.substring(0, 50)}...`);
});
// Check for pending approvals
console.log('\n4. Checking for pending approvals:');
const hasPending = messages.some(m => hasPendingApprovals(m));
console.log(` Has pending approvals: ${hasPending}`);
// Convert model messages back to UI messages
console.log('\n5. Round-trip conversion:');
const uiMessages = convertToUIMessages(modelMessages);
console.log(` Converted ${uiMessages.length} messages back to UI format`);
console.log('\n✅ UI Message demo completed!');
}
main().catch(console.error);
+65
View File
@@ -0,0 +1,65 @@
/**
* Basic Embedding Example
*
* Demonstrates using Agent.embed() for text embeddings
*
* Run: npx ts-node examples/js/embeddings/embed-basic.ts
*/
import { Agent } from '../../../src/praisonai-ts/dist/agent/simple';
import { embed, embedMany, cosineSimilarity } from '../../../src/praisonai-ts/dist/llm/embeddings';
async function main() {
console.log('=== Basic Embedding Example ===\n');
// Method 1: Using Agent.embed()
console.log('1. Agent.embed() - Single text');
const agent = new Agent({
instructions: 'You are a helpful assistant',
llm: 'openai/gpt-4o-mini',
verbose: false
});
const singleEmbedding = await agent.embed('Hello world');
console.log(` Dimensions: ${(singleEmbedding as number[]).length}`);
console.log(` First 5 values: [${(singleEmbedding as number[]).slice(0, 5).map(v => v.toFixed(4)).join(', ')}...]`);
// Method 2: Using Agent.embed() with multiple texts
console.log('\n2. Agent.embed() - Multiple texts');
const multiEmbeddings = await agent.embed(['Hello', 'World', 'How are you?']);
console.log(` Count: ${(multiEmbeddings as number[][]).length}`);
console.log(` Each has ${(multiEmbeddings as number[][])[0].length} dimensions`);
// Method 3: Direct embed function
console.log('\n3. Direct embed() function');
const result = await embed('PraisonAI is an AI agent framework');
console.log(` Dimensions: ${result.embedding.length}`);
console.log(` Tokens used: ${result.usage?.tokens || 'N/A'}`);
// Method 4: Batch embeddings
console.log('\n4. Batch embedMany() function');
const texts = [
'Machine learning is a subset of AI',
'Deep learning uses neural networks',
'Natural language processing handles text'
];
const batchResult = await embedMany(texts);
console.log(` Embedded ${batchResult.embeddings.length} texts`);
console.log(` Total tokens: ${batchResult.usage?.tokens || 'N/A'}`);
// Method 5: Similarity comparison
console.log('\n5. Cosine Similarity');
const emb1 = await embed('I love programming');
const emb2 = await embed('Coding is my passion');
const emb3 = await embed('The weather is nice today');
const sim12 = cosineSimilarity(emb1.embedding, emb2.embedding);
const sim13 = cosineSimilarity(emb1.embedding, emb3.embedding);
console.log(` "I love programming" vs "Coding is my passion": ${(sim12 * 100).toFixed(1)}%`);
console.log(` "I love programming" vs "The weather is nice": ${(sim13 * 100).toFixed(1)}%`);
console.log('\n=== Complete ===');
}
main().catch(console.error);
+119
View File
@@ -0,0 +1,119 @@
/**
* Document Embedding Example
*
* Demonstrates chunking and embedding documents for retrieval
*
* Run: npx ts-node examples/js/embeddings/embed-docs.ts
*/
import { embedMany, cosineSimilarity } from '../../../src/praisonai-ts/dist/llm/embeddings';
// Sample documents
const documents = [
{
id: 'doc1',
title: 'Introduction to AI',
content: 'Artificial Intelligence (AI) is the simulation of human intelligence by machines. It includes machine learning, natural language processing, and computer vision.'
},
{
id: 'doc2',
title: 'Machine Learning Basics',
content: 'Machine learning is a subset of AI that enables systems to learn from data. It includes supervised learning, unsupervised learning, and reinforcement learning.'
},
{
id: 'doc3',
title: 'Deep Learning',
content: 'Deep learning uses neural networks with multiple layers. It excels at tasks like image recognition, speech processing, and natural language understanding.'
},
{
id: 'doc4',
title: 'Natural Language Processing',
content: 'NLP enables computers to understand human language. Applications include chatbots, translation, sentiment analysis, and text summarization.'
},
{
id: 'doc5',
title: 'Computer Vision',
content: 'Computer vision allows machines to interpret visual information. It powers applications like facial recognition, autonomous vehicles, and medical imaging.'
}
];
interface EmbeddedDocument {
id: string;
title: string;
content: string;
embedding: number[];
}
async function main() {
console.log('=== Document Embedding Example ===\n');
// Step 1: Embed all documents
console.log('1. Embedding documents...');
const contents = documents.map(d => d.content);
const { embeddings } = await embedMany(contents);
const embeddedDocs: EmbeddedDocument[] = documents.map((doc, i) => ({
...doc,
embedding: embeddings[i]
}));
console.log(` Embedded ${embeddedDocs.length} documents`);
console.log(` Embedding dimensions: ${embeddings[0].length}`);
// Step 2: Search function
async function search(query: string, topK: number = 3): Promise<Array<{ doc: EmbeddedDocument; score: number }>> {
const { embedding: queryEmbedding } = await (await import('../../../src/praisonai-ts/dist/llm/embeddings')).embed(query);
const results = embeddedDocs.map(doc => ({
doc,
score: cosineSimilarity(queryEmbedding, doc.embedding)
}));
return results
.sort((a, b) => b.score - a.score)
.slice(0, topK);
}
// Step 3: Test queries
console.log('\n2. Testing search queries...\n');
const queries = [
'How do machines learn from data?',
'What is image recognition?',
'Tell me about chatbots and translation'
];
for (const query of queries) {
console.log(`Query: "${query}"`);
const results = await search(query);
console.log('Top results:');
results.forEach((r, i) => {
console.log(` ${i + 1}. [${(r.score * 100).toFixed(1)}%] ${r.doc.title}`);
});
console.log('');
}
// Step 4: Show document similarity matrix
console.log('3. Document Similarity Matrix');
console.log('');
// Header
process.stdout.write(' ');
documents.forEach((_, i) => process.stdout.write(`Doc${i + 1} `));
console.log('');
// Matrix
for (let i = 0; i < embeddedDocs.length; i++) {
process.stdout.write(`Doc${i + 1} `);
for (let j = 0; j < embeddedDocs.length; j++) {
const sim = cosineSimilarity(embeddedDocs[i].embedding, embeddedDocs[j].embedding);
process.stdout.write(`${(sim * 100).toFixed(0).padStart(3)}% `);
}
console.log('');
}
console.log('\n=== Complete ===');
}
main().catch(console.error);
+131
View File
@@ -0,0 +1,131 @@
/**
* Retrieval Example - Simple Vector Search
*
* Demonstrates building a simple retrieval system with embeddings
*
* Run: npx ts-node examples/js/embeddings/retrieval.ts
*/
import { embed, embedMany, cosineSimilarity } from '../../../src/praisonai-ts/dist/llm/embeddings';
interface Document {
id: string;
content: string;
embedding?: number[];
}
interface SearchResult {
document: Document;
score: number;
}
// Simple in-memory vector store
class SimpleVectorStore {
private documents: Document[] = [];
async add(doc: Document): Promise<void> {
const { embedding } = await embed(doc.content);
this.documents.push({ ...doc, embedding });
}
async addBatch(docs: Document[]): Promise<void> {
const contents = docs.map(d => d.content);
const { embeddings } = await embedMany(contents);
docs.forEach((doc, i) => {
this.documents.push({ ...doc, embedding: embeddings[i] });
});
}
async search(query: string, topK: number = 3): Promise<SearchResult[]> {
const { embedding: queryEmb } = await embed(query);
const results = this.documents
.filter(doc => doc.embedding)
.map(doc => ({
document: doc,
score: cosineSimilarity(queryEmb, doc.embedding!)
}))
.sort((a, b) => b.score - a.score)
.slice(0, topK);
return results;
}
get(id: string): Document | undefined {
return this.documents.find(d => d.id === id);
}
delete(id: string): boolean {
const idx = this.documents.findIndex(d => d.id === id);
if (idx >= 0) {
this.documents.splice(idx, 1);
return true;
}
return false;
}
}
async function main() {
console.log('=== Simple Vector Retrieval ===\n');
// Create vector store
console.log('1. Creating vector store...');
const store = new SimpleVectorStore();
// Add documents
console.log('2. Adding documents...');
const docs: Document[] = [
{ id: '1', content: 'PraisonAI is an AI agent framework for building intelligent applications.' },
{ id: '2', content: 'Agents can use tools to interact with external systems and APIs.' },
{ id: '3', content: 'Workflows allow orchestrating multiple agents in sequence or parallel.' },
{ id: '4', content: 'Memory systems help agents maintain context across conversations.' },
{ id: '5', content: 'The AI SDK provides a unified interface for multiple LLM providers.' },
{ id: '6', content: 'Embeddings enable semantic search and similarity matching.' },
{ id: '7', content: 'Guardrails ensure agent outputs meet safety and quality standards.' },
{ id: '8', content: 'Sessions persist conversation state for multi-turn interactions.' }
];
await store.addBatch(docs);
console.log(` Added ${docs.length} documents`);
// Search queries
console.log('\n3. Running search queries...\n');
const queries = [
'How do I build an AI application?',
'Can agents call external APIs?',
'How to maintain conversation context?',
'What providers are supported?'
];
for (const query of queries) {
console.log(`Query: "${query}"`);
const results = await store.search(query);
if (results.length > 0) {
console.log('Results:');
results.forEach((r, i) => {
console.log(` ${i + 1}. [${(r.score * 100).toFixed(1)}%] ${r.document.content.substring(0, 60)}...`);
});
} else {
console.log(' No results found');
}
console.log('');
}
// Get specific document
console.log('4. Get document by ID...');
const doc = store.get('5');
if (doc) {
console.log(` Found: ${doc.content}`);
}
// Delete document
console.log('\n5. Delete document...');
const deleted = store.delete('8');
console.log(` Deleted doc 8: ${deleted}`);
console.log('\n=== Complete ===');
}
main().catch(console.error);
+49
View File
@@ -0,0 +1,49 @@
/**
* MCP Filesystem Server Example
*
* Demonstrates using MCP to connect to a filesystem server.
*
* Prerequisites:
* npm install praisonai-ts
* export OPENAI_API_KEY=your-api-key
*
* Run:
* npx ts-node filesystem-mcp.ts
*/
import { Agent, createMCP } from '../../../src/praisonai-ts/src';
async function main() {
console.log('=== MCP Filesystem Server Example ===\n');
// Connect to MCP filesystem server
const mcp = await createMCP({
transport: {
type: 'stdio',
command: 'npx',
args: ['-y', '@modelcontextprotocol/server-filesystem', '.'],
},
});
console.log('Connected to MCP filesystem server\n');
// Get tools from MCP server
const tools = await mcp.tools();
console.log('Available tools:', Object.keys(tools));
// Create agent with MCP tools
const agent = new Agent({
name: 'FileAgent',
instructions: 'You help manage files using the filesystem tools.',
tools: Object.values(tools),
});
// Ask the agent to list files
const response = await agent.chat('List all TypeScript files in the current directory');
console.log('\nResponse:', response);
// Cleanup
await mcp.close();
}
main().catch(console.error);
+41
View File
@@ -0,0 +1,41 @@
/**
* Multimodal Image Analysis Example
*
* Demonstrates analyzing images with a vision-capable model.
*
* Prerequisites:
* npm install praisonai-ts
* export OPENAI_API_KEY=your-api-key
*
* Run:
* npx ts-node image-analysis.ts
*/
import { Agent } from '../../../src/praisonai-ts/src';
async function main() {
console.log('=== Multimodal Image Analysis Example ===\n');
// Create an agent with a vision-capable model
const agent = new Agent({
name: 'VisionAgent',
llm: 'gpt-4o', // Vision-capable model
instructions: 'You analyze images and describe what you see in detail.',
});
// For multimodal, we describe the image in the prompt
// The agent will use its vision capabilities
const imageUrl = 'https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/PNG_transparency_demonstration_1.png/300px-PNG_transparency_demonstration_1.png';
console.log('Analyzing image...\n');
console.log('Image URL:', imageUrl, '\n');
// Ask the agent to analyze (in production, use multimodal message format)
const response = await agent.chat(
`Please analyze this image: ${imageUrl}\nDescribe what you see in detail.`
);
console.log('Analysis:', response);
}
main().catch(console.error);
+50
View File
@@ -0,0 +1,50 @@
# Observability Examples
Examples demonstrating how to use observability integrations with PraisonAI TypeScript.
## Setup
```bash
# Install dependencies
cd /path/to/praisonai-ts
pnpm install
# Set API keys for observability tools
export LANGFUSE_SECRET_KEY=sk-...
export LANGFUSE_PUBLIC_KEY=pk-...
# Or use built-in tools (no API key needed)
```
## Examples
| File | Description |
|------|-------------|
| `basic-tracing.ts` | Basic tracing with memory adapter |
| `langfuse-integration.ts` | Langfuse observability integration |
| `multi-agent-attribution.ts` | Track agent_id/run_id across agents |
| `switchable-tools.ts` | Switch between observability tools via env |
## Running Examples
```bash
# Run with ts-node or tsx
npx tsx examples/js/observability/basic-tracing.ts
npx tsx examples/js/observability/multi-agent-attribution.ts
# With Langfuse (requires API keys)
LANGFUSE_SECRET_KEY=sk-... npx tsx examples/js/observability/langfuse-integration.ts
```
## Supported Tools
| Tool | Env Variable | Package |
|------|--------------|---------|
| Langfuse | `LANGFUSE_SECRET_KEY` | `langfuse` |
| LangSmith | `LANGCHAIN_API_KEY` | `langsmith` |
| LangWatch | `LANGWATCH_API_KEY` | `langwatch` |
| Arize | `ARIZE_API_KEY` | `arize-phoenix` |
| Axiom | `AXIOM_TOKEN` | `@axiomhq/js` |
| Braintrust | `BRAINTRUST_API_KEY` | `braintrust` |
| Helicone | `HELICONE_API_KEY` | `@helicone/helicone` |
| Console | (none) | built-in |
| Memory | (none) | built-in |
@@ -0,0 +1,45 @@
/**
* Basic Tracing Example
*
* Demonstrates basic observability with the built-in memory adapter.
* No external API keys required.
*/
import { Agent, MemoryObservabilityAdapter, setObservabilityAdapter } from 'praisonai';
async function main() {
// Create memory adapter (stores traces in memory)
const obs = new MemoryObservabilityAdapter();
setObservabilityAdapter(obs);
// Create agent with observability enabled
const agent = new Agent({
name: 'TracedAgent',
instructions: 'You are a helpful assistant. Answer questions concisely.'
});
console.log('Running agent with tracing...\n');
// Run some interactions
await agent.chat('What is 2 + 2?');
await agent.chat('What is the capital of France?');
// View traces
const traces = obs.getAllTraces();
console.log('\n=== Trace Summary ===');
console.log(`Total traces: ${traces.length}`);
for (const trace of traces) {
console.log(`\nTrace: ${trace.name}`);
console.log(` Status: ${trace.status}`);
console.log(` Spans: ${trace.spans.length}`);
for (const span of trace.spans) {
const duration = span.endTime ? span.endTime - span.startTime : 'running';
console.log(` - ${span.kind}: ${span.name} (${duration}ms)`);
}
}
}
main().catch(console.error);
@@ -0,0 +1,60 @@
/**
* Multi-Agent Attribution Example
*
* Demonstrates tracking agent_id, run_id, and trace_id across
* multi-agent workflows for proper attribution.
*/
import { Agent, Agents, MemoryObservabilityAdapter, setObservabilityAdapter } from 'praisonai';
async function main() {
// Create memory adapter
const obs = new MemoryObservabilityAdapter();
setObservabilityAdapter(obs);
// Create agents
const researcher = new Agent({
name: 'Researcher',
instructions: 'Research topics and gather information.'
});
const writer = new Agent({
name: 'Writer',
instructions: 'Write clear summaries based on research.'
});
// Create multi-agent workflow
const agents = new Agents({
agents: [researcher, writer],
tasks: [
{ agent: researcher, description: 'Research: {topic}' },
{ agent: writer, description: 'Summarize the research findings' }
]
});
console.log('Running multi-agent workflow with attribution...\n');
await agents.start({ topic: 'TypeScript best practices' });
// View traces with attribution
const traces = obs.getAllTraces();
console.log('\n=== Attribution Report ===');
for (const trace of traces) {
console.log(`\nTrace ID: ${trace.id}`);
console.log(` Name: ${trace.name}`);
if (trace.attribution) {
console.log(` Agent ID: ${trace.attribution.agentId || 'N/A'}`);
console.log(` Run ID: ${trace.attribution.runId || 'N/A'}`);
console.log(` Session ID: ${trace.attribution.sessionId || 'N/A'}`);
}
console.log(` Spans: ${trace.spans.length}`);
for (const span of trace.spans) {
console.log(` - [${span.kind}] ${span.name}`);
}
}
}
main().catch(console.error);
@@ -0,0 +1,40 @@
/**
* Switchable Observability Tools Example
*
* Demonstrates switching between observability tools via environment variable.
* Set OBS_TOOL env var to: langfuse, langsmith, console, memory, noop
*
* Usage:
* OBS_TOOL=console npx tsx switchable-tools.ts
* OBS_TOOL=langfuse LANGFUSE_SECRET_KEY=sk-... npx tsx switchable-tools.ts
*/
import { Agent, createObservabilityAdapter, setObservabilityAdapter } from 'praisonai';
async function main() {
// Get tool from environment (default to memory)
const toolName = process.env.OBS_TOOL || 'memory';
console.log(`Using observability tool: ${toolName}\n`);
// Create adapter dynamically
const obs = await createObservabilityAdapter(toolName as any);
setObservabilityAdapter(obs);
// Create agent
const agent = new Agent({
name: 'ObservedAgent',
instructions: 'You are a helpful assistant.'
});
// Run interaction
console.log('Running agent...\n');
const response = await agent.chat('Hello! What is 2 + 2?');
console.log('Response:', response);
// Flush traces to backend
await obs.flush();
console.log('\nTraces flushed to', toolName);
}
main().catch(console.error);
+47
View File
@@ -0,0 +1,47 @@
# Provider Examples
Examples demonstrating how to use different LLM providers with PraisonAI TypeScript.
## Setup
```bash
# Install dependencies
cd /path/to/praisonai-ts
pnpm install
# Set API keys
export OPENAI_API_KEY=sk-...
export ANTHROPIC_API_KEY=sk-ant-...
export GOOGLE_API_KEY=AIza...
```
## Examples
| File | Description |
|------|-------------|
| `multi-provider.ts` | Use multiple providers in one workflow |
| `openai-compatible.ts` | Use OpenAI-compatible APIs (LM Studio, etc.) |
| `local-ollama.ts` | Use local Ollama models |
| `provider-selection.ts` | Dynamic provider selection based on task |
## Running Examples
```bash
# Run with ts-node or tsx
npx tsx examples/js/providers/multi-provider.ts
npx tsx examples/js/providers/openai-compatible.ts
npx tsx examples/js/providers/local-ollama.ts
```
## Environment Variables
| Provider | Variable | Example |
|----------|----------|---------|
| OpenAI | `OPENAI_API_KEY` | `sk-...` |
| Anthropic | `ANTHROPIC_API_KEY` | `sk-ant-...` |
| Google | `GOOGLE_API_KEY` | `AIza...` |
| xAI | `XAI_API_KEY` | `xai-...` |
| Groq | `GROQ_API_KEY` | `gsk_...` |
| Mistral | `MISTRAL_API_KEY` | `...` |
| DeepSeek | `DEEPSEEK_API_KEY` | `...` |
| Ollama | `OLLAMA_BASE_URL` | `http://localhost:11434` |
+40
View File
@@ -0,0 +1,40 @@
/**
* Local Ollama Provider Example
*
* Demonstrates using local Ollama models with PraisonAI.
*
* Setup:
* 1. Install Ollama: https://ollama.ai
* 2. Pull a model: ollama pull llama3.2
* 3. Run this example
*
* Optional env var:
* - OLLAMA_BASE_URL (default: http://localhost:11434)
*/
import { Agent } from 'praisonai';
async function main() {
// Ollama with llama3.2
const ollamaAgent = new Agent({
name: 'LocalLlama',
instructions: 'You are a helpful local AI assistant running on Ollama.',
llm: 'ollama/llama3.2',
llmConfig: {
baseUrl: process.env.OLLAMA_BASE_URL || 'http://localhost:11434'
}
});
console.log('Testing Ollama local model...\n');
try {
const response = await ollamaAgent.chat('Hello! What model are you running?');
console.log('Response:', response);
} catch (error: any) {
console.log('Note: Ollama not running or model not pulled.');
console.log('Start Ollama and run: ollama pull llama3.2');
console.log('Error:', error.message);
}
}
main().catch(console.error);
+57
View File
@@ -0,0 +1,57 @@
/**
* Multi-Provider Example
*
* Demonstrates using multiple LLM providers in a single workflow.
* Each agent uses a different provider optimized for its task.
*
* Required env vars:
* - OPENAI_API_KEY
* - ANTHROPIC_API_KEY
* - GOOGLE_API_KEY
*/
import { Agent, Agents } from 'praisonai';
async function main() {
// Fast model for quick categorization
const triageAgent = new Agent({
name: 'Triage',
instructions: 'Quickly categorize the user request into: technical, creative, or general.',
llm: 'openai/gpt-4o-mini' // Fast and cheap
});
// Powerful model for complex reasoning
const analysisAgent = new Agent({
name: 'Analyst',
instructions: 'Analyze the request in depth and provide detailed insights.',
llm: 'anthropic/claude-3-5-sonnet' // Best reasoning
});
// Creative model for content generation
const writerAgent = new Agent({
name: 'Writer',
instructions: 'Write a clear, engaging response based on the analysis.',
llm: 'google/gemini-2.0-flash' // Good for creative tasks
});
// Create multi-agent workflow
const agents = new Agents({
agents: [triageAgent, analysisAgent, writerAgent],
tasks: [
{ agent: triageAgent, description: 'Categorize: {input}' },
{ agent: analysisAgent, description: 'Analyze the categorized request' },
{ agent: writerAgent, description: 'Write the final response' }
]
});
console.log('Starting multi-provider workflow...\n');
const result = await agents.start({
input: 'Explain how neural networks learn'
});
console.log('\n=== Final Result ===');
console.log(result);
}
main().catch(console.error);
@@ -0,0 +1,48 @@
/**
* OpenAI-Compatible Provider Example
*
* Demonstrates using any OpenAI-compatible API endpoint.
* Works with LM Studio, vLLM, LocalAI, and other compatible servers.
*
* Required env vars:
* - OPENAI_COMPATIBLE_BASE_URL (e.g., http://localhost:1234/v1)
* - OPENAI_COMPATIBLE_API_KEY (optional, use 'not-needed' for local)
*/
import { Agent } from 'praisonai';
async function main() {
// LM Studio example (local)
const lmStudioAgent = new Agent({
name: 'LocalAssistant',
instructions: 'You are a helpful local AI assistant.',
llm: 'openai-compatible/local-model',
llmConfig: {
baseUrl: process.env.LM_STUDIO_BASE_URL || 'http://localhost:1234/v1',
apiKey: 'not-needed'
}
});
// Custom OpenAI-compatible endpoint
const customAgent = new Agent({
name: 'CustomAssistant',
instructions: 'You are a helpful assistant.',
llm: 'openai-compatible/my-model',
llmConfig: {
baseUrl: process.env.OPENAI_COMPATIBLE_BASE_URL || 'http://localhost:8000/v1',
apiKey: process.env.OPENAI_COMPATIBLE_API_KEY || 'not-needed'
}
});
console.log('Testing OpenAI-compatible endpoint...\n');
try {
const response = await lmStudioAgent.chat('Hello! What model are you?');
console.log('Response:', response);
} catch (error: any) {
console.log('Note: LM Studio not running. Start it with a model loaded.');
console.log('Error:', error.message);
}
}
main().catch(console.error);
+57
View File
@@ -0,0 +1,57 @@
/**
* Basic RAG Agent Example
*
* Demonstrates building a RAG agent with context injection.
*
* Prerequisites:
* npm install praisonai-ts
* export OPENAI_API_KEY=your-api-key
*
* Run:
* npx ts-node basic-rag.ts
*/
import { Agent } from '../../../src/praisonai-ts/src';
// Simulated knowledge base (in production, use a vector store)
const knowledgeBase = [
{ id: '1', content: 'PraisonAI is an AI agent framework for building intelligent applications.' },
{ id: '2', content: 'Agents can use tools to accomplish tasks like web search and calculations.' },
{ id: '3', content: 'RAG combines retrieval with generation for accurate, grounded responses.' },
];
// Simple search function (in production, use vector similarity)
function searchKnowledge(query: string): string[] {
const queryLower = query.toLowerCase();
return knowledgeBase
.filter(doc => doc.content.toLowerCase().includes(queryLower.split(' ')[0]))
.map(doc => doc.content);
}
async function main() {
console.log('=== Basic RAG Agent Example ===\n');
// Create an agent with RAG-style instructions
const agent = new Agent({
name: 'RAGAgent',
instructions: `You are a helpful assistant that answers questions using provided context.
Always base your answers on the context provided.
If the context doesn't contain relevant information, say so.`,
});
// Query with context injection
const query = 'What is PraisonAI?';
console.log(`Q: ${query}`);
// Retrieve relevant context
const context = searchKnowledge(query).join('\n\n');
// Ask the agent with context
const response = await agent.chat(
`Context:\n${context}\n\nQuestion: ${query}`
);
console.log(`A: ${response}\n`);
}
main().catch(console.error);
@@ -0,0 +1,77 @@
/**
* Express Server Adapter Example
*
* Demonstrates deploying an AI agent as an Express HTTP server.
*
* Prerequisites:
* npm install praisonai-ts express
* export OPENAI_API_KEY=your-api-key
*
* Run:
* npx ts-node express-server.ts
*
* Test:
* curl -X POST http://localhost:3000/api/chat \
* -H "Content-Type: application/json" \
* -d '{"message": "Hello!"}'
*/
import { Agent } from '../../../src/praisonai-ts/src';
// Simple HTTP server without Express dependency for this example
import http from 'http';
async function main() {
console.log('=== Express Server Adapter Example ===\n');
// Create an agent
const agent = new Agent({
name: 'APIAgent',
instructions: 'You are a helpful API assistant. Keep responses concise.',
});
// Create HTTP server
const server = http.createServer(async (req, res) => {
// CORS headers
res.setHeader('Access-Control-Allow-Origin', '*');
res.setHeader('Access-Control-Allow-Methods', 'POST, OPTIONS');
res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
if (req.method === 'OPTIONS') {
res.writeHead(200);
res.end();
return;
}
if (req.method === 'POST' && req.url === '/api/chat') {
let body = '';
req.on('data', chunk => body += chunk);
req.on('end', async () => {
try {
const { message } = JSON.parse(body);
const response = await agent.chat(message);
res.writeHead(200, { 'Content-Type': 'application/json' });
res.end(JSON.stringify({ response }));
} catch (error: any) {
res.writeHead(500, { 'Content-Type': 'application/json' });
res.end(JSON.stringify({ error: error.message }));
}
});
} else {
res.writeHead(404);
res.end('Not Found');
}
});
const PORT = process.env.PORT || 3000;
server.listen(PORT, () => {
console.log(`Server running on http://localhost:${PORT}`);
console.log('\nTest with:');
console.log(` curl -X POST http://localhost:${PORT}/api/chat \\`);
console.log(' -H "Content-Type: application/json" \\');
console.log(' -d \'{"message": "Hello!"}\'');
});
}
main().catch(console.error);
+130
View File
@@ -0,0 +1,130 @@
"""
LSP Code Intelligence Example
This example demonstrates how to use LSP-powered tools through an Agent
for semantic code analysis (listing symbols, finding definitions, etc.).
How to run:
export OPENAI_API_KEY=your_key
python examples/lsp/lsp_code_analysis_example.py
Prerequisites:
pip install praisonai praisonaiagents
pip install python-lsp-server # Optional, for full LSP support
"""
import asyncio
import os
import sys
# Add paths for development
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai-agents'))
async def main():
print("=" * 60)
print("LSP CODE INTELLIGENCE EXAMPLE")
print("=" * 60)
print()
# Import after path setup
from praisonai.cli.features import (
create_agent_centric_tools,
InteractiveRuntime,
RuntimeConfig
)
from praisonaiagents import Agent
# Use the praisonaiagents source as workspace for analysis
workspace = os.path.join(
os.path.dirname(__file__),
'../../src/praisonai-agents/praisonaiagents'
)
workspace = os.path.abspath(workspace)
print(f"Workspace: {workspace}")
print()
# 1. Create runtime with LSP enabled
print("1. Creating InteractiveRuntime with LSP...")
config = RuntimeConfig(
workspace=workspace,
lsp_enabled=True, # Enable LSP for code intelligence
acp_enabled=False, # Disable ACP for this example
approval_mode="auto"
)
runtime = InteractiveRuntime(config)
status = await runtime.start()
print(f" LSP enabled: {status['lsp']['enabled']}")
print(f" LSP ready: {runtime.lsp_ready}")
print(f" LSP status: {status['lsp']['status']}")
if not runtime.lsp_ready:
print(" Note: LSP not available, will use regex fallback")
print()
# 2. Create agent-centric tools
print("2. Creating agent-centric tools...")
tools = create_agent_centric_tools(runtime)
# Filter to show only LSP tools
lsp_tools = [t for t in tools if t.__name__.startswith('lsp_')]
print(f" LSP tools available:")
for tool in lsp_tools:
print(f" - {tool.__name__}")
print()
# 3. Create Agent with LSP-powered tools
print("3. Creating Agent with LSP tools...")
agent = Agent(
name="CodeAnalyzer",
instructions="""You are a code analysis assistant.
Use lsp_list_symbols to list functions and classes in files.
Use lsp_find_definition to find where symbols are defined.
Use lsp_find_references to find where symbols are used.
Provide clear, organized summaries of code structure.""",
tools=tools,
verbose=True
)
print()
# 4. Test 1: List symbols in a file
print("4. TEST 1: List symbols in agent/agent.py")
print("-" * 50)
result1 = agent.start("List all the main classes and public methods in agent/agent.py. Give me a brief summary.")
print("-" * 50)
print(f"Summary: {str(result1)[:300]}...")
print()
# 5. Test 2: Find definition of a symbol
print("5. TEST 2: Find where 'Agent' class is defined")
print("-" * 50)
result2 = agent.start("Where is the Agent class defined? Use lsp_find_definition to find it.")
print("-" * 50)
print(f"Result: {str(result2)[:300]}...")
print()
# 6. Test 3: Analyze main.py
print("6. TEST 3: Analyze main.py structure")
print("-" * 50)
result3 = agent.start("List all the functions in main.py and briefly describe what the file does.")
print("-" * 50)
print(f"Analysis: {str(result3)[:300]}...")
print()
# Cleanup
await runtime.stop()
print("=" * 60)
print("EXAMPLE COMPLETE")
print("=" * 60)
print()
print("Key takeaways:")
print("- LSP tools work through the Agent (agentic LSP)")
print("- Falls back to regex if LSP server not available")
print("- Returns structured data with citations")
if __name__ == "__main__":
asyncio.run(main())
+133
View File
@@ -0,0 +1,133 @@
"""
PraisonAI Output UX Examples
This example demonstrates the various output modes and configurations
available in PraisonAI agents.
Features demonstrated:
- verbose=True: Full output with agent info, task, and response panels
- verbose=False: Clean output with just the result
- Tool call display with consolidated output
- Multi-agent output with clear attribution
"""
import sys
sys.path.insert(0, '/Users/praison/praisonai-package/src/praisonai-agents')
from praisonaiagents import Agent, Task, PraisonAIAgents
def get_weather(city: str) -> str:
"""Get the current weather for a city.
Args:
city: The name of the city to get weather for
Returns:
Weather information string
"""
return f"The weather in {city} is sunny and 72°F"
def example_verbose_true():
"""Example with verbose=True (default) - shows full output panels."""
print("\n" + "=" * 60)
print("EXAMPLE 1: verbose=True (Full Output)")
print("=" * 60)
agent = Agent(
name="Assistant",
instructions="You are a helpful assistant. Be concise.",
verbose=True # Default - shows Agent Info, Task, and Response panels
)
result = agent.start("What is the capital of France? Answer in one sentence.")
print(f"\nReturned value: {result}")
def example_verbose_false():
"""Example with verbose=False - clean output, just the result."""
print("\n" + "=" * 60)
print("EXAMPLE 2: verbose=False (Clean Output)")
print("=" * 60)
agent = Agent(
name="Assistant",
instructions="You are a helpful assistant. Be concise.",
verbose=False # No panels, just returns the result
)
result = agent.start("What is 2+2? Answer in one word.")
print(f"Result: {result}")
def example_with_tools():
"""Example showing tool call output."""
print("\n" + "=" * 60)
print("EXAMPLE 3: Agent with Tools")
print("=" * 60)
agent = Agent(
name="Weather Bot",
instructions="You help users check the weather. Use the get_weather tool.",
tools=[get_weather],
verbose=True
)
result = agent.start("What's the weather in Tokyo?")
print(f"\nReturned value: {result}")
def example_multi_agent():
"""Example showing multi-agent output with clear attribution."""
print("\n" + "=" * 60)
print("EXAMPLE 4: Multi-Agent Output")
print("=" * 60)
researcher = Agent(
name="Researcher",
role="Research Analyst",
instructions="You research topics and provide factual information.",
verbose=True
)
writer = Agent(
name="Writer",
role="Content Writer",
instructions="You write clear summaries based on research.",
verbose=True
)
task1 = Task(
description="Research what AI is in one sentence",
expected_output="A brief factual description",
agent=researcher
)
task2 = Task(
description="Write a one-sentence summary about AI",
expected_output="A clear summary",
agent=writer
)
agents = PraisonAIAgents(
agents=[researcher, writer],
tasks=[task1, task2],
process="sequential",
verbose=1
)
result = agents.start()
print(f"\nFinal result: {result}")
if __name__ == "__main__":
print("=" * 60)
print("PRAISONAI OUTPUT UX EXAMPLES")
print("=" * 60)
# Run examples
example_verbose_true()
example_verbose_false()
example_with_tools()
example_multi_agent()
print("\n" + "=" * 60)
print("ALL EXAMPLES COMPLETE")
print("=" * 60)
@@ -0,0 +1,197 @@
"""
Isolated Registry Example - Multi-Agent Safety
This example demonstrates how to use isolated registries to prevent
provider registration collisions in multi-agent scenarios.
Key concepts:
- Each agent/context can have its own registry
- No global state mutation
- Safe for concurrent/parallel runs
- Different agents can use different provider implementations
Requirements:
- pip install praisonai
"""
import sys
sys.path.insert(0, '../../src/praisonai')
from praisonai.llm import (
LLMProviderRegistry,
create_llm_provider,
register_llm_provider,
get_default_llm_registry,
_reset_default_registry
)
# Define two different provider implementations with the same name
class Agent1Provider:
"""Provider implementation for Agent 1."""
provider_id = "custom"
def __init__(self, model_id, config=None):
self.model_id = model_id
self.config = config or {}
self.agent_name = "Agent1"
def generate(self, prompt):
return f"[{self.agent_name}] Response to: {prompt}"
class Agent2Provider:
"""Provider implementation for Agent 2 - different behavior."""
provider_id = "custom"
def __init__(self, model_id, config=None):
self.model_id = model_id
self.config = config or {}
self.agent_name = "Agent2"
def generate(self, prompt):
return f"[{self.agent_name}] Different response to: {prompt}"
def demonstrate_collision_problem():
"""Show what happens with global registry collisions."""
print("=== Problem: Global Registry Collision ===\n")
# Reset to clean state
_reset_default_registry()
# Agent 1 registers "custom" provider
register_llm_provider("custom", Agent1Provider)
print("Agent 1 registered 'custom' provider")
# Agent 2 tries to register same name - ERROR!
try:
register_llm_provider("custom", Agent2Provider)
print("Agent 2 registered 'custom' provider")
except ValueError as e:
print(f"Agent 2 FAILED: {e}")
print()
print("Problem: Both agents want to use 'custom' but with different implementations.")
print("Solution: Use isolated registries.\n")
def demonstrate_isolated_solution():
"""Show how isolated registries solve the collision problem."""
print("=== Solution: Isolated Registries ===\n")
# Create isolated registries for each agent
agent1_registry = LLMProviderRegistry()
agent2_registry = LLMProviderRegistry()
# Each agent registers its own "custom" provider
agent1_registry.register("custom", Agent1Provider)
agent2_registry.register("custom", Agent2Provider)
print("Agent 1 registry providers:", agent1_registry.list())
print("Agent 2 registry providers:", agent2_registry.list())
print()
# Create providers from each registry
provider1 = create_llm_provider("custom/model-v1", registry=agent1_registry)
provider2 = create_llm_provider("custom/model-v1", registry=agent2_registry)
# They have different implementations!
print("Provider 1 response:", provider1.generate("Hello"))
print("Provider 2 response:", provider2.generate("Hello"))
print()
# Verify isolation
print("Verification:")
print(f" provider1.agent_name = {provider1.agent_name}")
print(f" provider2.agent_name = {provider2.agent_name}")
print(f" Same class? {type(provider1) == type(provider2)}")
print()
def demonstrate_parallel_agents():
"""Simulate parallel agent execution with isolated registries."""
print("=== Parallel Agent Simulation ===\n")
def run_agent(agent_id, registry):
"""Simulate an agent run with its own registry."""
provider = create_llm_provider("custom/model", registry=registry)
result = provider.generate(f"Query from agent {agent_id}")
return result
# Setup registries
registries = {}
for i in range(3):
reg = LLMProviderRegistry()
# Each agent gets a slightly different provider
class DynamicProvider:
provider_id = "custom"
def __init__(self, model_id, config=None):
self.model_id = model_id
self.agent_id = i # Capture agent ID
def generate(self, prompt):
return f"Agent-{self.agent_id} processed: {prompt}"
reg.register("custom", DynamicProvider)
registries[i] = reg
# Simulate parallel execution
print("Simulating parallel agent runs:")
for agent_id, registry in registries.items():
result = run_agent(agent_id, registry)
print(f" {result}")
print()
print("Each agent used its own isolated registry - no collisions!")
print()
def demonstrate_shared_vs_isolated():
"""Compare shared (global) vs isolated registry patterns."""
print("=== Shared vs Isolated Patterns ===\n")
# Reset global registry
_reset_default_registry()
# Pattern 1: Shared (global) registry
print("Pattern 1: Shared Registry (use for app-wide providers)")
register_llm_provider("shared-provider", Agent1Provider)
# All code sees the same provider
global_reg = get_default_llm_registry()
print(f" Global registry has 'shared-provider': {global_reg.has('shared-provider')}")
# Pattern 2: Isolated registry
print("\nPattern 2: Isolated Registry (use for per-agent/per-run providers)")
isolated = LLMProviderRegistry()
isolated.register("isolated-provider", Agent2Provider)
# Only code with reference to this registry sees it
print(f" Isolated registry has 'isolated-provider': {isolated.has('isolated-provider')}")
print(f" Global registry has 'isolated-provider': {global_reg.has('isolated-provider')}")
print()
print("Recommendation:")
print(" - Use global registry for app-wide, stable providers")
print(" - Use isolated registries for per-agent or per-run customization")
print()
def main():
print("=" * 60)
print("Isolated Registry Example - Multi-Agent Safety")
print("=" * 60)
print()
demonstrate_collision_problem()
demonstrate_isolated_solution()
demonstrate_parallel_agents()
demonstrate_shared_vs_isolated()
print("=" * 60)
print("Example Complete")
print("=" * 60)
if __name__ == "__main__":
main()
+83
View File
@@ -0,0 +1,83 @@
# PraisonAI Profiling Examples
This directory contains examples for profiling PraisonAI agent performance.
## Prerequisites
```bash
pip install praisonai
export OPENAI_API_KEY=your_key_here
```
## Examples
### 1. Basic Profiling (`basic_profiling.py`)
Demonstrates programmatic profiling of a single query:
```bash
python basic_profiling.py
```
### 2. Suite Profiling (`suite_profiling.py`)
Runs a comprehensive profiling suite with multiple scenarios:
```bash
python suite_profiling.py
```
### 3. Optimization Demo (`optimization_example.py`)
Demonstrates Tier 0/1/2 performance optimizations:
```bash
python optimization_example.py
```
## CLI Commands
You can also profile directly from the command line:
```bash
# Profile a query
praisonai profile query "What is 2+2?"
# Profile with file grouping
praisonai profile query "Hello" --show-files --limit 20
# Profile imports
praisonai profile imports
# Profile startup time
praisonai profile startup
# Run comprehensive suite
praisonai profile suite --quick
# Create performance baseline
praisonai profile snapshot --baseline
# Compare against baseline
praisonai profile snapshot current --compare
# Show optimization status
praisonai profile optimize --show
```
## Output
- **Text output**: Human-readable timing breakdown
- **JSON output**: Machine-readable for CI/CD (`--format json`)
- **Artifacts**: Binary cProfile data and reports (`--save ./results`)
- **Snapshots**: Performance baselines for regression detection
## Performance Optimizations
Enable opt-in optimizations via environment variables:
```bash
export PRAISONAI_LITE_MODE=1
export PRAISONAI_SKIP_TYPE_VALIDATION=1
export PRAISONAI_MINIMAL_IMPORTS=1
```
@@ -0,0 +1,71 @@
"""
Basic Profiling Example for PraisonAI
This example demonstrates how to use the unified profiler programmatically
to measure agent performance.
Run with:
python basic_profiling.py
"""
import os
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
from praisonai.cli.execution import (
ExecutionRequest,
Profiler,
ProfilerConfig,
)
def main():
"""Run a profiled query and display results."""
# Configure profiler (Layer 1 = basic profiling with function stats)
config = ProfilerConfig(
layer=1, # 0=minimal, 1=basic, 2=deep
limit=15, # Show top 15 functions
show_callers=False, # Set True for caller info (layer 2)
show_callees=False, # Set True for callee info (layer 2)
)
# Create execution request
request = ExecutionRequest(
prompt="What is 2+2?",
agent_name="ProfiledAgent",
)
# Run profiled query
print("🔬 Running profiled query...")
profiler = Profiler(config)
result, report = profiler.profile_sync(request)
# Print formatted report
print("\n" + report.to_text())
# Access timing data programmatically
print("\n📊 Timing Summary:")
print(f" Imports: {report.timing.imports_ms:>10.2f} ms")
print(f" Agent Init: {report.timing.agent_init_ms:>10.2f} ms")
print(f" Execution: {report.timing.execution_ms:>10.2f} ms")
print(f" Total: {report.timing.total_ms:>10.2f} ms")
# Show top 5 functions
if report.functions:
print("\n🔥 Top 5 Functions by Cumulative Time:")
for i, func in enumerate(report.functions[:5], 1):
print(f" {i}. {func.name}: {func.cumulative_time_ms:.2f}ms")
# Show response
print(f"\n💬 Response: {result.output}")
# Export to JSON
print("\n📄 JSON Export available via: report.to_json()")
if __name__ == "__main__":
main()
@@ -0,0 +1,223 @@
"""
Performance Optimization Example for PraisonAI
This example demonstrates how to use the Tier 0/1/2 optimization features
programmatically.
Run with:
python optimization_example.py
"""
import os
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
def demo_provider_cache():
"""Demonstrate provider caching (Tier 0)."""
print("\n" + "=" * 60)
print("Tier 0: Provider Cache Demo")
print("=" * 60)
from praisonai.cli.features.profiler import (
get_provider_cache,
resolve_model_provider,
)
cache = get_provider_cache()
cache.invalidate() # Start fresh
# Resolve some models
models = ['gpt-4o', 'claude-3-opus', 'gemini-pro', 'gpt-4o-mini']
print("\nResolving models:")
for model in models:
provider, name = resolve_model_provider(model)
print(f" {model} -> {provider}/{name}")
# Check cache stats
stats = cache.get_stats()
print("\nCache Stats:")
print(f" Hits: {stats['hits']}")
print(f" Misses: {stats['misses']}")
print(f" Size: {stats['size']}")
def demo_lazy_importer():
"""Demonstrate lazy imports (Tier 0)."""
print("\n" + "=" * 60)
print("Tier 0: Lazy Importer Demo")
print("=" * 60)
from praisonai.cli.features.profiler import LazyImporter
import time
LazyImporter.clear()
# Time lazy import
start = time.perf_counter()
json_module = LazyImporter.get('json')
first_import = (time.perf_counter() - start) * 1000
# Second import (cached)
start = time.perf_counter()
json_module2 = LazyImporter.get('json')
cached_import = (time.perf_counter() - start) * 1000
print(f"\nFirst import: {first_import:.3f}ms")
print(f"Cached import: {cached_import:.3f}ms")
print(f"Same object: {json_module is json_module2}")
def demo_client_pool():
"""Demonstrate client pooling (Tier 1)."""
print("\n" + "=" * 60)
print("Tier 1: Client Pool Demo")
print("=" * 60)
from praisonai.cli.features.profiler import ClientPool
pool = ClientPool(max_size=3)
# Create some mock clients
def create_client(name):
return {'name': name, 'created': True}
# Get or create clients
client1 = pool.get_or_create('openai', lambda: create_client('openai'))
_ = pool.get_or_create('anthropic', lambda: create_client('anthropic'))
# Reuse existing client
client1_again = pool.get_or_create('openai', lambda: create_client('openai_new'))
print(f"\nPool size: {pool.size()}")
print(f"Client reused: {client1 is client1_again}")
print(f"Client1 name: {client1['name']}")
def demo_prewarm():
"""Demonstrate pre-warming (Tier 1)."""
print("\n" + "=" * 60)
print("Tier 1: Pre-warming Demo")
print("=" * 60)
from praisonai.cli.features.profiler import PrewarmManager
import time
PrewarmManager.reset()
print(f"\nPre-warming enabled: {PrewarmManager.is_enabled()}")
# Enable and prewarm
PrewarmManager.enable()
print(f"After enable: {PrewarmManager.is_enabled()}")
# Prewarm OpenAI (runs in background)
PrewarmManager.prewarm_provider('openai')
print("Pre-warming OpenAI in background...")
# Wait a bit for background thread
time.sleep(0.5)
print(f"OpenAI prewarmed: {PrewarmManager.is_prewarmed('openai')}")
PrewarmManager.reset()
def demo_lite_mode():
"""Demonstrate lite mode (Tier 2)."""
print("\n" + "=" * 60)
print("Tier 2: Lite Mode Demo")
print("=" * 60)
from praisonai.cli.features.profiler import (
LiteModeConfig,
is_lite_mode,
)
# Default config
config = LiteModeConfig()
print(f"\nDefault lite mode: {config.enabled}")
# Config from env
print("\nTo enable lite mode, set environment variables:")
print(" export PRAISONAI_LITE_MODE=1")
print(" export PRAISONAI_SKIP_TYPE_VALIDATION=1")
print(f"\nCurrent is_lite_mode(): {is_lite_mode()}")
def demo_perf_snapshot():
"""Demonstrate performance snapshots (Tier 2)."""
print("\n" + "=" * 60)
print("Tier 2: Performance Snapshot Demo")
print("=" * 60)
from praisonai.cli.features.profiler import (
PerfSnapshot,
PerfComparison,
PerfSnapshotManager,
)
import tempfile
from pathlib import Path
# Create manager with temp directory
with tempfile.TemporaryDirectory() as tmpdir:
manager = PerfSnapshotManager(snapshot_dir=Path(tmpdir))
# Create baseline
baseline = PerfSnapshot(
timestamp='2025-01-01T00:00:00Z',
name='baseline',
startup_cold_ms=100.0,
startup_warm_ms=80.0,
import_time_ms=500.0,
query_time_ms=2000.0,
first_token_ms=300.0,
)
manager.save_baseline(baseline)
print("\nBaseline saved")
# Create current snapshot (simulated improvement)
current = PerfSnapshot(
timestamp='2025-01-02T00:00:00Z',
name='current',
startup_cold_ms=95.0, # 5% faster
startup_warm_ms=78.0,
import_time_ms=480.0, # 4% faster
query_time_ms=1900.0, # 5% faster
first_token_ms=280.0,
)
# Compare
comparison = PerfComparison(baseline=baseline, current=current)
print(f"\nStartup improvement: {-comparison.startup_cold_diff_pct:.1f}%")
print(f"Import improvement: {-comparison.import_time_diff_pct:.1f}%")
print(f"Query improvement: {-comparison.query_time_diff_pct:.1f}%")
print(f"Is regression: {comparison.is_regression()}")
def main():
"""Run all demos."""
print("=" * 60)
print("PraisonAI Performance Optimization Demo")
print("=" * 60)
demo_provider_cache()
demo_lazy_importer()
demo_client_pool()
demo_prewarm()
demo_lite_mode()
demo_perf_snapshot()
print("\n" + "=" * 60)
print("All demos complete!")
print("=" * 60)
if __name__ == "__main__":
main()
@@ -0,0 +1,227 @@
{
"timestamp": "2025-12-31T18:45:38.187539Z",
"metadata": {
"python_version": "3.12.11",
"platform": "macOS-15.7.4-arm64-arm-64bit",
"praisonai_version": "2.9.2",
"timestamp": "2025-12-31T18:45:38.226558Z"
},
"startup": {
"cold_ms": 80.9305000002496,
"warm_ms": 55.266541996388696
},
"import_analysis": [
{
"module": "praisonaiagents",
"self_us": 2419,
"cumulative_us": 1645390,
"self_ms": 2.419,
"cumulative_ms": 1645.39
},
{
"module": "praisonaiagents.workflows",
"self_us": 670,
"cumulative_us": 1375919,
"self_ms": 0.67,
"cumulative_ms": 1375.919
},
{
"module": "praisonaiagents.workflows.yaml_parser",
"self_us": 775,
"cumulative_us": 1369435,
"self_ms": 0.775,
"cumulative_ms": 1369.435
},
{
"module": "praisonaiagents.agent.agent",
"self_us": 19,
"cumulative_us": 1349174,
"self_ms": 0.019,
"cumulative_ms": 1349.174
},
{
"module": "praisonaiagents.agent",
"self_us": 537,
"cumulative_us": 1349155,
"self_ms": 0.537,
"cumulative_ms": 1349.155
},
{
"module": "praisonaiagents.agent.agent",
"self_us": 1453,
"cumulative_us": 1285299,
"self_ms": 1.453,
"cumulative_ms": 1285.299
},
{
"module": "praisonaiagents.llm.openai_client",
"self_us": 7887,
"cumulative_us": 1123578,
"self_ms": 7.887,
"cumulative_ms": 1123.578
},
{
"module": "openai",
"self_us": 1083,
"cumulative_us": 1115692,
"self_ms": 1.083,
"cumulative_ms": 1115.692
},
{
"module": "openai.types",
"self_us": 2606,
"cumulative_us": 957990,
"self_ms": 2.606,
"cumulative_ms": 957.99
},
{
"module": "openai.types.batch",
"self_us": 1787,
"cumulative_us": 485667,
"self_ms": 1.787,
"cumulative_ms": 485.667
},
{
"module": "openai._models",
"self_us": 13423,
"cumulative_us": 450798,
"self_ms": 13.423,
"cumulative_ms": 450.798
},
{
"module": "openai.types.eval_create_params",
"self_us": 644,
"cumulative_us": 393275,
"self_ms": 0.644,
"cumulative_ms": 393.275
},
{
"module": "openai.types.graders.grader_inputs_param",
"self_us": 17,
"cumulative_us": 359515,
"self_ms": 0.017,
"cumulative_ms": 359.515
},
{
"module": "openai.types.graders",
"self_us": 392,
"cumulative_us": 359498,
"self_ms": 0.392,
"cumulative_ms": 359.498
},
{
"module": "openai.types.graders.multi_grader",
"self_us": 3847,
"cumulative_us": 351251,
"self_ms": 3.847,
"cumulative_ms": 351.251
},
{
"module": "openai.types.graders.label_model_grader",
"self_us": 5450,
"cumulative_us": 337291,
"self_ms": 5.45,
"cumulative_ms": 337.291
},
{
"module": "openai.types.graders.grader_inputs",
"self_us": 545,
"cumulative_us": 331841,
"self_ms": 0.545,
"cumulative_ms": 331.841
},
{
"module": "openai.types.responses.response_input_text",
"self_us": 21,
"cumulative_us": 331297,
"self_ms": 0.021,
"cumulative_ms": 331.297
},
{
"module": "openai.types.responses",
"self_us": 1301,
"cumulative_us": 331277,
"self_ms": 1.301,
"cumulative_ms": 331.277
},
{
"module": "rich.logging",
"self_us": 1066,
"cumulative_us": 238256,
"self_ms": 1.066,
"cumulative_ms": 238.256
}
],
"scenarios": [
{
"name": "simple_query",
"config": {
"prompt": "What is 2+2?",
"model": "default",
"stream": false,
"iterations": 2
},
"import_time_stats": {
"mean": 0.008833496394800022,
"min": 0.005916997906751931,
"max": 0.011749994882848114,
"stdev": 0.0041245517164382375,
"median": 0.008833496394800022
},
"total_time_stats": {
"mean": 1727.4090415012324,
"min": 1367.2264170018025,
"max": 2087.5916660006624,
"stdev": 509.3751524982296,
"median": 1727.4090415012324
},
"first_token_stats": null,
"raw_import_times_ms": [
0.011749994882848114,
0.005916997906751931
],
"raw_total_times_ms": [
1367.2264170018025,
2087.5916660006624
]
},
{
"name": "simple_stream",
"config": {
"prompt": "What is 2+2?",
"model": "default",
"stream": true,
"iterations": 2
},
"import_time_stats": {
"mean": 0.008792001608526334,
"min": 0.006375004886649549,
"max": 0.01120899833040312,
"stdev": 0.0034181495442894613,
"median": 0.008792001608526334
},
"total_time_stats": {
"mean": 2321.0982290001994,
"min": 1697.1089580038097,
"max": 2945.087499996589,
"stdev": 882.4540898183949,
"median": 2321.0982290001994
},
"first_token_stats": {
"mean": 2312.8511880022415,
"min": 1681.2440420035273,
"max": 2944.458334000956,
"stdev": 893.2273919631452,
"median": 2312.8511880022415
},
"raw_import_times_ms": [
0.01120899833040312,
0.006375004886649549
],
"raw_total_times_ms": [
2945.087499996589,
1697.1089580038097
]
}
]
}
@@ -0,0 +1,80 @@
"""
Profile Suite Example for PraisonAI
This example demonstrates how to run a comprehensive profiling suite
programmatically with multiple scenarios.
Run with:
python suite_profiling.py
"""
import os
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
from praisonai.cli.features.profiler import (
ScenarioConfig,
run_profile_suite,
)
def main():
"""Run a profiling suite with custom scenarios."""
# Define custom scenarios
scenarios = [
ScenarioConfig(
name="simple_query",
prompt="What is 2+2?",
stream=False,
iterations=2,
warmup=1,
),
ScenarioConfig(
name="simple_stream",
prompt="What is 2+2?",
stream=True,
iterations=2,
warmup=1,
),
]
# Run suite with custom output directory
print("🔬 Running Profile Suite...")
result = run_profile_suite(
output_dir="./profile_results",
scenarios=scenarios,
verbose=True,
)
# Print summary
print("\n" + "=" * 60)
print("Profile Suite Results")
print("=" * 60)
print("\n📊 Startup Times:")
print(f" Cold Start: {result.startup_cold_ms:.2f}ms")
print(f" Warm Start: {result.startup_warm_ms:.2f}ms")
if result.import_analysis:
print("\n📦 Top Import:")
print(f" {result.import_analysis[0]['module']}: {result.import_analysis[0]['cumulative_ms']:.2f}ms")
print("\n🎯 Scenario Results:")
for scenario in result.scenarios:
stats = scenario.get_stats(scenario.total_times)
print(f" {scenario.name}:")
print(f" Mean: {stats['mean']:.2f}ms")
print(f" Min: {stats['min']:.2f}ms")
print(f" Max: {stats['max']:.2f}ms")
print("\n✅ Results saved to: ./profile_results/")
print(" - suite_results.json (machine-readable)")
print(" - suite_report.txt (human-readable)")
if __name__ == "__main__":
main()
@@ -0,0 +1,165 @@
"""
Session Persistence Example
Demonstrates automatic session persistence with zero configuration.
When you provide a session_id to an Agent, conversation history is
automatically saved to disk and restored on subsequent runs.
Usage:
python session_persistence_example.py
Requirements:
- praisonaiagents package installed
- OPENAI_API_KEY environment variable set
"""
from praisonaiagents import Agent
def example_basic_persistence():
"""Basic session persistence example."""
print("=" * 60)
print("EXAMPLE 1: Basic Session Persistence")
print("=" * 60)
session_id = "demo-session-basic"
# Create agent with session_id
agent = Agent(
name="Assistant",
instructions="You are a helpful assistant. Be very brief.",
session_id=session_id,
verbose=False,
)
# First message
print("\n[User]: My favorite color is blue and my lucky number is 7.")
response1 = agent.chat("My favorite color is blue and my lucky number is 7.")
print(f"[Assistant]: {response1}")
# Second message - agent remembers
print("\n[User]: What is my favorite color?")
response2 = agent.chat("What is my favorite color?")
print(f"[Assistant]: {response2}")
print(f"\n✅ Session saved to: ~/.praison/sessions/{session_id}.json")
print(f" Chat history length: {len(agent.chat_history)} messages")
def example_session_restoration():
"""Demonstrates session restoration across Agent instances."""
print("\n" + "=" * 60)
print("EXAMPLE 2: Session Restoration")
print("=" * 60)
session_id = "demo-session-restore"
# First agent instance
print("\n--- First Agent Instance ---")
agent1 = Agent(
name="MemoryBot",
instructions="You are a helpful assistant. Be very brief.",
session_id=session_id,
verbose=False,
)
print("[User]: Remember that my pet's name is Max.")
response1 = agent1.chat("Remember that my pet's name is Max.")
print(f"[Assistant]: {response1}")
# Simulate new process by creating new agent instance
print("\n--- New Agent Instance (simulating restart) ---")
agent2 = Agent(
name="MemoryBot",
instructions="You are a helpful assistant. Be very brief.",
session_id=session_id,
verbose=False,
)
print("[User]: What is my pet's name?")
response2 = agent2.chat("What is my pet's name?")
print(f"[Assistant]: {response2}")
if "max" in response2.lower():
print("\n✅ Session restored successfully! Agent remembered 'Max'.")
else:
print("\n⚠️ Agent may not have remembered correctly.")
def example_in_memory_only():
"""Demonstrates in-memory memory without persistence."""
print("\n" + "=" * 60)
print("EXAMPLE 3: In-Memory Only (No session_id)")
print("=" * 60)
# Agent without session_id - in-memory only
agent = Agent(
name="TempBot",
instructions="You are a helpful assistant. Be very brief.",
verbose=False,
)
print("\n[User]: My name is Alice.")
response1 = agent.chat("My name is Alice.")
print(f"[Assistant]: {response1}")
print("\n[User]: What is my name?")
response2 = agent.chat("What is my name?")
print(f"[Assistant]: {response2}")
print(f"\n✅ In-memory chat history: {len(agent.chat_history)} messages")
print(" (This will be lost when the Agent instance is garbage collected)")
def example_direct_store_access():
"""Demonstrates direct access to the session store."""
print("\n" + "=" * 60)
print("EXAMPLE 4: Direct Session Store Access")
print("=" * 60)
from praisonaiagents.session import get_default_session_store
store = get_default_session_store()
session_id = "demo-direct-access"
# Add messages directly
store.add_user_message(session_id, "Hello from direct access!")
store.add_assistant_message(session_id, "Hi! I received your message.")
# Get history
history = store.get_chat_history(session_id)
print(f"\n✅ Added {len(history)} messages directly to session store")
for msg in history:
print(f" [{msg['role']}]: {msg['content']}")
# List all sessions
sessions = store.list_sessions()
print(f"\n📋 Total sessions in store: {len(sessions)}")
# Clean up demo session
store.delete_session(session_id)
print(f" Deleted demo session: {session_id}")
def main():
"""Run all examples."""
print("\n" + "=" * 60)
print("SESSION PERSISTENCE EXAMPLES")
print("=" * 60)
# Run examples
example_basic_persistence()
example_session_restoration()
example_in_memory_only()
example_direct_store_access()
print("\n" + "=" * 60)
print("ALL EXAMPLES COMPLETE")
print("=" * 60)
print("\nSession files are stored in: ~/.praison/sessions/")
print("You can inspect them with: cat ~/.praison/sessions/<session_id>.json")
if __name__ == "__main__":
main()
+30
View File
@@ -0,0 +1,30 @@
"""
Basic TUI Usage Example for PraisonAI.
This example demonstrates how to launch the interactive TUI.
"""
import os
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
def main():
"""Launch the TUI."""
try:
from praisonai.cli.features.tui import run_tui
# Launch with default settings
run_tui()
except ImportError:
print("Textual is required for TUI. Install with:")
print(" pip install praisonai[tui]")
exit(1)
if __name__ == "__main__":
main()
+134
View File
@@ -0,0 +1,134 @@
"""
Headless TUI Simulation Example for PraisonAI.
This example demonstrates how to run TUI simulations without
launching the interactive UI - useful for testing and CI.
"""
import asyncio
from praisonai.cli.features.tui import (
TuiOrchestrator,
SimulationRunner,
MockProvider,
)
from praisonai.cli.features.tui.orchestrator import OutputMode
from praisonai.cli.features.queue import QueueConfig
async def run_simulation():
"""Run a headless simulation."""
# Create orchestrator with mock provider (no real API calls)
config = QueueConfig(enable_persistence=False)
orchestrator = TuiOrchestrator(
queue_config=config,
output_mode=OutputMode.PRETTY,
debug=True,
)
# Create simulation runner
runner = SimulationRunner(orchestrator, assert_mode=True)
# Define simulation script
script = {
"session_id": "sim-example",
"model": "gpt-4o-mini",
"steps": [
# Navigate to different screens
{"action": "navigate", "args": {"screen": "main"}},
{"action": "focus", "args": {"widget": "composer"}},
# Change model
{
"action": "model",
"args": {"model": "gpt-4"},
"expected": {"model": "gpt-4"}
},
# Navigate to queue screen
{"action": "navigate", "args": {"screen": "queue"}},
# Take a snapshot
{"action": "snapshot"},
# Navigate back
{"action": "navigate", "args": {"screen": "main"}},
# Small delay
{"action": "sleep", "args": {"seconds": 0.5}},
]
}
# Run the simulation
print("Running headless TUI simulation...")
print("=" * 50)
success = await runner.run_script(script)
# Print results
print("=" * 50)
summary = runner.get_summary()
print(f"\nSimulation {'PASSED' if success else 'FAILED'}")
print(f" Assertions passed: {summary['assertions_passed']}")
print(f" Assertions failed: {summary['assertions_failed']}")
if summary['errors']:
print("\nErrors:")
for error in summary['errors']:
print(f" - {error}")
return success
async def run_with_events():
"""Run simulation and capture all events."""
config = QueueConfig(enable_persistence=False)
orchestrator = TuiOrchestrator(
queue_config=config,
output_mode=OutputMode.SILENT,
)
# Capture events
events = []
orchestrator.add_event_callback(lambda e: events.append(e))
await orchestrator.start(session_id="event-capture")
# Perform some actions
orchestrator.set_model("gpt-4")
orchestrator.navigate_screen("queue")
orchestrator.set_focus("queue-panel")
orchestrator.navigate_screen("main")
await orchestrator.stop()
# Print captured events
print("\nCaptured Events:")
print("-" * 40)
for event in events:
print(f" {event.event_type.value}")
return events
def main():
"""Main entry point."""
print("PraisonAI TUI Headless Simulation Example")
print("=" * 50)
# Run simulation
success = asyncio.run(run_simulation())
# Run event capture
print("\n" + "=" * 50)
events = asyncio.run(run_with_events())
print(f"\nTotal events captured: {len(events)}")
return 0 if success else 1
if __name__ == "__main__":
exit(main())
+122
View File
@@ -0,0 +1,122 @@
"""
Queue Management Example for PraisonAI.
This example demonstrates how to use the queue system programmatically.
"""
import asyncio
from praisonai.cli.features.queue import (
QueueManager,
QueueConfig,
QueuedRun,
RunState,
RunPriority,
)
async def main():
"""Demonstrate queue management."""
print("PraisonAI Queue Management Example")
print("=" * 50)
# Create queue configuration
config = QueueConfig(
max_concurrent_global=2,
max_concurrent_per_agent=1,
max_queue_size=10,
enable_persistence=False, # Use in-memory for demo
)
# Track outputs
outputs = {}
async def on_output(run_id: str, chunk: str):
if run_id not in outputs:
outputs[run_id] = ""
outputs[run_id] += chunk
print(f" [{run_id[:8]}] {chunk}", end="", flush=True)
async def on_complete(run_id: str, run: QueuedRun):
print(f"\n [{run_id[:8]}] Completed: {run.state.value}")
async def on_error(run_id: str, error: Exception):
print(f"\n [{run_id[:8]}] Error: {error}")
# Create queue manager
manager = QueueManager(
config=config,
on_output=on_output,
on_complete=on_complete,
on_error=on_error,
)
# Start the manager
await manager.start(recover=False)
print("Queue manager started")
# Submit some runs with different priorities
print("\nSubmitting runs...")
run1 = await manager.submit(
input_content="Low priority task",
agent_name="Agent1",
priority=RunPriority.LOW,
)
print(f" Submitted run1: {run1} (LOW priority)")
run2 = await manager.submit(
input_content="High priority task",
agent_name="Agent2",
priority=RunPriority.HIGH,
)
print(f" Submitted run2: {run2} (HIGH priority)")
run3 = await manager.submit(
input_content="Normal priority task",
agent_name="Agent1",
priority=RunPriority.NORMAL,
)
print(f" Submitted run3: {run3} (NORMAL priority)")
# Check queue status
print(f"\nQueue status:")
print(f" Queued: {manager.queued_count}")
print(f" Running: {manager.running_count}")
# List runs
runs = manager.list_runs()
print(f"\nAll runs ({len(runs)}):")
for run in runs:
print(f" {run.run_id[:8]}: {run.agent_name} - {run.state.value} ({run.priority.name})")
# Wait a bit for processing
print("\nWaiting for processing...")
await asyncio.sleep(2)
# Cancel a run
if runs:
cancel_id = runs[-1].run_id
print(f"\nCancelling run: {cancel_id[:8]}")
await manager.cancel(cancel_id)
# Get stats
stats = manager.get_stats()
print(f"\nQueue statistics:")
print(f" Total runs: {stats.total_runs}")
print(f" Queued: {stats.queued_count}")
print(f" Running: {stats.running_count}")
print(f" Succeeded: {stats.succeeded_count}")
print(f" Failed: {stats.failed_count}")
print(f" Cancelled: {stats.cancelled_count}")
# Stop the manager
await manager.stop()
print("\nQueue manager stopped")
print("\n" + "=" * 50)
print("Example complete!")
if __name__ == "__main__":
asyncio.run(main())
+84
View File
@@ -0,0 +1,84 @@
/**
* Exa Web Search Tool Example
*
* Demonstrates semantic web search with Exa AI.
*
* Prerequisites:
* npm install praisonai-ts @exalabs/ai-sdk
* export EXA_API_KEY=your-api-key
* export OPENAI_API_KEY=your-api-key
*/
// For local development, use relative imports
// For production, use: import { Agent, exaSearch } from 'praisonai-ts'
import { Agent } from '../../../src/praisonai-ts/src';
import { exaSearch } from '../../../src/praisonai-ts/src/tools/builtins';
async function main() {
console.log('=== Exa Web Search Tool Example ===\n');
// Example 1: Basic Semantic Search
console.log('1. Basic Semantic Search');
console.log('------------------------');
const searchAgent = new Agent({
name: 'SemanticSearcher',
instructions: 'You search the web using semantic search and provide detailed summaries.',
tools: [exaSearch({ numResults: 5, type: 'auto' })],
});
const searchResult = await searchAgent.chat('Find the latest AI research breakthroughs');
console.log('Search Result:', searchResult);
console.log();
// Example 2: Company Research
console.log('2. Company Research');
console.log('-------------------');
const companyAgent = new Agent({
name: 'CompanyResearcher',
instructions: 'You research companies and provide analysis.',
tools: [
exaSearch({
type: 'auto',
numResults: 5,
category: 'company',
contents: {
text: { maxCharacters: 1000 },
summary: true,
livecrawl: 'preferred',
},
}),
],
});
const companyResult = await companyAgent.chat(
'Find AI startups in Europe founded after 2020'
);
console.log('Company Result:', companyResult);
console.log();
// Example 3: Academic Paper Search
console.log('3. Academic Paper Search');
console.log('------------------------');
const academicAgent = new Agent({
name: 'AcademicResearcher',
instructions: 'You search for academic papers and summarize findings.',
tools: [
exaSearch({
type: 'neural',
category: 'research paper',
numResults: 5,
includeDomains: ['arxiv.org', 'nature.com', 'science.org'],
}),
],
});
const academicResult = await academicAgent.chat(
'Find recent papers on transformer architectures and attention mechanisms'
);
console.log('Academic Result:', academicResult);
}
main().catch(console.error);
@@ -0,0 +1,86 @@
/**
* Tavily Search Tool Example
*
* Demonstrates web search, content extraction, and crawling with Tavily.
*
* Prerequisites:
* npm install praisonai @tavily/ai-sdk
* export TAVILY_API_KEY=your-api-key
* export OPENAI_API_KEY=your-api-key
*/
// For local development, use relative imports
// For production, use: import { Agent, tavilySearch } from 'praisonai-ts'
import { Agent } from '../../../src/praisonai-ts/src';
import { tavilySearch, tavilyExtract, tavilyCrawl } from '../../../src/praisonai-ts/src/tools/builtins';
async function main() {
console.log('=== Tavily Search Tool Example ===\n');
// Example 1: Basic Web Search
console.log('1. Basic Web Search');
console.log('-------------------');
const searchAgent = new Agent({
name: 'WebSearcher',
instructions: 'You search the web for information and provide concise summaries with sources.',
tools: [tavilySearch({ maxResults: 5, includeAnswer: true })],
});
const searchResult = await searchAgent.chat('What are the latest developments in AI agents?');
console.log('Search Result:', searchResult);
console.log();
// Example 2: News Search with Time Filter
console.log('2. News Search with Time Filter');
console.log('-------------------------------');
const newsAgent = new Agent({
name: 'NewsSearcher',
instructions: 'You search for recent news and summarize key points.',
tools: [
tavilySearch({
topic: 'news',
timeRange: 'week',
maxResults: 5,
}),
],
});
const newsResult = await newsAgent.chat('Latest news about OpenAI');
console.log('News Result:', newsResult);
console.log();
// Example 3: Content Extraction
console.log('3. Content Extraction');
console.log('--------------------');
const extractAgent = new Agent({
name: 'ContentExtractor',
instructions: 'You extract and summarize content from web pages.',
tools: [tavilyExtract()],
});
const extractResult = await extractAgent.chat(
'Extract the main content from https://docs.tavily.com and summarize it'
);
console.log('Extract Result:', extractResult);
console.log();
// Example 4: Website Crawling
console.log('4. Website Crawling');
console.log('------------------');
const crawlAgent = new Agent({
name: 'WebCrawler',
instructions: 'You crawl websites and analyze their structure and content.',
tools: [tavilyCrawl({ maxDepth: 2 })],
});
const crawlResult = await crawlAgent.chat(
'Crawl https://docs.tavily.com and list the main documentation sections'
);
console.log('Crawl Result:', crawlResult);
}
main().catch(console.error);
@@ -189,7 +189,7 @@ def __getattr__(name):
_lazy_cache[name] = AutoAgents
return AutoAgents
# Session (imports requests)
# Session (imports requests lazily via session/api.py)
elif name == "Session":
from .session import Session
_lazy_cache[name] = Session
@@ -309,7 +309,9 @@ class Agent:
session_id: Optional[str] = None,
hooks: Optional[List[Any]] = None,
llm_config: Optional[Dict[str, Any]] = None,
rate_limiter: Optional[Any] = None
rate_limiter: Optional[Any] = None,
auto_summarize: bool = False,
summarize_threshold: float = 0.8
):
"""Initialize an Agent instance.
@@ -461,6 +463,12 @@ class Agent:
# Store rate limiter (optional, zero overhead when None)
self._rate_limiter = rate_limiter
# Auto-summarization configuration
if auto_summarize and not (0.0 < summarize_threshold <= 1.0):
raise ValueError("summarize_threshold must be between 0.0 and 1.0")
self.auto_summarize = auto_summarize
self.summarize_threshold = summarize_threshold
# Store OpenAI client parameters for lazy initialization
self._openai_api_key = api_key
self._openai_base_url = base_url
@@ -687,6 +695,11 @@ Your Goal: {self.goal}
self._db = db
self._session_id = session_id
self._db_initialized = False
# Session store for JSON-based persistence (lazy - no imports until used)
# Used when session_id is provided but no DB adapter
self._session_store = None
self._session_store_initialized = False
@property
def console(self):
@@ -2094,6 +2107,49 @@ Your Goal: {self.goal}"""
self._db_initialized = True
self._current_run_id = None # Track current run
def _init_session_store(self):
"""
Initialize session store for JSON-based persistence (lazy, first chat only).
This is used when session_id is provided but no DB adapter.
Enables automatic session persistence with zero configuration.
"""
if self._session_store_initialized:
return
# Only initialize if session_id is provided and no DB adapter
if self._session_id is None or self._db is not None:
self._session_store_initialized = True
return
try:
from ..session import get_default_session_store
self._session_store = get_default_session_store()
# Restore chat history from previous session
history = self._session_store.get_chat_history(self._session_id)
if history:
# Only restore if chat_history is empty (avoid duplicates)
if not self.chat_history:
for msg in history:
self.chat_history.append({
"role": msg["role"],
"content": msg["content"]
})
logging.info(f"Restored session {self._session_id} with {len(history)} messages from JSON store")
# Set agent info
self._session_store.set_agent_info(
self._session_id,
agent_name=self.name,
user_id=self.user_id,
)
except Exception as e:
logging.warning(f"Failed to initialize session store: {e}")
self._session_store = None
self._session_store_initialized = True
def _start_run(self, input_content: str):
"""Start a new run (turn) for persistence tracking."""
@@ -2135,17 +2191,27 @@ Your Goal: {self.goal}"""
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:
"""Persist a message to the DB or session store."""
# Try DB adapter first
if self._db is not None:
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 to DB: {e}")
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}")
# Fall back to session store (JSON-based)
if self._session_store is not None and self._session_id is not None:
try:
if role == "user":
self._session_store.add_user_message(self._session_id, content)
elif role == "assistant":
self._session_store.add_assistant_message(self._session_id, content)
except Exception as e:
logging.warning(f"Failed to persist message to session store: {e}")
@property
def session_id(self) -> Optional[str]:
@@ -2160,6 +2226,10 @@ Your Goal: {self.goal}"""
# Initialize DB session on first chat (lazy)
self._init_db_session()
# Initialize session store for JSON-based persistence (lazy)
# This enables automatic session persistence when session_id is provided
self._init_session_store()
# Start a new run for this chat turn
prompt_str = prompt if isinstance(prompt, str) else str(prompt)
self._start_run(prompt_str)
@@ -3061,6 +3131,31 @@ Write the complete compiled report:"""
# Return the single result if only one step
return results[0]["result"] if results else None
def switch_model(self, new_model: str) -> None:
"""
Switch the agent's LLM model while preserving conversation history.
Args:
new_model: The new model name to switch to (e.g., "gpt-4o", "claude-3-sonnet")
"""
# Store the new model name
self.llm = new_model
# Recreate the LLM instance with the new model
try:
from ..llm.llm import LLM
self._llm_instance = LLM(
model=new_model,
base_url=self._openai_base_url,
api_key=self._openai_api_key,
)
self._using_custom_llm = True
except ImportError:
# If LLM class not available, just update the model string
pass
# Chat history is preserved in self.chat_history (no action needed)
def start(self, prompt: str, **kwargs):
"""Start the agent with a prompt. This is a convenience method that wraps chat()."""
# Load history from past sessions if history_in_context is set
@@ -0,0 +1,202 @@
"""
Message Queue for PraisonAI Agents.
Provides queuing for agent prompts when agent is busy.
Thread-safe implementation for concurrent access.
"""
import threading
import heapq
import time
import logging
from typing import Optional, Any, List, Tuple
from dataclasses import dataclass, field
from enum import Enum
logger = logging.getLogger(__name__)
class MessagePriority(Enum):
"""Message priority levels."""
LOW = 1
NORMAL = 5
HIGH = 10
URGENT = 20
@dataclass(order=True)
class QueuedMessage:
"""A message in the queue with priority and timestamp."""
priority: int
timestamp: float = field(compare=False)
content: Any = field(compare=False)
metadata: dict = field(default_factory=dict, compare=False)
def __post_init__(self):
# Negate priority for max-heap behavior (higher priority = lower number in heap)
self.priority = -self.priority
class AgentMessageQueue:
"""
Thread-safe message queue for agent prompts.
Supports priority-based ordering and FIFO within same priority.
Usage:
queue = AgentMessageQueue()
queue.enqueue("Hello", priority=5)
queue.enqueue("Urgent!", priority=10)
msg = queue.dequeue() # Returns "Urgent!" first
"""
def __init__(self, max_size: int = 100):
"""
Initialize message queue.
Args:
max_size: Maximum number of messages in queue (0 = unlimited)
"""
self._queue: List[QueuedMessage] = []
self._lock = threading.RLock()
self._not_empty = threading.Condition(self._lock)
self._max_size = max_size
self._message_count = 0
def enqueue(
self,
content: Any,
priority: int = MessagePriority.NORMAL.value,
metadata: Optional[dict] = None
) -> bool:
"""
Add a message to the queue.
Args:
content: Message content
priority: Message priority (higher = more urgent)
metadata: Optional metadata dict
Returns:
True if message was added, False if queue is full
"""
with self._lock:
if self._max_size > 0 and len(self._queue) >= self._max_size:
logger.warning(f"Message queue full (max_size={self._max_size})")
return False
message = QueuedMessage(
priority=priority,
timestamp=time.time(),
content=content,
metadata=metadata or {}
)
heapq.heappush(self._queue, message)
self._message_count += 1
self._not_empty.notify()
logger.debug(f"Message enqueued (priority={priority}, queue_size={len(self._queue)})")
return True
def dequeue(self, timeout: Optional[float] = None) -> Optional[Any]:
"""
Remove and return the highest priority message.
Args:
timeout: Maximum time to wait for a message (None = don't wait)
Returns:
Message content or None if queue is empty
"""
with self._not_empty:
if not self._queue:
if timeout is not None:
self._not_empty.wait(timeout)
if not self._queue:
return None
message = heapq.heappop(self._queue)
logger.debug(f"Message dequeued (queue_size={len(self._queue)})")
return message.content
def peek(self) -> Optional[Any]:
"""Return the highest priority message without removing it."""
with self._lock:
if not self._queue:
return None
return self._queue[0].content
def is_empty(self) -> bool:
"""Check if queue is empty."""
with self._lock:
return len(self._queue) == 0
def size(self) -> int:
"""Get current queue size."""
with self._lock:
return len(self._queue)
def clear(self) -> int:
"""Clear all messages from queue. Returns number of cleared messages."""
with self._lock:
count = len(self._queue)
self._queue.clear()
return count
def get_all(self) -> List[Any]:
"""Get all messages in priority order without removing them."""
with self._lock:
# Sort by priority (remember priority is negated)
sorted_queue = sorted(self._queue)
return [msg.content for msg in sorted_queue]
@property
def total_messages_processed(self) -> int:
"""Total number of messages ever enqueued."""
return self._message_count
class AsyncAgentMessageQueue:
"""
Async-compatible message queue for agent prompts.
Wraps AgentMessageQueue with async interface.
"""
def __init__(self, max_size: int = 100):
self._sync_queue = AgentMessageQueue(max_size=max_size)
async def enqueue(
self,
content: Any,
priority: int = MessagePriority.NORMAL.value,
metadata: Optional[dict] = None
) -> bool:
"""Async enqueue."""
return self._sync_queue.enqueue(content, priority, metadata)
async def dequeue(self, timeout: Optional[float] = None) -> Optional[Any]:
"""Async dequeue."""
import asyncio
if timeout is None:
return self._sync_queue.dequeue(timeout=0)
# Poll with async sleep for timeout
start = time.time()
while time.time() - start < timeout:
result = self._sync_queue.dequeue(timeout=0)
if result is not None:
return result
await asyncio.sleep(0.01)
return None
def is_empty(self) -> bool:
return self._sync_queue.is_empty()
def size(self) -> int:
return self._sync_queue.size()
@@ -0,0 +1,226 @@
"""
Automatic Conversation Summarization for PraisonAI Agents.
Provides automatic summarization when context window fills up.
Reuses existing telemetry/token_collector for token tracking.
"""
import logging
from typing import Optional, List, Dict, Any, Callable
from dataclasses import dataclass, field
logger = logging.getLogger(__name__)
@dataclass
class SummarizationConfig:
"""Configuration for auto-summarization."""
enabled: bool = False
threshold: float = 0.8 # 80% of context window
context_window: int = 128000 # Default for GPT-4
summary_model: Optional[str] = None # Use smaller model for summaries
preserve_system: bool = True # Keep system message
preserve_recent: int = 2 # Keep last N message pairs
class SummarizationManager:
"""
Manages automatic conversation summarization.
Tracks token usage and triggers summarization when threshold is reached.
Thread-safe for multi-agent use.
"""
# Context window sizes for common models
MODEL_CONTEXT_WINDOWS = {
"gpt-4o": 128000,
"gpt-4o-mini": 128000,
"gpt-4-turbo": 128000,
"gpt-4": 8192,
"gpt-3.5-turbo": 16385,
"claude-3-opus": 200000,
"claude-3-sonnet": 200000,
"claude-3-haiku": 200000,
"claude-3-5-sonnet": 200000,
"gemini-pro": 32000,
"gemini-1.5-pro": 1000000,
"gemini-1.5-flash": 1000000,
}
def __init__(
self,
context_window: int = 128000,
threshold: float = 0.8,
summary_model: Optional[str] = None,
preserve_system: bool = True,
preserve_recent: int = 2,
):
"""
Initialize summarization manager.
Args:
context_window: Maximum context window size in tokens
threshold: Percentage of context window to trigger summarization (0.0-1.0)
summary_model: Model to use for generating summaries (None = use same model)
preserve_system: Whether to preserve system message after summarization
preserve_recent: Number of recent message pairs to preserve
"""
if not 0.0 < threshold <= 1.0:
raise ValueError("threshold must be between 0.0 and 1.0")
self.context_window = context_window
self.threshold = threshold
self.summary_model = summary_model
self.preserve_system = preserve_system
self.preserve_recent = preserve_recent
self.current_tokens = 0
self._summarization_count = 0
@classmethod
def for_model(cls, model: str, threshold: float = 0.8, **kwargs) -> "SummarizationManager":
"""Create a SummarizationManager configured for a specific model."""
# Normalize model name
model_lower = model.lower()
# Find matching context window
context_window = 128000 # Default
for model_key, window in cls.MODEL_CONTEXT_WINDOWS.items():
if model_key in model_lower:
context_window = window
break
return cls(context_window=context_window, threshold=threshold, **kwargs)
def add_tokens(self, tokens: int) -> None:
"""Add tokens to the current count."""
self.current_tokens += tokens
def set_tokens(self, tokens: int) -> None:
"""Set the current token count."""
self.current_tokens = tokens
def reset_tokens(self) -> None:
"""Reset token count after summarization."""
self.current_tokens = 0
def should_summarize(self) -> bool:
"""Check if summarization should be triggered."""
threshold_tokens = int(self.context_window * self.threshold)
return self.current_tokens >= threshold_tokens
def get_usage_percentage(self) -> float:
"""Get current context window usage as percentage."""
return self.current_tokens / self.context_window
def generate_summary_prompt(self, messages: List[Dict[str, str]]) -> str:
"""Generate prompt for summarizing conversation."""
conversation = "\n".join([
f"{msg['role'].upper()}: {msg['content']}"
for msg in messages
if msg['role'] != 'system'
])
return f"""Summarize the following conversation concisely, preserving:
1. Key decisions and conclusions
2. Important context and facts
3. Any pending tasks or questions
4. Technical details that may be needed later
Keep the summary under 500 words.
CONVERSATION:
{conversation}
SUMMARY:"""
def prepare_summarized_history(
self,
messages: List[Dict[str, str]],
summary: str
) -> List[Dict[str, str]]:
"""
Prepare new message history after summarization.
Args:
messages: Original message history
summary: Generated summary
Returns:
New message history with summary replacing old messages
"""
new_history = []
# Preserve system message if configured
if self.preserve_system and messages and messages[0].get('role') == 'system':
new_history.append(messages[0])
# Add summary as assistant message
new_history.append({
"role": "assistant",
"content": f"[Previous conversation summary]\n{summary}"
})
# Preserve recent messages
if self.preserve_recent > 0:
# Get last N pairs (user + assistant)
recent_count = self.preserve_recent * 2
recent_messages = [
m for m in messages[-recent_count:]
if m.get('role') != 'system'
]
new_history.extend(recent_messages)
self._summarization_count += 1
logger.info(f"Conversation summarized (count: {self._summarization_count})")
return new_history
@property
def summarization_count(self) -> int:
"""Number of times summarization has been performed."""
return self._summarization_count
async def summarize_conversation(
messages: List[Dict[str, str]],
llm_call: Callable,
manager: SummarizationManager,
) -> List[Dict[str, str]]:
"""
Summarize conversation using LLM.
Args:
messages: Current message history
llm_call: Async function to call LLM
manager: SummarizationManager instance
Returns:
New message history with summary
"""
prompt = manager.generate_summary_prompt(messages)
# Call LLM to generate summary
summary = await llm_call(prompt)
# Prepare new history
new_history = manager.prepare_summarized_history(messages, summary)
# Reset token count (will be recalculated on next message)
manager.reset_tokens()
return new_history
def summarize_conversation_sync(
messages: List[Dict[str, str]],
llm_call: Callable,
manager: SummarizationManager,
) -> List[Dict[str, str]]:
"""
Synchronous version of summarize_conversation.
"""
prompt = manager.generate_summary_prompt(messages)
summary = llm_call(prompt)
new_history = manager.prepare_summarized_history(messages, summary)
manager.reset_tokens()
return new_history
@@ -0,0 +1,238 @@
"""
Agent Profiles and Modes for PraisonAI Agents.
Provides built-in agent profiles and mode configurations
similar to OpenCode's agent system.
"""
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional
class AgentMode(str, Enum):
"""Agent execution modes."""
PRIMARY = "primary" # Main agent, can spawn subagents
SUBAGENT = "subagent" # Spawned by another agent
ALL = "all" # Can be used in any context
@dataclass
class AgentProfile:
"""
A built-in agent profile configuration.
Attributes:
name: Unique agent name
description: What this agent does
mode: Execution mode (primary/subagent/all)
system_prompt: System prompt for the agent
tools: List of tool names this agent can use
temperature: LLM temperature setting
top_p: LLM top_p setting
max_steps: Maximum execution steps
hidden: Whether to hide from agent list
color: Display color (for UI)
"""
name: str
description: str = ""
mode: AgentMode = AgentMode.ALL
system_prompt: str = ""
tools: List[str] = field(default_factory=list)
temperature: float = 0.7
top_p: float = 1.0
max_steps: int = 50
hidden: bool = False
color: str = "#3B82F6"
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert profile to dictionary."""
return {
"name": self.name,
"description": self.description,
"mode": self.mode.value,
"system_prompt": self.system_prompt,
"tools": self.tools,
"temperature": self.temperature,
"top_p": self.top_p,
"max_steps": self.max_steps,
"hidden": self.hidden,
"color": self.color,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "AgentProfile":
"""Create profile from dictionary."""
return cls(
name=data.get("name", ""),
description=data.get("description", ""),
mode=AgentMode(data.get("mode", "all")),
system_prompt=data.get("system_prompt", ""),
tools=data.get("tools", []),
temperature=data.get("temperature", 0.7),
top_p=data.get("top_p", 1.0),
max_steps=data.get("max_steps", 50),
hidden=data.get("hidden", False),
color=data.get("color", "#3B82F6"),
metadata=data.get("metadata", {}),
)
# Built-in agent profiles
BUILTIN_PROFILES: Dict[str, AgentProfile] = {
"general": AgentProfile(
name="general",
description="General-purpose coding assistant",
mode=AgentMode.PRIMARY,
system_prompt=(
"You are a helpful coding assistant. You can help with a wide variety "
"of programming tasks including writing code, debugging, explaining "
"concepts, and answering questions about software development."
),
tools=["read_file", "write_file", "bash", "search"],
temperature=0.7,
color="#3B82F6",
),
"coder": AgentProfile(
name="coder",
description="Focused code implementation agent",
mode=AgentMode.ALL,
system_prompt=(
"You are an expert code implementation agent. Your focus is on writing "
"clean, efficient, and well-documented code. You follow best practices "
"and design patterns. When implementing features, you consider edge cases "
"and error handling."
),
tools=["read_file", "write_file", "bash"],
temperature=0.3,
color="#10B981",
),
"planner": AgentProfile(
name="planner",
description="Task planning and decomposition agent",
mode=AgentMode.SUBAGENT,
system_prompt=(
"You are a planning agent. Your role is to break down complex tasks "
"into smaller, manageable steps. You create clear, actionable plans "
"with specific milestones and success criteria. You do not execute "
"tasks yourself, only plan them."
),
tools=["read_file", "search"],
temperature=0.5,
max_steps=20,
color="#F59E0B",
),
"reviewer": AgentProfile(
name="reviewer",
description="Code review and quality agent",
mode=AgentMode.SUBAGENT,
system_prompt=(
"You are a code review agent. Your role is to review code for quality, "
"correctness, security, and best practices. You provide constructive "
"feedback and suggest improvements. You look for bugs, security issues, "
"and opportunities for optimization."
),
tools=["read_file", "search"],
temperature=0.3,
max_steps=30,
color="#EF4444",
),
"explorer": AgentProfile(
name="explorer",
description="Codebase exploration and understanding agent",
mode=AgentMode.SUBAGENT,
system_prompt=(
"You are a codebase exploration agent. Your role is to understand and "
"navigate codebases. You can find relevant files, understand code "
"structure, and identify patterns. You help locate where changes need "
"to be made."
),
tools=["read_file", "search", "grep"],
temperature=0.5,
max_steps=40,
color="#8B5CF6",
),
"debugger": AgentProfile(
name="debugger",
description="Debugging and troubleshooting agent",
mode=AgentMode.SUBAGENT,
system_prompt=(
"You are a debugging agent. Your role is to identify and fix bugs. "
"You analyze error messages, trace code execution, and identify root "
"causes. You suggest minimal fixes that address the issue without "
"introducing new problems."
),
tools=["read_file", "bash", "search"],
temperature=0.3,
max_steps=30,
color="#EC4899",
),
"compaction": AgentProfile(
name="compaction",
description="Context compaction agent (internal)",
mode=AgentMode.SUBAGENT,
system_prompt=(
"You are a context compaction agent. Your role is to summarize "
"conversation history while preserving essential information. "
"You create concise summaries that maintain context for continued "
"work."
),
tools=[],
temperature=0.3,
max_steps=5,
hidden=True,
color="#6B7280",
),
"title": AgentProfile(
name="title",
description="Session title generation agent (internal)",
mode=AgentMode.SUBAGENT,
system_prompt=(
"You are a title generation agent. Your role is to create short, "
"descriptive titles for conversations based on their content. "
"Titles should be concise (3-7 words) and capture the main topic."
),
tools=[],
temperature=0.7,
max_steps=1,
hidden=True,
color="#6B7280",
),
}
def get_profile(name: str) -> Optional[AgentProfile]:
"""Get a built-in agent profile by name."""
return BUILTIN_PROFILES.get(name)
def list_profiles(include_hidden: bool = False) -> List[AgentProfile]:
"""List all built-in agent profiles."""
profiles = list(BUILTIN_PROFILES.values())
if not include_hidden:
profiles = [p for p in profiles if not p.hidden]
return profiles
def register_profile(profile: AgentProfile) -> None:
"""Register a custom agent profile."""
BUILTIN_PROFILES[profile.name] = profile
def get_profiles_by_mode(mode: AgentMode) -> List[AgentProfile]:
"""Get all profiles that can be used in a specific mode."""
return [
p for p in BUILTIN_PROFILES.values()
if p.mode == mode or p.mode == AgentMode.ALL
]
@@ -7,9 +7,12 @@ It extends the existing callback system to require human approval before executi
import logging
import asyncio
from typing import Dict, Set, Optional, Callable, Any, Literal
import json
import os
from typing import Dict, Set, Optional, Callable, Any, Literal, List
from functools import wraps
from contextvars import ContextVar
from dataclasses import dataclass, field
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
@@ -28,6 +31,172 @@ approval_callback: Optional[Callable] = None
# Context variable to track if we're in an approved execution context
_approved_context: ContextVar[Set[str]] = ContextVar('approved_context', default=set())
# Global permission allowlist
_permission_allowlist: Optional["PermissionAllowlist"] = None
@dataclass
class ToolPermission:
"""Permission entry for a tool."""
tool_name: str
allowed_paths: List[str] = field(default_factory=list)
session_only: bool = False # If True, permission expires with session
class PermissionAllowlist:
"""
Persistent permission allowlist for tools.
Allows pre-approving tools and paths to skip interactive approval.
Can be saved/loaded from JSON for persistence across sessions.
Usage:
allowlist = PermissionAllowlist()
allowlist.add_tool("read_file")
allowlist.add_tool("write_file", paths=["./src", "./tests"])
if allowlist.is_allowed("read_file"):
# Skip approval prompt
pass
"""
def __init__(self):
self._tools: Dict[str, ToolPermission] = {}
self._session_tools: Set[str] = set()
def add_tool(
self,
tool_name: str,
paths: Optional[List[str]] = None,
session_only: bool = False
) -> None:
"""
Add a tool to the allowlist.
Args:
tool_name: Name of the tool to allow
paths: Optional list of allowed paths (empty = all paths)
session_only: If True, permission expires with session
"""
self._tools[tool_name] = ToolPermission(
tool_name=tool_name,
allowed_paths=paths or [],
session_only=session_only
)
if session_only:
self._session_tools.add(tool_name)
def remove_tool(self, tool_name: str) -> bool:
"""Remove a tool from the allowlist."""
if tool_name in self._tools:
del self._tools[tool_name]
self._session_tools.discard(tool_name)
return True
return False
def is_allowed(self, tool_name: str, path: Optional[str] = None) -> bool:
"""
Check if a tool is allowed.
Args:
tool_name: Name of the tool
path: Optional path to check against allowed paths
Returns:
True if tool is allowed (optionally for the given path)
"""
if tool_name not in self._tools:
return False
permission = self._tools[tool_name]
# If no paths specified, tool is allowed for all paths
if not permission.allowed_paths:
return True
# If path not provided, check if tool has any permissions
if path is None:
return True
# Check if path matches any allowed path
for allowed_path in permission.allowed_paths:
# Normalize paths for comparison
norm_allowed = os.path.normpath(allowed_path)
norm_path = os.path.normpath(path)
# Check if path starts with allowed path
if norm_path.startswith(norm_allowed):
return True
return False
def is_empty(self) -> bool:
"""Check if allowlist is empty."""
return len(self._tools) == 0
def list_tools(self) -> List[str]:
"""List all allowed tools."""
return list(self._tools.keys())
def clear_session_permissions(self) -> None:
"""Clear session-only permissions."""
for tool_name in list(self._session_tools):
if tool_name in self._tools:
del self._tools[tool_name]
self._session_tools.clear()
def save(self, filepath: str) -> None:
"""Save allowlist to JSON file."""
data = {
"tools": {
name: {
"allowed_paths": perm.allowed_paths,
"session_only": perm.session_only
}
for name, perm in self._tools.items()
if not perm.session_only # Don't persist session-only permissions
}
}
os.makedirs(os.path.dirname(filepath) or ".", exist_ok=True)
with open(filepath, "w") as f:
json.dump(data, f, indent=2)
@classmethod
def load(cls, filepath: str) -> "PermissionAllowlist":
"""Load allowlist from JSON file."""
allowlist = cls()
if not os.path.exists(filepath):
return allowlist
with open(filepath, "r") as f:
data = json.load(f)
for tool_name, tool_data in data.get("tools", {}).items():
allowlist.add_tool(
tool_name,
paths=tool_data.get("allowed_paths", []),
session_only=tool_data.get("session_only", False)
)
return allowlist
def get_permission_allowlist() -> PermissionAllowlist:
"""Get the global permission allowlist."""
global _permission_allowlist
if _permission_allowlist is None:
_permission_allowlist = PermissionAllowlist()
return _permission_allowlist
def set_permission_allowlist(allowlist: PermissionAllowlist) -> None:
"""Set the global permission allowlist."""
global _permission_allowlist
_permission_allowlist = allowlist
class ApprovalDecision:
"""Result of an approval request"""
def __init__(self, approved: bool, modified_args: Optional[Dict[str, Any]] = None, reason: str = ""):
@@ -0,0 +1,286 @@
"""
Background Job Manager for PraisonAI Agents.
Provides job lifecycle tracking, auto-backgrounding based on threshold,
and status inspection for long-running agent tasks.
"""
import threading
import time
import uuid
from enum import Enum
from typing import Any, Callable, Dict, Optional
from dataclasses import dataclass, field
from concurrent.futures import ThreadPoolExecutor, Future
class JobStatus(str, Enum):
"""Status of a background job."""
PENDING = "pending"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
CANCELLED = "cancelled"
@dataclass
class JobInfo:
"""Information about a background job."""
job_id: str
status: JobStatus
created_at: float = field(default_factory=time.time)
started_at: Optional[float] = None
completed_at: Optional[float] = None
result: Any = None
error: Optional[str] = None
@property
def duration(self) -> Optional[float]:
"""Get job duration in seconds."""
if self.started_at is None:
return None
end_time = self.completed_at or time.time()
return end_time - self.started_at
class BackgroundJobManager:
"""
Manages background jobs for agent tasks.
Features:
- Job lifecycle tracking (pending, running, completed, failed)
- Auto-backgrounding based on execution time threshold
- Thread-safe job status inspection
- Configurable thread pool for concurrent execution
"""
def __init__(
self,
max_workers: int = 4,
auto_background_threshold: float = 5.0,
):
"""
Initialize the background job manager.
Args:
max_workers: Maximum number of concurrent background jobs
auto_background_threshold: Time in seconds after which a job
should be automatically backgrounded
"""
self.max_workers = max_workers
self.auto_background_threshold = auto_background_threshold
self._executor = ThreadPoolExecutor(max_workers=max_workers)
self._jobs: Dict[str, JobInfo] = {}
self._futures: Dict[str, Future] = {}
self._lock = threading.Lock()
@property
def active_jobs(self) -> int:
"""Get the number of active (pending or running) jobs."""
with self._lock:
return sum(
1 for job in self._jobs.values()
if job.status in (JobStatus.PENDING, JobStatus.RUNNING)
)
def start_job(
self,
func: Callable[[], Any],
job_id: Optional[str] = None,
) -> str:
"""
Start a new background job.
Args:
func: The function to execute in the background
job_id: Optional custom job ID (auto-generated if not provided)
Returns:
The job ID
"""
job_id = job_id or str(uuid.uuid4())[:8]
job_info = JobInfo(
job_id=job_id,
status=JobStatus.PENDING,
)
with self._lock:
self._jobs[job_id] = job_info
def _run_job():
with self._lock:
self._jobs[job_id].status = JobStatus.RUNNING
self._jobs[job_id].started_at = time.time()
try:
result = func()
with self._lock:
self._jobs[job_id].status = JobStatus.COMPLETED
self._jobs[job_id].completed_at = time.time()
self._jobs[job_id].result = result
return result
except Exception as e:
with self._lock:
self._jobs[job_id].status = JobStatus.FAILED
self._jobs[job_id].completed_at = time.time()
self._jobs[job_id].error = str(e)
raise
future = self._executor.submit(_run_job)
with self._lock:
self._futures[job_id] = future
return job_id
def get_status(self, job_id: str) -> JobStatus:
"""
Get the status of a job.
Args:
job_id: The job ID to check
Returns:
The job status
Raises:
KeyError: If job_id is not found
"""
with self._lock:
if job_id not in self._jobs:
raise KeyError(f"Job {job_id} not found")
return self._jobs[job_id].status
def get_job_info(self, job_id: str) -> JobInfo:
"""
Get full information about a job.
Args:
job_id: The job ID to check
Returns:
The job information
Raises:
KeyError: If job_id is not found
"""
with self._lock:
if job_id not in self._jobs:
raise KeyError(f"Job {job_id} not found")
return self._jobs[job_id]
def get_result(self, job_id: str, timeout: Optional[float] = None) -> Any:
"""
Get the result of a completed job.
Args:
job_id: The job ID
timeout: Maximum time to wait for completion
Returns:
The job result
Raises:
KeyError: If job_id is not found
TimeoutError: If timeout is reached before completion
RuntimeError: If job failed
"""
with self._lock:
if job_id not in self._futures:
raise KeyError(f"Job {job_id} not found")
future = self._futures[job_id]
try:
return future.result(timeout=timeout)
except TimeoutError:
raise TimeoutError(f"Job {job_id} did not complete within {timeout} seconds")
def cancel_job(self, job_id: str) -> bool:
"""
Attempt to cancel a job.
Args:
job_id: The job ID to cancel
Returns:
True if cancellation was successful, False otherwise
"""
with self._lock:
if job_id not in self._futures:
return False
future = self._futures[job_id]
cancelled = future.cancel()
if cancelled:
self._jobs[job_id].status = JobStatus.CANCELLED
self._jobs[job_id].completed_at = time.time()
return cancelled
def list_jobs(self, status: Optional[JobStatus] = None) -> Dict[str, JobInfo]:
"""
List all jobs, optionally filtered by status.
Args:
status: Optional status filter
Returns:
Dictionary of job_id -> JobInfo
"""
with self._lock:
if status is None:
return dict(self._jobs)
return {
job_id: info
for job_id, info in self._jobs.items()
if info.status == status
}
def cleanup_completed(self, max_age: float = 3600.0) -> int:
"""
Remove completed jobs older than max_age seconds.
Args:
max_age: Maximum age in seconds for completed jobs
Returns:
Number of jobs removed
"""
now = time.time()
removed = 0
with self._lock:
to_remove = []
for job_id, info in self._jobs.items():
if info.status in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
if info.completed_at and (now - info.completed_at) > max_age:
to_remove.append(job_id)
for job_id in to_remove:
del self._jobs[job_id]
if job_id in self._futures:
del self._futures[job_id]
removed += 1
return removed
def shutdown(self, wait: bool = True) -> None:
"""
Shutdown the job manager.
Args:
wait: Whether to wait for pending jobs to complete
"""
self._executor.shutdown(wait=wait)
# Global job manager instance (lazy initialized)
_global_job_manager: Optional[BackgroundJobManager] = None
_global_lock = threading.Lock()
def get_job_manager() -> BackgroundJobManager:
"""Get the global job manager instance."""
global _global_job_manager
with _global_lock:
if _global_job_manager is None:
_global_job_manager = BackgroundJobManager()
return _global_job_manager
@@ -0,0 +1,71 @@
"""
Event Bus Module for PraisonAI Agents.
Provides a typed publish/subscribe event system for real-time communication
between components. Extends the existing hooks system with a more general
event-driven architecture.
Features:
- Typed event definitions with dataclass payloads
- Sync and async subscribers
- Event filtering by type
- Global and scoped event buses
- SSE-compatible event streaming
Zero Performance Impact:
- All imports are lazy loaded
- No overhead when not subscribed
- Optional dependency for server features
Usage:
from praisonaiagents.bus import EventBus, Event
# Create event bus
bus = EventBus()
# Subscribe to events
@bus.on("session.created")
def handle_session(event):
print(f"Session created: {event.data}")
# Publish events
bus.publish("session.created", {"session_id": "abc123"})
"""
__all__ = [
"EventBus",
"Event",
"EventType",
"Subscriber",
"get_default_bus",
"set_default_bus",
]
def __getattr__(name: str):
"""Lazy load module components."""
if name == "EventBus":
from .bus import EventBus
return EventBus
if name == "Event":
from .event import Event
return Event
if name == "EventType":
from .event import EventType
return EventType
if name == "Subscriber":
from .bus import Subscriber
return Subscriber
if name == "get_default_bus":
from .bus import get_default_bus
return get_default_bus
if name == "set_default_bus":
from .bus import set_default_bus
return set_default_bus
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,346 @@
"""
Event Bus implementation for PraisonAI Agents.
Provides publish/subscribe functionality for typed events.
"""
import asyncio
import logging
import threading
from typing import Any, Callable, Dict, List, Optional, Set, Union
from dataclasses import dataclass, field
from .event import Event, EventType
logger = logging.getLogger(__name__)
@dataclass
class Subscriber:
"""
A subscriber to the event bus.
Attributes:
callback: Function to call when event is received
event_types: Set of event types to subscribe to (None = all)
is_async: Whether the callback is async
id: Unique subscriber identifier
"""
callback: Callable[[Event], Any]
event_types: Optional[Set[str]] = None
is_async: bool = False
id: str = field(default_factory=lambda: str(__import__("uuid").uuid4()))
def matches(self, event_type: str) -> bool:
"""Check if subscriber should receive this event type."""
if self.event_types is None:
return True
return event_type in self.event_types
class EventBus:
"""
Event bus for publish/subscribe communication.
Thread-safe implementation supporting both sync and async subscribers.
Example:
bus = EventBus()
# Subscribe with decorator
@bus.on(EventType.SESSION_CREATED)
def handle_session(event):
print(f"Session: {event.data}")
# Subscribe to all events
bus.subscribe(lambda e: print(e))
# Publish event
bus.publish(EventType.SESSION_CREATED, {"session_id": "abc"})
"""
def __init__(self):
"""Initialize the event bus."""
self._subscribers: List[Subscriber] = []
self._lock = threading.RLock()
self._event_history: List[Event] = []
self._max_history = 1000
def subscribe(
self,
callback: Callable[[Event], Any],
event_types: Optional[Union[str, List[str], Set[str]]] = None,
) -> str:
"""
Subscribe to events.
Args:
callback: Function to call when event is received
event_types: Optional event type(s) to filter by
Returns:
Subscriber ID for unsubscribing
"""
# Normalize event_types to a set
if event_types is None:
types_set = None
elif isinstance(event_types, str):
types_set = {event_types}
else:
types_set = set(event_types)
# Check if callback is async
is_async = asyncio.iscoroutinefunction(callback)
subscriber = Subscriber(
callback=callback,
event_types=types_set,
is_async=is_async,
)
with self._lock:
self._subscribers.append(subscriber)
logger.debug(f"Subscribed {subscriber.id} to {types_set or 'all events'}")
return subscriber.id
def unsubscribe(self, subscriber_id: str) -> bool:
"""
Unsubscribe from events.
Args:
subscriber_id: The subscriber ID returned from subscribe()
Returns:
True if found and removed, False otherwise
"""
with self._lock:
for i, sub in enumerate(self._subscribers):
if sub.id == subscriber_id:
self._subscribers.pop(i)
logger.debug(f"Unsubscribed {subscriber_id}")
return True
return False
def on(
self,
*event_types: Union[str, EventType],
) -> Callable:
"""
Decorator to subscribe a function to events.
Args:
event_types: Event type(s) to subscribe to
Returns:
Decorator function
Example:
@bus.on(EventType.SESSION_CREATED)
def handle_session(event):
print(event.data)
"""
# Convert EventType enums to strings
types = {
t.value if isinstance(t, EventType) else t
for t in event_types
} if event_types else None
def decorator(func: Callable[[Event], Any]) -> Callable:
self.subscribe(func, types)
return func
return decorator
def publish(
self,
event_type: Union[str, EventType],
data: Optional[Dict[str, Any]] = None,
source: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> Event:
"""
Publish an event to all matching subscribers.
Args:
event_type: The event type
data: Event payload data
source: Optional source identifier
metadata: Optional additional metadata
Returns:
The published Event object
"""
# Convert EventType enum to string
type_str = event_type.value if isinstance(event_type, EventType) else event_type
event = Event(
type=type_str,
data=data or {},
source=source,
metadata=metadata or {},
)
return self.publish_event(event)
def publish_event(self, event: Event) -> Event:
"""
Publish a pre-constructed event.
Args:
event: The event to publish
Returns:
The published event
"""
logger.debug(f"Publishing event: {event.type}")
# Store in history
with self._lock:
self._event_history.append(event)
if len(self._event_history) > self._max_history:
self._event_history = self._event_history[-self._max_history:]
# Get matching subscribers
subscribers = [
sub for sub in self._subscribers
if sub.matches(event.type)
]
# Dispatch to subscribers
for sub in subscribers:
try:
if sub.is_async:
# Schedule async callback
try:
loop = asyncio.get_running_loop()
loop.create_task(sub.callback(event))
except RuntimeError:
# No running loop, run synchronously
asyncio.run(sub.callback(event))
else:
sub.callback(event)
except Exception as e:
logger.error(f"Error in subscriber {sub.id}: {e}")
return event
async def publish_async(
self,
event_type: Union[str, EventType],
data: Optional[Dict[str, Any]] = None,
source: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> Event:
"""
Async version of publish that awaits all async subscribers.
Args:
event_type: The event type
data: Event payload data
source: Optional source identifier
metadata: Optional additional metadata
Returns:
The published Event object
"""
type_str = event_type.value if isinstance(event_type, EventType) else event_type
event = Event(
type=type_str,
data=data or {},
source=source,
metadata=metadata or {},
)
logger.debug(f"Publishing async event: {event.type}")
# Store in history
with self._lock:
self._event_history.append(event)
if len(self._event_history) > self._max_history:
self._event_history = self._event_history[-self._max_history:]
subscribers = [
sub for sub in self._subscribers
if sub.matches(event.type)
]
# Dispatch to subscribers
tasks = []
for sub in subscribers:
try:
if sub.is_async:
tasks.append(sub.callback(event))
else:
sub.callback(event)
except Exception as e:
logger.error(f"Error in subscriber {sub.id}: {e}")
# Await all async tasks
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
return event
def get_history(
self,
event_type: Optional[str] = None,
limit: int = 100,
) -> List[Event]:
"""
Get recent event history.
Args:
event_type: Optional filter by event type
limit: Maximum number of events to return
Returns:
List of recent events
"""
with self._lock:
events = self._event_history.copy()
if event_type:
events = [e for e in events if e.type == event_type]
return events[-limit:]
def clear_history(self):
"""Clear the event history."""
with self._lock:
self._event_history.clear()
def clear_subscribers(self):
"""Remove all subscribers."""
with self._lock:
self._subscribers.clear()
@property
def subscriber_count(self) -> int:
"""Get the number of subscribers."""
with self._lock:
return len(self._subscribers)
def __repr__(self) -> str:
return f"EventBus(subscribers={self.subscriber_count})"
# Global default bus
_default_bus: Optional[EventBus] = None
_bus_lock = threading.Lock()
def get_default_bus() -> EventBus:
"""Get the default global event bus."""
global _default_bus
with _bus_lock:
if _default_bus is None:
_default_bus = EventBus()
return _default_bus
def set_default_bus(bus: EventBus):
"""Set the default global event bus."""
global _default_bus
with _bus_lock:
_default_bus = bus
@@ -0,0 +1,111 @@
"""
Event definitions for the Event Bus.
Provides typed event structures for the publish/subscribe system.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, Optional
from enum import Enum
import time
import uuid
class EventType(str, Enum):
"""Standard event types for the agent system."""
# Session events
SESSION_CREATED = "session.created"
SESSION_UPDATED = "session.updated"
SESSION_DELETED = "session.deleted"
SESSION_FORKED = "session.forked"
SESSION_REVERTED = "session.reverted"
# Message events
MESSAGE_CREATED = "message.created"
MESSAGE_UPDATED = "message.updated"
MESSAGE_PART_CREATED = "message.part.created"
MESSAGE_PART_UPDATED = "message.part.updated"
# Permission events
PERMISSION_ASKED = "permission.asked"
PERMISSION_REPLIED = "permission.replied"
# Agent events
AGENT_STARTED = "agent.started"
AGENT_COMPLETED = "agent.completed"
AGENT_ERROR = "agent.error"
# Tool events
TOOL_STARTED = "tool.started"
TOOL_COMPLETED = "tool.completed"
TOOL_ERROR = "tool.error"
# Snapshot events
SNAPSHOT_CREATED = "snapshot.created"
SNAPSHOT_RESTORED = "snapshot.restored"
# Server events
SERVER_STARTED = "server.started"
SERVER_STOPPED = "server.stopped"
CLIENT_CONNECTED = "client.connected"
CLIENT_DISCONNECTED = "client.disconnected"
# Compaction events
COMPACTION_STARTED = "compaction.started"
COMPACTION_COMPLETED = "compaction.completed"
# Custom event (for user-defined events)
CUSTOM = "custom"
@dataclass
class Event:
"""
A typed event for the event bus.
Attributes:
type: The event type (from EventType enum or custom string)
data: The event payload data
id: Unique event identifier
timestamp: Unix timestamp when event was created
source: Optional source identifier (e.g., session_id, agent_name)
metadata: Optional additional metadata
"""
type: str
data: Dict[str, Any] = field(default_factory=dict)
id: str = field(default_factory=lambda: str(uuid.uuid4()))
timestamp: float = field(default_factory=time.time)
source: Optional[str] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert event to dictionary for serialization."""
return {
"id": self.id,
"type": self.type,
"data": self.data,
"timestamp": self.timestamp,
"source": self.source,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "Event":
"""Create event from dictionary."""
return cls(
id=data.get("id", str(uuid.uuid4())),
type=data.get("type", EventType.CUSTOM.value),
data=data.get("data", {}),
timestamp=data.get("timestamp", time.time()),
source=data.get("source"),
metadata=data.get("metadata", {}),
)
def to_sse(self) -> str:
"""Format event for Server-Sent Events."""
import json
return f"id: {self.id}\nevent: {self.type}\ndata: {json.dumps(self.to_dict())}\n\n"
def __repr__(self) -> str:
return f"Event(type={self.type!r}, id={self.id[:8]}..., source={self.source!r})"
@@ -109,6 +109,10 @@ class ContextCompactor:
compacted = self._summarize(messages)
elif self.strategy == CompactionStrategy.SMART:
compacted = self._smart_compact(messages)
elif self.strategy == CompactionStrategy.PRUNE:
compacted = self._prune(messages)
elif self.strategy == CompactionStrategy.LLM_SUMMARIZE:
compacted = self._llm_summarize(messages)
else:
compacted = self._truncate(messages)
@@ -223,6 +227,89 @@ class ContextCompactor:
# For now, use sliding window as base
return self._sliding_window(messages)
def _prune(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Prune old tool outputs while keeping tool calls.
This reduces token usage by removing verbose tool outputs
from older messages while preserving the context of what
tools were called.
"""
result = []
# Separate system and other 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 intact
recent = other_msgs[-self.preserve_recent:]
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
# Prune older messages
for msg in older:
role = msg.get("role", "")
content = msg.get("content", "")
# If this is a tool result, truncate it
if role == "tool" or msg.get("tool_call_id"):
if isinstance(content, str) and len(content) > 500:
pruned_msg = msg.copy()
pruned_msg["content"] = content[:200] + "\n...[output truncated]..."
result.append(pruned_msg)
else:
result.append(msg)
else:
result.append(msg)
result.extend(recent)
return result
def _llm_summarize(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Use LLM to summarize older messages.
Note: This is a placeholder that returns a structured summary.
Actual LLM integration should be done at the agent level.
"""
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:]
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
if older:
# Create structured summary for LLM to process
summary_parts = []
for msg in older:
role = msg.get("role", "unknown")
content = msg.get("content", "")
if isinstance(content, str) and content:
# Extract key information
summary_parts.append(f"[{role}]: {content[:150]}...")
if summary_parts:
summary = (
"[Compacted conversation history - summarize key points]\n"
+ "\n".join(summary_parts[:10])
)
result.append({
"role": "system",
"content": summary,
"_compacted": True,
"_original_count": len(older),
})
result.extend(recent)
return result
def get_stats(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
"""Get statistics about messages."""
total_tokens = self.count_total_tokens(messages)
@@ -1,13 +1,16 @@
"""
Compaction Strategies for PraisonAI Agents.
Compaction strategies for context management.
"""
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
"""Available compaction strategies."""
TRUNCATE = "truncate"
SLIDING = "sliding"
SUMMARIZE = "summarize"
SMART = "smart"
LLM_SUMMARIZE = "llm_summarize" # Use LLM for summarization
PRUNE = "prune" # Remove old tool outputs
+22 -13
View File
@@ -58,6 +58,12 @@ class LLM:
# Class-level flag for one-time logging configuration
_logging_configured = False
# Class-level cache for LiteLLM module (avoids repeated import overhead)
_litellm_module = None
# Class-level flag for LiteLLM global settings (set once)
_litellm_configured = False
# Default window sizes for different models (75% of actual to be safe)
MODEL_WINDOWS = {
# OpenAI
@@ -272,14 +278,16 @@ Respond with ONLY a valid JSON tool call in this format:
if not LLM._logging_configured:
LLM._configure_logging()
# Import litellm after logging is configured
try:
import litellm
except ImportError:
raise ImportError(
"LiteLLM is required but not installed. "
"Please install with: pip install 'praisonaiagents[llm]'"
)
# Use cached LiteLLM module or import once (avoids repeated import overhead)
if LLM._litellm_module is None:
try:
import litellm
LLM._litellm_module = litellm
except ImportError:
raise ImportError(
"LiteLLM is required but not installed. "
"Please install with: pip install 'praisonaiagents[llm]'"
)
self.model = model
self.timeout = timeout
@@ -341,11 +349,12 @@ Respond with ONLY a valid JSON tool call in this format:
self._formatted_tools_cache = {}
self._max_cache_size = 100
# Enable error dropping for cleaner output
import litellm
litellm.drop_params = True
# Enable parameter modification for providers like Anthropic
litellm.modify_params = True
# Configure LiteLLM global settings only once (avoids repeated setup cost)
if not LLM._litellm_configured:
litellm = LLM._litellm_module
litellm.drop_params = True
litellm.modify_params = True
LLM._litellm_configured = True
self._setup_event_tracking(events)
# Log all initialization parameters when in debug mode or verbose >= 10
+32 -8
View File
@@ -150,7 +150,10 @@ def _clean_display_content(content: str, max_length: int = 20000) -> str:
return content.strip()
def display_interaction(message, response, markdown=True, generation_time=None, console=None, agent_name=None, agent_role=None, agent_tools=None, task_name=None, task_description=None, task_id=None):
"""Synchronous version of display_interaction."""
"""Synchronous version of display_interaction.
Displays the task/message and response in clean panels with response time in title.
"""
if console is None:
console = Console()
@@ -179,16 +182,18 @@ def display_interaction(message, response, markdown=True, generation_time=None,
task_description=task_description,
task_id=task_id
)
# Rest of the display logic...
# Build response title with time (Agno-style: "Response (1.2s)")
response_title = "Response"
if generation_time:
console.print(Text(f"Response generated in {generation_time:.1f}s", style="dim"))
response_title = f"Response ({generation_time:.1f}s)"
if markdown:
console.print(Panel.fit(Markdown(message), title="Task", border_style="cyan"))
console.print(Panel.fit(Markdown(response), title="Response", border_style="cyan"))
console.print(Panel.fit(Markdown(response), title=response_title, border_style="cyan"))
else:
console.print(Panel.fit(Text(message, style="bold green"), title="Task", border_style="cyan"))
console.print(Panel.fit(Text(response, style="bold blue"), title="Response", border_style="cyan"))
console.print(Panel.fit(Text(response, style="bold blue"), title=response_title, border_style="cyan"))
def display_self_reflection(message: str, console=None):
if not message or not message.strip():
@@ -226,7 +231,16 @@ def display_instruction(message: str, console=None, agent_name: str = None, agen
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
console.print(Panel.fit(Text(message, style="bold blue"), title="Instruction", border_style="cyan"))
def display_tool_call(message: str, console=None):
def display_tool_call(message: str, console=None, tool_name: str = None, tool_input: dict = None, tool_output: str = None):
"""Display tool call information in a consolidated panel.
Args:
message: The tool call message (legacy format)
console: Rich console for output
tool_name: Name of the tool being called
tool_input: Input arguments to the tool
tool_output: Output from the tool (if available)
"""
logging.debug(f"display_tool_call called with message: {repr(message)}")
if not message or not message.strip():
logging.debug("Empty message in display_tool_call, returning early")
@@ -235,11 +249,21 @@ def display_tool_call(message: str, console=None):
logging.debug(f"Cleaned message in display_tool_call: {repr(message)}")
# Execute synchronous callbacks (always, even when console is None)
execute_sync_callback('tool_call', message=message)
execute_sync_callback('tool_call', message=message, tool_name=tool_name, tool_input=tool_input, tool_output=tool_output)
# Only print panel if console is provided (verbose mode)
if console is not None:
console.print(Panel.fit(Text(message, style="bold cyan"), title="Tool Call", border_style="green"))
# Build consolidated tool call display (Agno-style)
if tool_name and tool_input is not None:
# Format: tool_name(arg1=val1, arg2=val2)
args_str = ", ".join(f"{k}={v}" for k, v in tool_input.items()) if tool_input else ""
tool_display = f"• {tool_name}({args_str})"
if tool_output:
tool_display += f"\n → {tool_output}"
console.print(Panel.fit(Text(tool_display, style="bold cyan"), title="Tool Call", border_style="green"))
else:
# Legacy format
console.print(Panel.fit(Text(message, style="bold cyan"), title="Tool Call", border_style="green"))
def display_error(message: str, console=None):
if not message or not message.strip():
@@ -11,8 +11,15 @@ import platform
from typing import Any, List, Optional, Callable, Iterable, Union
from functools import wraps, partial
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
try:
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
MCP_AVAILABLE = True
except ImportError:
MCP_AVAILABLE = False
ClientSession = None
StdioServerParameters = None
stdio_client = None
class MCPToolRunner(threading.Thread):
"""A dedicated thread for running MCP operations."""
@@ -196,10 +203,17 @@ class MCP:
debug: Enable debug logging for MCP operations (default: False)
**kwargs: Additional parameters for StdioServerParameters
"""
# Check if MCP is available
if not MCP_AVAILABLE:
raise ImportError(
"MCP (Model Context Protocol) package is not installed. "
"Install it with: pip install praisonaiagents[mcp]"
)
# Handle backward compatibility with named parameter 'command'
if command_or_string is None and command is not None:
command_or_string = command
# Set up logging - default to WARNING level to hide INFO messages
if debug:
logging.getLogger("mcp-wrapper").setLevel(logging.DEBUG)
@@ -16,11 +16,17 @@ import weakref
from typing import List, Dict, Any, Optional, Callable, Iterable, Union
from urllib.parse import urlparse, urljoin
from mcp import ClientSession
try:
from mcp import ClientSession
MCP_AVAILABLE = True
except ImportError:
MCP_AVAILABLE = False
ClientSession = None
try:
import aiohttp
except ImportError:
raise ImportError("aiohttp is required for HTTP Stream transport. Install with: pip install praisonaiagents[mcp]")
aiohttp = None
logger = logging.getLogger("mcp-http-stream")
@@ -507,9 +513,23 @@ class HTTPStreamMCPClient:
timeout: Timeout in seconds for operations (default: 60)
options: Additional configuration options for the transport
"""
# Check if MCP is available
if not MCP_AVAILABLE:
raise ImportError(
"MCP (Model Context Protocol) package is not installed. "
"Install it with: pip install praisonaiagents[mcp]"
)
# Check if aiohttp is available
if aiohttp is None:
raise ImportError(
"aiohttp is required for HTTP Stream transport. "
"Install it with: pip install praisonaiagents[mcp]"
)
# Parse URL to extract base URL and endpoint
parsed = urlparse(server_url)
# If the URL already has a path, use it; otherwise use default /mcp endpoint
if parsed.path and parsed.path != '/':
self.base_url = server_url
@@ -11,8 +11,14 @@ import inspect
import json
from typing import List, Dict, Any, Optional, Callable, Iterable
from mcp import ClientSession
from mcp.client.sse import sse_client
try:
from mcp import ClientSession
from mcp.client.sse import sse_client
MCP_AVAILABLE = True
except ImportError:
MCP_AVAILABLE = False
ClientSession = None
sse_client = None
logger = logging.getLogger("mcp-sse")
@@ -169,12 +175,19 @@ class SSEMCPClient:
debug: Whether to enable debug logging
timeout: Timeout in seconds for operations (default: 60)
"""
# Check if MCP is available
if not MCP_AVAILABLE:
raise ImportError(
"MCP (Model Context Protocol) package is not installed. "
"Install it with: pip install praisonaiagents[mcp]"
)
self.server_url = server_url
self.debug = debug
self.timeout = timeout
self.session = None
self.tools = []
# Set up logging
if debug:
logger.setLevel(logging.DEBUG)
@@ -236,3 +236,46 @@ def functions_to_mcp_tools(functions: List[Callable]) -> List[Dict[str, Any]]:
List of MCP tool schemas
"""
return [function_to_mcp_schema(func) for func in functions]
def filter_disabled_tools(
tools: List[Dict[str, Any]],
disabled_tools: List[str]
) -> List[Dict[str, Any]]:
"""
Filter out disabled tools from a tool list.
Args:
tools: List of tool definitions (dicts with 'name' key)
disabled_tools: List of tool names to disable
Returns:
Filtered list of tools
"""
if not disabled_tools:
return tools
disabled_set = set(disabled_tools)
return [t for t in tools if t.get("name") not in disabled_set]
def filter_tools_by_allowlist(
tools: List[Dict[str, Any]],
allowed_tools: List[str]
) -> List[Dict[str, Any]]:
"""
Filter tools to only include those in allowlist.
Args:
tools: List of tool definitions (dicts with 'name' key)
allowed_tools: List of tool names to allow (empty = allow all)
Returns:
Filtered list of tools
"""
if not allowed_tools:
return tools
allowed_set = set(allowed_tools)
return [t for t in tools if t.get("name") in allowed_set]
@@ -0,0 +1,62 @@
"""
Permissions Module for PraisonAI Agents.
Provides pattern-based permission rules, persistent approvals,
and doom loop detection for safe agent execution.
Features:
- Pattern-based permission rules (allow, deny, ask)
- Persistent approval storage
- Per-agent permission rulesets
- Doom loop detection and prevention
- Integration with existing approval system
Usage:
from praisonaiagents.permissions import PermissionManager, PermissionRule
# Create permission manager
manager = PermissionManager()
# Add rules
manager.add_rule(PermissionRule(
pattern="bash:*",
action="ask",
description="Require approval for shell commands"
))
# Check permission
result = manager.check("bash:rm -rf /tmp/test")
"""
__all__ = [
"PermissionManager",
"PermissionRule",
"PermissionAction",
"PermissionResult",
"DoomLoopDetector",
]
def __getattr__(name: str):
"""Lazy load module components."""
if name == "PermissionManager":
from .manager import PermissionManager
return PermissionManager
if name == "PermissionRule":
from .rules import PermissionRule
return PermissionRule
if name == "PermissionAction":
from .rules import PermissionAction
return PermissionAction
if name == "PermissionResult":
from .rules import PermissionResult
return PermissionResult
if name == "DoomLoopDetector":
from .doom_loop import DoomLoopDetector
return DoomLoopDetector
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,256 @@
"""
Doom Loop Detection for PraisonAI Agents.
Detects and prevents agents from getting stuck in repetitive loops.
"""
import logging
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
logger = logging.getLogger(__name__)
# Default thresholds
DEFAULT_LOOP_THRESHOLD = 3
DEFAULT_WINDOW_SECONDS = 60
DEFAULT_MAX_TOOL_CALLS = 50
@dataclass
class ToolCallRecord:
"""Record of a tool call for loop detection."""
tool_name: str
arguments_hash: str
timestamp: float = field(default_factory=time.time)
session_id: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
return {
"tool_name": self.tool_name,
"arguments_hash": self.arguments_hash,
"timestamp": self.timestamp,
"session_id": self.session_id,
}
@dataclass
class DoomLoopResult:
"""Result of a doom loop check."""
is_loop: bool
reason: str = ""
loop_count: int = 0
tool_name: Optional[str] = None
recommendation: str = ""
def to_dict(self) -> Dict[str, Any]:
return {
"is_loop": self.is_loop,
"reason": self.reason,
"loop_count": self.loop_count,
"tool_name": self.tool_name,
"recommendation": self.recommendation,
}
class DoomLoopDetector:
"""
Detects and prevents doom loops in agent execution.
A doom loop occurs when an agent repeatedly calls the same
tool with the same arguments, indicating it's stuck.
Example:
detector = DoomLoopDetector()
# Record tool calls
detector.record("bash", {"command": "ls"})
detector.record("bash", {"command": "ls"})
detector.record("bash", {"command": "ls"})
# Check for loop
result = detector.check("bash", {"command": "ls"})
if result.is_loop:
print(f"Doom loop detected: {result.reason}")
"""
def __init__(
self,
loop_threshold: int = DEFAULT_LOOP_THRESHOLD,
window_seconds: float = DEFAULT_WINDOW_SECONDS,
max_tool_calls: int = DEFAULT_MAX_TOOL_CALLS,
):
"""
Initialize the doom loop detector.
Args:
loop_threshold: Number of identical calls to trigger loop detection
window_seconds: Time window for counting calls
max_tool_calls: Maximum total tool calls before warning
"""
self.loop_threshold = loop_threshold
self.window_seconds = window_seconds
self.max_tool_calls = max_tool_calls
self._records: List[ToolCallRecord] = []
self._call_counts: Dict[str, int] = defaultdict(int)
self._session_id: Optional[str] = None
def set_session(self, session_id: str):
"""Set the current session ID."""
self._session_id = session_id
def _hash_arguments(self, arguments: Dict[str, Any]) -> str:
"""Create a hash of tool arguments."""
import hashlib
import json
try:
arg_str = json.dumps(arguments, sort_keys=True, default=str)
return hashlib.md5(arg_str.encode()).hexdigest()[:16]
except (TypeError, ValueError):
return "unhashable"
def _cleanup_old_records(self):
"""Remove records outside the time window."""
cutoff = time.time() - self.window_seconds
self._records = [r for r in self._records if r.timestamp > cutoff]
def record(
self,
tool_name: str,
arguments: Optional[Dict[str, Any]] = None,
) -> ToolCallRecord:
"""
Record a tool call.
Args:
tool_name: Name of the tool
arguments: Tool arguments
Returns:
The recorded ToolCallRecord
"""
self._cleanup_old_records()
args_hash = self._hash_arguments(arguments or {})
record = ToolCallRecord(
tool_name=tool_name,
arguments_hash=args_hash,
session_id=self._session_id,
)
self._records.append(record)
self._call_counts[tool_name] += 1
return record
def check(
self,
tool_name: str,
arguments: Optional[Dict[str, Any]] = None,
) -> DoomLoopResult:
"""
Check if calling this tool would indicate a doom loop.
Args:
tool_name: Name of the tool
arguments: Tool arguments
Returns:
DoomLoopResult indicating if a loop is detected
"""
self._cleanup_old_records()
args_hash = self._hash_arguments(arguments or {})
# Count identical calls in window
identical_count = sum(
1 for r in self._records
if r.tool_name == tool_name and r.arguments_hash == args_hash
)
# Check for loop
if identical_count >= self.loop_threshold:
return DoomLoopResult(
is_loop=True,
reason=f"Tool '{tool_name}' called {identical_count} times with same arguments",
loop_count=identical_count,
tool_name=tool_name,
recommendation="Consider a different approach or ask for user guidance",
)
# Check for excessive tool calls
total_calls = len(self._records)
if total_calls >= self.max_tool_calls:
return DoomLoopResult(
is_loop=True,
reason=f"Excessive tool calls ({total_calls}) in time window",
loop_count=total_calls,
recommendation="Reduce tool usage or break task into smaller steps",
)
# Check for tool-specific excessive calls
tool_count = self._call_counts.get(tool_name, 0)
if tool_count >= self.max_tool_calls // 2:
return DoomLoopResult(
is_loop=True,
reason=f"Tool '{tool_name}' called excessively ({tool_count} times)",
loop_count=tool_count,
tool_name=tool_name,
recommendation=f"Consider alternative to '{tool_name}'",
)
return DoomLoopResult(
is_loop=False,
loop_count=identical_count,
tool_name=tool_name,
)
def record_and_check(
self,
tool_name: str,
arguments: Optional[Dict[str, Any]] = None,
) -> DoomLoopResult:
"""
Record a tool call and check for doom loop.
Convenience method that combines record() and check().
Args:
tool_name: Name of the tool
arguments: Tool arguments
Returns:
DoomLoopResult indicating if a loop is detected
"""
result = self.check(tool_name, arguments)
if not result.is_loop:
self.record(tool_name, arguments)
return result
def get_stats(self) -> Dict[str, Any]:
"""Get statistics about tool calls."""
self._cleanup_old_records()
return {
"total_calls": len(self._records),
"unique_tools": len(set(r.tool_name for r in self._records)),
"call_counts": dict(self._call_counts),
"window_seconds": self.window_seconds,
"loop_threshold": self.loop_threshold,
}
def reset(self):
"""Reset all records."""
self._records.clear()
self._call_counts.clear()
def reset_tool(self, tool_name: str):
"""Reset records for a specific tool."""
self._records = [r for r in self._records if r.tool_name != tool_name]
self._call_counts.pop(tool_name, None)
@@ -0,0 +1,359 @@
"""
Permission Manager for PraisonAI Agents.
Manages permission rules, persistent approvals, and integrates
with doom loop detection.
"""
import json
import logging
import os
import threading
from typing import Any, Callable, Dict, List, Optional
from .rules import PermissionRule, PermissionAction, PermissionResult, PersistentApproval
from .doom_loop import DoomLoopDetector
logger = logging.getLogger(__name__)
# Default storage path
DEFAULT_PERMISSIONS_DIR = os.path.expanduser("~/.praison/permissions")
class PermissionManager:
"""
Manages permission rules and approvals.
Provides pattern-based permission checking with persistent
approval storage and doom loop detection.
Example:
manager = PermissionManager()
# Add rules
manager.add_rule(PermissionRule(
pattern="bash:rm *",
action=PermissionAction.ASK,
description="Require approval for rm commands"
))
manager.add_rule(PermissionRule(
pattern="read:*",
action=PermissionAction.ALLOW,
description="Allow all read operations"
))
# Check permission
result = manager.check("bash:rm -rf /tmp/test")
if result.needs_approval:
# Ask user
approved = ask_user(result.target)
manager.approve(result.target, approved)
"""
def __init__(
self,
storage_dir: Optional[str] = None,
agent_name: Optional[str] = None,
approval_callback: Optional[Callable[[str, str], bool]] = None,
):
"""
Initialize the permission manager.
Args:
storage_dir: Directory for persistent storage
agent_name: Optional agent name for filtering rules
approval_callback: Optional callback for user approval
"""
self.storage_dir = storage_dir or DEFAULT_PERMISSIONS_DIR
self.agent_name = agent_name
self.approval_callback = approval_callback
self._rules: List[PermissionRule] = []
self._approvals: List[PersistentApproval] = []
self._doom_detector = DoomLoopDetector()
self._lock = threading.RLock()
# Ensure storage directory exists
os.makedirs(self.storage_dir, exist_ok=True)
# Load persistent data
self._load_rules()
self._load_approvals()
def _get_rules_path(self) -> str:
"""Get path to rules file."""
return os.path.join(self.storage_dir, "rules.json")
def _get_approvals_path(self) -> str:
"""Get path to approvals file."""
return os.path.join(self.storage_dir, "approvals.json")
def _load_rules(self):
"""Load rules from disk."""
path = self._get_rules_path()
if os.path.exists(path):
try:
with open(path, "r") as f:
data = json.load(f)
self._rules = [PermissionRule.from_dict(r) for r in data]
except (json.JSONDecodeError, IOError) as e:
logger.warning(f"Failed to load rules: {e}")
def _save_rules(self):
"""Save rules to disk."""
path = self._get_rules_path()
try:
with open(path, "w") as f:
json.dump([r.to_dict() for r in self._rules], f, indent=2)
except IOError as e:
logger.error(f"Failed to save rules: {e}")
def _load_approvals(self):
"""Load approvals from disk."""
path = self._get_approvals_path()
if os.path.exists(path):
try:
with open(path, "r") as f:
data = json.load(f)
self._approvals = [PersistentApproval.from_dict(a) for a in data]
# Filter out expired approvals
self._approvals = [a for a in self._approvals if a.is_valid()]
except (json.JSONDecodeError, IOError) as e:
logger.warning(f"Failed to load approvals: {e}")
def _save_approvals(self):
"""Save approvals to disk."""
path = self._get_approvals_path()
try:
# Filter out expired and one-time approvals
persistent = [
a for a in self._approvals
if a.is_valid() and a.scope in ("session", "always")
]
with open(path, "w") as f:
json.dump([a.to_dict() for a in persistent], f, indent=2)
except IOError as e:
logger.error(f"Failed to save approvals: {e}")
def add_rule(self, rule: PermissionRule) -> str:
"""
Add a permission rule.
Args:
rule: The rule to add
Returns:
The rule ID
"""
with self._lock:
self._rules.append(rule)
# Sort by priority (higher first)
self._rules.sort(key=lambda r: r.priority, reverse=True)
self._save_rules()
return rule.id
def remove_rule(self, rule_id: str) -> bool:
"""
Remove a permission rule.
Args:
rule_id: The rule ID to remove
Returns:
True if found and removed
"""
with self._lock:
for i, rule in enumerate(self._rules):
if rule.id == rule_id:
self._rules.pop(i)
self._save_rules()
return True
return False
def get_rules(self, agent_name: Optional[str] = None) -> List[PermissionRule]:
"""
Get all rules, optionally filtered by agent.
Args:
agent_name: Optional agent name to filter by
Returns:
List of matching rules
"""
with self._lock:
if agent_name is None:
return self._rules.copy()
return [
r for r in self._rules
if r.agent_name is None or r.agent_name == agent_name
]
def check(self, target: str, agent_name: Optional[str] = None) -> PermissionResult:
"""
Check permission for a target.
Args:
target: The target to check (e.g., "bash:rm -rf /tmp")
agent_name: Optional agent name
Returns:
PermissionResult with the decision
"""
agent = agent_name or self.agent_name
with self._lock:
# Check persistent approvals first
for approval in self._approvals:
if approval.matches(target, agent):
return PermissionResult(
action=PermissionAction.ALLOW if approval.approved else PermissionAction.DENY,
target=target,
reason=f"Persistent approval: {'approved' if approval.approved else 'denied'}",
approved=approval.approved,
)
# Check rules
for rule in self._rules:
if rule.agent_name and agent and rule.agent_name != agent:
continue
if rule.matches(target):
return PermissionResult(
action=rule.action,
rule=rule,
target=target,
reason=rule.description or f"Matched rule: {rule.pattern}",
)
# Default: ask
return PermissionResult(
action=PermissionAction.ASK,
target=target,
reason="No matching rule, requires approval",
)
def approve(
self,
target: str,
approved: bool,
scope: str = "once",
agent_name: Optional[str] = None,
) -> PersistentApproval:
"""
Record an approval decision.
Args:
target: The target that was approved/denied
approved: Whether it was approved
scope: Scope of approval (once, session, always)
agent_name: Optional agent name
Returns:
The created PersistentApproval
"""
approval = PersistentApproval(
pattern=target,
approved=approved,
scope=scope,
agent_name=agent_name or self.agent_name,
)
with self._lock:
self._approvals.append(approval)
if scope in ("session", "always"):
self._save_approvals()
return approval
def check_and_approve(
self,
target: str,
agent_name: Optional[str] = None,
) -> PermissionResult:
"""
Check permission and request approval if needed.
Uses the approval_callback if set.
Args:
target: The target to check
agent_name: Optional agent name
Returns:
PermissionResult with final decision
"""
result = self.check(target, agent_name)
if result.needs_approval and self.approval_callback:
approved = self.approval_callback(target, result.reason)
result.approved = approved
self.approve(target, approved, scope="once", agent_name=agent_name)
return result
def check_doom_loop(
self,
tool_name: str,
arguments: Optional[Dict[str, Any]] = None,
):
"""
Check for doom loop and record tool call.
Args:
tool_name: Name of the tool
arguments: Tool arguments
Returns:
DoomLoopResult
"""
return self._doom_detector.record_and_check(tool_name, arguments)
def get_doom_stats(self) -> Dict[str, Any]:
"""Get doom loop detector statistics."""
return self._doom_detector.get_stats()
def reset_doom_detector(self):
"""Reset the doom loop detector."""
self._doom_detector.reset()
def clear_approvals(self, scope: Optional[str] = None):
"""
Clear approvals.
Args:
scope: Optional scope to clear (None = all)
"""
with self._lock:
if scope is None:
self._approvals.clear()
else:
self._approvals = [a for a in self._approvals if a.scope != scope]
self._save_approvals()
def clear_rules(self):
"""Clear all rules."""
with self._lock:
self._rules.clear()
self._save_rules()
def set_approval_callback(self, callback: Callable[[str, str], bool]):
"""Set the approval callback."""
self.approval_callback = callback
def to_dict(self) -> Dict[str, Any]:
"""Export manager state to dictionary."""
with self._lock:
return {
"rules": [r.to_dict() for r in self._rules],
"approvals": [a.to_dict() for a in self._approvals],
"agent_name": self.agent_name,
}
def from_dict(self, data: Dict[str, Any]):
"""Import manager state from dictionary."""
with self._lock:
self._rules = [PermissionRule.from_dict(r) for r in data.get("rules", [])]
self._approvals = [PersistentApproval.from_dict(a) for a in data.get("approvals", [])]
if "agent_name" in data:
self.agent_name = data["agent_name"]
@@ -0,0 +1,211 @@
"""
Permission rules for PraisonAI Agents.
Provides pattern-based permission rules with allow/deny/ask actions.
"""
import fnmatch
import re
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional
class PermissionAction(str, Enum):
"""Actions for permission rules."""
ALLOW = "allow"
DENY = "deny"
ASK = "ask"
@dataclass
class PermissionRule:
"""
A pattern-based permission rule.
Attributes:
pattern: Glob or regex pattern to match against
action: What to do when pattern matches (allow/deny/ask)
description: Human-readable description
is_regex: Whether pattern is a regex (default: glob)
priority: Higher priority rules are checked first
agent_name: Optional agent name this rule applies to
enabled: Whether the rule is active
"""
pattern: str
action: PermissionAction = PermissionAction.ASK
description: str = ""
is_regex: bool = False
priority: int = 0
agent_name: Optional[str] = None
enabled: bool = True
id: str = field(default_factory=lambda: str(uuid.uuid4()))
created_at: float = field(default_factory=time.time)
def matches(self, target: str) -> bool:
"""
Check if this rule matches the target.
Args:
target: The string to match against (e.g., "bash:rm -rf /tmp")
Returns:
True if the pattern matches
"""
if not self.enabled:
return False
if self.is_regex:
try:
return bool(re.match(self.pattern, target))
except re.error:
return False
else:
return fnmatch.fnmatch(target, self.pattern)
def to_dict(self) -> Dict[str, Any]:
"""Convert rule to dictionary."""
return {
"id": self.id,
"pattern": self.pattern,
"action": self.action.value,
"description": self.description,
"is_regex": self.is_regex,
"priority": self.priority,
"agent_name": self.agent_name,
"enabled": self.enabled,
"created_at": self.created_at,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "PermissionRule":
"""Create rule from dictionary."""
return cls(
id=data.get("id", str(uuid.uuid4())),
pattern=data.get("pattern", "*"),
action=PermissionAction(data.get("action", "ask")),
description=data.get("description", ""),
is_regex=data.get("is_regex", False),
priority=data.get("priority", 0),
agent_name=data.get("agent_name"),
enabled=data.get("enabled", True),
created_at=data.get("created_at", time.time()),
)
@dataclass
class PermissionResult:
"""
Result of a permission check.
Attributes:
action: The resulting action (allow/deny/ask)
rule: The rule that matched (if any)
target: The target that was checked
reason: Human-readable reason
approved: Whether the action was approved (for ask actions)
"""
action: PermissionAction
rule: Optional[PermissionRule] = None
target: str = ""
reason: str = ""
approved: Optional[bool] = None
@property
def is_allowed(self) -> bool:
"""Check if the action is allowed."""
if self.action == PermissionAction.ALLOW:
return True
if self.action == PermissionAction.ASK:
return self.approved is True
return False
@property
def is_denied(self) -> bool:
"""Check if the action is denied."""
if self.action == PermissionAction.DENY:
return True
if self.action == PermissionAction.ASK:
return self.approved is False
return False
@property
def needs_approval(self) -> bool:
"""Check if the action needs user approval."""
return self.action == PermissionAction.ASK and self.approved is None
def to_dict(self) -> Dict[str, Any]:
"""Convert result to dictionary."""
return {
"action": self.action.value,
"rule_id": self.rule.id if self.rule else None,
"target": self.target,
"reason": self.reason,
"approved": self.approved,
"is_allowed": self.is_allowed,
"is_denied": self.is_denied,
"needs_approval": self.needs_approval,
}
@dataclass
class PersistentApproval:
"""
A persistent approval for a pattern.
Used to remember user decisions for future requests.
"""
pattern: str
approved: bool
scope: str = "once" # once, session, always
created_at: float = field(default_factory=time.time)
expires_at: Optional[float] = None
agent_name: Optional[str] = None
id: str = field(default_factory=lambda: str(uuid.uuid4()))
def is_valid(self) -> bool:
"""Check if the approval is still valid."""
if self.expires_at and time.time() > self.expires_at:
return False
return True
def matches(self, target: str, agent_name: Optional[str] = None) -> bool:
"""Check if this approval applies to the target."""
if not self.is_valid():
return False
if self.agent_name and agent_name and self.agent_name != agent_name:
return False
return fnmatch.fnmatch(target, self.pattern)
def to_dict(self) -> Dict[str, Any]:
"""Convert approval to dictionary."""
return {
"id": self.id,
"pattern": self.pattern,
"approved": self.approved,
"scope": self.scope,
"created_at": self.created_at,
"expires_at": self.expires_at,
"agent_name": self.agent_name,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "PersistentApproval":
"""Create approval from dictionary."""
return cls(
id=data.get("id", str(uuid.uuid4())),
pattern=data.get("pattern", "*"),
approved=data.get("approved", False),
scope=data.get("scope", "once"),
created_at=data.get("created_at", time.time()),
expires_at=data.get("expires_at"),
agent_name=data.get("agent_name"),
)
@@ -0,0 +1,47 @@
"""
Plugin Module for PraisonAI Agents.
Provides dynamic plugin loading and hook-based extension system.
Features:
- Dynamic plugin discovery and loading
- Hook-based extension points
- Tool plugins
- Auth plugins
- Configuration plugins
Usage:
from praisonaiagents.plugins import PluginManager
# Create plugin manager
manager = PluginManager()
# Load plugins from directory
manager.load_from_directory("./plugins")
# Get loaded plugins
plugins = manager.list_plugins()
"""
__all__ = [
"PluginManager",
"Plugin",
"PluginHook",
]
def __getattr__(name: str):
"""Lazy load module components."""
if name == "PluginManager":
from .manager import PluginManager
return PluginManager
if name == "Plugin":
from .plugin import Plugin
return Plugin
if name == "PluginHook":
from .plugin import PluginHook
return PluginHook
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,299 @@
"""
Plugin Manager for PraisonAI Agents.
Provides dynamic plugin discovery, loading, and lifecycle management.
"""
import importlib.util
import logging
import os
import sys
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Type
from .plugin import Plugin, PluginHook, PluginInfo, FunctionPlugin
logger = logging.getLogger(__name__)
class PluginManager:
"""
Manages plugin discovery, loading, and execution.
Example:
manager = PluginManager()
# Load plugins from directory
manager.load_from_directory("./plugins")
# Register a plugin directly
manager.register(MyPlugin())
# Execute hooks
args = manager.execute_hook(PluginHook.BEFORE_TOOL, "bash", {"cmd": "ls"})
"""
def __init__(self):
self._plugins: Dict[str, Plugin] = {}
self._enabled: Dict[str, bool] = {}
def register(self, plugin: Plugin) -> bool:
"""
Register a plugin.
Args:
plugin: The plugin to register
Returns:
True if registered successfully
"""
try:
info = plugin.info
if info.name in self._plugins:
logger.warning(f"Plugin '{info.name}' already registered")
return False
self._plugins[info.name] = plugin
self._enabled[info.name] = True
# Initialize plugin
plugin.on_init({})
logger.info(f"Registered plugin: {info.name} v{info.version}")
return True
except Exception as e:
logger.error(f"Failed to register plugin: {e}")
return False
def unregister(self, name: str) -> bool:
"""
Unregister a plugin.
Args:
name: The plugin name
Returns:
True if unregistered successfully
"""
if name not in self._plugins:
return False
try:
plugin = self._plugins[name]
plugin.on_shutdown()
except Exception as e:
logger.warning(f"Error during plugin shutdown: {e}")
del self._plugins[name]
del self._enabled[name]
logger.info(f"Unregistered plugin: {name}")
return True
def enable(self, name: str) -> bool:
"""Enable a plugin."""
if name in self._enabled:
self._enabled[name] = True
return True
return False
def disable(self, name: str) -> bool:
"""Disable a plugin."""
if name in self._enabled:
self._enabled[name] = False
return True
return False
def is_enabled(self, name: str) -> bool:
"""Check if a plugin is enabled."""
return self._enabled.get(name, False)
def get_plugin(self, name: str) -> Optional[Plugin]:
"""Get a plugin by name."""
return self._plugins.get(name)
def list_plugins(self) -> List[PluginInfo]:
"""List all registered plugins."""
return [p.info for p in self._plugins.values()]
def load_from_directory(self, directory: str) -> int:
"""
Load plugins from a directory.
Looks for Python files with a `create_plugin` function
or classes that inherit from Plugin.
Args:
directory: Path to the plugins directory
Returns:
Number of plugins loaded
"""
directory = Path(directory)
if not directory.exists():
logger.warning(f"Plugin directory does not exist: {directory}")
return 0
loaded = 0
for file_path in directory.glob("*.py"):
if file_path.name.startswith("_"):
continue
try:
plugin = self._load_plugin_file(file_path)
if plugin and self.register(plugin):
loaded += 1
except Exception as e:
logger.error(f"Failed to load plugin from {file_path}: {e}")
return loaded
def _load_plugin_file(self, file_path: Path) -> Optional[Plugin]:
"""Load a plugin from a Python file."""
module_name = f"praison_plugin_{file_path.stem}"
spec = importlib.util.spec_from_file_location(module_name, file_path)
if spec is None or spec.loader is None:
return None
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
try:
spec.loader.exec_module(module)
except Exception as e:
logger.error(f"Error loading module {file_path}: {e}")
del sys.modules[module_name]
return None
# Look for create_plugin function
if hasattr(module, "create_plugin"):
return module.create_plugin()
# Look for Plugin subclass
for name in dir(module):
obj = getattr(module, name)
if (
isinstance(obj, type)
and issubclass(obj, Plugin)
and obj is not Plugin
and obj is not FunctionPlugin
):
return obj()
return None
def load_from_module(self, module_path: str) -> bool:
"""
Load a plugin from a module path.
Args:
module_path: Dotted module path (e.g., "mypackage.plugins.myplugin")
Returns:
True if loaded successfully
"""
try:
module = importlib.import_module(module_path)
if hasattr(module, "create_plugin"):
plugin = module.create_plugin()
return self.register(plugin)
# Look for Plugin subclass
for name in dir(module):
obj = getattr(module, name)
if (
isinstance(obj, type)
and issubclass(obj, Plugin)
and obj is not Plugin
and obj is not FunctionPlugin
):
return self.register(obj())
return False
except Exception as e:
logger.error(f"Failed to load plugin from module {module_path}: {e}")
return False
def execute_hook(
self,
hook: PluginHook,
*args,
**kwargs
) -> Any:
"""
Execute a hook across all enabled plugins.
Args:
hook: The hook to execute
*args: Arguments to pass to hook
**kwargs: Keyword arguments to pass to hook
Returns:
The result (may be modified by plugins)
"""
result = args[0] if args else None
for name, plugin in self._plugins.items():
if not self._enabled.get(name, False):
continue
if hook not in plugin.info.hooks:
continue
try:
method = getattr(plugin, hook.value, None)
if method:
if hook == PluginHook.BEFORE_TOOL:
result = method(args[0], args[1] if len(args) > 1 else {})
elif hook == PluginHook.AFTER_TOOL:
result = method(args[0], args[1] if len(args) > 1 else None)
elif hook == PluginHook.BEFORE_AGENT:
result = method(args[0], kwargs.get("context", {}))
elif hook == PluginHook.AFTER_AGENT:
result = method(args[0], kwargs.get("context", {}))
elif hook == PluginHook.BEFORE_LLM:
result = method(args[0], args[1] if len(args) > 1 else {})
elif hook == PluginHook.AFTER_LLM:
result = method(args[0], kwargs.get("usage", {}))
else:
result = method(*args, **kwargs)
except Exception as e:
logger.error(f"Error executing hook {hook.value} in plugin {name}: {e}")
return result
def get_all_tools(self) -> List[Dict[str, Any]]:
"""Get all tools from all enabled plugins."""
tools = []
for name, plugin in self._plugins.items():
if not self._enabled.get(name, False):
continue
try:
plugin_tools = plugin.get_tools()
tools.extend(plugin_tools)
except Exception as e:
logger.error(f"Error getting tools from plugin {name}: {e}")
return tools
def shutdown(self):
"""Shutdown all plugins."""
for name in list(self._plugins.keys()):
self.unregister(name)
# Global plugin manager instance
_default_manager: Optional[PluginManager] = None
def get_plugin_manager() -> PluginManager:
"""Get the global plugin manager."""
global _default_manager
if _default_manager is None:
_default_manager = PluginManager()
return _default_manager
@@ -0,0 +1,210 @@
"""
Plugin definitions for PraisonAI Agents.
Provides base plugin class and hook definitions.
"""
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Dict, List, Optional
logger = logging.getLogger(__name__)
class PluginHook(str, Enum):
"""Available plugin hooks."""
# Lifecycle hooks
ON_INIT = "on_init"
ON_SHUTDOWN = "on_shutdown"
# Agent hooks
BEFORE_AGENT = "before_agent"
AFTER_AGENT = "after_agent"
# Tool hooks
BEFORE_TOOL = "before_tool"
AFTER_TOOL = "after_tool"
# Message hooks
BEFORE_MESSAGE = "before_message"
AFTER_MESSAGE = "after_message"
# LLM hooks
BEFORE_LLM = "before_llm"
AFTER_LLM = "after_llm"
# Permission hooks
ON_PERMISSION_ASK = "on_permission_ask"
# Config hooks
ON_CONFIG = "on_config"
# Auth hooks
ON_AUTH = "on_auth"
@dataclass
class PluginInfo:
"""Plugin metadata."""
name: str
version: str = "1.0.0"
description: str = ""
author: str = ""
hooks: List[PluginHook] = field(default_factory=list)
dependencies: List[str] = field(default_factory=list)
def to_dict(self) -> Dict[str, Any]:
return {
"name": self.name,
"version": self.version,
"description": self.description,
"author": self.author,
"hooks": [h.value for h in self.hooks],
"dependencies": self.dependencies,
}
class Plugin(ABC):
"""
Base class for PraisonAI plugins.
Plugins can extend agent functionality by implementing
hook methods that are called at various points in the
agent lifecycle.
Example:
class MyPlugin(Plugin):
@property
def info(self) -> PluginInfo:
return PluginInfo(
name="my_plugin",
version="1.0.0",
hooks=[PluginHook.BEFORE_TOOL]
)
def before_tool(self, tool_name: str, args: dict) -> dict:
# Modify args before tool execution
return args
"""
@property
@abstractmethod
def info(self) -> PluginInfo:
"""Get plugin information."""
pass
def on_init(self, context: Dict[str, Any]) -> None:
"""Called when plugin is initialized."""
pass
def on_shutdown(self) -> None:
"""Called when plugin is shutting down."""
pass
def before_agent(self, prompt: str, context: Dict[str, Any]) -> str:
"""Called before agent execution. Can modify prompt."""
return prompt
def after_agent(self, response: str, context: Dict[str, Any]) -> str:
"""Called after agent execution. Can modify response."""
return response
def before_tool(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
"""Called before tool execution. Can modify args."""
return args
def after_tool(self, tool_name: str, result: Any) -> Any:
"""Called after tool execution. Can modify result."""
return result
def before_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
"""Called before message is processed. Can modify message."""
return message
def after_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
"""Called after message is processed. Can modify message."""
return message
def before_llm(self, messages: List[Dict], params: Dict[str, Any]) -> tuple:
"""Called before LLM call. Can modify messages and params."""
return messages, params
def after_llm(self, response: str, usage: Dict[str, Any]) -> str:
"""Called after LLM call. Can modify response."""
return response
def on_permission_ask(self, target: str, reason: str) -> Optional[bool]:
"""Called when permission is requested. Return True/False to auto-approve/deny."""
return None
def on_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
"""Called to modify configuration."""
return config
def on_auth(self, auth_type: str, credentials: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Called for authentication. Return credentials or None."""
return None
def get_tools(self) -> List[Dict[str, Any]]:
"""Return additional tools provided by this plugin."""
return []
class FunctionPlugin(Plugin):
"""
A simple plugin that wraps functions for specific hooks.
Example:
def my_before_tool(tool_name, args):
return args
plugin = FunctionPlugin(
name="my_plugin",
hooks={PluginHook.BEFORE_TOOL: my_before_tool}
)
"""
def __init__(
self,
name: str,
hooks: Optional[Dict[PluginHook, Callable]] = None,
version: str = "1.0.0",
description: str = "",
):
self._name = name
self._hooks = hooks or {}
self._version = version
self._description = description
@property
def info(self) -> PluginInfo:
return PluginInfo(
name=self._name,
version=self._version,
description=self._description,
hooks=list(self._hooks.keys()),
)
def before_tool(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
if PluginHook.BEFORE_TOOL in self._hooks:
return self._hooks[PluginHook.BEFORE_TOOL](tool_name, args)
return args
def after_tool(self, tool_name: str, result: Any) -> Any:
if PluginHook.AFTER_TOOL in self._hooks:
return self._hooks[PluginHook.AFTER_TOOL](tool_name, result)
return result
def before_agent(self, prompt: str, context: Dict[str, Any]) -> str:
if PluginHook.BEFORE_AGENT in self._hooks:
return self._hooks[PluginHook.BEFORE_AGENT](prompt, context)
return prompt
def after_agent(self, response: str, context: Dict[str, Any]) -> str:
if PluginHook.AFTER_AGENT in self._hooks:
return self._hooks[PluginHook.AFTER_AGENT](response, context)
return response
@@ -0,0 +1,43 @@
"""
Server Module for PraisonAI Agents.
Provides an optional HTTP server with SSE event streaming
for real-time agent communication.
Features:
- REST API for agent operations
- Server-Sent Events (SSE) for real-time updates
- Event bus integration
- Multi-project support
- CORS handling
Usage:
from praisonaiagents.server import AgentServer
# Create and start server
server = AgentServer(port=8080)
server.start()
# Or use as context manager
with AgentServer(port=8080) as server:
# Server is running
pass
"""
__all__ = [
"AgentServer",
"ServerConfig",
]
def __getattr__(name: str):
"""Lazy load module components."""
if name == "AgentServer":
from .server import AgentServer
return AgentServer
if name == "ServerConfig":
from .server import ServerConfig
return ServerConfig
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,331 @@
"""
HTTP Server implementation for PraisonAI Agents.
Provides REST API and SSE event streaming.
"""
import asyncio
import json
import logging
import queue
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Generator, List, Optional
logger = logging.getLogger(__name__)
# Default configuration
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8765
@dataclass
class ServerConfig:
"""Server configuration."""
host: str = DEFAULT_HOST
port: int = DEFAULT_PORT
cors_origins: List[str] = field(default_factory=lambda: ["*"])
auth_token: Optional[str] = None
max_connections: int = 100
def to_dict(self) -> Dict[str, Any]:
return {
"host": self.host,
"port": self.port,
"cors_origins": self.cors_origins,
"auth_token": "***" if self.auth_token else None,
"max_connections": self.max_connections,
}
class SSEClient:
"""A Server-Sent Events client connection."""
def __init__(self, client_id: str):
self.client_id = client_id
self.queue: queue.Queue = queue.Queue()
self.connected = True
self.created_at = time.time()
def send(self, event_type: str, data: Dict[str, Any]):
"""Send an event to this client."""
if self.connected:
self.queue.put({
"event": event_type,
"data": data,
"id": str(uuid.uuid4()),
})
def disconnect(self):
"""Disconnect this client."""
self.connected = False
def events(self) -> Generator[str, None, None]:
"""Generate SSE formatted events."""
while self.connected:
try:
event = self.queue.get(timeout=30)
yield f"id: {event['id']}\n"
yield f"event: {event['event']}\n"
yield f"data: {json.dumps(event['data'])}\n\n"
except queue.Empty:
# Send keepalive
yield ": keepalive\n\n"
class AgentServer:
"""
HTTP server for PraisonAI Agents.
Provides REST API endpoints and SSE event streaming
for real-time agent communication.
Example:
server = AgentServer(port=8080)
# Register event handler
@server.on_event("message")
def handle_message(data):
print(f"Message: {data}")
# Start server
server.start()
"""
def __init__(
self,
host: str = DEFAULT_HOST,
port: int = DEFAULT_PORT,
config: Optional[ServerConfig] = None,
):
"""
Initialize the server.
Args:
host: Host to bind to
port: Port to bind to
config: Optional server configuration
"""
self.config = config or ServerConfig(host=host, port=port)
self.host = self.config.host
self.port = self.config.port
self._clients: Dict[str, SSEClient] = {}
self._event_handlers: Dict[str, List[Callable]] = {}
self._running = False
self._server_thread: Optional[threading.Thread] = None
self._app = None
self._server = None
def _create_app(self):
"""Create the ASGI application."""
try:
from starlette.applications import Starlette
from starlette.responses import JSONResponse, StreamingResponse
from starlette.routing import Route
from starlette.middleware.cors import CORSMiddleware
except ImportError:
logger.warning("Starlette not installed. Server features unavailable.")
return None
async def health(request):
return JSONResponse({
"status": "ok",
"timestamp": time.time(),
"clients": len(self._clients),
})
async def events(request):
client_id = str(uuid.uuid4())
client = SSEClient(client_id)
self._clients[client_id] = client
logger.info(f"SSE client connected: {client_id}")
async def event_generator():
try:
for event in client.events():
yield event
finally:
client.disconnect()
self._clients.pop(client_id, None)
logger.info(f"SSE client disconnected: {client_id}")
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
},
)
async def publish(request):
try:
data = await request.json()
event_type = data.get("type", "message")
event_data = data.get("data", {})
self.broadcast(event_type, event_data)
return JSONResponse({
"success": True,
"clients": len(self._clients),
})
except Exception as e:
return JSONResponse(
{"success": False, "error": str(e)},
status_code=400
)
async def info(request):
return JSONResponse({
"name": "PraisonAI Agent Server",
"version": "1.0.0",
"clients": len(self._clients),
"config": self.config.to_dict(),
})
routes = [
Route("/health", health, methods=["GET"]),
Route("/events", events, methods=["GET"]),
Route("/publish", publish, methods=["POST"]),
Route("/info", info, methods=["GET"]),
]
app = Starlette(routes=routes)
# Add CORS middleware
app = CORSMiddleware(
app,
allow_origins=self.config.cors_origins,
allow_methods=["*"],
allow_headers=["*"],
)
return app
def broadcast(self, event_type: str, data: Dict[str, Any]):
"""
Broadcast an event to all connected clients.
Args:
event_type: The event type
data: The event data
"""
for client in list(self._clients.values()):
try:
client.send(event_type, data)
except Exception as e:
logger.error(f"Error broadcasting to client: {e}")
def on_event(self, event_type: str) -> Callable:
"""
Decorator to register an event handler.
Args:
event_type: The event type to handle
Returns:
Decorator function
"""
def decorator(func: Callable) -> Callable:
if event_type not in self._event_handlers:
self._event_handlers[event_type] = []
self._event_handlers[event_type].append(func)
return func
return decorator
def start(self, blocking: bool = False):
"""
Start the server.
Args:
blocking: If True, block until server stops
"""
if self._running:
logger.warning("Server already running")
return
self._app = self._create_app()
if self._app is None:
logger.error("Failed to create app - missing dependencies")
return
self._running = True
if blocking:
self._run_server()
else:
self._server_thread = threading.Thread(
target=self._run_server,
daemon=True
)
self._server_thread.start()
# Give server time to start
time.sleep(0.5)
logger.info(f"Server started on {self.host}:{self.port}")
def _run_server(self):
"""Run the server in an event loop."""
try:
import uvicorn
config = uvicorn.Config(
self._app,
host=self.host,
port=self.port,
log_level="warning",
)
self._server = uvicorn.Server(config)
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(self._server.serve())
except ImportError:
logger.error("uvicorn not installed. Cannot start server.")
except Exception as e:
logger.error(f"Server error: {e}")
finally:
self._running = False
def stop(self):
"""Stop the server."""
if not self._running:
return
self._running = False
# Disconnect all clients
for client in list(self._clients.values()):
client.disconnect()
self._clients.clear()
# Stop the server
if self._server:
self._server.should_exit = True
logger.info("Server stopped")
@property
def is_running(self) -> bool:
"""Check if server is running."""
return self._running
@property
def client_count(self) -> int:
"""Get number of connected clients."""
return len(self._clients)
def __enter__(self):
"""Context manager entry."""
self.start()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit."""
self.stop()
@@ -0,0 +1,185 @@
# Session Persistence for PraisonAI Agents
This module provides automatic session persistence with zero configuration.
## Quick Start
```python
from praisonaiagents import Agent
# With session persistence (auto-enabled)
agent = Agent(
name="Assistant",
session_id="my-session-123"
)
agent.start("Hello, my name is Alice")
# Later, in a new process - history is restored automatically
agent = Agent(
name="Assistant",
session_id="my-session-123"
)
agent.start("What is my name?") # Agent remembers: "Alice"
```
## How It Works
When you provide a `session_id` to an Agent:
1. **Automatic Persistence**: Conversation history is automatically saved to disk
2. **Automatic Restoration**: When a new Agent is created with the same `session_id`, history is restored
3. **Zero Configuration**: No database setup required - uses JSON files by default
### Default Storage Location
Sessions are stored in: `~/.praison/sessions/{session_id}.json`
### Session File Format
```json
{
"session_id": "my-session-123",
"messages": [
{"role": "user", "content": "Hello", "timestamp": 1704153600.0},
{"role": "assistant", "content": "Hi there!", "timestamp": 1704153601.5}
],
"created_at": "2026-01-02T04:00:00+00:00",
"updated_at": "2026-01-02T04:01:00+00:00",
"agent_name": "Assistant"
}
```
## Behavior Matrix
| Scenario | Behavior |
|----------|----------|
| `session_id` provided, no DB | JSON persistence (auto) |
| `session_id` provided, with DB | DB adapter used |
| No `session_id`, same Agent instance | In-memory only |
| No `session_id`, new Agent instance | No history |
## Advanced Usage
### Direct Session Store Access
```python
from praisonaiagents.session import get_default_session_store
store = get_default_session_store()
# Add messages
store.add_user_message("session-123", "Hello")
store.add_assistant_message("session-123", "Hi there!")
# Get history
history = store.get_chat_history("session-123")
# [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]
# List all sessions
sessions = store.list_sessions()
# Delete a session
store.delete_session("session-123")
```
### Custom Session Directory
```python
from praisonaiagents.session import DefaultSessionStore
store = DefaultSessionStore(
session_dir="/custom/path/sessions",
max_messages=200, # Default: 100
lock_timeout=10.0, # Default: 5.0 seconds
)
```
### Using with DB Adapter
When a DB adapter is provided, it takes precedence over JSON persistence:
```python
from praisonaiagents import Agent
from praisonai.db import PostgresAdapter
agent = Agent(
name="Assistant",
session_id="my-session",
db=PostgresAdapter(connection_string="...")
)
```
## Multi-Process Safety
The session store uses file locking to ensure safe concurrent access:
- **Unix**: Uses `fcntl.flock()` for file locking
- **Windows**: Uses `msvcrt.locking()` for file locking
- **Atomic Writes**: Uses temp file + rename to prevent corruption
Multiple processes can safely read/write to the same session.
## Context Caching
For cost optimization, use `prompt_caching=True` with Anthropic models:
```python
agent = Agent(
name="Assistant",
session_id="my-session",
prompt_caching=True, # Enables Anthropic prompt caching
)
```
This caches the system prompt, reducing token costs for repeated conversations.
## API Reference
### DefaultSessionStore
```python
class DefaultSessionStore:
def __init__(
self,
session_dir: Optional[str] = None, # Default: ~/.praison/sessions/
max_messages: int = 100,
lock_timeout: float = 5.0,
): ...
def add_message(self, session_id: str, role: str, content: str) -> bool: ...
def add_user_message(self, session_id: str, content: str) -> bool: ...
def add_assistant_message(self, session_id: str, content: str) -> bool: ...
def get_chat_history(self, session_id: str, max_messages: int = None) -> List[Dict]: ...
def get_session(self, session_id: str) -> SessionData: ...
def clear_session(self, session_id: str) -> bool: ...
def delete_session(self, session_id: str) -> bool: ...
def list_sessions(self, limit: int = 50) -> List[Dict]: ...
def session_exists(self, session_id: str) -> bool: ...
```
### SessionData
```python
@dataclass
class SessionData:
session_id: str
messages: List[SessionMessage]
created_at: str
updated_at: str
agent_name: Optional[str]
user_id: Optional[str]
metadata: Dict[str, Any]
def get_chat_history(self, max_messages: int = None) -> List[Dict[str, str]]: ...
```
### SessionMessage
```python
@dataclass
class SessionMessage:
role: str # "user", "assistant", "system"
content: str
timestamp: float
metadata: Dict[str, Any]
```
@@ -0,0 +1,76 @@
"""
Session persistence module for PraisonAI Agents.
Provides automatic session persistence with zero configuration.
When a session_id is provided to an Agent, conversation history
is automatically persisted to disk and restored on subsequent runs.
Usage:
from praisonaiagents import Agent
# With session persistence (auto-enabled)
agent = Agent(
name="Assistant",
session_id="my-session-123"
)
agent.start("Hello")
# Later, new process - history is restored
agent = Agent(
name="Assistant",
session_id="my-session-123"
)
agent.start("What did I say before?") # Remembers!
Default storage: ~/.praison/sessions/{session_id}.json
"""
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from .store import DefaultSessionStore, SessionMessage, SessionData
# Lazy loading for zero import overhead
_module_cache = {}
def __getattr__(name: str):
"""Lazy load session components."""
if name in _module_cache:
return _module_cache[name]
if name == "Session":
from .api import Session
_module_cache[name] = Session
return Session
if name == "DefaultSessionStore":
from .store import DefaultSessionStore
_module_cache[name] = DefaultSessionStore
return DefaultSessionStore
if name == "SessionMessage":
from .store import SessionMessage
_module_cache[name] = SessionMessage
return SessionMessage
if name == "SessionData":
from .store import SessionData
_module_cache[name] = SessionData
return SessionData
if name == "get_default_session_store":
from .store import get_default_session_store
_module_cache[name] = get_default_session_store
return get_default_session_store
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
"Session",
"DefaultSessionStore",
"SessionMessage",
"SessionData",
"get_default_session_store",
]
@@ -0,0 +1,532 @@
"""
Session Management for PraisonAI Agents
A simple wrapper around existing stateful capabilities to provide a unified
session API for developers building stateful agent applications.
This module uses deferred imports to ensure `requests` is NOT loaded at import time.
"""
from __future__ import annotations
import os
import uuid
import json
from typing import Any, Dict, List, Optional, TYPE_CHECKING
if TYPE_CHECKING:
from ..memory.memory import Memory
from ..knowledge.knowledge import Knowledge
from ..agent import Agent
def _get_requests():
"""Deferred import of requests library."""
import requests
return requests
class Session:
"""
A simple wrapper around PraisonAI's existing stateful capabilities.
Provides a unified API for:
- Session management with persistent state
- Memory operations (short-term, long-term, user-specific)
- Knowledge base operations
- Agent state management
- Remote agent connectivity
Examples:
# Local session with agent
session = Session(session_id="chat_123", user_id="user_456")
agent = session.Agent(name="Assistant", role="Helpful AI")
# Remote agent session (similar to Google ADK)
session = Session(agent_url="192.168.1.10:8000/agent")
response = session.chat("Hello from remote client!")
# Save session state
session.save_state({"conversation_topic": "AI research"})
"""
def __init__(
self,
session_id: Optional[str] = None,
user_id: Optional[str] = None,
agent_url: Optional[str] = None,
memory_config: Optional[Dict[str, Any]] = None,
knowledge_config: Optional[Dict[str, Any]] = None,
timeout: int = 30
):
"""
Initialize a new session with optional persistence or remote agent connectivity.
Args:
session_id: Unique session identifier. Auto-generated if None.
user_id: User identifier for user-specific memory operations.
agent_url: URL of remote agent for direct connectivity (e.g., "192.168.1.10:8000/agent")
memory_config: Configuration for memory system (defaults to RAG)
knowledge_config: Configuration for knowledge base system
timeout: HTTP timeout for remote agent calls (default: 30 seconds)
"""
self.session_id = session_id or str(uuid.uuid4())[:8]
self.user_id = user_id or "default_user"
self.agent_url = agent_url
self.timeout = timeout
self.is_remote = agent_url is not None
# Validate agent_url format
if self.is_remote:
if not self.agent_url.startswith(('http://', 'https://')):
self.agent_url = f"http://{self.agent_url}"
self._test_remote_connection()
# Initialize memory with sensible defaults (only for local sessions)
if not self.is_remote:
default_memory_config = {
"provider": "rag",
"use_embedding": False,
"rag_db_path": f".praison/sessions/{self.session_id}/chroma_db"
}
if memory_config:
default_memory_config.update(memory_config)
self.memory_config = default_memory_config
default_knowledge_config = knowledge_config or {}
self.knowledge_config = default_knowledge_config
os.makedirs(f".praison/sessions/{self.session_id}", exist_ok=True)
self._memory = None
self._knowledge = None
self._agents_instance = None
self._agents = {}
else:
self.memory_config = {}
self.knowledge_config = {}
self._memory = None
self._knowledge = None
self._agents_instance = None
self._agents = {}
@property
def memory(self) -> "Memory":
"""Lazy-loaded memory instance"""
if self.is_remote:
raise ValueError("Memory operations are not available for remote agent sessions")
if self._memory is None:
from ..memory.memory import Memory
self._memory = Memory(config=self.memory_config)
return self._memory
@property
def knowledge(self) -> "Knowledge":
"""Lazy-loaded knowledge instance"""
if self.is_remote:
raise ValueError("Knowledge operations are not available for remote agent sessions")
if self._knowledge is None:
from ..knowledge.knowledge import Knowledge
self._knowledge = Knowledge(config=self.knowledge_config)
return self._knowledge
def Agent(
self,
name: str,
role: str = "Assistant",
instructions: Optional[str] = None,
tools: Optional[List[Any]] = None,
memory: bool = True,
knowledge: Optional[List[str]] = None,
**kwargs
) -> "Agent":
"""
Create an agent with session context.
Args:
name: Agent name
role: Agent role
instructions: Agent instructions
tools: List of tools for the agent
memory: Enable memory for the agent
knowledge: Knowledge sources for the agent
**kwargs: Additional agent parameters
Returns:
Configured Agent instance
Raises:
ValueError: If this is a remote session (use chat() instead)
"""
if self.is_remote:
raise ValueError("Cannot create local agents in remote sessions. Use chat() to communicate with the remote agent.")
from ..agent import Agent as AgentClass
agent_kwargs = {
"name": name,
"role": role,
"user_id": self.user_id,
**kwargs
}
if instructions:
agent_kwargs["instructions"] = instructions
if tools:
agent_kwargs["tools"] = tools
if memory:
agent_kwargs["memory"] = self.memory
if knowledge:
agent_kwargs["knowledge"] = knowledge
agent_kwargs["knowledge_config"] = self.knowledge_config
agent = AgentClass(**agent_kwargs)
agent_key = f"{name}:{role}"
if agent_key in self._agents:
agent.chat_history = self._agents[agent_key].get("chat_history", [])
else:
restored_history = self._restore_agent_chat_history(agent_key)
agent.chat_history = restored_history
self._agents[agent_key] = {"agent": agent, "chat_history": agent.chat_history}
return agent
def create_agent(self, *args, **kwargs) -> "Agent":
"""Backward compatibility wrapper for Agent method"""
return self.Agent(*args, **kwargs)
def save_state(self, state_data: Dict[str, Any]) -> None:
"""
Save session state data to memory.
Args:
state_data: Dictionary of state data to save
Raises:
ValueError: If this is a remote session
"""
if self.is_remote:
raise ValueError("State operations are not available for remote agent sessions")
self._save_agent_chat_histories()
state_text = f"Session state: {state_data}"
self.memory.store_short_term(
text=state_text,
metadata={
"type": "session_state",
"session_id": self.session_id,
"user_id": self.user_id,
**state_data
}
)
def restore_state(self) -> Dict[str, Any]:
"""
Restore session state from memory.
Returns:
Dictionary of restored state data
Raises:
ValueError: If this is a remote session
"""
if self.is_remote:
raise ValueError("State operations are not available for remote agent sessions")
results = self.memory.search_short_term(
query="Session state:",
limit=10
)
self._restore_agent_chat_histories()
for result in results:
metadata = result.get("metadata", {})
if (metadata.get("type") == "session_state" and
metadata.get("session_id") == self.session_id):
state_data = {k: v for k, v in metadata.items()
if k not in ["type", "session_id", "user_id"]}
return state_data
return {}
def _restore_agent_chat_history(self, agent_key: str) -> List[Dict[str, Any]]:
"""
Restore agent chat history from memory.
Args:
agent_key: Unique identifier for the agent
Returns:
List of chat history messages or empty list if not found
"""
if self.is_remote:
return []
results = self.memory.search_short_term(
query="Agent chat history for",
limit=10
)
for result in results:
metadata = result.get("metadata", {})
if (metadata.get("type") == "agent_chat_history" and
metadata.get("session_id") == self.session_id and
metadata.get("agent_key") == agent_key):
chat_history = metadata.get("chat_history", [])
return chat_history
return []
def _restore_agent_chat_histories(self) -> None:
"""Restore all agent chat histories from memory."""
if self.is_remote:
return
results = self.memory.search_short_term(
query="Agent chat history for",
limit=50
)
for result in results:
metadata = result.get("metadata", {})
if (metadata.get("type") == "agent_chat_history" and
metadata.get("session_id") == self.session_id):
agent_key = metadata.get("agent_key")
chat_history = metadata.get("chat_history", [])
if agent_key and chat_history:
self._agents[agent_key] = {
"agent": None,
"chat_history": chat_history
}
def _save_agent_chat_histories(self) -> None:
"""Save all agent chat histories to memory."""
if self.is_remote:
return
for agent_key, agent_data in self._agents.items():
agent = agent_data.get("agent")
chat_history = None
if agent and hasattr(agent, 'chat_history'):
chat_history = agent.chat_history
agent_data["chat_history"] = chat_history
else:
chat_history = agent_data.get("chat_history")
if chat_history is not None:
history_text = f"Agent chat history for {agent_key}"
self.memory.store_short_term(
text=history_text,
metadata={
"type": "agent_chat_history",
"session_id": self.session_id,
"user_id": self.user_id,
"agent_key": agent_key,
"chat_history": chat_history
}
)
def get_state(self, key: str, default: Any = None) -> Any:
"""Get a specific state value"""
state = self.restore_state()
return state.get(key, default)
def set_state(self, key: str, value: Any) -> None:
"""Set a specific state value"""
current_state = self.restore_state()
current_state[key] = value
self.save_state(current_state)
def increment_state(self, key: str, increment: int = 1, default: int = 0) -> None:
"""Increment a numeric state value"""
current_value = self.get_state(key, default)
self.set_state(key, current_value + increment)
def add_memory(self, text: str, memory_type: str = "long", **metadata) -> None:
"""
Add information to session memory.
Args:
text: Text to store
memory_type: "short" or "long" term memory
**metadata: Additional metadata
"""
metadata.update({
"session_id": self.session_id,
"user_id": self.user_id
})
if memory_type == "short":
self.memory.store_short_term(text, metadata=metadata)
else:
self.memory.store_long_term(text, metadata=metadata)
def search_memory(self, query: str, memory_type: str = "long", limit: int = 5) -> List[Dict[str, Any]]:
"""
Search session memory.
Args:
query: Search query
memory_type: "short" or "long" term memory
limit: Maximum results to return
Returns:
List of memory results
"""
if memory_type == "short":
return self.memory.search_short_term(query, limit=limit)
return self.memory.search_long_term(query, limit=limit)
def add_knowledge(self, source: str) -> None:
"""
Add knowledge source to session.
Args:
source: File path, URL, or text content
"""
self.knowledge.add(source, user_id=self.user_id, agent_id=self.session_id)
def search_knowledge(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
"""
Search session knowledge base.
Args:
query: Search query
limit: Maximum results to return
Returns:
List of knowledge results
"""
return self.knowledge.search(query, agent_id=self.session_id, limit=limit)
def clear_memory(self, memory_type: str = "all") -> None:
"""
Clear session memory.
Args:
memory_type: "short", "long", or "all"
"""
if memory_type in ["short", "all"]:
self.memory.reset_short_term()
if memory_type in ["long", "all"]:
self.memory.reset_long_term()
def get_context(self, query: str, max_items: int = 3) -> str:
"""
Build context from session memory and knowledge.
Args:
query: Query to build context for
max_items: Maximum items per section
Returns:
Formatted context string
"""
return self.memory.build_context_for_task(
task_descr=query,
user_id=self.user_id,
max_items=max_items
)
def _test_remote_connection(self) -> None:
"""
Test connectivity to the remote agent.
Raises:
ConnectionError: If unable to connect to the remote agent
"""
requests = _get_requests()
try:
test_url = self.agent_url.rstrip('/') + '/health' if '/health' not in self.agent_url else self.agent_url
response = requests.get(test_url, timeout=self.timeout)
if response.status_code != 200:
response = requests.head(self.agent_url, timeout=self.timeout)
if response.status_code not in [200, 405]:
raise ConnectionError(f"Remote agent returned status code: {response.status_code}")
print(f"✅ Successfully connected to remote agent at {self.agent_url}")
except requests.exceptions.Timeout:
raise ConnectionError(f"Timeout connecting to remote agent at {self.agent_url}")
except requests.exceptions.ConnectionError:
raise ConnectionError(f"Failed to connect to remote agent at {self.agent_url}")
except Exception as e:
raise ConnectionError(f"Error connecting to remote agent: {str(e)}")
def chat(self, message: str, **kwargs) -> str:
"""
Send a message to the remote agent or handle local session.
Args:
message: The message to send to the agent
**kwargs: Additional parameters for the request
Returns:
The agent's response
Raises:
ValueError: If this is not a remote session
ConnectionError: If unable to communicate with remote agent
"""
if not self.is_remote:
raise ValueError("chat() method is only available for remote agent sessions. Use Agent.chat() for local agents.")
requests = _get_requests()
try:
payload = {
"query": message,
"session_id": self.session_id,
"user_id": self.user_id,
**kwargs
}
response = requests.post(
self.agent_url,
json=payload,
headers={"Content-Type": "application/json"},
timeout=self.timeout
)
response.raise_for_status()
result = response.json()
if isinstance(result, dict):
return result.get("response", str(result))
else:
return str(result)
except requests.exceptions.Timeout:
raise ConnectionError(f"Timeout communicating with remote agent at {self.agent_url}")
except requests.exceptions.ConnectionError:
raise ConnectionError(f"Failed to communicate with remote agent at {self.agent_url}")
except requests.exceptions.HTTPError as e:
raise ConnectionError(f"HTTP error from remote agent: {e}")
except json.JSONDecodeError:
return response.text
except Exception as e:
raise ConnectionError(f"Error communicating with remote agent: {str(e)}")
def send_message(self, message: str, **kwargs) -> str:
"""
Alias for chat() method to match Google ADK pattern.
Args:
message: The message to send to the agent
**kwargs: Additional parameters for the request
Returns:
The agent's response
"""
return self.chat(message, **kwargs)
def __str__(self) -> str:
if self.is_remote:
return f"Session(id='{self.session_id}', user='{self.user_id}', remote_agent='{self.agent_url}')"
return f"Session(id='{self.session_id}', user='{self.user_id}')"
def __repr__(self) -> str:
return self.__str__()
@@ -0,0 +1,545 @@
"""
Session Hierarchy for PraisonAI Agents.
Extends the session system with parent-child relationships,
forking, and revert capabilities.
"""
import copy
import json
import logging
import os
import threading
import time
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from .store import SessionData, SessionMessage, DefaultSessionStore, FileLock
logger = logging.getLogger(__name__)
@dataclass
class SessionSnapshot:
"""A snapshot of session state at a point in time."""
id: str = field(default_factory=lambda: str(uuid.uuid4()))
session_id: str = ""
message_index: int = 0 # Index of last message in snapshot
created_at: float = field(default_factory=time.time)
label: Optional[str] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
return {
"id": self.id,
"session_id": self.session_id,
"message_index": self.message_index,
"created_at": self.created_at,
"label": self.label,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "SessionSnapshot":
return cls(
id=data.get("id", str(uuid.uuid4())),
session_id=data.get("session_id", ""),
message_index=data.get("message_index", 0),
created_at=data.get("created_at", time.time()),
label=data.get("label"),
metadata=data.get("metadata", {}),
)
@dataclass
class ExtendedSessionData(SessionData):
"""Extended session data with hierarchy support."""
parent_id: Optional[str] = None
forked_from_message_id: Optional[str] = None
children_ids: List[str] = field(default_factory=list)
snapshots: List[SessionSnapshot] = field(default_factory=list)
is_shared: bool = False
title: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d.update({
"parent_id": self.parent_id,
"forked_from_message_id": self.forked_from_message_id,
"children_ids": self.children_ids,
"snapshots": [s.to_dict() for s in self.snapshots],
"is_shared": self.is_shared,
"title": self.title,
})
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "ExtendedSessionData":
messages = [
SessionMessage.from_dict(m)
for m in data.get("messages", [])
]
snapshots = [
SessionSnapshot.from_dict(s)
for s in data.get("snapshots", [])
]
return cls(
session_id=data.get("session_id", ""),
messages=messages,
created_at=data.get("created_at", datetime.now(timezone.utc).isoformat()),
updated_at=data.get("updated_at", datetime.now(timezone.utc).isoformat()),
agent_name=data.get("agent_name"),
user_id=data.get("user_id"),
metadata=data.get("metadata", {}),
parent_id=data.get("parent_id"),
forked_from_message_id=data.get("forked_from_message_id"),
children_ids=data.get("children_ids", []),
snapshots=snapshots,
is_shared=data.get("is_shared", False),
title=data.get("title"),
)
@classmethod
def from_session_data(cls, session: SessionData) -> "ExtendedSessionData":
"""Convert a basic SessionData to ExtendedSessionData."""
return cls(
session_id=session.session_id,
messages=session.messages,
created_at=session.created_at,
updated_at=session.updated_at,
agent_name=session.agent_name,
user_id=session.user_id,
metadata=session.metadata,
)
class HierarchicalSessionStore(DefaultSessionStore):
"""
Session store with hierarchy, forking, and revert support.
Extends DefaultSessionStore with:
- Parent-child session relationships
- Session forking from any message
- Snapshot creation and revert
- Session sharing
Usage:
store = HierarchicalSessionStore()
# Create parent session
parent_id = store.create_session(title="Main conversation")
# Fork from a message
child_id = store.fork_session(parent_id, from_message_index=5)
# Create snapshot
snapshot_id = store.create_snapshot(parent_id, label="Before refactor")
# Revert to snapshot
store.revert_to_snapshot(parent_id, snapshot_id)
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._extended_cache: Dict[str, ExtendedSessionData] = {}
def add_message(
self,
session_id: str,
role: str,
content: str,
metadata: Optional[Dict[str, Any]] = None,
) -> bool:
"""
Add a message to a session, preserving extended fields.
Overrides parent to preserve extended session data.
"""
from .store import SessionMessage
# Load extended session (force reload to get latest)
session = self._load_extended_session(session_id, force_reload=True)
message = SessionMessage(
role=role,
content=content,
timestamp=time.time(),
metadata=metadata or {},
)
session.messages.append(message)
# Trim messages if over limit
if len(session.messages) > self.max_messages:
session.messages = session.messages[-self.max_messages:]
return self._save_extended_session(session)
def _load_extended_session(self, session_id: str, force_reload: bool = False) -> ExtendedSessionData:
"""Load extended session from disk."""
filepath = self._get_session_path(session_id)
# Check cache first (unless force reload)
if not force_reload:
with self._lock:
if session_id in self._extended_cache:
return self._extended_cache[session_id]
# Load from disk
if not os.path.exists(filepath):
session = ExtendedSessionData(session_id=session_id)
with self._lock:
self._extended_cache[session_id] = session
return session
with FileLock(filepath, self.lock_timeout):
try:
with open(filepath, "r", encoding="utf-8") as f:
data = json.load(f)
session = ExtendedSessionData.from_dict(data)
except (json.JSONDecodeError, IOError) as e:
logger.warning(f"Failed to load session {session_id}: {e}")
session = ExtendedSessionData(session_id=session_id)
with self._lock:
self._extended_cache[session_id] = session
return session
def _save_extended_session(self, session: ExtendedSessionData) -> bool:
"""Save extended session to disk."""
filepath = self._get_session_path(session.session_id)
session.updated_at = datetime.now(timezone.utc).isoformat()
# Trim messages if over limit
if len(session.messages) > self.max_messages:
session.messages = session.messages[-self.max_messages:]
with FileLock(filepath, self.lock_timeout):
try:
import tempfile
dir_path = os.path.dirname(filepath)
os.makedirs(dir_path, exist_ok=True)
with tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
dir=dir_path,
delete=False,
suffix=".tmp"
) as f:
json.dump(session.to_dict(), f, indent=2, ensure_ascii=False)
temp_path = f.name
os.replace(temp_path, filepath)
with self._lock:
self._extended_cache[session.session_id] = session
return True
except (IOError, OSError) as e:
logger.error(f"Failed to save session {session.session_id}: {e}")
try:
if 'temp_path' in locals():
os.remove(temp_path)
except (IOError, OSError):
pass
return False
def create_session(
self,
session_id: Optional[str] = None,
title: Optional[str] = None,
parent_id: Optional[str] = None,
agent_name: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> str:
"""
Create a new session.
Args:
session_id: Optional custom session ID
title: Optional session title
parent_id: Optional parent session ID
agent_name: Optional agent name
metadata: Optional metadata
Returns:
The session ID
"""
sid = session_id or str(uuid.uuid4())
session = ExtendedSessionData(
session_id=sid,
title=title,
parent_id=parent_id,
agent_name=agent_name,
metadata=metadata or {},
)
# Update parent's children list
if parent_id:
parent = self._load_extended_session(parent_id)
if sid not in parent.children_ids:
parent.children_ids.append(sid)
self._save_extended_session(parent)
self._save_extended_session(session)
return sid
def fork_session(
self,
session_id: str,
from_message_index: Optional[int] = None,
title: Optional[str] = None,
) -> str:
"""
Fork a session from a specific message.
Args:
session_id: The session to fork from
from_message_index: Message index to fork from (None = all messages)
title: Optional title for the forked session
Returns:
The new forked session ID
"""
# Force reload to get latest messages from disk
parent = self._load_extended_session(session_id, force_reload=True)
# Determine which messages to copy
if from_message_index is None:
messages_to_copy = copy.deepcopy(parent.messages)
fork_msg_id = None
else:
messages_to_copy = copy.deepcopy(parent.messages[:from_message_index + 1])
fork_msg_id = str(from_message_index)
# Create new session
new_id = str(uuid.uuid4())
forked_title = title or f"Fork of {parent.title or session_id}"
forked = ExtendedSessionData(
session_id=new_id,
messages=messages_to_copy,
parent_id=session_id,
forked_from_message_id=fork_msg_id,
title=forked_title,
agent_name=parent.agent_name,
metadata=copy.deepcopy(parent.metadata),
)
# Update parent's children list
parent.children_ids.append(new_id)
self._save_extended_session(forked)
self._save_extended_session(parent)
return new_id
def get_children(self, session_id: str) -> List[str]:
"""Get all child session IDs."""
session = self._load_extended_session(session_id)
return session.children_ids.copy()
def get_parent(self, session_id: str) -> Optional[str]:
"""Get parent session ID."""
session = self._load_extended_session(session_id)
return session.parent_id
def get_session_tree(self, session_id: str) -> Dict[str, Any]:
"""
Get the full session tree starting from a session.
Returns a nested dictionary representing the tree structure.
"""
session = self._load_extended_session(session_id)
tree = {
"session_id": session.session_id,
"title": session.title,
"message_count": len(session.messages),
"children": []
}
for child_id in session.children_ids:
tree["children"].append(self.get_session_tree(child_id))
return tree
def create_snapshot(
self,
session_id: str,
label: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> str:
"""
Create a snapshot of the current session state.
Args:
session_id: The session to snapshot
label: Optional label for the snapshot
metadata: Optional metadata
Returns:
The snapshot ID
"""
# Force reload to get latest messages
session = self._load_extended_session(session_id, force_reload=True)
snapshot = SessionSnapshot(
session_id=session_id,
message_index=len(session.messages) - 1 if session.messages else -1,
label=label,
metadata=metadata or {},
)
session.snapshots.append(snapshot)
self._save_extended_session(session)
return snapshot.id
def get_snapshots(self, session_id: str) -> List[SessionSnapshot]:
"""Get all snapshots for a session."""
session = self._load_extended_session(session_id)
return session.snapshots.copy()
def revert_to_snapshot(self, session_id: str, snapshot_id: str) -> bool:
"""
Revert a session to a snapshot.
Args:
session_id: The session to revert
snapshot_id: The snapshot to revert to
Returns:
True if successful
"""
session = self._load_extended_session(session_id)
# Find the snapshot
snapshot = None
for s in session.snapshots:
if s.id == snapshot_id:
snapshot = s
break
if snapshot is None:
logger.warning(f"Snapshot {snapshot_id} not found")
return False
# Revert messages
if snapshot.message_index >= 0:
session.messages = session.messages[:snapshot.message_index + 1]
else:
session.messages = []
return self._save_extended_session(session)
def revert_to_message(self, session_id: str, message_index: int) -> bool:
"""
Revert a session to a specific message index.
Args:
session_id: The session to revert
message_index: The message index to revert to
Returns:
True if successful
"""
# Force reload to get latest messages
session = self._load_extended_session(session_id, force_reload=True)
if message_index < 0 or message_index >= len(session.messages):
logger.warning(f"Invalid message index {message_index}")
return False
session.messages = session.messages[:message_index + 1]
return self._save_extended_session(session)
def share_session(self, session_id: str) -> bool:
"""Mark a session as shared."""
session = self._load_extended_session(session_id)
session.is_shared = True
return self._save_extended_session(session)
def unshare_session(self, session_id: str) -> bool:
"""Mark a session as not shared."""
session = self._load_extended_session(session_id)
session.is_shared = False
return self._save_extended_session(session)
def is_shared(self, session_id: str) -> bool:
"""Check if a session is shared."""
session = self._load_extended_session(session_id)
return session.is_shared
def set_title(self, session_id: str, title: str) -> bool:
"""Set session title."""
session = self._load_extended_session(session_id)
session.title = title
return self._save_extended_session(session)
def get_extended_session(self, session_id: str) -> ExtendedSessionData:
"""Get extended session data."""
return self._load_extended_session(session_id)
def export_session(self, session_id: str) -> Dict[str, Any]:
"""
Export a session to a portable format.
Returns a dictionary that can be serialized to JSON.
"""
# Force reload to get latest data
session = self._load_extended_session(session_id, force_reload=True)
return session.to_dict()
def import_session(
self,
data: Dict[str, Any],
new_session_id: Optional[str] = None,
) -> str:
"""
Import a session from exported data.
Args:
data: The exported session data
new_session_id: Optional new session ID (generates one if not provided)
Returns:
The imported session ID
"""
session = ExtendedSessionData.from_dict(data)
if new_session_id:
session.session_id = new_session_id
else:
session.session_id = str(uuid.uuid4())
# Clear hierarchy references since this is an import
session.parent_id = None
session.children_ids = []
session.forked_from_message_id = None
self._save_extended_session(session)
return session.session_id
# Global hierarchical store instance
_hierarchical_store: Optional[HierarchicalSessionStore] = None
_hierarchical_lock = threading.Lock()
def get_hierarchical_session_store() -> HierarchicalSessionStore:
"""Get the global hierarchical session store instance."""
global _hierarchical_store
if _hierarchical_store is None:
with _hierarchical_lock:
if _hierarchical_store is None:
_hierarchical_store = HierarchicalSessionStore()
return _hierarchical_store
@@ -0,0 +1,457 @@
"""
Message Parts for PraisonAI Agents.
Provides structured message parts for better streaming and display,
similar to OpenCode's message-v2 parts system.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union
from enum import Enum
import time
import uuid
class PartType(str, Enum):
"""Types of message parts."""
TEXT = "text"
TOOL = "tool"
REASONING = "reasoning"
FILE = "file"
IMAGE = "image"
COMPACTION = "compaction"
ERROR = "error"
class ToolStatus(str, Enum):
"""Status of a tool execution."""
PENDING = "pending"
RUNNING = "running"
COMPLETED = "completed"
ERROR = "error"
@dataclass
class BasePart:
"""Base class for all message parts."""
id: str = field(default_factory=lambda: str(uuid.uuid4()))
type: PartType = PartType.TEXT
created_at: float = field(default_factory=time.time)
updated_at: float = field(default_factory=time.time)
def to_dict(self) -> Dict[str, Any]:
"""Convert part to dictionary."""
return {
"id": self.id,
"type": self.type.value,
"created_at": self.created_at,
"updated_at": self.updated_at,
}
def update(self):
"""Update the updated_at timestamp."""
self.updated_at = time.time()
@dataclass
class TextPart(BasePart):
"""A text content part."""
type: PartType = field(default=PartType.TEXT)
text: str = ""
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d["text"] = self.text
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "TextPart":
return cls(
id=data.get("id", str(uuid.uuid4())),
text=data.get("text", ""),
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
)
@dataclass
class ToolPart(BasePart):
"""A tool execution part."""
type: PartType = field(default=PartType.TOOL)
tool_name: str = ""
tool_call_id: str = ""
status: ToolStatus = ToolStatus.PENDING
input: Dict[str, Any] = field(default_factory=dict)
output: Optional[str] = None
error: Optional[str] = None
title: Optional[str] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d.update({
"tool_name": self.tool_name,
"tool_call_id": self.tool_call_id,
"status": self.status.value,
"input": self.input,
"output": self.output,
"error": self.error,
"title": self.title,
"metadata": self.metadata,
})
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "ToolPart":
return cls(
id=data.get("id", str(uuid.uuid4())),
tool_name=data.get("tool_name", ""),
tool_call_id=data.get("tool_call_id", ""),
status=ToolStatus(data.get("status", "pending")),
input=data.get("input", {}),
output=data.get("output"),
error=data.get("error"),
title=data.get("title"),
metadata=data.get("metadata", {}),
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
)
def start(self):
"""Mark tool as running."""
self.status = ToolStatus.RUNNING
self.update()
def complete(self, output: str, title: Optional[str] = None, metadata: Optional[Dict] = None):
"""Mark tool as completed."""
self.status = ToolStatus.COMPLETED
self.output = output
if title:
self.title = title
if metadata:
self.metadata.update(metadata)
self.update()
def fail(self, error: str):
"""Mark tool as failed."""
self.status = ToolStatus.ERROR
self.error = error
self.update()
@dataclass
class ReasoningPart(BasePart):
"""A reasoning/thinking content part."""
type: PartType = field(default=PartType.REASONING)
text: str = ""
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d["text"] = self.text
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "ReasoningPart":
return cls(
id=data.get("id", str(uuid.uuid4())),
text=data.get("text", ""),
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
)
@dataclass
class FilePart(BasePart):
"""A file attachment part."""
type: PartType = field(default=PartType.FILE)
path: str = ""
name: str = ""
mime_type: Optional[str] = None
size: Optional[int] = None
content: Optional[str] = None # Base64 encoded for binary
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d.update({
"path": self.path,
"name": self.name,
"mime_type": self.mime_type,
"size": self.size,
"content": self.content,
})
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "FilePart":
return cls(
id=data.get("id", str(uuid.uuid4())),
path=data.get("path", ""),
name=data.get("name", ""),
mime_type=data.get("mime_type"),
size=data.get("size"),
content=data.get("content"),
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
)
@dataclass
class ImagePart(BasePart):
"""An image attachment part."""
type: PartType = field(default=PartType.IMAGE)
url: Optional[str] = None
path: Optional[str] = None
base64: Optional[str] = None
mime_type: str = "image/png"
alt_text: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d.update({
"url": self.url,
"path": self.path,
"base64": self.base64,
"mime_type": self.mime_type,
"alt_text": self.alt_text,
})
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "ImagePart":
return cls(
id=data.get("id", str(uuid.uuid4())),
url=data.get("url"),
path=data.get("path"),
base64=data.get("base64"),
mime_type=data.get("mime_type", "image/png"),
alt_text=data.get("alt_text"),
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
)
@dataclass
class CompactionPart(BasePart):
"""A compaction summary part."""
type: PartType = field(default=PartType.COMPACTION)
summary: str = ""
original_tokens: int = 0
compacted_tokens: int = 0
messages_removed: int = 0
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d.update({
"summary": self.summary,
"original_tokens": self.original_tokens,
"compacted_tokens": self.compacted_tokens,
"messages_removed": self.messages_removed,
})
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "CompactionPart":
return cls(
id=data.get("id", str(uuid.uuid4())),
summary=data.get("summary", ""),
original_tokens=data.get("original_tokens", 0),
compacted_tokens=data.get("compacted_tokens", 0),
messages_removed=data.get("messages_removed", 0),
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
)
@dataclass
class ErrorPart(BasePart):
"""An error part."""
type: PartType = field(default=PartType.ERROR)
error_type: str = ""
message: str = ""
stack_trace: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
d = super().to_dict()
d.update({
"error_type": self.error_type,
"message": self.message,
"stack_trace": self.stack_trace,
})
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "ErrorPart":
return cls(
id=data.get("id", str(uuid.uuid4())),
error_type=data.get("error_type", ""),
message=data.get("message", ""),
stack_trace=data.get("stack_trace"),
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
)
# Type alias for any part
MessagePart = Union[TextPart, ToolPart, ReasoningPart, FilePart, ImagePart, CompactionPart, ErrorPart]
def part_from_dict(data: Dict[str, Any]) -> MessagePart:
"""Create a part from a dictionary based on its type."""
part_type = data.get("type", "text")
if part_type == PartType.TEXT.value:
return TextPart.from_dict(data)
elif part_type == PartType.TOOL.value:
return ToolPart.from_dict(data)
elif part_type == PartType.REASONING.value:
return ReasoningPart.from_dict(data)
elif part_type == PartType.FILE.value:
return FilePart.from_dict(data)
elif part_type == PartType.IMAGE.value:
return ImagePart.from_dict(data)
elif part_type == PartType.COMPACTION.value:
return CompactionPart.from_dict(data)
elif part_type == PartType.ERROR.value:
return ErrorPart.from_dict(data)
else:
# Default to text part
return TextPart.from_dict(data)
@dataclass
class StructuredMessage:
"""
A message with structured parts.
This extends the basic message format with support for
multiple parts of different types.
"""
id: str = field(default_factory=lambda: str(uuid.uuid4()))
role: str = "assistant" # user, assistant, system
parts: List[MessagePart] = field(default_factory=list)
created_at: float = field(default_factory=time.time)
updated_at: float = field(default_factory=time.time)
session_id: Optional[str] = None
agent_name: Optional[str] = None
model_id: Optional[str] = None
provider_id: Optional[str] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def add_part(self, part: MessagePart) -> MessagePart:
"""Add a part to the message."""
self.parts.append(part)
self.updated_at = time.time()
return part
def add_text(self, text: str) -> TextPart:
"""Add a text part."""
part = TextPart(text=text)
return self.add_part(part)
def add_tool(
self,
tool_name: str,
tool_call_id: str,
input: Optional[Dict] = None,
) -> ToolPart:
"""Add a tool part."""
part = ToolPart(
tool_name=tool_name,
tool_call_id=tool_call_id,
input=input or {},
)
return self.add_part(part)
def add_reasoning(self, text: str) -> ReasoningPart:
"""Add a reasoning part."""
part = ReasoningPart(text=text)
return self.add_part(part)
def add_file(self, path: str, name: str, **kwargs) -> FilePart:
"""Add a file part."""
part = FilePart(path=path, name=name, **kwargs)
return self.add_part(part)
def add_image(self, **kwargs) -> ImagePart:
"""Add an image part."""
part = ImagePart(**kwargs)
return self.add_part(part)
def add_error(self, error_type: str, message: str, stack_trace: Optional[str] = None) -> ErrorPart:
"""Add an error part."""
part = ErrorPart(error_type=error_type, message=message, stack_trace=stack_trace)
return self.add_part(part)
def get_part(self, part_id: str) -> Optional[MessagePart]:
"""Get a part by ID."""
for part in self.parts:
if part.id == part_id:
return part
return None
def get_parts_by_type(self, part_type: PartType) -> List[MessagePart]:
"""Get all parts of a specific type."""
return [p for p in self.parts if p.type == part_type]
def get_text_content(self) -> str:
"""Get combined text content from all text parts."""
text_parts = self.get_parts_by_type(PartType.TEXT)
return "\n".join(p.text for p in text_parts if hasattr(p, 'text'))
def get_tool_parts(self) -> List[ToolPart]:
"""Get all tool parts."""
return [p for p in self.parts if isinstance(p, ToolPart)]
def to_dict(self) -> Dict[str, Any]:
"""Convert message to dictionary."""
return {
"id": self.id,
"role": self.role,
"parts": [p.to_dict() for p in self.parts],
"created_at": self.created_at,
"updated_at": self.updated_at,
"session_id": self.session_id,
"agent_name": self.agent_name,
"model_id": self.model_id,
"provider_id": self.provider_id,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "StructuredMessage":
"""Create message from dictionary."""
parts = [part_from_dict(p) for p in data.get("parts", [])]
return cls(
id=data.get("id", str(uuid.uuid4())),
role=data.get("role", "assistant"),
parts=parts,
created_at=data.get("created_at", time.time()),
updated_at=data.get("updated_at", time.time()),
session_id=data.get("session_id"),
agent_name=data.get("agent_name"),
model_id=data.get("model_id"),
provider_id=data.get("provider_id"),
metadata=data.get("metadata", {}),
)
def to_llm_message(self) -> Dict[str, Any]:
"""Convert to LLM-compatible message format."""
content = self.get_text_content()
return {
"role": self.role,
"content": content,
}
def __repr__(self) -> str:
return f"StructuredMessage(id={self.id[:8]}..., role={self.role}, parts={len(self.parts)})"
@@ -0,0 +1,521 @@
"""
Default Session Store for PraisonAI Agents.
JSON-based session persistence with file locking and atomic writes.
Zero dependencies beyond stdlib.
"""
import fcntl
import json
import logging
import os
import sys
import tempfile
import threading
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
logger = logging.getLogger(__name__)
# Default session directory
DEFAULT_SESSION_DIR = os.path.expanduser("~/.praison/sessions")
# Default limits
DEFAULT_MAX_MESSAGES = 100
DEFAULT_LOCK_TIMEOUT = 5.0 # seconds
@dataclass
class SessionMessage:
"""A single message in a session."""
role: str # "user", "assistant", "system"
content: str
timestamp: float = field(default_factory=time.time)
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"role": self.role,
"content": self.content,
"timestamp": self.timestamp,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "SessionMessage":
"""Create from dictionary."""
return cls(
role=data.get("role", "user"),
content=data.get("content", ""),
timestamp=data.get("timestamp", time.time()),
metadata=data.get("metadata", {}),
)
@dataclass
class SessionData:
"""Complete session data structure."""
session_id: str
messages: List[SessionMessage] = field(default_factory=list)
created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
updated_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
agent_name: Optional[str] = None
user_id: Optional[str] = None
metadata: 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,
"messages": [m.to_dict() for m in self.messages],
"created_at": self.created_at,
"updated_at": self.updated_at,
"agent_name": self.agent_name,
"user_id": self.user_id,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "SessionData":
"""Create from dictionary."""
messages = [
SessionMessage.from_dict(m)
for m in data.get("messages", [])
]
return cls(
session_id=data.get("session_id", ""),
messages=messages,
created_at=data.get("created_at", datetime.now(timezone.utc).isoformat()),
updated_at=data.get("updated_at", datetime.now(timezone.utc).isoformat()),
agent_name=data.get("agent_name"),
user_id=data.get("user_id"),
metadata=data.get("metadata", {}),
)
def get_chat_history(self, max_messages: Optional[int] = None) -> List[Dict[str, str]]:
"""
Get chat history in LLM-compatible format.
Returns list of {"role": "user/assistant", "content": "..."} dicts.
"""
messages = self.messages
if max_messages and len(messages) > max_messages:
messages = messages[-max_messages:]
return [{"role": m.role, "content": m.content} for m in messages]
class FileLock:
"""
Cross-platform file locking.
Uses fcntl on Unix and msvcrt on Windows.
"""
def __init__(self, filepath: str, timeout: float = DEFAULT_LOCK_TIMEOUT):
self.filepath = filepath
self.timeout = timeout
self._lock_file = None
self._lock_path = filepath + ".lock"
def __enter__(self):
self.acquire()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.release()
return False
def acquire(self) -> bool:
"""Acquire the file lock."""
# Ensure directory exists
os.makedirs(os.path.dirname(self._lock_path) or ".", exist_ok=True)
start_time = time.time()
while True:
try:
self._lock_file = open(self._lock_path, "w")
if sys.platform == "win32":
# Windows locking
import msvcrt
msvcrt.locking(self._lock_file.fileno(), msvcrt.LK_NBLCK, 1)
else:
# Unix locking
fcntl.flock(self._lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
return True
except (IOError, OSError, BlockingIOError):
if self._lock_file:
self._lock_file.close()
self._lock_file = None
if time.time() - start_time > self.timeout:
logger.warning(f"Failed to acquire lock for {self.filepath} after {self.timeout}s")
return False
time.sleep(0.05) # Wait 50ms before retry
def release(self) -> None:
"""Release the file lock."""
if self._lock_file:
try:
if sys.platform == "win32":
import msvcrt
msvcrt.locking(self._lock_file.fileno(), msvcrt.LK_UNLCK, 1)
else:
fcntl.flock(self._lock_file.fileno(), fcntl.LOCK_UN)
except (IOError, OSError):
pass
finally:
self._lock_file.close()
self._lock_file = None
# Clean up lock file
try:
os.remove(self._lock_path)
except (IOError, OSError):
pass
class DefaultSessionStore:
"""
JSON-based session persistence with file locking.
Features:
- Zero configuration required
- Automatic file locking for multi-process safety
- Atomic writes to prevent corruption
- Configurable message limits
- Thread-safe operations
Usage:
store = DefaultSessionStore()
# Add messages
store.add_message("session-123", "user", "Hello")
store.add_message("session-123", "assistant", "Hi there!")
# Get history
history = store.get_chat_history("session-123")
# [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]
# Restore in new process
store2 = DefaultSessionStore()
history = store2.get_chat_history("session-123") # Same history!
"""
def __init__(
self,
session_dir: Optional[str] = None,
max_messages: int = DEFAULT_MAX_MESSAGES,
lock_timeout: float = DEFAULT_LOCK_TIMEOUT,
):
"""
Initialize session store.
Args:
session_dir: Directory for session files. Defaults to ~/.praison/sessions/
max_messages: Maximum messages to keep per session.
lock_timeout: Timeout for file lock acquisition.
"""
self.session_dir = session_dir or DEFAULT_SESSION_DIR
self.max_messages = max_messages
self.lock_timeout = lock_timeout
self._lock = threading.RLock()
self._cache: Dict[str, SessionData] = {}
# Ensure session directory exists
os.makedirs(self.session_dir, exist_ok=True)
def _get_session_path(self, session_id: str) -> str:
"""Get the file path for a session."""
# Sanitize session_id for filesystem
safe_id = "".join(c if c.isalnum() or c in "-_" else "_" for c in session_id)
return os.path.join(self.session_dir, f"{safe_id}.json")
def _load_session(self, session_id: str) -> SessionData:
"""Load session from disk with file locking."""
filepath = self._get_session_path(session_id)
# Check cache first
with self._lock:
if session_id in self._cache:
return self._cache[session_id]
# Load from disk
if not os.path.exists(filepath):
session = SessionData(session_id=session_id)
with self._lock:
self._cache[session_id] = session
return session
with FileLock(filepath, self.lock_timeout):
try:
with open(filepath, "r", encoding="utf-8") as f:
data = json.load(f)
session = SessionData.from_dict(data)
except (json.JSONDecodeError, IOError) as e:
logger.warning(f"Failed to load session {session_id}: {e}")
session = SessionData(session_id=session_id)
with self._lock:
self._cache[session_id] = session
return session
def _save_session(self, session: SessionData) -> bool:
"""Save session to disk with atomic write."""
filepath = self._get_session_path(session.session_id)
session.updated_at = datetime.now(timezone.utc).isoformat()
# Trim messages if over limit
if len(session.messages) > self.max_messages:
session.messages = session.messages[-self.max_messages:]
with FileLock(filepath, self.lock_timeout):
try:
# Atomic write: write to temp file, then rename
dir_path = os.path.dirname(filepath)
with tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
dir=dir_path,
delete=False,
suffix=".tmp"
) as f:
json.dump(session.to_dict(), f, indent=2, ensure_ascii=False)
temp_path = f.name
# Atomic rename
os.replace(temp_path, filepath)
return True
except (IOError, OSError) as e:
logger.error(f"Failed to save session {session.session_id}: {e}")
# Clean up temp file if it exists
try:
if 'temp_path' in locals():
os.remove(temp_path)
except (IOError, OSError):
pass
return False
def add_message(
self,
session_id: str,
role: str,
content: str,
metadata: Optional[Dict[str, Any]] = None,
) -> bool:
"""
Add a message to a session.
Args:
session_id: The session ID.
role: Message role ("user", "assistant", "system").
content: Message content.
metadata: Optional metadata.
Returns:
True if saved successfully.
"""
filepath = self._get_session_path(session_id)
message = SessionMessage(
role=role,
content=content,
timestamp=time.time(),
metadata=metadata or {},
)
# Use file lock for atomic read-modify-write
with FileLock(filepath, self.lock_timeout):
# Always reload from disk inside lock to avoid race conditions
if os.path.exists(filepath):
try:
with open(filepath, "r", encoding="utf-8") as f:
data = json.load(f)
session = SessionData.from_dict(data)
except (json.JSONDecodeError, IOError):
session = SessionData(session_id=session_id)
else:
session = SessionData(session_id=session_id)
# Add message
session.messages.append(message)
session.updated_at = datetime.now(timezone.utc).isoformat()
# Trim messages if over limit
if len(session.messages) > self.max_messages:
session.messages = session.messages[-self.max_messages:]
# Write atomically
try:
dir_path = os.path.dirname(filepath) or "."
os.makedirs(dir_path, exist_ok=True)
with tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
dir=dir_path,
delete=False,
suffix=".tmp"
) as f:
json.dump(session.to_dict(), f, indent=2, ensure_ascii=False)
temp_path = f.name
os.replace(temp_path, filepath)
# Update cache
with self._lock:
self._cache[session_id] = session
return True
except (IOError, OSError) as e:
logger.error(f"Failed to save session {session_id}: {e}")
try:
if 'temp_path' in locals():
os.remove(temp_path)
except (IOError, OSError):
pass
return False
def add_user_message(
self,
session_id: str,
content: str,
metadata: Optional[Dict[str, Any]] = None,
) -> bool:
"""Add a user message to a session."""
return self.add_message(session_id, "user", content, metadata)
def add_assistant_message(
self,
session_id: str,
content: str,
metadata: Optional[Dict[str, Any]] = None,
) -> bool:
"""Add an assistant message to a session."""
return self.add_message(session_id, "assistant", content, metadata)
def get_chat_history(
self,
session_id: str,
max_messages: Optional[int] = None,
) -> List[Dict[str, str]]:
"""
Get chat history for a session in LLM-compatible format.
Args:
session_id: The session ID.
max_messages: Maximum messages to return (defaults to store limit).
Returns:
List of {"role": "user/assistant", "content": "..."} dicts.
"""
session = self._load_session(session_id)
limit = max_messages or self.max_messages
return session.get_chat_history(limit)
def get_session(self, session_id: str) -> SessionData:
"""Get full session data."""
return self._load_session(session_id)
def set_agent_info(
self,
session_id: str,
agent_name: Optional[str] = None,
user_id: Optional[str] = None,
) -> bool:
"""Set agent info for a session."""
session = self._load_session(session_id)
with self._lock:
if agent_name:
session.agent_name = agent_name
if user_id:
session.user_id = user_id
self._cache[session_id] = session
return self._save_session(session)
def clear_session(self, session_id: str) -> bool:
"""Clear all messages from a session."""
session = self._load_session(session_id)
with self._lock:
session.messages.clear()
self._cache[session_id] = session
return self._save_session(session)
def delete_session(self, session_id: str) -> bool:
"""Delete a session completely."""
filepath = self._get_session_path(session_id)
with self._lock:
if session_id in self._cache:
del self._cache[session_id]
try:
if os.path.exists(filepath):
os.remove(filepath)
return True
except (IOError, OSError) as e:
logger.error(f"Failed to delete session {session_id}: {e}")
return False
def list_sessions(self, limit: int = 50) -> List[Dict[str, Any]]:
"""List all sessions with metadata."""
sessions = []
try:
for filename in os.listdir(self.session_dir):
if filename.endswith(".json"):
filepath = os.path.join(self.session_dir, filename)
try:
with open(filepath, "r", encoding="utf-8") as f:
data = json.load(f)
sessions.append({
"session_id": data.get("session_id", filename[:-5]),
"agent_name": data.get("agent_name"),
"created_at": data.get("created_at"),
"updated_at": data.get("updated_at"),
"message_count": len(data.get("messages", [])),
})
except (json.JSONDecodeError, IOError):
continue
except (IOError, OSError):
pass
# Sort by updated_at descending
sessions.sort(key=lambda x: x.get("updated_at", ""), reverse=True)
return sessions[:limit]
def session_exists(self, session_id: str) -> bool:
"""Check if a session exists."""
filepath = self._get_session_path(session_id)
return os.path.exists(filepath)
def invalidate_cache(self, session_id: Optional[str] = None) -> None:
"""Invalidate cache for a session or all sessions."""
with self._lock:
if session_id:
self._cache.pop(session_id, None)
else:
self._cache.clear()
# Global session store instance (lazy initialized)
_default_store: Optional[DefaultSessionStore] = None
_store_lock = threading.Lock()
def get_default_session_store() -> DefaultSessionStore:
"""Get the global default session store instance."""
global _default_store
if _default_store is None:
with _store_lock:
if _default_store is None:
_default_store = DefaultSessionStore()
return _default_store
@@ -0,0 +1,52 @@
"""
Snapshot Module for PraisonAI Agents.
Provides file change tracking using a shadow git repository,
enabling undo/restore capabilities without affecting the user's
actual git repository.
Features:
- Shadow git repository for tracking changes
- File diff generation
- Snapshot creation and restoration
- Session-based change tracking
- Zero interference with user's git repos
Usage:
from praisonaiagents.snapshot import FileSnapshot
# Initialize for a project
snapshot = FileSnapshot("/path/to/project")
# Track current state
hash = snapshot.track()
# Get diff from a snapshot
diff = snapshot.diff(hash)
# Restore to a snapshot
snapshot.restore(hash)
"""
__all__ = [
"FileSnapshot",
"SnapshotInfo",
"FileDiff",
]
def __getattr__(name: str):
"""Lazy load module components."""
if name == "FileSnapshot":
from .snapshot import FileSnapshot
return FileSnapshot
if name == "SnapshotInfo":
from .snapshot import SnapshotInfo
return SnapshotInfo
if name == "FileDiff":
from .snapshot import FileDiff
return FileDiff
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,490 @@
"""
File Snapshot implementation for PraisonAI Agents.
Provides file change tracking using a shadow git repository.
"""
import hashlib
import logging
import os
import shutil
import subprocess
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
logger = logging.getLogger(__name__)
# Default snapshot directory
DEFAULT_SNAPSHOT_DIR = os.path.expanduser("~/.praison/snapshots")
@dataclass
class FileDiff:
"""Represents a diff for a single file."""
path: str
additions: int = 0
deletions: int = 0
diff_content: str = ""
status: str = "modified" # added, modified, deleted
def to_dict(self) -> Dict[str, Any]:
return {
"path": self.path,
"additions": self.additions,
"deletions": self.deletions,
"diff_content": self.diff_content,
"status": self.status,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "FileDiff":
return cls(
path=data.get("path", ""),
additions=data.get("additions", 0),
deletions=data.get("deletions", 0),
diff_content=data.get("diff_content", ""),
status=data.get("status", "modified"),
)
@dataclass
class SnapshotInfo:
"""Information about a snapshot."""
id: str = field(default_factory=lambda: str(uuid.uuid4()))
commit_hash: str = ""
session_id: Optional[str] = None
message: str = ""
created_at: float = field(default_factory=time.time)
files_changed: int = 0
additions: int = 0
deletions: int = 0
def to_dict(self) -> Dict[str, Any]:
return {
"id": self.id,
"commit_hash": self.commit_hash,
"session_id": self.session_id,
"message": self.message,
"created_at": self.created_at,
"files_changed": self.files_changed,
"additions": self.additions,
"deletions": self.deletions,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "SnapshotInfo":
return cls(
id=data.get("id", str(uuid.uuid4())),
commit_hash=data.get("commit_hash", ""),
session_id=data.get("session_id"),
message=data.get("message", ""),
created_at=data.get("created_at", time.time()),
files_changed=data.get("files_changed", 0),
additions=data.get("additions", 0),
deletions=data.get("deletions", 0),
)
class FileSnapshot:
"""
File snapshot manager using a shadow git repository.
Creates and manages snapshots of file changes without
interfering with the user's actual git repository.
Example:
snapshot = FileSnapshot("/path/to/project")
# Track current state
info = snapshot.track(message="Before refactor")
# Make changes...
# Get diff
diffs = snapshot.diff(info.commit_hash)
# Restore if needed
snapshot.restore(info.commit_hash)
"""
def __init__(
self,
project_path: str,
snapshot_dir: Optional[str] = None,
session_id: Optional[str] = None,
):
"""
Initialize the file snapshot manager.
Args:
project_path: Path to the project to track
snapshot_dir: Optional custom snapshot directory
session_id: Optional session ID for grouping snapshots
"""
self.project_path = os.path.abspath(project_path)
self.session_id = session_id
# Create unique shadow repo path based on project path hash
project_hash = hashlib.md5(self.project_path.encode()).hexdigest()[:12]
base_dir = snapshot_dir or DEFAULT_SNAPSHOT_DIR
self.shadow_path = os.path.join(base_dir, project_hash)
self._initialized = False
def _run_git(self, *args, cwd: Optional[str] = None, check: bool = True) -> subprocess.CompletedProcess:
"""Run a git command in the shadow repository."""
cmd = ["git"] + list(args)
try:
result = subprocess.run(
cmd,
cwd=cwd or self.shadow_path,
capture_output=True,
text=True,
check=check,
)
return result
except subprocess.CalledProcessError as e:
logger.error(f"Git command failed: {' '.join(cmd)}")
logger.error(f"stderr: {e.stderr}")
raise
def _ensure_initialized(self):
"""Ensure the shadow repository is initialized."""
if self._initialized:
return
if not os.path.exists(self.shadow_path):
os.makedirs(self.shadow_path, exist_ok=True)
self._run_git("init", cwd=self.shadow_path)
self._run_git("config", "user.email", "[email protected]")
self._run_git("config", "user.name", "PraisonAI Snapshot")
# Create initial empty commit
self._run_git("commit", "--allow-empty", "-m", "Initial snapshot")
self._initialized = True
def _sync_files(self):
"""Sync project files to shadow repository."""
self._ensure_initialized()
# Get list of files to track (respecting .gitignore if exists)
files_to_track = []
gitignore_path = os.path.join(self.project_path, ".gitignore")
ignore_patterns = set()
if os.path.exists(gitignore_path):
with open(gitignore_path, "r") as f:
for line in f:
line = line.strip()
if line and not line.startswith("#"):
ignore_patterns.add(line)
# Always ignore common patterns
ignore_patterns.update([
".git",
"__pycache__",
"*.pyc",
".DS_Store",
"node_modules",
".env",
"venv",
".venv",
])
# Walk project and copy files
for root, dirs, files in os.walk(self.project_path):
# Skip ignored directories
dirs[:] = [d for d in dirs if not self._should_ignore(d, ignore_patterns)]
rel_root = os.path.relpath(root, self.project_path)
for file in files:
if self._should_ignore(file, ignore_patterns):
continue
src_path = os.path.join(root, file)
if rel_root == ".":
rel_path = file
else:
rel_path = os.path.join(rel_root, file)
dst_path = os.path.join(self.shadow_path, rel_path)
# Create directory if needed
os.makedirs(os.path.dirname(dst_path), exist_ok=True)
# Copy file
try:
shutil.copy2(src_path, dst_path)
files_to_track.append(rel_path)
except (IOError, OSError) as e:
logger.warning(f"Failed to copy {src_path}: {e}")
return files_to_track
def _should_ignore(self, name: str, patterns: set) -> bool:
"""Check if a file/directory should be ignored."""
for pattern in patterns:
if pattern.startswith("*."):
if name.endswith(pattern[1:]):
return True
elif name == pattern:
return True
return False
def track(self, message: Optional[str] = None) -> SnapshotInfo:
"""
Track the current state of the project.
Args:
message: Optional commit message
Returns:
SnapshotInfo with the commit details
"""
self._ensure_initialized()
# Sync files
files = self._sync_files()
# Stage all changes
self._run_git("add", "-A")
# Check if there are changes
status = self._run_git("status", "--porcelain")
if not status.stdout.strip():
# No changes, return current HEAD
head = self._run_git("rev-parse", "HEAD")
return SnapshotInfo(
commit_hash=head.stdout.strip(),
session_id=self.session_id,
message="No changes",
files_changed=0,
)
# Commit changes
commit_msg = message or f"Snapshot at {time.strftime('%Y-%m-%d %H:%M:%S')}"
self._run_git("commit", "-m", commit_msg)
# Get commit hash
head = self._run_git("rev-parse", "HEAD")
commit_hash = head.stdout.strip()
# Get stats
stats = self._run_git("diff", "--stat", "HEAD~1", "HEAD", check=False)
additions = 0
deletions = 0
files_changed = 0
for line in stats.stdout.split("\n"):
if "insertion" in line or "deletion" in line:
parts = line.split(",")
for part in parts:
if "insertion" in part:
try:
additions = int(part.strip().split()[0])
except (ValueError, IndexError):
pass
elif "deletion" in part:
try:
deletions = int(part.strip().split()[0])
except (ValueError, IndexError):
pass
elif "file" in part:
try:
files_changed = int(part.strip().split()[0])
except (ValueError, IndexError):
pass
return SnapshotInfo(
commit_hash=commit_hash,
session_id=self.session_id,
message=commit_msg,
files_changed=files_changed,
additions=additions,
deletions=deletions,
)
def diff(self, from_hash: str, to_hash: Optional[str] = None) -> List[FileDiff]:
"""
Get diff between two snapshots.
Args:
from_hash: Starting commit hash
to_hash: Ending commit hash (defaults to HEAD)
Returns:
List of FileDiff objects
"""
self._ensure_initialized()
to_hash = to_hash or "HEAD"
# Get list of changed files
result = self._run_git("diff", "--name-status", from_hash, to_hash, check=False)
diffs = []
for line in result.stdout.strip().split("\n"):
if not line:
continue
parts = line.split("\t")
if len(parts) < 2:
continue
status_code = parts[0]
file_path = parts[1]
# Map status code
if status_code.startswith("A"):
status = "added"
elif status_code.startswith("D"):
status = "deleted"
else:
status = "modified"
# Get detailed diff for this file
file_diff = self._run_git(
"diff", from_hash, to_hash, "--", file_path,
check=False
)
# Count additions/deletions
additions = 0
deletions = 0
for diff_line in file_diff.stdout.split("\n"):
if diff_line.startswith("+") and not diff_line.startswith("+++"):
additions += 1
elif diff_line.startswith("-") and not diff_line.startswith("---"):
deletions += 1
diffs.append(FileDiff(
path=file_path,
additions=additions,
deletions=deletions,
diff_content=file_diff.stdout,
status=status,
))
return diffs
def restore(self, commit_hash: str, files: Optional[List[str]] = None) -> bool:
"""
Restore files from a snapshot.
Args:
commit_hash: The commit hash to restore from
files: Optional list of specific files to restore
Returns:
True if successful
"""
self._ensure_initialized()
try:
if files:
# Restore specific files
for file_path in files:
# Checkout file from shadow repo
self._run_git("checkout", commit_hash, "--", file_path)
# Copy to project
src = os.path.join(self.shadow_path, file_path)
dst = os.path.join(self.project_path, file_path)
if os.path.exists(src):
os.makedirs(os.path.dirname(dst), exist_ok=True)
shutil.copy2(src, dst)
else:
# Restore all files
self._run_git("checkout", commit_hash, "--", ".")
# Copy all files back to project
for root, dirs, files_list in os.walk(self.shadow_path):
# Skip .git directory
if ".git" in root:
continue
rel_root = os.path.relpath(root, self.shadow_path)
for file in files_list:
if rel_root == ".":
rel_path = file
else:
rel_path = os.path.join(rel_root, file)
src = os.path.join(root, file)
dst = os.path.join(self.project_path, rel_path)
os.makedirs(os.path.dirname(dst), exist_ok=True)
shutil.copy2(src, dst)
return True
except Exception as e:
logger.error(f"Failed to restore: {e}")
return False
def revert(self, file_path: str) -> bool:
"""
Revert a single file to its last tracked state.
Args:
file_path: Path to the file to revert
Returns:
True if successful
"""
return self.restore("HEAD", files=[file_path])
def list_snapshots(self, limit: int = 50) -> List[SnapshotInfo]:
"""
List recent snapshots.
Args:
limit: Maximum number of snapshots to return
Returns:
List of SnapshotInfo objects
"""
self._ensure_initialized()
result = self._run_git(
"log", f"-{limit}", "--format=%H|%s|%at",
check=False
)
snapshots = []
for line in result.stdout.strip().split("\n"):
if not line:
continue
parts = line.split("|")
if len(parts) >= 3:
try:
snapshots.append(SnapshotInfo(
commit_hash=parts[0],
message=parts[1],
created_at=float(parts[2]),
session_id=self.session_id,
))
except (ValueError, IndexError):
pass
return snapshots
def get_current_hash(self) -> str:
"""Get the current HEAD commit hash."""
self._ensure_initialized()
result = self._run_git("rev-parse", "HEAD")
return result.stdout.strip()
def cleanup(self):
"""Remove the shadow repository."""
if os.path.exists(self.shadow_path):
shutil.rmtree(self.shadow_path)
self._initialized = False
@@ -35,13 +35,18 @@ class FileTools:
Raises:
ValueError: If path contains suspicious patterns
"""
# Expand ~ to user home directory FIRST (before any validation)
if filepath.startswith('~'):
filepath = os.path.expanduser(filepath)
# Normalize the path
normalized = os.path.normpath(filepath)
absolute = os.path.abspath(normalized)
# Check for suspicious patterns
if '..' in filepath or filepath.startswith('~'):
raise ValueError(f"Suspicious path pattern detected: {filepath}")
# Check for path traversal attempts (.. after normalization)
# We check the original input for '..' to catch traversal attempts
if '..' in normalized:
raise ValueError(f"Path traversal detected: {filepath}")
# Additional check: ensure the resolved path doesn't escape expected boundaries
# This is a basic check - in production, you'd want to define allowed directories
@@ -243,6 +248,76 @@ class FileTools:
logging.error(error_msg)
return False
@staticmethod
@require_approval(risk_level="medium")
def download_file(
url: str,
destination: str,
timeout: int = 30,
progress_callback: Optional[callable] = None,
) -> Dict[str, Union[bool, str, int]]:
"""
Download a file from a URL.
Args:
url: URL to download from
destination: Local path to save the file
timeout: Request timeout in seconds
progress_callback: Optional callback for progress updates
Called with (bytes_downloaded, total_bytes)
Returns:
Dict with keys: success, path, size, error
"""
try:
import httpx
except ImportError:
return {
"success": False,
"path": "",
"size": 0,
"error": "httpx not installed. Install with: pip install httpx"
}
try:
# Validate destination path
safe_path = FileTools._validate_path(destination)
# Create directory if needed
os.makedirs(os.path.dirname(safe_path) or ".", exist_ok=True)
# Download file
with httpx.stream("GET", url, timeout=timeout, follow_redirects=True) as response:
response.raise_for_status()
total_size = int(response.headers.get("content-length", 0))
downloaded = 0
with open(safe_path, "wb") as f:
for chunk in response.iter_bytes(chunk_size=8192):
f.write(chunk)
downloaded += len(chunk)
if progress_callback and total_size:
progress_callback(downloaded, total_size)
return {
"success": True,
"path": safe_path,
"size": downloaded,
"error": ""
}
except Exception as e:
error_msg = f"Error downloading {url}: {str(e)}"
logging.error(error_msg)
return {
"success": False,
"path": "",
"size": 0,
"error": error_msg
}
# Create instance for direct function access
_file_tools = FileTools()
read_file = _file_tools.read_file
@@ -252,6 +327,7 @@ get_file_info = _file_tools.get_file_info
copy_file = _file_tools.copy_file
move_file = _file_tools.move_file
delete_file = _file_tools.delete_file
download_file = _file_tools.download_file
if __name__ == "__main__":
# Example usage
@@ -0,0 +1,250 @@
"""
Sourcegraph Integration Tool for PraisonAI Agents.
Provides code search across repositories using Sourcegraph API.
"""
import os
import logging
from typing import List, Optional, Dict, Any
from dataclasses import dataclass
logger = logging.getLogger(__name__)
SOURCEGRAPH_API_URL = "https://sourcegraph.com/.api"
@dataclass
class CodeResult:
"""A code search result."""
repository: str
file_path: str
line_number: int
content: str
url: str
def to_dict(self) -> dict:
return {
"repository": self.repository,
"file_path": self.file_path,
"line_number": self.line_number,
"content": self.content,
"url": self.url,
}
class SourcegraphTools:
"""
Sourcegraph code search integration.
Provides code search across public and private repositories
using the Sourcegraph API.
Usage:
tools = SourcegraphTools()
results = tools.search("function handleError")
for result in results:
print(f"{result.repository}: {result.file_path}:{result.line_number}")
"""
def __init__(
self,
api_url: Optional[str] = None,
access_token: Optional[str] = None,
):
self.api_url = api_url or os.environ.get("SOURCEGRAPH_URL", SOURCEGRAPH_API_URL)
self.access_token = access_token or os.environ.get("SOURCEGRAPH_ACCESS_TOKEN")
def _get_headers(self) -> Dict[str, str]:
"""Get request headers."""
headers = {"Content-Type": "application/json"}
if self.access_token:
headers["Authorization"] = f"token {self.access_token}"
return headers
def search(
self,
query: str,
repo: Optional[str] = None,
file_pattern: Optional[str] = None,
max_results: int = 20,
) -> List[CodeResult]:
"""
Search code using Sourcegraph.
Args:
query: Search query
repo: Optional repository filter (e.g., "github.com/owner/repo")
file_pattern: Optional file pattern filter (e.g., "*.py")
max_results: Maximum number of results
Returns:
List of CodeResult objects
"""
try:
import httpx
except ImportError:
logger.error("httpx not installed. Install with: pip install httpx")
return []
# Build search query
search_query = query
if repo:
search_query = f"repo:{repo} {search_query}"
if file_pattern:
search_query = f"file:{file_pattern} {search_query}"
# GraphQL query
graphql_query = """
query Search($query: String!, $first: Int!) {
search(query: $query, version: V2) {
results {
results {
... on FileMatch {
repository {
name
}
file {
path
url
}
lineMatches {
lineNumber
preview
}
}
}
}
}
}
"""
try:
response = httpx.post(
f"{self.api_url}/graphql",
headers=self._get_headers(),
json={
"query": graphql_query,
"variables": {
"query": search_query,
"first": max_results,
}
},
timeout=30.0,
)
response.raise_for_status()
data = response.json()
results = []
search_results = data.get("data", {}).get("search", {}).get("results", {}).get("results", [])
for match in search_results[:max_results]:
if "repository" not in match:
continue
repo_name = match["repository"]["name"]
file_path = match.get("file", {}).get("path", "")
file_url = match.get("file", {}).get("url", "")
for line_match in match.get("lineMatches", []):
results.append(CodeResult(
repository=repo_name,
file_path=file_path,
line_number=line_match.get("lineNumber", 0),
content=line_match.get("preview", ""),
url=f"{self.api_url.replace('/.api', '')}{file_url}",
))
return results
except Exception as e:
logger.error(f"Sourcegraph search failed: {e}")
return []
def get_file_content(
self,
repo: str,
file_path: str,
revision: str = "HEAD",
) -> Optional[str]:
"""
Get file content from a repository.
Args:
repo: Repository name (e.g., "github.com/owner/repo")
file_path: Path to file in repository
revision: Git revision (default: HEAD)
Returns:
File content or None if not found
"""
try:
import httpx
except ImportError:
logger.error("httpx not installed")
return None
graphql_query = """
query GetFile($repo: String!, $path: String!, $rev: String!) {
repository(name: $repo) {
commit(rev: $rev) {
file(path: $path) {
content
}
}
}
}
"""
try:
response = httpx.post(
f"{self.api_url}/graphql",
headers=self._get_headers(),
json={
"query": graphql_query,
"variables": {
"repo": repo,
"path": file_path,
"rev": revision,
}
},
timeout=30.0,
)
response.raise_for_status()
data = response.json()
content = (
data.get("data", {})
.get("repository", {})
.get("commit", {})
.get("file", {})
.get("content")
)
return content
except Exception as e:
logger.error(f"Failed to get file content: {e}")
return None
def get_sourcegraph_tools() -> List[callable]:
"""Get Sourcegraph tools as callable functions for agent use."""
tools = SourcegraphTools()
def search_code(query: str, repo: str = None, max_results: int = 10) -> str:
"""Search code across repositories using Sourcegraph."""
results = tools.search(query, repo=repo, max_results=max_results)
if not results:
return "No results found"
output = []
for r in results:
output.append(f"{r.repository}/{r.file_path}:{r.line_number}")
output.append(f" {r.content.strip()}")
output.append("")
return "\n".join(output)
return [search_code]
@@ -0,0 +1,373 @@
"""
Subagent Tool for PraisonAI Agents.
Provides a tool for spawning subagents to handle specific tasks,
enabling hierarchical task delegation and multi-agent coordination.
"""
import logging
from typing import Any, Callable, Dict, List, Optional
logger = logging.getLogger(__name__)
def create_subagent_tool(
agent_factory: Optional[Callable[..., Any]] = None,
allowed_agents: Optional[List[str]] = None,
max_depth: int = 3,
) -> Dict[str, Any]:
"""
Create a subagent tool for task delegation.
This tool allows an agent to spawn a subagent to handle
a specific task, enabling hierarchical task execution.
Args:
agent_factory: Optional factory function to create agents
allowed_agents: Optional list of allowed agent names
max_depth: Maximum nesting depth for subagents
Returns:
Tool definition dictionary
"""
_current_depth = 0
def spawn_subagent(
task: str,
agent_name: Optional[str] = None,
context: Optional[str] = None,
tools: Optional[List[str]] = None,
) -> Dict[str, Any]:
"""
Spawn a subagent to handle a specific task.
Args:
task: The task description for the subagent
agent_name: Optional specific agent to use
context: Optional additional context
tools: Optional list of tools to give the subagent
Returns:
Result from the subagent execution
"""
nonlocal _current_depth
# Check depth limit
if _current_depth >= max_depth:
return {
"success": False,
"error": f"Maximum subagent depth ({max_depth}) exceeded",
"output": None,
}
# Check allowed agents
if allowed_agents and agent_name and agent_name not in allowed_agents:
return {
"success": False,
"error": f"Agent '{agent_name}' not in allowed list",
"output": None,
}
try:
_current_depth += 1
# If we have an agent factory, use it
if agent_factory:
subagent = agent_factory(
name=agent_name or "subagent",
tools=tools,
)
# Build prompt with context
prompt = task
if context:
prompt = f"Context: {context}\n\nTask: {task}"
# Execute the subagent
result = subagent.chat(prompt)
return {
"success": True,
"output": result,
"agent_name": agent_name or "subagent",
"task": task,
}
else:
# Without factory, return a placeholder
return {
"success": True,
"output": f"[Subagent would execute: {task}]",
"agent_name": agent_name or "subagent",
"task": task,
"note": "No agent factory provided - simulation mode",
}
except Exception as e:
logger.error(f"Subagent execution failed: {e}")
return {
"success": False,
"error": str(e),
"output": None,
}
finally:
_current_depth -= 1
return {
"name": "spawn_subagent",
"description": (
"Spawn a subagent to handle a specific task. Use this when a task "
"is complex and would benefit from dedicated focus, or when you need "
"specialized capabilities. The subagent will execute the task and "
"return the result."
),
"function": spawn_subagent,
"parameters": {
"type": "object",
"properties": {
"task": {
"type": "string",
"description": "The task description for the subagent to execute",
},
"agent_name": {
"type": "string",
"description": "Optional name of a specific agent type to use",
},
"context": {
"type": "string",
"description": "Optional additional context for the task",
},
"tools": {
"type": "array",
"items": {"type": "string"},
"description": "Optional list of tool names to give the subagent",
},
},
"required": ["task"],
},
}
def create_batch_tool() -> Dict[str, Any]:
"""
Create a batch tool for executing multiple operations.
This tool allows executing multiple similar operations
in a single call, reducing overhead.
Returns:
Tool definition dictionary
"""
def batch_execute(
operations: List[Dict[str, Any]],
parallel: bool = False,
) -> Dict[str, Any]:
"""
Execute multiple operations in batch.
Args:
operations: List of operations to execute
parallel: Whether to execute in parallel (not implemented)
Returns:
Results from all operations
"""
results = []
errors = []
for i, op in enumerate(operations):
try:
op_type = op.get("type", "unknown")
op_args = op.get("args", {})
# Placeholder for actual operation execution
results.append({
"index": i,
"type": op_type,
"success": True,
"result": f"[Would execute {op_type} with {op_args}]",
})
except Exception as e:
errors.append({
"index": i,
"error": str(e),
})
return {
"total": len(operations),
"successful": len(results),
"failed": len(errors),
"results": results,
"errors": errors,
}
return {
"name": "batch_execute",
"description": (
"Execute multiple operations in a single batch. Use this when you "
"need to perform many similar operations to reduce overhead. "
"Each operation should specify a type and arguments."
),
"function": batch_execute,
"parameters": {
"type": "object",
"properties": {
"operations": {
"type": "array",
"items": {
"type": "object",
"properties": {
"type": {"type": "string"},
"args": {"type": "object"},
},
"required": ["type"],
},
"description": "List of operations to execute",
},
"parallel": {
"type": "boolean",
"description": "Whether to execute operations in parallel",
"default": False,
},
},
"required": ["operations"],
},
}
def create_todo_tools() -> List[Dict[str, Any]]:
"""
Create todo management tools.
Returns:
List of todo tool definitions
"""
_todos: Dict[str, Dict[str, Any]] = {}
_counter = 0
def add_todo(
content: str,
priority: str = "medium",
status: str = "pending",
) -> Dict[str, Any]:
"""Add a new todo item."""
nonlocal _counter
_counter += 1
todo_id = f"todo_{_counter}"
_todos[todo_id] = {
"id": todo_id,
"content": content,
"priority": priority,
"status": status,
}
return {"success": True, "id": todo_id, "todo": _todos[todo_id]}
def update_todo(
todo_id: str,
content: Optional[str] = None,
priority: Optional[str] = None,
status: Optional[str] = None,
) -> Dict[str, Any]:
"""Update an existing todo item."""
if todo_id not in _todos:
return {"success": False, "error": f"Todo {todo_id} not found"}
todo = _todos[todo_id]
if content is not None:
todo["content"] = content
if priority is not None:
todo["priority"] = priority
if status is not None:
todo["status"] = status
return {"success": True, "todo": todo}
def list_todos(
status: Optional[str] = None,
priority: Optional[str] = None,
) -> Dict[str, Any]:
"""List todo items with optional filtering."""
todos = list(_todos.values())
if status:
todos = [t for t in todos if t["status"] == status]
if priority:
todos = [t for t in todos if t["priority"] == priority]
return {"success": True, "todos": todos, "count": len(todos)}
def delete_todo(todo_id: str) -> Dict[str, Any]:
"""Delete a todo item."""
if todo_id not in _todos:
return {"success": False, "error": f"Todo {todo_id} not found"}
del _todos[todo_id]
return {"success": True, "deleted": todo_id}
return [
{
"name": "add_todo",
"description": "Add a new todo item to track tasks",
"function": add_todo,
"parameters": {
"type": "object",
"properties": {
"content": {"type": "string", "description": "Todo content"},
"priority": {
"type": "string",
"enum": ["low", "medium", "high"],
"description": "Priority level",
},
"status": {
"type": "string",
"enum": ["pending", "in_progress", "completed"],
"description": "Current status",
},
},
"required": ["content"],
},
},
{
"name": "update_todo",
"description": "Update an existing todo item",
"function": update_todo,
"parameters": {
"type": "object",
"properties": {
"todo_id": {"type": "string", "description": "Todo ID"},
"content": {"type": "string", "description": "New content"},
"priority": {"type": "string", "description": "New priority"},
"status": {"type": "string", "description": "New status"},
},
"required": ["todo_id"],
},
},
{
"name": "list_todos",
"description": "List todo items with optional filtering",
"function": list_todos,
"parameters": {
"type": "object",
"properties": {
"status": {"type": "string", "description": "Filter by status"},
"priority": {"type": "string", "description": "Filter by priority"},
},
},
},
{
"name": "delete_todo",
"description": "Delete a todo item",
"function": delete_todo,
"parameters": {
"type": "object",
"properties": {
"todo_id": {"type": "string", "description": "Todo ID to delete"},
},
"required": ["todo_id"],
},
},
]
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "praisonaiagents"
version = "0.5.1"
version = "0.6.5"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
readme = "README.md"
requires-python = ">=3.10"
@@ -0,0 +1,149 @@
"""
Test that MCP module imports are optional and don't break basic functionality.
This test verifies the fix for issue #1091.
"""
import pytest
class TestMCPOptionalImport:
"""Test that MCP is properly optional."""
def test_praisonaiagents_import_without_mcp(self):
"""Test that praisonaiagents can be imported without mcp installed."""
# This should not raise any ImportError
import praisonaiagents
assert praisonaiagents is not None
def test_agent_import_without_mcp(self):
"""Test that Agent can be imported without mcp installed."""
from praisonaiagents import Agent
assert Agent is not None
def test_mcp_lazy_loading(self):
"""Test that MCP is lazily loaded."""
from praisonaiagents import MCP # noqa: F401
# MCP should either be available or None (not raise ImportError on access)
# The actual ImportError should only be raised when trying to instantiate
# if mcp package is not installed
def test_mcp_module_import_no_crash(self):
"""Test that mcp module can be imported without crashing."""
# This tests that the try/except in mcp.py works correctly
try:
from praisonaiagents.mcp import mcp as mcp_module
assert hasattr(mcp_module, 'MCP_AVAILABLE')
except ImportError:
# If import fails, it should be because of missing mcp,
# not because of a syntax error
pass
def test_mcp_sse_module_import_no_crash(self):
"""Test that mcp_sse module can be imported without crashing."""
try:
from praisonaiagents.mcp import mcp_sse
assert hasattr(mcp_sse, 'MCP_AVAILABLE')
except ImportError:
pass
def test_mcp_http_stream_module_import_no_crash(self):
"""Test that mcp_http_stream module can be imported without crashing."""
try:
from praisonaiagents.mcp import mcp_http_stream
assert hasattr(mcp_http_stream, 'MCP_AVAILABLE')
except ImportError:
pass
def test_basic_agent_creation_without_mcp_tools(self):
"""Test that agents can be created without MCP tools."""
from praisonaiagents import Agent
# Creating an agent without MCP tools should work
agent = Agent(
name="Test Agent",
instructions="You are a test agent.",
llm="gpt-4o-mini"
)
assert agent is not None
def test_mcp_raises_import_error_when_unavailable(self):
"""Test that MCP raises ImportError with install instructions when mcp package unavailable."""
import praisonaiagents.mcp.mcp as mcp_module
# Save original value
original_available = mcp_module.MCP_AVAILABLE
# Simulate MCP not installed
mcp_module.MCP_AVAILABLE = False
try:
from praisonaiagents.mcp.mcp import MCP
with pytest.raises(ImportError) as exc_info:
MCP('test-command')
error_msg = str(exc_info.value)
assert 'pip install praisonaiagents[mcp]' in error_msg
assert 'MCP' in error_msg or 'Model Context Protocol' in error_msg
finally:
# Restore original value
mcp_module.MCP_AVAILABLE = original_available
def test_sse_client_raises_import_error_when_unavailable(self):
"""Test that SSEMCPClient raises ImportError with install instructions when mcp unavailable."""
import praisonaiagents.mcp.mcp_sse as sse_module
original_available = sse_module.MCP_AVAILABLE
sse_module.MCP_AVAILABLE = False
try:
from praisonaiagents.mcp.mcp_sse import SSEMCPClient
with pytest.raises(ImportError) as exc_info:
SSEMCPClient('http://test.com/sse')
error_msg = str(exc_info.value)
assert 'pip install praisonaiagents[mcp]' in error_msg
finally:
sse_module.MCP_AVAILABLE = original_available
def test_http_stream_client_raises_import_error_when_mcp_unavailable(self):
"""Test that HTTPStreamMCPClient raises ImportError when mcp unavailable."""
import praisonaiagents.mcp.mcp_http_stream as http_module
original_available = http_module.MCP_AVAILABLE
http_module.MCP_AVAILABLE = False
try:
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamMCPClient
with pytest.raises(ImportError) as exc_info:
HTTPStreamMCPClient('http://test.com/mcp')
error_msg = str(exc_info.value)
assert 'pip install praisonaiagents[mcp]' in error_msg
finally:
http_module.MCP_AVAILABLE = original_available
def test_http_stream_client_raises_import_error_when_aiohttp_unavailable(self):
"""Test that HTTPStreamMCPClient raises ImportError when aiohttp unavailable."""
import praisonaiagents.mcp.mcp_http_stream as http_module
# Ensure MCP is available but aiohttp is not
original_available = http_module.MCP_AVAILABLE
original_aiohttp = http_module.aiohttp
http_module.MCP_AVAILABLE = True
http_module.aiohttp = None
try:
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamMCPClient
with pytest.raises(ImportError) as exc_info:
HTTPStreamMCPClient('http://test.com/mcp')
error_msg = str(exc_info.value)
assert 'aiohttp' in error_msg.lower()
assert 'pip install praisonaiagents[mcp]' in error_msg
finally:
http_module.MCP_AVAILABLE = original_available
http_module.aiohttp = original_aiohttp
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,215 @@
"""
Tests for Agent Profiles and Modes.
TDD: Tests for built-in agent profiles and mode configurations.
"""
from praisonaiagents.agents.profiles import (
AgentMode,
AgentProfile,
get_profile,
list_profiles,
register_profile,
get_profiles_by_mode,
BUILTIN_PROFILES,
)
class TestAgentMode:
"""Tests for AgentMode enum."""
def test_mode_values(self):
"""Test mode enum values."""
assert AgentMode.PRIMARY.value == "primary"
assert AgentMode.SUBAGENT.value == "subagent"
assert AgentMode.ALL.value == "all"
class TestAgentProfile:
"""Tests for AgentProfile."""
def test_profile_creation(self):
"""Test basic profile creation."""
profile = AgentProfile(
name="test",
description="Test agent",
mode=AgentMode.PRIMARY
)
assert profile.name == "test"
assert profile.description == "Test agent"
assert profile.mode == AgentMode.PRIMARY
def test_profile_defaults(self):
"""Test profile default values."""
profile = AgentProfile(name="test")
assert profile.mode == AgentMode.ALL
assert profile.temperature == 0.7
assert profile.top_p == 1.0
assert profile.max_steps == 50
assert profile.hidden is False
def test_profile_serialization(self):
"""Test profile round-trip."""
profile = AgentProfile(
name="test",
description="Test",
mode=AgentMode.SUBAGENT,
tools=["read_file", "bash"],
temperature=0.5,
)
d = profile.to_dict()
restored = AgentProfile.from_dict(d)
assert restored.name == profile.name
assert restored.mode == profile.mode
assert restored.tools == profile.tools
assert restored.temperature == profile.temperature
class TestBuiltinProfiles:
"""Tests for built-in profiles."""
def test_general_profile_exists(self):
"""Test general profile exists."""
profile = get_profile("general")
assert profile is not None
assert profile.name == "general"
assert profile.mode == AgentMode.PRIMARY
def test_coder_profile_exists(self):
"""Test coder profile exists."""
profile = get_profile("coder")
assert profile is not None
assert profile.name == "coder"
def test_planner_profile_exists(self):
"""Test planner profile exists."""
profile = get_profile("planner")
assert profile is not None
assert profile.mode == AgentMode.SUBAGENT
def test_reviewer_profile_exists(self):
"""Test reviewer profile exists."""
profile = get_profile("reviewer")
assert profile is not None
def test_explorer_profile_exists(self):
"""Test explorer profile exists."""
profile = get_profile("explorer")
assert profile is not None
def test_debugger_profile_exists(self):
"""Test debugger profile exists."""
profile = get_profile("debugger")
assert profile is not None
def test_hidden_profiles(self):
"""Test hidden profiles exist."""
compaction = get_profile("compaction")
title = get_profile("title")
assert compaction is not None
assert compaction.hidden is True
assert title is not None
assert title.hidden is True
def test_nonexistent_profile(self):
"""Test getting non-existent profile."""
profile = get_profile("nonexistent")
assert profile is None
class TestListProfiles:
"""Tests for listing profiles."""
def test_list_profiles(self):
"""Test listing profiles."""
profiles = list_profiles()
assert len(profiles) > 0
names = [p.name for p in profiles]
assert "general" in names
assert "coder" in names
def test_list_profiles_excludes_hidden(self):
"""Test that hidden profiles are excluded by default."""
profiles = list_profiles(include_hidden=False)
names = [p.name for p in profiles]
assert "compaction" not in names
assert "title" not in names
def test_list_profiles_includes_hidden(self):
"""Test including hidden profiles."""
profiles = list_profiles(include_hidden=True)
names = [p.name for p in profiles]
assert "compaction" in names
assert "title" in names
class TestRegisterProfile:
"""Tests for registering custom profiles."""
def test_register_profile(self):
"""Test registering a custom profile."""
custom = AgentProfile(
name="custom_test",
description="Custom test agent",
mode=AgentMode.PRIMARY,
)
register_profile(custom)
retrieved = get_profile("custom_test")
assert retrieved is not None
assert retrieved.name == "custom_test"
# Cleanup
if "custom_test" in BUILTIN_PROFILES:
del BUILTIN_PROFILES["custom_test"]
class TestGetProfilesByMode:
"""Tests for getting profiles by mode."""
def test_get_primary_profiles(self):
"""Test getting primary mode profiles."""
profiles = get_profiles_by_mode(AgentMode.PRIMARY)
assert len(profiles) > 0
# Should include PRIMARY and ALL mode profiles
for p in profiles:
assert p.mode in (AgentMode.PRIMARY, AgentMode.ALL)
def test_get_subagent_profiles(self):
"""Test getting subagent mode profiles."""
profiles = get_profiles_by_mode(AgentMode.SUBAGENT)
assert len(profiles) > 0
# Should include SUBAGENT and ALL mode profiles
for p in profiles:
assert p.mode in (AgentMode.SUBAGENT, AgentMode.ALL)
def test_all_mode_in_both(self):
"""Test that ALL mode profiles appear in both."""
primary = get_profiles_by_mode(AgentMode.PRIMARY)
subagent = get_profiles_by_mode(AgentMode.SUBAGENT)
primary_names = {p.name for p in primary}
subagent_names = {p.name for p in subagent}
# Coder is ALL mode, should be in both
assert "coder" in primary_names
assert "coder" in subagent_names
@@ -0,0 +1 @@
"""Tests for the Event Bus module."""
@@ -0,0 +1,384 @@
"""
Tests for the Event Bus module.
TDD: Tests written before implementation verification.
"""
import pytest
import asyncio
from unittest.mock import Mock, AsyncMock
from praisonaiagents.bus import EventBus, Event, EventType, get_default_bus, set_default_bus
class TestEvent:
"""Tests for the Event class."""
def test_event_creation(self):
"""Test basic event creation."""
event = Event(type="test.event", data={"key": "value"})
assert event.type == "test.event"
assert event.data == {"key": "value"}
assert event.id is not None
assert event.timestamp > 0
def test_event_with_source(self):
"""Test event with source identifier."""
event = Event(
type=EventType.SESSION_CREATED.value,
data={"session_id": "abc123"},
source="agent_1"
)
assert event.source == "agent_1"
def test_event_to_dict(self):
"""Test event serialization."""
event = Event(
type="test.event",
data={"key": "value"},
source="test"
)
d = event.to_dict()
assert d["type"] == "test.event"
assert d["data"] == {"key": "value"}
assert d["source"] == "test"
assert "id" in d
assert "timestamp" in d
def test_event_from_dict(self):
"""Test event deserialization."""
d = {
"type": "test.event",
"data": {"key": "value"},
"source": "test",
"id": "test-id",
"timestamp": 1234567890.0
}
event = Event.from_dict(d)
assert event.type == "test.event"
assert event.data == {"key": "value"}
assert event.source == "test"
assert event.id == "test-id"
assert event.timestamp == 1234567890.0
def test_event_to_sse(self):
"""Test SSE formatting."""
event = Event(type="test.event", data={"key": "value"})
sse = event.to_sse()
assert f"id: {event.id}" in sse
assert "event: test.event" in sse
assert "data:" in sse
class TestEventBus:
"""Tests for the EventBus class."""
def test_bus_creation(self):
"""Test basic bus creation."""
bus = EventBus()
assert bus.subscriber_count == 0
def test_subscribe_and_publish(self):
"""Test basic subscribe and publish."""
bus = EventBus()
received = []
def handler(event):
received.append(event)
bus.subscribe(handler)
bus.publish("test.event", {"key": "value"})
assert len(received) == 1
assert received[0].type == "test.event"
assert received[0].data == {"key": "value"}
def test_subscribe_with_filter(self):
"""Test subscribing to specific event types."""
bus = EventBus()
received = []
def handler(event):
received.append(event)
bus.subscribe(handler, event_types=["test.event"])
bus.publish("test.event", {"key": "value"})
bus.publish("other.event", {"key": "other"})
assert len(received) == 1
assert received[0].type == "test.event"
def test_subscribe_multiple_types(self):
"""Test subscribing to multiple event types."""
bus = EventBus()
received = []
def handler(event):
received.append(event)
bus.subscribe(handler, event_types=["type.a", "type.b"])
bus.publish("type.a", {})
bus.publish("type.b", {})
bus.publish("type.c", {})
assert len(received) == 2
def test_unsubscribe(self):
"""Test unsubscribing."""
bus = EventBus()
received = []
def handler(event):
received.append(event)
sub_id = bus.subscribe(handler)
bus.publish("test.event", {})
assert len(received) == 1
result = bus.unsubscribe(sub_id)
assert result is True
bus.publish("test.event", {})
assert len(received) == 1 # No new events
def test_unsubscribe_invalid_id(self):
"""Test unsubscribing with invalid ID."""
bus = EventBus()
result = bus.unsubscribe("invalid-id")
assert result is False
def test_on_decorator(self):
"""Test the @bus.on() decorator."""
bus = EventBus()
received = []
@bus.on(EventType.SESSION_CREATED)
def handler(event):
received.append(event)
bus.publish(EventType.SESSION_CREATED, {"session_id": "abc"})
assert len(received) == 1
def test_on_decorator_multiple_types(self):
"""Test decorator with multiple event types."""
bus = EventBus()
received = []
@bus.on(EventType.SESSION_CREATED, EventType.SESSION_UPDATED)
def handler(event):
received.append(event)
bus.publish(EventType.SESSION_CREATED, {})
bus.publish(EventType.SESSION_UPDATED, {})
bus.publish(EventType.SESSION_DELETED, {})
assert len(received) == 2
def test_publish_returns_event(self):
"""Test that publish returns the event."""
bus = EventBus()
event = bus.publish("test.event", {"key": "value"})
assert isinstance(event, Event)
assert event.type == "test.event"
def test_publish_with_enum(self):
"""Test publishing with EventType enum."""
bus = EventBus()
received = []
bus.subscribe(lambda e: received.append(e))
bus.publish(EventType.AGENT_STARTED, {"agent": "test"})
assert len(received) == 1
assert received[0].type == EventType.AGENT_STARTED.value
def test_event_history(self):
"""Test event history tracking."""
bus = EventBus()
bus.publish("event.1", {})
bus.publish("event.2", {})
bus.publish("event.3", {})
history = bus.get_history()
assert len(history) == 3
assert history[0].type == "event.1"
assert history[2].type == "event.3"
def test_event_history_with_filter(self):
"""Test filtered event history."""
bus = EventBus()
bus.publish("type.a", {})
bus.publish("type.b", {})
bus.publish("type.a", {})
history = bus.get_history(event_type="type.a")
assert len(history) == 2
def test_event_history_limit(self):
"""Test event history limit."""
bus = EventBus()
for i in range(10):
bus.publish(f"event.{i}", {})
history = bus.get_history(limit=5)
assert len(history) == 5
assert history[0].type == "event.5"
def test_clear_history(self):
"""Test clearing event history."""
bus = EventBus()
bus.publish("test.event", {})
bus.clear_history()
assert len(bus.get_history()) == 0
def test_clear_subscribers(self):
"""Test clearing all subscribers."""
bus = EventBus()
bus.subscribe(lambda e: None)
bus.subscribe(lambda e: None)
assert bus.subscriber_count == 2
bus.clear_subscribers()
assert bus.subscriber_count == 0
def test_subscriber_error_handling(self):
"""Test that subscriber errors don't break the bus."""
bus = EventBus()
received = []
def bad_handler(event):
raise ValueError("Test error")
def good_handler(event):
received.append(event)
bus.subscribe(bad_handler)
bus.subscribe(good_handler)
# Should not raise
bus.publish("test.event", {})
# Good handler should still receive event
assert len(received) == 1
class TestAsyncEventBus:
"""Tests for async event bus functionality."""
@pytest.mark.asyncio
async def test_async_subscriber(self):
"""Test async subscriber callback."""
bus = EventBus()
received = []
async def async_handler(event):
await asyncio.sleep(0.01)
received.append(event)
bus.subscribe(async_handler)
await bus.publish_async("test.event", {"key": "value"})
assert len(received) == 1
@pytest.mark.asyncio
async def test_mixed_sync_async_subscribers(self):
"""Test mix of sync and async subscribers."""
bus = EventBus()
sync_received = []
async_received = []
def sync_handler(event):
sync_received.append(event)
async def async_handler(event):
await asyncio.sleep(0.01)
async_received.append(event)
bus.subscribe(sync_handler)
bus.subscribe(async_handler)
await bus.publish_async("test.event", {})
assert len(sync_received) == 1
assert len(async_received) == 1
class TestDefaultBus:
"""Tests for the global default bus."""
def test_get_default_bus(self):
"""Test getting the default bus."""
bus = get_default_bus()
assert isinstance(bus, EventBus)
def test_get_default_bus_singleton(self):
"""Test that default bus is a singleton."""
bus1 = get_default_bus()
bus2 = get_default_bus()
assert bus1 is bus2
def test_set_default_bus(self):
"""Test setting a custom default bus."""
original = get_default_bus()
custom = EventBus()
set_default_bus(custom)
assert get_default_bus() is custom
# Restore original
set_default_bus(original)
class TestEventType:
"""Tests for EventType enum."""
def test_event_type_values(self):
"""Test that event types have correct values."""
assert EventType.SESSION_CREATED.value == "session.created"
assert EventType.MESSAGE_PART_UPDATED.value == "message.part.updated"
assert EventType.PERMISSION_ASKED.value == "permission.asked"
def test_event_type_in_subscribe(self):
"""Test using EventType in subscribe."""
bus = EventBus()
received = []
bus.subscribe(
lambda e: received.append(e),
event_types=[EventType.SESSION_CREATED.value]
)
bus.publish(EventType.SESSION_CREATED, {})
assert len(received) == 1
@@ -0,0 +1 @@
"""Tests for the Permissions module."""
Loaded 100 of 402 files, more files were not shown because too many files have changed in this diff. Show more