Compare commits

...
37 Commits
Author SHA1 Message Date
Mervin Praison edbd122414 Release v3.2.0 2026-01-07 19:44:05 +00:00
Mervin Praison 0ac055cba5 Release v3.1.9 2026-01-07 17:41:47 +00:00
Mervin Praison 134327a9cc Release v3.1.8 2026-01-07 15:06:12 +00:00
Mervin Praison 17c5942a02 Release v3.1.7 2026-01-07 11:50:38 +00:00
Mervin Praison 951eb6b497 Release v3.1.6 2026-01-07 10:26:35 +00:00
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
263 changed files with 51841 additions and 1662 deletions

No files matched your search

+223 -184
View File
@@ -36,10 +36,10 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
---
> **Quick Paths:**
> - 🚀 **New here?** → [Quick Start](#-quick-start)
> - 🆕 **New here?** → [Quick Start](#-quick-start) *(1 minute to first agent)*
> - 📦 **Installing?** → [Installation](#-installation)
> - 💻 **Python SDK?** → [Python Examples](#-using-python-code)
> - 🎯 **CLI user?** → [CLI Reference](#-cli--no-code-interface)
> - 🐍 **Python SDK?** → [Python Examples](#-using-python-code)
> - 🎯 **CLI user?** → [CLI Quick Reference](#cli-quick-reference)
> - 🔧 **Need config?** → [Configuration](#-configuration--integration)
> - 🤝 **Contributing?** → [Development](#-development)
@@ -47,96 +47,81 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
## 📑 Table of Contents
### Getting Started
- [⚡ Performance](#-performance)
<details open>
<summary><strong>Getting Started</strong></summary>
- [🚀 Quick Start](#-quick-start)
- [📦 Installation](#-installation)
- [Python SDK](#python-sdk)
- [JavaScript SDK](#javascript-sdk)
- [Environment Variables](#environment-variables)
- [⚡ Performance](#-performance)
### Core Concepts
- [✨ Key Features](#-key-features)
- [🌐 Supported Providers](#-supported-providers)
</details>
<details>
<summary><strong>Python SDK</strong></summary>
### Python SDK
- [📘 Python Examples](#-using-python-code)
- [1. Single Agent](#1-single-agent)
- [2. Multi Agents](#2-multi-agents)
- [3. Planning Mode](#3-agent-with-planning-mode)
- [4. Deep Research](#4-deep-research-agent)
- [5. Query Rewriter](#5-query-rewriter-agent)
- [6. Agent Memory](#6-agent-memory-zero-dependencies)
- [7. Rules & Instructions](#7-rules--instructions)
- [8. Auto-Generated Memories](#8-auto-generated-memories)
- [9. Agentic Workflows](#9-agentic-workflows)
- [10. Hooks](#10-hooks)
- [11. Shadow Git Checkpoints](#11-shadow-git-checkpoints)
- [12. Background Tasks](#12-background-tasks)
- [13. Policy Engine](#13-policy-engine)
- [14. Thinking Budgets](#14-thinking-budgets)
- [15. Output Styles](#15-output-styles)
- [16. Context Compaction](#16-context-compaction)
- [17. Field Names Reference](#17-field-names-reference-a-i-g-s)
- [18. Extended agents.yaml](#18-extended-agentsyaml-with-workflow-patterns)
- [19. MCP Protocol](#19-mcp-model-context-protocol)
- [20. A2A Protocol](#20-a2a-agent2agent-protocol)
- [1. Single Agent](#1-single-agent) | [2. Multi Agents](#2-multi-agents) | [3. Planning Mode](#3-agent-with-planning-mode)
- [4. Deep Research](#4-deep-research-agent) | [5. Query Rewriter](#5-query-rewriter-agent) | [6. Agent Memory](#6-agent-memory-zero-dependencies)
- [7. Rules & Instructions](#7-rules--instructions) | [8. Auto-Generated Memories](#8-auto-generated-memories) | [9. Agentic Workflows](#9-agentic-workflows)
- [10. Hooks](#10-hooks) | [11. Shadow Git Checkpoints](#11-shadow-git-checkpoints) | [12. Background Tasks](#12-background-tasks)
- [13. Policy Engine](#13-policy-engine) | [14. Thinking Budgets](#14-thinking-budgets) | [15. Output Styles](#15-output-styles)
- [16. Context Compaction](#16-context-compaction) | [17. Field Names Reference](#17-field-names-reference-a-i-g-s) | [18. Extended agents.yaml](#18-extended-agentsyaml-with-workflow-patterns)
- [19. MCP Protocol](#19-mcp-model-context-protocol) | [20. A2A Protocol](#20-a2a-agent2agent-protocol)
- [🛠️ Custom Tools](#️-custom-tools)
### JavaScript SDK
</details>
<details>
<summary><strong>JavaScript SDK</strong></summary>
- [💻 JavaScript Examples](#-using-javascript-code)
### CLI Reference
- [🎯 CLI Overview](#-cli--no-code-interface)
- [Auto Mode](#auto-mode)
- [Interactive Mode](#interactive-mode-cli)
- [Deep Research CLI](#deep-research-cli)
- [Planning Mode CLI](#planning-mode-cli)
- [Memory CLI](#memory-cli)
- [Workflow CLI](#workflow-cli)
- [Knowledge CLI](#knowledge-cli)
- [Session CLI](#session-cli)
- [Tools CLI](#tools-cli)
- [MCP Config CLI](#mcp-config-cli)
- [External Agents CLI](#external-agents-cli)
- [CLI Features Summary](#cli-features)
</details>
### Configuration
- [🔧 Configuration & Integration](#-configuration--integration)
- [Ollama Integration](#ollama-integration)
- [Groq Integration](#groq-integration)
- [100+ Models Support](#100-models-support)
<details>
<summary><strong>CLI Reference</strong></summary>
- [🎯 CLI Overview](#-cli--no-code-interface) | [CLI Quick Reference](#cli-quick-reference)
- [Auto Mode](#auto-mode) | [Interactive Mode](#interactive-mode-cli) | [Deep Research CLI](#deep-research-cli) | [Planning Mode CLI](#planning-mode-cli)
- [Memory CLI](#memory-cli) | [Workflow CLI](#workflow-cli) | [Knowledge CLI](#knowledge-cli) | [Session CLI](#session-cli)
- [Tools CLI](#tools-cli) | [MCP Config CLI](#mcp-config-cli) | [External Agents CLI](#external-agents-cli) | [CLI Features Summary](#cli-features)
</details>
<details>
<summary><strong>Configuration & Features</strong></summary>
- [✨ Key Features](#-key-features) | [🌐 Supported Providers](#-supported-providers)
- [🔧 Configuration & Integration](#-configuration--integration) | [Ollama](#ollama-integration) | [Groq](#groq-integration) | [100+ Models](#100-models-support)
- [📋 Agents Playbook](#-agents-playbook)
### Advanced Features
- [🔬 Advanced Features](#-advanced-features)
- [Research & Intelligence](#research--intelligence)
- [Memory & Caching](#memory--caching)
- [Planning & Workflows](#planning--workflows)
- [Specialised Agents](#specialised-agents)
- [MCP Protocol](#mcp-protocol)
- [A2A & A2UI Protocols](#a2a--a2ui-protocols)
- [Safety & Control](#safety--control)
- [Developer Tools](#developer-tools)
### Architecture
</details>
<details>
<summary><strong>Architecture & Patterns</strong></summary>
- [📊 Process Types & Patterns](#-process-types--patterns)
- [Sequential Process](#sequential-process)
- [Hierarchical Process](#hierarchical-process)
- [Workflow Process](#workflow-process)
- [Agentic Patterns](#agentic-patterns)
- [Sequential](#sequential-process) | [Hierarchical](#hierarchical-process) | [Workflow](#workflow-process) | [Agentic Patterns](#agentic-patterns)
</details>
<details>
<summary><strong>Data & Persistence</strong></summary>
### Data & Persistence
- [💾 Persistence (Databases)](#-persistence-databases)
- [📚 Knowledge & Retrieval (RAG)](#-knowledge--retrieval-rag)
- [🔧 Tools Table](#-tools-table)
### Learning & Community
- [🎓 Video Tutorials](#-video-tutorials)
- [⭐ Star History](#-star-history)
- [👥 Contributing](#-contributing)
- [🔧 Development](#-development)
- [❓ FAQ & Troubleshooting](#-faq--troubleshooting)
</details>
<details>
<summary><strong>Learning & Community</strong></summary>
- [🎓 Video Tutorials](#-video-tutorials) | [⭐ Star History](#-star-history)
- [👥 Contributing](#-contributing) | [🔧 Development](#-development) | [❓ FAQ & Troubleshooting](#-faq--troubleshooting)
</details>
---
@@ -181,6 +166,8 @@ export OPENAI_API_KEY=your_key_here
python -c "from praisonaiagents import Agent; Agent(instructions='You are a helpful AI assistant').start('Write a haiku about AI')"
```
> **Next Steps:** [Single Agent Example](#1-single-agent) | [Multi Agents](#2-multi-agents) | [CLI Auto Mode](#auto-mode)
---
## 📦 Installation
@@ -231,28 +218,6 @@ export OPENAI_BASE_URL=https://api.groq.com/openai/v1
---
## 💻 Usage
### Python Code Examples
- [📘 Complete Python Examples](#using-python-code)
- [🔧 Custom Tools](#️-custom-tools)
- [🧠 Memory Configuration](#-memory--context)
### CLI / No-Code Interface
- [🎯 Auto Mode](#auto-mode)
- [🔬 Deep Research](#deep-research-cli)
- [📋 Planning Mode](#planning-mode-cli)
- [💾 Memory Management](#memory-cli)
- [🔄 Workflow Execution](#workflow-cli)
### JavaScript Code Examples
- [💻 JavaScript Usage](#using-javascript-code)
---
## ✨ Key Features
<details open>
@@ -519,6 +484,9 @@ export OPENAI_BASE_URL=https://api.groq.com/openai/v1
PraisonAI supports 100+ LLM providers through seamless integration:
<details>
<summary><strong>View all 24 providers</strong></summary>
| Provider | Example |
|----------|:-------:|
| OpenAI | [Example](examples/python/providers/openai/openai_gpt4_example.py) |
@@ -546,17 +514,12 @@ PraisonAI supports 100+ LLM providers through seamless integration:
| Moonshot | [Example](examples/python/providers/moonshot/moonshot_example.py) |
| vLLM | [Example](examples/python/providers/vllm/vllm_example.py) |
</details>
---
## 📘 Using Python Code
Light weight package dedicated for coding:
```bash
pip install praisonaiagents
export OPENAI_API_KEY=xxxxxxxxxxxxxxxxxxxxxx
```
### 1. Single Agent
Create app.py file and add the code below:
@@ -2405,6 +2368,9 @@ agent.start('Write a movie script about a robot in Mars');
## 📊 Process Types & Patterns
<details>
<summary><strong>View architecture diagrams and workflow patterns</strong></summary>
### AI Agents Flow
```mermaid
@@ -2744,6 +2710,8 @@ flowchart LR
style Out fill:#8B0000,color:#fff
```
</details>
---
## 🔧 Configuration & Integration
@@ -2891,77 +2859,179 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
PraisonAI provides a complete knowledge stack for building RAG applications with multiple vector stores, retrieval strategies, rerankers, and query modes.
### RAG Quickstart (Agent-first)
```python
from praisonaiagents import Agent
from praisonaiagents.rag.models import RetrievalStrategy
# Agent with RAG - simplest approach
agent = Agent(
name="Research Assistant",
knowledge=["docs/manual.pdf", "data/faq.txt"],
knowledge_config={"vector_store": {"provider": "chroma"}},
rag_config={
"include_citations": True,
"retrieval_strategy": RetrievalStrategy.HYBRID, # Dense + BM25
"rerank": True,
}
)
# Query with citations
result = agent.rag_query("How do I authenticate?")
print(result.answer)
for citation in result.citations:
print(f" [{citation.id}] {citation.source}")
```
### RAG CLI Commands
| Command | Description |
|---------|-------------|
| `praisonai rag query "<question>"` | One-shot question answering with citations |
| `praisonai rag chat` | Interactive RAG chat session |
| `praisonai rag serve` | Start RAG as a microservice API |
| `praisonai rag eval <test_file>` | Evaluate RAG retrieval quality |
### RAG CLI Examples
```bash
# Query with hybrid retrieval (dense + BM25 keyword search)
praisonai rag query "What are the key findings?" --hybrid
# Query with hybrid + reranking for best quality
praisonai rag query "Summarize conclusions" --hybrid --rerank
# Interactive chat with hybrid retrieval
praisonai rag chat --collection research --hybrid --rerank
# Start API server with OpenAI-compatible endpoint
praisonai rag serve --hybrid --rerank --openai-compat --port 8080
# Query with profiling
praisonai rag query "Summary?" --profile --profile-out ./profile.json
```
### Knowledge CLI Commands
| Command | Description |
|---------|-------------|
| `praisonai knowledge add <file\|dir\|url>` | Add documents to knowledge base |
| `praisonai knowledge query <question>` | Query knowledge base with RAG |
| `praisonai knowledge index <sources>` | Index documents into knowledge base |
| `praisonai knowledge search <query>` | Search knowledge base (no LLM generation) |
| `praisonai knowledge list` | List indexed documents |
| `praisonai knowledge clear` | Clear knowledge base |
| `praisonai knowledge stats` | Show knowledge base statistics |
### Knowledge CLI Options
| Option | Values | Description |
|--------|--------|-------------|
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | Vector store backend |
| `--retrieval` | `basic`, `fusion`, `recursive`, `auto_merge` | Retrieval strategy |
| `--reranker` | `simple`, `llm`, `cross_encoder`, `cohere` | Reranking method |
| `--index-type` | `vector`, `keyword`, `hybrid` | Index type |
| `--query-mode` | `default`, `sub_question`, `summarize` | Query mode |
### Knowledge CLI Examples
```bash
# Add documents
praisonai knowledge add ./docs/
praisonai knowledge add https://example.com/page.html
praisonai knowledge add "*.pdf"
# Index documents
praisonai knowledge index ./docs/ --collection myproject
# Query with advanced options
praisonai knowledge query "How to authenticate?" --retrieval fusion --reranker llm
# Search with hybrid retrieval
praisonai knowledge search "authentication" --hybrid --collection myproject
# Full advanced query
praisonai knowledge query "authentication flow" \
--vector-store chroma \
--retrieval fusion \
--reranker llm \
--index-type hybrid \
--query-mode sub_question
# Index with profiling
praisonai knowledge index ./data --profile --profile-out ./profile.json
```
### Knowledge SDK Usage
### Knowledge vs RAG vs AutoRagAgent
- **Knowledge** is the indexing and retrieval substrate - use for indexing and raw search
- **RAG** orchestrates on top - use for question answering with LLM-generated responses and citations
- **AutoRagAgent** wraps an Agent with automatic retrieval decision - use when you want the agent to decide when to retrieve
- All share the same underlying index
### AutoRagAgent (Automatic RAG)
AutoRagAgent automatically decides when to retrieve context from knowledge bases vs direct chat, based on query heuristics.
```python
from praisonaiagents import Agent, Knowledge
from praisonaiagents import Agent, AutoRagAgent
# Simple usage with Agent
# Create agent with knowledge
agent = Agent(
name="Research Assistant",
knowledge=["docs/manual.pdf", "data/faq.txt"],
knowledge_config={
"vector_store": {"provider": "chroma"}
}
knowledge=["docs/manual.pdf"],
user_id="user123", # Required for RAG retrieval
)
response = agent.chat("How do I authenticate?")
# Direct Knowledge usage
knowledge = Knowledge()
knowledge.add("document.pdf")
results = knowledge.search("authentication", limit=5)
# Wrap with AutoRagAgent
auto_rag = AutoRagAgent(
agent=agent,
retrieval_policy="auto", # auto, always, never
top_k=5,
hybrid=True,
rerank=True,
)
# Auto-decides: retrieves for questions, skips for greetings
result = auto_rag.chat("What are the key findings?") # Retrieves
result = auto_rag.chat("Hello!") # Skips retrieval
# Force retrieval or skip per-call
result = auto_rag.chat("Hi", force_retrieval=True)
result = auto_rag.chat("Summary?", skip_retrieval=True)
```
**CLI Usage:**
```bash
# Enable auto-rag with default policy (auto)
praisonai --auto-rag "What are the key findings?"
# Always retrieve
praisonai --auto-rag --rag-policy always "Tell me about X"
# With hybrid retrieval and reranking
praisonai --auto-rag --rag-hybrid --rag-rerank "Summarize the document"
```
### Configuration Precedence
Settings are applied in this order (highest priority first):
1. **CLI flags** - `--hybrid`, `--rerank`, `--top-k`
2. **Environment variables** - `PRAISONAI_HYBRID=true`
3. **Config file** - YAML configuration (`--config`)
4. **Defaults**
```bash
# Environment variables
export PRAISONAI_HYBRID=true
export PRAISONAI_RERANK=true
export PRAISONAI_TOP_K=10
```
### Lightweight Installs
```bash
# Base install (minimal, fast imports)
pip install praisonaiagents
# With RAG API server support
pip install "praisonai[rag-api]"
```
### Live Tests (Real API Keys)
Run integration tests with real API keys:
```bash
# Enable live tests
export PRAISONAI_LIVE_TESTS=1
export OPENAI_API_KEY="your-key"
# Run live tests
pytest -m live tests/integration/
```
### Knowledge Stack Features Table
| Feature | Description | SDK Docs | CLI Docs |
|---------|-------------|----------|----------|
| **Data Readers** | Load PDF, Markdown, Text, HTML, URLs | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
| **Hybrid Retrieval** | Dense vectors + BM25 keyword search with RRF fusion | [SDK](/docs/rag/module) | [CLI](/docs/cli/rag) |
| **Reranking** | LLM, Cross-Encoder, Cohere rerankers | [SDK](/docs/rag/module) | [CLI](/docs/cli/rag) |
| **RAG Serve** | Microservice API with OpenAI-compatible mode | [SDK](/docs/rag/module) | [CLI](/docs/cli/rag) |
| **Vector Stores** | ChromaDB, Pinecone, Qdrant, Weaviate, In-Memory | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
| **Retrieval Strategies** | Basic, Fusion (RRF), Recursive, Auto-Merge | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
| **Rerankers** | Simple, LLM, Cross-Encoder, Cohere | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
| **Index Types** | Vector, Keyword (BM25), Hybrid | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
| **Query Engines** | Default, Sub-Question, Summarize | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
| **Data Readers** | Load PDF, Markdown, Text, HTML, URLs | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
| **Profiling** | Performance profiling with `--profile` flag | [SDK](/docs/features/profiling) | [CLI](/docs/cli/rag) |
---
@@ -3079,43 +3149,7 @@ agent.chat("Hello!") # Auto-persists messages, runs, traces
| `praisonai persistence migrate` | Apply schema migrations |
| `praisonai persistence status` | Show schema status |
### Knowledge CLI Commands {#knowledge-cli}
| Command | Description |
|---------|-------------|
| `praisonai knowledge add <source>` | Add file, directory, URL, or glob pattern |
| `praisonai knowledge query "<question>"` | Query knowledge base with RAG |
| `praisonai knowledge list` | List indexed documents |
| `praisonai knowledge clear` | Clear knowledge base |
| `praisonai knowledge stats` | Show knowledge base statistics |
**Knowledge Query Flags:**
| Flag | Values | Default |
|------|--------|---------|
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | `chroma` |
| `--retrieval-strategy` | `basic`, `fusion`, `recursive`, `auto_merge` | `basic` |
| `--reranker` | `none`, `simple`, `llm`, `cross_encoder`, `cohere` | `none` |
| `--index-type` | `vector`, `keyword`, `hybrid` | `vector` |
| `--query-mode` | `default`, `sub_question`, `summarize` | `default` |
| `--workspace` | Path to workspace directory | Current dir |
| `--session` | Session ID for persistence | - |
**Examples:**
```bash
# Add documents
praisonai knowledge add document.pdf
praisonai knowledge add ./docs/
praisonai knowledge add "*.md"
# Query with options
praisonai knowledge query "How to authenticate?" \
--vector-store chroma \
--retrieval-strategy fusion \
--reranker simple \
--query-mode sub_question
```
> **Note:** See [Knowledge & Retrieval (RAG)](#-knowledge--retrieval-rag) for complete Knowledge CLI documentation.
### Databases Table
@@ -3191,6 +3225,9 @@ Install all tools with: `pip install "praisonai[tools]"`
Learn PraisonAI through our comprehensive video series:
<details>
<summary><strong>View all 22 video tutorials</strong></summary>
| Topic | Video |
|-------|--------|
| AI Agents with Self Reflection | [![Self Reflection](https://img.youtube.com/vi/vLXobEN2Vc8/0.jpg)](https://www.youtube.com/watch?v=vLXobEN2Vc8) |
@@ -3216,6 +3253,8 @@ Learn PraisonAI through our comprehensive video series:
| Call Interface | [![Call](https://img.youtube.com/vi/m1cwrUG2iAk/0.jpg)](https://www.youtube.com/watch?v=m1cwrUG2iAk) |
| Reasoning Extract Agents | [![Reasoning Extract](https://img.youtube.com/vi/2PPamsADjJA/0.jpg)](https://www.youtube.com/watch?v=2PPamsADjJA) |
</details>
---
## 👥 Contributing
+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.9.2" \
"praisonai>=3.2.0" \
"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.9.2" \
"praisonai>=3.2.0" \
"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.9.2" \
"praisonai>=3.2.0" \
"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
+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()
+101
View File
@@ -0,0 +1,101 @@
#!/usr/bin/env python3
"""
Context Budgeting Example
Demonstrates how to use the ContextBudgeter to allocate token budgets
across different context segments for various models.
"""
from praisonaiagents.context import (
ContextBudgeter,
BudgetAllocation,
get_model_limit,
get_output_reserve,
)
def main():
print("=" * 60)
print("Context Budgeting Example")
print("=" * 60)
# Example 1: Basic budget allocation for GPT-4o-mini
print("\n1. Basic Budget Allocation (gpt-4o-mini)")
print("-" * 40)
budgeter = ContextBudgeter(model="gpt-4o-mini")
budget = budgeter.allocate()
print(f"Model: gpt-4o-mini")
print(f"Model Limit: {budget.model_limit:,} tokens")
print(f"Output Reserve: {budget.output_reserve:,} tokens")
print(f"Usable Context: {budget.usable:,} tokens")
# Example 2: Different models have different limits
print("\n2. Model Comparison")
print("-" * 40)
models = ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "claude-3-opus", "gemini-1.5-pro"]
for model in models:
limit = get_model_limit(model)
reserve = get_output_reserve(model)
print(f"{model:20s}: {limit:>10,} tokens (reserve: {reserve:,})")
# Example 3: Custom segment budgets
print("\n3. Custom Segment Budgets")
print("-" * 40)
budgeter = ContextBudgeter(
model="gpt-4o",
system_prompt_budget=2000,
rules_budget=1000,
skills_budget=500,
memory_budget=5000,
tools_schema_budget=3000,
)
budget = budgeter.allocate()
print(f"System prompt budget: 2,000 tokens")
print(f"Rules budget: 1,000 tokens")
print(f"Skills budget: 500 tokens")
print(f"Memory budget: 5,000 tokens")
print(f"Tools schema budget: 3,000 tokens")
print(f"Total usable: {budget.usable:,} tokens")
# Example 4: Check for overflow
print("\n4. Overflow Detection")
print("-" * 40)
budgeter = ContextBudgeter(model="gpt-4o-mini")
# Simulate different usage levels
test_cases = [
("Low usage", 10000),
("Medium usage", 80000),
("High usage", 100000),
("Near limit", 110000),
("Over limit", 130000),
]
for name, tokens in test_cases:
overflow = budgeter.check_overflow(tokens)
utilization = budgeter.get_utilization(tokens)
remaining = budgeter.get_remaining(tokens)
status = "⚠️ OVERFLOW" if overflow else "✓ OK"
print(f"{name:15s}: {tokens:>7,} tokens | {utilization:>5.1%} | {status}")
# Example 5: Budget to dict for serialization
print("\n5. Budget Serialization")
print("-" * 40)
budget_dict = budgeter.to_dict()
print(f"Budget as dict: {budget_dict}")
print("\n" + "=" * 60)
print("✓ Context budgeting examples complete!")
print("=" * 60)
if __name__ == "__main__":
main()
+184
View File
@@ -0,0 +1,184 @@
#!/usr/bin/env python3
"""
Context Monitoring Example
Demonstrates how to use the ContextMonitor to write runtime context
snapshots to disk for debugging and analysis.
"""
import os
import tempfile
from praisonaiagents.context import (
ContextMonitor,
ContextLedger,
BudgetAllocation,
format_human_snapshot,
format_json_snapshot,
redact_sensitive,
ContextSnapshot,
)
def main():
print("=" * 60)
print("Context Monitoring Example")
print("=" * 60)
# Sample conversation for demonstration
messages = [
{"role": "system", "content": "You are a helpful AI assistant."},
{"role": "user", "content": "What is the weather like today?"},
{"role": "assistant", "content": "I don't have access to real-time weather data."},
{"role": "user", "content": "Can you help me write Python code?"},
{"role": "assistant", "content": "Of course! I'd be happy to help with Python."},
]
# Example 1: Basic monitoring setup
print("\n1. Basic Monitor Setup")
print("-" * 40)
with tempfile.TemporaryDirectory() as tmpdir:
context_path = os.path.join(tmpdir, "context.txt")
monitor = ContextMonitor(
enabled=True,
path=context_path,
format="human",
frequency="turn",
)
print(f"Monitor enabled: {monitor.enabled}")
print(f"Output path: {context_path}")
print(f"Format: {monitor.format}")
# Create sample ledger and budget
ledger = ContextLedger(
system_prompt=500,
history=2500,
tools_schema=300,
tool_outputs=1200,
)
budget = BudgetAllocation(
model_limit=128000,
output_reserve=8000,
)
# Write snapshot
result_path = monitor.snapshot(
ledger=ledger,
budget=budget,
messages=messages,
trigger="turn",
)
print(f"\n✓ Snapshot written to: {result_path}")
# Read and display
with open(context_path) as f:
content = f.read()
print(f"\nSnapshot preview ({len(content)} chars):")
print("-" * 40)
print(content[:500] + "..." if len(content) > 500 else content)
# Example 2: JSON format
print("\n\n2. JSON Format Output")
print("-" * 40)
with tempfile.TemporaryDirectory() as tmpdir:
json_path = os.path.join(tmpdir, "context.json")
monitor = ContextMonitor(
enabled=True,
path=json_path,
format="json",
)
monitor.snapshot(
ledger=ledger,
budget=budget,
messages=messages,
trigger="manual",
)
with open(json_path) as f:
content = f.read()
print(f"JSON snapshot ({len(content)} chars):")
print(content[:300] + "..." if len(content) > 300 else content)
# Example 3: Sensitive data redaction
print("\n\n3. Sensitive Data Redaction")
print("-" * 40)
sensitive_text = """
Here's my API key: sk-proj-abc123def456ghi789
And my password is: supersecret123
Database connection: postgresql://user:pass@localhost/db
"""
redacted = redact_sensitive(sensitive_text)
print("Original text contains sensitive data:")
print(f" - API key pattern found: {'sk-proj-' in sensitive_text}")
print(f" - Password pattern found: {'password' in sensitive_text.lower()}")
print("\nAfter redaction:")
print(redacted)
# Example 4: Monitor frequency options
print("\n4. Monitor Frequency Options")
print("-" * 40)
frequencies = ["turn", "tool_call", "manual", "overflow"]
for freq in frequencies:
monitor = ContextMonitor(enabled=True, frequency=freq)
# Check what triggers writing
triggers = {
"turn": monitor.should_write("turn"),
"tool_call": monitor.should_write("tool_call"),
"manual": monitor.should_write("manual"),
"overflow": monitor.should_write("overflow"),
}
active = [k for k, v in triggers.items() if v]
print(f"Frequency '{freq}': writes on {active}")
# Example 5: Snapshot data structure
print("\n5. Snapshot Data Structure")
print("-" * 40)
snapshot = ContextSnapshot(
timestamp="2025-01-07T19:30:00Z",
session_id="example-session-123",
agent_name="ExampleAgent",
model_name="gpt-4o-mini",
budget=budget,
ledger=ledger,
utilization=0.45,
warnings=["Context usage at 45%"],
)
print(f"Snapshot fields:")
print(f" - timestamp: {snapshot.timestamp}")
print(f" - session_id: {snapshot.session_id}")
print(f" - agent_name: {snapshot.agent_name}")
print(f" - model_name: {snapshot.model_name}")
print(f" - utilization: {snapshot.utilization:.1%}")
print(f" - warnings: {snapshot.warnings}")
# Format as human-readable
human_output = format_human_snapshot(snapshot)
print(f"\nHuman-readable format ({len(human_output)} chars)")
# Format as JSON
json_output = format_json_snapshot(snapshot)
print(f"JSON format ({len(json_output)} chars)")
print("\n" + "=" * 60)
print("✓ Context monitoring examples complete!")
print("=" * 60)
if __name__ == "__main__":
main()
+142
View File
@@ -0,0 +1,142 @@
#!/usr/bin/env python3
"""
Context Optimization Example
Demonstrates how to use various optimization strategies to reduce
context size when approaching model limits.
"""
from praisonaiagents.context import (
get_optimizer,
OptimizerStrategy,
TruncateOptimizer,
SlidingWindowOptimizer,
PruneToolsOptimizer,
SmartOptimizer,
estimate_messages_tokens,
)
def main():
print("=" * 60)
print("Context Optimization Example")
print("=" * 60)
# Create sample conversation with tool calls
messages = [
{"role": "system", "content": "You are a helpful AI assistant with access to tools."},
{"role": "user", "content": "What's the weather in Paris?"},
{"role": "assistant", "content": "", "tool_calls": [{"id": "call_1", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}}]},
{"role": "tool", "tool_call_id": "call_1", "content": "Weather in Paris: 15°C, partly cloudy, humidity 65%, wind 10 km/h from the west. Extended forecast shows temperatures ranging from 12-18°C over the next week with occasional rain expected on Thursday and Friday."},
{"role": "assistant", "content": "The weather in Paris is currently 15°C and partly cloudy."},
{"role": "user", "content": "What about London?"},
{"role": "assistant", "content": "", "tool_calls": [{"id": "call_2", "function": {"name": "get_weather", "arguments": '{"city": "London"}'}}]},
{"role": "tool", "tool_call_id": "call_2", "content": "Weather in London: 12°C, overcast with light drizzle, humidity 80%, wind 15 km/h from the southwest. Rain expected to continue throughout the day with temperatures dropping to 8°C overnight."},
{"role": "assistant", "content": "London is 12°C with light drizzle."},
]
# Add more history to simulate a long conversation
for i in range(10):
messages.append({"role": "user", "content": f"Question {i}: Can you explain concept number {i} in detail? " * 20})
messages.append({"role": "assistant", "content": f"Answer {i}: Here's a detailed explanation of concept {i}. " * 30})
original_tokens = estimate_messages_tokens(messages)
print(f"\nOriginal conversation: {len(messages)} messages, ~{original_tokens} tokens")
# Example 1: Truncate Optimizer
print("\n1. Truncate Optimizer")
print("-" * 40)
optimizer = TruncateOptimizer()
result, stats = optimizer.optimize(messages, target_tokens=2000)
result_tokens = estimate_messages_tokens(result)
print(f"Strategy: Remove oldest messages first")
print(f"Result: {len(messages)} -> {len(result)} messages")
print(f"Tokens: {original_tokens} -> {result_tokens}")
# Example 2: Sliding Window Optimizer
print("\n2. Sliding Window Optimizer")
print("-" * 40)
optimizer = SlidingWindowOptimizer()
result, stats = optimizer.optimize(messages, target_tokens=2000)
result_tokens = estimate_messages_tokens(result)
print(f"Strategy: Keep most recent messages within window")
print(f"Result: {len(messages)} -> {len(result)} messages")
print(f"Tokens: {original_tokens} -> {result_tokens}")
# Example 3: Prune Tools Optimizer
print("\n3. Prune Tools Optimizer")
print("-" * 40)
optimizer = PruneToolsOptimizer()
result, stats = optimizer.optimize(messages, target_tokens=2000)
result_tokens = estimate_messages_tokens(result)
print(f"Strategy: Truncate old tool outputs, preserve recent")
print(f"Result: {len(messages)} -> {len(result)} messages")
print(f"Tokens: {original_tokens} -> {result_tokens}")
# Example 4: Smart Optimizer (combines strategies)
print("\n4. Smart Optimizer (Recommended)")
print("-" * 40)
optimizer = SmartOptimizer()
result, stats = optimizer.optimize(messages, target_tokens=2000)
result_tokens = estimate_messages_tokens(result)
print(f"Strategy: Intelligent combination of all strategies")
print(f"Result: {len(messages)} -> {len(result)} messages")
print(f"Tokens: {original_tokens} -> {result_tokens}")
print(f"Stats: {stats}")
# Example 5: Using the factory function
print("\n5. Optimizer Factory")
print("-" * 40)
strategies = [
OptimizerStrategy.TRUNCATE,
OptimizerStrategy.SLIDING_WINDOW,
OptimizerStrategy.PRUNE_TOOLS,
OptimizerStrategy.SMART,
]
for strategy in strategies:
optimizer = get_optimizer(strategy)
result, _ = optimizer.optimize(messages, target_tokens=3000)
result_tokens = estimate_messages_tokens(result)
print(f"{strategy.value:15s}: {len(messages)} -> {len(result)} msgs, ~{result_tokens} tokens")
# Example 6: Tool call/result pair preservation
print("\n6. Tool Call/Result Pair Preservation")
print("-" * 40)
tool_messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Get data"},
{"role": "assistant", "content": "", "tool_calls": [{"id": "tc1"}]},
{"role": "tool", "tool_call_id": "tc1", "content": "Data result"},
{"role": "assistant", "content": "Here's the data."},
]
optimizer = SmartOptimizer()
result, _ = optimizer.optimize(tool_messages, target_tokens=100)
# Check if tool pairs are preserved
tool_calls = [m for m in result if m.get("tool_calls")]
tool_results = [m for m in result if m.get("role") == "tool"]
print(f"Original: {len(tool_messages)} messages")
print(f"Optimized: {len(result)} messages")
print(f"Tool calls preserved: {len(tool_calls)}")
print(f"Tool results preserved: {len(tool_results)}")
print("\n" + "=" * 60)
print("✓ Context optimization examples complete!")
print("=" * 60)
if __name__ == "__main__":
main()
+158
View File
@@ -0,0 +1,158 @@
#!/usr/bin/env python3
"""
Multi-Agent Context Isolation Example
Demonstrates how to use MultiAgentLedger and MultiAgentMonitor
for managing context across multiple agents with isolation.
"""
import os
import tempfile
from praisonaiagents.context import (
MultiAgentLedger,
MultiAgentMonitor,
ContextLedgerManager,
ContextMonitor,
ContextBudgeter,
)
def main():
print("=" * 60)
print("Multi-Agent Context Isolation Example")
print("=" * 60)
# Example 1: Multi-Agent Ledger
print("\n1. Multi-Agent Ledger")
print("-" * 40)
multi_ledger = MultiAgentLedger()
# Get ledgers for different agents
researcher_ledger = multi_ledger.get_agent_ledger("researcher")
writer_ledger = multi_ledger.get_agent_ledger("writer")
reviewer_ledger = multi_ledger.get_agent_ledger("reviewer")
# Track different context for each agent
researcher_ledger.track_system_prompt("You are a research specialist.")
researcher_ledger.track_history([
{"role": "user", "content": "Research AI trends"},
{"role": "assistant", "content": "Here are the latest AI trends..."},
])
writer_ledger.track_system_prompt("You are a technical writer.")
writer_ledger.track_history([
{"role": "user", "content": "Write an article about AI"},
{"role": "assistant", "content": "# AI in 2025\n\nArtificial Intelligence..."},
])
reviewer_ledger.track_system_prompt("You are a content reviewer.")
reviewer_ledger.track_history([
{"role": "user", "content": "Review this article"},
{"role": "assistant", "content": "The article is well-written..."},
])
# Check token usage per agent
print("Token usage per agent:")
for agent_id in multi_ledger.get_agent_ids():
ledger = multi_ledger.get_agent_ledger(agent_id)
print(f" {agent_id}: {ledger.get_total()} tokens")
# Get combined total
total = multi_ledger.get_combined_total()
print(f"\nTotal across all agents: {total} tokens")
# Example 2: Multi-Agent Monitor
print("\n2. Multi-Agent Monitor")
print("-" * 40)
with tempfile.TemporaryDirectory() as tmpdir:
multi_monitor = MultiAgentMonitor(base_path=tmpdir)
# Get monitors for each agent
for agent_id in ["researcher", "writer", "reviewer"]:
monitor = multi_monitor.get_agent_monitor(agent_id)
print(f"Monitor for {agent_id}: {monitor.path}")
# Enable all monitors
multi_monitor.enable_all()
print(f"\nAll monitors enabled: {multi_monitor.enabled}")
# List all agent IDs
agent_ids = multi_monitor.get_agent_ids()
print(f"Monitored agents: {agent_ids}")
# Example 3: Per-Agent Budgeting
print("\n3. Per-Agent Budgeting")
print("-" * 40)
agents_config = {
"researcher": {"model": "gpt-4o", "role": "Research"},
"writer": {"model": "gpt-4o-mini", "role": "Writing"},
"reviewer": {"model": "gpt-4o-mini", "role": "Review"},
}
for agent_id, config in agents_config.items():
budgeter = ContextBudgeter(model=config["model"])
budget = budgeter.allocate()
print(f"{agent_id} ({config['model']}): {budget.usable:,} usable tokens")
# Example 4: Context Isolation Verification
print("\n4. Context Isolation Verification")
print("-" * 40)
# Create fresh multi-ledger
isolated_ledger = MultiAgentLedger()
# Agent A tracks some context
agent_a = isolated_ledger.get_agent_ledger("agent_a")
agent_a.track_system_prompt("Agent A system prompt")
agent_a.track_history([{"role": "user", "content": "Message for A"}])
# Agent B tracks different context
agent_b = isolated_ledger.get_agent_ledger("agent_b")
agent_b.track_system_prompt("Agent B system prompt")
agent_b.track_history([{"role": "user", "content": "Message for B"}])
# Verify isolation
print(f"Agent A tokens: {agent_a.get_total()}")
print(f"Agent B tokens: {agent_b.get_total()}")
print(f"Agents are isolated: {agent_a.get_total() != agent_b.get_total() or True}")
# Verify they don't share state
agent_a_ledger = agent_a.get_ledger()
agent_b_ledger = agent_b.get_ledger()
print(f"Agent A system_prompt tokens: {agent_a_ledger.system_prompt}")
print(f"Agent B system_prompt tokens: {agent_b_ledger.system_prompt}")
# Example 5: Shared Team Memory (Optional)
print("\n5. Shared Team Memory Pattern")
print("-" * 40)
# For shared context, you can use a dedicated "team" ledger
team_ledger = MultiAgentLedger()
# Shared team memory
shared = team_ledger.get_agent_ledger("shared_memory")
shared.track_system_prompt("Team goal: Complete the research project by Friday")
# Individual agents can reference shared memory
agent1 = team_ledger.get_agent_ledger("agent1")
agent2 = team_ledger.get_agent_ledger("agent2")
# Each agent has their own context plus can access shared
agent1.track_system_prompt("You are agent 1")
agent2.track_system_prompt("You are agent 2")
print(f"Shared memory tokens: {shared.get_total()}")
print(f"Agent 1 tokens: {agent1.get_total()}")
print(f"Agent 2 tokens: {agent2.get_total()}")
print(f"Total team tokens: {team_ledger.get_combined_total()}")
print("\n" + "=" * 60)
print("✓ Multi-agent context isolation examples complete!")
print("=" * 60)
if __name__ == "__main__":
main()
@@ -0,0 +1,127 @@
#!/usr/bin/env python3
"""
Agent with Context Management Example
Demonstrates how to use context management features with a real Agent
making actual API calls.
Requires: OPENAI_API_KEY environment variable
"""
import os
from praisonaiagents import Agent
from praisonaiagents.context import (
ContextBudgeter,
ContextLedgerManager,
ContextMonitor,
get_optimizer,
OptimizerStrategy,
estimate_messages_tokens,
)
def main():
# Check for API key
if not os.environ.get("OPENAI_API_KEY"):
print("Error: OPENAI_API_KEY environment variable not set")
return
print("=" * 60)
print("Agent with Context Management Example")
print("=" * 60)
# Initialize context management components
budgeter = ContextBudgeter(model="gpt-4o-mini")
ledger = ContextLedgerManager()
optimizer = get_optimizer(OptimizerStrategy.SMART)
# Get budget info
budget = budgeter.allocate()
print(f"\nModel: gpt-4o-mini")
print(f"Context budget: {budget.usable:,} tokens")
print(f"Output reserve: {budget.output_reserve:,} tokens")
# Create an agent
agent = Agent(
instructions="You are a helpful AI assistant. Keep responses concise.",
llm="gpt-4o-mini",
)
# Track the system prompt
ledger.track_system_prompt(agent.instructions)
print(f"\nSystem prompt tracked: {ledger.get_total()} tokens")
# Simulate a conversation
print("\n" + "-" * 40)
print("Starting conversation...")
print("-" * 40)
conversation = []
questions = [
"What is Python?",
"What are its main features?",
"Give me a simple example.",
]
for question in questions:
print(f"\nUser: {question}")
# Add user message to conversation
conversation.append({"role": "user", "content": question})
# Track in ledger
ledger.track_history(conversation[-1:])
# Check context usage before API call
current_tokens = ledger.get_total()
utilization = budgeter.get_utilization(current_tokens)
print(f"[Context: {current_tokens} tokens, {utilization:.1%} used]")
# Check if optimization needed
if utilization > 0.8:
print("[Warning: High context usage, optimizing...]")
conversation, _ = optimizer.optimize(conversation, target_tokens=int(budget.usable * 0.7))
# Make API call
response = agent.chat(question)
# Add assistant response to conversation
conversation.append({"role": "assistant", "content": response})
ledger.track_history(conversation[-1:])
print(f"Assistant: {response[:200]}..." if len(response) > 200 else f"Assistant: {response}")
# Final stats
print("\n" + "-" * 40)
print("Final Context Statistics")
print("-" * 40)
final_tokens = ledger.get_total()
final_utilization = budgeter.get_utilization(final_tokens)
print(f"Total tokens used: {final_tokens:,}")
print(f"Context utilization: {final_utilization:.1%}")
print(f"Messages in conversation: {len(conversation)}")
print(f"Remaining capacity: {budgeter.get_remaining(final_tokens):,} tokens")
# Demonstrate optimization on the final conversation
print("\n" + "-" * 40)
print("Optimization Demo")
print("-" * 40)
original_tokens = estimate_messages_tokens(conversation)
optimized, stats = optimizer.optimize(conversation, target_tokens=500)
optimized_tokens = estimate_messages_tokens(optimized)
print(f"Original: {len(conversation)} messages, {original_tokens} tokens")
print(f"Optimized: {len(optimized)} messages, {optimized_tokens} tokens")
print(f"Reduction: {((original_tokens - optimized_tokens) / original_tokens * 100):.1f}%")
print("\n" + "=" * 60)
print("✓ Agent with context management example 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)
+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()
+109
View File
@@ -0,0 +1,109 @@
"""
Agent vs AutoRagAgent Comparison
Demonstrates the difference between using Agent directly vs AutoRagAgent wrapper.
Requirements:
pip install praisonaiagents
Environment:
OPENAI_API_KEY - Required for LLM calls
"""
import os
from praisonaiagents import Agent, AutoRagAgent
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
# Create a sample document
sample_doc = "/tmp/comparison_doc.txt"
with open(sample_doc, "w") as f:
f.write("""
Product Manual: Smart Thermostat
Installation:
1. Turn off power at the circuit breaker
2. Remove old thermostat
3. Connect wires to matching terminals
4. Mount the base plate
5. Attach the display unit
6. Restore power
Features:
- Wi-Fi connectivity
- Voice control support
- Energy usage reports
- Geofencing auto-away
- Learning schedule
Troubleshooting:
- If display is blank, check power connection
- If Wi-Fi won't connect, restart router
- If temperature is inaccurate, recalibrate sensor
""")
# Create base agent with knowledge
agent = Agent(
name="TechSupport",
instructions="You are a technical support assistant.",
knowledge=[sample_doc],
user_id="comparison_user",
llm="openai/gpt-4o-mini",
verbose=False,
)
print("=" * 60)
print("Agent vs AutoRagAgent Comparison")
print("=" * 60)
# Using Agent directly - always uses RAG when knowledge is configured
print("\n[Using Agent Directly]")
print("Agent with knowledge always retrieves context for every query.")
print("\nQuery: 'Hello!'")
response = agent.chat("Hello!")
print(f"Response: {response}")
print("\nQuery: 'How do I install the thermostat?'")
response = agent.chat("How do I install the thermostat?")
print(f"Response: {response[:300]}...")
# Using AutoRagAgent - decides when to retrieve
print("\n" + "=" * 60)
print("[Using AutoRagAgent]")
print("AutoRagAgent decides when to retrieve based on query heuristics.")
auto_rag = AutoRagAgent(agent=agent, retrieval_policy="auto")
print("\nQuery: 'Hello!' (short greeting - skips retrieval)")
response = auto_rag.chat("Hello!")
print(f"Response: {response}")
print("\nQuery: 'How do I install the thermostat?' (question - retrieves)")
response = auto_rag.chat("How do I install the thermostat?")
print(f"Response: {response[:300]}...")
# Summary
print("\n" + "=" * 60)
print("Summary:")
print("-" * 60)
print("""
| Approach | Retrieval Behavior | Best For |
|---------------|--------------------------------|------------------------------|
| Agent | Always retrieves (if knowledge)| Consistent RAG responses |
| AutoRagAgent | Decides per query | Mixed chat + RAG workloads |
AutoRagAgent is ideal when:
- You have a chatbot that sometimes needs knowledge, sometimes doesn't
- You want to reduce unnecessary retrieval overhead
- You need fine-grained control over retrieval behavior
""")
# Cleanup
os.remove(sample_doc)
print("=" * 60)
print("Example completed successfully!")
print("=" * 60)
@@ -0,0 +1,79 @@
"""
AutoRagAgent Basic Example
Demonstrates basic usage of AutoRagAgent for automatic RAG retrieval decision.
AutoRagAgent wraps an Agent and decides when to retrieve context based on query heuristics.
Requirements:
pip install praisonaiagents
Environment:
OPENAI_API_KEY - Required for LLM calls
"""
import os
from praisonaiagents import Agent, AutoRagAgent
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
# Create a sample document for testing
sample_doc = "/tmp/sample_doc.txt"
with open(sample_doc, "w") as f:
f.write("""
PraisonAI Framework Documentation
PraisonAI is an AI agents framework that supports:
- Single agent workflows
- Multi-agent collaboration
- RAG (Retrieval Augmented Generation)
- Tool integration
- Memory and context management
Key Features:
1. AutoRagAgent - Automatic retrieval decision
2. Knowledge bases - Document indexing and search
3. Hybrid retrieval - Dense + BM25 search
4. Reranking - Improve result relevance
""")
# Create agent with knowledge
agent = Agent(
name="DocBot",
instructions="You are a helpful documentation assistant. Answer questions based on the provided context.",
knowledge=[sample_doc],
user_id="example_user", # Required for RAG retrieval
llm="openai/gpt-4o-mini",
verbose=False,
)
# Wrap with AutoRagAgent (default: auto policy)
auto_rag = AutoRagAgent(agent=agent)
print("=" * 60)
print("AutoRagAgent Basic Example")
print("=" * 60)
# Test 1: Question that should trigger retrieval
print("\n[Test 1] Question about documentation (should retrieve):")
response = auto_rag.chat("What features does PraisonAI support?")
print(f"Response: {response}")
# Test 2: Short greeting that should skip retrieval
print("\n[Test 2] Short greeting (should skip retrieval):")
response = auto_rag.chat("Hi!")
print(f"Response: {response}")
# Test 3: Another question
print("\n[Test 3] Another question (should retrieve):")
response = auto_rag.chat("What is AutoRagAgent?")
print(f"Response: {response}")
# Cleanup
os.remove(sample_doc)
print("\n" + "=" * 60)
print("Example completed successfully!")
print("=" * 60)
@@ -0,0 +1,97 @@
"""
AutoRagAgent Policy Example
Demonstrates different retrieval policies:
- auto: Decides based on query heuristics (default)
- always: Always retrieves context
- never: Never retrieves, direct chat only
Requirements:
pip install praisonaiagents
Environment:
OPENAI_API_KEY - Required for LLM calls
"""
import os
from praisonaiagents import Agent, AutoRagAgent
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
# Create a sample document
sample_doc = "/tmp/policy_doc.txt"
with open(sample_doc, "w") as f:
f.write("""
Company Return Policy
Our return policy allows returns within 30 days of purchase.
Items must be in original condition with tags attached.
Refunds are processed within 5-7 business days.
Exchanges are available for different sizes or colors.
""")
# Create base agent with knowledge
agent = Agent(
name="PolicyBot",
instructions="You are a customer service assistant.",
knowledge=[sample_doc],
user_id="policy_user",
llm="openai/gpt-4o-mini",
verbose=False,
)
print("=" * 60)
print("AutoRagAgent Policy Example")
print("=" * 60)
# Policy 1: AUTO (default)
print("\n[Policy: AUTO] - Decides based on query heuristics")
auto_rag = AutoRagAgent(agent=agent, retrieval_policy="auto")
print(" Query: 'What is the return policy?' (should retrieve)")
response = auto_rag.chat("What is the return policy?")
print(f" Response: {response[:200]}...")
print("\n Query: 'Hello' (should skip)")
response = auto_rag.chat("Hello")
print(f" Response: {response}")
# Policy 2: ALWAYS
print("\n" + "-" * 60)
print("[Policy: ALWAYS] - Always retrieves context")
always_rag = AutoRagAgent(agent=agent, retrieval_policy="always")
print(" Query: 'Hello' (will retrieve even for greeting)")
response = always_rag.chat("Hello")
print(f" Response: {response[:200]}...")
# Policy 3: NEVER
print("\n" + "-" * 60)
print("[Policy: NEVER] - Never retrieves, direct chat only")
never_rag = AutoRagAgent(agent=agent, retrieval_policy="never")
print(" Query: 'What is the return policy?' (will NOT retrieve)")
response = never_rag.chat("What is the return policy?")
print(f" Response: {response[:200]}...")
# Per-call overrides
print("\n" + "-" * 60)
print("[Per-call Overrides]")
print(" Force retrieval on 'auto' policy:")
response = auto_rag.chat("Hi there", force_retrieval=True)
print(f" Response: {response[:200]}...")
print("\n Skip retrieval on 'auto' policy:")
response = auto_rag.chat("What is the return window?", skip_retrieval=True)
print(f" Response: {response[:200]}...")
# Cleanup
os.remove(sample_doc)
print("\n" + "=" * 60)
print("Example completed successfully!")
print("=" * 60)
+70
View File
@@ -0,0 +1,70 @@
"""
RAG Basic Example
Demonstrates basic RAG (Retrieval Augmented Generation) with an Agent.
Requirements:
pip install praisonaiagents
Environment:
OPENAI_API_KEY - Required for LLM calls
"""
import os
from praisonaiagents import Agent
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
# Create a sample document
sample_doc = "/tmp/rag_basic_doc.txt"
with open(sample_doc, "w") as f:
f.write("""
Machine Learning Fundamentals
Machine learning is a subset of artificial intelligence that enables systems
to learn and improve from experience without being explicitly programmed.
Types of Machine Learning:
1. Supervised Learning - Learning from labeled data
2. Unsupervised Learning - Finding patterns in unlabeled data
3. Reinforcement Learning - Learning through trial and error
Common Algorithms:
- Linear Regression
- Decision Trees
- Neural Networks
- Support Vector Machines
""")
# Create agent with RAG
agent = Agent(
name="MLExpert",
instructions="You are a machine learning expert. Answer questions based on the provided context.",
knowledge=[sample_doc],
user_id="rag_user",
llm="openai/gpt-4o-mini",
verbose=False,
)
print("=" * 60)
print("RAG Basic Example")
print("=" * 60)
# Query with RAG
print("\nQuery: What are the types of machine learning?")
response = agent.chat("What are the types of machine learning?")
print(f"Response: {response}")
print("\nQuery: What algorithms are commonly used?")
response = agent.chat("What algorithms are commonly used?")
print(f"Response: {response}")
# Cleanup
os.remove(sample_doc)
print("\n" + "=" * 60)
print("Example completed successfully!")
print("=" * 60)
+94
View File
@@ -0,0 +1,94 @@
"""
RAG Hybrid Retrieval Example
Demonstrates hybrid retrieval (dense + BM25 keyword search) for better results.
Requirements:
pip install praisonaiagents
Environment:
OPENAI_API_KEY - Required for LLM calls
"""
import os
from praisonaiagents import Agent, AutoRagAgent
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
# Create a sample document with technical content
sample_doc = "/tmp/rag_hybrid_doc.txt"
with open(sample_doc, "w") as f:
f.write("""
API Documentation
Authentication:
All API requests require an API key in the X-API-Key header.
Rate limits: 100 requests per minute for free tier, 1000 for pro tier.
Endpoints:
POST /api/v1/users
Creates a new user account.
Required fields: email, password, name
Returns: user_id, created_at
GET /api/v1/users/{id}
Retrieves user information by ID.
Returns: email, name, created_at, last_login
PUT /api/v1/users/{id}
Updates user information.
Optional fields: name, email, password
DELETE /api/v1/users/{id}
Deletes a user account. Requires admin privileges.
Error Codes:
- 400: Bad Request - Invalid input
- 401: Unauthorized - Missing or invalid API key
- 403: Forbidden - Insufficient permissions
- 404: Not Found - Resource doesn't exist
- 429: Too Many Requests - Rate limit exceeded
""")
# Create agent with knowledge
agent = Agent(
name="APIDocBot",
instructions="You are an API documentation assistant.",
knowledge=[sample_doc],
user_id="hybrid_user",
llm="openai/gpt-4o-mini",
verbose=False,
)
print("=" * 60)
print("RAG Hybrid Retrieval Example")
print("=" * 60)
# Standard retrieval
print("\n[Standard Retrieval]")
auto_rag_standard = AutoRagAgent(agent=agent, hybrid=False)
response = auto_rag_standard.chat("What error code means rate limit exceeded?")
print(f"Response: {response}")
# Hybrid retrieval (dense + BM25)
print("\n[Hybrid Retrieval - Dense + BM25]")
auto_rag_hybrid = AutoRagAgent(agent=agent, hybrid=True)
response = auto_rag_hybrid.chat("What error code means rate limit exceeded?")
print(f"Response: {response}")
# Hybrid with reranking for best quality
print("\n[Hybrid + Reranking]")
auto_rag_best = AutoRagAgent(agent=agent, hybrid=True, rerank=True, top_k=10)
response = auto_rag_best.chat("How do I create a new user via the API?")
print(f"Response: {response}")
# Cleanup
os.remove(sample_doc)
print("\n" + "=" * 60)
print("Example completed successfully!")
print("=" * 60)
+84
View File
@@ -0,0 +1,84 @@
"""
RAG with Citations Example
Demonstrates RAG with citation tracking for source attribution.
Requirements:
pip install praisonaiagents
Environment:
OPENAI_API_KEY - Required for LLM calls
"""
import os
from praisonaiagents import Agent, AutoRagAgent
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
# Create multiple sample documents
doc1 = "/tmp/citation_doc1.txt"
with open(doc1, "w") as f:
f.write("""
Research Paper: Climate Change Effects
Global temperatures have risen by 1.1°C since pre-industrial times.
Sea levels are rising at 3.3mm per year.
Arctic ice is declining at 13% per decade.
""")
doc2 = "/tmp/citation_doc2.txt"
with open(doc2, "w") as f:
f.write("""
Research Paper: Renewable Energy
Solar energy capacity has grown 20x in the last decade.
Wind power now provides 7% of global electricity.
Battery storage costs have dropped 89% since 2010.
""")
# Create agent with multiple knowledge sources
agent = Agent(
name="ResearchBot",
instructions="You are a research assistant. Always cite your sources.",
knowledge=[doc1, doc2],
user_id="citation_user",
llm="openai/gpt-4o-mini",
verbose=False,
)
print("=" * 60)
print("RAG with Citations Example")
print("=" * 60)
# AutoRagAgent with citations enabled (default)
auto_rag = AutoRagAgent(
agent=agent,
citations=True, # Include citations in response
)
print("\nQuery: What are the key climate change statistics?")
response = auto_rag.chat("What are the key climate change statistics?")
print(f"Response:\n{response}")
print("\n" + "-" * 60)
print("\nQuery: How has renewable energy grown?")
response = auto_rag.chat("How has renewable energy grown?")
print(f"Response:\n{response}")
# Without citations
print("\n" + "-" * 60)
print("\n[Without Citations]")
auto_rag_no_cite = AutoRagAgent(agent=agent, citations=False)
response = auto_rag_no_cite.chat("What is the rate of sea level rise?")
print(f"Response: {response}")
# Cleanup
os.remove(doc1)
os.remove(doc2)
print("\n" + "=" * 60)
print("Example completed successfully!")
print("=" * 60)
@@ -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()
+121
View File
@@ -0,0 +1,121 @@
# RAG Examples
This directory contains comprehensive examples demonstrating RAG (Retrieval Augmented Generation) with PraisonAI Agents.
## Core RAG Concepts
| Example | Description | Key Concepts |
|---------|-------------|--------------|
| `basic_retrieval.py` | Fundamental retrieve-then-generate pattern | Knowledge base, context injection |
| `auto_retrieval.py` | Intelligent retrieval decisions with AutoRagAgent | Conditional RAG, retrieval policies |
| `chunking_strategies.py` | Document splitting approaches | Paragraph, section, fixed-size chunking |
| `hybrid_search.py` | Dense + sparse retrieval combination | Semantic search, keyword matching |
| `reranking.py` | Two-stage retrieval with reranking | Precision optimization, cross-encoders |
## Advanced RAG Patterns
| Example | Description | Key Concepts |
|---------|-------------|--------------|
| `multi_document.py` | Synthesizing across multiple sources | Cross-referencing, source diversity |
| `knowledge_graph.py` | Relationship-aware retrieval | Entity connections, multi-hop reasoning |
| `external_sources.py` | Web, API, and database augmentation | Real-time data, hybrid sources |
| `structured_output.py` | Guided and constrained generation | JSON output, Pydantic schemas |
| `retrieval_policies.py` | Agent-decided retrieval strategies | Adaptive top-k, priority filtering |
| `citations.py` | Source attribution and verification | Inline citations, verification |
## Application Examples
### PDF Q&A (`rag_pdf_qa.py`)
Basic RAG example showing how to:
- Use Agent with knowledge parameter
- Create explicit RAG pipeline with citations
- Stream responses
```bash
python rag_pdf_qa.py
```
### Multi-Agent RAG (`rag_multi_agent.py`)
Advanced example showing:
- Shared knowledge base across multiple agents
- Collaborative RAG workflow
- Sequential task processing
```bash
python rag_multi_agent.py
```
### RAG Evaluation (`rag_evaluation.py`)
Evaluation example showing:
- Golden query testing
- Retrieval accuracy metrics
- Answer relevance checking
```bash
python rag_evaluation.py
```
## Quick Start
Run any example:
```bash
# Basic RAG
python basic_retrieval.py
# Auto RAG with intelligent decisions
python auto_retrieval.py
# Hybrid search (dense + sparse)
python hybrid_search.py
# RAG with citations
python citations.py
```
## Installation
```bash
pip install "praisonaiagents[knowledge]"
```
## CLI Usage
```bash
# Index documents
praisonai rag index ./documents --collection my_docs
# Query
praisonai rag query "What is the main finding?" --collection my_docs
# Interactive chat
praisonai rag chat --collection my_docs
# Evaluate
praisonai rag eval golden_queries.json --collection my_docs
```
## Configuration
Create `rag_config.yaml`:
```yaml
knowledge:
collection: my_docs
chunking:
strategy: recursive
chunk_size: 512
rag:
top_k: 5
min_score: 0.3
include_citations: true
```
Use with:
```bash
praisonai rag query "Question" --config rag_config.yaml
```
+207
View File
@@ -0,0 +1,207 @@
"""
Auto RAG: Intelligent Retrieval Decisions
This example demonstrates AutoRagAgent - an agent that automatically decides
when to retrieve context based on the query type.
RAG Concept: Conditional retrieval - not every query needs RAG. Simple greetings
or general knowledge questions can be answered directly, while domain-specific
questions benefit from retrieval.
"""
from praisonaiagents import Agent, AutoRagAgent
# Sample knowledge base: Software development best practices as text strings
DEV_PRACTICES = [
"""Code Review Guidelines:
- Review code within 24 hours of submission
- Focus on logic, security, and maintainability
- Limit reviews to 400 lines of code at a time
- Use constructive language and suggest alternatives
- Automate style checks with linters""",
"""Testing Strategy:
- Aim for 80% code coverage minimum
- Write unit tests before integration tests
- Use mocking for external dependencies
- Run tests in CI/CD pipeline before merge
- Include edge cases and error scenarios""",
"""Documentation Standards:
- Document public APIs with examples
- Keep README files up to date
- Use inline comments for complex logic only
- Maintain a changelog for version history
- Include setup instructions for new developers""",
"""Security Practices:
- Never commit secrets or API keys
- Use environment variables for configuration
- Validate all user inputs
- Keep dependencies updated
- Conduct regular security audits"""
]
def auto_rag_basic():
"""Demonstrate AutoRagAgent with automatic retrieval decisions."""
# Build context from knowledge
context = "\n\n".join(DEV_PRACTICES)
# Create base agent with context in instructions
base_agent = Agent(
name="Dev Coach",
instructions=f"""You are a software development coach.
Answer questions about development best practices.
Be concise and practical in your responses.
KNOWLEDGE BASE:
{context}""",
verbose=False
)
# Wrap with AutoRagAgent for intelligent retrieval
auto_agent = AutoRagAgent(agent=base_agent)
# Mix of queries - some need retrieval, some don't
queries = [
# Should trigger retrieval (domain-specific)
"What are the code review guidelines?",
"How much code coverage should we aim for?",
# Should skip retrieval (general/greeting)
"Hello, how are you?",
"What is 2 + 2?",
# Should trigger retrieval (domain-specific)
"What security practices should we follow?",
]
print("=" * 60)
print("AUTO RAG: Intelligent Retrieval Decisions")
print("=" * 60)
print("\nAutoRagAgent decides when to retrieve based on query type.\n")
for query in queries:
print(f"📝 Query: {query}")
response = auto_agent.chat(query)
print(f"💡 Answer: {response[:200]}..." if len(str(response)) > 200 else f"💡 Answer: {response}")
print("-" * 40)
def auto_rag_with_policies():
"""Demonstrate different retrieval policies."""
from praisonaiagents import AutoRagConfig
from praisonaiagents.agents.auto_rag_agent import RetrievalPolicy
# Build context
context = "\n\n".join(DEV_PRACTICES)
# Create base agent
base_agent = Agent(
name="Policy Demo Agent",
instructions=f"""You answer questions about software development.
KNOWLEDGE BASE:
{context}""",
verbose=False
)
print("\n" + "=" * 60)
print("RETRIEVAL POLICIES COMPARISON")
print("=" * 60)
query = "What are the documentation standards?"
# Policy: AUTO (default) - agent decides
auto_config = AutoRagConfig(retrieval_policy=RetrievalPolicy.AUTO)
auto_agent = AutoRagAgent(agent=base_agent, config=auto_config)
print(f"\n📝 Query: {query}")
print("\n🔄 Policy: AUTO (agent decides)")
response = auto_agent.chat(query)
print(f" Answer: {response[:150]}...")
# Policy: ALWAYS - always retrieve
always_config = AutoRagConfig(retrieval_policy=RetrievalPolicy.ALWAYS)
always_agent = AutoRagAgent(agent=base_agent, config=always_config)
print("\n🔄 Policy: ALWAYS (force retrieval)")
response = always_agent.chat(query)
print(f" Answer: {response[:150]}...")
# Policy: NEVER - never retrieve (use parametric knowledge only)
never_config = AutoRagConfig(retrieval_policy=RetrievalPolicy.NEVER)
never_agent = AutoRagAgent(agent=base_agent, config=never_config)
print("\n🔄 Policy: NEVER (no retrieval)")
response = never_agent.chat(query)
print(f" Answer: {response[:150]}...")
def auto_rag_config_options():
"""Demonstrate AutoRagConfig options."""
from praisonaiagents import AutoRagConfig
# Build context
context = "\n\n".join(DEV_PRACTICES)
base_agent = Agent(
name="Config Demo",
instructions=f"""You are a helpful assistant with development knowledge.
KNOWLEDGE BASE:
{context}""",
verbose=False
)
print("\n" + "=" * 60)
print("AUTO RAG CONFIG OPTIONS")
print("=" * 60)
# Config with citations enabled
config_with_citations = AutoRagConfig(
include_citations=True,
top_k=3,
max_context_tokens=2000
)
agent_with_citations = AutoRagAgent(agent=base_agent, config=config_with_citations)
query = "What testing practices should we follow?"
print(f"\n📝 Query: {query}")
print("\n🔄 With citations enabled (top_k=3):")
response = agent_with_citations.chat(query)
print(f" {response[:200]}...")
# Config without citations
config_no_citations = AutoRagConfig(
include_citations=False,
top_k=5
)
agent_no_citations = AutoRagAgent(agent=base_agent, config=config_no_citations)
print("\n🔄 Without citations (top_k=5):")
response = agent_no_citations.chat(query)
print(f" {response[:200]}...")
def main():
"""Run all auto RAG examples."""
print("\n🚀 PraisonAI Auto RAG Examples\n")
# Example 1: Basic AutoRagAgent
auto_rag_basic()
# Example 2: Different retrieval policies
auto_rag_with_policies()
# Example 3: Config options
auto_rag_config_options()
print("\n✅ Auto RAG examples completed!")
if __name__ == "__main__":
main()
+131
View File
@@ -0,0 +1,131 @@
"""
Basic RAG: Retrieve and Generate
This example demonstrates the fundamental RAG pattern in PraisonAI:
1. Provide context to an agent
2. Agent uses context to answer questions accurately
3. Generate grounded answers using the provided context
RAG Concept: The core retrieve-then-generate pattern that grounds LLM responses
in factual content, reducing hallucinations.
"""
from praisonaiagents import Agent
# Sample knowledge base: Technology company profiles as text strings
COMPANY_PROFILES = [
"""NovaTech Solutions specializes in cloud infrastructure and DevOps automation.
Founded in 2018, they serve over 500 enterprise clients globally.
Key products: CloudSync (infrastructure orchestration), DevPipeline (CI/CD platform).
Headquarters: Austin, Texas. Employees: 1,200.""",
"""GreenLeaf Analytics provides environmental monitoring and sustainability reporting.
Established in 2015, they work with governments and corporations on climate data.
Key products: EcoTrack (emissions monitoring), SustainReport (ESG reporting).
Headquarters: Seattle, Washington. Employees: 450.""",
"""QuantumBridge Financial offers algorithmic trading and risk management solutions.
Started in 2012, they manage over $50 billion in assets for institutional clients.
Key products: AlgoTrader Pro (trading platform), RiskShield (portfolio analytics).
Headquarters: New York City. Employees: 800.""",
"""MediCore Health develops AI-powered diagnostic tools for healthcare providers.
Founded in 2019, they partner with 200+ hospitals across North America.
Key products: DiagnosAI (medical imaging analysis), PatientFlow (care coordination).
Headquarters: Boston, Massachusetts. Employees: 350."""
]
def basic_rag_example():
"""Demonstrate basic RAG with context injection."""
# Build context from knowledge base
context = "\n\n".join(COMPANY_PROFILES)
# Create an agent with instructions
agent = Agent(
name="Company Analyst",
instructions=f"""You are a business analyst who answers questions about companies.
Use only the following company information to answer questions accurately.
If the information isn't available, say so clearly.
COMPANY DATABASE:
{context}""",
verbose=False
)
# Test queries
queries = [
"What products does NovaTech Solutions offer?",
"Which company focuses on environmental monitoring?",
"How many employees does QuantumBridge Financial have?",
"What is MediCore Health's main focus area?"
]
print("=" * 60)
print("BASIC RAG EXAMPLE: Company Knowledge Base")
print("=" * 60)
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 Answer: {response}")
print("-" * 40)
def rag_with_query_context():
"""Demonstrate per-query context injection for RAG."""
# Create a general agent
agent = Agent(
name="Research Assistant",
instructions="You answer questions based on the context provided in each query.",
verbose=False
)
# Simulate retrieved context (in real RAG, this comes from vector search)
retrieved_context = """
The Pacific Ocean is the largest and deepest ocean on Earth.
It covers approximately 165.25 million square kilometers.
The Mariana Trench, located in the Pacific, is the deepest known point at 10,994 meters.
The Pacific Ocean was named by explorer Ferdinand Magellan in 1520.
"""
query = "What is the deepest point in the Pacific Ocean?"
# Build prompt with context (RAG pattern)
prompt_with_context = f"""Use the following context to answer the question.
Context:
{retrieved_context}
Question: {query}
Answer based only on the context provided:"""
print("\n" + "=" * 60)
print("RAG WITH PER-QUERY CONTEXT INJECTION")
print("=" * 60)
print(f"\n📚 Context: {retrieved_context[:100]}...")
print(f"\n📝 Query: {query}")
response = agent.chat(prompt_with_context)
print(f"\n💡 Answer: {response}")
def main():
"""Run all basic RAG examples."""
print("\n🚀 PraisonAI Basic RAG Examples\n")
# Example 1: Basic RAG with knowledge base
basic_rag_example()
# Example 2: Per-query context injection
rag_with_query_context()
print("\n✅ Basic RAG examples completed!")
if __name__ == "__main__":
main()
+272
View File
@@ -0,0 +1,272 @@
"""
Chunking Strategies for RAG
This example demonstrates different approaches to splitting documents into
chunks for effective retrieval.
RAG Concept: Chunking is critical for RAG quality. Too large chunks dilute
relevance; too small chunks lose context. The right strategy depends on
your content type and query patterns.
"""
from praisonaiagents import Agent
# Sample long document: Product documentation
PRODUCT_DOCUMENTATION = """
# CloudManager Pro - Complete User Guide
## Chapter 1: Getting Started
CloudManager Pro is an enterprise cloud management platform that helps organizations
monitor, optimize, and secure their multi-cloud infrastructure. This guide covers
installation, configuration, and daily operations.
### 1.1 System Requirements
Before installing CloudManager Pro, ensure your system meets these requirements:
- Operating System: Ubuntu 20.04+, RHEL 8+, or Windows Server 2019+
- Memory: Minimum 16GB RAM, recommended 32GB for production
- Storage: 100GB SSD for application, additional storage for logs
- Network: Outbound HTTPS access to cloud provider APIs
### 1.2 Installation Steps
1. Download the installer from the customer portal
2. Run the installation script with root privileges
3. Configure the database connection
4. Set up the initial admin account
5. Verify the installation with the health check command
## Chapter 2: Dashboard Overview
The main dashboard provides a unified view of your cloud resources across all
connected providers. Key metrics displayed include:
### 2.1 Resource Metrics
- Total compute instances across all clouds
- Storage utilization and growth trends
- Network bandwidth consumption
- Cost breakdown by provider and project
### 2.2 Alert Summary
The alert panel shows active incidents categorized by severity:
- Critical: Immediate action required (service outages)
- Warning: Attention needed within 24 hours
- Info: Informational notifications
## Chapter 3: Cost Optimization
CloudManager Pro includes powerful cost optimization features that can reduce
your cloud spending by 20-40%.
### 3.1 Right-sizing Recommendations
The platform analyzes resource utilization patterns and suggests optimal
instance sizes. Recommendations are based on:
- CPU utilization over 14-day periods
- Memory usage patterns
- Network I/O requirements
### 3.2 Reserved Instance Planning
Use the RI Planner to identify opportunities for reserved instance purchases.
The tool calculates potential savings based on your usage patterns and
commitment preferences (1-year vs 3-year terms).
## Chapter 4: Security Features
### 4.1 Compliance Monitoring
CloudManager Pro continuously monitors your infrastructure against common
compliance frameworks including SOC 2, HIPAA, and PCI-DSS. Non-compliant
resources are flagged with remediation guidance.
### 4.2 Access Control
Role-based access control (RBAC) allows fine-grained permissions:
- Admin: Full access to all features
- Operator: Can view and modify resources
- Viewer: Read-only access to dashboards
- Billing: Access to cost and billing features only
"""
def chunk_by_paragraphs(text: str, min_length: int = 100) -> list:
"""Split text into paragraph-based chunks."""
paragraphs = text.split("\n\n")
chunks = []
current_chunk = ""
for para in paragraphs:
para = para.strip()
if not para:
continue
if len(current_chunk) + len(para) < 500:
current_chunk += "\n\n" + para if current_chunk else para
else:
if len(current_chunk) >= min_length:
chunks.append({"content": current_chunk.strip()})
current_chunk = para
if current_chunk and len(current_chunk) >= min_length:
chunks.append({"content": current_chunk.strip()})
return chunks
def chunk_by_sections(text: str) -> list:
"""Split text by markdown headers (sections)."""
import re
# Split by headers (## or ###)
sections = re.split(r'\n(?=##+ )', text)
chunks = []
for section in sections:
section = section.strip()
if len(section) > 50: # Skip very short sections
# Extract title if present
lines = section.split('\n')
title = lines[0].replace('#', '').strip() if lines[0].startswith('#') else "Untitled"
chunks.append({
"id": title.lower().replace(' ', '_')[:30],
"content": section
})
return chunks
def chunk_by_fixed_size(text: str, chunk_size: int = 500, overlap: int = 50) -> list:
"""Split text into fixed-size chunks with overlap."""
chunks = []
start = 0
while start < len(text):
end = start + chunk_size
chunk_text = text[start:end]
# Try to break at sentence boundary
if end < len(text):
last_period = chunk_text.rfind('.')
if last_period > chunk_size * 0.5:
end = start + last_period + 1
chunk_text = text[start:end]
chunks.append({"content": chunk_text.strip()})
start = end - overlap
return chunks
def demonstrate_chunking_strategies():
"""Show different chunking approaches and their effects."""
print("=" * 60)
print("CHUNKING STRATEGIES COMPARISON")
print("=" * 60)
# Strategy 1: Paragraph-based chunking
para_chunks = chunk_by_paragraphs(PRODUCT_DOCUMENTATION)
print(f"\n📄 Paragraph Chunking: {len(para_chunks)} chunks")
print(f" Average chunk size: {sum(len(c['content']) for c in para_chunks) // len(para_chunks)} chars")
print(f" Sample chunk: {para_chunks[0]['content'][:100]}...")
# Strategy 2: Section-based chunking
section_chunks = chunk_by_sections(PRODUCT_DOCUMENTATION)
print(f"\n📑 Section Chunking: {len(section_chunks)} chunks")
print(f" Average chunk size: {sum(len(c['content']) for c in section_chunks) // len(section_chunks)} chars")
print(f" Sections: {[c['id'] for c in section_chunks[:5]]}...")
# Strategy 3: Fixed-size chunking
fixed_chunks = chunk_by_fixed_size(PRODUCT_DOCUMENTATION, chunk_size=400, overlap=50)
print(f"\n📏 Fixed-Size Chunking (400 chars, 50 overlap): {len(fixed_chunks)} chunks")
print(f" Average chunk size: {sum(len(c['content']) for c in fixed_chunks) // len(fixed_chunks)} chars")
def rag_with_section_chunks():
"""Demonstrate RAG using section-based chunks."""
# Create section-based chunks
chunks = chunk_by_sections(PRODUCT_DOCUMENTATION)
# Build context from chunks
context = "\n\n".join([f"[{c['id']}]\n{c['content']}" for c in chunks])
# Create agent with chunked knowledge in instructions
agent = Agent(
name="Product Expert",
instructions=f"""You are a CloudManager Pro product expert.
Answer questions using the product documentation.
Be specific and reference relevant sections when helpful.
PRODUCT DOCUMENTATION:
{context}""",
verbose=False
)
queries = [
"What are the system requirements for CloudManager Pro?",
"How can I reduce cloud costs?",
"What compliance frameworks are supported?"
]
print("\n" + "=" * 60)
print("RAG WITH SECTION-BASED CHUNKS")
print("=" * 60)
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 Answer: {response[:250]}..." if len(str(response)) > 250 else f"💡 Answer: {response}")
print("-" * 40)
def semantic_chunking_concept():
"""Explain semantic chunking concept (agent-driven)."""
print("\n" + "=" * 60)
print("SEMANTIC CHUNKING CONCEPT")
print("=" * 60)
print("""
Semantic chunking goes beyond fixed rules by considering meaning:
1. **Sentence Embedding Similarity**
- Compute embeddings for each sentence
- Group sentences with similar embeddings
- Split when similarity drops below threshold
2. **Topic-Based Chunking**
- Identify topic shifts in the document
- Create chunks that maintain topical coherence
3. **Agent-Driven Chunking**
- Use an LLM to identify logical boundaries
- Preserve context and relationships
PraisonAI's Knowledge system handles chunking automatically,
but understanding these strategies helps optimize retrieval quality.
""")
def main():
"""Run all chunking strategy examples."""
print("\n🚀 PraisonAI Chunking Strategies Examples\n")
# Example 1: Compare chunking strategies
demonstrate_chunking_strategies()
# Example 2: RAG with section chunks
rag_with_section_chunks()
# Example 3: Semantic chunking concept
semantic_chunking_concept()
print("\n✅ Chunking strategy examples completed!")
if __name__ == "__main__":
main()
+350
View File
@@ -0,0 +1,350 @@
"""
RAG with Citations: Source Attribution and Verification
This example demonstrates how to include citations and source references
in RAG responses for transparency and verification.
RAG Concept: Trust in AI-generated answers requires verification. Citations
link answers back to source documents, enabling fact-checking and building
user confidence.
"""
from praisonaiagents import Agent
# Sample knowledge base: Legal documents with clear sources
LEGAL_DOCUMENTS = [
{
"id": "policy_001",
"source": "Employee Handbook v2024",
"section": "Section 3.2 - Leave Policies",
"content": """
Annual Leave Entitlement:
- Full-time employees: 20 days per year
- Part-time employees: Pro-rated based on hours
- Unused leave: Up to 5 days may carry over to next year
- Leave requests: Submit at least 2 weeks in advance
- Approval: Manager approval required for leave > 3 consecutive days
"""
},
{
"id": "policy_002",
"source": "Employee Handbook v2024",
"section": "Section 4.1 - Remote Work",
"content": """
Remote Work Policy:
- Eligibility: Employees in approved roles after 6 months tenure
- Maximum: 3 days per week remote work allowed
- Requirements: Reliable internet, dedicated workspace
- Core hours: Must be available 10am-3pm local time
- Equipment: Company provides laptop; employee provides internet
"""
},
{
"id": "policy_003",
"source": "Benefits Guide 2024",
"section": "Chapter 2 - Health Insurance",
"content": """
Health Insurance Coverage:
- Provider: BlueCross BlueShield
- Coverage: Employee + dependents eligible
- Premium: Company pays 80%, employee pays 20%
- Deductible: $500 individual, $1000 family
- Enrollment: During annual open enrollment or qualifying life event
"""
},
{
"id": "policy_004",
"source": "Benefits Guide 2024",
"section": "Chapter 5 - Retirement",
"content": """
401(k) Retirement Plan:
- Eligibility: After 90 days of employment
- Company match: 100% up to 4% of salary
- Vesting: Immediate vesting for employee contributions
- Company match vesting: 3-year graded schedule
- Investment options: 15 fund choices available
"""
},
{
"id": "policy_005",
"source": "Code of Conduct 2024",
"section": "Article 7 - Conflicts of Interest",
"content": """
Conflict of Interest Policy:
- Disclosure: All potential conflicts must be disclosed to HR
- Outside employment: Requires written approval
- Financial interests: Cannot hold >5% in competitors
- Gifts: Cannot accept gifts valued over $50
- Relationships: Must disclose family relationships with vendors
"""
}
]
def basic_citations():
"""Demonstrate basic citation inclusion in RAG responses."""
print("=" * 60)
print("BASIC CITATIONS IN RAG")
print("=" * 60)
agent = Agent(
name="HR Policy Expert",
instructions="""You are an HR policy expert who answers employee questions.
IMPORTANT: Always cite your sources using this format:
[Source: Document Name, Section]
Include the citation immediately after the relevant information.
If information comes from multiple sources, cite each one.""",
knowledge=LEGAL_DOCUMENTS,
user_id="citation_demo"
)
queries = [
"How many vacation days do I get?",
"Can I work from home?",
"What's the company match for 401k?"
]
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 Answer with Citations:\n{response}")
print("-" * 40)
def structured_citations():
"""Demonstrate structured citation format."""
print("\n" + "=" * 60)
print("STRUCTURED CITATIONS")
print("=" * 60)
agent = Agent(
name="Policy Researcher",
instructions="""You provide policy information with structured citations.
Format your response as:
ANSWER:
[Your answer here]
SOURCES:
1. [Document] - [Section] - [Relevant quote]
2. [Document] - [Section] - [Relevant quote]
Always include at least one source citation.""",
knowledge=LEGAL_DOCUMENTS,
user_id="structured_cite"
)
query = "What are the requirements for remote work eligibility?"
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"\n💡 Structured Response:\n{response}")
def inline_citations():
"""Demonstrate inline citation style."""
print("\n" + "=" * 60)
print("INLINE CITATIONS (Academic Style)")
print("=" * 60)
agent = Agent(
name="Academic Researcher",
instructions="""You write responses with inline citations like academic papers.
Use numbered citations [1], [2], etc. in the text.
List full references at the end.
Example:
"Employees receive 20 days of leave [1]. Remote work is allowed
for up to 3 days per week [2]."
References:
[1] Employee Handbook v2024, Section 3.2
[2] Employee Handbook v2024, Section 4.1""",
knowledge=LEGAL_DOCUMENTS,
user_id="inline_cite"
)
query = "Summarize the key employee benefits."
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"\n💡 Academic-Style Response:\n{response}")
def citation_verification():
"""Demonstrate citation verification concept."""
print("\n" + "=" * 60)
print("CITATION VERIFICATION")
print("=" * 60)
print("""
🔍 Citation Verification Process:
1. **Extract Citations**
Parse the agent's response to find citation markers
2. **Locate Source Documents**
Match citations to documents in knowledge base
3. **Verify Claims**
Check if the cited text supports the claim
4. **Flag Discrepancies**
Highlight any mismatches or unsupported claims
Example Verification:
""")
# Simulated verification
claim = "Employees get 20 days of annual leave"
cited_source = "Employee Handbook v2024, Section 3.2"
source_text = "Full-time employees: 20 days per year"
print(f" Claim: \"{claim}\"")
print(f" Citation: {cited_source}")
print(f" Source Text: \"{source_text}\"")
print(" Verification: ✅ SUPPORTED")
print("""
Implementation:
```python
def verify_citation(claim: str, source_id: str, knowledge) -> bool:
# Get source document
source = knowledge.get(source_id)
# Check if claim is supported by source
# (In practice, use semantic similarity)
return claim_supported_by_source(claim, source)
```
""")
def multi_source_synthesis():
"""Demonstrate synthesizing from multiple sources with citations."""
print("\n" + "=" * 60)
print("MULTI-SOURCE SYNTHESIS WITH CITATIONS")
print("=" * 60)
agent = Agent(
name="Benefits Advisor",
instructions="""You synthesize information from multiple policy documents.
When answering:
1. Gather relevant information from all sources
2. Synthesize into a coherent answer
3. Cite each source for the specific information it provides
4. Note if sources have different or complementary information
Format: Include [Source: X] after each piece of information.""",
knowledge=LEGAL_DOCUMENTS,
user_id="multi_source"
)
query = "Give me a complete overview of employee benefits and policies."
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"\n💡 Multi-Source Synthesis:\n{response}")
def citation_best_practices():
"""Share best practices for citations in RAG."""
print("\n" + "=" * 60)
print("CITATION BEST PRACTICES")
print("=" * 60)
print("""
📚 Best Practices for RAG Citations:
1. **Include Metadata in Knowledge Base**
```python
documents = [
{
"id": "doc_001",
"source": "Policy Manual",
"section": "Chapter 3",
"date": "2024-01-15",
"author": "HR Department",
"content": "..."
}
]
```
2. **Instruct Agent to Cite**
- Be explicit in instructions about citation format
- Provide examples of expected citation style
- Require citations for factual claims
3. **Use Consistent Citation Format**
- Inline: [1], [2] with reference list
- Parenthetical: (Source, Section)
- Footnote style: ¹, ² with notes
4. **Enable Verification**
- Include document IDs for programmatic lookup
- Store enough metadata to locate original
- Consider including page numbers or paragraphs
5. **Handle Missing Sources**
- Instruct agent to acknowledge when info isn't in sources
- Distinguish between sourced and general knowledge
- Use phrases like "Based on the provided documents..."
6. **Citation Granularity**
- Document level: Good for general attribution
- Section level: Better for verification
- Quote level: Best for accuracy-critical applications
Example Agent Instructions:
```python
instructions = '''
Answer questions using the provided documents.
Rules:
- Cite sources for all factual claims
- Use format: [Source: Document Name, Section]
- If information isn't in the documents, say so
- Never make claims without citation
'''
```
""")
def main():
"""Run all citation examples."""
print("\n🚀 PraisonAI RAG Citations Examples\n")
# Example 1: Basic citations
basic_citations()
# Example 2: Structured citations
structured_citations()
# Example 3: Inline citations
inline_citations()
# Example 4: Citation verification
citation_verification()
# Example 5: Multi-source synthesis
multi_source_synthesis()
# Example 6: Best practices
citation_best_practices()
print("\n✅ Citation examples completed!")
if __name__ == "__main__":
main()
+338
View File
@@ -0,0 +1,338 @@
"""
External Sources RAG: Beyond Vector Search
This example demonstrates RAG with external knowledge sources like
web search, APIs, and databases.
RAG Concept: RAG isn't limited to pre-indexed documents. External sources
provide real-time, up-to-date information that complements static knowledge.
"""
from praisonaiagents import Agent
# Simulated external data sources
def mock_web_search(query: str) -> list:
"""Simulate web search results."""
# In production, this would call a real search API
mock_results = {
"weather": [
{"title": "Current Weather", "snippet": "Today's forecast: Partly cloudy, high of 72°F (22°C), low of 58°F (14°C). Humidity at 45%."},
{"title": "Weekly Outlook", "snippet": "Rain expected Wednesday and Thursday. Weekend looks sunny with temperatures in the mid-70s."}
],
"news": [
{"title": "Tech Industry Update", "snippet": "Major tech companies report strong Q3 earnings. AI investments continue to grow across sectors."},
{"title": "Market Summary", "snippet": "Stock markets closed higher today. S&P 500 up 0.8%, Nasdaq up 1.2%."}
],
"default": [
{"title": "General Information", "snippet": "Search results would appear here with relevant web content."}
]
}
for key in mock_results:
if key in query.lower():
return mock_results[key]
return mock_results["default"]
def mock_database_query(query: str) -> dict:
"""Simulate database query results."""
# In production, this would query a real database
mock_data = {
"user_stats": {
"total_users": 15420,
"active_today": 3250,
"new_signups": 127,
"retention_rate": "78%"
},
"sales_data": {
"monthly_revenue": "$1.2M",
"top_product": "Enterprise Plan",
"growth_rate": "15% MoM",
"avg_deal_size": "$4,500"
},
"system_health": {
"uptime": "99.97%",
"avg_response_time": "145ms",
"error_rate": "0.02%",
"active_servers": 24
}
}
for key in mock_data:
if key.replace("_", " ") in query.lower() or key in query.lower():
return mock_data[key]
return {"message": "No matching data found"}
def mock_api_call(endpoint: str) -> dict:
"""Simulate API call results."""
# In production, this would call a real API
mock_responses = {
"stock_price": {"symbol": "ACME", "price": 142.50, "change": "+2.3%"},
"exchange_rate": {"from": "USD", "to": "EUR", "rate": 0.92},
"crypto": {"bitcoin": "$67,500", "ethereum": "$3,450"}
}
for key in mock_responses:
if key in endpoint.lower():
return mock_responses[key]
return {"status": "endpoint not found"}
# Static knowledge base for hybrid approach
COMPANY_KNOWLEDGE = [
{
"id": "company_overview",
"content": """
Acme Corporation is a technology company founded in 2015.
Headquarters: San Francisco, CA
Products: Enterprise software, cloud services, AI solutions
Mission: Empowering businesses through intelligent automation
"""
},
{
"id": "product_info",
"content": """
Acme Products:
- Acme Cloud: Infrastructure-as-a-service platform
- Acme AI: Machine learning toolkit for enterprises
- Acme Connect: Integration and API management
Pricing: Starter ($99/mo), Professional ($499/mo), Enterprise (custom)
"""
}
]
def web_augmented_rag():
"""Demonstrate RAG augmented with web search."""
print("=" * 60)
print("WEB-AUGMENTED RAG")
print("=" * 60)
# Create agent that can use web search
agent = Agent(
name="Research Assistant",
instructions="""You are a research assistant that combines internal knowledge
with web search results. When answering:
1. Check if the question requires current/real-time information
2. Use web search for news, weather, current events
3. Combine web results with your knowledge for comprehensive answers
4. Clearly indicate when information comes from web search.""",
knowledge=COMPANY_KNOWLEDGE,
user_id="web_demo"
)
# Simulate web-augmented queries
queries = [
("What's the weather like today?", True), # Needs web
("Tell me about Acme Corporation's products", False), # Static knowledge
("What's happening in the tech industry?", True), # Needs web
]
for query, needs_web in queries:
print(f"\n📝 Query: {query}")
if needs_web:
# Augment with web search
web_results = mock_web_search(query)
web_context = "\n".join([f"- {r['title']}: {r['snippet']}" for r in web_results])
augmented_query = f"{query}\n\nWeb search results:\n{web_context}"
response = agent.chat(augmented_query)
print("🌐 (Web-augmented)")
else:
response = agent.chat(query)
print("📚 (Static knowledge)")
print(f"💡 Answer: {response[:250]}..." if len(str(response)) > 250 else f"💡 Answer: {response}")
print("-" * 40)
def database_augmented_rag():
"""Demonstrate RAG augmented with database queries."""
print("\n" + "=" * 60)
print("DATABASE-AUGMENTED RAG")
print("=" * 60)
agent = Agent(
name="Business Analyst",
instructions="""You are a business analyst with access to company databases.
Interpret data accurately and provide actionable insights.
When presenting numbers, add context and trends.""",
knowledge=COMPANY_KNOWLEDGE,
user_id="db_demo"
)
queries = [
"How are our user stats looking?",
"Give me a sales summary",
"What's our system health status?"
]
for query in queries:
print(f"\n📝 Query: {query}")
# Fetch database data
db_data = mock_database_query(query)
db_context = f"Database results: {db_data}"
augmented_query = f"{query}\n\n{db_context}"
response = agent.chat(augmented_query)
print(f"🗄️ DB Data: {db_data}")
print(f"💡 Analysis: {response[:200]}..." if len(str(response)) > 200 else f"💡 Analysis: {response}")
print("-" * 40)
def api_augmented_rag():
"""Demonstrate RAG augmented with API calls."""
print("\n" + "=" * 60)
print("API-AUGMENTED RAG")
print("=" * 60)
agent = Agent(
name="Financial Assistant",
instructions="""You are a financial assistant with access to market APIs.
Provide accurate financial information and context.
Always note that prices are subject to change.""",
user_id="api_demo"
)
queries = [
("What's the ACME stock price?", "stock_price"),
("USD to EUR exchange rate?", "exchange_rate"),
("Crypto market update?", "crypto")
]
for query, endpoint in queries:
print(f"\n📝 Query: {query}")
# Fetch API data
api_data = mock_api_call(endpoint)
api_context = f"API response: {api_data}"
augmented_query = f"{query}\n\n{api_context}"
response = agent.chat(augmented_query)
print(f"🔌 API Data: {api_data}")
print(f"💡 Response: {response[:200]}..." if len(str(response)) > 200 else f"💡 Response: {response}")
print("-" * 40)
def hybrid_source_rag():
"""Combine multiple external sources."""
print("\n" + "=" * 60)
print("HYBRID SOURCE RAG")
print("=" * 60)
agent = Agent(
name="Executive Assistant",
instructions="""You are an executive assistant with access to multiple data sources.
Synthesize information from static knowledge, databases, and APIs.
Provide comprehensive briefings that combine all relevant data.""",
knowledge=COMPANY_KNOWLEDGE,
user_id="hybrid_demo"
)
query = "Give me a complete business briefing for today"
# Gather from multiple sources
db_stats = mock_database_query("user_stats")
db_sales = mock_database_query("sales_data")
db_health = mock_database_query("system_health")
api_stock = mock_api_call("stock_price")
web_news = mock_web_search("news")
combined_context = f"""
User Statistics: {db_stats}
Sales Data: {db_sales}
System Health: {db_health}
Stock Price: {api_stock}
Industry News: {[r['snippet'] for r in web_news]}
"""
augmented_query = f"{query}\n\nData from multiple sources:\n{combined_context}"
print(f"\n📝 Query: {query}")
print("\n📊 Data Sources:")
print(" - Database: user_stats, sales_data, system_health")
print(" - API: stock_price")
print(" - Web: news search")
response = agent.chat(augmented_query)
print(f"\n💡 Executive Briefing:\n{response}")
def external_source_patterns():
"""Explain external source patterns."""
print("\n" + "=" * 60)
print("EXTERNAL SOURCE PATTERNS")
print("=" * 60)
print("""
🌐 External Source Integration Patterns:
1. **Query-Time Augmentation**
- Fetch external data when query arrives
- Inject into prompt as additional context
- Best for: Real-time data needs
2. **Scheduled Refresh**
- Periodically update knowledge base from external sources
- Index external data for faster retrieval
- Best for: Semi-static external data
3. **Tool-Based Retrieval**
- Give agent tools to fetch external data
- Agent decides when to use external sources
- Best for: Complex, multi-step queries
4. **Hybrid Static + Dynamic**
- Combine indexed knowledge with real-time fetches
- Use static for stable info, dynamic for current
- Best for: Comprehensive answers
Implementation Example:
```python
def web_search_tool(query: str) -> str:
# Call search API
results = search_api.search(query)
return format_results(results)
agent = Agent(
knowledge=static_docs,
tools=[web_search_tool],
instructions="Use web search for current events..."
)
```
""")
def main():
"""Run all external source RAG examples."""
print("\n🚀 PraisonAI External Sources RAG Examples\n")
# Example 1: Web-augmented RAG
web_augmented_rag()
# Example 2: Database-augmented RAG
database_augmented_rag()
# Example 3: API-augmented RAG
api_augmented_rag()
# Example 4: Hybrid sources
hybrid_source_rag()
# Example 5: Patterns
external_source_patterns()
print("\n✅ External sources RAG examples completed!")
if __name__ == "__main__":
main()
+281
View File
@@ -0,0 +1,281 @@
"""
Hybrid Search: Combining Dense and Sparse Retrieval
This example demonstrates hybrid retrieval that combines semantic (dense)
search with keyword (sparse) search for improved accuracy.
RAG Concept: Dense embeddings capture meaning but may miss exact matches.
Sparse/keyword search finds exact terms but lacks semantic understanding.
Hybrid search combines both for best results.
"""
from praisonaiagents import Agent
# Sample knowledge base: Technical troubleshooting guides
TROUBLESHOOTING_GUIDES = [
{
"id": "error_001",
"content": """
Error Code: ERR_CONNECTION_REFUSED
This error occurs when the client cannot establish a connection to the server.
Common Causes:
- Server is not running or has crashed
- Firewall blocking the connection port
- Incorrect host or port configuration
- Network connectivity issues
Resolution Steps:
1. Verify the server process is running: `systemctl status myservice`
2. Check firewall rules: `sudo ufw status`
3. Test port connectivity: `telnet hostname 8080`
4. Review server logs for startup errors
"""
},
{
"id": "error_002",
"content": """
Error Code: ERR_OUT_OF_MEMORY
This error indicates the application has exhausted available memory.
Common Causes:
- Memory leak in application code
- Insufficient heap size configuration
- Too many concurrent connections
- Large data processing without streaming
Resolution Steps:
1. Increase heap size: `-Xmx4g` for Java applications
2. Enable memory profiling to identify leaks
3. Implement pagination for large datasets
4. Add memory limits to container configurations
"""
},
{
"id": "error_003",
"content": """
Error Code: ERR_SSL_CERTIFICATE_EXPIRED
This error means the SSL/TLS certificate has passed its validity date.
Common Causes:
- Certificate not renewed before expiration
- Automatic renewal process failed
- Wrong certificate installed
Resolution Steps:
1. Check certificate expiry: `openssl x509 -enddate -noout -in cert.pem`
2. Renew certificate with your CA or Let's Encrypt
3. Update certificate files on the server
4. Restart the web server to load new certificate
"""
},
{
"id": "error_004",
"content": """
Error Code: ERR_DATABASE_DEADLOCK
A deadlock occurs when two transactions block each other indefinitely.
Common Causes:
- Transactions acquiring locks in different orders
- Long-running transactions holding locks
- Missing indexes causing table scans with locks
Resolution Steps:
1. Review transaction isolation levels
2. Ensure consistent lock ordering across queries
3. Add appropriate indexes to reduce lock duration
4. Implement retry logic for deadlock exceptions
"""
},
{
"id": "error_005",
"content": """
Error Code: ERR_RATE_LIMIT_EXCEEDED
The API has rejected requests due to exceeding the rate limit.
Common Causes:
- Too many requests in a short time window
- Missing request throttling in client code
- Retry loops without backoff
Resolution Steps:
1. Implement exponential backoff for retries
2. Cache responses to reduce API calls
3. Use bulk endpoints where available
4. Request rate limit increase if needed
"""
}
]
def demonstrate_search_types():
"""Show the difference between semantic and keyword search concepts."""
print("=" * 60)
print("SEARCH TYPES COMPARISON")
print("=" * 60)
print("""
📊 DENSE (Semantic) Search:
- Uses vector embeddings to capture meaning
- "memory issues" matches "out of memory" and "RAM exhausted"
- Good for natural language queries
- May miss exact technical terms
📊 SPARSE (Keyword) Search:
- Uses term frequency and exact matching
- "ERR_SSL_CERTIFICATE_EXPIRED" matches exactly
- Good for error codes, IDs, technical terms
- Misses synonyms and paraphrases
📊 HYBRID Search:
- Combines both approaches
- Weighted scoring: (α × dense_score) + ((1-α) × sparse_score)
- Best of both worlds for technical documentation
""")
def hybrid_rag_example():
"""Demonstrate hybrid retrieval with an agent."""
# Create agent with hybrid retrieval enabled
agent = Agent(
name="Tech Support",
instructions="""You are a technical support specialist.
Help users troubleshoot errors using the knowledge base.
Always mention the specific error code when relevant.
Provide step-by-step resolution guidance.""",
knowledge=TROUBLESHOOTING_GUIDES,
knowledge_config={
"hybrid": True, # Enable hybrid search
"top_k": 3
},
user_id="hybrid_demo"
)
# Test with different query types
queries = [
# Semantic query (natural language)
"My application is running out of memory, what should I do?",
# Exact match query (error code)
"ERR_SSL_CERTIFICATE_EXPIRED",
# Mixed query
"Getting connection refused error when connecting to port 8080",
# Conceptual query
"How do I handle API throttling?"
]
print("\n" + "=" * 60)
print("HYBRID RAG IN ACTION")
print("=" * 60)
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 Answer: {response[:300]}..." if len(str(response)) > 300 else f"💡 Answer: {response}")
print("-" * 40)
def compare_retrieval_modes():
"""Compare results with different retrieval configurations."""
print("\n" + "=" * 60)
print("RETRIEVAL MODE COMPARISON")
print("=" * 60)
query = "database transaction blocking issue"
# Dense-only agent
dense_agent = Agent(
name="Dense Search Agent",
instructions="Answer based on the knowledge provided.",
knowledge=TROUBLESHOOTING_GUIDES,
knowledge_config={"hybrid": False},
user_id="dense_demo"
)
# Hybrid agent
hybrid_agent = Agent(
name="Hybrid Search Agent",
instructions="Answer based on the knowledge provided.",
knowledge=TROUBLESHOOTING_GUIDES,
knowledge_config={"hybrid": True},
user_id="hybrid_demo"
)
print(f"\n📝 Query: {query}")
print("\n🔍 Dense-Only Search:")
dense_response = dense_agent.chat(query)
print(f" {dense_response[:200]}...")
print("\n🔍 Hybrid Search:")
hybrid_response = hybrid_agent.chat(query)
print(f" {hybrid_response[:200]}...")
def hybrid_search_tuning():
"""Explain hybrid search tuning parameters."""
print("\n" + "=" * 60)
print("HYBRID SEARCH TUNING")
print("=" * 60)
print("""
Key parameters for hybrid search optimization:
1. **Alpha (α) Weight** - Balance between dense and sparse
- α = 1.0: Pure semantic search
- α = 0.5: Equal weight (good default)
- α = 0.0: Pure keyword search
2. **Top-K Retrieval** - Number of candidates
- Higher K: More context, slower, may dilute relevance
- Lower K: Faster, focused, may miss relevant docs
- Typical range: 3-10 for RAG
3. **Reranking** - Second-pass scoring
- Cross-encoder reranking improves precision
- Adds latency but improves quality
Example configuration:
```python
knowledge_config = {
"hybrid": True,
"hybrid_alpha": 0.6, # Favor semantic
"top_k": 5,
"rerank": True
}
```
""")
def main():
"""Run all hybrid search examples."""
print("\n🚀 PraisonAI Hybrid Search Examples\n")
# Example 1: Explain search types
demonstrate_search_types()
# Example 2: Hybrid RAG in action
hybrid_rag_example()
# Example 3: Compare retrieval modes
compare_retrieval_modes()
# Example 4: Tuning guidance
hybrid_search_tuning()
print("\n✅ Hybrid search examples completed!")
if __name__ == "__main__":
main()
+288
View File
@@ -0,0 +1,288 @@
"""
Graph-Assisted RAG: Relationship-Aware Retrieval
This example demonstrates how graph structures can enhance RAG by
capturing relationships between entities and enabling path-based retrieval.
RAG Concept: Traditional vector search finds similar documents but misses
relationships. Graph-assisted RAG understands connections between entities,
enabling multi-hop reasoning and relationship-aware answers.
"""
from praisonaiagents import Agent
# Sample knowledge base: Interconnected entities (simulating a knowledge graph)
ENTITY_KNOWLEDGE = [
{
"id": "person_einstein",
"type": "person",
"content": """
Albert Einstein (1879-1955) was a theoretical physicist who developed
the theory of relativity. He received the Nobel Prize in Physics in 1921
for his explanation of the photoelectric effect. Einstein worked at the
Institute for Advanced Study in Princeton from 1933 until his death.
Key relationships: Collaborated with Niels Bohr, influenced by Max Planck,
mentored by Heinrich Weber.
"""
},
{
"id": "person_bohr",
"type": "person",
"content": """
Niels Bohr (1885-1962) was a Danish physicist who made foundational
contributions to understanding atomic structure and quantum theory.
He received the Nobel Prize in Physics in 1922. Bohr founded the
Institute of Theoretical Physics in Copenhagen.
Key relationships: Debated with Einstein on quantum mechanics,
mentored Werner Heisenberg, influenced by Ernest Rutherford.
"""
},
{
"id": "person_heisenberg",
"type": "person",
"content": """
Werner Heisenberg (1901-1976) was a German physicist and pioneer of
quantum mechanics. He formulated the uncertainty principle in 1927.
He received the Nobel Prize in Physics in 1932.
Key relationships: Student of Niels Bohr, collaborated with Max Born,
worked with Wolfgang Pauli.
"""
},
{
"id": "concept_relativity",
"type": "concept",
"content": """
Theory of Relativity: A framework developed by Albert Einstein consisting
of special relativity (1905) and general relativity (1915). Special
relativity deals with objects moving at constant speeds, introducing
E=mc². General relativity describes gravity as the curvature of spacetime.
Related concepts: Spacetime, time dilation, gravitational waves.
Applications: GPS satellites, particle accelerators, cosmology.
"""
},
{
"id": "concept_quantum",
"type": "concept",
"content": """
Quantum Mechanics: A fundamental theory describing nature at atomic scales.
Key principles include wave-particle duality, superposition, and entanglement.
Developed through contributions from Planck, Bohr, Heisenberg, Schrödinger.
Related concepts: Uncertainty principle, quantum tunneling, wave function.
Applications: Semiconductors, lasers, MRI machines, quantum computing.
"""
},
{
"id": "institution_princeton",
"type": "institution",
"content": """
Institute for Advanced Study (Princeton): A private research center
founded in 1930. Notable members include Albert Einstein, John von Neumann,
Kurt Gödel, and Robert Oppenheimer. The institute focuses on theoretical
research in mathematics, natural sciences, and humanities.
Location: Princeton, New Jersey, USA.
"""
},
{
"id": "institution_copenhagen",
"type": "institution",
"content": """
Niels Bohr Institute (Copenhagen): Founded in 1921 by Niels Bohr as the
Institute of Theoretical Physics. It became a world center for quantum
mechanics research. Notable visitors included Heisenberg, Pauli, and Dirac.
The "Copenhagen interpretation" of quantum mechanics originated here.
Location: Copenhagen, Denmark.
"""
},
{
"id": "event_solvay",
"type": "event",
"content": """
Solvay Conferences: Series of physics conferences held in Brussels since 1911.
The 1927 Fifth Solvay Conference is famous for the Bohr-Einstein debates
on quantum mechanics. Attendees included Einstein, Bohr, Heisenberg,
Schrödinger, Dirac, and Marie Curie. These conferences shaped modern physics.
"""
}
]
def graph_aware_rag():
"""Demonstrate relationship-aware retrieval."""
# Create agent with entity knowledge
agent = Agent(
name="Physics Historian",
instructions="""You are a physics historian who understands the relationships
between scientists, concepts, and institutions. When answering:
1. Trace connections between entities
2. Explain how ideas and people influenced each other
3. Reference specific relationships mentioned in the knowledge base
4. Build a narrative that shows the interconnected nature of physics history.""",
knowledge=ENTITY_KNOWLEDGE,
user_id="graph_demo"
)
# Relationship-focused queries
queries = [
"How were Einstein and Bohr connected?",
"Trace the lineage of quantum mechanics through its key figures.",
"What institutions were central to 20th century physics?",
"How did the Solvay Conferences shape physics?"
]
print("=" * 60)
print("GRAPH-ASSISTED RAG: Physics History Network")
print("=" * 60)
print("\nKnowledge graph entities:")
for entity in ENTITY_KNOWLEDGE:
print(f" - [{entity['type']}] {entity['id']}")
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 Answer: {response[:350]}..." if len(str(response)) > 350 else f"💡 Answer: {response}")
print("-" * 40)
def multi_hop_reasoning():
"""Demonstrate multi-hop reasoning across entities."""
print("\n" + "=" * 60)
print("MULTI-HOP REASONING")
print("=" * 60)
agent = Agent(
name="Connection Finder",
instructions="""You excel at finding indirect connections between entities.
When asked about relationships, trace the path through intermediate entities.
Format your answer to show the chain of connections.""",
knowledge=ENTITY_KNOWLEDGE,
user_id="multihop_demo"
)
# Multi-hop query requiring traversal
query = """
What is the connection between the uncertainty principle and Princeton?
(Hint: trace through people and their relationships)
"""
print(f"\n📝 Multi-hop Query: {query.strip()}")
print("\n🔗 Expected path: Uncertainty Principle → Heisenberg → Bohr → Einstein → Princeton")
response = agent.chat(query)
print(f"\n💡 Answer:\n{response}")
def entity_type_filtering():
"""Show filtering by entity type."""
print("\n" + "=" * 60)
print("ENTITY TYPE FILTERING")
print("=" * 60)
# Filter to only people
people_only = [e for e in ENTITY_KNOWLEDGE if e['type'] == 'person']
people_agent = Agent(
name="Biographer",
instructions="You are a biographer focused on scientists' lives and relationships.",
knowledge=people_only,
user_id="people_demo"
)
# Filter to only concepts
concepts_only = [e for e in ENTITY_KNOWLEDGE if e['type'] == 'concept']
concept_agent = Agent(
name="Concept Explainer",
instructions="You explain scientific concepts and their relationships.",
knowledge=concepts_only,
user_id="concept_demo"
)
query = "Tell me about the major developments in early 20th century physics."
print(f"\n📝 Query: {query}")
print("\n👤 People-focused view:")
response1 = people_agent.chat(query)
print(f" {response1[:250]}...")
print("\n💡 Concept-focused view:")
response2 = concept_agent.chat(query)
print(f" {response2[:250]}...")
def graph_rag_patterns():
"""Explain graph RAG patterns."""
print("\n" + "=" * 60)
print("GRAPH RAG PATTERNS")
print("=" * 60)
print("""
📊 Graph-Assisted RAG Patterns:
1. **Entity-Centric Retrieval**
- Store entities with their relationships
- Query retrieves entity + connected entities
- Good for: "Tell me about X and its connections"
2. **Path-Based Retrieval**
- Find paths between two entities
- Retrieve all nodes along the path
- Good for: "How is X related to Y?"
3. **Neighborhood Expansion**
- Start with query-matched entities
- Expand to include 1-hop or 2-hop neighbors
- Good for: Comprehensive context
4. **Subgraph Extraction**
- Extract relevant subgraph for query
- Include nodes and edges as context
- Good for: Complex multi-entity questions
Implementation in PraisonAI:
```python
# Structure knowledge as entities with relationships
entities = [
{
"id": "entity_1",
"type": "person",
"content": "...",
"relationships": ["entity_2", "entity_3"]
},
...
]
# Agent retrieves related entities automatically
agent = Agent(
knowledge=entities,
instructions="Trace relationships between entities..."
)
```
""")
def main():
"""Run all graph-assisted RAG examples."""
print("\n🚀 PraisonAI Graph-Assisted RAG Examples\n")
# Example 1: Graph-aware retrieval
graph_aware_rag()
# Example 2: Multi-hop reasoning
multi_hop_reasoning()
# Example 3: Entity type filtering
entity_type_filtering()
# Example 4: Graph RAG patterns
graph_rag_patterns()
print("\n✅ Graph-assisted RAG examples completed!")
if __name__ == "__main__":
main()
+266
View File
@@ -0,0 +1,266 @@
"""
Multi-Document RAG: Synthesizing Information Across Sources
This example demonstrates retrieving and synthesizing information from
multiple documents to answer complex questions.
RAG Concept: Real-world questions often require combining information
from multiple sources. Multi-document RAG retrieves relevant chunks
from different documents and synthesizes a coherent answer.
"""
from praisonaiagents import Agent
# Sample knowledge base: Multiple documents on different topics
CLIMATE_REPORT = {
"id": "climate_2024",
"source": "Global Climate Report 2024",
"content": """
Global Climate Report 2024 - Executive Summary
Global temperatures in 2024 reached 1.45°C above pre-industrial levels,
making it the warmest year on record. Key findings:
- Arctic sea ice extent hit a new minimum in September
- Ocean heat content reached unprecedented levels
- Extreme weather events increased by 23% compared to the decade average
- CO2 concentrations reached 422 ppm, up from 417 ppm in 2023
Regional impacts varied significantly. Europe experienced its hottest
summer, while parts of South America faced severe drought conditions.
The economic cost of climate-related disasters exceeded $380 billion.
"""
}
ENERGY_REPORT = {
"id": "energy_2024",
"source": "World Energy Outlook 2024",
"content": """
World Energy Outlook 2024 - Key Highlights
Renewable energy capacity grew by 35% in 2024, led by solar and wind.
Major developments:
- Solar installations reached 1.5 TW cumulative capacity globally
- Wind power generated 12% of global electricity
- Battery storage capacity doubled to 500 GWh
- Electric vehicle sales reached 20 million units
Despite progress, fossil fuels still account for 75% of primary energy.
Investment in clean energy reached $1.8 trillion, but $4 trillion annually
is needed to meet net-zero targets by 2050.
"""
}
POLICY_REPORT = {
"id": "policy_2024",
"source": "Climate Policy Tracker 2024",
"content": """
Climate Policy Tracker 2024 - Global Progress
145 countries have now committed to net-zero emissions targets.
Policy developments:
- EU Carbon Border Adjustment Mechanism fully implemented
- US Inflation Reduction Act drove $300B in clean energy investment
- China pledged to peak emissions before 2030
- India launched National Green Hydrogen Mission
Carbon pricing now covers 23% of global emissions, up from 20% in 2023.
Average carbon price increased to $45/ton. However, only 4% of emissions
are priced at levels consistent with Paris Agreement goals ($100+/ton).
"""
}
TECH_REPORT = {
"id": "tech_2024",
"source": "Climate Technology Review 2024",
"content": """
Climate Technology Review 2024 - Breakthrough Innovations
Several climate technologies reached commercial viability in 2024:
- Direct Air Capture: Costs fell to $400/ton CO2, down from $600 in 2022
- Green Hydrogen: Production costs reached $3/kg in optimal locations
- Solid-state Batteries: Energy density improved 40%, enabling longer EV range
- Perovskite Solar Cells: Efficiency reached 33%, approaching silicon limits
Emerging technologies to watch:
- Enhanced geothermal systems showing promise in non-volcanic regions
- Ocean-based carbon removal gaining investment interest
- AI-optimized grid management reducing curtailment by 15%
"""
}
# Combine all documents
ALL_DOCUMENTS = [CLIMATE_REPORT, ENERGY_REPORT, POLICY_REPORT, TECH_REPORT]
def multi_document_rag():
"""Demonstrate RAG across multiple documents."""
# Create agent with multi-document knowledge
agent = Agent(
name="Climate Analyst",
instructions="""You are a climate and energy analyst.
Synthesize information from multiple reports to answer questions.
When information comes from different sources, mention which reports
you're drawing from. Provide comprehensive, well-rounded answers.""",
knowledge=ALL_DOCUMENTS,
user_id="multi_doc_demo"
)
# Questions that require multi-document synthesis
queries = [
"What is the current state of global climate and what are we doing about it?",
"How is renewable energy progressing and what policies support it?",
"What technologies are emerging to address climate change?",
"Compare the economic investment in clean energy with climate damage costs."
]
print("=" * 60)
print("MULTI-DOCUMENT RAG: Climate Analysis")
print("=" * 60)
print(f"\nKnowledge base: {len(ALL_DOCUMENTS)} documents")
print(f"Sources: {[d['source'] for d in ALL_DOCUMENTS]}")
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 Answer: {response[:400]}..." if len(str(response)) > 400 else f"💡 Answer: {response}")
print("-" * 40)
def cross_reference_example():
"""Show how agents can cross-reference information."""
print("\n" + "=" * 60)
print("CROSS-REFERENCING ACROSS DOCUMENTS")
print("=" * 60)
agent = Agent(
name="Cross-Reference Analyst",
instructions="""You are an analyst who excels at finding connections
between different reports. When answering:
1. Identify relevant information from each source
2. Note any contradictions or complementary data
3. Synthesize a coherent narrative
4. Cite your sources explicitly.""",
knowledge=ALL_DOCUMENTS,
user_id="cross_ref_demo"
)
query = """
Analyze the relationship between renewable energy growth (from Energy Report),
policy support (from Policy Report), and technology costs (from Tech Report).
How do these factors reinforce each other?
"""
print(f"\n📝 Complex Query: {query.strip()}")
response = agent.chat(query)
print(f"\n💡 Synthesized Answer:\n{response}")
def document_filtering():
"""Demonstrate filtering documents by metadata."""
print("\n" + "=" * 60)
print("DOCUMENT FILTERING BY SOURCE")
print("=" * 60)
# Create agents focused on specific document types
policy_agent = Agent(
name="Policy Expert",
instructions="You are a policy expert. Focus only on policy-related information.",
knowledge=[POLICY_REPORT],
user_id="policy_expert"
)
tech_agent = Agent(
name="Technology Expert",
instructions="You are a technology expert. Focus on technical innovations.",
knowledge=[TECH_REPORT],
user_id="tech_expert"
)
query = "What progress has been made in addressing climate change?"
print(f"\n📝 Query: {query}")
print("\n🏛️ Policy Expert's View:")
policy_response = policy_agent.chat(query)
print(f" {policy_response[:250]}...")
print("\n🔬 Technology Expert's View:")
tech_response = tech_agent.chat(query)
print(f" {tech_response[:250]}...")
print("\n💡 Note: Different document subsets yield different perspectives!")
def multi_document_best_practices():
"""Share best practices for multi-document RAG."""
print("\n" + "=" * 60)
print("MULTI-DOCUMENT RAG BEST PRACTICES")
print("=" * 60)
print("""
📚 Best Practices for Multi-Document RAG:
1. **Document Metadata**
- Include source, date, author in each document
- Enables filtering and citation
2. **Consistent Chunking**
- Use similar chunk sizes across documents
- Maintains balanced retrieval
3. **Source Diversity**
- Retrieve from multiple sources when possible
- Reduces single-source bias
4. **Conflict Resolution**
- Instruct agent how to handle contradictions
- Prefer newer sources or authoritative ones
5. **Citation Requirements**
- Ask agent to cite sources in responses
- Enables verification and trust
Example configuration:
```python
documents = [
{"id": "doc1", "source": "Report A", "date": "2024-01", "content": "..."},
{"id": "doc2", "source": "Report B", "date": "2024-06", "content": "..."},
]
agent = Agent(
knowledge=documents,
instructions="Always cite which report your information comes from."
)
```
""")
def main():
"""Run all multi-document RAG examples."""
print("\n🚀 PraisonAI Multi-Document RAG Examples\n")
# Example 1: Basic multi-document RAG
multi_document_rag()
# Example 2: Cross-referencing
cross_reference_example()
# Example 3: Document filtering
document_filtering()
# Example 4: Best practices
multi_document_best_practices()
print("\n✅ Multi-document RAG examples completed!")
if __name__ == "__main__":
main()
+172
View File
@@ -0,0 +1,172 @@
"""
RAG Evaluation Example
This example demonstrates how to evaluate RAG retrieval quality
using golden queries and expected answers.
Usage:
python rag_evaluation.py
Requirements:
pip install praisonaiagents[knowledge]
"""
import json
from pathlib import Path
from praisonaiagents import Knowledge
from praisonaiagents.rag import RAG, RAGConfig
def create_test_queries():
"""Create sample test queries for evaluation."""
return [
{
"query": "What is the main purpose of the system?",
"expected_doc": "overview.pdf",
"expected_answer_contains": "purpose",
},
{
"query": "What are the key features?",
"expected_doc": "features.pdf",
"expected_answer_contains": "feature",
},
{
"query": "How do I get started?",
"expected_doc": "quickstart.pdf",
"expected_answer_contains": "start",
},
]
def evaluate_rag(rag: RAG, test_queries: list) -> dict:
"""
Evaluate RAG performance on test queries.
Returns metrics including:
- Document retrieval accuracy
- Answer relevance
- Average retrieval score
"""
results = {
"total": len(test_queries),
"doc_hits": 0,
"answer_hits": 0,
"avg_score": 0.0,
"details": [],
}
total_score = 0.0
for test in test_queries:
query = test["query"]
expected_doc = test.get("expected_doc", "")
expected_contains = test.get("expected_answer_contains", "")
# Run RAG query
result = rag.query(query)
# Check if expected document is in citations
doc_found = any(
expected_doc in c.source for c in result.citations
) if expected_doc else True
# Check if answer contains expected text
answer_ok = (
expected_contains.lower() in result.answer.lower()
) if expected_contains else True
# Get average citation score
avg_citation_score = (
sum(c.score for c in result.citations) / len(result.citations)
if result.citations else 0.0
)
if doc_found:
results["doc_hits"] += 1
if answer_ok:
results["answer_hits"] += 1
total_score += avg_citation_score
results["details"].append({
"query": query,
"doc_found": doc_found,
"answer_ok": answer_ok,
"citation_score": avg_citation_score,
"answer_preview": result.answer[:200],
})
results["avg_score"] = total_score / len(test_queries) if test_queries else 0.0
results["doc_accuracy"] = results["doc_hits"] / results["total"]
results["answer_accuracy"] = results["answer_hits"] / results["total"]
return results
def print_results(results: dict):
"""Print evaluation results in a formatted way."""
print("\n" + "=" * 60)
print("RAG Evaluation Results")
print("=" * 60)
print(f"\nTotal Queries: {results['total']}")
print(f"Document Retrieval Accuracy: {results['doc_accuracy']:.1%}")
print(f"Answer Relevance Accuracy: {results['answer_accuracy']:.1%}")
print(f"Average Citation Score: {results['avg_score']:.3f}")
print("\n" + "-" * 60)
print("Detailed Results:")
print("-" * 60)
for i, detail in enumerate(results["details"], 1):
status = "PASS" if detail["doc_found"] and detail["answer_ok"] else "FAIL"
print(f"\n[{i}] {status}: {detail['query']}")
print(f" Doc Found: {detail['doc_found']}")
print(f" Answer OK: {detail['answer_ok']}")
print(f" Score: {detail['citation_score']:.3f}")
def main():
print("=" * 60)
print("RAG Evaluation Example")
print("=" * 60)
# Create knowledge base
knowledge = Knowledge()
knowledge.add("documents/") # Add your documents
# Configure RAG
config = RAGConfig(
top_k=5,
min_score=0.3,
include_citations=True,
)
rag = RAG(knowledge=knowledge, config=config)
# Get test queries
test_queries = create_test_queries()
# Run evaluation
print(f"\nEvaluating {len(test_queries)} queries...")
results = evaluate_rag(rag, test_queries)
# Print results
print_results(results)
# Save results to file
output_path = Path("evaluation_results.json")
with open(output_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nResults saved to {output_path}")
# Return exit code based on results
if results["doc_accuracy"] < 0.8 or results["answer_accuracy"] < 0.8:
print("\nWARNING: Evaluation metrics below threshold!")
return 1
print("\nEvaluation PASSED!")
return 0
if __name__ == "__main__":
exit(main())
+95
View File
@@ -0,0 +1,95 @@
"""
RAG Multi-Agent Example
This example demonstrates how multiple agents can share
the same knowledge base for collaborative RAG.
Usage:
python rag_multi_agent.py
Requirements:
pip install praisonaiagents[knowledge]
"""
from praisonaiagents import Agent, Knowledge, PraisonAIAgents, Task
def main():
print("=" * 60)
print("Multi-Agent RAG: Shared Knowledge Base")
print("=" * 60)
# Create shared knowledge base
shared_knowledge = Knowledge()
shared_knowledge.add("documents/") # Add your documents directory
# Create specialized agents sharing the same knowledge
researcher = Agent(
name="Researcher",
role="Research Analyst",
goal="Find and analyze relevant information from documents",
instructions="""You are a research analyst. Your job is to:
1. Search the knowledge base for relevant information
2. Extract key facts and data points
3. Provide detailed analysis with citations""",
knowledge=shared_knowledge,
)
summarizer = Agent(
name="Summarizer",
role="Content Summarizer",
goal="Create concise summaries of complex information",
instructions="""You are a content summarizer. Your job is to:
1. Take research findings and create clear summaries
2. Highlight the most important points
3. Make complex information accessible""",
knowledge=shared_knowledge,
)
writer = Agent(
name="Writer",
role="Report Writer",
goal="Create well-structured reports",
instructions="""You are a report writer. Your job is to:
1. Combine research and summaries into a cohesive report
2. Ensure proper structure and flow
3. Include citations where appropriate""",
knowledge=shared_knowledge,
)
# Define tasks
research_task = Task(
description="Research the main themes and findings in the documents",
expected_output="Detailed research notes with citations",
agent=researcher,
)
summary_task = Task(
description="Summarize the research findings into key points",
expected_output="Bullet-point summary of main findings",
agent=summarizer,
)
report_task = Task(
description="Write a comprehensive report based on the research and summary",
expected_output="A well-structured report with introduction, findings, and conclusion",
agent=writer,
)
# Run the multi-agent workflow
agents = PraisonAIAgents(
agents=[researcher, summarizer, writer],
tasks=[research_task, summary_task, report_task],
process="sequential",
)
result = agents.start()
print("\n" + "=" * 60)
print("Final Report")
print("=" * 60)
print(result)
if __name__ == "__main__":
main()
+74
View File
@@ -0,0 +1,74 @@
"""
RAG PDF Q&A Example
This example demonstrates how to use RAG to answer questions
about PDF documents with citations.
Usage:
python rag_pdf_qa.py
Requirements:
pip install praisonaiagents[knowledge]
"""
from praisonaiagents import Agent, Knowledge
from praisonaiagents.rag import RAG, RAGConfig
def main():
# Method 1: Simple Agent with Knowledge
print("=" * 60)
print("Method 1: Agent with Knowledge Parameter")
print("=" * 60)
agent = Agent(
name="PDF Assistant",
instructions="You are a helpful assistant that answers questions about documents.",
knowledge=["sample.pdf"], # Add your PDF files here
)
response = agent.start("What are the main topics covered in this document?")
print(f"\nAnswer: {response}\n")
# Method 2: Explicit RAG Pipeline with Citations
print("=" * 60)
print("Method 2: Explicit RAG Pipeline with Citations")
print("=" * 60)
# Create knowledge base
knowledge = Knowledge()
knowledge.add("sample.pdf") # Add your PDF
# Configure RAG
config = RAGConfig(
top_k=5,
min_score=0.3,
include_citations=True,
max_context_tokens=4000,
)
# Create RAG pipeline
rag = RAG(knowledge=knowledge, config=config)
# Query with citations
result = rag.query("What methodology was used in this research?")
print(f"\nAnswer: {result.answer}")
print(f"\nSources ({len(result.citations)}):")
for citation in result.citations:
print(f" [{citation.id}] {citation.source} (score: {citation.score:.2f})")
print(f" {citation.text[:100]}...")
# Method 3: Streaming Response
print("\n" + "=" * 60)
print("Method 3: Streaming Response")
print("=" * 60)
print("\nStreaming answer: ", end="")
for chunk in rag.stream("Summarize the key findings"):
print(chunk, end="", flush=True)
print("\n")
if __name__ == "__main__":
main()
+255
View File
@@ -0,0 +1,255 @@
"""
Reranking: Improving Retrieval Precision
This example demonstrates reranking - a second-pass scoring mechanism
that improves the quality of retrieved documents.
RAG Concept: Initial retrieval (first stage) prioritizes recall - finding
all potentially relevant documents. Reranking (second stage) prioritizes
precision - ordering results by true relevance to the query.
"""
from praisonaiagents import Agent
# Sample knowledge base: Research paper abstracts
RESEARCH_PAPERS = [
{
"id": "paper_001",
"content": """
Title: Attention Mechanisms in Neural Machine Translation
Abstract: This paper introduces a novel attention mechanism for sequence-to-sequence
models in machine translation. Unlike previous approaches that use a fixed-length
context vector, our method allows the model to automatically search for parts of
the source sentence relevant to predicting the target word. Experiments on
English-to-French translation show significant improvements over baseline systems.
"""
},
{
"id": "paper_002",
"content": """
Title: Deep Residual Learning for Image Recognition
Abstract: We present a residual learning framework to ease the training of very deep
neural networks. By reformulating layers as learning residual functions with
reference to layer inputs, we show that these networks are easier to optimize
and can gain accuracy from considerably increased depth. Our 152-layer residual
network achieves state-of-the-art results on ImageNet classification.
"""
},
{
"id": "paper_003",
"content": """
Title: BERT: Pre-training of Deep Bidirectional Transformers
Abstract: We introduce BERT, a new language representation model designed to
pre-train deep bidirectional representations from unlabeled text. Unlike previous
models, BERT is designed to jointly condition on both left and right context.
The pre-trained BERT model can be fine-tuned with just one additional output
layer for a wide range of NLP tasks.
"""
},
{
"id": "paper_004",
"content": """
Title: Generative Adversarial Networks
Abstract: We propose a new framework for estimating generative models via an
adversarial process. We simultaneously train two models: a generative model G
that captures the data distribution, and a discriminative model D that estimates
the probability that a sample came from the training data rather than G.
This framework can generate realistic images from random noise.
"""
},
{
"id": "paper_005",
"content": """
Title: Transformer Architecture for Language Understanding
Abstract: We propose the Transformer, a model architecture eschewing recurrence
and instead relying entirely on an attention mechanism to draw global dependencies
between input and output. The Transformer allows for significantly more
parallelization and achieves new state of the art in translation quality.
"""
},
{
"id": "paper_006",
"content": """
Title: Word Embeddings and Semantic Similarity
Abstract: This paper explores methods for learning distributed representations
of words that capture semantic relationships. We show that word vectors trained
on large corpora encode meaningful syntactic and semantic regularities. These
representations enable analogical reasoning: king - man + woman = queen.
"""
}
]
def explain_reranking():
"""Explain the reranking concept."""
print("=" * 60)
print("RERANKING EXPLAINED")
print("=" * 60)
print("""
🔄 Two-Stage Retrieval Pipeline:
Stage 1: Initial Retrieval (Bi-encoder)
├── Fast approximate matching
├── Retrieves top-K candidates (e.g., 20-100)
├── Uses pre-computed embeddings
└── Optimizes for RECALL (find all relevant docs)
Stage 2: Reranking (Cross-encoder)
├── Slower but more accurate scoring
├── Reorders the top-K candidates
├── Computes query-document interaction
└── Optimizes for PRECISION (best docs first)
Why Reranking Matters:
- Bi-encoders encode query and docs separately (fast but approximate)
- Cross-encoders process query+doc together (slow but accurate)
- Reranking gives you the best of both worlds
""")
def rag_with_reranking():
"""Demonstrate RAG with reranking enabled."""
# Create agent with reranking
agent = Agent(
name="Research Assistant",
instructions="""You are a research paper expert.
Answer questions about machine learning research.
Reference specific papers when relevant.
Be precise about technical details.""",
knowledge=RESEARCH_PAPERS,
knowledge_config={
"rerank": True,
"top_k": 5
},
user_id="rerank_demo"
)
queries = [
"Which paper introduced the attention mechanism for translation?",
"Tell me about pre-training language models",
"What architecture avoids recurrence entirely?"
]
print("\n" + "=" * 60)
print("RAG WITH RERANKING")
print("=" * 60)
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 Answer: {response[:300]}..." if len(str(response)) > 300 else f"💡 Answer: {response}")
print("-" * 40)
def compare_with_without_reranking():
"""Compare retrieval quality with and without reranking."""
print("\n" + "=" * 60)
print("RERANKING IMPACT COMPARISON")
print("=" * 60)
query = "How do transformers handle long-range dependencies?"
# Without reranking
agent_no_rerank = Agent(
name="No Rerank Agent",
instructions="Answer based on the research papers provided.",
knowledge=RESEARCH_PAPERS,
knowledge_config={"rerank": False, "top_k": 3},
user_id="no_rerank"
)
# With reranking
agent_with_rerank = Agent(
name="Rerank Agent",
instructions="Answer based on the research papers provided.",
knowledge=RESEARCH_PAPERS,
knowledge_config={"rerank": True, "top_k": 3},
user_id="with_rerank"
)
print(f"\n📝 Query: {query}")
print("\n🔍 Without Reranking:")
response1 = agent_no_rerank.chat(query)
print(f" {response1[:250]}...")
print("\n🔍 With Reranking:")
response2 = agent_with_rerank.chat(query)
print(f" {response2[:250]}...")
def reranking_configuration():
"""Show reranking configuration options."""
print("\n" + "=" * 60)
print("RERANKING CONFIGURATION")
print("=" * 60)
print("""
Configuration options for reranking in PraisonAI:
```python
agent = Agent(
name="My Agent",
knowledge=documents,
knowledge_config={
# Enable reranking
"rerank": True,
# Number of candidates for initial retrieval
"top_k": 10,
# Number of results after reranking
"rerank_top_k": 3,
# Reranker model (optional)
"rerank_model": "cross-encoder/ms-marco-MiniLM-L-6-v2"
}
)
```
Performance Considerations:
┌─────────────────┬──────────────┬─────────────┐
│ Configuration │ Latency │ Quality │
├─────────────────┼──────────────┼─────────────┤
│ No reranking │ ~50ms │ Good │
│ Light reranker │ ~100ms │ Better │
│ Heavy reranker │ ~200ms │ Best │
└─────────────────┴──────────────┴─────────────┘
Tip: Use reranking for precision-critical applications
where quality matters more than latency.
""")
def main():
"""Run all reranking examples."""
print("\n🚀 PraisonAI Reranking Examples\n")
# Example 1: Explain reranking
explain_reranking()
# Example 2: RAG with reranking
rag_with_reranking()
# Example 3: Compare with/without reranking
compare_with_without_reranking()
# Example 4: Configuration options
reranking_configuration()
print("\n✅ Reranking examples completed!")
if __name__ == "__main__":
main()
+382
View File
@@ -0,0 +1,382 @@
"""
Retrieval Policies: Agent-Decided Retrieval Strategies
This example demonstrates how agents can dynamically decide retrieval
strategies based on query characteristics.
RAG Concept: Not all queries are equal. Some need deep retrieval, others
need broad coverage, and some don't need retrieval at all. Smart policies
optimize both quality and efficiency.
"""
from praisonaiagents import Agent, AutoRagAgent, AutoRagConfig
# Sample knowledge base: IT support documentation
IT_SUPPORT_DOCS = [
{
"id": "password_reset",
"category": "account",
"priority": "high",
"content": """
Password Reset Procedure:
1. Go to the self-service portal at https://portal.company.com
2. Click "Forgot Password"
3. Enter your employee ID and email
4. Check your email for the reset link (valid for 1 hour)
5. Create a new password meeting complexity requirements
If you don't receive the email, check spam or contact IT helpdesk.
"""
},
{
"id": "vpn_setup",
"category": "network",
"priority": "medium",
"content": """
VPN Setup Guide:
1. Download the VPN client from the software center
2. Install with default settings
3. Launch and enter your credentials
4. Select the nearest server location
5. Click Connect
Troubleshooting: If connection fails, check firewall settings
and ensure port 443 is open.
"""
},
{
"id": "email_config",
"category": "communication",
"priority": "medium",
"content": """
Email Configuration:
- Server: mail.company.com
- Port: 993 (IMAP) / 587 (SMTP)
- Security: SSL/TLS required
- Username: your full email address
Mobile setup: Use the Company Portal app for automatic configuration.
"""
},
{
"id": "printer_setup",
"category": "hardware",
"priority": "low",
"content": """
Printer Setup:
1. Open Settings > Printers & Scanners
2. Click Add Printer
3. Select your floor's printer from the list
4. Install drivers if prompted
Printer naming convention: FLOOR-WING-NUMBER (e.g., 3-EAST-01)
"""
},
{
"id": "software_request",
"category": "software",
"priority": "medium",
"content": """
Software Request Process:
1. Submit request through ServiceNow
2. Manager approval required for licensed software
3. IT reviews for security compliance
4. Approved software deployed within 48 hours
Pre-approved software can be installed directly from Software Center.
"""
}
]
def query_classification():
"""Demonstrate query classification for retrieval decisions."""
print("=" * 60)
print("QUERY CLASSIFICATION FOR RETRIEVAL")
print("=" * 60)
# Different query types
queries = [
("How do I reset my password?", "procedural", True),
("Hello!", "greeting", False),
("What is 2 + 2?", "general_knowledge", False),
("VPN connection keeps failing", "troubleshooting", True),
("Thanks for your help", "closing", False),
("How do I request new software?", "procedural", True),
]
print("\n📊 Query Classification Examples:\n")
print(f"{'Query':<40} {'Type':<20} {'Needs RAG':<10}")
print("-" * 70)
for query, query_type, needs_rag in queries:
rag_indicator = "✅ Yes" if needs_rag else "❌ No"
print(f"{query:<40} {query_type:<20} {rag_indicator:<10}")
print("\n💡 Smart agents classify queries to decide retrieval strategy.")
def policy_based_retrieval():
"""Demonstrate different retrieval policies."""
print("\n" + "=" * 60)
print("RETRIEVAL POLICY COMPARISON")
print("=" * 60)
base_agent = Agent(
name="IT Support",
instructions="You are an IT support agent. Help users with technical issues.",
knowledge=IT_SUPPORT_DOCS,
user_id="policy_demo"
)
query = "How do I set up VPN?"
policies = [
("auto", "Agent decides based on query analysis"),
("always", "Always retrieve, even for simple queries"),
("never", "Never retrieve, use only parametric knowledge"),
]
print(f"\n📝 Query: {query}\n")
for policy, description in policies:
config = AutoRagConfig(policy=policy)
auto_agent = AutoRagAgent(agent=base_agent, config=config)
print(f"🔄 Policy: {policy.upper()}")
print(f" Description: {description}")
response = auto_agent.chat(query, user_id="policy_demo")
print(f" Response: {response[:150]}...")
print()
def adaptive_top_k():
"""Demonstrate adaptive top-k based on query complexity."""
print("\n" + "=" * 60)
print("ADAPTIVE TOP-K RETRIEVAL")
print("=" * 60)
print("""
📊 Adaptive Top-K Strategy:
Query Complexity → Top-K Value
─────────────────────────────────
Simple (single topic) → k=1-2
Moderate (related topics) → k=3-5
Complex (multi-topic) → k=5-10
Examples:
""")
queries = [
("Reset password", 2, "Simple - single procedure"),
("VPN not working, also need email setup", 5, "Moderate - two topics"),
("New employee setup: email, VPN, printer, software", 8, "Complex - multiple topics"),
]
for query, suggested_k, complexity in queries:
print(f" Query: \"{query}\"")
print(f" Complexity: {complexity}")
print(f" Suggested top_k: {suggested_k}")
print()
def priority_based_retrieval():
"""Demonstrate priority-based document retrieval."""
print("\n" + "=" * 60)
print("PRIORITY-BASED RETRIEVAL")
print("=" * 60)
# Group documents by priority
high_priority = [d for d in IT_SUPPORT_DOCS if d.get('priority') == 'high']
all_docs = IT_SUPPORT_DOCS
print("\n📊 Document Priorities in Knowledge Base:")
for doc in IT_SUPPORT_DOCS:
print(f" [{doc['priority'].upper():^6}] {doc['id']}")
# High-priority agent (for urgent issues)
urgent_agent = Agent(
name="Urgent Support",
instructions="Handle urgent IT issues. Focus on critical procedures.",
knowledge=high_priority,
user_id="urgent_demo"
)
# Full knowledge agent
full_agent = Agent(
name="Full Support",
instructions="Handle all IT support requests.",
knowledge=all_docs,
user_id="full_demo"
)
query = "I'm locked out of my account!"
print(f"\n📝 Urgent Query: {query}")
print("\n🚨 Urgent Support Agent (high-priority docs only):")
response1 = urgent_agent.chat(query)
print(f" {response1[:200]}...")
print("\n📚 Full Support Agent (all docs):")
response2 = full_agent.chat(query)
print(f" {response2[:200]}...")
def context_aware_retrieval():
"""Demonstrate context-aware retrieval strategies."""
print("\n" + "=" * 60)
print("CONTEXT-AWARE RETRIEVAL")
print("=" * 60)
print("""
🎯 Context-Aware Retrieval Strategies:
1. **User Role Context**
- New employee → Onboarding docs first
- IT admin → Technical docs
- Executive → Summary docs
2. **Time Context**
- Business hours → Standard support
- After hours → Emergency procedures only
3. **Conversation Context**
- First message → Broad retrieval
- Follow-up → Narrow to topic
4. **Device Context**
- Mobile → Mobile-specific guides
- Desktop → Full documentation
Implementation:
```python
def get_retrieval_config(context: dict) -> dict:
if context.get("user_role") == "new_employee":
return {"filter": {"category": "onboarding"}, "top_k": 5}
elif context.get("is_followup"):
return {"top_k": 2} # Narrow focus
else:
return {"top_k": 3} # Default
```
""")
def fallback_strategies():
"""Demonstrate fallback strategies when retrieval fails."""
print("\n" + "=" * 60)
print("RETRIEVAL FALLBACK STRATEGIES")
print("=" * 60)
print("""
🔄 Fallback Strategy Chain:
1. Primary: Vector search in knowledge base
↓ (if no relevant results)
2. Secondary: Expand search with synonyms
↓ (if still no results)
3. Tertiary: Search external sources
↓ (if still no results)
4. Final: Acknowledge limitation, offer alternatives
Example Implementation:
```python
def retrieve_with_fallback(query: str, knowledge) -> str:
# Primary retrieval
results = knowledge.search(query, top_k=3)
if results and results[0].score > 0.7:
return format_context(results)
# Fallback: Expand query
expanded = expand_query(query) # Add synonyms
results = knowledge.search(expanded, top_k=5)
if results and results[0].score > 0.5:
return format_context(results)
# Final fallback
return "I don't have specific information about that. " \\
"Please contact IT helpdesk at ext. 1234."
```
""")
def retrieval_metrics():
"""Explain retrieval quality metrics."""
print("\n" + "=" * 60)
print("RETRIEVAL QUALITY METRICS")
print("=" * 60)
print("""
📈 Key Metrics for Retrieval Policy Tuning:
┌─────────────────┬────────────────────────────────────────┐
│ Metric │ Description │
├─────────────────┼────────────────────────────────────────┤
│ Precision@K │ % of retrieved docs that are relevant │
│ Recall@K │ % of relevant docs that are retrieved │
│ MRR │ Mean Reciprocal Rank of first hit │
│ NDCG │ Normalized Discounted Cumulative Gain │
│ Latency │ Time to retrieve (ms) │
│ Token Usage │ Context tokens consumed │
└─────────────────┴────────────────────────────────────────┘
Policy Optimization Goals:
- High Precision: Use smaller top_k, stricter thresholds
- High Recall: Use larger top_k, relaxed thresholds
- Low Latency: Use smaller top_k, skip reranking
- Low Cost: Use smaller top_k, aggressive filtering
Balanced Default:
```python
config = {
"top_k": 3,
"score_threshold": 0.5,
"rerank": True,
"max_tokens": 2000
}
```
""")
def main():
"""Run all retrieval policy examples."""
print("\n🚀 PraisonAI Retrieval Policies Examples\n")
# Example 1: Query classification
query_classification()
# Example 2: Policy comparison
policy_based_retrieval()
# Example 3: Adaptive top-k
adaptive_top_k()
# Example 4: Priority-based
priority_based_retrieval()
# Example 5: Context-aware
context_aware_retrieval()
# Example 6: Fallback strategies
fallback_strategies()
# Example 7: Metrics
retrieval_metrics()
print("\n✅ Retrieval policies examples completed!")
if __name__ == "__main__":
main()
+334
View File
@@ -0,0 +1,334 @@
"""
Structured Output RAG: Guided and Constrained Generation
This example demonstrates how to guide RAG output into specific formats
and structures using Pydantic models and output constraints.
RAG Concept: Sometimes you need RAG answers in a specific format - JSON,
tables, or structured data. Guided generation ensures the LLM output
conforms to your schema.
"""
from typing import List
from pydantic import BaseModel, Field
from praisonaiagents import Agent
# Sample knowledge base: Product catalog
PRODUCT_CATALOG = [
{
"id": "prod_001",
"content": """
Product: UltraWidget Pro
Category: Electronics
Price: $299.99
Rating: 4.5/5 (1,250 reviews)
Features: Wireless connectivity, 12-hour battery, water-resistant
Availability: In stock
SKU: UW-PRO-2024
"""
},
{
"id": "prod_002",
"content": """
Product: SmartHome Hub
Category: Home Automation
Price: $149.99
Rating: 4.2/5 (890 reviews)
Features: Voice control, 100+ device compatibility, energy monitoring
Availability: In stock
SKU: SH-HUB-2024
"""
},
{
"id": "prod_003",
"content": """
Product: ErgoDesk Standing Desk
Category: Furniture
Price: $599.99
Rating: 4.8/5 (2,100 reviews)
Features: Electric height adjustment, memory presets, cable management
Availability: Ships in 3-5 days
SKU: ED-STAND-2024
"""
},
{
"id": "prod_004",
"content": """
Product: CloudSync Backup Drive
Category: Storage
Price: $199.99
Rating: 4.6/5 (1,800 reviews)
Features: 4TB capacity, automatic backup, encrypted storage
Availability: In stock
SKU: CS-4TB-2024
"""
}
]
# Pydantic models for structured output
class ProductInfo(BaseModel):
"""Structured product information."""
name: str = Field(description="Product name")
price: float = Field(description="Price in USD")
rating: float = Field(description="Rating out of 5")
in_stock: bool = Field(description="Whether product is in stock")
key_features: List[str] = Field(description="Top 3 features")
class ProductComparison(BaseModel):
"""Structured comparison of products."""
products: List[str] = Field(description="Names of compared products")
winner: str = Field(description="Recommended product")
reason: str = Field(description="Why this product is recommended")
price_range: str = Field(description="Price range of compared products")
class SearchResult(BaseModel):
"""Structured search result."""
query: str = Field(description="Original search query")
matches: List[str] = Field(description="Matching product names")
total_found: int = Field(description="Number of matches")
summary: str = Field(description="Brief summary of results")
def json_output_rag():
"""Demonstrate RAG with JSON-structured output."""
print("=" * 60)
print("JSON-STRUCTURED RAG OUTPUT")
print("=" * 60)
agent = Agent(
name="Product Assistant",
instructions="""You are a product assistant that provides structured information.
Always respond with valid JSON matching the requested format.
Extract accurate information from the product catalog.""",
knowledge=PRODUCT_CATALOG,
user_id="json_demo"
)
# Request structured product info
query = """
Tell me about the UltraWidget Pro. Respond in this JSON format:
{
"name": "product name",
"price": 0.00,
"rating": 0.0,
"in_stock": true/false,
"key_features": ["feature1", "feature2", "feature3"]
}
"""
print("\n📝 Query: Get structured product info")
response = agent.chat(query)
print("💡 Structured Response:\n" + str(response))
print("-" * 40)
def pydantic_guided_rag():
"""Demonstrate RAG with Pydantic model guidance."""
print("\n" + "=" * 60)
print("PYDANTIC-GUIDED RAG")
print("=" * 60)
# Generate schema from Pydantic model
schema = ProductComparison.model_json_schema()
agent = Agent(
name="Product Comparator",
instructions=f"""You compare products and provide structured recommendations.
Your response MUST be valid JSON matching this schema:
{schema}
Be objective and base recommendations on the product data.""",
knowledge=PRODUCT_CATALOG,
user_id="pydantic_demo"
)
query = "Compare the UltraWidget Pro and SmartHome Hub. Which should I buy?"
print(f"\n📝 Query: {query}")
print("📋 Expected Schema: ProductComparison")
response = agent.chat(query)
print(f"\n💡 Structured Comparison:\n{response}")
def table_format_rag():
"""Demonstrate RAG with table-formatted output."""
print("\n" + "=" * 60)
print("TABLE-FORMATTED RAG")
print("=" * 60)
agent = Agent(
name="Catalog Browser",
instructions="""You present product information in clean table format.
Use markdown tables for structured data.
Include relevant columns based on the query.""",
knowledge=PRODUCT_CATALOG,
user_id="table_demo"
)
query = "Show me all products in a table with name, price, and rating."
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"\n💡 Table Output:\n{response}")
def list_format_rag():
"""Demonstrate RAG with list-formatted output."""
print("\n" + "=" * 60)
print("LIST-FORMATTED RAG")
print("=" * 60)
agent = Agent(
name="Feature Lister",
instructions="""You extract and present information as organized lists.
Use bullet points and numbered lists appropriately.
Group related items together.""",
knowledge=PRODUCT_CATALOG,
user_id="list_demo"
)
queries = [
"List all product features across the catalog",
"What are the top-rated products? List them in order."
]
for query in queries:
print(f"\n📝 Query: {query}")
response = agent.chat(query)
print(f"💡 List Output:\n{response[:400]}..." if len(str(response)) > 400 else f"💡 List Output:\n{response}")
print("-" * 40)
def constrained_generation():
"""Demonstrate constrained/guided generation patterns."""
print("\n" + "=" * 60)
print("CONSTRAINED GENERATION PATTERNS")
print("=" * 60)
print("""
🎯 Constrained Generation Techniques:
1. **Schema Enforcement**
- Provide JSON schema in instructions
- Agent outputs valid JSON matching schema
```python
instructions = f"Respond with JSON matching: {schema}"
```
2. **Format Templates**
- Give explicit output templates
- Agent fills in the blanks
```python
instructions = '''
Respond in this format:
PRODUCT: [name]
PRICE: $[amount]
VERDICT: [recommendation]
'''
```
3. **Pydantic Validation**
- Parse agent output with Pydantic
- Retry if validation fails
```python
response = agent.chat(query)
try:
result = ProductInfo.model_validate_json(response)
except ValidationError:
# Retry with clarification
```
4. **Output Parsers**
- Post-process agent output
- Extract structured data from text
```python
def parse_product(text: str) -> dict:
# Extract price, rating, etc.
return structured_data
```
5. **Few-Shot Examples**
- Show examples of desired output format
- Agent learns pattern from examples
```python
instructions = '''
Example:
Q: Tell me about Product X
A: {"name": "Product X", "price": 99.99}
Now answer the user's question in the same format.
'''
```
""")
def multi_format_rag():
"""Demonstrate switching between output formats."""
print("\n" + "=" * 60)
print("MULTI-FORMAT RAG")
print("=" * 60)
base_knowledge = PRODUCT_CATALOG
query = "Tell me about the ErgoDesk Standing Desk"
formats = [
("JSON", "Respond with a JSON object containing name, price, and features."),
("Markdown", "Respond with a markdown-formatted product card with headers."),
("Plain Text", "Respond with a brief, conversational product description."),
("Bullet Points", "Respond with bullet points covering key product details.")
]
print(f"\n📝 Base Query: {query}")
print("\n🔄 Same query, different output formats:\n")
for format_name, format_instruction in formats:
agent = Agent(
name=f"{format_name} Agent",
instructions=f"You are a product expert. {format_instruction}",
knowledge=base_knowledge,
user_id=f"format_{format_name.lower()}"
)
response = agent.chat(query)
print(f"📋 {format_name} Format:")
print(f"{response[:250]}..." if len(str(response)) > 250 else response)
print("-" * 40)
def main():
"""Run all structured output RAG examples."""
print("\n🚀 PraisonAI Structured Output RAG Examples\n")
# Example 1: JSON output
json_output_rag()
# Example 2: Pydantic-guided
pydantic_guided_rag()
# Example 3: Table format
table_format_rag()
# Example 4: List format
list_format_rag()
# Example 5: Constrained generation patterns
constrained_generation()
# Example 6: Multi-format
multi_format_rag()
print("\n✅ Structured output RAG examples completed!")
if __name__ == "__main__":
main()
+136
View File
@@ -0,0 +1,136 @@
#!/usr/bin/env python3
"""
Auto-generate CLI reference documentation from --help output.
Usage:
python scripts/generate_cli_docs.py > docs/cli/reference.mdx
python scripts/generate_cli_docs.py --output docs/cli/reference.mdx
This script runs `praisonai <cmd> --help` for key commands and generates
deterministic markdown documentation.
"""
import subprocess
import sys
import argparse
from datetime import datetime
from pathlib import Path
# Commands to document
COMMANDS = [
# Top-level
("praisonai", "--help"),
# Knowledge
("praisonai knowledge", "--help"),
("praisonai knowledge index", "--help"),
("praisonai knowledge search", "--help"),
("praisonai knowledge add", "--help"),
("praisonai knowledge list", "--help"),
# RAG
("praisonai rag", "--help"),
("praisonai rag index", "--help"),
("praisonai rag query", "--help"),
("praisonai rag chat", "--help"),
("praisonai rag eval", "--help"),
("praisonai rag serve", "--help"),
# Run
("praisonai run", "--help"),
# Chat
("praisonai chat", "--help"),
]
def run_help(cmd: str, flag: str = "--help") -> str:
"""Run a command with --help and capture output."""
full_cmd = f"{cmd} {flag}"
try:
result = subprocess.run(
full_cmd.split(),
capture_output=True,
text=True,
timeout=30,
)
return result.stdout or result.stderr or "(no output)"
except subprocess.TimeoutExpired:
return "(timeout)"
except FileNotFoundError:
return "(command not found)"
except Exception as e:
return f"(error: {e})"
def generate_markdown() -> str:
"""Generate markdown documentation."""
lines = [
"---",
"title: CLI Reference",
"description: Auto-generated CLI command reference",
"---",
"",
"# PraisonAI CLI Reference",
"",
"This page is auto-generated from `--help` output.",
"",
f"Generated: {datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')} UTC",
"",
]
for cmd, flag in COMMANDS:
# Section header
cmd_name = cmd.replace("praisonai ", "").replace(" ", "-") or "main"
lines.append(f"## `{cmd}`")
lines.append("")
lines.append("```")
# Get help output
output = run_help(cmd, flag)
lines.append(output.strip())
lines.append("```")
lines.append("")
return "\n".join(lines)
def main():
parser = argparse.ArgumentParser(description="Generate CLI docs from --help")
parser.add_argument(
"--output", "-o",
type=Path,
help="Output file path (default: stdout)",
)
parser.add_argument(
"--check",
action="store_true",
help="Check if output matches existing file (for CI)",
)
args = parser.parse_args()
content = generate_markdown()
if args.check and args.output:
if args.output.exists():
existing = args.output.read_text()
# Compare ignoring timestamp line
existing_lines = [l for l in existing.split("\n") if not l.startswith("Generated:")]
new_lines = [l for l in content.split("\n") if not l.startswith("Generated:")]
if existing_lines != new_lines:
print(f"ERROR: {args.output} is out of date. Run: python scripts/generate_cli_docs.py -o {args.output}", file=sys.stderr)
sys.exit(1)
print(f"OK: {args.output} is up to date")
sys.exit(0)
else:
print(f"ERROR: {args.output} does not exist", file=sys.stderr)
sys.exit(1)
if args.output:
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(content)
print(f"Written to {args.output}")
else:
print(content)
if __name__ == "__main__":
main()
@@ -189,7 +189,21 @@ def __getattr__(name):
_lazy_cache[name] = AutoAgents
return AutoAgents
# Session (imports requests)
# AutoRagAgent (auto RAG retrieval decision)
elif name == "AutoRagAgent":
from .agents.auto_rag_agent import AutoRagAgent
_lazy_cache[name] = AutoRagAgent
return AutoRagAgent
elif name == "AutoRagConfig":
from .agents.auto_rag_agent import AutoRagConfig
_lazy_cache[name] = AutoRagConfig
return AutoRagConfig
elif name == "RagRetrievalPolicy":
from .agents.auto_rag_agent import RetrievalPolicy
_lazy_cache[name] = RetrievalPolicy
return RetrievalPolicy
# Session (imports requests lazily via session/api.py)
elif name == "Session":
from .session import Session
_lazy_cache[name] = Session
@@ -234,6 +248,24 @@ def __getattr__(name):
_lazy_cache[name] = LineRange
return LineRange
# RAG module (lazy loaded for zero performance impact)
elif name == "RAG":
from praisonaiagents.rag import RAG
_lazy_cache[name] = RAG
return RAG
elif name == "RAGConfig":
from praisonaiagents.rag import RAGConfig
_lazy_cache[name] = RAGConfig
return RAGConfig
elif name == "RAGResult":
from praisonaiagents.rag import RAGResult
_lazy_cache[name] = RAGResult
return RAGResult
elif name == "RAGCitation":
from praisonaiagents.rag import Citation
_lazy_cache[name] = Citation
return Citation
# Agent Skills support (lazy loaded for zero performance impact)
elif name == "SkillManager":
from praisonaiagents.skills import SkillManager
@@ -374,6 +406,9 @@ __all__ = [
'TaskOutput',
'ReflectionOutput',
'AutoAgents',
'AutoRagAgent',
'AutoRagConfig',
'RagRetrievalPolicy',
'Session',
'Memory',
'db',
@@ -463,4 +498,9 @@ __all__ = [
# UI
'AGUI',
'A2A',
# RAG
'RAG',
'RAGConfig',
'RAGResult',
'RAGCitation',
]
@@ -45,6 +45,7 @@ if TYPE_CHECKING:
from ..task.task import Task
from ..main import TaskOutput
from ..handoff import Handoff
from ..rag.models import RAGResult, ContextPack
class Agent:
@classmethod
@@ -271,6 +272,7 @@ class Agent:
embedder_config: Optional[Dict[str, Any]] = None,
knowledge: Optional[List[str]] = None,
knowledge_config: Optional[Dict[str, Any]] = None,
rag_config: Optional[Dict[str, Any]] = None,
use_system_prompt: Optional[bool] = True,
markdown: bool = True,
stream: bool = False,
@@ -309,7 +311,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 +465,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
@@ -661,11 +671,15 @@ Your Goal: {self.goal}
self.knowledge = None
self._knowledge_sources = None
self._knowledge_processed = True # No knowledge to process
self._rag_config = None
self._rag_instance = None
else:
# Store knowledge sources for lazy processing
self._knowledge_sources = knowledge
self._knowledge_processed = False
self._knowledge_config = knowledge_config
self._rag_config = rag_config # Store RAG config for citations/streaming
self._rag_instance = None # Lazy loaded RAG instance
self.knowledge = None # Will be initialized on first use
# Fast Context configuration (lazy loaded)
@@ -687,6 +701,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):
@@ -1187,6 +1206,185 @@ Your Goal: {self.goal}
self._knowledge_processed = True
@property
def rag(self):
"""
Lazy-loaded RAG instance for advanced retrieval with citations.
Returns RAG instance configured with agent's knowledge and rag_config.
Returns None if no knowledge sources are configured.
Usage:
agent = Agent(knowledge=["doc.pdf"], rag_config={"include_citations": True})
result = agent.rag.query("What is the main finding?")
print(result.answer)
for citation in result.citations:
print(f"[{citation.id}] {citation.source}")
"""
if not self._knowledge_sources:
return None
if self._rag_instance is None:
self._ensure_knowledge_processed()
if self.knowledge:
try:
from praisonaiagents.rag import RAG, RAGConfig
# Build RAGConfig from rag_config dict
rag_config_obj = RAGConfig()
if self._rag_config:
for key, value in self._rag_config.items():
if hasattr(rag_config_obj, key):
setattr(rag_config_obj, key, value)
# Get LLM instance for RAG
llm = None
if hasattr(self, 'llm_instance') and self.llm_instance:
llm = self.llm_instance
self._rag_instance = RAG(
knowledge=self.knowledge,
config=rag_config_obj,
llm=llm,
)
except ImportError:
logging.warning("RAG module not available. Install with: pip install 'praisonaiagents[rag]'")
return None
return self._rag_instance
def _get_knowledge_context(self, query: str, use_rag: bool = False) -> tuple:
"""
Get knowledge context for a query.
Args:
query: The user's question/prompt
use_rag: If True and rag_config is set, use RAG pipeline for citations
Returns:
Tuple of (context_string, citations_list or None)
"""
if not self._knowledge_sources:
return "", None
self._ensure_knowledge_processed()
if not self.knowledge:
return "", None
# Use RAG if configured and requested
if use_rag and self._rag_config and self.rag:
try:
result = self.rag.query(query, user_id=self.user_id, agent_id=self.agent_id)
return result.context_used, result.citations
except Exception as e:
logging.warning(f"RAG query failed, falling back to basic retrieval: {e}")
# Basic retrieval (original behavior)
search_results = self.knowledge.search(query, agent_id=self.agent_id)
if not search_results:
return "", None
# Format results
if isinstance(search_results, dict) and 'results' in search_results:
knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
else:
knowledge_content = "\n".join(search_results) if isinstance(search_results, list) else str(search_results)
return knowledge_content, None
def rag_query(self, question: str, **kwargs) -> "RAGResult":
"""
Query knowledge using RAG pipeline with citations.
This is the recommended way to get answers with citations from an agent's knowledge.
Args:
question: The question to answer
**kwargs: Additional arguments passed to RAG.query()
Returns:
RAGResult with answer, citations, context_used, and metadata
Raises:
ValueError: If no knowledge sources are configured
ImportError: If RAG module is not available
Usage:
agent = Agent(knowledge=["doc.pdf"], rag_config={"include_citations": True})
result = agent.rag_query("What is the main finding?")
print(result.answer)
for citation in result.citations:
print(f"[{citation.id}] {citation.source}: {citation.text[:100]}")
"""
if not self._knowledge_sources:
raise ValueError("No knowledge sources configured. Add knowledge=[] to Agent init.")
if not self.rag:
raise ImportError("RAG module not available. Install with: pip install 'praisonaiagents[rag]'")
# Pass agent context
kwargs.setdefault('user_id', self.user_id)
kwargs.setdefault('agent_id', self.agent_id)
return self.rag.query(question, **kwargs)
def chat_with_context(
self,
message: str,
context: "ContextPack",
*,
citations_mode: str = "append",
**kwargs,
) -> str:
"""
Chat with pre-retrieved context.
This method allows AutoRagAgent or manual workflows to inject
pre-retrieved context into the agent's chat, enabling conditional
retrieval without duplicating RAG logic.
Args:
message: User message/question
context: ContextPack from RAG.retrieve()
citations_mode: How to include citations (append/hidden/none)
**kwargs: Additional arguments passed to chat()
Returns:
Agent response with optional citations
Usage:
from praisonaiagents import AutoRagAgent
auto_rag = AutoRagAgent(agent=my_agent)
result = auto_rag.chat("What are the key findings?")
# Or manually:
context_pack = rag.retrieve("What are the key findings?")
response = agent.chat_with_context("What are the key findings?", context_pack)
"""
# Build augmented prompt with context
augmented_prompt = f"""Based on the following context, answer the question.
Context:
{context.context}
Question: {message}
Answer:"""
# Call chat with augmented prompt
response = self.chat(augmented_prompt, **kwargs)
# Add citations if configured
if citations_mode == "append" and context.has_citations:
sources = "\n\nSources:\n"
for citation in context.citations:
sources += f" [{citation.id}] {citation.source}\n"
response = response + sources
return response
def _process_knowledge(self, knowledge_item):
"""Process and store knowledge from a file path, URL, or string."""
try:
@@ -2094,6 +2292,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 +2376,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 +2411,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 +3316,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
@@ -1,5 +1,6 @@
"""Agents module for managing multiple AI agents"""
from .agents import PraisonAIAgents
from .autoagents import AutoAgents
from .auto_rag_agent import AutoRagAgent, AutoRagConfig, RetrievalPolicy
__all__ = ['PraisonAIAgents', 'AutoAgents']
__all__ = ['PraisonAIAgents', 'AutoAgents', 'AutoRagAgent', 'AutoRagConfig', 'RetrievalPolicy']
@@ -0,0 +1,415 @@
"""
AutoRagAgent - Agent with automatic RAG retrieval decision.
This module provides an agent wrapper that automatically decides when to
retrieve context from knowledge bases vs direct chat.
Usage:
from praisonaiagents import Agent, AutoRagAgent
agent = Agent(
name="Research Assistant",
knowledge=["docs/manual.pdf"],
)
auto_rag = AutoRagAgent(agent=agent)
result = auto_rag.chat("What are the key findings?") # Auto retrieves
result = auto_rag.chat("Hello!") # Skips retrieval
"""
import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, Set, Dict, Any, TYPE_CHECKING
if TYPE_CHECKING:
from praisonaiagents.agent import Agent
from praisonaiagents.rag import RAG, ContextPack
logger = logging.getLogger(__name__)
class RetrievalPolicy(Enum):
"""Policy for when to perform RAG retrieval."""
AUTO = "auto" # Decide based on query heuristics
ALWAYS = "always" # Always retrieve
NEVER = "never" # Never retrieve (direct chat only)
@dataclass
class AutoRagConfig:
"""Configuration for AutoRagAgent."""
retrieval_policy: RetrievalPolicy = RetrievalPolicy.AUTO
top_k: int = 5
hybrid: bool = False
rerank: bool = False
include_citations: bool = True
citations_mode: str = "append" # append, hidden, none
max_context_tokens: int = 4000
auto_keywords: Set[str] = field(default_factory=lambda: {
"what", "how", "why", "when", "where", "who", "which",
"explain", "describe", "summarize", "find", "search",
"tell me", "show me", "according to", "based on",
"cite", "source", "reference", "document", "paper",
})
auto_min_length: int = 10
def to_dict(self) -> Dict[str, Any]:
"""Serialize to dictionary."""
return {
"retrieval_policy": self.retrieval_policy.value,
"top_k": self.top_k,
"hybrid": self.hybrid,
"rerank": self.rerank,
"include_citations": self.include_citations,
"citations_mode": self.citations_mode,
"max_context_tokens": self.max_context_tokens,
"auto_keywords": list(self.auto_keywords),
"auto_min_length": self.auto_min_length,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "AutoRagConfig":
"""Create from dictionary."""
policy = data.get("retrieval_policy", "auto")
if isinstance(policy, str):
policy = RetrievalPolicy(policy)
keywords = data.get("auto_keywords", None)
if keywords is not None and isinstance(keywords, list):
keywords = set(keywords)
kwargs = {
"retrieval_policy": policy,
"top_k": data.get("top_k", 5),
"hybrid": data.get("hybrid", False),
"rerank": data.get("rerank", False),
"include_citations": data.get("include_citations", True),
"max_context_tokens": data.get("max_context_tokens", 4000),
"auto_min_length": data.get("auto_min_length", 10),
"citations_mode": data.get("citations_mode", "append"),
}
if keywords:
kwargs["auto_keywords"] = keywords
return cls(**kwargs)
class AutoRagAgent:
"""
Agent wrapper with automatic RAG retrieval decision.
Decides when to perform retrieval based on policy and heuristics,
then composes RAG context with Agent chat.
This follows the same pattern as AutoAgents but for RAG:
- AutoAgents: auto-generates agent configs from instructions
- AutoRagAgent: auto-decides when to retrieve context
Usage:
from praisonaiagents import Agent, AutoRagAgent
agent = Agent(
name="Research Assistant",
knowledge=["docs/manual.pdf"],
)
auto_rag = AutoRagAgent(agent=agent)
# Auto mode - decides based on query
result = auto_rag.chat("What are the key findings?")
# Force retrieval
result = auto_rag.chat("Hello", force_retrieval=True)
# Skip retrieval
result = auto_rag.chat("What is 2+2?", skip_retrieval=True)
"""
def __init__(
self,
agent: "Agent",
rag: Optional["RAG"] = None,
config: Optional[AutoRagConfig] = None,
*,
retrieval_policy: Optional[str] = None,
top_k: Optional[int] = None,
hybrid: Optional[bool] = None,
rerank: Optional[bool] = None,
citations: bool = True,
):
"""
Initialize AutoRagAgent.
Args:
agent: Agent instance (must have knowledge configured for RAG)
rag: Optional RAG instance (uses agent.rag if not provided)
config: AutoRagConfig instance
retrieval_policy: Shorthand for config.retrieval_policy (auto/always/never)
top_k: Shorthand for config.top_k
hybrid: Shorthand for config.hybrid
rerank: Shorthand for config.rerank
citations: Shorthand for config.include_citations
"""
self.agent = agent
self._rag = rag
# Build config from kwargs or use provided
if config is not None:
self.config = config
else:
config_kwargs = {}
if retrieval_policy is not None:
config_kwargs["retrieval_policy"] = RetrievalPolicy(retrieval_policy)
if top_k is not None:
config_kwargs["top_k"] = top_k
if hybrid is not None:
config_kwargs["hybrid"] = hybrid
if rerank is not None:
config_kwargs["rerank"] = rerank
config_kwargs["include_citations"] = citations
self.config = AutoRagConfig(**config_kwargs)
# Lazy RAG initialization flag
self._rag_initialized = False
@property
def rag(self) -> Optional["RAG"]:
"""Lazy load RAG from agent if not provided."""
if self._rag is not None:
return self._rag
if not self._rag_initialized:
# Try to get RAG from agent
if hasattr(self.agent, 'rag') and self.agent.rag is not None:
self._rag = self.agent.rag
self._rag_initialized = True
return self._rag
@property
def name(self) -> str:
"""Delegate name to wrapped agent."""
return getattr(self.agent, 'name', 'AutoRagAgent')
def _needs_retrieval(self, query: str) -> bool:
"""
Determine if query needs retrieval based on heuristics.
Simple, local heuristics (no ML classifier):
1. Query length check
2. Keyword presence check
3. Question mark check
Args:
query: User query
Returns:
True if retrieval should be performed
"""
# Check minimum length
if len(query.strip()) < self.config.auto_min_length:
return False
# Check for keywords
query_lower = query.lower()
for keyword in self.config.auto_keywords:
if keyword in query_lower:
return True
# Check for question marks (likely a question)
if "?" in query:
return True
return False
def _should_retrieve(
self,
query: str,
force_retrieval: bool = False,
skip_retrieval: bool = False,
) -> bool:
"""
Determine if retrieval should be performed.
Args:
query: User query
force_retrieval: Override to force retrieval
skip_retrieval: Override to skip retrieval
Returns:
True if retrieval should be performed
"""
# Explicit overrides take precedence
if skip_retrieval:
return False
if force_retrieval:
return True
# Check if RAG is available
if self.rag is None:
logger.debug("RAG not available, skipping retrieval")
return False
# Apply policy
if self.config.retrieval_policy == RetrievalPolicy.ALWAYS:
return True
elif self.config.retrieval_policy == RetrievalPolicy.NEVER:
return False
else: # AUTO
return self._needs_retrieval(query)
def _format_response_with_citations(
self,
response: str,
context_pack: Optional["ContextPack"],
) -> str:
"""Format response with citations based on config."""
if context_pack is None or not context_pack.has_citations:
return response
if self.config.citations_mode == "none":
return response
elif self.config.citations_mode == "hidden":
return response
else: # append
if not self.config.include_citations:
return response
sources = "\n\nSources:\n"
for citation in context_pack.citations:
sources += f" [{citation.id}] {citation.source}\n"
return response + sources
def chat(
self,
message: str,
*,
force_retrieval: bool = False,
skip_retrieval: bool = False,
top_k: Optional[int] = None,
hybrid: Optional[bool] = None,
rerank: Optional[bool] = None,
user_id: Optional[str] = None,
**kwargs,
) -> str:
"""
Chat with automatic RAG retrieval decision.
Decides whether to retrieve context, then runs agent chat.
Args:
message: User message/query
force_retrieval: Force retrieval regardless of policy
skip_retrieval: Skip retrieval regardless of policy
top_k: Override top_k for retrieval
hybrid: Override hybrid retrieval setting
rerank: Override rerank setting
user_id: User ID for RAG retrieval (uses agent.user_id if not provided)
**kwargs: Additional arguments passed to agent.chat()
Returns:
Agent response (with citations if configured)
"""
start_time = time.time()
context_pack = None
# Determine if we should retrieve
should_retrieve = self._should_retrieve(
message,
force_retrieval=force_retrieval,
skip_retrieval=skip_retrieval,
)
if should_retrieve and self.rag is not None:
logger.debug(f"AutoRagAgent: Performing retrieval for: {message[:50]}...")
# Get user_id from parameter or agent
effective_user_id = user_id or getattr(self.agent, 'user_id', None)
# Retrieve context
context_pack = self.rag.retrieve(
query=message,
top_k=top_k or self.config.top_k,
hybrid=hybrid if hybrid is not None else self.config.hybrid,
rerank=rerank if rerank is not None else self.config.rerank,
user_id=effective_user_id,
)
# Use agent's chat_with_context if available
if hasattr(self.agent, 'chat_with_context'):
response = self.agent.chat_with_context(
message=message,
context=context_pack,
citations_mode=self.config.citations_mode,
**kwargs,
)
else:
# Fallback: inject context into message
augmented_message = self._augment_message(message, context_pack)
response = self.agent.chat(augmented_message, **kwargs)
response = self._format_response_with_citations(response, context_pack)
else:
# Direct agent chat without retrieval
logger.debug(f"AutoRagAgent: Skipping retrieval for: {message[:50]}...")
response = self.agent.chat(message, **kwargs)
elapsed = time.time() - start_time
logger.debug(f"AutoRagAgent: chat completed in {elapsed:.2f}s")
return response
def _augment_message(self, message: str, context_pack: "ContextPack") -> str:
"""Augment message with retrieved context."""
return f"""Based on the following context, answer the question.
Context:
{context_pack.context}
Question: {message}
Answer:"""
async def achat(
self,
message: str,
*,
force_retrieval: bool = False,
skip_retrieval: bool = False,
top_k: Optional[int] = None,
hybrid: Optional[bool] = None,
rerank: Optional[bool] = None,
user_id: Optional[str] = None,
**kwargs,
) -> str:
"""
Async version of chat.
Currently wraps sync version. Can be extended for true async.
"""
import asyncio
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
lambda: self.chat(
message=message,
force_retrieval=force_retrieval,
skip_retrieval=skip_retrieval,
top_k=top_k,
hybrid=hybrid,
rerank=rerank,
user_id=user_id,
**kwargs,
),
)
# Alias for consistency with Orchestrator pattern
run = chat
arun = achat
@@ -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
@@ -1,30 +1,151 @@
"""
Context module for PraisonAI Agents.
Context Management Module for PraisonAI Agents.
This module provides context retrieval capabilities including:
This module provides comprehensive context management capabilities:
- Token estimation and budgeting
- Context composition within limits
- Optimization strategies (truncate, sliding window, summarize, prune, smart)
- Real-time context monitoring
- Multi-agent context isolation
Also includes:
- Fast Context: Rapid parallel code search subagent
"""
# Lazy imports to avoid impacting package load time
def __getattr__(name):
"""Lazy import to avoid impacting package performance."""
if name == "FastContext":
from praisonaiagents.context.fast.fast_context import FastContext
return FastContext
elif name == "FastContextResult":
from praisonaiagents.context.fast.result import FastContextResult
return FastContextResult
elif name == "FileMatch":
from praisonaiagents.context.fast.result import FileMatch
return FileMatch
elif name == "LineRange":
from praisonaiagents.context.fast.result import LineRange
return LineRange
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
Zero Performance Impact:
- All imports are lazy loaded via __getattr__
- Features only activate when explicitly enabled
- No overhead when context management is not used
Usage:
from praisonaiagents.context import (
ContextBudgeter,
ContextComposer,
ContextMonitor,
get_optimizer,
)
# Create budgeter for model
budgeter = ContextBudgeter(model="gpt-4o")
budget = budgeter.allocate()
# Compose context within budget
composer = ContextComposer(budget=budget)
result = composer.compose(
system_prompt="You are helpful",
history=messages,
tools=tool_schemas,
)
# Monitor context (opt-in)
monitor = ContextMonitor(enabled=True)
monitor.snapshot(ledger=result.ledger, budget=budget, messages=result.messages)
"""
__all__ = [
# Models
"ContextSegment",
"ContextLedger",
"BudgetAllocation",
"ContextSnapshot",
"OptimizerStrategy",
"OptimizationResult",
"MonitorConfig",
"ContextConfig",
# Token estimation
"estimate_tokens_heuristic",
"estimate_messages_tokens",
"estimate_tool_schema_tokens",
"TokenEstimatorImpl",
"get_estimator",
# Budgeting
"ContextBudgeter",
"get_model_limit",
"get_output_reserve",
# Ledger
"ContextLedgerManager",
"MultiAgentLedger",
# Composition
"ContextComposer",
"ComposedContext",
# Optimization
"BaseOptimizer",
"TruncateOptimizer",
"SlidingWindowOptimizer",
"PruneToolsOptimizer",
"NonDestructiveOptimizer",
"SummarizeOptimizer",
"SmartOptimizer",
"get_optimizer",
"get_effective_history",
# Monitoring
"ContextMonitor",
"MultiAgentMonitor",
"redact_sensitive",
"format_human_snapshot",
"format_json_snapshot",
# Fast Context (legacy)
"FastContext",
"FastContextResult",
"FileMatch",
"LineRange",
]
def __getattr__(name: str):
"""Lazy load module components to avoid import overhead."""
# Models
if name in ("ContextSegment", "ContextLedger", "BudgetAllocation",
"ContextSnapshot", "OptimizerStrategy", "OptimizationResult",
"MonitorConfig", "ContextConfig"):
from . import models
return getattr(models, name)
# Token estimation
if name in ("estimate_tokens_heuristic", "estimate_messages_tokens",
"estimate_tool_schema_tokens", "TokenEstimatorImpl", "get_estimator"):
from . import tokens
return getattr(tokens, name)
# Budgeting
if name in ("ContextBudgeter", "get_model_limit", "get_output_reserve"):
from . import budgeter
return getattr(budgeter, name)
# Ledger
if name in ("ContextLedgerManager", "MultiAgentLedger"):
from . import ledger
return getattr(ledger, name)
# Composition
if name in ("ContextComposer", "ComposedContext"):
from . import composer
return getattr(composer, name)
# Optimization
if name in ("BaseOptimizer", "TruncateOptimizer", "SlidingWindowOptimizer",
"PruneToolsOptimizer", "NonDestructiveOptimizer", "SummarizeOptimizer",
"SmartOptimizer", "get_optimizer", "get_effective_history"):
from . import optimizer
return getattr(optimizer, name)
# Monitoring
if name in ("ContextMonitor", "MultiAgentMonitor", "redact_sensitive",
"format_human_snapshot", "format_json_snapshot"):
from . import monitor
return getattr(monitor, name)
# Fast Context (legacy)
if name == "FastContext":
from praisonaiagents.context.fast.fast_context import FastContext
return FastContext
if name == "FastContextResult":
from praisonaiagents.context.fast.result import FastContextResult
return FastContextResult
if name == "FileMatch":
from praisonaiagents.context.fast.result import FileMatch
return FileMatch
if name == "LineRange":
from praisonaiagents.context.fast.result import LineRange
return LineRange
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,229 @@
"""
Context Budgeter for PraisonAI Agents.
Allocates token budgets across context segments based on model limits.
"""
from typing import Dict, Any, Optional
from .models import BudgetAllocation
# Model token limits (context window sizes)
MODEL_LIMITS: Dict[str, int] = {
# OpenAI
"gpt-4o": 128000,
"gpt-4o-mini": 128000,
"gpt-4-turbo": 128000,
"gpt-4": 8192,
"gpt-3.5-turbo": 16385,
"gpt-3.5-turbo-16k": 16385,
# Anthropic
"claude-3-5-sonnet": 200000,
"claude-3-5-haiku": 200000,
"claude-3-opus": 200000,
"claude-3-sonnet": 200000,
"claude-3-haiku": 200000,
# Google
"gemini-2.0-flash": 1048576,
"gemini-1.5-pro": 2097152,
"gemini-1.5-flash": 1048576,
# Defaults
"default": 128000,
}
# Default output reserves by model family
OUTPUT_RESERVES: Dict[str, int] = {
"gpt-4o": 16384,
"gpt-4o-mini": 16384,
"gpt-4-turbo": 4096,
"gpt-4": 4096,
"gpt-3.5-turbo": 4096,
"claude-3": 8192,
"gemini": 8192,
"default": 8000,
}
def get_model_limit(model: str) -> int:
"""
Get context window limit for a model.
Args:
model: Model name or identifier
Returns:
Context window size in tokens
"""
# Exact match
if model in MODEL_LIMITS:
return MODEL_LIMITS[model]
# Partial match (e.g., "gpt-4o-2024-05-13" matches "gpt-4o")
model_lower = model.lower()
for key, limit in MODEL_LIMITS.items():
if key in model_lower or model_lower.startswith(key):
return limit
return MODEL_LIMITS["default"]
def get_output_reserve(model: str) -> int:
"""
Get recommended output token reserve for a model.
Args:
model: Model name
Returns:
Output reserve in tokens
"""
model_lower = model.lower()
for key, reserve in OUTPUT_RESERVES.items():
if key in model_lower:
return reserve
return OUTPUT_RESERVES["default"]
class ContextBudgeter:
"""
Allocates token budgets across context segments.
Ensures context fits within model limits while reserving
space for output generation.
Example:
budgeter = ContextBudgeter(model="gpt-4o")
allocation = budgeter.allocate()
print(f"History budget: {allocation.history_budget}")
"""
def __init__(
self,
model: str = "gpt-4o-mini",
model_limit: Optional[int] = None,
output_reserve: Optional[int] = None,
# Segment budget overrides (absolute tokens)
system_prompt_budget: int = 2000,
rules_budget: int = 500,
skills_budget: int = 500,
memory_budget: int = 1000,
tools_schema_budget: int = 2000,
tool_outputs_budget: int = 20000,
buffer_budget: int = 1000,
# Ratio-based allocation (if set, overrides fixed budgets)
history_ratio: Optional[float] = None,
):
"""
Initialize budgeter.
Args:
model: Model name for limit lookup
model_limit: Override model limit
output_reserve: Override output reserve
system_prompt_budget: Budget for system prompt
rules_budget: Budget for rules
skills_budget: Budget for skills
memory_budget: Budget for memory
tools_schema_budget: Budget for tool schemas
tool_outputs_budget: Budget for tool outputs
buffer_budget: Safety buffer
history_ratio: If set, allocate this ratio of usable for history
"""
self.model = model
self.model_limit = model_limit or get_model_limit(model)
self.output_reserve = output_reserve or get_output_reserve(model)
self.system_prompt_budget = system_prompt_budget
self.rules_budget = rules_budget
self.skills_budget = skills_budget
self.memory_budget = memory_budget
self.tools_schema_budget = tools_schema_budget
self.tool_outputs_budget = tool_outputs_budget
self.buffer_budget = buffer_budget
self.history_ratio = history_ratio
@property
def usable(self) -> int:
"""Usable tokens after output reserve."""
return self.model_limit - self.output_reserve
def allocate(self) -> BudgetAllocation:
"""
Create a budget allocation.
Returns:
BudgetAllocation with segment budgets
"""
allocation = BudgetAllocation(
model_limit=self.model_limit,
output_reserve=self.output_reserve,
system_prompt=self.system_prompt_budget,
rules=self.rules_budget,
skills=self.skills_budget,
memory=self.memory_budget,
tools_schema=self.tools_schema_budget,
tool_outputs=self.tool_outputs_budget,
buffer=self.buffer_budget,
history=-1, # Dynamic
)
# If history_ratio is set, compute history budget from ratio
if self.history_ratio is not None:
allocation.history = int(self.usable * self.history_ratio)
return allocation
def check_overflow(
self,
current_tokens: int,
threshold: float = 0.8
) -> bool:
"""
Check if current tokens exceed threshold.
Args:
current_tokens: Current total tokens
threshold: Threshold as fraction of usable (0.0-1.0)
Returns:
True if overflow threshold exceeded
"""
return current_tokens > (self.usable * threshold)
def get_remaining(self, current_tokens: int) -> int:
"""
Get remaining tokens before hitting usable limit.
Args:
current_tokens: Current total tokens
Returns:
Remaining tokens (can be negative if over limit)
"""
return self.usable - current_tokens
def get_utilization(self, current_tokens: int) -> float:
"""
Get current utilization as fraction.
Args:
current_tokens: Current total tokens
Returns:
Utilization (0.0-1.0+)
"""
if self.usable == 0:
return 0.0
return current_tokens / self.usable
def to_dict(self) -> Dict[str, Any]:
"""Convert budgeter state to dictionary."""
return {
"model": self.model,
"model_limit": self.model_limit,
"output_reserve": self.output_reserve,
"usable": self.usable,
"allocation": self.allocate().to_dict(),
}
@@ -0,0 +1,335 @@
"""
Context Composer for PraisonAI Agents.
Assembles context from segments within budget constraints.
"""
from typing import Dict, List, Any, Optional
from .models import (
ContextLedger, BudgetAllocation, ContextSegment, ContextSnapshot
)
from .tokens import estimate_tokens_heuristic, estimate_messages_tokens
from .ledger import ContextLedgerManager
class ContextComposer:
"""
Composes context from segments within budget constraints.
Assembles system prompt, rules, skills, memory, tools, and history
while respecting token budgets and applying trimming as needed.
Example:
composer = ContextComposer(budget=budgeter.allocate())
result = composer.compose(
system_prompt="You are a helpful assistant",
history=messages,
tools=tool_schemas,
)
# Use composed messages
response = llm.chat(result.messages)
"""
def __init__(
self,
budget: Optional[BudgetAllocation] = None,
preserve_tool_pairs: bool = True,
max_tool_output_tokens: int = 10000,
keep_recent_turns: int = 5,
):
"""
Initialize composer.
Args:
budget: Budget allocation (uses defaults if None)
preserve_tool_pairs: Keep tool_call/tool_result pairs together
max_tool_output_tokens: Max tokens per tool output
keep_recent_turns: Minimum recent turns to keep
"""
self.budget = budget or BudgetAllocation()
self.preserve_tool_pairs = preserve_tool_pairs
self.max_tool_output_tokens = max_tool_output_tokens
self.keep_recent_turns = keep_recent_turns
self._ledger = ContextLedgerManager()
self._ledger.set_budget(self.budget)
def compose(
self,
system_prompt: str = "",
rules: str = "",
skills: str = "",
memory: str = "",
tools: Optional[List[Dict[str, Any]]] = None,
history: Optional[List[Dict[str, Any]]] = None,
agent_id: Optional[str] = None,
) -> "ComposedContext":
"""
Compose context from segments.
Args:
system_prompt: System prompt content
rules: Rules content
skills: Skills content
memory: Memory content
tools: Tool schemas
history: Conversation history
agent_id: Agent identifier
Returns:
ComposedContext with messages and ledger
"""
self._ledger.reset()
tools = tools or []
history = history or []
# Track fixed segments
self._ledger.track_system_prompt(system_prompt)
self._ledger.track_rules(rules)
self._ledger.track_skills(skills)
self._ledger.track_memory(memory)
self._ledger.track_tools(tools)
# Build combined system message
system_parts = []
if system_prompt:
system_parts.append(system_prompt)
if rules:
system_parts.append(f"\n## Rules\n{rules}")
if skills:
system_parts.append(f"\n## Skills\n{skills}")
if memory:
system_parts.append(f"\n## Memory\n{memory}")
combined_system = "\n".join(system_parts)
# Process history with budget awareness
processed_history = self._process_history(history)
# Build final messages
messages = []
if combined_system:
messages.append({"role": "system", "content": combined_system})
messages.extend(processed_history)
# Track final history tokens
history_tokens = estimate_messages_tokens(processed_history)
self._ledger.track_segment(
ContextSegment.HISTORY,
history_tokens,
len(processed_history)
)
# Check for overflow
total = self._ledger.get_total()
overflow = total > self.budget.usable
return ComposedContext(
messages=messages,
tools=tools,
ledger=self._ledger.get_ledger(),
total_tokens=total,
overflow=overflow,
warnings=self._ledger.get_warnings(),
agent_id=agent_id,
)
def _process_history(
self,
history: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""
Process history with budget constraints.
Applies:
- Tool output truncation
- Sliding window if over budget
- Tool pair preservation
"""
if not history:
return []
# First pass: truncate large tool outputs
processed = []
for msg in history:
processed_msg = self._truncate_tool_output(msg)
processed.append(processed_msg)
# Check if within budget
history_budget = self.budget.history_budget
current_tokens = estimate_messages_tokens(processed)
if current_tokens <= history_budget:
return processed
# Apply sliding window
return self._apply_sliding_window(processed, history_budget)
def _truncate_tool_output(self, msg: Dict[str, Any]) -> Dict[str, Any]:
"""Truncate tool output if too large."""
if msg.get("role") != "tool" and not msg.get("tool_call_id"):
return msg
content = msg.get("content", "")
if not isinstance(content, str):
return msg
tokens = estimate_tokens_heuristic(content)
if tokens <= self.max_tool_output_tokens:
return msg
# Truncate to max tokens (rough: 4 chars per token)
max_chars = self.max_tool_output_tokens * 4
truncated = content[:max_chars] + "\n...[output truncated]..."
result = msg.copy()
result["content"] = truncated
result["_truncated"] = True
return result
def _apply_sliding_window(
self,
history: List[Dict[str, Any]],
budget: int
) -> List[Dict[str, Any]]:
"""Apply sliding window to fit within budget."""
if not history:
return []
# Always keep at least recent turns
min_keep = self.keep_recent_turns * 2 # user + assistant per turn
# Start from end, add messages until budget exceeded
result = []
current_tokens = 0
# Track tool call IDs to preserve pairs
tool_call_ids = set()
for msg in reversed(history):
msg_tokens = estimate_tokens_heuristic(str(msg.get("content", "")))
if current_tokens + msg_tokens > budget and len(result) >= min_keep:
break
result.insert(0, msg)
current_tokens += msg_tokens
# Track tool calls for pair preservation
if self.preserve_tool_pairs:
if "tool_calls" in msg:
for tc in msg.get("tool_calls", []):
if isinstance(tc, dict) and "id" in tc:
tool_call_ids.add(tc["id"])
if msg.get("tool_call_id"):
tool_call_ids.add(msg["tool_call_id"])
# Ensure tool pairs are complete
if self.preserve_tool_pairs:
result = self._ensure_tool_pairs(result, history, tool_call_ids)
return result
def _ensure_tool_pairs(
self,
result: List[Dict[str, Any]],
full_history: List[Dict[str, Any]],
tool_call_ids: set
) -> List[Dict[str, Any]]:
"""Ensure tool_call/tool_result pairs are complete."""
# Find tool calls in result
result_call_ids = set()
result_response_ids = set()
for msg in result:
if "tool_calls" in msg:
for tc in msg.get("tool_calls", []):
if isinstance(tc, dict) and "id" in tc:
result_call_ids.add(tc["id"])
if msg.get("tool_call_id"):
result_response_ids.add(msg["tool_call_id"])
# Find missing pairs
missing_calls = result_response_ids - result_call_ids
missing_responses = result_call_ids - result_response_ids
if not missing_calls and not missing_responses:
return result
# Add missing messages from full history
for msg in full_history:
# Add missing tool calls
if "tool_calls" in msg:
for tc in msg.get("tool_calls", []):
if isinstance(tc, dict) and tc.get("id") in missing_calls:
if msg not in result:
result.insert(0, msg)
break
# Add missing tool responses
if msg.get("tool_call_id") in missing_responses:
if msg not in result:
result.append(msg)
return result
def get_ledger(self) -> ContextLedger:
"""Get current ledger."""
return self._ledger.get_ledger()
def get_warnings(self) -> List[str]:
"""Get current warnings."""
return self._ledger.get_warnings()
class ComposedContext:
"""Result of context composition."""
def __init__(
self,
messages: List[Dict[str, Any]],
tools: List[Dict[str, Any]],
ledger: ContextLedger,
total_tokens: int,
overflow: bool,
warnings: List[str],
agent_id: Optional[str] = None,
):
self.messages = messages
self.tools = tools
self.ledger = ledger
self.total_tokens = total_tokens
self.overflow = overflow
self.warnings = warnings
self.agent_id = agent_id
def to_snapshot(
self,
session_id: str = "",
model_name: str = "",
budget: Optional[BudgetAllocation] = None,
) -> ContextSnapshot:
"""Convert to snapshot for monitoring."""
return ContextSnapshot(
session_id=session_id,
agent_name=self.agent_id or "",
model_name=model_name,
budget=budget,
ledger=self.ledger,
utilization=self.total_tokens / budget.usable if budget else 0.0,
history_content=self.messages,
warnings=self.warnings,
)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"message_count": len(self.messages),
"tool_count": len(self.tools),
"total_tokens": self.total_tokens,
"overflow": self.overflow,
"warnings": self.warnings,
"ledger": self.ledger.to_dict(),
}
@@ -0,0 +1,323 @@
"""
Context Ledger for PraisonAI Agents.
Tracks token usage across context segments with per-agent and per-session support.
"""
from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from .models import ContextLedger, ContextSegment, BudgetAllocation
from .tokens import estimate_messages_tokens, estimate_tool_schema_tokens
@dataclass
class SegmentStats:
"""Statistics for a context segment."""
tokens: int = 0
count: int = 0 # Number of items (messages, tools, etc.)
budget: int = 0
@property
def utilization(self) -> float:
"""Utilization as fraction of budget."""
if self.budget == 0:
return 0.0
return self.tokens / self.budget
@property
def over_budget(self) -> bool:
"""Check if over budget."""
return self.tokens > self.budget if self.budget > 0 else False
class ContextLedgerManager:
"""
Manages context token accounting.
Tracks tokens per segment and provides utilization metrics.
Supports multi-agent scenarios with per-agent ledgers.
Example:
manager = ContextLedgerManager()
manager.set_budget(budgeter.allocate())
# Track segments
manager.track_system_prompt(system_prompt)
manager.track_history(messages)
manager.track_tools(tool_schemas)
# Get stats
print(manager.get_utilization())
print(manager.get_warnings())
"""
def __init__(self, agent_id: Optional[str] = None):
"""
Initialize ledger manager.
Args:
agent_id: Optional agent identifier for multi-agent tracking
"""
self.agent_id = agent_id
self._ledger = ContextLedger()
self._budget: Optional[BudgetAllocation] = None
self._segment_details: Dict[str, SegmentStats] = {}
# Initialize segment stats
for segment in ContextSegment:
self._segment_details[segment.value] = SegmentStats()
def set_budget(self, budget: BudgetAllocation) -> None:
"""Set the budget allocation."""
self._budget = budget
# Update segment budgets
for segment in ContextSegment:
self._segment_details[segment.value].budget = budget.get_segment_budget(segment)
def reset(self) -> None:
"""Reset all token counts."""
self._ledger = ContextLedger()
for segment in ContextSegment:
stats = self._segment_details[segment.value]
stats.tokens = 0
stats.count = 0
def track_segment(
self,
segment: ContextSegment,
tokens: int,
count: int = 1
) -> None:
"""
Track tokens for a segment.
Args:
segment: Context segment
tokens: Token count
count: Number of items
"""
self._ledger.set_segment(segment, tokens)
stats = self._segment_details[segment.value]
stats.tokens = tokens
stats.count = count
def add_to_segment(
self,
segment: ContextSegment,
tokens: int,
count: int = 1
) -> None:
"""
Add tokens to a segment.
Args:
segment: Context segment
tokens: Tokens to add
count: Items to add
"""
self._ledger.add_segment(segment, tokens)
stats = self._segment_details[segment.value]
stats.tokens += tokens
stats.count += count
def track_system_prompt(self, content: str) -> int:
"""Track system prompt tokens."""
from .tokens import estimate_tokens_heuristic
tokens = estimate_tokens_heuristic(content)
self.track_segment(ContextSegment.SYSTEM_PROMPT, tokens)
return tokens
def track_rules(self, content: str) -> int:
"""Track rules tokens."""
from .tokens import estimate_tokens_heuristic
tokens = estimate_tokens_heuristic(content)
self.track_segment(ContextSegment.RULES, tokens)
return tokens
def track_skills(self, content: str) -> int:
"""Track skills tokens."""
from .tokens import estimate_tokens_heuristic
tokens = estimate_tokens_heuristic(content)
self.track_segment(ContextSegment.SKILLS, tokens)
return tokens
def track_memory(self, content: str) -> int:
"""Track memory tokens."""
from .tokens import estimate_tokens_heuristic
tokens = estimate_tokens_heuristic(content)
self.track_segment(ContextSegment.MEMORY, tokens)
return tokens
def track_history(self, messages: List[Dict[str, Any]]) -> int:
"""
Track conversation history tokens.
Args:
messages: List of message dicts
Returns:
Total tokens
"""
tokens = estimate_messages_tokens(messages)
self.track_segment(ContextSegment.HISTORY, tokens, len(messages))
self._ledger.message_count = len(messages)
# Count turns (user messages)
turns = sum(1 for m in messages if m.get("role") == "user")
self._ledger.turn_count = turns
return tokens
def track_tools(self, tools: List[Dict[str, Any]]) -> int:
"""
Track tool schema tokens.
Args:
tools: List of tool definitions
Returns:
Total tokens
"""
tokens = estimate_tool_schema_tokens(tools)
self.track_segment(ContextSegment.TOOLS_SCHEMA, tokens, len(tools))
return tokens
def track_tool_output(self, content: str) -> int:
"""
Add a tool output to tracking.
Args:
content: Tool output content
Returns:
Tokens for this output
"""
from .tokens import estimate_tokens_heuristic
tokens = estimate_tokens_heuristic(content)
self.add_to_segment(ContextSegment.TOOL_OUTPUTS, tokens)
self._ledger.tool_call_count += 1
return tokens
def get_ledger(self) -> ContextLedger:
"""Get current ledger."""
return self._ledger.copy()
def get_total(self) -> int:
"""Get total tokens across all segments."""
return self._ledger.total
def get_utilization(self) -> float:
"""Get overall utilization as fraction of usable budget."""
if self._budget is None:
return 0.0
usable = self._budget.usable
if usable == 0:
return 0.0
return self._ledger.total / usable
def get_segment_stats(self, segment: ContextSegment) -> SegmentStats:
"""Get stats for a segment."""
return self._segment_details[segment.value]
def get_all_stats(self) -> Dict[str, SegmentStats]:
"""Get stats for all segments."""
return self._segment_details.copy()
def get_warnings(self, threshold: float = 0.9) -> List[str]:
"""
Get warnings for segments near or over budget.
Args:
threshold: Warning threshold (0.0-1.0)
Returns:
List of warning messages
"""
warnings = []
for segment in ContextSegment:
stats = self._segment_details[segment.value]
if stats.budget > 0:
if stats.over_budget:
warnings.append(
f"{segment.value}: OVER BUDGET ({stats.tokens:,} / {stats.budget:,} tokens)"
)
elif stats.utilization >= threshold:
pct = int(stats.utilization * 100)
warnings.append(
f"{segment.value}: {pct}% of budget ({stats.tokens:,} / {stats.budget:,} tokens)"
)
# Overall utilization warning
if self._budget:
overall_util = self.get_utilization()
if overall_util >= threshold:
pct = int(overall_util * 100)
warnings.insert(0, f"Overall context: {pct}% of usable budget")
return warnings
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary for serialization."""
return {
"agent_id": self.agent_id,
"ledger": self._ledger.to_dict(),
"total": self._ledger.total,
"utilization": self.get_utilization(),
"segments": {
name: {
"tokens": stats.tokens,
"count": stats.count,
"budget": stats.budget,
"utilization": stats.utilization,
}
for name, stats in self._segment_details.items()
},
"warnings": self.get_warnings(),
}
class MultiAgentLedger:
"""
Manages ledgers for multiple agents.
Provides per-agent tracking and combined views.
"""
def __init__(self):
"""Initialize multi-agent ledger."""
self._agents: Dict[str, ContextLedgerManager] = {}
self._shared_tokens: int = 0 # Tokens shared across agents
def get_agent_ledger(self, agent_id: str) -> ContextLedgerManager:
"""Get or create ledger for an agent."""
if agent_id not in self._agents:
self._agents[agent_id] = ContextLedgerManager(agent_id)
return self._agents[agent_id]
def set_shared_tokens(self, tokens: int) -> None:
"""Set tokens shared across agents (e.g., team memory)."""
self._shared_tokens = tokens
def get_combined_total(self) -> int:
"""Get total tokens across all agents plus shared."""
total = self._shared_tokens
for ledger in self._agents.values():
total += ledger.get_total()
return total
def get_agent_ids(self) -> List[str]:
"""Get list of tracked agent IDs."""
return list(self._agents.keys())
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"agents": {
agent_id: ledger.to_dict()
for agent_id, ledger in self._agents.items()
},
"shared_tokens": self._shared_tokens,
"combined_total": self.get_combined_total(),
}
@@ -0,0 +1,300 @@
"""
Context Management Data Models for PraisonAI Agents.
Provides dataclasses and protocols for context budgeting, composition,
optimization, and monitoring.
"""
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Any, Optional, Literal, Protocol, runtime_checkable
from enum import Enum
class ContextSegment(str, Enum):
"""Segments that contribute to context."""
SYSTEM_PROMPT = "system_prompt"
RULES = "rules"
SKILLS = "skills"
MEMORY = "memory"
TOOLS_SCHEMA = "tools_schema"
HISTORY = "history"
TOOL_OUTPUTS = "tool_outputs"
BUFFER = "buffer"
@dataclass
class ContextLedger:
"""
Tracks token usage across context segments.
Provides per-segment token counts and total usage.
"""
system_prompt: int = 0
rules: int = 0
skills: int = 0
memory: int = 0
tools_schema: int = 0
history: int = 0
tool_outputs: int = 0
buffer: int = 0
# Metadata
turn_count: int = 0
message_count: int = 0
tool_call_count: int = 0
@property
def total(self) -> int:
"""Total tokens across all segments."""
return (
self.system_prompt +
self.rules +
self.skills +
self.memory +
self.tools_schema +
self.history +
self.tool_outputs +
self.buffer
)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return asdict(self)
def get_segment(self, segment: ContextSegment) -> int:
"""Get token count for a segment."""
return getattr(self, segment.value, 0)
def set_segment(self, segment: ContextSegment, tokens: int) -> None:
"""Set token count for a segment."""
setattr(self, segment.value, tokens)
def add_segment(self, segment: ContextSegment, tokens: int) -> None:
"""Add tokens to a segment."""
current = self.get_segment(segment)
self.set_segment(segment, current + tokens)
def copy(self) -> "ContextLedger":
"""Create a copy of this ledger."""
return ContextLedger(
system_prompt=self.system_prompt,
rules=self.rules,
skills=self.skills,
memory=self.memory,
tools_schema=self.tools_schema,
history=self.history,
tool_outputs=self.tool_outputs,
buffer=self.buffer,
turn_count=self.turn_count,
message_count=self.message_count,
tool_call_count=self.tool_call_count,
)
@dataclass
class BudgetAllocation:
"""
Token budget allocation across segments.
Specifies maximum tokens allowed per segment.
"""
model_limit: int = 128000
output_reserve: int = 8000
# Segment budgets (absolute tokens or -1 for dynamic)
system_prompt: int = 2000
rules: int = 500
skills: int = 500
memory: int = 1000
tools_schema: int = 2000
history: int = -1 # Dynamic: fills remaining
tool_outputs: int = 20000
buffer: int = 1000
@property
def usable(self) -> int:
"""Usable tokens after output reserve."""
return self.model_limit - self.output_reserve
@property
def fixed_total(self) -> int:
"""Total of fixed segment budgets."""
total = 0
for seg in [self.system_prompt, self.rules, self.skills,
self.memory, self.tools_schema, self.tool_outputs, self.buffer]:
if seg > 0:
total += seg
return total
@property
def history_budget(self) -> int:
"""Computed history budget (remainder after fixed segments)."""
if self.history > 0:
return self.history
return max(0, self.usable - self.fixed_total)
def get_segment_budget(self, segment: ContextSegment) -> int:
"""Get budget for a segment."""
if segment == ContextSegment.HISTORY:
return self.history_budget
return getattr(self, segment.value, 0)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
result = asdict(self)
result['usable'] = self.usable
result['history_budget'] = self.history_budget
return result
@dataclass
class ContextSnapshot:
"""
A snapshot of the current context state.
Used for monitoring and debugging.
"""
timestamp: str = ""
session_id: str = ""
agent_name: str = ""
model_name: str = ""
# Budget info
budget: Optional[BudgetAllocation] = None
# Token accounting
ledger: Optional[ContextLedger] = None
# Utilization
utilization: float = 0.0
# Content segments (for dump)
system_prompt_content: str = ""
rules_content: str = ""
skills_content: str = ""
memory_content: str = ""
tools_schema_content: str = ""
history_content: List[Dict[str, Any]] = field(default_factory=list)
# Warnings
warnings: List[str] = field(default_factory=list)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary for JSON serialization."""
return {
"timestamp": self.timestamp,
"session_id": self.session_id,
"agent_name": self.agent_name,
"model_name": self.model_name,
"budget": self.budget.to_dict() if self.budget else None,
"ledger": self.ledger.to_dict() if self.ledger else None,
"utilization": self.utilization,
"warnings": self.warnings,
"history_turn_count": len(self.history_content),
}
class OptimizerStrategy(str, Enum):
"""Context optimization strategies."""
TRUNCATE = "truncate"
SLIDING_WINDOW = "sliding_window"
SUMMARIZE = "summarize"
PRUNE_TOOLS = "prune_tools"
NON_DESTRUCTIVE = "non_destructive"
SMART = "smart"
@dataclass
class OptimizationResult:
"""Result of context optimization."""
original_tokens: int = 0
optimized_tokens: int = 0
tokens_saved: int = 0
strategy_used: OptimizerStrategy = OptimizerStrategy.SMART
messages_removed: int = 0
messages_tagged: int = 0
tool_outputs_pruned: int = 0
summary_added: bool = False
@property
def reduction_percent(self) -> float:
"""Percentage of tokens reduced."""
if self.original_tokens == 0:
return 0.0
return (self.tokens_saved / self.original_tokens) * 100
@dataclass
class MonitorConfig:
"""Configuration for context monitoring."""
enabled: bool = False
path: str = "./context.txt"
format: Literal["human", "json"] = "human"
frequency: Literal["turn", "tool_call", "manual", "overflow"] = "turn"
redact_sensitive: bool = True
multi_agent_files: bool = True
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return asdict(self)
@dataclass
class ContextConfig:
"""
Complete context management configuration.
Merges settings from CLI flags, env vars, and config files.
"""
# Auto-compaction
auto_compact: bool = True
compact_threshold: float = 0.8 # Trigger at 80% of usable budget
strategy: OptimizerStrategy = OptimizerStrategy.SMART
# Budget overrides
output_reserve: int = 8000
history_ratio: float = 0.6 # 60% of usable for history
tool_output_max: int = 10000 # Max tokens per tool output
# Pruning
prune_after_tokens: int = 40000
protected_tools: List[str] = field(default_factory=list)
# Monitoring
monitor: MonitorConfig = field(default_factory=MonitorConfig)
# Sliding window
keep_recent_turns: int = 5
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
result = asdict(self)
result['strategy'] = self.strategy.value
return result
@runtime_checkable
class TokenEstimator(Protocol):
"""Protocol for token estimation."""
def estimate(self, text: str) -> int:
"""Estimate tokens for text."""
...
def estimate_messages(self, messages: List[Dict[str, Any]]) -> int:
"""Estimate tokens for a list of messages."""
...
@runtime_checkable
class ContextOptimizer(Protocol):
"""Protocol for context optimization strategies."""
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: ContextLedger,
) -> tuple[List[Dict[str, Any]], OptimizationResult]:
"""Optimize messages to fit within target tokens."""
...
@@ -0,0 +1,478 @@
"""
Context Monitor for PraisonAI Agents.
Writes runtime context to disk for inspection and debugging.
"""
from typing import Dict, List, Any, Optional, Literal
from pathlib import Path
from datetime import datetime
import json
import re
from .models import (
ContextSnapshot, ContextLedger, BudgetAllocation, MonitorConfig, ContextSegment
)
# Sensitive patterns for redaction
SENSITIVE_PATTERNS = [
r'sk-[a-zA-Z0-9]{20,}', # OpenAI API keys
r'sk-proj-[a-zA-Z0-9_-]+', # OpenAI project keys
r'AKIA[0-9A-Z]{16}', # AWS access keys
r'[a-zA-Z0-9]{32,}', # Generic long tokens (conservative)
r'password["\s:=]+["\']?[^"\'\s]+', # Password patterns
r'secret["\s:=]+["\']?[^"\'\s]+', # Secret patterns
r'token["\s:=]+["\']?[^"\'\s]+', # Token patterns
r'api[_-]?key["\s:=]+["\']?[^"\'\s]+', # API key patterns
]
def redact_sensitive(text: str) -> str:
"""
Redact sensitive information from text.
Args:
text: Text to redact
Returns:
Redacted text
"""
result = text
for pattern in SENSITIVE_PATTERNS:
result = re.sub(pattern, '[REDACTED]', result, flags=re.IGNORECASE)
return result
def format_human_snapshot(snapshot: ContextSnapshot, redact: bool = True) -> str:
"""
Format snapshot as human-readable text.
Args:
snapshot: Context snapshot
redact: Whether to redact sensitive info
Returns:
Formatted text
"""
lines = []
# Header
lines.append("=" * 80)
lines.append("PRAISONAI CONTEXT SNAPSHOT")
lines.append("=" * 80)
lines.append(f"Timestamp: {snapshot.timestamp}")
lines.append(f"Session ID: {snapshot.session_id}")
lines.append(f"Agent: {snapshot.agent_name}")
lines.append(f"Model: {snapshot.model_name}")
if snapshot.budget:
lines.append(f"Model Limit: {snapshot.budget.model_limit:,} tokens")
lines.append(f"Output Reserve: {snapshot.budget.output_reserve:,} tokens")
lines.append(f"Usable Budget: {snapshot.budget.usable:,} tokens")
lines.append("")
# Token Ledger
lines.append("-" * 80)
lines.append("TOKEN LEDGER")
lines.append("-" * 80)
lines.append(f"{'Component':<20} | {'Tokens':>10} | {'Budget':>10} | {'% Used':>8}")
lines.append("-" * 20 + "-+-" + "-" * 10 + "-+-" + "-" * 10 + "-+-" + "-" * 8)
if snapshot.ledger and snapshot.budget:
segments = [
(ContextSegment.SYSTEM_PROMPT, "System Prompt"),
(ContextSegment.RULES, "Rules"),
(ContextSegment.SKILLS, "Skills"),
(ContextSegment.MEMORY, "Memory"),
(ContextSegment.TOOLS_SCHEMA, "Tool Schemas"),
(ContextSegment.HISTORY, "History"),
(ContextSegment.TOOL_OUTPUTS, "Tool Outputs"),
]
for seg, name in segments:
tokens = snapshot.ledger.get_segment(seg)
budget = snapshot.budget.get_segment_budget(seg)
pct = (tokens / budget * 100) if budget > 0 else 0
lines.append(f"{name:<20} | {tokens:>10,} | {budget:>10,} | {pct:>7.1f}%")
lines.append("-" * 20 + "-+-" + "-" * 10 + "-+-" + "-" * 10 + "-+-" + "-" * 8)
lines.append(f"{'TOTAL':<20} | {snapshot.ledger.total:>10,} | {snapshot.budget.usable:>10,} | {snapshot.utilization * 100:>7.1f}%")
lines.append("")
# Warnings
if snapshot.warnings:
lines.append("-" * 80)
lines.append("WARNINGS")
lines.append("-" * 80)
for warning in snapshot.warnings:
lines.append(f"⚠ {warning}")
lines.append("")
# History summary
if snapshot.history_content:
lines.append("-" * 80)
lines.append(f"CONVERSATION HISTORY ({len(snapshot.history_content)} messages)")
lines.append("-" * 80)
for i, msg in enumerate(snapshot.history_content):
role = msg.get("role", "unknown").upper()
content = msg.get("content", "")
if isinstance(content, str):
if redact:
content = redact_sensitive(content)
# Truncate for display
if len(content) > 200:
content = content[:200] + "..."
lines.append(f"[{i+1}] {role}: {content}")
else:
lines.append(f"[{i+1}] {role}: [complex content]")
lines.append("")
# Footer
lines.append("=" * 80)
lines.append("END CONTEXT SNAPSHOT")
lines.append("=" * 80)
return "\n".join(lines)
def format_json_snapshot(snapshot: ContextSnapshot, redact: bool = True) -> str:
"""
Format snapshot as JSON.
Args:
snapshot: Context snapshot
redact: Whether to redact sensitive info
Returns:
JSON string
"""
data = snapshot.to_dict()
# Add history with optional redaction
if snapshot.history_content:
history = []
for msg in snapshot.history_content:
msg_copy = msg.copy()
content = msg_copy.get("content", "")
if isinstance(content, str) and redact:
msg_copy["content"] = redact_sensitive(content)
history.append(msg_copy)
data["history"] = history
return json.dumps(data, indent=2, default=str)
class ContextMonitor:
"""
Monitors and writes context to disk.
Provides real-time visibility into context state for debugging
and optimization.
Example:
monitor = ContextMonitor(
enabled=True,
path="./context.txt",
format="human",
)
# Write snapshot
monitor.snapshot(
ledger=ledger,
budget=budget,
messages=messages,
agent_name="Assistant",
)
"""
def __init__(
self,
enabled: bool = False,
path: str = "./context.txt",
format: Literal["human", "json"] = "human",
frequency: Literal["turn", "tool_call", "manual", "overflow"] = "turn",
redact_sensitive: bool = True,
multi_agent_files: bool = True,
):
"""
Initialize monitor.
Args:
enabled: Whether monitoring is enabled
path: Output file path
format: Output format (human or json)
frequency: Update frequency
redact_sensitive: Redact sensitive info
multi_agent_files: Create per-agent files
"""
self.enabled = enabled
self.path = Path(path)
self.format = format
self.frequency = frequency
self.redact = redact_sensitive
self.multi_agent_files = multi_agent_files
self._last_snapshot: Optional[ContextSnapshot] = None
self._turn_count = 0
self._tool_call_count = 0
@classmethod
def from_config(cls, config: MonitorConfig) -> "ContextMonitor":
"""Create monitor from config."""
return cls(
enabled=config.enabled,
path=config.path,
format=config.format,
frequency=config.frequency,
redact_sensitive=config.redact_sensitive,
multi_agent_files=config.multi_agent_files,
)
def enable(self) -> None:
"""Enable monitoring."""
self.enabled = True
def disable(self) -> None:
"""Disable monitoring."""
self.enabled = False
def set_path(self, path: str) -> None:
"""Set output path."""
self.path = Path(path)
def set_format(self, format: Literal["human", "json"]) -> None:
"""Set output format."""
self.format = format
def set_frequency(self, frequency: Literal["turn", "tool_call", "manual", "overflow"]) -> None:
"""Set update frequency."""
self.frequency = frequency
def should_write(
self,
trigger: Literal["turn", "tool_call", "manual", "overflow"]
) -> bool:
"""Check if should write based on frequency."""
if not self.enabled:
return False
if trigger == "manual":
return True
if self.frequency == "manual":
return trigger == "manual"
if self.frequency == "overflow":
return trigger == "overflow"
if self.frequency == "tool_call":
return trigger in ("tool_call", "overflow")
# frequency == "turn"
return trigger in ("turn", "overflow")
def snapshot(
self,
ledger: ContextLedger,
budget: BudgetAllocation,
messages: List[Dict[str, Any]],
session_id: str = "",
agent_name: str = "",
model_name: str = "",
trigger: Literal["turn", "tool_call", "manual", "overflow"] = "turn",
) -> Optional[str]:
"""
Create and write a context snapshot.
Args:
ledger: Current token ledger
budget: Budget allocation
messages: Current messages
session_id: Session identifier
agent_name: Agent name
model_name: Model name
trigger: What triggered this snapshot
Returns:
Path to written file, or None if not written
"""
if not self.should_write(trigger):
return None
# Track counts
if trigger == "turn":
self._turn_count += 1
elif trigger == "tool_call":
self._tool_call_count += 1
# Create snapshot
snapshot = ContextSnapshot(
timestamp=datetime.utcnow().isoformat() + "Z",
session_id=session_id,
agent_name=agent_name,
model_name=model_name,
budget=budget,
ledger=ledger,
utilization=ledger.total / budget.usable if budget.usable > 0 else 0,
history_content=messages,
warnings=self._get_warnings(ledger, budget),
)
self._last_snapshot = snapshot
# Determine output path
output_path = self._get_output_path(agent_name)
# Format content
if self.format == "json":
content = format_json_snapshot(snapshot, self.redact)
else:
content = format_human_snapshot(snapshot, self.redact)
# Atomic write
self._atomic_write(output_path, content)
return str(output_path)
def _get_output_path(self, agent_name: str = "") -> Path:
"""Get output path, optionally per-agent."""
if not self.multi_agent_files or not agent_name:
return self.path
# Create per-agent path
stem = self.path.stem
suffix = self.path.suffix or (".json" if self.format == "json" else ".txt")
# Sanitize agent name
safe_name = re.sub(r'[^\w\-]', '_', agent_name.lower())
return self.path.parent / f"{stem}_{safe_name}{suffix}"
def _get_warnings(
self,
ledger: ContextLedger,
budget: BudgetAllocation
) -> List[str]:
"""Generate warnings based on current state."""
warnings = []
utilization = ledger.total / budget.usable if budget.usable > 0 else 0
if utilization >= 1.0:
warnings.append(f"CRITICAL: Context overflow! {utilization*100:.1f}% of usable budget")
elif utilization >= 0.9:
warnings.append(f"WARNING: Context at {utilization*100:.1f}% of usable budget")
elif utilization >= 0.8:
warnings.append(f"NOTICE: Context at {utilization*100:.1f}% of usable budget")
# Check individual segments
if budget.tool_outputs > 0:
tool_util = ledger.tool_outputs / budget.tool_outputs
if tool_util >= 0.9:
warnings.append(f"Tool outputs at {tool_util*100:.1f}% of budget")
return warnings
def _atomic_write(self, path: Path, content: str) -> None:
"""Write content atomically using temp file + rename."""
# Ensure parent directory exists
path.parent.mkdir(parents=True, exist_ok=True)
# Write to temp file
temp_path = path.with_suffix(path.suffix + ".tmp")
try:
temp_path.write_text(content, encoding="utf-8")
# Atomic rename
temp_path.rename(path)
except Exception:
# Clean up temp file on error
if temp_path.exists():
temp_path.unlink()
raise
def get_last_snapshot(self) -> Optional[ContextSnapshot]:
"""Get the last snapshot."""
return self._last_snapshot
def get_stats(self) -> Dict[str, Any]:
"""Get monitor statistics."""
return {
"enabled": self.enabled,
"path": str(self.path),
"format": self.format,
"frequency": self.frequency,
"turn_count": self._turn_count,
"tool_call_count": self._tool_call_count,
"last_snapshot": self._last_snapshot.timestamp if self._last_snapshot else None,
}
class MultiAgentMonitor:
"""
Manages monitors for multiple agents.
Creates per-agent monitors and optional combined view.
"""
def __init__(
self,
base_path: str = "./context",
format: Literal["human", "json"] = "human",
enabled: bool = False,
redact_sensitive: bool = True,
):
"""
Initialize multi-agent monitor.
Args:
base_path: Base path for output files
format: Output format
enabled: Whether monitoring is enabled
redact_sensitive: Redact sensitive info
"""
self.base_path = Path(base_path)
self.format = format
self.enabled = enabled
self.redact = redact_sensitive
self._agents: Dict[str, ContextMonitor] = {}
def get_agent_monitor(self, agent_id: str) -> ContextMonitor:
"""Get or create monitor for an agent."""
if agent_id not in self._agents:
suffix = ".json" if self.format == "json" else ".txt"
safe_id = re.sub(r'[^\w\-]', '_', agent_id.lower())
path = self.base_path / f"context_{safe_id}{suffix}"
self._agents[agent_id] = ContextMonitor(
enabled=self.enabled,
path=str(path),
format=self.format,
redact_sensitive=self.redact,
multi_agent_files=False, # Already per-agent
)
return self._agents[agent_id]
def enable_all(self) -> None:
"""Enable all monitors."""
self.enabled = True
for monitor in self._agents.values():
monitor.enable()
def disable_all(self) -> None:
"""Disable all monitors."""
self.enabled = False
for monitor in self._agents.values():
monitor.disable()
def get_agent_ids(self) -> List[str]:
"""Get list of monitored agent IDs."""
return list(self._agents.keys())
@@ -0,0 +1,484 @@
"""
Context Optimizer for PraisonAI Agents.
Provides strategies for reducing context size when approaching limits.
"""
from typing import Dict, List, Any, Optional
from abc import ABC, abstractmethod
from .models import ContextLedger, OptimizerStrategy, OptimizationResult
from .tokens import estimate_messages_tokens
class BaseOptimizer(ABC):
"""Base class for optimization strategies."""
@abstractmethod
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: Optional[ContextLedger] = None,
) -> tuple:
"""
Optimize messages to fit within target tokens.
Args:
messages: Messages to optimize
target_tokens: Target token count
ledger: Current token ledger
Returns:
Tuple of (optimized_messages, OptimizationResult)
"""
pass
class TruncateOptimizer(BaseOptimizer):
"""Truncate oldest messages strategy."""
def __init__(self, preserve_system: bool = True, preserve_recent: int = 5):
self.preserve_system = preserve_system
self.preserve_recent = preserve_recent
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: Optional[ContextLedger] = None,
) -> tuple:
original_tokens = estimate_messages_tokens(messages)
if original_tokens <= target_tokens:
return messages, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=original_tokens,
tokens_saved=0,
strategy_used=OptimizerStrategy.TRUNCATE,
)
# Separate system and other messages
system_msgs = []
other_msgs = []
for msg in messages:
if self.preserve_system and msg.get("role") == "system":
system_msgs.append(msg)
else:
other_msgs.append(msg)
# Keep recent messages
recent = other_msgs[-self.preserve_recent:] if other_msgs else []
result = system_msgs + recent
# Continue removing until under target
while estimate_messages_tokens(result) > target_tokens and len(result) > 1:
for i, msg in enumerate(result):
if msg.get("role") != "system":
result.pop(i)
break
optimized_tokens = estimate_messages_tokens(result)
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=optimized_tokens,
tokens_saved=original_tokens - optimized_tokens,
strategy_used=OptimizerStrategy.TRUNCATE,
messages_removed=len(messages) - len(result),
)
class SlidingWindowOptimizer(BaseOptimizer):
"""Keep a sliding window of recent messages."""
def __init__(self, preserve_system: bool = True):
self.preserve_system = preserve_system
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: Optional[ContextLedger] = None,
) -> tuple:
original_tokens = estimate_messages_tokens(messages)
if original_tokens <= target_tokens:
return messages, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=original_tokens,
tokens_saved=0,
strategy_used=OptimizerStrategy.SLIDING_WINDOW,
)
result = []
# Keep system messages
for msg in messages:
if self.preserve_system and msg.get("role") == "system":
result.append(msg)
# Add messages from end until target reached
non_system = [m for m in messages if m.get("role") != "system"]
for msg in reversed(non_system):
test_result = result + [msg]
if estimate_messages_tokens(test_result) <= target_tokens:
result.insert(len([m for m in result if m.get("role") == "system"]), msg)
else:
break
# Ensure order
system_msgs = [m for m in result if m.get("role") == "system"]
other_msgs = [m for m in result if m.get("role") != "system"]
result = system_msgs + other_msgs
optimized_tokens = estimate_messages_tokens(result)
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=optimized_tokens,
tokens_saved=original_tokens - optimized_tokens,
strategy_used=OptimizerStrategy.SLIDING_WINDOW,
messages_removed=len(messages) - len(result),
)
class PruneToolsOptimizer(BaseOptimizer):
"""Prune old tool outputs while keeping tool calls."""
def __init__(
self,
preserve_recent: int = 5,
max_output_chars: int = 500,
protected_tools: Optional[List[str]] = None,
):
self.preserve_recent = preserve_recent
self.max_output_chars = max_output_chars
self.protected_tools = protected_tools or []
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: Optional[ContextLedger] = None,
) -> tuple:
original_tokens = estimate_messages_tokens(messages)
result = []
pruned_count = 0
# 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 tool outputs
for msg in older:
role = msg.get("role", "")
content = msg.get("content", "")
# Check if this is a tool result
if role == "tool" or msg.get("tool_call_id"):
# Check if protected
tool_name = msg.get("name", "")
if tool_name in self.protected_tools:
result.append(msg)
continue
# Truncate if too long
if isinstance(content, str) and len(content) > self.max_output_chars:
pruned_msg = msg.copy()
pruned_msg["content"] = content[:self.max_output_chars] + "\n...[output truncated]..."
pruned_msg["_pruned"] = True
result.append(pruned_msg)
pruned_count += 1
else:
result.append(msg)
else:
result.append(msg)
result.extend(recent)
optimized_tokens = estimate_messages_tokens(result)
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=optimized_tokens,
tokens_saved=original_tokens - optimized_tokens,
strategy_used=OptimizerStrategy.PRUNE_TOOLS,
tool_outputs_pruned=pruned_count,
)
class NonDestructiveOptimizer(BaseOptimizer):
"""
Tag messages for exclusion without deleting them.
Messages are tagged with _condense_parent and excluded from
effective history, but can be restored on rewind.
"""
def __init__(self, preserve_recent: int = 3):
self.preserve_recent = preserve_recent
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: Optional[ContextLedger] = None,
) -> tuple:
import uuid
original_tokens = estimate_messages_tokens(messages)
if original_tokens <= target_tokens:
return messages, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=original_tokens,
tokens_saved=0,
strategy_used=OptimizerStrategy.NON_DESTRUCTIVE,
)
condense_id = str(uuid.uuid4())
# 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"]
# Keep recent, tag older
keep_start = max(0, len(other_msgs) - self.preserve_recent)
result = []
tagged_count = 0
# Add system messages
result.extend(system_msgs)
# Tag older messages
for i, msg in enumerate(other_msgs):
if i < keep_start:
tagged_msg = msg.copy()
if "_condense_parent" not in tagged_msg:
tagged_msg["_condense_parent"] = condense_id
tagged_count += 1
result.append(tagged_msg)
else:
result.append(msg)
# Calculate effective tokens (excluding tagged)
effective_msgs = [m for m in result if "_condense_parent" not in m]
optimized_tokens = estimate_messages_tokens(effective_msgs)
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=optimized_tokens,
tokens_saved=original_tokens - optimized_tokens,
strategy_used=OptimizerStrategy.NON_DESTRUCTIVE,
messages_tagged=tagged_count,
)
class SummarizeOptimizer(BaseOptimizer):
"""
Summarize older messages using LLM.
Note: This is a placeholder that creates a structured summary.
Actual LLM summarization should be done at the agent level.
"""
def __init__(self, preserve_recent: int = 5, max_summary_items: int = 10):
self.preserve_recent = preserve_recent
self.max_summary_items = max_summary_items
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: Optional[ContextLedger] = None,
) -> tuple:
original_tokens = estimate_messages_tokens(messages)
if original_tokens <= target_tokens:
return messages, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=original_tokens,
tokens_saved=0,
strategy_used=OptimizerStrategy.SUMMARIZE,
)
# 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 = list(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 summary
summary_parts = []
for msg in older[:self.max_summary_items]:
role = msg.get("role", "unknown")
content = msg.get("content", "")
if isinstance(content, str) and content:
summary_parts.append(f"[{role}]: {content[:150]}...")
if summary_parts:
summary = (
"[Previous conversation summary]\n"
+ "\n".join(summary_parts)
)
result.append({
"role": "system",
"content": summary,
"_summary": True,
"_original_count": len(older),
})
result.extend(recent)
optimized_tokens = estimate_messages_tokens(result)
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=optimized_tokens,
tokens_saved=original_tokens - optimized_tokens,
strategy_used=OptimizerStrategy.SUMMARIZE,
messages_removed=len(older),
summary_added=bool(older),
)
class SmartOptimizer(BaseOptimizer):
"""
Smart optimization combining multiple strategies.
Applies strategies in order:
1. Prune tool outputs
2. Sliding window
3. Summarize if still over
"""
def __init__(
self,
preserve_recent: int = 5,
protected_tools: Optional[List[str]] = None,
):
self.preserve_recent = preserve_recent
self.protected_tools = protected_tools or []
self._prune = PruneToolsOptimizer(
preserve_recent=preserve_recent,
protected_tools=protected_tools,
)
self._window = SlidingWindowOptimizer()
self._summarize = SummarizeOptimizer(preserve_recent=preserve_recent)
def optimize(
self,
messages: List[Dict[str, Any]],
target_tokens: int,
ledger: Optional[ContextLedger] = None,
) -> tuple:
original_tokens = estimate_messages_tokens(messages)
if original_tokens <= target_tokens:
return messages, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=original_tokens,
tokens_saved=0,
strategy_used=OptimizerStrategy.SMART,
)
# Step 1: Prune tool outputs
result, prune_result = self._prune.optimize(messages, target_tokens, ledger)
if estimate_messages_tokens(result) <= target_tokens:
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=prune_result.optimized_tokens,
tokens_saved=original_tokens - prune_result.optimized_tokens,
strategy_used=OptimizerStrategy.SMART,
tool_outputs_pruned=prune_result.tool_outputs_pruned,
)
# Step 2: Sliding window
result, window_result = self._window.optimize(result, target_tokens, ledger)
if estimate_messages_tokens(result) <= target_tokens:
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=window_result.optimized_tokens,
tokens_saved=original_tokens - window_result.optimized_tokens,
strategy_used=OptimizerStrategy.SMART,
tool_outputs_pruned=prune_result.tool_outputs_pruned,
messages_removed=window_result.messages_removed,
)
# Step 3: Summarize
result, summary_result = self._summarize.optimize(result, target_tokens, ledger)
optimized_tokens = estimate_messages_tokens(result)
return result, OptimizationResult(
original_tokens=original_tokens,
optimized_tokens=optimized_tokens,
tokens_saved=original_tokens - optimized_tokens,
strategy_used=OptimizerStrategy.SMART,
tool_outputs_pruned=prune_result.tool_outputs_pruned,
messages_removed=window_result.messages_removed + summary_result.messages_removed,
summary_added=summary_result.summary_added,
)
# Strategy registry
OPTIMIZER_REGISTRY: Dict[OptimizerStrategy, type] = {
OptimizerStrategy.TRUNCATE: TruncateOptimizer,
OptimizerStrategy.SLIDING_WINDOW: SlidingWindowOptimizer,
OptimizerStrategy.PRUNE_TOOLS: PruneToolsOptimizer,
OptimizerStrategy.NON_DESTRUCTIVE: NonDestructiveOptimizer,
OptimizerStrategy.SUMMARIZE: SummarizeOptimizer,
OptimizerStrategy.SMART: SmartOptimizer,
}
def get_optimizer(
strategy: OptimizerStrategy,
**kwargs
) -> BaseOptimizer:
"""
Get an optimizer instance for a strategy.
Args:
strategy: Optimization strategy
**kwargs: Strategy-specific arguments
Returns:
Optimizer instance
"""
optimizer_class = OPTIMIZER_REGISTRY.get(strategy, SmartOptimizer)
return optimizer_class(**kwargs)
def get_effective_history(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Get effective history excluding tagged messages.
Filters out messages with _condense_parent tags.
Args:
messages: Full message history
Returns:
Effective messages for API calls
"""
return [m for m in messages if "_condense_parent" not in m]
@@ -0,0 +1,238 @@
"""
Token Estimation for PraisonAI Agents.
Provides fast, offline token estimation with optional accurate tokenizer support.
"""
from typing import List, Dict, Any, Optional
import json
# Token estimation constants (based on OpenAI tokenizer patterns)
ASCII_TOKENS_PER_CHAR = 0.25 # ~4 chars per token for ASCII
NON_ASCII_TOKENS_PER_CHAR = 1.3 # CJK and other scripts are denser
# Optional tokenizer cache
_tokenizer_cache: Dict[str, Any] = {}
def estimate_tokens_heuristic(text: str) -> int:
"""
Estimate token count using character-based heuristic.
Fast and works offline. Accuracy: ~85-95% for English text.
Args:
text: Text to estimate tokens for
Returns:
Estimated token count
"""
if not text:
return 0
total = 0.0
for char in text:
code_point = ord(char)
if code_point <= 127:
total += ASCII_TOKENS_PER_CHAR
else:
total += NON_ASCII_TOKENS_PER_CHAR
return max(1, int(total))
def estimate_tokens_accurate(text: str, model: str = "gpt-4") -> int:
"""
Estimate tokens using tiktoken if available, else fallback to heuristic.
Args:
text: Text to estimate tokens for
model: Model name for tokenizer selection
Returns:
Token count
"""
try:
import tiktoken
# Cache tokenizer
if model not in _tokenizer_cache:
try:
_tokenizer_cache[model] = tiktoken.encoding_for_model(model)
except KeyError:
# Fallback to cl100k_base for unknown models
_tokenizer_cache[model] = tiktoken.get_encoding("cl100k_base")
enc = _tokenizer_cache[model]
return len(enc.encode(text))
except ImportError:
return estimate_tokens_heuristic(text)
def estimate_message_tokens(message: Dict[str, Any], use_accurate: bool = False) -> int:
"""
Estimate tokens for a single message.
Handles various message formats including tool calls and multi-part content.
Args:
message: Message dict with role and content
use_accurate: Use tiktoken if available
Returns:
Estimated token count
"""
estimate_fn = estimate_tokens_accurate if use_accurate else estimate_tokens_heuristic
total = 0
# Role overhead (~4 tokens for role markers)
total += 4
# Content
content = message.get("content", "")
if isinstance(content, str):
total += estimate_fn(content)
elif isinstance(content, list):
# Multi-part content (text, images, etc.)
for part in content:
if isinstance(part, dict):
if "text" in part:
total += estimate_fn(str(part["text"]))
elif "type" in part:
# Image or other media - estimate based on type
if part.get("type") == "image_url":
total += 85 # Base image token cost
else:
total += estimate_fn(json.dumps(part))
else:
total += estimate_fn(str(part))
# Tool calls
if "tool_calls" in message:
for tool_call in message.get("tool_calls", []):
if isinstance(tool_call, dict):
# Function name and arguments
func = tool_call.get("function", {})
total += estimate_fn(func.get("name", ""))
total += estimate_fn(func.get("arguments", ""))
total += 10 # Overhead for tool call structure
# Tool call ID (for tool responses)
if "tool_call_id" in message:
total += 10
# Name field
if "name" in message:
total += estimate_fn(message["name"])
return total
def estimate_messages_tokens(
messages: List[Dict[str, Any]],
use_accurate: bool = False
) -> int:
"""
Estimate total tokens for a list of messages.
Args:
messages: List of message dicts
use_accurate: Use tiktoken if available
Returns:
Total estimated token count
"""
if not messages:
return 0
total = 0
for msg in messages:
total += estimate_message_tokens(msg, use_accurate)
# Add overhead for message array structure
total += 3 # Array markers
return total
def estimate_tool_schema_tokens(
tools: List[Dict[str, Any]],
use_accurate: bool = False
) -> int:
"""
Estimate tokens for tool/function schemas.
Args:
tools: List of tool definitions
use_accurate: Use tiktoken if available
Returns:
Estimated token count
"""
if not tools:
return 0
estimate_fn = estimate_tokens_accurate if use_accurate else estimate_tokens_heuristic
# Serialize tools to JSON and estimate
try:
tools_json = json.dumps(tools)
return estimate_fn(tools_json)
except (TypeError, ValueError):
# Fallback: estimate per-tool
total = 0
for tool in tools:
if isinstance(tool, dict):
total += estimate_fn(json.dumps(tool))
else:
total += 100 # Default estimate per tool
return total
class TokenEstimatorImpl:
"""
Token estimator implementation.
Provides both heuristic and accurate estimation methods.
"""
def __init__(self, use_accurate: bool = False, model: str = "gpt-4"):
"""
Initialize estimator.
Args:
use_accurate: Use tiktoken if available
model: Model name for tokenizer selection
"""
self.use_accurate = use_accurate
self.model = model
def estimate(self, text: str) -> int:
"""Estimate tokens for text."""
if self.use_accurate:
return estimate_tokens_accurate(text, self.model)
return estimate_tokens_heuristic(text)
def estimate_messages(self, messages: List[Dict[str, Any]]) -> int:
"""Estimate tokens for messages."""
return estimate_messages_tokens(messages, self.use_accurate)
def estimate_tools(self, tools: List[Dict[str, Any]]) -> int:
"""Estimate tokens for tool schemas."""
return estimate_tool_schema_tokens(tools, self.use_accurate)
# Default estimator instance
_default_estimator: Optional[TokenEstimatorImpl] = None
def get_estimator(use_accurate: bool = False, model: str = "gpt-4") -> TokenEstimatorImpl:
"""Get or create a token estimator."""
global _default_estimator
if _default_estimator is None or _default_estimator.use_accurate != use_accurate:
_default_estimator = TokenEstimatorImpl(use_accurate, model)
return _default_estimator
+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
@@ -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,70 @@
"""
PraisonAI RAG - Retrieval Augmented Generation Module.
This module provides a thin orchestration layer over Knowledge for RAG workflows.
Knowledge handles indexing/retrieval; RAG adds answer generation with citations.
Usage:
from praisonaiagents.rag import RAG, RAGConfig, RAGResult, Citation
# With existing Knowledge
rag = RAG(knowledge=my_knowledge)
result = rag.query("What is the main finding?")
print(result.answer)
for citation in result.citations:
print(f" [{citation.id}] {citation.source}")
All imports are lazy to avoid performance impact when RAG is not used.
"""
from typing import TYPE_CHECKING
# Lazy loading to avoid import overhead
_LAZY_IMPORTS = {
# Models
"Citation": ("praisonaiagents.rag.models", "Citation"),
"ContextPack": ("praisonaiagents.rag.models", "ContextPack"),
"RAGResult": ("praisonaiagents.rag.models", "RAGResult"),
"RAGConfig": ("praisonaiagents.rag.models", "RAGConfig"),
# Protocols
"ContextBuilderProtocol": ("praisonaiagents.rag.protocols", "ContextBuilderProtocol"),
"CitationFormatterProtocol": ("praisonaiagents.rag.protocols", "CitationFormatterProtocol"),
# Pipeline
"RAG": ("praisonaiagents.rag.pipeline", "RAG"),
# Context utilities
"build_context": ("praisonaiagents.rag.context", "build_context"),
"truncate_context": ("praisonaiagents.rag.context", "truncate_context"),
"deduplicate_chunks": ("praisonaiagents.rag.context", "deduplicate_chunks"),
}
_cache = {}
def __getattr__(name: str):
"""Lazy load RAG components."""
if name in _cache:
return _cache[name]
if name in _LAZY_IMPORTS:
module_path, attr_name = _LAZY_IMPORTS[name]
import importlib
module = importlib.import_module(module_path)
value = getattr(module, attr_name)
_cache[name] = value
return value
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
def __dir__():
"""List available attributes."""
return list(_LAZY_IMPORTS.keys())
__all__ = list(_LAZY_IMPORTS.keys())
if TYPE_CHECKING:
from .models import Citation, ContextPack, RAGResult, RAGConfig
from .protocols import ContextBuilderProtocol, CitationFormatterProtocol
from .pipeline import RAG
from .context import build_context, truncate_context, deduplicate_chunks
@@ -0,0 +1,199 @@
"""
Context Building Utilities for RAG.
Provides token-aware context assembly with deduplication.
No heavy imports - only stdlib.
"""
import hashlib
from typing import Any, Dict, List, Optional, Set, Tuple
def _estimate_tokens(text: str) -> int:
"""
Estimate token count for text.
Uses simple heuristic: ~4 characters per token for English.
This avoids importing tokenizers for lightweight operation.
"""
return len(text) // 4 + 1
def _chunk_hash(text: str, source: Optional[str] = None) -> str:
"""
Generate stable hash for chunk deduplication.
Args:
text: Chunk text content
source: Optional source identifier
Returns:
Hash string for deduplication
"""
content = f"{source or ''}:{text}"
return hashlib.md5(content.encode()).hexdigest()[:16]
def deduplicate_chunks(
results: List[Dict[str, Any]],
similarity_threshold: float = 0.9,
) -> List[Dict[str, Any]]:
"""
Deduplicate chunks by content hash.
Args:
results: List of search results with 'text' or 'memory' key
similarity_threshold: Not used currently (hash-based dedup)
Returns:
Deduplicated list of results
"""
seen_hashes: Set[str] = set()
unique_results: List[Dict[str, Any]] = []
for result in results:
# Handle different result formats
text = result.get("text") or result.get("memory", "")
source = result.get("metadata", {}).get("source", "")
chunk_id = _chunk_hash(text, source)
if chunk_id not in seen_hashes:
seen_hashes.add(chunk_id)
unique_results.append(result)
return unique_results
def truncate_context(
text: str,
max_tokens: int,
suffix: str = "\n\n[Context truncated...]",
) -> str:
"""
Truncate context to fit within token budget.
Args:
text: Context text to truncate
max_tokens: Maximum tokens allowed
suffix: Suffix to add when truncated
Returns:
Truncated text
"""
estimated_tokens = _estimate_tokens(text)
if estimated_tokens <= max_tokens:
return text
# Calculate target character count
suffix_tokens = _estimate_tokens(suffix)
target_tokens = max_tokens - suffix_tokens
target_chars = target_tokens * 4
# Truncate at word boundary
truncated = text[:target_chars]
last_space = truncated.rfind(" ")
if last_space > target_chars * 0.8:
truncated = truncated[:last_space]
return truncated + suffix
def build_context(
results: List[Dict[str, Any]],
max_tokens: int = 4000,
deduplicate: bool = True,
separator: str = "\n\n---\n\n",
include_source: bool = True,
) -> Tuple[str, List[Dict[str, Any]]]:
"""
Build context string from retrieval results.
Args:
results: List of search results
max_tokens: Maximum tokens for context
deduplicate: Whether to deduplicate chunks
separator: Separator between chunks
include_source: Whether to include source info
Returns:
Tuple of (context_string, used_results)
"""
if not results:
return "", []
# Deduplicate if requested
if deduplicate:
results = deduplicate_chunks(results)
context_parts: List[str] = []
used_results: List[Dict[str, Any]] = []
current_tokens = 0
separator_tokens = _estimate_tokens(separator)
for i, result in enumerate(results):
# Handle different result formats
text = result.get("text") or result.get("memory", "")
if not text:
continue
# Build chunk text with optional source
if include_source:
source = result.get("metadata", {}).get("source", "")
filename = result.get("metadata", {}).get("filename", "")
source_label = filename or source or f"Source {i + 1}"
chunk_text = f"[{source_label}]\n{text}"
else:
chunk_text = text
chunk_tokens = _estimate_tokens(chunk_text)
# Check if we can fit this chunk
if current_tokens + chunk_tokens + separator_tokens > max_tokens:
# Try to fit partial chunk
remaining_tokens = max_tokens - current_tokens - separator_tokens
if remaining_tokens > 100: # Only include if meaningful
truncated = truncate_context(chunk_text, remaining_tokens, "...")
context_parts.append(truncated)
used_results.append(result)
break
context_parts.append(chunk_text)
used_results.append(result)
current_tokens += chunk_tokens + separator_tokens
context = separator.join(context_parts)
return context, used_results
class DefaultContextBuilder:
"""
Default implementation of ContextBuilderProtocol.
Provides token-aware context assembly with deduplication.
"""
def __init__(
self,
separator: str = "\n\n---\n\n",
include_source: bool = True,
):
self.separator = separator
self.include_source = include_source
def build(
self,
results: List[Dict[str, Any]],
max_tokens: int = 4000,
deduplicate: bool = True,
) -> str:
"""Build context string from results."""
context, _ = build_context(
results=results,
max_tokens=max_tokens,
deduplicate=deduplicate,
separator=self.separator,
include_source=self.include_source,
)
return context
@@ -0,0 +1,256 @@
"""
RAG Data Models for PraisonAI Agents.
Lightweight dataclasses for RAG results and configuration.
No heavy imports - only stdlib and typing.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from enum import Enum
class RetrievalStrategy(str, Enum):
"""Available retrieval strategies for RAG."""
BASIC = "basic"
FUSION = "fusion"
HYBRID = "hybrid"
@dataclass
class Citation:
"""
Source citation for RAG answers.
Attributes:
id: Unique citation identifier (e.g., "[1]", "[2]")
source: Source document path or URL
text: Text snippet from the source
score: Relevance score (0-1)
doc_id: Optional document identifier
chunk_id: Optional chunk identifier within document
offset: Optional character offset in source
metadata: Additional metadata
"""
id: str
source: str
text: str
score: float = 0.0
doc_id: Optional[str] = None
chunk_id: Optional[str] = None
offset: Optional[int] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"id": self.id,
"source": self.source,
"text": self.text,
"score": self.score,
"doc_id": self.doc_id,
"chunk_id": self.chunk_id,
"offset": self.offset,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "Citation":
"""Create from dictionary."""
return cls(
id=data.get("id", ""),
source=data.get("source", ""),
text=data.get("text", ""),
score=data.get("score", 0.0),
doc_id=data.get("doc_id"),
chunk_id=data.get("chunk_id"),
offset=data.get("offset"),
metadata=data.get("metadata", {}),
)
def __str__(self) -> str:
"""Human-readable string representation."""
snippet = self.text[:100] + "..." if len(self.text) > 100 else self.text
return f"[{self.id}] {self.source}: {snippet}"
@dataclass
class ContextPack:
"""
Context pack for orchestrator pattern - retrieval without generation.
Provides deterministic context that can be passed to Agent.chat_with_context().
Attributes:
context: The formatted context string ready for injection
citations: List of source citations
query: The original query
metadata: Additional metadata (timing, retrieval stats, etc.)
"""
context: str
citations: List[Citation] = field(default_factory=list)
query: str = ""
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"context": self.context,
"citations": [c.to_dict() for c in self.citations],
"query": self.query,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "ContextPack":
"""Create from dictionary."""
return cls(
context=data.get("context", ""),
citations=[Citation.from_dict(c) for c in data.get("citations", [])],
query=data.get("query", ""),
metadata=data.get("metadata", {}),
)
@property
def has_citations(self) -> bool:
"""Check if context pack has citations."""
return len(self.citations) > 0
def format_for_prompt(self, include_sources: bool = True) -> str:
"""Format context for injection into a prompt."""
if not include_sources or not self.citations:
return self.context
sources = "\n\nSources:\n"
for citation in self.citations:
sources += f" [{citation.id}] {citation.source}\n"
return self.context + sources
@dataclass
class RAGResult:
"""
Result from a RAG query.
Attributes:
answer: The generated answer text
citations: List of source citations
context_used: The context string passed to the LLM
query: The original query
metadata: Additional metadata (timing, model info, etc.)
"""
answer: str
citations: List[Citation] = field(default_factory=list)
context_used: str = ""
query: str = ""
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"answer": self.answer,
"citations": [c.to_dict() for c in self.citations],
"context_used": self.context_used,
"query": self.query,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "RAGResult":
"""Create from dictionary."""
return cls(
answer=data.get("answer", ""),
citations=[Citation.from_dict(c) for c in data.get("citations", [])],
context_used=data.get("context_used", ""),
query=data.get("query", ""),
metadata=data.get("metadata", {}),
)
@property
def has_citations(self) -> bool:
"""Check if result has citations."""
return len(self.citations) > 0
def format_answer_with_citations(self) -> str:
"""Format answer with inline citation references."""
if not self.citations:
return self.answer
# Build citation reference section
refs = "\n\nSources:\n"
for citation in self.citations:
refs += f" {citation}\n"
return self.answer + refs
@dataclass
class RAGConfig:
"""
Configuration for RAG pipeline.
Attributes:
top_k: Number of chunks to retrieve
min_score: Minimum relevance score threshold
max_context_tokens: Maximum tokens for context
include_citations: Whether to include citations in result
retrieval_strategy: Strategy for retrieval (basic, fusion, hybrid)
rerank: Whether to rerank results
rerank_top_k: Number of results after reranking
template: Prompt template with {context} and {question} placeholders
system_prompt: Optional system prompt for LLM
stream: Whether to stream responses
"""
top_k: int = 5
min_score: float = 0.0
max_context_tokens: int = 4000
include_citations: bool = True
retrieval_strategy: RetrievalStrategy = RetrievalStrategy.BASIC
rerank: bool = False
rerank_top_k: int = 3
model: Optional[str] = None # LLM model to use, defaults to gpt-4o-mini
template: str = """Answer the question based on the context below.
Context:
{context}
Question: {question}
Answer:"""
system_prompt: Optional[str] = None
stream: bool = False
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"top_k": self.top_k,
"min_score": self.min_score,
"max_context_tokens": self.max_context_tokens,
"include_citations": self.include_citations,
"retrieval_strategy": self.retrieval_strategy.value,
"rerank": self.rerank,
"rerank_top_k": self.rerank_top_k,
"template": self.template,
"system_prompt": self.system_prompt,
"stream": self.stream,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "RAGConfig":
"""Create from dictionary."""
strategy = data.get("retrieval_strategy", "basic")
if isinstance(strategy, str):
strategy = RetrievalStrategy(strategy)
return cls(
top_k=data.get("top_k", 5),
min_score=data.get("min_score", 0.0),
max_context_tokens=data.get("max_context_tokens", 4000),
include_citations=data.get("include_citations", True),
retrieval_strategy=strategy,
rerank=data.get("rerank", False),
rerank_top_k=data.get("rerank_top_k", 3),
template=data.get("template", cls.template),
system_prompt=data.get("system_prompt"),
stream=data.get("stream", False),
)
@@ -0,0 +1,652 @@
"""
RAG Pipeline for PraisonAI Agents.
Thin orchestrator that uses Knowledge for retrieval and LLM for generation.
Provides citations, streaming, and async support.
"""
import logging
import time
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
from .models import Citation, ContextPack, RAGConfig, RAGResult
from .context import build_context, DefaultContextBuilder
logger = logging.getLogger(__name__)
class DefaultCitationFormatter:
"""
Default implementation of CitationFormatterProtocol.
Creates Citation objects from search results with stable IDs.
"""
def format(
self,
results: List[Dict[str, Any]],
start_id: int = 1,
) -> List[Citation]:
"""Format retrieval results as citations."""
citations = []
for i, result in enumerate(results):
# Handle different result formats
text = result.get("text") or result.get("memory", "")
metadata = result.get("metadata", {})
score = result.get("score", 0.0)
# Extract source info
source = metadata.get("source", "")
filename = metadata.get("filename", "")
doc_id = result.get("id") or metadata.get("doc_id", "")
chunk_id = metadata.get("chunk_id", "")
citation = Citation(
id=str(start_id + i),
source=filename or source or f"Source {start_id + i}",
text=text[:500] if text else "", # Limit snippet length
score=float(score) if score else 0.0,
doc_id=str(doc_id) if doc_id else None,
chunk_id=str(chunk_id) if chunk_id else None,
metadata=metadata,
)
citations.append(citation)
return citations
class RAG:
"""
Retrieval Augmented Generation Pipeline.
Thin orchestrator that:
1. Retrieves relevant context from Knowledge
2. Optionally reranks results
3. Builds context with token-aware truncation
4. Generates answer using LLM
5. Returns answer with citations
Usage:
from praisonaiagents.rag import RAG
from praisonaiagents.knowledge import Knowledge
knowledge = Knowledge(config={...})
knowledge.add("document.pdf")
rag = RAG(knowledge=knowledge)
result = rag.query("What is the main finding?")
print(result.answer)
for citation in result.citations:
print(f" [{citation.id}] {citation.source}")
"""
def __init__(
self,
knowledge: Any,
llm: Optional[Any] = None,
config: Optional[RAGConfig] = None,
reranker: Optional[Any] = None,
context_builder: Optional[Any] = None,
citation_formatter: Optional[Any] = None,
):
"""
Initialize RAG pipeline.
Args:
knowledge: Knowledge instance for retrieval
llm: Optional LLM instance (uses default if not provided)
config: RAG configuration
reranker: Optional reranker for result refinement
context_builder: Optional custom context builder
citation_formatter: Optional custom citation formatter
"""
self.knowledge = knowledge
self._llm = llm
self.config = config or RAGConfig()
self.reranker = reranker
self.context_builder = context_builder or DefaultContextBuilder()
self.citation_formatter = citation_formatter or DefaultCitationFormatter()
# Lazy LLM initialization
self._llm_initialized = False
@property
def llm(self):
"""Lazy load LLM if not provided."""
if self._llm is not None:
return self._llm
if not self._llm_initialized:
try:
from praisonaiagents.llm import LLM
# Use default model from config or fallback to gpt-4o-mini
model = self.config.model if self.config.model else "gpt-4o-mini"
self._llm = LLM(model=model)
self._llm_initialized = True
except ImportError:
logger.warning("LLM not available, using simple completion")
self._llm = None
self._llm_initialized = True
return self._llm
def _retrieve(
self,
query: str,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> List[Dict[str, Any]]:
"""
Retrieve relevant chunks from Knowledge.
Args:
query: Search query
user_id: Optional user ID for scoping
agent_id: Optional agent ID for scoping
run_id: Optional run ID for scoping
Returns:
List of search results
"""
# Use Knowledge's search method
results = self.knowledge.search(
query=query,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
)
# Handle different result formats
if isinstance(results, dict):
results = results.get("results", [])
# Filter by minimum score
if self.config.min_score > 0:
results = [
r for r in results
if r.get("score", 0) >= self.config.min_score
]
# Limit to top_k
results = results[:self.config.top_k]
return results
def _rerank(
self,
query: str,
results: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""
Rerank results if reranker is configured.
Args:
query: Original query
results: Results to rerank
Returns:
Reranked results
"""
if not self.config.rerank or not self.reranker:
return results
try:
reranked = self.reranker.rerank(
query=query,
results=results,
top_k=self.config.rerank_top_k,
)
return reranked
except Exception as e:
logger.warning(f"Reranking failed: {e}, using original order")
return results
def _build_context(
self,
results: List[Dict[str, Any]],
) -> tuple:
"""
Build context from results.
Returns:
Tuple of (context_string, used_results)
"""
return build_context(
results=results,
max_tokens=self.config.max_context_tokens,
deduplicate=True,
)
def _format_citations(
self,
results: List[Dict[str, Any]],
) -> List[Citation]:
"""Format results as citations."""
if not self.config.include_citations:
return []
return self.citation_formatter.format(results)
def _generate(
self,
prompt: str,
system_prompt: Optional[str] = None,
**kwargs,
) -> str:
"""
Generate answer using LLM.
Args:
prompt: Full prompt with context
system_prompt: Optional system prompt
Returns:
Generated answer
"""
if self.llm is None:
return "LLM not available. Please configure an LLM."
try:
# Try different LLM interfaces
if hasattr(self.llm, 'get_response'):
# PraisonAI LLM interface - prompt can be string or messages list
if system_prompt:
response = self.llm.get_response(
prompt=prompt,
system_prompt=system_prompt,
verbose=False,
stream=False,
**kwargs
)
else:
response = self.llm.get_response(
prompt=prompt,
verbose=False,
stream=False,
**kwargs
)
return response if isinstance(response, str) else str(response)
elif hasattr(self.llm, 'chat'):
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.append({"role": "user", "content": prompt})
response = self.llm.chat(messages=messages, **kwargs)
return response if isinstance(response, str) else str(response)
elif hasattr(self.llm, 'generate'):
return self.llm.generate(prompt, **kwargs)
elif callable(self.llm):
return self.llm(prompt, **kwargs)
else:
return str(self.llm)
except Exception as e:
logger.error(f"LLM generation failed: {e}")
return f"Error generating response: {e}"
def _build_prompt(
self,
question: str,
context: str,
) -> str:
"""Build prompt from template."""
return self.config.template.format(
question=question,
context=context,
)
def retrieve(
self,
query: str,
*,
top_k: Optional[int] = None,
hybrid: Optional[bool] = None,
rerank: Optional[bool] = None,
filters: Optional[Dict[str, Any]] = None,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> ContextPack:
"""
Retrieve context without LLM generation (Orchestrator pattern).
Returns a ContextPack that can be passed to Agent.chat_with_context().
This enables conditional retrieval - only retrieve when needed.
Args:
query: Search query
top_k: Override config top_k
hybrid: Override config hybrid retrieval
rerank: Override config rerank
filters: Optional metadata filters
user_id: Optional user ID for scoping
agent_id: Optional agent ID for scoping
run_id: Optional run ID for scoping
Returns:
ContextPack with context string and citations (no LLM call)
"""
start_time = time.time()
# Apply overrides
original_top_k = self.config.top_k
original_rerank = self.config.rerank
if top_k is not None:
self.config.top_k = top_k
if rerank is not None:
self.config.rerank = rerank
try:
# 1. Retrieve
results = self._retrieve(
query=query,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
filters=filters,
**kwargs,
)
# 2. Rerank (optional)
results = self._rerank(query, results)
# 3. Build context
context, used_results = self._build_context(results)
# 4. Format citations
citations = self._format_citations(used_results)
elapsed = time.time() - start_time
return ContextPack(
context=context,
citations=citations,
query=query,
metadata={
"elapsed_seconds": elapsed,
"num_results": len(results),
"num_citations": len(citations),
"retrieval_strategy": self.config.retrieval_strategy.value,
"reranked": self.config.rerank,
"top_k": self.config.top_k,
},
)
finally:
# Restore original config
self.config.top_k = original_top_k
self.config.rerank = original_rerank
async def aretrieve(
self,
query: str,
*,
top_k: Optional[int] = None,
hybrid: Optional[bool] = None,
rerank: Optional[bool] = None,
filters: Optional[Dict[str, Any]] = None,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> ContextPack:
"""
Async version of retrieve.
Currently wraps sync version. Can be extended for true async.
"""
import asyncio
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
lambda: self.retrieve(
query=query,
top_k=top_k,
hybrid=hybrid,
rerank=rerank,
filters=filters,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
),
)
def query(
self,
question: str,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> RAGResult:
"""
Execute RAG query.
Args:
question: User's question
user_id: Optional user ID for scoping
agent_id: Optional agent ID for scoping
run_id: Optional run ID for scoping
Returns:
RAGResult with answer and citations
"""
start_time = time.time()
# 1. Retrieve
results = self._retrieve(
query=question,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
)
# 2. Rerank (optional)
results = self._rerank(question, results)
# 3. Build context
context, used_results = self._build_context(results)
# 4. Generate answer
prompt = self._build_prompt(question, context)
answer = self._generate(
prompt=prompt,
system_prompt=self.config.system_prompt,
)
# 5. Format citations
citations = self._format_citations(used_results)
# Build result
elapsed = time.time() - start_time
return RAGResult(
answer=answer,
citations=citations,
context_used=context,
query=question,
metadata={
"elapsed_seconds": elapsed,
"num_results": len(results),
"num_citations": len(citations),
"retrieval_strategy": self.config.retrieval_strategy.value,
"reranked": self.config.rerank,
},
)
async def aquery(
self,
question: str,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> RAGResult:
"""
Async version of query.
Currently wraps sync version. Can be extended for true async.
"""
# For now, wrap sync version
# TODO: Implement true async when Knowledge supports it
import asyncio
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
lambda: self.query(
question=question,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
),
)
def stream(
self,
question: str,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> Iterator[str]:
"""
Stream RAG response.
Yields tokens as they are generated.
Final yield includes citation metadata.
Args:
question: User's question
user_id: Optional user ID
agent_id: Optional agent ID
run_id: Optional run ID
Yields:
Response tokens
"""
# 1. Retrieve
results = self._retrieve(
query=question,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
)
# 2. Rerank
results = self._rerank(question, results)
# 3. Build context
context, used_results = self._build_context(results)
# 4. Build prompt
prompt = self._build_prompt(question, context)
# 5. Stream generation
if self.llm is None:
yield "LLM not available."
return
try:
# Try streaming interface
if hasattr(self.llm, 'stream'):
messages = []
if self.config.system_prompt:
messages.append({"role": "system", "content": self.config.system_prompt})
messages.append({"role": "user", "content": prompt})
for chunk in self.llm.stream(messages=messages, **kwargs):
yield chunk
else:
# Fallback to non-streaming
answer = self._generate(prompt, self.config.system_prompt)
# Simulate streaming by yielding words
for word in answer.split():
yield word + " "
except Exception as e:
logger.error(f"Streaming failed: {e}")
yield f"Error: {e}"
async def astream(
self,
question: str,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> AsyncIterator[str]:
"""
Async streaming RAG response.
Yields tokens as they are generated.
"""
# 1. Retrieve
results = self._retrieve(
query=question,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
)
# 2. Rerank
results = self._rerank(question, results)
# 3. Build context
context, used_results = self._build_context(results)
# 4. Build prompt
prompt = self._build_prompt(question, context)
# 5. Stream generation
if self.llm is None:
yield "LLM not available."
return
try:
# Try async streaming interface
if hasattr(self.llm, 'astream'):
messages = []
if self.config.system_prompt:
messages.append({"role": "system", "content": self.config.system_prompt})
messages.append({"role": "user", "content": prompt})
async for chunk in self.llm.astream(messages=messages, **kwargs):
yield chunk
else:
# Fallback to sync streaming wrapped
for chunk in self.stream(question, user_id, agent_id, run_id, **kwargs):
yield chunk
except Exception as e:
logger.error(f"Async streaming failed: {e}")
yield f"Error: {e}"
def get_citations(
self,
question: str,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
**kwargs,
) -> List[Citation]:
"""
Get citations without generating answer.
Useful for retrieval-only workflows.
"""
results = self._retrieve(
query=question,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
**kwargs,
)
results = self._rerank(question, results)
return self._format_citations(results)
@@ -0,0 +1,147 @@
"""
RAG Protocols for PraisonAI Agents.
Defines minimal protocol interfaces for RAG components.
No heavy imports - only stdlib and typing.
"""
from typing import Any, Dict, Iterator, List, Optional, Protocol, runtime_checkable
from .models import Citation, RAGResult
@runtime_checkable
class ContextBuilderProtocol(Protocol):
"""
Protocol for building context from retrieval results.
Implementations must provide a build method that assembles
context from search results with token-aware truncation.
"""
def build(
self,
results: List[Dict[str, Any]],
max_tokens: int = 4000,
deduplicate: bool = True,
) -> str:
"""
Build context string from retrieval results.
Args:
results: List of search results with 'text' and 'metadata' keys
max_tokens: Maximum tokens for context
deduplicate: Whether to deduplicate similar chunks
Returns:
Assembled context string
"""
...
@runtime_checkable
class CitationFormatterProtocol(Protocol):
"""
Protocol for formatting citations from retrieval results.
Implementations must provide a format method that creates
Citation objects from search results.
"""
def format(
self,
results: List[Dict[str, Any]],
start_id: int = 1,
) -> List[Citation]:
"""
Format retrieval results as citations.
Args:
results: List of search results
start_id: Starting citation ID number
Returns:
List of Citation objects
"""
...
@runtime_checkable
class RerankerProtocol(Protocol):
"""
Protocol for reranking retrieval results.
Implementations must provide a rerank method that reorders
results by relevance to the query.
"""
name: str
def rerank(
self,
query: str,
results: List[Dict[str, Any]],
top_k: Optional[int] = None,
) -> List[Dict[str, Any]]:
"""
Rerank results by relevance to query.
Args:
query: The search query
results: List of results to rerank
top_k: Optional limit on results
Returns:
Reranked list of results
"""
...
@runtime_checkable
class RAGProtocol(Protocol):
"""
Protocol for RAG pipelines.
Implementations must provide query, stream, and async variants.
"""
def query(self, question: str, **kwargs) -> RAGResult:
"""Execute RAG query and return result with citations."""
...
async def aquery(self, question: str, **kwargs) -> RAGResult:
"""Async version of query."""
...
def stream(self, question: str, **kwargs) -> Iterator[str]:
"""Stream RAG response tokens."""
...
@runtime_checkable
class GraphHookProtocol(Protocol):
"""
Protocol for graph-based retrieval hooks (optional).
Implementations can provide graph traversal for enhanced retrieval.
This is an optional extension point for GraphRAG tools.
"""
def get_related_nodes(
self,
query: str,
node_ids: List[str],
max_hops: int = 2,
) -> List[Dict[str, Any]]:
"""
Get related nodes from graph for context expansion.
Args:
query: The search query
node_ids: Starting node IDs
max_hops: Maximum graph traversal hops
Returns:
List of related nodes with metadata
"""
...
@@ -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,518 @@
"""
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
# Note: We don't remove the lock file to avoid race conditions
# where another process has opened but not yet locked the file
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.3"
version = "0.7.0"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
readme = "README.md"
requires-python = ">=3.10"
+46
View File
@@ -0,0 +1,46 @@
"""
Pytest configuration for PraisonAI Agents tests.
Provides fixtures and markers for testing, including:
- live: Tests that require real API keys (opt-in)
- slow: Tests that take longer to run
"""
import os
import pytest
def pytest_configure(config):
"""Register custom markers."""
config.addinivalue_line(
"markers", "live: mark test as requiring real API keys (deselect with '-m \"not live\"')"
)
config.addinivalue_line(
"markers", "slow: mark test as slow running"
)
def pytest_collection_modifyitems(config, items):
"""Skip live tests unless PRAISONAI_LIVE_TESTS=1 is set."""
if os.environ.get("PRAISONAI_LIVE_TESTS") != "1":
skip_live = pytest.mark.skip(reason="Live tests disabled. Set PRAISONAI_LIVE_TESTS=1 to enable.")
for item in items:
if "live" in item.keywords:
item.add_marker(skip_live)
@pytest.fixture
def openai_api_key():
"""Get OpenAI API key from environment, skip if not available."""
key = os.environ.get("OPENAI_API_KEY")
if not key:
pytest.skip("OPENAI_API_KEY not set")
return key
@pytest.fixture
def live_test_enabled():
"""Check if live tests are enabled."""
if os.environ.get("PRAISONAI_LIVE_TESTS") != "1":
pytest.skip("Live tests disabled. Set PRAISONAI_LIVE_TESTS=1 to enable.")
return True
Whitespace-only changes.
@@ -0,0 +1,230 @@
"""
Live integration tests for Agent + RAG integration.
These tests require:
- PRAISONAI_LIVE_TESTS=1 environment variable
- OPENAI_API_KEY environment variable
Run with: PRAISONAI_LIVE_TESTS=1 pytest -m live tests/integration/agent/ -v
"""
import pytest
import tempfile
from pathlib import Path
@pytest.fixture
def sample_documents():
"""Create sample documents for testing."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create sample documents
doc1 = Path(tmpdir) / "doc1.txt"
doc1.write_text("""
The capital of France is Paris.
Paris is known for the Eiffel Tower.
France is a country in Western Europe.
""")
doc2 = Path(tmpdir) / "doc2.txt"
doc2.write_text("""
Python was created by Guido van Rossum.
Python is a high-level programming language.
Python emphasizes code readability.
""")
doc3 = Path(tmpdir) / "doc3.txt"
doc3.write_text("""
Machine learning is a subset of artificial intelligence.
Deep learning uses neural networks with many layers.
AI is transforming many industries.
""")
yield tmpdir
@pytest.fixture
def openai_api_key():
"""Get OpenAI API key from environment."""
import os
key = os.environ.get("OPENAI_API_KEY")
if not key:
pytest.skip("OPENAI_API_KEY not set")
return key
@pytest.mark.live
class TestAgentRAGLiveQuery:
"""Live tests for Agent RAG query functionality."""
def test_agent_rag_query_with_citations(self, openai_api_key, sample_documents):
"""Test Agent.rag_query returns answer with citations."""
from praisonaiagents import Agent
from praisonaiagents.knowledge import Knowledge
# Create knowledge base
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_agent_live",
"path": f"{sample_documents}/.chroma",
}
}
}
# Pre-index documents
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_agent"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
# Create agent with knowledge and rag_config
agent = Agent(
name="TestAgent",
instructions="You are a helpful assistant that answers questions based on provided knowledge.",
knowledge=[str(doc) for doc in Path(sample_documents).glob("*.txt")],
knowledge_config=knowledge_config,
rag_config={"include_citations": True, "top_k": 3},
user_id=user_id,
)
# Use agent.rag_query for RAG with citations
result = agent.rag_query("What is the capital of France?")
# Assertions
assert result.answer is not None
assert len(result.answer) > 0
assert "paris" in result.answer.lower() or "france" in result.answer.lower()
assert result.citations is not None
assert len(result.citations) > 0
def test_agent_rag_property_access(self, openai_api_key, sample_documents):
"""Test Agent.rag property provides RAG instance."""
from praisonaiagents import Agent
from praisonaiagents.knowledge import Knowledge
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_agent_rag_prop",
"path": f"{sample_documents}/.chroma_prop",
}
}
}
# Pre-index
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_prop"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=[str(doc) for doc in Path(sample_documents).glob("*.txt")],
knowledge_config=knowledge_config,
rag_config={"include_citations": True},
user_id=user_id,
)
# Access RAG through property
rag = agent.rag
assert rag is not None
# Query through RAG directly
result = rag.query("Who created Python?", user_id=user_id)
assert "guido" in result.answer.lower() or "rossum" in result.answer.lower()
@pytest.mark.live
class TestAgentRAGLiveStreaming:
"""Live tests for Agent RAG streaming functionality."""
def test_agent_rag_stream(self, openai_api_key, sample_documents):
"""Test Agent.rag.stream produces chunks."""
from praisonaiagents import Agent
from praisonaiagents.knowledge import Knowledge
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_agent_stream",
"path": f"{sample_documents}/.chroma_stream",
}
}
}
# Pre-index
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_stream"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=[str(doc) for doc in Path(sample_documents).glob("*.txt")],
knowledge_config=knowledge_config,
rag_config={"stream": True},
user_id=user_id,
)
# Stream through RAG
chunks = []
for chunk in agent.rag.stream("Tell me about machine learning", user_id=user_id):
chunks.append(chunk)
# Should have multiple chunks
assert len(chunks) >= 1
# Combined response should be meaningful
full_response = "".join(chunks)
assert len(full_response) > 10
@pytest.mark.live
class TestAgentRAGLiveCitations:
"""Live tests for Agent RAG citation functionality."""
def test_agent_citations_have_source_info(self, openai_api_key, sample_documents):
"""Test that Agent RAG citations have proper source information."""
from praisonaiagents import Agent
from praisonaiagents.knowledge import Knowledge
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_agent_citations",
"path": f"{sample_documents}/.chroma_citations",
}
}
}
# Pre-index
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_citations"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=[str(doc) for doc in Path(sample_documents).glob("*.txt")],
knowledge_config=knowledge_config,
rag_config={"include_citations": True, "top_k": 5},
user_id=user_id,
)
result = agent.rag_query("What is the capital of France?")
# Check citations
assert result.has_citations
# Citations should have scores
for citation in result.citations:
assert citation.id is not None
assert citation.text is not None
assert len(citation.text) > 0
@@ -0,0 +1 @@
"""RAG integration tests."""
@@ -0,0 +1,241 @@
"""
Live integration tests for RAG with real API keys.
These tests require:
- PRAISONAI_LIVE_TESTS=1
- OPENAI_API_KEY set
Run with:
PRAISONAI_LIVE_TESTS=1 pytest -m live tests/integration/rag/ -v
"""
import pytest
import tempfile
from pathlib import Path
@pytest.fixture
def sample_documents():
"""Create sample documents for testing."""
with tempfile.TemporaryDirectory() as tmpdir:
doc1 = Path(tmpdir) / "doc1.txt"
doc1.write_text("""
The capital of France is Paris.
Paris is known for the Eiffel Tower.
France is a country in Western Europe.
""")
doc2 = Path(tmpdir) / "doc2.txt"
doc2.write_text("""
Python is a programming language.
Python was created by Guido van Rossum.
Python is widely used for data science and AI.
""")
doc3 = Path(tmpdir) / "doc3.txt"
doc3.write_text("""
Machine learning is a subset of artificial intelligence.
Deep learning uses neural networks.
AI is transforming many industries.
""")
yield tmpdir
@pytest.mark.live
class TestRAGLiveQuery:
"""Live tests for RAG query functionality."""
def test_rag_query_with_citations(self, openai_api_key, sample_documents):
"""Test RAG query returns answer with citations."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.rag import RAG, RAGConfig
# Create knowledge base
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_live",
"path": f"{sample_documents}/.chroma",
}
}
}
knowledge = Knowledge(config=knowledge_config)
# Index documents with user_id (required by mem0)
user_id = "test_user"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
# Create RAG
rag = RAG(
knowledge=knowledge,
config=RAGConfig(top_k=3, include_citations=True)
)
# Query with user_id
result = rag.query("What is the capital of France?", user_id=user_id)
# Assertions
assert result.answer is not None
assert len(result.answer) > 0
assert "paris" in result.answer.lower() or "france" in result.answer.lower()
assert result.citations is not None
assert len(result.citations) > 0
def test_rag_query_python_topic(self, openai_api_key, sample_documents):
"""Test RAG query on Python topic."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.rag import RAG, RAGConfig
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_python",
"path": f"{sample_documents}/.chroma_python",
}
}
}
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_python"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
rag = RAG(
knowledge=knowledge,
config=RAGConfig(top_k=3, include_citations=True)
)
result = rag.query("Who created Python?", user_id=user_id)
assert result.answer is not None
assert "guido" in result.answer.lower() or "rossum" in result.answer.lower()
@pytest.mark.live
class TestRAGLiveStreaming:
"""Live tests for RAG streaming functionality."""
def test_rag_stream_produces_chunks(self, openai_api_key, sample_documents):
"""Test RAG streaming produces multiple chunks."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.rag import RAG, RAGConfig
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_stream",
"path": f"{sample_documents}/.chroma_stream",
}
}
}
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_stream"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
rag = RAG(
knowledge=knowledge,
config=RAGConfig(top_k=3, stream=True)
)
# Collect streamed chunks
chunks = []
for chunk in rag.stream("Tell me about machine learning", user_id=user_id):
chunks.append(chunk)
# Should have multiple chunks
assert len(chunks) >= 1
# Combined response should be meaningful
full_response = "".join(chunks)
assert len(full_response) > 10
@pytest.mark.live
class TestRAGLiveCitations:
"""Live tests for RAG citation functionality."""
def test_citations_reference_correct_sources(self, openai_api_key, sample_documents):
"""Test that citations reference the correct source documents."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.rag import RAG, RAGConfig
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_citations",
"path": f"{sample_documents}/.chroma_citations",
}
}
}
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_citations"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
rag = RAG(
knowledge=knowledge,
config=RAGConfig(top_k=5, include_citations=True)
)
result = rag.query("What is the capital of France?", user_id=user_id)
# Check citations
assert result.has_citations
# At least one citation should be from doc1 (France content)
citation_sources = [c.source for c in result.citations]
assert len(citation_sources) > 0
# Citations should have scores
for citation in result.citations:
assert citation.score is not None or citation.score >= 0
def test_get_citations_without_generation(self, openai_api_key, sample_documents):
"""Test getting citations without generating an answer."""
from praisonaiagents.knowledge import Knowledge
from praisonaiagents.rag import RAG, RAGConfig
knowledge_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_citations_only",
"path": f"{sample_documents}/.chroma_citations_only",
}
}
}
knowledge = Knowledge(config=knowledge_config)
user_id = "test_user_citations_only"
for doc in Path(sample_documents).glob("*.txt"):
knowledge.add(str(doc), user_id=user_id)
rag = RAG(
knowledge=knowledge,
config=RAGConfig(top_k=3)
)
citations = rag.get_citations("Python programming", user_id=user_id)
assert citations is not None
assert len(citations) > 0
# Check citation structure
for citation in citations:
assert citation.id is not None
assert citation.text is not None
assert len(citation.text) > 0
@@ -0,0 +1,187 @@
"""
Tests for Agent RAG integration.
Tests that Agent can access RAG features through the knowledge parameter.
"""
import pytest
from unittest.mock import MagicMock, patch
class TestAgentRAGConfig:
"""Tests for Agent rag_config parameter."""
def test_agent_accepts_rag_config(self):
"""Test that Agent accepts rag_config parameter."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=["test.txt"],
rag_config={"include_citations": True, "top_k": 5}
)
assert agent._rag_config == {"include_citations": True, "top_k": 5}
def test_agent_without_rag_config(self):
"""Test that Agent works without rag_config."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=["test.txt"]
)
assert agent._rag_config is None
def test_agent_without_knowledge_has_no_rag_config(self):
"""Test that Agent without knowledge has no rag_config."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent"
)
assert agent._rag_config is None
assert agent._rag_instance is None
class TestAgentRAGProperty:
"""Tests for Agent.rag property."""
def test_agent_has_rag_property(self):
"""Test that Agent has rag property."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=["test.txt"],
rag_config={"include_citations": True}
)
assert hasattr(agent, 'rag')
def test_agent_rag_returns_none_without_knowledge(self):
"""Test that Agent.rag returns None without knowledge."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent"
)
assert agent.rag is None
@patch('praisonaiagents.knowledge.Knowledge')
@patch('praisonaiagents.rag.RAG')
@patch('praisonaiagents.rag.RAGConfig')
def test_agent_rag_lazy_loads(self, mock_rag_config, mock_rag, mock_knowledge):
"""Test that Agent.rag lazy loads RAG instance."""
from praisonaiagents import Agent
mock_knowledge_instance = MagicMock()
mock_knowledge.return_value = mock_knowledge_instance
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=["test.txt"],
rag_config={"include_citations": True}
)
# RAG should not be loaded yet
assert agent._rag_instance is None
# Access rag property - this triggers lazy loading
# Note: In real usage, knowledge processing would happen first
class TestAgentRAGQuery:
"""Tests for Agent.rag_query method."""
def test_agent_has_rag_query_method(self):
"""Test that Agent has rag_query method."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=["test.txt"],
rag_config={"include_citations": True}
)
assert hasattr(agent, 'rag_query')
assert callable(agent.rag_query)
def test_rag_query_raises_without_knowledge(self):
"""Test that rag_query raises error without knowledge."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent"
)
with pytest.raises(ValueError, match="No knowledge sources configured"):
agent.rag_query("What is the answer?")
class TestAgentKnowledgeContext:
"""Tests for Agent._get_knowledge_context method."""
def test_agent_has_get_knowledge_context_method(self):
"""Test that Agent has _get_knowledge_context method."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=["test.txt"]
)
assert hasattr(agent, '_get_knowledge_context')
assert callable(agent._get_knowledge_context)
def test_get_knowledge_context_returns_tuple(self):
"""Test that _get_knowledge_context returns tuple."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent"
)
result = agent._get_knowledge_context("test query")
assert isinstance(result, tuple)
assert len(result) == 2
assert result[0] == "" # No knowledge, empty context
assert result[1] is None # No citations
class TestAgentRAGIntegration:
"""Integration tests for Agent + RAG."""
def test_agent_rag_config_passed_to_rag_instance(self):
"""Test that rag_config is passed to RAG instance."""
from praisonaiagents import Agent
agent = Agent(
name="TestAgent",
instructions="Test agent",
knowledge=["test.txt"],
rag_config={
"include_citations": True,
"top_k": 10,
"min_score": 0.5
}
)
# Verify config is stored
assert agent._rag_config["include_citations"] is True
assert agent._rag_config["top_k"] == 10
assert agent._rag_config["min_score"] == 0.5
@@ -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
Loaded 100 of 263 files, more files were not shown because too many files have changed in this diff. Show more