mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 01:08:39 -10:00
Compare commits
55
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
17c5942a02 | ||
|
|
951eb6b497 | ||
|
|
1d0f08b20a | ||
|
|
07ff93ed7f | ||
|
|
2d5c982893 | ||
|
|
f5d4db4262 | ||
|
|
87d691c616 | ||
|
|
8ea15dee0f | ||
|
|
4f56842e08 | ||
|
|
a24836aaee | ||
|
|
be9ce27615 | ||
|
|
0159a1e082 | ||
|
|
71a5076c07 | ||
|
|
9815a9f8a8 | ||
|
|
2ae44a05b0 | ||
|
|
3386d25022 | ||
|
|
6dca2ace8a | ||
|
|
63548739e8 | ||
|
|
3e046537f4 | ||
|
|
4a34677567 | ||
|
|
19e91cb669 | ||
|
|
5a6125021a | ||
|
|
da26887f24 | ||
|
|
1b2f7139f9 | ||
|
|
aafbc86180 | ||
|
|
3d4ca1b7c0 | ||
|
|
1180aa38bd | ||
|
|
65f65e5c7b | ||
|
|
f81a12d5e8 | ||
|
|
cf6f5fbefe | ||
|
|
908d7710db | ||
|
|
7d5b41e31f | ||
|
|
e46a0bc9f3 | ||
|
|
911483c92c | ||
|
|
4ffb9d6e3f | ||
|
|
b686e5b00f | ||
|
|
d8d84ee3e3 | ||
|
|
9f2c4d10b5 | ||
|
|
42a86c831f | ||
|
|
48443b11d6 | ||
|
|
71745aa14a | ||
|
|
4f309ec16d | ||
|
|
e48afc8993 | ||
|
|
90950f8889 | ||
|
|
56f36e2593 | ||
|
|
7aa67ad27c | ||
|
|
307df443ef | ||
|
|
6434a87b60 | ||
|
|
97d5b6db0d | ||
|
|
4996a9f9fa | ||
|
|
54f4a96cba | ||
|
|
e59be8f0fd | ||
|
|
b52708f3f6 | ||
|
|
ba37ce39d3 | ||
|
|
98c5e8f175 |
No files matched your search
@@ -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
|
||||
@@ -3079,43 +3047,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 +3123,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 | [](https://www.youtube.com/watch?v=vLXobEN2Vc8) |
|
||||
@@ -3216,6 +3151,8 @@ Learn PraisonAI through our comprehensive video series:
|
||||
| Call Interface | [](https://www.youtube.com/watch?v=m1cwrUG2iAk) |
|
||||
| Reasoning Extract Agents | [](https://www.youtube.com/watch?v=2PPamsADjJA) |
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
## 👥 Contributing
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
Auto-generated API server for PraisonAI agents.
|
||||
"""
|
||||
from flask import Flask, request, jsonify
|
||||
from flask_cors import CORS
|
||||
from praisonai import PraisonAI
|
||||
import os
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# CORS configuration
|
||||
CORS(app)
|
||||
|
||||
# Authentication
|
||||
AUTH_ENABLED = False
|
||||
AUTH_TOKEN = None
|
||||
|
||||
def check_auth():
|
||||
"""Check authentication if enabled."""
|
||||
if not AUTH_ENABLED:
|
||||
return True
|
||||
|
||||
token = request.headers.get('Authorization', '').replace('Bearer ', '')
|
||||
return token == AUTH_TOKEN
|
||||
|
||||
@app.route('/health', methods=['GET'])
|
||||
def health():
|
||||
"""Health check endpoint."""
|
||||
return jsonify({"status": "ok", "service": "praisonai-api"})
|
||||
|
||||
@app.route('/chat', methods=['POST'])
|
||||
def chat():
|
||||
"""Chat endpoint for agent interaction."""
|
||||
if not check_auth():
|
||||
return jsonify({"error": "Unauthorized"}), 401
|
||||
|
||||
data = request.get_json()
|
||||
if not data or 'message' not in data:
|
||||
return jsonify({"error": "Message required"}), 400
|
||||
|
||||
try:
|
||||
praisonai = PraisonAI(agent_file="agents.yaml")
|
||||
result = praisonai.run()
|
||||
|
||||
return jsonify({
|
||||
"response": result,
|
||||
"status": "success"
|
||||
})
|
||||
except Exception as e:
|
||||
return jsonify({
|
||||
"error": str(e),
|
||||
"status": "error"
|
||||
}), 500
|
||||
|
||||
@app.route('/agents', methods=['GET'])
|
||||
def list_agents():
|
||||
"""List available agents."""
|
||||
if not check_auth():
|
||||
return jsonify({"error": "Unauthorized"}), 401
|
||||
|
||||
return jsonify({
|
||||
"agents": ["default"],
|
||||
"agent_file": "agents.yaml"
|
||||
})
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(
|
||||
host='0.0.0.0',
|
||||
port=8080,
|
||||
debug=False
|
||||
)
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.7.9" \
|
||||
"praisonai>=3.1.7" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.7.9" \
|
||||
"praisonai>=3.1.7" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.7.9" \
|
||||
"praisonai>=3.1.7" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# Profiling CLI
|
||||
|
||||
PraisonAI provides unified profiling capabilities for performance analysis and debugging, with **full visibility from ENTER to first response**.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Basic profiling
|
||||
praisonai "What is 2+2?" --profile
|
||||
|
||||
# Deep profiling with call graph, decision trace, and module breakdown
|
||||
praisonai "What is 2+2?" --profile --profile-deep
|
||||
|
||||
# JSON output format
|
||||
praisonai "What is 2+2?" --profile --profile-format json
|
||||
```
|
||||
|
||||
## Time to First Response (TTFR)
|
||||
|
||||
**TTFR** is the key user-perceived latency metric. It measures the time from when you press ENTER until the first visible output appears.
|
||||
|
||||
### What contributes to TTFR?
|
||||
|
||||
1. **CLI Entry** - Time to reach the CLI entry point
|
||||
2. **CLI Parse** - Argparse/Typer routing time
|
||||
3. **Routing** - Command routing to handler
|
||||
4. **Imports** - Loading required modules
|
||||
5. **Agent Init** - Initializing the agent
|
||||
6. **Network Request** - Sending request to LLM API
|
||||
7. **First Token** - Receiving first token (streaming)
|
||||
8. **First Output** - Rendering first visible output
|
||||
|
||||
### How to reduce TTFR
|
||||
|
||||
1. **Warm runs** - Imports are cached after first run
|
||||
2. **Streaming mode** - See output as tokens arrive
|
||||
3. **Lighter models** - Faster response times
|
||||
4. **Local models** - Eliminate network latency
|
||||
|
||||
## CLI Flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--profile` | Enable profiling with timing breakdown | Disabled |
|
||||
| `--profile-deep` | Enable deep profiling with call graph (higher overhead) | Disabled |
|
||||
| `--profile-format` | Output format: `text` or `json` | `text` |
|
||||
|
||||
## Profile Command
|
||||
|
||||
For more detailed profiling options, use the `profile` subcommand:
|
||||
|
||||
```bash
|
||||
# Profile a query
|
||||
praisonai profile query "What is 2+2?"
|
||||
|
||||
# With deep call tracing
|
||||
praisonai profile query "Hello" --deep
|
||||
|
||||
# Limit functions shown
|
||||
praisonai profile query "Test" --limit 20
|
||||
|
||||
# Save artifacts to directory
|
||||
praisonai profile query "Test" --save ./profile_results
|
||||
|
||||
# JSON output
|
||||
praisonai profile query "Test" --format json
|
||||
```
|
||||
|
||||
### Profile Query Options
|
||||
|
||||
| Option | Description | Default |
|
||||
|--------|-------------|---------|
|
||||
| `--model, -m` | Model to use | Default model |
|
||||
| `--stream/--no-stream` | Use streaming mode | No stream |
|
||||
| `--deep` | Enable deep call tracing | Disabled |
|
||||
| `--limit, -n` | Top N functions to show | 30 |
|
||||
| `--sort, -s` | Sort by: `cumulative` or `tottime` | `cumulative` |
|
||||
| `--show-callers` | Show caller functions | Disabled |
|
||||
| `--show-callees` | Show callee functions | Disabled |
|
||||
| `--save` | Save artifacts to path | None |
|
||||
| `--format, -f` | Output format: `text` or `json` | `text` |
|
||||
|
||||
## Other Profile Commands
|
||||
|
||||
```bash
|
||||
# Profile module import times
|
||||
praisonai profile imports
|
||||
|
||||
# Profile CLI startup time
|
||||
praisonai profile startup
|
||||
|
||||
# Run comprehensive profiling suite
|
||||
praisonai profile suite
|
||||
|
||||
# Create performance snapshots for comparison
|
||||
praisonai profile snapshot --baseline
|
||||
praisonai profile snapshot current --compare
|
||||
|
||||
# Configure performance optimizations
|
||||
praisonai profile optimize --show
|
||||
```
|
||||
|
||||
## Output Format
|
||||
|
||||
### Text Report
|
||||
|
||||
```
|
||||
======================================================================
|
||||
PraisonAI Profile Report
|
||||
======================================================================
|
||||
|
||||
Run ID: abc12345
|
||||
Timestamp: 2024-01-02T12:00:00.000000Z
|
||||
Method: cli_direct
|
||||
Version: 3.0.1
|
||||
|
||||
## Execution Timeline
|
||||
---------------------------------------------
|
||||
Imports : 500.00 ms
|
||||
Agent Init : 0.50 ms
|
||||
Execution : 2000.00 ms
|
||||
───────────────────────────────────────────
|
||||
⏱ Time to First Response : 2500.50 ms
|
||||
TOTAL : 2500.50 ms
|
||||
|
||||
## Timing Breakdown
|
||||
----------------------------------------
|
||||
Imports: 500.00 ms
|
||||
Agent Init: 0.50 ms
|
||||
Execution: 2000.00 ms
|
||||
─────────────────────────────────────
|
||||
TOTAL: 2500.50 ms
|
||||
|
||||
## Decision Trace (deep profile only)
|
||||
----------------------------------------
|
||||
Agent Config: default
|
||||
Model: gpt-4
|
||||
Streaming: disabled
|
||||
Profile Layer: 2
|
||||
|
||||
## Top Functions by Cumulative Time
|
||||
----------------------------------------------------------------------
|
||||
Function Calls Cumul (ms)
|
||||
----------------------------------------------------------------------
|
||||
request 1 1800.00
|
||||
send 1 1750.00
|
||||
...
|
||||
```
|
||||
|
||||
### JSON Report
|
||||
|
||||
```json
|
||||
{
|
||||
"schema_version": "1.0",
|
||||
"run_id": "abc12345",
|
||||
"timestamp": "2024-01-02T12:00:00.000000Z",
|
||||
"invocation": {
|
||||
"method": "cli_direct",
|
||||
"flags": {"layer": 1, "stream": false, "model": null},
|
||||
"praisonai_version": "3.0.1",
|
||||
"python_version": "3.12.0"
|
||||
},
|
||||
"timing": {
|
||||
"total_ms": 2500.5,
|
||||
"imports_ms": 500.0,
|
||||
"agent_init_ms": 0.5,
|
||||
"execution_ms": 2000.0
|
||||
},
|
||||
"functions": [...],
|
||||
"response_preview": "four"
|
||||
}
|
||||
```
|
||||
|
||||
## Profiling Layers
|
||||
|
||||
| Layer | Overhead | Features |
|
||||
|-------|----------|----------|
|
||||
| 0 | < 1ms | Wall-clock phases only |
|
||||
| 1 | ~5% | + cProfile function stats |
|
||||
| 2 | ~15% | + Call graph (callers/callees) |
|
||||
|
||||
- `--profile` uses Layer 1
|
||||
- `--profile --profile-deep` uses Layer 2
|
||||
|
||||
## Performance Tips
|
||||
|
||||
1. **Use Layer 0 for production monitoring** - minimal overhead
|
||||
2. **Use Layer 1 for development profiling** - good balance
|
||||
3. **Use Layer 2 only for deep debugging** - higher overhead
|
||||
|
||||
## Observational Equivalence
|
||||
|
||||
The following commands produce equivalent profile reports:
|
||||
|
||||
```bash
|
||||
# CLI direct
|
||||
praisonai "prompt" --profile
|
||||
|
||||
# Profile command
|
||||
praisonai profile query "prompt"
|
||||
```
|
||||
|
||||
Both use the same unified profiling architecture and produce identical schema output.
|
||||
@@ -0,0 +1,257 @@
|
||||
# Profiling Module
|
||||
|
||||
PraisonAI provides a unified profiling module for programmatic performance analysis.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import (
|
||||
ExecutionRequest,
|
||||
Profiler,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
# Create execution request
|
||||
request = ExecutionRequest(prompt="What is 2+2?")
|
||||
|
||||
# Configure profiler
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
# Run with profiling
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request)
|
||||
|
||||
# Access results
|
||||
print(f"Output: {result.output}")
|
||||
print(f"Total time: {report.timing.total_ms}ms")
|
||||
```
|
||||
|
||||
## ExecutionRequest
|
||||
|
||||
Immutable request descriptor for execution.
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import ExecutionRequest
|
||||
|
||||
# Basic request
|
||||
request = ExecutionRequest(prompt="Hello")
|
||||
|
||||
# With model
|
||||
request = ExecutionRequest(
|
||||
prompt="Hello",
|
||||
model="gpt-4",
|
||||
agent_name="MyAgent",
|
||||
)
|
||||
|
||||
# With streaming
|
||||
request = ExecutionRequest(
|
||||
prompt="Hello",
|
||||
stream=True,
|
||||
)
|
||||
```
|
||||
|
||||
### Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `prompt` | str | Required | The prompt to execute |
|
||||
| `agent_name` | str | "Assistant" | Name of the agent |
|
||||
| `agent_instructions` | str | None | Custom instructions |
|
||||
| `model` | str | None | Model to use |
|
||||
| `stream` | bool | False | Enable streaming |
|
||||
| `tools` | tuple | () | Tools to provide |
|
||||
|
||||
## ProfilerConfig
|
||||
|
||||
Configuration for profiling behavior.
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import ProfilerConfig
|
||||
|
||||
# Layer 0: Wall-clock only (< 1ms overhead)
|
||||
config = ProfilerConfig(layer=0)
|
||||
|
||||
# Layer 1: + cProfile stats (~5% overhead)
|
||||
config = ProfilerConfig(layer=1)
|
||||
|
||||
# Layer 2: + Call graph (~15% overhead)
|
||||
config = ProfilerConfig(
|
||||
layer=2,
|
||||
show_callers=True,
|
||||
show_callees=True,
|
||||
)
|
||||
|
||||
# From CLI flags
|
||||
config = ProfilerConfig.from_flags(
|
||||
profile=True,
|
||||
deep=True,
|
||||
network=False,
|
||||
)
|
||||
```
|
||||
|
||||
### Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `layer` | int | 1 | Profiling layer (0, 1, or 2) |
|
||||
| `limit` | int | 30 | Max functions to show |
|
||||
| `sort_by` | str | "cumulative" | Sort by: cumulative, time, calls |
|
||||
| `show_callers` | bool | False | Show caller functions |
|
||||
| `show_callees` | bool | False | Show callee functions |
|
||||
| `track_network` | bool | False | Track network requests |
|
||||
| `output_format` | str | "text" | Output format: text, json |
|
||||
| `save_path` | str | None | Path to save artifacts |
|
||||
|
||||
## Profiler
|
||||
|
||||
The main profiler class that wraps execution.
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import Profiler, ProfilerConfig, ExecutionRequest
|
||||
|
||||
config = ProfilerConfig(layer=1)
|
||||
profiler = Profiler(config)
|
||||
|
||||
request = ExecutionRequest(prompt="Hello")
|
||||
result, report = profiler.profile_sync(request, invocation_method="api")
|
||||
```
|
||||
|
||||
### Methods
|
||||
|
||||
#### `profile_sync(request, invocation_method="profile_command")`
|
||||
|
||||
Profile synchronous execution.
|
||||
|
||||
**Parameters:**
|
||||
- `request`: ExecutionRequest - The execution request
|
||||
- `invocation_method`: str - How this was invoked (for report)
|
||||
|
||||
**Returns:** Tuple of (ExecutionResult, ProfileReport)
|
||||
|
||||
## ProfileReport
|
||||
|
||||
The canonical profile report with timing and statistics.
|
||||
|
||||
```python
|
||||
# Access timing breakdown
|
||||
print(f"Total: {report.timing.total_ms}ms")
|
||||
print(f"Imports: {report.timing.imports_ms}ms")
|
||||
print(f"Agent Init: {report.timing.agent_init_ms}ms")
|
||||
print(f"Execution: {report.timing.execution_ms}ms")
|
||||
|
||||
# Access function stats (layer 1+)
|
||||
if report.functions:
|
||||
for func in report.functions[:5]:
|
||||
print(f"{func.name}: {func.cumulative_time_ms}ms")
|
||||
|
||||
# Access call graph (layer 2)
|
||||
if report.call_graph:
|
||||
print(f"Edges: {len(report.call_graph.edges)}")
|
||||
|
||||
# Export formats
|
||||
json_str = report.to_json()
|
||||
text_str = report.to_text()
|
||||
dict_data = report.to_dict()
|
||||
```
|
||||
|
||||
### Schema Version
|
||||
|
||||
The report uses schema version `1.0`. Schema is frozen after release - new fields are always optional with defaults.
|
||||
|
||||
### Mandatory Fields
|
||||
|
||||
- `schema_version`: str
|
||||
- `run_id`: str
|
||||
- `timestamp`: str (ISO 8601)
|
||||
- `invocation`: InvocationInfo
|
||||
- `timing`: TimingBreakdown
|
||||
- `response_preview`: str
|
||||
|
||||
### Optional Fields
|
||||
|
||||
- `functions`: List[FunctionStat] (layer 1+)
|
||||
- `call_graph`: CallGraph (layer 2)
|
||||
- `network`: List[RequestTiming] (if track_network)
|
||||
|
||||
## ExecutionResult
|
||||
|
||||
Result of execution with output and metadata.
|
||||
|
||||
```python
|
||||
# Access output
|
||||
print(result.output)
|
||||
print(result.success)
|
||||
print(result.run_id)
|
||||
|
||||
# Check for errors
|
||||
if result.error:
|
||||
print(f"Error: {result.error}")
|
||||
|
||||
# Duration
|
||||
print(f"Duration: {result.duration_ms}ms")
|
||||
```
|
||||
|
||||
## Data Bounds
|
||||
|
||||
Profile data is strictly bounded to prevent memory issues:
|
||||
|
||||
| Data Type | Maximum |
|
||||
|-----------|---------|
|
||||
| Function stats | 1000 |
|
||||
| Call graph edges | 5000 |
|
||||
| Network requests | 100 |
|
||||
|
||||
## Example: Full Profiling Workflow
|
||||
|
||||
```python
|
||||
from praisonai.cli.execution import (
|
||||
ExecutionRequest,
|
||||
Profiler,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
# Configure deep profiling
|
||||
config = ProfilerConfig(
|
||||
layer=2,
|
||||
limit=50,
|
||||
show_callers=True,
|
||||
show_callees=True,
|
||||
)
|
||||
|
||||
# Create request
|
||||
request = ExecutionRequest(
|
||||
prompt="Explain quantum computing in simple terms",
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
# Profile execution
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request)
|
||||
|
||||
# Print text report
|
||||
print(report.to_text())
|
||||
|
||||
# Save JSON report
|
||||
with open("profile.json", "w") as f:
|
||||
f.write(report.to_json())
|
||||
|
||||
# Analyze timing
|
||||
print(f"\nTiming Analysis:")
|
||||
print(f" Import overhead: {report.timing.imports_ms / report.timing.total_ms * 100:.1f}%")
|
||||
print(f" Execution time: {report.timing.execution_ms / report.timing.total_ms * 100:.1f}%")
|
||||
|
||||
# Top 5 hotspots
|
||||
if report.functions:
|
||||
print(f"\nTop 5 Hotspots:")
|
||||
for func in report.functions[:5]:
|
||||
print(f" {func.name}: {func.cumulative_time_ms:.2f}ms ({func.calls} calls)")
|
||||
```
|
||||
|
||||
## Invariants
|
||||
|
||||
The profiling module guarantees:
|
||||
|
||||
1. **Output Equivalence**: Profiled execution produces identical output to non-profiled execution
|
||||
2. **Schema Stability**: Schema version 1.0 is frozen; new fields are always optional
|
||||
3. **Bounded Data**: All profile data respects maximum bounds
|
||||
4. **Zero Overhead When Disabled**: No profiling code runs unless explicitly enabled
|
||||
@@ -0,0 +1,137 @@
|
||||
# ACP (Agent Client Protocol) Examples
|
||||
|
||||
This directory contains examples for using PraisonAI with the Agent Client Protocol (ACP).
|
||||
|
||||
## What is ACP?
|
||||
|
||||
The Agent Client Protocol (ACP) is a standardized JSON-RPC 2.0 protocol that allows code editors and IDEs to communicate with AI coding agents. It enables seamless integration between PraisonAI and editors like:
|
||||
|
||||
- **Zed**
|
||||
- **JetBrains IDEs** (IntelliJ, PyCharm, WebStorm, etc.)
|
||||
- **VSCode** (via extensions)
|
||||
- **Toad**
|
||||
|
||||
## Quick Start
|
||||
|
||||
### CLI Usage
|
||||
|
||||
```bash
|
||||
# Start ACP server with defaults
|
||||
praisonai acp
|
||||
|
||||
# Enable debug logging
|
||||
praisonai acp --debug
|
||||
|
||||
# Allow file writes
|
||||
praisonai acp --allow-write
|
||||
|
||||
# Use a specific model
|
||||
praisonai acp -m gpt-4o
|
||||
|
||||
# Resume last session
|
||||
praisonai acp --resume --last
|
||||
```
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from praisonai.acp import serve
|
||||
|
||||
serve(
|
||||
workspace=".",
|
||||
debug=True,
|
||||
allow_write=True,
|
||||
)
|
||||
```
|
||||
|
||||
## Editor Configuration
|
||||
|
||||
### Zed
|
||||
|
||||
Add to `~/.config/zed/settings.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"agent_servers": {
|
||||
"PraisonAI": {
|
||||
"command": "praisonai",
|
||||
"args": ["acp"],
|
||||
"env": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### JetBrains
|
||||
|
||||
Add to `~/.jetbrains/acp.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"agent_servers": {
|
||||
"PraisonAI": {
|
||||
"command": "praisonai",
|
||||
"args": ["acp"],
|
||||
"env": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Toad
|
||||
|
||||
```bash
|
||||
toad acp "praisonai acp"
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
- `basic_acp_server.py` - Simple ACP server with default settings
|
||||
- `custom_agent_acp.py` - ACP server with a custom agent
|
||||
|
||||
## CLI Options
|
||||
|
||||
| Option | Description |
|
||||
|--------|-------------|
|
||||
| `-w, --workspace` | Workspace root directory |
|
||||
| `-a, --agent` | Agent name or config file |
|
||||
| `-m, --model` | LLM model to use |
|
||||
| `-r, --resume` | Resume session by ID |
|
||||
| `--last` | Resume the last session |
|
||||
| `--approve` | Approval mode: manual, auto, scoped |
|
||||
| `--read-only` | Read-only mode (default) |
|
||||
| `--allow-write` | Allow file writes |
|
||||
| `--allow-shell` | Allow shell commands |
|
||||
| `--debug` | Enable debug logging |
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install praisonai[acp]
|
||||
```
|
||||
|
||||
## Protocol Details
|
||||
|
||||
ACP uses JSON-RPC 2.0 over stdio:
|
||||
- **stdin**: Receives JSON-RPC requests from the client
|
||||
- **stdout**: Sends JSON-RPC responses to the client
|
||||
- **stderr**: Debug/log output (never pollutes stdout)
|
||||
|
||||
### Supported Methods
|
||||
|
||||
**Agent Methods (client → agent):**
|
||||
- `initialize` - Negotiate protocol version and capabilities
|
||||
- `authenticate` - Optional authentication
|
||||
- `session/new` - Create a new session
|
||||
- `session/load` - Load an existing session
|
||||
- `session/prompt` - Send user message
|
||||
- `session/cancel` - Cancel ongoing operations
|
||||
- `session/set_mode` - Change operating mode
|
||||
|
||||
**Client Methods (agent → client):**
|
||||
- `session/update` - Send progress updates
|
||||
- `session/request_permission` - Request user approval
|
||||
- `fs/read_text_file` - Read file contents
|
||||
- `fs/write_text_file` - Write file contents
|
||||
- `terminal/create` - Create terminal
|
||||
- `terminal/output` - Get terminal output
|
||||
@@ -0,0 +1,27 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic ACP Server Example
|
||||
|
||||
This example demonstrates how to run PraisonAI as an ACP server
|
||||
that can be connected to by IDEs like Zed, JetBrains, VSCode, or Toad.
|
||||
|
||||
Usage:
|
||||
# Run from command line
|
||||
praisonai acp
|
||||
|
||||
# Or run this script directly
|
||||
python basic_acp_server.py
|
||||
"""
|
||||
|
||||
from praisonai.acp import serve
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Start ACP server with default settings
|
||||
# This will listen on stdin/stdout for JSON-RPC messages
|
||||
serve(
|
||||
workspace=".", # Current directory as workspace
|
||||
agent="default", # Use default agent
|
||||
debug=False, # Set to True for debug logging to stderr
|
||||
read_only=True, # Safe by default
|
||||
approval_mode="manual", # Require approval for actions
|
||||
)
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Custom Agent ACP Server Example
|
||||
|
||||
This example shows how to run a custom PraisonAI agent as an ACP server.
|
||||
|
||||
Usage:
|
||||
python custom_agent_acp.py
|
||||
"""
|
||||
|
||||
from praisonai.acp import serve, ACPServer, ACPConfig
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Create a custom agent
|
||||
agent = Agent(
|
||||
name="CodeAssistant",
|
||||
instructions="""You are an expert coding assistant. You help users:
|
||||
- Write clean, efficient code
|
||||
- Debug issues
|
||||
- Explain complex concepts
|
||||
- Review code for best practices
|
||||
|
||||
Always provide clear explanations with your code suggestions.""",
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create config
|
||||
config = ACPConfig(
|
||||
workspace=".",
|
||||
debug=True,
|
||||
allow_write=True, # Allow file modifications
|
||||
approval_mode="manual",
|
||||
)
|
||||
|
||||
# Create server with custom agent
|
||||
server = ACPServer(config=config, agent=agent)
|
||||
|
||||
# Run the server
|
||||
serve(
|
||||
workspace=".",
|
||||
debug=True,
|
||||
allow_write=True,
|
||||
approval_mode="manual",
|
||||
)
|
||||
@@ -0,0 +1,135 @@
|
||||
"""
|
||||
Agent-Centric Tools Example
|
||||
|
||||
This example demonstrates how to use LSP/ACP-powered tools that make
|
||||
the Agent the central orchestrator for file operations.
|
||||
|
||||
How to run:
|
||||
export OPENAI_API_KEY=your_key
|
||||
python examples/agent_tools/agent_centric_example.py
|
||||
|
||||
Prerequisites:
|
||||
pip install praisonai praisonaiagents
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import shutil
|
||||
|
||||
# Add paths for development
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai-agents'))
|
||||
|
||||
|
||||
async def main():
|
||||
print("=" * 60)
|
||||
print("AGENT-CENTRIC TOOLS EXAMPLE")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
# Create temporary workspace
|
||||
workspace = tempfile.mkdtemp(prefix="praisonai_example_")
|
||||
print(f"Workspace: {workspace}")
|
||||
print()
|
||||
|
||||
try:
|
||||
# Import after path setup
|
||||
from praisonai.cli.features import (
|
||||
create_agent_centric_tools,
|
||||
InteractiveRuntime,
|
||||
RuntimeConfig
|
||||
)
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# 1. Create runtime with ACP enabled
|
||||
print("1. Creating InteractiveRuntime with ACP...")
|
||||
config = RuntimeConfig(
|
||||
workspace=workspace,
|
||||
lsp_enabled=False, # Skip LSP for this example
|
||||
acp_enabled=True,
|
||||
approval_mode="auto" # Auto-approve for demo
|
||||
)
|
||||
runtime = InteractiveRuntime(config)
|
||||
status = await runtime.start()
|
||||
|
||||
print(f" ACP ready: {runtime.acp_ready}")
|
||||
print(f" Read-only: {runtime.read_only}")
|
||||
print()
|
||||
|
||||
# 2. Create agent-centric tools
|
||||
print("2. Creating agent-centric tools...")
|
||||
tools = create_agent_centric_tools(runtime)
|
||||
print(f" Tools created: {len(tools)}")
|
||||
for tool in tools:
|
||||
print(f" - {tool.__name__}")
|
||||
print()
|
||||
|
||||
# 3. Create Agent with ACP-powered tools
|
||||
print("3. Creating Agent with ACP-powered tools...")
|
||||
agent = Agent(
|
||||
name="FileAgent",
|
||||
instructions="""You help create and manage files.
|
||||
Use acp_create_file to create files.
|
||||
Use read_file to read files.
|
||||
Use list_files to list directory contents.""",
|
||||
tools=tools,
|
||||
verbose=True
|
||||
)
|
||||
print()
|
||||
|
||||
# 4. Ask agent to create a file
|
||||
print("4. Asking agent to create a Python file...")
|
||||
print("-" * 40)
|
||||
result = agent.start(
|
||||
"Create a Python file called calculator.py with add and subtract functions"
|
||||
)
|
||||
print("-" * 40)
|
||||
print(f"Agent response: {result[:200]}..." if len(str(result)) > 200 else f"Agent response: {result}")
|
||||
print()
|
||||
|
||||
# 5. Verify file was created
|
||||
print("5. Verifying file creation...")
|
||||
calc_path = os.path.join(workspace, "calculator.py")
|
||||
if os.path.exists(calc_path):
|
||||
print(f" ✓ File created at {calc_path}")
|
||||
with open(calc_path) as f:
|
||||
content = f.read()
|
||||
print(f" Content preview:")
|
||||
for line in content.split('\n')[:10]:
|
||||
print(f" {line}")
|
||||
else:
|
||||
# Check what files exist
|
||||
files = os.listdir(workspace)
|
||||
print(f" Files in workspace: {files}")
|
||||
for f in files:
|
||||
fp = os.path.join(workspace, f)
|
||||
if os.path.isfile(fp):
|
||||
print(f" ✓ Found: {f}")
|
||||
print()
|
||||
|
||||
# 6. Ask agent to read the file
|
||||
print("6. Asking agent to read the file...")
|
||||
print("-" * 40)
|
||||
result2 = agent.start("Read the calculator.py file and tell me what functions it has")
|
||||
print("-" * 40)
|
||||
print(f"Agent response: {result2[:300]}..." if len(str(result2)) > 300 else f"Agent response: {result2}")
|
||||
print()
|
||||
|
||||
# Cleanup
|
||||
await runtime.stop()
|
||||
|
||||
finally:
|
||||
# Clean up workspace
|
||||
shutil.rmtree(workspace, ignore_errors=True)
|
||||
print(f"Cleaned up workspace: {workspace}")
|
||||
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("EXAMPLE COMPLETE")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -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()
|
||||
@@ -0,0 +1,228 @@
|
||||
"""
|
||||
Debug Interactive Mode Example
|
||||
|
||||
This example demonstrates how to use the debug CLI commands to test
|
||||
the Interactive Coding Assistant without entering TUI mode.
|
||||
|
||||
Features demonstrated:
|
||||
- LSP symbol listing
|
||||
- LSP definition lookup
|
||||
- ACP action planning
|
||||
- ACP action execution
|
||||
- Trace capture and JSON output
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add the src directory to path for development
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src" / "praisonai"))
|
||||
|
||||
from praisonai.cli.features.interactive_runtime import create_runtime
|
||||
from praisonai.cli.features.code_intelligence import CodeIntelligenceRouter
|
||||
from praisonai.cli.features.action_orchestrator import ActionOrchestrator
|
||||
|
||||
|
||||
async def demo_lsp_queries(workspace: str):
|
||||
"""Demonstrate LSP-based code queries."""
|
||||
print("\n" + "=" * 60)
|
||||
print("LSP CODE QUERIES DEMO")
|
||||
print("=" * 60)
|
||||
|
||||
runtime = create_runtime(
|
||||
workspace=workspace,
|
||||
lsp=True,
|
||||
acp=False,
|
||||
trace=True
|
||||
)
|
||||
|
||||
try:
|
||||
await runtime.start()
|
||||
router = CodeIntelligenceRouter(runtime)
|
||||
|
||||
# Query 1: List symbols
|
||||
print("\n[Query 1] List all symbols in app.py")
|
||||
result = await router.handle_query("list all functions in app.py", "app.py")
|
||||
print(f" Success: {result.success}")
|
||||
print(f" LSP Used: {result.lsp_used}")
|
||||
print(f" Fallback Used: {result.fallback_used}")
|
||||
if result.data:
|
||||
print(f" Symbols found: {len(result.data)}")
|
||||
for symbol in result.data[:5]:
|
||||
print(f" - {symbol['name']} ({symbol['kind']}) at line {symbol['line']}")
|
||||
|
||||
# Query 2: Find definition
|
||||
print("\n[Query 2] Find definition of Calculator")
|
||||
result = await router.handle_query("where is Calculator defined", "app.py")
|
||||
print(f" Success: {result.success}")
|
||||
if result.citations:
|
||||
print(f" Found at: {result.citations[0]}")
|
||||
|
||||
# Get trace
|
||||
trace = runtime.get_trace()
|
||||
print(f"\n[Trace] {len(trace.entries)} entries captured")
|
||||
|
||||
finally:
|
||||
await runtime.stop()
|
||||
|
||||
|
||||
async def demo_acp_actions(workspace: str):
|
||||
"""Demonstrate ACP-based code modifications."""
|
||||
print("\n" + "=" * 60)
|
||||
print("ACP CODE MODIFICATIONS DEMO")
|
||||
print("=" * 60)
|
||||
|
||||
runtime = create_runtime(
|
||||
workspace=workspace,
|
||||
lsp=False,
|
||||
acp=True,
|
||||
approval="auto",
|
||||
trace=True
|
||||
)
|
||||
|
||||
try:
|
||||
await runtime.start()
|
||||
print(f"\n ACP Ready: {runtime.acp_ready}")
|
||||
print(f" Read Only: {runtime.read_only}")
|
||||
|
||||
orchestrator = ActionOrchestrator(runtime)
|
||||
|
||||
# Action 1: Create a plan
|
||||
print("\n[Action 1] Create plan for new file")
|
||||
result = await orchestrator.create_plan("create a new file called demo_output.py")
|
||||
print(f" Success: {result.success}")
|
||||
if result.plan:
|
||||
print(f" Plan ID: {result.plan.id}")
|
||||
print(f" Steps: {len(result.plan.steps)}")
|
||||
for step in result.plan.steps:
|
||||
print(f" - {step.action_type.value}: {step.target}")
|
||||
|
||||
# Action 2: Execute with auto-approve
|
||||
print("\n[Action 2] Execute file creation")
|
||||
exec_result = await orchestrator.execute(
|
||||
"create a new file called demo_output.py with a hello function",
|
||||
auto_approve=True
|
||||
)
|
||||
print(f" Success: {exec_result.success}")
|
||||
print(f" Applied Steps: {exec_result.applied_steps}")
|
||||
|
||||
# Verify file was created
|
||||
output_file = Path(workspace) / "demo_output.py"
|
||||
if output_file.exists():
|
||||
print(f" File created: {output_file}")
|
||||
print(f" Content preview: {output_file.read_text()[:100]}...")
|
||||
|
||||
# Get trace
|
||||
trace = runtime.get_trace()
|
||||
print(f"\n[Trace] {len(trace.entries)} entries captured")
|
||||
for entry in trace.entries:
|
||||
print(f" [{entry.category}] {entry.action}")
|
||||
|
||||
finally:
|
||||
await runtime.stop()
|
||||
|
||||
|
||||
async def demo_full_interactive_turn(workspace: str):
|
||||
"""Demonstrate full interactive turn simulation."""
|
||||
print("\n" + "=" * 60)
|
||||
print("FULL INTERACTIVE TURN SIMULATION")
|
||||
print("=" * 60)
|
||||
|
||||
runtime = create_runtime(
|
||||
workspace=workspace,
|
||||
lsp=True,
|
||||
acp=True,
|
||||
approval="auto",
|
||||
trace=True
|
||||
)
|
||||
|
||||
try:
|
||||
status = await runtime.start()
|
||||
print("\n Runtime Status:")
|
||||
print(f" LSP: {'ready' if status['lsp']['ready'] else 'fallback'}")
|
||||
print(f" ACP: {'ready' if status['acp']['ready'] else 'unavailable'}")
|
||||
|
||||
router = CodeIntelligenceRouter(runtime)
|
||||
orchestrator = ActionOrchestrator(runtime)
|
||||
|
||||
# Simulate user prompt: code query
|
||||
prompt1 = "What functions are in app.py?"
|
||||
print(f"\n[Turn 1] User: {prompt1}")
|
||||
|
||||
# Classify and route
|
||||
result = await router.handle_query(prompt1, "app.py")
|
||||
print(f" Intent: {result.intent.value}")
|
||||
print(f" Response: Found {len(result.data) if result.data else 0} symbols")
|
||||
|
||||
# Simulate user prompt: code modification
|
||||
prompt2 = "Create a new test file"
|
||||
print(f"\n[Turn 2] User: {prompt2}")
|
||||
|
||||
exec_result = await orchestrator.execute(
|
||||
"create a new file called test_demo.py",
|
||||
auto_approve=True
|
||||
)
|
||||
print(" Action: file_create")
|
||||
print(f" Result: {'success' if exec_result.success else 'failed'}")
|
||||
|
||||
# Export trace as JSON
|
||||
trace = runtime.get_trace()
|
||||
trace_json = trace.to_dict()
|
||||
print("[Trace Export]")
|
||||
print(json.dumps(trace_json, indent=2, default=str)[:500] + "...")
|
||||
|
||||
finally:
|
||||
await runtime.stop()
|
||||
|
||||
|
||||
def main():
|
||||
"""Run the debug demo."""
|
||||
# Create a temporary workspace
|
||||
workspace = Path(__file__).parent / "_demo_workspace"
|
||||
workspace.mkdir(exist_ok=True)
|
||||
|
||||
# Create a sample app.py
|
||||
app_file = workspace / "app.py"
|
||||
app_file.write_text('''"""Sample application."""
|
||||
|
||||
def greet(name: str) -> str:
|
||||
"""Greet someone."""
|
||||
return f"Hello, {name}!"
|
||||
|
||||
def add(a: int, b: int) -> int:
|
||||
"""Add two numbers."""
|
||||
return a + b
|
||||
|
||||
class Calculator:
|
||||
"""Simple calculator."""
|
||||
|
||||
def __init__(self):
|
||||
self.value = 0
|
||||
|
||||
def add(self, x):
|
||||
self.value += x
|
||||
return self.value
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(greet("World"))
|
||||
''')
|
||||
|
||||
print("=" * 60)
|
||||
print("PRAISONAI DEBUG MODE DEMO")
|
||||
print("=" * 60)
|
||||
print(f"Workspace: {workspace}")
|
||||
|
||||
# Run demos
|
||||
asyncio.run(demo_lsp_queries(str(workspace)))
|
||||
asyncio.run(demo_acp_actions(str(workspace)))
|
||||
asyncio.run(demo_full_interactive_turn(str(workspace)))
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("DEMO COMPLETE")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
Interactive Default Tools Example
|
||||
|
||||
This example demonstrates how PraisonAI's interactive modes (TUI and prompt)
|
||||
include ACP and LSP tools by default.
|
||||
|
||||
Tool Groups:
|
||||
- ACP: acp_create_file, acp_edit_file, acp_delete_file, acp_execute_command
|
||||
- LSP: lsp_list_symbols, lsp_find_definition, lsp_find_references, lsp_get_diagnostics
|
||||
- Basic: read_file, write_file, list_files, execute_command, internet_search
|
||||
|
||||
Usage:
|
||||
python interactive_tools_example.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add paths for development
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai-agents'))
|
||||
|
||||
|
||||
def example_get_default_tools():
|
||||
"""Example: Get all default interactive tools."""
|
||||
from praisonai.cli.features.interactive_tools import (
|
||||
get_interactive_tools,
|
||||
ToolConfig,
|
||||
TOOL_GROUPS,
|
||||
)
|
||||
|
||||
print("=== Default Interactive Tools ===\n")
|
||||
|
||||
# Show tool groups
|
||||
print("Tool Groups:")
|
||||
for group, tools in TOOL_GROUPS.items():
|
||||
if group != "interactive":
|
||||
print(f" {group}: {tools}")
|
||||
print()
|
||||
|
||||
# Get all tools
|
||||
config = ToolConfig(workspace=os.getcwd())
|
||||
tools = get_interactive_tools(config=config)
|
||||
|
||||
print(f"Total tools loaded: {len(tools)}")
|
||||
print(f"Tool names: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def example_disable_acp():
|
||||
"""Example: Disable ACP tools."""
|
||||
from praisonai.cli.features.interactive_tools import (
|
||||
get_interactive_tools,
|
||||
ToolConfig,
|
||||
)
|
||||
|
||||
print("=== Disable ACP Tools (--no-acp) ===\n")
|
||||
|
||||
config = ToolConfig(workspace=os.getcwd(), enable_acp=False)
|
||||
tools = get_interactive_tools(config=config, disable=["acp"])
|
||||
|
||||
acp_tools = [t for t in tools if t.__name__.startswith("acp_")]
|
||||
print(f"ACP tools after --no-acp: {len(acp_tools)}")
|
||||
print(f"Remaining tools: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def example_disable_lsp():
|
||||
"""Example: Disable LSP tools."""
|
||||
from praisonai.cli.features.interactive_tools import (
|
||||
get_interactive_tools,
|
||||
ToolConfig,
|
||||
)
|
||||
|
||||
print("=== Disable LSP Tools (--no-lsp) ===\n")
|
||||
|
||||
config = ToolConfig(workspace=os.getcwd(), enable_lsp=False)
|
||||
tools = get_interactive_tools(config=config, disable=["lsp"])
|
||||
|
||||
lsp_tools = [t for t in tools if t.__name__.startswith("lsp_")]
|
||||
print(f"LSP tools after --no-lsp: {len(lsp_tools)}")
|
||||
print(f"Remaining tools: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
def example_agent_with_default_tools():
|
||||
"""Example: Create an Agent with default interactive tools."""
|
||||
from praisonai.cli.features.interactive_tools import get_interactive_tools, ToolConfig
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("=== Agent with Default Interactive Tools ===\n")
|
||||
|
||||
# Get default tools
|
||||
config = ToolConfig(workspace=os.getcwd())
|
||||
tools = get_interactive_tools(config=config)
|
||||
|
||||
# Create agent with tools
|
||||
agent = Agent(
|
||||
name="InteractiveAgent",
|
||||
instructions="You are a helpful coding assistant with file and code analysis capabilities.",
|
||||
tools=tools,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
print(f"Agent created with {len(tools)} tools")
|
||||
print(f"Tools: {[t.__name__ for t in tools]}")
|
||||
print()
|
||||
|
||||
# Simple test - list files
|
||||
print("Testing agent with list_files tool...")
|
||||
result = agent.start("List the Python files in the current directory using list_files tool.")
|
||||
print(f"Result: {result[:200]}..." if len(str(result)) > 200 else f"Result: {result}")
|
||||
print()
|
||||
|
||||
return agent
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("PraisonAI Interactive Default Tools Example")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
# Example 1: Get default tools
|
||||
example_get_default_tools()
|
||||
|
||||
# Example 2: Disable ACP
|
||||
example_disable_acp()
|
||||
|
||||
# Example 3: Disable LSP
|
||||
example_disable_lsp()
|
||||
|
||||
# Example 4: Agent with default tools (requires API key)
|
||||
if os.environ.get("OPENAI_API_KEY"):
|
||||
example_agent_with_default_tools()
|
||||
else:
|
||||
print("=== Agent Example Skipped (no OPENAI_API_KEY) ===")
|
||||
|
||||
print("=" * 60)
|
||||
print("Examples Complete")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,41 @@
|
||||
# AI SDK Examples
|
||||
|
||||
Examples demonstrating the AI SDK integration in PraisonAI TypeScript.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
npm install praisonai ai @ai-sdk/openai @ai-sdk/anthropic @ai-sdk/google zod
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-...
|
||||
export ANTHROPIC_API_KEY=sk-ant-...
|
||||
export GOOGLE_API_KEY=AIza...
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `basic-chat.ts` | Simple text generation |
|
||||
| `streaming.ts` | Streaming text output |
|
||||
| `tool-calling.ts` | Function/tool calling |
|
||||
| `multi-provider.ts` | Using multiple providers |
|
||||
| `structured-output.ts` | JSON schema output |
|
||||
| `multi-agent.ts` | Multi-agent attribution |
|
||||
|
||||
## Running Examples
|
||||
|
||||
```bash
|
||||
npx ts-node basic-chat.ts
|
||||
npx ts-node streaming.ts
|
||||
npx ts-node tool-calling.ts
|
||||
```
|
||||
|
||||
## More Information
|
||||
|
||||
- [AI SDK Documentation](https://docs.praisonai.com/docs/js/ai-sdk)
|
||||
- [AI SDK CLI](https://docs.praisonai.com/docs/js/ai-sdk-cli)
|
||||
@@ -0,0 +1,39 @@
|
||||
/**
|
||||
* AI SDK Basic Chat Example
|
||||
*
|
||||
* Demonstrates simple text generation using the AI SDK backend.
|
||||
*
|
||||
* Usage:
|
||||
* npx ts-node basic-chat.ts
|
||||
*
|
||||
* Environment:
|
||||
* OPENAI_API_KEY - Your OpenAI API key
|
||||
*/
|
||||
|
||||
import { createAISDKBackend } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
// Create backend with model string
|
||||
const backend = createAISDKBackend('openai/gpt-4o-mini');
|
||||
|
||||
console.log('Sending message to OpenAI...\n');
|
||||
|
||||
// Generate text
|
||||
const result = await backend.generateText({
|
||||
messages: [
|
||||
{ role: 'system', content: 'You are a helpful assistant.' },
|
||||
{ role: 'user', content: 'What are the three primary colors?' }
|
||||
],
|
||||
maxTokens: 200,
|
||||
temperature: 0.7,
|
||||
});
|
||||
|
||||
console.log('Response:', result.text);
|
||||
console.log('\nFinish reason:', result.finishReason);
|
||||
|
||||
if (result.usage) {
|
||||
console.log('Token usage:', result.usage);
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,32 @@
|
||||
/**
|
||||
* Generate Text Example
|
||||
*
|
||||
* Demonstrates basic text generation using AI SDK v6 via praisonai-ts.
|
||||
*
|
||||
* Run: npx ts-node generate-text.ts
|
||||
* Required: OPENAI_API_KEY
|
||||
*/
|
||||
|
||||
// Import from local praisonai-ts package
|
||||
// In production: import { Agent } from 'praisonai';
|
||||
import { Agent } from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
// Simple text generation with Agent
|
||||
const agent = new Agent({
|
||||
instructions: 'You are a helpful assistant that provides concise answers.',
|
||||
llm: 'openai/gpt-4o-mini',
|
||||
});
|
||||
|
||||
console.log('=== Generate Text Example ===\n');
|
||||
|
||||
// Basic chat
|
||||
const response = await agent.chat('What is the capital of France?');
|
||||
console.log('Response:', response);
|
||||
|
||||
// With context
|
||||
const response2 = await agent.chat('What is its population?');
|
||||
console.log('\nFollow-up:', response2);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,91 @@
|
||||
/**
|
||||
* MCP Client Example (AI SDK v6 Parity)
|
||||
*
|
||||
* Demonstrates Model Context Protocol client with HTTP/SSE transports.
|
||||
*
|
||||
* Run: npx ts-node mcp-client.ts
|
||||
* Note: Requires an MCP server running
|
||||
*/
|
||||
|
||||
import {
|
||||
createMCP,
|
||||
closeMCPClient,
|
||||
type MCPConfig,
|
||||
} from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
console.log('=== MCP Client Example ===\n');
|
||||
|
||||
// Example configurations for different transports
|
||||
const configs: Record<string, MCPConfig> = {
|
||||
// HTTP transport (recommended for production)
|
||||
http: {
|
||||
transport: {
|
||||
type: 'http',
|
||||
url: 'https://mcp-server.example.com/mcp',
|
||||
headers: { Authorization: 'Bearer token' },
|
||||
},
|
||||
name: 'http-client',
|
||||
},
|
||||
// SSE transport
|
||||
sse: {
|
||||
transport: {
|
||||
type: 'sse',
|
||||
url: 'https://mcp-server.example.com/sse',
|
||||
},
|
||||
name: 'sse-client',
|
||||
},
|
||||
// Stdio transport (for local servers)
|
||||
stdio: {
|
||||
transport: {
|
||||
type: 'stdio',
|
||||
command: 'npx',
|
||||
args: ['-y', '@modelcontextprotocol/server-everything'],
|
||||
},
|
||||
name: 'stdio-client',
|
||||
},
|
||||
};
|
||||
|
||||
console.log('MCP Transport Configurations:');
|
||||
console.log('');
|
||||
console.log('1. HTTP Transport (recommended):');
|
||||
console.log(` ${JSON.stringify(configs.http.transport, null, 2)}`);
|
||||
console.log('');
|
||||
console.log('2. SSE Transport:');
|
||||
console.log(` ${JSON.stringify(configs.sse.transport, null, 2)}`);
|
||||
console.log('');
|
||||
console.log('3. Stdio Transport (local):');
|
||||
console.log(` ${JSON.stringify(configs.stdio.transport, null, 2)}`);
|
||||
|
||||
console.log('\n--- Example Usage ---\n');
|
||||
console.log(`
|
||||
// Create MCP client
|
||||
const client = await createMCP({
|
||||
transport: {
|
||||
type: 'http',
|
||||
url: 'https://your-mcp-server.com/mcp',
|
||||
headers: { Authorization: 'Bearer your-token' },
|
||||
}
|
||||
});
|
||||
|
||||
// Get available tools
|
||||
const tools = await client.tools();
|
||||
console.log('Available tools:', Object.keys(tools));
|
||||
|
||||
// List resources
|
||||
const resources = await client.listResources();
|
||||
console.log('Resources:', resources);
|
||||
|
||||
// List prompts
|
||||
const prompts = await client.listPrompts();
|
||||
console.log('Prompts:', prompts);
|
||||
|
||||
// Close when done
|
||||
await closeMCPClient('default');
|
||||
`);
|
||||
|
||||
console.log('\n✅ MCP Client demo completed!');
|
||||
console.log('Note: To run with a real MCP server, uncomment the connection code.');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,96 @@
|
||||
/**
|
||||
* AI SDK Multi-Agent Attribution Example
|
||||
*
|
||||
* Demonstrates how to propagate attribution context (agentId, runId, traceId)
|
||||
* to LLM calls for multi-agent observability and safety.
|
||||
*
|
||||
* Usage:
|
||||
* npx ts-node multi-agent.ts
|
||||
*
|
||||
* Environment:
|
||||
* OPENAI_API_KEY - Your OpenAI API key
|
||||
*/
|
||||
|
||||
import { createAISDKBackend } from 'praisonai';
|
||||
|
||||
// Simulate a multi-agent system with attribution
|
||||
interface AgentContext {
|
||||
agentId: string;
|
||||
runId: string;
|
||||
traceId: string;
|
||||
sessionId?: string;
|
||||
}
|
||||
|
||||
async function runAgent(name: string, task: string, context: AgentContext) {
|
||||
console.log(`\n[${name}] Starting task: ${task}`);
|
||||
console.log(` Attribution: agentId=${context.agentId}, runId=${context.runId}`);
|
||||
|
||||
const backend = createAISDKBackend('openai/gpt-4o-mini', {
|
||||
attribution: context,
|
||||
timeout: 30000,
|
||||
});
|
||||
|
||||
const result = await backend.generateText({
|
||||
messages: [
|
||||
{ role: 'system', content: `You are ${name}, a specialized AI agent.` },
|
||||
{ role: 'user', content: task }
|
||||
],
|
||||
maxTokens: 100,
|
||||
});
|
||||
|
||||
console.log(` Response: ${result.text.slice(0, 100)}...`);
|
||||
return result.text;
|
||||
}
|
||||
|
||||
async function main() {
|
||||
// Create a shared trace ID for the entire workflow
|
||||
const traceId = `trace-${Date.now().toString(36)}`;
|
||||
const sessionId = 'session-user123';
|
||||
|
||||
console.log('Multi-Agent Workflow with Attribution');
|
||||
console.log('=====================================');
|
||||
console.log(`Trace ID: ${traceId}`);
|
||||
console.log(`Session ID: ${sessionId}`);
|
||||
|
||||
// Agent 1: Research Agent
|
||||
const researchResult = await runAgent(
|
||||
'ResearchAgent',
|
||||
'What are the key benefits of TypeScript?',
|
||||
{
|
||||
agentId: 'agent-research',
|
||||
runId: `run-${Date.now().toString(36)}-1`,
|
||||
traceId,
|
||||
sessionId,
|
||||
}
|
||||
);
|
||||
|
||||
// Agent 2: Summary Agent
|
||||
const summaryResult = await runAgent(
|
||||
'SummaryAgent',
|
||||
`Summarize this in one sentence: ${researchResult}`,
|
||||
{
|
||||
agentId: 'agent-summary',
|
||||
runId: `run-${Date.now().toString(36)}-2`,
|
||||
traceId,
|
||||
sessionId,
|
||||
}
|
||||
);
|
||||
|
||||
// Agent 3: Review Agent
|
||||
await runAgent(
|
||||
'ReviewAgent',
|
||||
`Review this summary for accuracy: ${summaryResult}`,
|
||||
{
|
||||
agentId: 'agent-review',
|
||||
runId: `run-${Date.now().toString(36)}-3`,
|
||||
traceId,
|
||||
sessionId,
|
||||
}
|
||||
);
|
||||
|
||||
console.log('\n\nWorkflow complete!');
|
||||
console.log('All requests were tagged with the same trace ID for observability.');
|
||||
console.log('Check your LLM provider dashboard to see the correlated requests.');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,59 @@
|
||||
/**
|
||||
* AI SDK Multi-Provider Example
|
||||
*
|
||||
* Demonstrates using multiple LLM providers with the same interface.
|
||||
*
|
||||
* Usage:
|
||||
* npx ts-node multi-provider.ts
|
||||
*
|
||||
* Environment:
|
||||
* OPENAI_API_KEY - Your OpenAI API key
|
||||
* ANTHROPIC_API_KEY - Your Anthropic API key
|
||||
* GOOGLE_API_KEY - Your Google API key
|
||||
*/
|
||||
|
||||
import { createAISDKBackend } from 'praisonai';
|
||||
|
||||
const PROMPT = 'Explain what an API is in one sentence.';
|
||||
|
||||
async function testProvider(modelString: string) {
|
||||
try {
|
||||
const backend = createAISDKBackend(modelString, {
|
||||
timeout: 30000,
|
||||
});
|
||||
|
||||
const startTime = Date.now();
|
||||
const result = await backend.generateText({
|
||||
messages: [{ role: 'user', content: PROMPT }],
|
||||
maxTokens: 100,
|
||||
});
|
||||
const duration = Date.now() - startTime;
|
||||
|
||||
console.log(`\n${modelString}:`);
|
||||
console.log(` Response: ${result.text}`);
|
||||
console.log(` Duration: ${duration}ms`);
|
||||
return true;
|
||||
} catch (error: any) {
|
||||
console.log(`\n${modelString}:`);
|
||||
console.log(` Error: ${error.message}`);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async function main() {
|
||||
console.log('Testing multiple providers with the same prompt...');
|
||||
console.log(`Prompt: "${PROMPT}"`);
|
||||
|
||||
const providers = [
|
||||
'openai/gpt-4o-mini',
|
||||
'anthropic/claude-3-haiku-20240307',
|
||||
'google/gemini-1.5-flash',
|
||||
];
|
||||
|
||||
const results = await Promise.all(providers.map(testProvider));
|
||||
|
||||
const successful = results.filter(Boolean).length;
|
||||
console.log(`\n\nSummary: ${successful}/${providers.length} providers succeeded`);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* Speech & Transcription Example (AI SDK v6 Parity)
|
||||
*
|
||||
* Demonstrates speech generation and transcription wrappers.
|
||||
*
|
||||
* Run: npx ts-node speech-transcription.ts
|
||||
* Required: OPENAI_API_KEY or ELEVENLABS_API_KEY
|
||||
*/
|
||||
|
||||
import {
|
||||
generateSpeech,
|
||||
transcribe,
|
||||
SPEECH_MODELS,
|
||||
TRANSCRIPTION_MODELS,
|
||||
} from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Speech & Transcription Example ===\n');
|
||||
|
||||
// Show available models
|
||||
console.log('1. Available Speech Models:');
|
||||
Object.entries(SPEECH_MODELS).forEach(([key, value]) => {
|
||||
console.log(` ${key}: ${value}`);
|
||||
});
|
||||
|
||||
console.log('\n2. Available Transcription Models:');
|
||||
Object.entries(TRANSCRIPTION_MODELS).forEach(([key, value]) => {
|
||||
console.log(` ${key}: ${value}`);
|
||||
});
|
||||
|
||||
// Example: Generate speech (requires API key)
|
||||
console.log('\n3. Speech Generation Example:');
|
||||
console.log(`
|
||||
// Generate speech from text
|
||||
const result = await generateSpeech({
|
||||
text: 'Hello, this is a test of speech synthesis.',
|
||||
model: 'openai/tts-1',
|
||||
voice: 'alloy',
|
||||
});
|
||||
|
||||
// Save to file
|
||||
await fs.writeFile('output.mp3', result.audio);
|
||||
console.log('Audio saved to output.mp3');
|
||||
`);
|
||||
|
||||
// Example: Transcribe audio (requires API key)
|
||||
console.log('4. Transcription Example:');
|
||||
console.log(`
|
||||
// Transcribe audio file
|
||||
const result = await transcribe({
|
||||
audio: await fs.readFile('audio.mp3'),
|
||||
model: 'openai/whisper-1',
|
||||
});
|
||||
|
||||
console.log('Transcription:', result.text);
|
||||
console.log('Language:', result.language);
|
||||
console.log('Duration:', result.durationInSeconds, 'seconds');
|
||||
|
||||
// With segments
|
||||
if (result.segments) {
|
||||
result.segments.forEach(seg => {
|
||||
console.log(\`[\${seg.start}s - \${seg.end}s]: \${seg.text}\`);
|
||||
});
|
||||
}
|
||||
`);
|
||||
|
||||
// Live test if API key is available
|
||||
if (process.env.OPENAI_API_KEY) {
|
||||
console.log('\n5. Live Speech Generation Test:');
|
||||
try {
|
||||
const result = await generateSpeech({
|
||||
text: 'Hello from PraisonAI!',
|
||||
model: 'openai/tts-1',
|
||||
voice: 'alloy',
|
||||
});
|
||||
console.log(` ✅ Generated ${result.audio.length} bytes of audio`);
|
||||
} catch (error: any) {
|
||||
console.log(` ⚠️ Speech generation failed: ${error.message}`);
|
||||
}
|
||||
} else {
|
||||
console.log('\n5. Live Test: Skipped (OPENAI_API_KEY not set)');
|
||||
}
|
||||
|
||||
console.log('\n✅ Speech & Transcription demo completed!');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,27 @@
|
||||
/**
|
||||
* Stream Text Example
|
||||
*
|
||||
* Demonstrates streaming text generation using AI SDK v6 via praisonai-ts.
|
||||
*
|
||||
* Run: npx ts-node stream-text.ts
|
||||
* Required: OPENAI_API_KEY
|
||||
*/
|
||||
|
||||
import { Agent } from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
const agent = new Agent({
|
||||
instructions: 'You are a helpful assistant that provides detailed answers.',
|
||||
llm: 'openai/gpt-4o-mini',
|
||||
stream: true,
|
||||
});
|
||||
|
||||
console.log('=== Stream Text Example ===\n');
|
||||
console.log('Streaming response:\n');
|
||||
|
||||
// Stream the response
|
||||
const response = await agent.chat('Explain quantum computing in 3 paragraphs.');
|
||||
console.log('\n\nFinal response:', response);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,39 @@
|
||||
/**
|
||||
* AI SDK Streaming Example
|
||||
*
|
||||
* Demonstrates streaming text generation with real-time output.
|
||||
*
|
||||
* Usage:
|
||||
* npx ts-node streaming.ts
|
||||
*
|
||||
* Environment:
|
||||
* OPENAI_API_KEY - Your OpenAI API key
|
||||
*/
|
||||
|
||||
import { createAISDKBackend } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
const backend = createAISDKBackend('openai/gpt-4o-mini');
|
||||
|
||||
console.log('Streaming response from OpenAI...\n');
|
||||
|
||||
const stream = await backend.streamText({
|
||||
messages: [
|
||||
{ role: 'user', content: 'Write a short poem about coding.' }
|
||||
],
|
||||
maxTokens: 200,
|
||||
});
|
||||
|
||||
// Stream chunks as they arrive
|
||||
for await (const chunk of stream) {
|
||||
if (chunk.text) {
|
||||
process.stdout.write(chunk.text);
|
||||
}
|
||||
|
||||
if (chunk.finishReason) {
|
||||
console.log('\n\nFinish reason:', chunk.finishReason);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,72 @@
|
||||
/**
|
||||
* AI SDK Structured Output Example
|
||||
*
|
||||
* Demonstrates generating typed JSON objects with Zod schemas.
|
||||
*
|
||||
* Usage:
|
||||
* npx ts-node structured-output.ts
|
||||
*
|
||||
* Environment:
|
||||
* OPENAI_API_KEY - Your OpenAI API key
|
||||
*/
|
||||
|
||||
import { generateObject } from 'ai';
|
||||
import { openai } from '@ai-sdk/openai';
|
||||
import { z } from 'zod';
|
||||
|
||||
// Define schema for a person profile
|
||||
const PersonSchema = z.object({
|
||||
name: z.string().describe('Full name'),
|
||||
age: z.number().describe('Age in years'),
|
||||
occupation: z.string().describe('Job title'),
|
||||
city: z.string().describe('City of residence'),
|
||||
interests: z.array(z.string()).describe('List of hobbies/interests'),
|
||||
bio: z.string().describe('Short biography'),
|
||||
});
|
||||
|
||||
// Define schema for a recipe
|
||||
const RecipeSchema = z.object({
|
||||
name: z.string().describe('Recipe name'),
|
||||
cuisine: z.string().describe('Type of cuisine'),
|
||||
prepTime: z.number().describe('Preparation time in minutes'),
|
||||
cookTime: z.number().describe('Cooking time in minutes'),
|
||||
servings: z.number().describe('Number of servings'),
|
||||
ingredients: z.array(z.object({
|
||||
item: z.string(),
|
||||
amount: z.string(),
|
||||
})).describe('List of ingredients'),
|
||||
instructions: z.array(z.string()).describe('Step-by-step instructions'),
|
||||
});
|
||||
|
||||
async function main() {
|
||||
console.log('Generating structured output with Zod schemas...\n');
|
||||
|
||||
// Generate a person profile
|
||||
console.log('1. Generating a fictional person profile...');
|
||||
const personResult = await generateObject({
|
||||
model: openai('gpt-4o-mini'),
|
||||
schema: PersonSchema,
|
||||
prompt: 'Generate a fictional software engineer profile.',
|
||||
});
|
||||
|
||||
console.log('Person:', JSON.stringify(personResult.object, null, 2));
|
||||
|
||||
// Generate a recipe
|
||||
console.log('\n2. Generating a recipe...');
|
||||
const recipeResult = await generateObject({
|
||||
model: openai('gpt-4o-mini'),
|
||||
schema: RecipeSchema,
|
||||
prompt: 'Generate a simple pasta recipe.',
|
||||
});
|
||||
|
||||
console.log('Recipe:', JSON.stringify(recipeResult.object, null, 2));
|
||||
|
||||
// Type-safe access
|
||||
console.log('\n3. Type-safe access:');
|
||||
console.log(` Person name: ${personResult.object.name}`);
|
||||
console.log(` Person age: ${personResult.object.age}`);
|
||||
console.log(` Recipe name: ${recipeResult.object.name}`);
|
||||
console.log(` Total time: ${recipeResult.object.prepTime + recipeResult.object.cookTime} minutes`);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,78 @@
|
||||
/**
|
||||
* Telemetry Example (AI SDK v6 Parity)
|
||||
*
|
||||
* Demonstrates telemetry and tracing with optional OpenTelemetry integration.
|
||||
*
|
||||
* Run: npx ts-node telemetry.ts
|
||||
* Required: OPENAI_API_KEY
|
||||
*/
|
||||
|
||||
import {
|
||||
configureTelemetry,
|
||||
isTelemetryEnabled,
|
||||
createAISpan,
|
||||
recordEvent,
|
||||
getEvents,
|
||||
clearEvents,
|
||||
Agent,
|
||||
} from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Telemetry Example ===\n');
|
||||
|
||||
// Configure telemetry
|
||||
configureTelemetry({
|
||||
isEnabled: true,
|
||||
functionId: 'telemetry-demo',
|
||||
metadata: {
|
||||
version: '1.0.0',
|
||||
environment: 'development',
|
||||
},
|
||||
});
|
||||
|
||||
console.log('1. Telemetry Status:');
|
||||
console.log(` Enabled: ${isTelemetryEnabled()}`);
|
||||
|
||||
// Create a span for an operation
|
||||
console.log('\n2. Creating AI Span:');
|
||||
const span = createAISpan('generateText', {
|
||||
model: 'gpt-4o-mini',
|
||||
provider: 'openai',
|
||||
});
|
||||
|
||||
try {
|
||||
// Simulate AI operation
|
||||
console.log(' Performing AI operation...');
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
|
||||
span.setAttribute('tokens', 150);
|
||||
span.setStatus({ code: 'ok' });
|
||||
console.log(' Span completed successfully');
|
||||
} catch (error: any) {
|
||||
span.recordException(error);
|
||||
span.setStatus({ code: 'error', message: error.message });
|
||||
} finally {
|
||||
span.end();
|
||||
}
|
||||
|
||||
// Record custom events
|
||||
console.log('\n3. Recording Events:');
|
||||
recordEvent('agent.started', { agentName: 'demo-agent' });
|
||||
recordEvent('tool.called', { toolName: 'search', duration: 250 });
|
||||
recordEvent('agent.completed', { success: true });
|
||||
|
||||
const events = getEvents();
|
||||
console.log(` Recorded ${events.length} events:`);
|
||||
events.forEach(e => {
|
||||
console.log(` - ${e.name}: ${JSON.stringify(e.attributes)}`);
|
||||
});
|
||||
|
||||
// Clear events
|
||||
console.log('\n4. Clearing Events:');
|
||||
clearEvents();
|
||||
console.log(` Events after clear: ${getEvents().length}`);
|
||||
|
||||
console.log('\n✅ Telemetry demo completed!');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,71 @@
|
||||
/**
|
||||
* Tool Approval Example (AI SDK v6 Parity)
|
||||
*
|
||||
* Demonstrates human-in-the-loop tool approval using needsApproval.
|
||||
*
|
||||
* Run: npx ts-node tool-approval.ts
|
||||
* Required: OPENAI_API_KEY
|
||||
*/
|
||||
|
||||
import {
|
||||
Agent,
|
||||
withApproval,
|
||||
getApprovalManager,
|
||||
createDangerousPatternChecker,
|
||||
} from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Tool Approval Example ===\n');
|
||||
|
||||
// Set up approval handler
|
||||
const manager = getApprovalManager();
|
||||
manager.onApprovalRequest(async (request) => {
|
||||
console.log(`\n🔐 APPROVAL REQUESTED`);
|
||||
console.log(` Tool: ${request.toolName}`);
|
||||
console.log(` Input: ${JSON.stringify(request.input)}`);
|
||||
console.log(` Reason: ${request.reason || 'No reason provided'}`);
|
||||
|
||||
// Auto-approve for demo (in production, prompt user)
|
||||
console.log(` ✅ Auto-approved for demo\n`);
|
||||
return true;
|
||||
});
|
||||
|
||||
// Create a tool that always requires approval
|
||||
const deleteFile = withApproval({
|
||||
name: 'deleteFile',
|
||||
description: 'Delete a file from the filesystem',
|
||||
needsApproval: true, // Always require approval
|
||||
execute: async (input: { path: string }) => {
|
||||
console.log(` Executing: Would delete ${input.path}`);
|
||||
return { success: true, deleted: input.path };
|
||||
},
|
||||
});
|
||||
|
||||
// Create a tool with conditional approval (dangerous patterns only)
|
||||
const runCommand = withApproval({
|
||||
name: 'runCommand',
|
||||
description: 'Run a shell command',
|
||||
needsApproval: createDangerousPatternChecker(),
|
||||
execute: async (input: { command: string }) => {
|
||||
console.log(` Executing: ${input.command}`);
|
||||
return { output: 'command executed' };
|
||||
},
|
||||
});
|
||||
|
||||
// Test the tools
|
||||
console.log('1. Testing deleteFile (always requires approval):');
|
||||
const result1 = await deleteFile({ path: '/tmp/test.txt' });
|
||||
console.log(' Result:', result1);
|
||||
|
||||
console.log('\n2. Testing runCommand with safe command (no approval needed):');
|
||||
const result2 = await runCommand({ command: 'ls -la' });
|
||||
console.log(' Result:', result2);
|
||||
|
||||
console.log('\n3. Testing runCommand with dangerous command (requires approval):');
|
||||
const result3 = await runCommand({ command: 'rm -rf /tmp/test' });
|
||||
console.log(' Result:', result3);
|
||||
|
||||
console.log('\n✅ Tool approval demo completed!');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,75 @@
|
||||
/**
|
||||
* AI SDK Tool Calling Example
|
||||
*
|
||||
* Demonstrates function/tool calling with Zod schemas.
|
||||
*
|
||||
* Usage:
|
||||
* npx ts-node tool-calling.ts
|
||||
*
|
||||
* Environment:
|
||||
* OPENAI_API_KEY - Your OpenAI API key
|
||||
*/
|
||||
|
||||
import { generateText, tool } from 'ai';
|
||||
import { openai } from '@ai-sdk/openai';
|
||||
import { z } from 'zod';
|
||||
|
||||
// Define tools with Zod schemas
|
||||
const tools = {
|
||||
get_weather: tool({
|
||||
description: 'Get the current weather for a city',
|
||||
inputSchema: z.object({
|
||||
city: z.string().describe('The city name'),
|
||||
unit: z.enum(['celsius', 'fahrenheit']).optional().describe('Temperature unit'),
|
||||
}),
|
||||
execute: async ({ city, unit = 'celsius' }) => {
|
||||
// Mock weather data
|
||||
const weather = {
|
||||
city,
|
||||
temperature: unit === 'celsius' ? 22 : 72,
|
||||
unit,
|
||||
condition: 'sunny',
|
||||
humidity: 45,
|
||||
};
|
||||
return weather;
|
||||
},
|
||||
}),
|
||||
|
||||
get_time: tool({
|
||||
description: 'Get the current time in a timezone',
|
||||
inputSchema: z.object({
|
||||
timezone: z.string().describe('Timezone like "America/New_York"'),
|
||||
}),
|
||||
execute: async ({ timezone }) => {
|
||||
const time = new Date().toLocaleString('en-US', { timeZone: timezone });
|
||||
return { timezone, time };
|
||||
},
|
||||
}),
|
||||
};
|
||||
|
||||
async function main() {
|
||||
console.log('Asking about weather with tool calling...\n');
|
||||
|
||||
const result = await generateText({
|
||||
model: openai('gpt-4o-mini'),
|
||||
messages: [
|
||||
{ role: 'user', content: 'What is the weather in Paris and what time is it in New York?' }
|
||||
],
|
||||
tools,
|
||||
maxSteps: 5, // Allow multiple tool calls
|
||||
});
|
||||
|
||||
console.log('Tool Calls:');
|
||||
for (const toolCall of result.toolCalls) {
|
||||
console.log(` - ${toolCall.toolName}:`, toolCall.args);
|
||||
}
|
||||
|
||||
console.log('\nTool Results:');
|
||||
for (const toolResult of result.toolResults) {
|
||||
console.log(` - ${toolResult.toolName}:`, toolResult.result);
|
||||
}
|
||||
|
||||
console.log('\nFinal Response:', result.text);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,63 @@
|
||||
/**
|
||||
* UI Message Example (AI SDK v6 Parity)
|
||||
*
|
||||
* Demonstrates UIMessage protocol for Next.js/React integration.
|
||||
*
|
||||
* Run: npx ts-node ui-message.ts
|
||||
* Required: OPENAI_API_KEY
|
||||
*/
|
||||
|
||||
import {
|
||||
createTextMessage,
|
||||
createSystemMessage,
|
||||
convertToModelMessages,
|
||||
convertToUIMessages,
|
||||
validateUIMessages,
|
||||
hasPendingApprovals,
|
||||
type UIMessage,
|
||||
} from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
console.log('=== UI Message Example ===\n');
|
||||
|
||||
// Create UI messages
|
||||
const messages: UIMessage[] = [
|
||||
createSystemMessage('You are a helpful assistant.'),
|
||||
createTextMessage('user', 'What is the capital of France?'),
|
||||
createTextMessage('assistant', 'The capital of France is Paris.'),
|
||||
createTextMessage('user', 'What is its population?'),
|
||||
];
|
||||
|
||||
console.log('1. Created UI Messages:');
|
||||
messages.forEach((msg, i) => {
|
||||
const text = msg.parts.find(p => p.type === 'text');
|
||||
console.log(` [${i}] ${msg.role}: ${text ? (text as any).text.substring(0, 50) : '...'}`);
|
||||
});
|
||||
|
||||
// Validate messages
|
||||
console.log('\n2. Validating messages:');
|
||||
const isValid = validateUIMessages(messages);
|
||||
console.log(` Valid: ${isValid}`);
|
||||
|
||||
// Convert to model messages (for AI SDK)
|
||||
console.log('\n3. Converting to Model Messages:');
|
||||
const modelMessages = await convertToModelMessages(messages);
|
||||
modelMessages.forEach((msg, i) => {
|
||||
const content = typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
|
||||
console.log(` [${i}] ${msg.role}: ${content.substring(0, 50)}...`);
|
||||
});
|
||||
|
||||
// Check for pending approvals
|
||||
console.log('\n4. Checking for pending approvals:');
|
||||
const hasPending = messages.some(m => hasPendingApprovals(m));
|
||||
console.log(` Has pending approvals: ${hasPending}`);
|
||||
|
||||
// Convert model messages back to UI messages
|
||||
console.log('\n5. Round-trip conversion:');
|
||||
const uiMessages = convertToUIMessages(modelMessages);
|
||||
console.log(` Converted ${uiMessages.length} messages back to UI format`);
|
||||
|
||||
console.log('\n✅ UI Message demo completed!');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,65 @@
|
||||
/**
|
||||
* Basic Embedding Example
|
||||
*
|
||||
* Demonstrates using Agent.embed() for text embeddings
|
||||
*
|
||||
* Run: npx ts-node examples/js/embeddings/embed-basic.ts
|
||||
*/
|
||||
|
||||
import { Agent } from '../../../src/praisonai-ts/dist/agent/simple';
|
||||
import { embed, embedMany, cosineSimilarity } from '../../../src/praisonai-ts/dist/llm/embeddings';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Basic Embedding Example ===\n');
|
||||
|
||||
// Method 1: Using Agent.embed()
|
||||
console.log('1. Agent.embed() - Single text');
|
||||
const agent = new Agent({
|
||||
instructions: 'You are a helpful assistant',
|
||||
llm: 'openai/gpt-4o-mini',
|
||||
verbose: false
|
||||
});
|
||||
|
||||
const singleEmbedding = await agent.embed('Hello world');
|
||||
console.log(` Dimensions: ${(singleEmbedding as number[]).length}`);
|
||||
console.log(` First 5 values: [${(singleEmbedding as number[]).slice(0, 5).map(v => v.toFixed(4)).join(', ')}...]`);
|
||||
|
||||
// Method 2: Using Agent.embed() with multiple texts
|
||||
console.log('\n2. Agent.embed() - Multiple texts');
|
||||
const multiEmbeddings = await agent.embed(['Hello', 'World', 'How are you?']);
|
||||
console.log(` Count: ${(multiEmbeddings as number[][]).length}`);
|
||||
console.log(` Each has ${(multiEmbeddings as number[][])[0].length} dimensions`);
|
||||
|
||||
// Method 3: Direct embed function
|
||||
console.log('\n3. Direct embed() function');
|
||||
const result = await embed('PraisonAI is an AI agent framework');
|
||||
console.log(` Dimensions: ${result.embedding.length}`);
|
||||
console.log(` Tokens used: ${result.usage?.tokens || 'N/A'}`);
|
||||
|
||||
// Method 4: Batch embeddings
|
||||
console.log('\n4. Batch embedMany() function');
|
||||
const texts = [
|
||||
'Machine learning is a subset of AI',
|
||||
'Deep learning uses neural networks',
|
||||
'Natural language processing handles text'
|
||||
];
|
||||
const batchResult = await embedMany(texts);
|
||||
console.log(` Embedded ${batchResult.embeddings.length} texts`);
|
||||
console.log(` Total tokens: ${batchResult.usage?.tokens || 'N/A'}`);
|
||||
|
||||
// Method 5: Similarity comparison
|
||||
console.log('\n5. Cosine Similarity');
|
||||
const emb1 = await embed('I love programming');
|
||||
const emb2 = await embed('Coding is my passion');
|
||||
const emb3 = await embed('The weather is nice today');
|
||||
|
||||
const sim12 = cosineSimilarity(emb1.embedding, emb2.embedding);
|
||||
const sim13 = cosineSimilarity(emb1.embedding, emb3.embedding);
|
||||
|
||||
console.log(` "I love programming" vs "Coding is my passion": ${(sim12 * 100).toFixed(1)}%`);
|
||||
console.log(` "I love programming" vs "The weather is nice": ${(sim13 * 100).toFixed(1)}%`);
|
||||
|
||||
console.log('\n=== Complete ===');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,119 @@
|
||||
/**
|
||||
* Document Embedding Example
|
||||
*
|
||||
* Demonstrates chunking and embedding documents for retrieval
|
||||
*
|
||||
* Run: npx ts-node examples/js/embeddings/embed-docs.ts
|
||||
*/
|
||||
|
||||
import { embedMany, cosineSimilarity } from '../../../src/praisonai-ts/dist/llm/embeddings';
|
||||
|
||||
// Sample documents
|
||||
const documents = [
|
||||
{
|
||||
id: 'doc1',
|
||||
title: 'Introduction to AI',
|
||||
content: 'Artificial Intelligence (AI) is the simulation of human intelligence by machines. It includes machine learning, natural language processing, and computer vision.'
|
||||
},
|
||||
{
|
||||
id: 'doc2',
|
||||
title: 'Machine Learning Basics',
|
||||
content: 'Machine learning is a subset of AI that enables systems to learn from data. It includes supervised learning, unsupervised learning, and reinforcement learning.'
|
||||
},
|
||||
{
|
||||
id: 'doc3',
|
||||
title: 'Deep Learning',
|
||||
content: 'Deep learning uses neural networks with multiple layers. It excels at tasks like image recognition, speech processing, and natural language understanding.'
|
||||
},
|
||||
{
|
||||
id: 'doc4',
|
||||
title: 'Natural Language Processing',
|
||||
content: 'NLP enables computers to understand human language. Applications include chatbots, translation, sentiment analysis, and text summarization.'
|
||||
},
|
||||
{
|
||||
id: 'doc5',
|
||||
title: 'Computer Vision',
|
||||
content: 'Computer vision allows machines to interpret visual information. It powers applications like facial recognition, autonomous vehicles, and medical imaging.'
|
||||
}
|
||||
];
|
||||
|
||||
interface EmbeddedDocument {
|
||||
id: string;
|
||||
title: string;
|
||||
content: string;
|
||||
embedding: number[];
|
||||
}
|
||||
|
||||
async function main() {
|
||||
console.log('=== Document Embedding Example ===\n');
|
||||
|
||||
// Step 1: Embed all documents
|
||||
console.log('1. Embedding documents...');
|
||||
const contents = documents.map(d => d.content);
|
||||
const { embeddings } = await embedMany(contents);
|
||||
|
||||
const embeddedDocs: EmbeddedDocument[] = documents.map((doc, i) => ({
|
||||
...doc,
|
||||
embedding: embeddings[i]
|
||||
}));
|
||||
|
||||
console.log(` Embedded ${embeddedDocs.length} documents`);
|
||||
console.log(` Embedding dimensions: ${embeddings[0].length}`);
|
||||
|
||||
// Step 2: Search function
|
||||
async function search(query: string, topK: number = 3): Promise<Array<{ doc: EmbeddedDocument; score: number }>> {
|
||||
const { embedding: queryEmbedding } = await (await import('../../../src/praisonai-ts/dist/llm/embeddings')).embed(query);
|
||||
|
||||
const results = embeddedDocs.map(doc => ({
|
||||
doc,
|
||||
score: cosineSimilarity(queryEmbedding, doc.embedding)
|
||||
}));
|
||||
|
||||
return results
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, topK);
|
||||
}
|
||||
|
||||
// Step 3: Test queries
|
||||
console.log('\n2. Testing search queries...\n');
|
||||
|
||||
const queries = [
|
||||
'How do machines learn from data?',
|
||||
'What is image recognition?',
|
||||
'Tell me about chatbots and translation'
|
||||
];
|
||||
|
||||
for (const query of queries) {
|
||||
console.log(`Query: "${query}"`);
|
||||
const results = await search(query);
|
||||
|
||||
console.log('Top results:');
|
||||
results.forEach((r, i) => {
|
||||
console.log(` ${i + 1}. [${(r.score * 100).toFixed(1)}%] ${r.doc.title}`);
|
||||
});
|
||||
console.log('');
|
||||
}
|
||||
|
||||
// Step 4: Show document similarity matrix
|
||||
console.log('3. Document Similarity Matrix');
|
||||
console.log('');
|
||||
|
||||
// Header
|
||||
process.stdout.write(' ');
|
||||
documents.forEach((_, i) => process.stdout.write(`Doc${i + 1} `));
|
||||
console.log('');
|
||||
|
||||
// Matrix
|
||||
for (let i = 0; i < embeddedDocs.length; i++) {
|
||||
process.stdout.write(`Doc${i + 1} `);
|
||||
for (let j = 0; j < embeddedDocs.length; j++) {
|
||||
const sim = cosineSimilarity(embeddedDocs[i].embedding, embeddedDocs[j].embedding);
|
||||
process.stdout.write(`${(sim * 100).toFixed(0).padStart(3)}% `);
|
||||
}
|
||||
console.log('');
|
||||
}
|
||||
|
||||
console.log('\n=== Complete ===');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,131 @@
|
||||
/**
|
||||
* Retrieval Example - Simple Vector Search
|
||||
*
|
||||
* Demonstrates building a simple retrieval system with embeddings
|
||||
*
|
||||
* Run: npx ts-node examples/js/embeddings/retrieval.ts
|
||||
*/
|
||||
|
||||
import { embed, embedMany, cosineSimilarity } from '../../../src/praisonai-ts/dist/llm/embeddings';
|
||||
|
||||
interface Document {
|
||||
id: string;
|
||||
content: string;
|
||||
embedding?: number[];
|
||||
}
|
||||
|
||||
interface SearchResult {
|
||||
document: Document;
|
||||
score: number;
|
||||
}
|
||||
|
||||
// Simple in-memory vector store
|
||||
class SimpleVectorStore {
|
||||
private documents: Document[] = [];
|
||||
|
||||
async add(doc: Document): Promise<void> {
|
||||
const { embedding } = await embed(doc.content);
|
||||
this.documents.push({ ...doc, embedding });
|
||||
}
|
||||
|
||||
async addBatch(docs: Document[]): Promise<void> {
|
||||
const contents = docs.map(d => d.content);
|
||||
const { embeddings } = await embedMany(contents);
|
||||
docs.forEach((doc, i) => {
|
||||
this.documents.push({ ...doc, embedding: embeddings[i] });
|
||||
});
|
||||
}
|
||||
|
||||
async search(query: string, topK: number = 3): Promise<SearchResult[]> {
|
||||
const { embedding: queryEmb } = await embed(query);
|
||||
|
||||
const results = this.documents
|
||||
.filter(doc => doc.embedding)
|
||||
.map(doc => ({
|
||||
document: doc,
|
||||
score: cosineSimilarity(queryEmb, doc.embedding!)
|
||||
}))
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, topK);
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
get(id: string): Document | undefined {
|
||||
return this.documents.find(d => d.id === id);
|
||||
}
|
||||
|
||||
delete(id: string): boolean {
|
||||
const idx = this.documents.findIndex(d => d.id === id);
|
||||
if (idx >= 0) {
|
||||
this.documents.splice(idx, 1);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async function main() {
|
||||
console.log('=== Simple Vector Retrieval ===\n');
|
||||
|
||||
// Create vector store
|
||||
console.log('1. Creating vector store...');
|
||||
const store = new SimpleVectorStore();
|
||||
|
||||
// Add documents
|
||||
console.log('2. Adding documents...');
|
||||
const docs: Document[] = [
|
||||
{ id: '1', content: 'PraisonAI is an AI agent framework for building intelligent applications.' },
|
||||
{ id: '2', content: 'Agents can use tools to interact with external systems and APIs.' },
|
||||
{ id: '3', content: 'Workflows allow orchestrating multiple agents in sequence or parallel.' },
|
||||
{ id: '4', content: 'Memory systems help agents maintain context across conversations.' },
|
||||
{ id: '5', content: 'The AI SDK provides a unified interface for multiple LLM providers.' },
|
||||
{ id: '6', content: 'Embeddings enable semantic search and similarity matching.' },
|
||||
{ id: '7', content: 'Guardrails ensure agent outputs meet safety and quality standards.' },
|
||||
{ id: '8', content: 'Sessions persist conversation state for multi-turn interactions.' }
|
||||
];
|
||||
|
||||
await store.addBatch(docs);
|
||||
console.log(` Added ${docs.length} documents`);
|
||||
|
||||
// Search queries
|
||||
console.log('\n3. Running search queries...\n');
|
||||
|
||||
const queries = [
|
||||
'How do I build an AI application?',
|
||||
'Can agents call external APIs?',
|
||||
'How to maintain conversation context?',
|
||||
'What providers are supported?'
|
||||
];
|
||||
|
||||
for (const query of queries) {
|
||||
console.log(`Query: "${query}"`);
|
||||
const results = await store.search(query);
|
||||
|
||||
if (results.length > 0) {
|
||||
console.log('Results:');
|
||||
results.forEach((r, i) => {
|
||||
console.log(` ${i + 1}. [${(r.score * 100).toFixed(1)}%] ${r.document.content.substring(0, 60)}...`);
|
||||
});
|
||||
} else {
|
||||
console.log(' No results found');
|
||||
}
|
||||
console.log('');
|
||||
}
|
||||
|
||||
// Get specific document
|
||||
console.log('4. Get document by ID...');
|
||||
const doc = store.get('5');
|
||||
if (doc) {
|
||||
console.log(` Found: ${doc.content}`);
|
||||
}
|
||||
|
||||
// Delete document
|
||||
console.log('\n5. Delete document...');
|
||||
const deleted = store.delete('8');
|
||||
console.log(` Deleted doc 8: ${deleted}`);
|
||||
|
||||
console.log('\n=== Complete ===');
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,49 @@
|
||||
/**
|
||||
* MCP Filesystem Server Example
|
||||
*
|
||||
* Demonstrates using MCP to connect to a filesystem server.
|
||||
*
|
||||
* Prerequisites:
|
||||
* npm install praisonai-ts
|
||||
* export OPENAI_API_KEY=your-api-key
|
||||
*
|
||||
* Run:
|
||||
* npx ts-node filesystem-mcp.ts
|
||||
*/
|
||||
|
||||
import { Agent, createMCP } from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
console.log('=== MCP Filesystem Server Example ===\n');
|
||||
|
||||
// Connect to MCP filesystem server
|
||||
const mcp = await createMCP({
|
||||
transport: {
|
||||
type: 'stdio',
|
||||
command: 'npx',
|
||||
args: ['-y', '@modelcontextprotocol/server-filesystem', '.'],
|
||||
},
|
||||
});
|
||||
|
||||
console.log('Connected to MCP filesystem server\n');
|
||||
|
||||
// Get tools from MCP server
|
||||
const tools = await mcp.tools();
|
||||
console.log('Available tools:', Object.keys(tools));
|
||||
|
||||
// Create agent with MCP tools
|
||||
const agent = new Agent({
|
||||
name: 'FileAgent',
|
||||
instructions: 'You help manage files using the filesystem tools.',
|
||||
tools: Object.values(tools),
|
||||
});
|
||||
|
||||
// Ask the agent to list files
|
||||
const response = await agent.chat('List all TypeScript files in the current directory');
|
||||
console.log('\nResponse:', response);
|
||||
|
||||
// Cleanup
|
||||
await mcp.close();
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,41 @@
|
||||
/**
|
||||
* Multimodal Image Analysis Example
|
||||
*
|
||||
* Demonstrates analyzing images with a vision-capable model.
|
||||
*
|
||||
* Prerequisites:
|
||||
* npm install praisonai-ts
|
||||
* export OPENAI_API_KEY=your-api-key
|
||||
*
|
||||
* Run:
|
||||
* npx ts-node image-analysis.ts
|
||||
*/
|
||||
|
||||
import { Agent } from '../../../src/praisonai-ts/src';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Multimodal Image Analysis Example ===\n');
|
||||
|
||||
// Create an agent with a vision-capable model
|
||||
const agent = new Agent({
|
||||
name: 'VisionAgent',
|
||||
llm: 'gpt-4o', // Vision-capable model
|
||||
instructions: 'You analyze images and describe what you see in detail.',
|
||||
});
|
||||
|
||||
// For multimodal, we describe the image in the prompt
|
||||
// The agent will use its vision capabilities
|
||||
const imageUrl = 'https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/PNG_transparency_demonstration_1.png/300px-PNG_transparency_demonstration_1.png';
|
||||
|
||||
console.log('Analyzing image...\n');
|
||||
console.log('Image URL:', imageUrl, '\n');
|
||||
|
||||
// Ask the agent to analyze (in production, use multimodal message format)
|
||||
const response = await agent.chat(
|
||||
`Please analyze this image: ${imageUrl}\nDescribe what you see in detail.`
|
||||
);
|
||||
|
||||
console.log('Analysis:', response);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,50 @@
|
||||
# Observability Examples
|
||||
|
||||
Examples demonstrating how to use observability integrations with PraisonAI TypeScript.
|
||||
|
||||
## Setup
|
||||
|
||||
```bash
|
||||
# Install dependencies
|
||||
cd /path/to/praisonai-ts
|
||||
pnpm install
|
||||
|
||||
# Set API keys for observability tools
|
||||
export LANGFUSE_SECRET_KEY=sk-...
|
||||
export LANGFUSE_PUBLIC_KEY=pk-...
|
||||
# Or use built-in tools (no API key needed)
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `basic-tracing.ts` | Basic tracing with memory adapter |
|
||||
| `langfuse-integration.ts` | Langfuse observability integration |
|
||||
| `multi-agent-attribution.ts` | Track agent_id/run_id across agents |
|
||||
| `switchable-tools.ts` | Switch between observability tools via env |
|
||||
|
||||
## Running Examples
|
||||
|
||||
```bash
|
||||
# Run with ts-node or tsx
|
||||
npx tsx examples/js/observability/basic-tracing.ts
|
||||
npx tsx examples/js/observability/multi-agent-attribution.ts
|
||||
|
||||
# With Langfuse (requires API keys)
|
||||
LANGFUSE_SECRET_KEY=sk-... npx tsx examples/js/observability/langfuse-integration.ts
|
||||
```
|
||||
|
||||
## Supported Tools
|
||||
|
||||
| Tool | Env Variable | Package |
|
||||
|------|--------------|---------|
|
||||
| Langfuse | `LANGFUSE_SECRET_KEY` | `langfuse` |
|
||||
| LangSmith | `LANGCHAIN_API_KEY` | `langsmith` |
|
||||
| LangWatch | `LANGWATCH_API_KEY` | `langwatch` |
|
||||
| Arize | `ARIZE_API_KEY` | `arize-phoenix` |
|
||||
| Axiom | `AXIOM_TOKEN` | `@axiomhq/js` |
|
||||
| Braintrust | `BRAINTRUST_API_KEY` | `braintrust` |
|
||||
| Helicone | `HELICONE_API_KEY` | `@helicone/helicone` |
|
||||
| Console | (none) | built-in |
|
||||
| Memory | (none) | built-in |
|
||||
@@ -0,0 +1,45 @@
|
||||
/**
|
||||
* Basic Tracing Example
|
||||
*
|
||||
* Demonstrates basic observability with the built-in memory adapter.
|
||||
* No external API keys required.
|
||||
*/
|
||||
|
||||
import { Agent, MemoryObservabilityAdapter, setObservabilityAdapter } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
// Create memory adapter (stores traces in memory)
|
||||
const obs = new MemoryObservabilityAdapter();
|
||||
setObservabilityAdapter(obs);
|
||||
|
||||
// Create agent with observability enabled
|
||||
const agent = new Agent({
|
||||
name: 'TracedAgent',
|
||||
instructions: 'You are a helpful assistant. Answer questions concisely.'
|
||||
});
|
||||
|
||||
console.log('Running agent with tracing...\n');
|
||||
|
||||
// Run some interactions
|
||||
await agent.chat('What is 2 + 2?');
|
||||
await agent.chat('What is the capital of France?');
|
||||
|
||||
// View traces
|
||||
const traces = obs.getAllTraces();
|
||||
|
||||
console.log('\n=== Trace Summary ===');
|
||||
console.log(`Total traces: ${traces.length}`);
|
||||
|
||||
for (const trace of traces) {
|
||||
console.log(`\nTrace: ${trace.name}`);
|
||||
console.log(` Status: ${trace.status}`);
|
||||
console.log(` Spans: ${trace.spans.length}`);
|
||||
|
||||
for (const span of trace.spans) {
|
||||
const duration = span.endTime ? span.endTime - span.startTime : 'running';
|
||||
console.log(` - ${span.kind}: ${span.name} (${duration}ms)`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,60 @@
|
||||
/**
|
||||
* Multi-Agent Attribution Example
|
||||
*
|
||||
* Demonstrates tracking agent_id, run_id, and trace_id across
|
||||
* multi-agent workflows for proper attribution.
|
||||
*/
|
||||
|
||||
import { Agent, Agents, MemoryObservabilityAdapter, setObservabilityAdapter } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
// Create memory adapter
|
||||
const obs = new MemoryObservabilityAdapter();
|
||||
setObservabilityAdapter(obs);
|
||||
|
||||
// Create agents
|
||||
const researcher = new Agent({
|
||||
name: 'Researcher',
|
||||
instructions: 'Research topics and gather information.'
|
||||
});
|
||||
|
||||
const writer = new Agent({
|
||||
name: 'Writer',
|
||||
instructions: 'Write clear summaries based on research.'
|
||||
});
|
||||
|
||||
// Create multi-agent workflow
|
||||
const agents = new Agents({
|
||||
agents: [researcher, writer],
|
||||
tasks: [
|
||||
{ agent: researcher, description: 'Research: {topic}' },
|
||||
{ agent: writer, description: 'Summarize the research findings' }
|
||||
]
|
||||
});
|
||||
|
||||
console.log('Running multi-agent workflow with attribution...\n');
|
||||
|
||||
await agents.start({ topic: 'TypeScript best practices' });
|
||||
|
||||
// View traces with attribution
|
||||
const traces = obs.getAllTraces();
|
||||
|
||||
console.log('\n=== Attribution Report ===');
|
||||
for (const trace of traces) {
|
||||
console.log(`\nTrace ID: ${trace.id}`);
|
||||
console.log(` Name: ${trace.name}`);
|
||||
|
||||
if (trace.attribution) {
|
||||
console.log(` Agent ID: ${trace.attribution.agentId || 'N/A'}`);
|
||||
console.log(` Run ID: ${trace.attribution.runId || 'N/A'}`);
|
||||
console.log(` Session ID: ${trace.attribution.sessionId || 'N/A'}`);
|
||||
}
|
||||
|
||||
console.log(` Spans: ${trace.spans.length}`);
|
||||
for (const span of trace.spans) {
|
||||
console.log(` - [${span.kind}] ${span.name}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,40 @@
|
||||
/**
|
||||
* Switchable Observability Tools Example
|
||||
*
|
||||
* Demonstrates switching between observability tools via environment variable.
|
||||
* Set OBS_TOOL env var to: langfuse, langsmith, console, memory, noop
|
||||
*
|
||||
* Usage:
|
||||
* OBS_TOOL=console npx tsx switchable-tools.ts
|
||||
* OBS_TOOL=langfuse LANGFUSE_SECRET_KEY=sk-... npx tsx switchable-tools.ts
|
||||
*/
|
||||
|
||||
import { Agent, createObservabilityAdapter, setObservabilityAdapter } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
// Get tool from environment (default to memory)
|
||||
const toolName = process.env.OBS_TOOL || 'memory';
|
||||
|
||||
console.log(`Using observability tool: ${toolName}\n`);
|
||||
|
||||
// Create adapter dynamically
|
||||
const obs = await createObservabilityAdapter(toolName as any);
|
||||
setObservabilityAdapter(obs);
|
||||
|
||||
// Create agent
|
||||
const agent = new Agent({
|
||||
name: 'ObservedAgent',
|
||||
instructions: 'You are a helpful assistant.'
|
||||
});
|
||||
|
||||
// Run interaction
|
||||
console.log('Running agent...\n');
|
||||
const response = await agent.chat('Hello! What is 2 + 2?');
|
||||
console.log('Response:', response);
|
||||
|
||||
// Flush traces to backend
|
||||
await obs.flush();
|
||||
console.log('\nTraces flushed to', toolName);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,47 @@
|
||||
# Provider Examples
|
||||
|
||||
Examples demonstrating how to use different LLM providers with PraisonAI TypeScript.
|
||||
|
||||
## Setup
|
||||
|
||||
```bash
|
||||
# Install dependencies
|
||||
cd /path/to/praisonai-ts
|
||||
pnpm install
|
||||
|
||||
# Set API keys
|
||||
export OPENAI_API_KEY=sk-...
|
||||
export ANTHROPIC_API_KEY=sk-ant-...
|
||||
export GOOGLE_API_KEY=AIza...
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `multi-provider.ts` | Use multiple providers in one workflow |
|
||||
| `openai-compatible.ts` | Use OpenAI-compatible APIs (LM Studio, etc.) |
|
||||
| `local-ollama.ts` | Use local Ollama models |
|
||||
| `provider-selection.ts` | Dynamic provider selection based on task |
|
||||
|
||||
## Running Examples
|
||||
|
||||
```bash
|
||||
# Run with ts-node or tsx
|
||||
npx tsx examples/js/providers/multi-provider.ts
|
||||
npx tsx examples/js/providers/openai-compatible.ts
|
||||
npx tsx examples/js/providers/local-ollama.ts
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
| Provider | Variable | Example |
|
||||
|----------|----------|---------|
|
||||
| OpenAI | `OPENAI_API_KEY` | `sk-...` |
|
||||
| Anthropic | `ANTHROPIC_API_KEY` | `sk-ant-...` |
|
||||
| Google | `GOOGLE_API_KEY` | `AIza...` |
|
||||
| xAI | `XAI_API_KEY` | `xai-...` |
|
||||
| Groq | `GROQ_API_KEY` | `gsk_...` |
|
||||
| Mistral | `MISTRAL_API_KEY` | `...` |
|
||||
| DeepSeek | `DEEPSEEK_API_KEY` | `...` |
|
||||
| Ollama | `OLLAMA_BASE_URL` | `http://localhost:11434` |
|
||||
@@ -0,0 +1,40 @@
|
||||
/**
|
||||
* Local Ollama Provider Example
|
||||
*
|
||||
* Demonstrates using local Ollama models with PraisonAI.
|
||||
*
|
||||
* Setup:
|
||||
* 1. Install Ollama: https://ollama.ai
|
||||
* 2. Pull a model: ollama pull llama3.2
|
||||
* 3. Run this example
|
||||
*
|
||||
* Optional env var:
|
||||
* - OLLAMA_BASE_URL (default: http://localhost:11434)
|
||||
*/
|
||||
|
||||
import { Agent } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
// Ollama with llama3.2
|
||||
const ollamaAgent = new Agent({
|
||||
name: 'LocalLlama',
|
||||
instructions: 'You are a helpful local AI assistant running on Ollama.',
|
||||
llm: 'ollama/llama3.2',
|
||||
llmConfig: {
|
||||
baseUrl: process.env.OLLAMA_BASE_URL || 'http://localhost:11434'
|
||||
}
|
||||
});
|
||||
|
||||
console.log('Testing Ollama local model...\n');
|
||||
|
||||
try {
|
||||
const response = await ollamaAgent.chat('Hello! What model are you running?');
|
||||
console.log('Response:', response);
|
||||
} catch (error: any) {
|
||||
console.log('Note: Ollama not running or model not pulled.');
|
||||
console.log('Start Ollama and run: ollama pull llama3.2');
|
||||
console.log('Error:', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,57 @@
|
||||
/**
|
||||
* Multi-Provider Example
|
||||
*
|
||||
* Demonstrates using multiple LLM providers in a single workflow.
|
||||
* Each agent uses a different provider optimized for its task.
|
||||
*
|
||||
* Required env vars:
|
||||
* - OPENAI_API_KEY
|
||||
* - ANTHROPIC_API_KEY
|
||||
* - GOOGLE_API_KEY
|
||||
*/
|
||||
|
||||
import { Agent, Agents } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
// Fast model for quick categorization
|
||||
const triageAgent = new Agent({
|
||||
name: 'Triage',
|
||||
instructions: 'Quickly categorize the user request into: technical, creative, or general.',
|
||||
llm: 'openai/gpt-4o-mini' // Fast and cheap
|
||||
});
|
||||
|
||||
// Powerful model for complex reasoning
|
||||
const analysisAgent = new Agent({
|
||||
name: 'Analyst',
|
||||
instructions: 'Analyze the request in depth and provide detailed insights.',
|
||||
llm: 'anthropic/claude-3-5-sonnet' // Best reasoning
|
||||
});
|
||||
|
||||
// Creative model for content generation
|
||||
const writerAgent = new Agent({
|
||||
name: 'Writer',
|
||||
instructions: 'Write a clear, engaging response based on the analysis.',
|
||||
llm: 'google/gemini-2.0-flash' // Good for creative tasks
|
||||
});
|
||||
|
||||
// Create multi-agent workflow
|
||||
const agents = new Agents({
|
||||
agents: [triageAgent, analysisAgent, writerAgent],
|
||||
tasks: [
|
||||
{ agent: triageAgent, description: 'Categorize: {input}' },
|
||||
{ agent: analysisAgent, description: 'Analyze the categorized request' },
|
||||
{ agent: writerAgent, description: 'Write the final response' }
|
||||
]
|
||||
});
|
||||
|
||||
console.log('Starting multi-provider workflow...\n');
|
||||
|
||||
const result = await agents.start({
|
||||
input: 'Explain how neural networks learn'
|
||||
});
|
||||
|
||||
console.log('\n=== Final Result ===');
|
||||
console.log(result);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,48 @@
|
||||
/**
|
||||
* OpenAI-Compatible Provider Example
|
||||
*
|
||||
* Demonstrates using any OpenAI-compatible API endpoint.
|
||||
* Works with LM Studio, vLLM, LocalAI, and other compatible servers.
|
||||
*
|
||||
* Required env vars:
|
||||
* - OPENAI_COMPATIBLE_BASE_URL (e.g., http://localhost:1234/v1)
|
||||
* - OPENAI_COMPATIBLE_API_KEY (optional, use 'not-needed' for local)
|
||||
*/
|
||||
|
||||
import { Agent } from 'praisonai';
|
||||
|
||||
async function main() {
|
||||
// LM Studio example (local)
|
||||
const lmStudioAgent = new Agent({
|
||||
name: 'LocalAssistant',
|
||||
instructions: 'You are a helpful local AI assistant.',
|
||||
llm: 'openai-compatible/local-model',
|
||||
llmConfig: {
|
||||
baseUrl: process.env.LM_STUDIO_BASE_URL || 'http://localhost:1234/v1',
|
||||
apiKey: 'not-needed'
|
||||
}
|
||||
});
|
||||
|
||||
// Custom OpenAI-compatible endpoint
|
||||
const customAgent = new Agent({
|
||||
name: 'CustomAssistant',
|
||||
instructions: 'You are a helpful assistant.',
|
||||
llm: 'openai-compatible/my-model',
|
||||
llmConfig: {
|
||||
baseUrl: process.env.OPENAI_COMPATIBLE_BASE_URL || 'http://localhost:8000/v1',
|
||||
apiKey: process.env.OPENAI_COMPATIBLE_API_KEY || 'not-needed'
|
||||
}
|
||||
});
|
||||
|
||||
console.log('Testing OpenAI-compatible endpoint...\n');
|
||||
|
||||
try {
|
||||
const response = await lmStudioAgent.chat('Hello! What model are you?');
|
||||
console.log('Response:', response);
|
||||
} catch (error: any) {
|
||||
console.log('Note: LM Studio not running. Start it with a model loaded.');
|
||||
console.log('Error:', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,57 @@
|
||||
/**
|
||||
* Basic RAG Agent Example
|
||||
*
|
||||
* Demonstrates building a RAG agent with context injection.
|
||||
*
|
||||
* Prerequisites:
|
||||
* npm install praisonai-ts
|
||||
* export OPENAI_API_KEY=your-api-key
|
||||
*
|
||||
* Run:
|
||||
* npx ts-node basic-rag.ts
|
||||
*/
|
||||
|
||||
import { Agent } from '../../../src/praisonai-ts/src';
|
||||
|
||||
// Simulated knowledge base (in production, use a vector store)
|
||||
const knowledgeBase = [
|
||||
{ id: '1', content: 'PraisonAI is an AI agent framework for building intelligent applications.' },
|
||||
{ id: '2', content: 'Agents can use tools to accomplish tasks like web search and calculations.' },
|
||||
{ id: '3', content: 'RAG combines retrieval with generation for accurate, grounded responses.' },
|
||||
];
|
||||
|
||||
// Simple search function (in production, use vector similarity)
|
||||
function searchKnowledge(query: string): string[] {
|
||||
const queryLower = query.toLowerCase();
|
||||
return knowledgeBase
|
||||
.filter(doc => doc.content.toLowerCase().includes(queryLower.split(' ')[0]))
|
||||
.map(doc => doc.content);
|
||||
}
|
||||
|
||||
async function main() {
|
||||
console.log('=== Basic RAG Agent Example ===\n');
|
||||
|
||||
// Create an agent with RAG-style instructions
|
||||
const agent = new Agent({
|
||||
name: 'RAGAgent',
|
||||
instructions: `You are a helpful assistant that answers questions using provided context.
|
||||
Always base your answers on the context provided.
|
||||
If the context doesn't contain relevant information, say so.`,
|
||||
});
|
||||
|
||||
// Query with context injection
|
||||
const query = 'What is PraisonAI?';
|
||||
console.log(`Q: ${query}`);
|
||||
|
||||
// Retrieve relevant context
|
||||
const context = searchKnowledge(query).join('\n\n');
|
||||
|
||||
// Ask the agent with context
|
||||
const response = await agent.chat(
|
||||
`Context:\n${context}\n\nQuestion: ${query}`
|
||||
);
|
||||
|
||||
console.log(`A: ${response}\n`);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,77 @@
|
||||
/**
|
||||
* Express Server Adapter Example
|
||||
*
|
||||
* Demonstrates deploying an AI agent as an Express HTTP server.
|
||||
*
|
||||
* Prerequisites:
|
||||
* npm install praisonai-ts express
|
||||
* export OPENAI_API_KEY=your-api-key
|
||||
*
|
||||
* Run:
|
||||
* npx ts-node express-server.ts
|
||||
*
|
||||
* Test:
|
||||
* curl -X POST http://localhost:3000/api/chat \
|
||||
* -H "Content-Type: application/json" \
|
||||
* -d '{"message": "Hello!"}'
|
||||
*/
|
||||
|
||||
import { Agent } from '../../../src/praisonai-ts/src';
|
||||
|
||||
// Simple HTTP server without Express dependency for this example
|
||||
import http from 'http';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Express Server Adapter Example ===\n');
|
||||
|
||||
// Create an agent
|
||||
const agent = new Agent({
|
||||
name: 'APIAgent',
|
||||
instructions: 'You are a helpful API assistant. Keep responses concise.',
|
||||
});
|
||||
|
||||
// Create HTTP server
|
||||
const server = http.createServer(async (req, res) => {
|
||||
// CORS headers
|
||||
res.setHeader('Access-Control-Allow-Origin', '*');
|
||||
res.setHeader('Access-Control-Allow-Methods', 'POST, OPTIONS');
|
||||
res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
|
||||
|
||||
if (req.method === 'OPTIONS') {
|
||||
res.writeHead(200);
|
||||
res.end();
|
||||
return;
|
||||
}
|
||||
|
||||
if (req.method === 'POST' && req.url === '/api/chat') {
|
||||
let body = '';
|
||||
req.on('data', chunk => body += chunk);
|
||||
req.on('end', async () => {
|
||||
try {
|
||||
const { message } = JSON.parse(body);
|
||||
const response = await agent.chat(message);
|
||||
|
||||
res.writeHead(200, { 'Content-Type': 'application/json' });
|
||||
res.end(JSON.stringify({ response }));
|
||||
} catch (error: any) {
|
||||
res.writeHead(500, { 'Content-Type': 'application/json' });
|
||||
res.end(JSON.stringify({ error: error.message }));
|
||||
}
|
||||
});
|
||||
} else {
|
||||
res.writeHead(404);
|
||||
res.end('Not Found');
|
||||
}
|
||||
});
|
||||
|
||||
const PORT = process.env.PORT || 3000;
|
||||
server.listen(PORT, () => {
|
||||
console.log(`Server running on http://localhost:${PORT}`);
|
||||
console.log('\nTest with:');
|
||||
console.log(` curl -X POST http://localhost:${PORT}/api/chat \\`);
|
||||
console.log(' -H "Content-Type: application/json" \\');
|
||||
console.log(' -d \'{"message": "Hello!"}\'');
|
||||
});
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,130 @@
|
||||
"""
|
||||
LSP Code Intelligence Example
|
||||
|
||||
This example demonstrates how to use LSP-powered tools through an Agent
|
||||
for semantic code analysis (listing symbols, finding definitions, etc.).
|
||||
|
||||
How to run:
|
||||
export OPENAI_API_KEY=your_key
|
||||
python examples/lsp/lsp_code_analysis_example.py
|
||||
|
||||
Prerequisites:
|
||||
pip install praisonai praisonaiagents
|
||||
pip install python-lsp-server # Optional, for full LSP support
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add paths for development
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src/praisonai-agents'))
|
||||
|
||||
|
||||
async def main():
|
||||
print("=" * 60)
|
||||
print("LSP CODE INTELLIGENCE EXAMPLE")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
# Import after path setup
|
||||
from praisonai.cli.features import (
|
||||
create_agent_centric_tools,
|
||||
InteractiveRuntime,
|
||||
RuntimeConfig
|
||||
)
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Use the praisonaiagents source as workspace for analysis
|
||||
workspace = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
'../../src/praisonai-agents/praisonaiagents'
|
||||
)
|
||||
workspace = os.path.abspath(workspace)
|
||||
|
||||
print(f"Workspace: {workspace}")
|
||||
print()
|
||||
|
||||
# 1. Create runtime with LSP enabled
|
||||
print("1. Creating InteractiveRuntime with LSP...")
|
||||
config = RuntimeConfig(
|
||||
workspace=workspace,
|
||||
lsp_enabled=True, # Enable LSP for code intelligence
|
||||
acp_enabled=False, # Disable ACP for this example
|
||||
approval_mode="auto"
|
||||
)
|
||||
runtime = InteractiveRuntime(config)
|
||||
status = await runtime.start()
|
||||
|
||||
print(f" LSP enabled: {status['lsp']['enabled']}")
|
||||
print(f" LSP ready: {runtime.lsp_ready}")
|
||||
print(f" LSP status: {status['lsp']['status']}")
|
||||
if not runtime.lsp_ready:
|
||||
print(" Note: LSP not available, will use regex fallback")
|
||||
print()
|
||||
|
||||
# 2. Create agent-centric tools
|
||||
print("2. Creating agent-centric tools...")
|
||||
tools = create_agent_centric_tools(runtime)
|
||||
|
||||
# Filter to show only LSP tools
|
||||
lsp_tools = [t for t in tools if t.__name__.startswith('lsp_')]
|
||||
print(f" LSP tools available:")
|
||||
for tool in lsp_tools:
|
||||
print(f" - {tool.__name__}")
|
||||
print()
|
||||
|
||||
# 3. Create Agent with LSP-powered tools
|
||||
print("3. Creating Agent with LSP tools...")
|
||||
agent = Agent(
|
||||
name="CodeAnalyzer",
|
||||
instructions="""You are a code analysis assistant.
|
||||
Use lsp_list_symbols to list functions and classes in files.
|
||||
Use lsp_find_definition to find where symbols are defined.
|
||||
Use lsp_find_references to find where symbols are used.
|
||||
Provide clear, organized summaries of code structure.""",
|
||||
tools=tools,
|
||||
verbose=True
|
||||
)
|
||||
print()
|
||||
|
||||
# 4. Test 1: List symbols in a file
|
||||
print("4. TEST 1: List symbols in agent/agent.py")
|
||||
print("-" * 50)
|
||||
result1 = agent.start("List all the main classes and public methods in agent/agent.py. Give me a brief summary.")
|
||||
print("-" * 50)
|
||||
print(f"Summary: {str(result1)[:300]}...")
|
||||
print()
|
||||
|
||||
# 5. Test 2: Find definition of a symbol
|
||||
print("5. TEST 2: Find where 'Agent' class is defined")
|
||||
print("-" * 50)
|
||||
result2 = agent.start("Where is the Agent class defined? Use lsp_find_definition to find it.")
|
||||
print("-" * 50)
|
||||
print(f"Result: {str(result2)[:300]}...")
|
||||
print()
|
||||
|
||||
# 6. Test 3: Analyze main.py
|
||||
print("6. TEST 3: Analyze main.py structure")
|
||||
print("-" * 50)
|
||||
result3 = agent.start("List all the functions in main.py and briefly describe what the file does.")
|
||||
print("-" * 50)
|
||||
print(f"Analysis: {str(result3)[:300]}...")
|
||||
print()
|
||||
|
||||
# Cleanup
|
||||
await runtime.stop()
|
||||
|
||||
print("=" * 60)
|
||||
print("EXAMPLE COMPLETE")
|
||||
print("=" * 60)
|
||||
print()
|
||||
print("Key takeaways:")
|
||||
print("- LSP tools work through the Agent (agentic LSP)")
|
||||
print("- Falls back to regex if LSP server not available")
|
||||
print("- Returns structured data with citations")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,133 @@
|
||||
"""
|
||||
PraisonAI Output UX Examples
|
||||
|
||||
This example demonstrates the various output modes and configurations
|
||||
available in PraisonAI agents.
|
||||
|
||||
Features demonstrated:
|
||||
- verbose=True: Full output with agent info, task, and response panels
|
||||
- verbose=False: Clean output with just the result
|
||||
- Tool call display with consolidated output
|
||||
- Multi-agent output with clear attribution
|
||||
"""
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, '/Users/praison/praisonai-package/src/praisonai-agents')
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
|
||||
def get_weather(city: str) -> str:
|
||||
"""Get the current weather for a city.
|
||||
|
||||
Args:
|
||||
city: The name of the city to get weather for
|
||||
|
||||
Returns:
|
||||
Weather information string
|
||||
"""
|
||||
return f"The weather in {city} is sunny and 72°F"
|
||||
|
||||
|
||||
def example_verbose_true():
|
||||
"""Example with verbose=True (default) - shows full output panels."""
|
||||
print("\n" + "=" * 60)
|
||||
print("EXAMPLE 1: verbose=True (Full Output)")
|
||||
print("=" * 60)
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant. Be concise.",
|
||||
verbose=True # Default - shows Agent Info, Task, and Response panels
|
||||
)
|
||||
result = agent.start("What is the capital of France? Answer in one sentence.")
|
||||
print(f"\nReturned value: {result}")
|
||||
|
||||
|
||||
def example_verbose_false():
|
||||
"""Example with verbose=False - clean output, just the result."""
|
||||
print("\n" + "=" * 60)
|
||||
print("EXAMPLE 2: verbose=False (Clean Output)")
|
||||
print("=" * 60)
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant. Be concise.",
|
||||
verbose=False # No panels, just returns the result
|
||||
)
|
||||
result = agent.start("What is 2+2? Answer in one word.")
|
||||
print(f"Result: {result}")
|
||||
|
||||
|
||||
def example_with_tools():
|
||||
"""Example showing tool call output."""
|
||||
print("\n" + "=" * 60)
|
||||
print("EXAMPLE 3: Agent with Tools")
|
||||
print("=" * 60)
|
||||
|
||||
agent = Agent(
|
||||
name="Weather Bot",
|
||||
instructions="You help users check the weather. Use the get_weather tool.",
|
||||
tools=[get_weather],
|
||||
verbose=True
|
||||
)
|
||||
result = agent.start("What's the weather in Tokyo?")
|
||||
print(f"\nReturned value: {result}")
|
||||
|
||||
|
||||
def example_multi_agent():
|
||||
"""Example showing multi-agent output with clear attribution."""
|
||||
print("\n" + "=" * 60)
|
||||
print("EXAMPLE 4: Multi-Agent Output")
|
||||
print("=" * 60)
|
||||
|
||||
researcher = Agent(
|
||||
name="Researcher",
|
||||
role="Research Analyst",
|
||||
instructions="You research topics and provide factual information.",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
writer = Agent(
|
||||
name="Writer",
|
||||
role="Content Writer",
|
||||
instructions="You write clear summaries based on research.",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
task1 = Task(
|
||||
description="Research what AI is in one sentence",
|
||||
expected_output="A brief factual description",
|
||||
agent=researcher
|
||||
)
|
||||
|
||||
task2 = Task(
|
||||
description="Write a one-sentence summary about AI",
|
||||
expected_output="A clear summary",
|
||||
agent=writer
|
||||
)
|
||||
|
||||
agents = PraisonAIAgents(
|
||||
agents=[researcher, writer],
|
||||
tasks=[task1, task2],
|
||||
process="sequential",
|
||||
verbose=1
|
||||
)
|
||||
|
||||
result = agents.start()
|
||||
print(f"\nFinal result: {result}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("PRAISONAI OUTPUT UX EXAMPLES")
|
||||
print("=" * 60)
|
||||
|
||||
# Run examples
|
||||
example_verbose_true()
|
||||
example_verbose_false()
|
||||
example_with_tools()
|
||||
example_multi_agent()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("ALL EXAMPLES COMPLETE")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,197 @@
|
||||
"""
|
||||
Isolated Registry Example - Multi-Agent Safety
|
||||
|
||||
This example demonstrates how to use isolated registries to prevent
|
||||
provider registration collisions in multi-agent scenarios.
|
||||
|
||||
Key concepts:
|
||||
- Each agent/context can have its own registry
|
||||
- No global state mutation
|
||||
- Safe for concurrent/parallel runs
|
||||
- Different agents can use different provider implementations
|
||||
|
||||
Requirements:
|
||||
- pip install praisonai
|
||||
"""
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, '../../src/praisonai')
|
||||
|
||||
from praisonai.llm import (
|
||||
LLMProviderRegistry,
|
||||
create_llm_provider,
|
||||
register_llm_provider,
|
||||
get_default_llm_registry,
|
||||
_reset_default_registry
|
||||
)
|
||||
|
||||
|
||||
# Define two different provider implementations with the same name
|
||||
class Agent1Provider:
|
||||
"""Provider implementation for Agent 1."""
|
||||
provider_id = "custom"
|
||||
|
||||
def __init__(self, model_id, config=None):
|
||||
self.model_id = model_id
|
||||
self.config = config or {}
|
||||
self.agent_name = "Agent1"
|
||||
|
||||
def generate(self, prompt):
|
||||
return f"[{self.agent_name}] Response to: {prompt}"
|
||||
|
||||
|
||||
class Agent2Provider:
|
||||
"""Provider implementation for Agent 2 - different behavior."""
|
||||
provider_id = "custom"
|
||||
|
||||
def __init__(self, model_id, config=None):
|
||||
self.model_id = model_id
|
||||
self.config = config or {}
|
||||
self.agent_name = "Agent2"
|
||||
|
||||
def generate(self, prompt):
|
||||
return f"[{self.agent_name}] Different response to: {prompt}"
|
||||
|
||||
|
||||
def demonstrate_collision_problem():
|
||||
"""Show what happens with global registry collisions."""
|
||||
print("=== Problem: Global Registry Collision ===\n")
|
||||
|
||||
# Reset to clean state
|
||||
_reset_default_registry()
|
||||
|
||||
# Agent 1 registers "custom" provider
|
||||
register_llm_provider("custom", Agent1Provider)
|
||||
print("Agent 1 registered 'custom' provider")
|
||||
|
||||
# Agent 2 tries to register same name - ERROR!
|
||||
try:
|
||||
register_llm_provider("custom", Agent2Provider)
|
||||
print("Agent 2 registered 'custom' provider")
|
||||
except ValueError as e:
|
||||
print(f"Agent 2 FAILED: {e}")
|
||||
|
||||
print()
|
||||
print("Problem: Both agents want to use 'custom' but with different implementations.")
|
||||
print("Solution: Use isolated registries.\n")
|
||||
|
||||
|
||||
def demonstrate_isolated_solution():
|
||||
"""Show how isolated registries solve the collision problem."""
|
||||
print("=== Solution: Isolated Registries ===\n")
|
||||
|
||||
# Create isolated registries for each agent
|
||||
agent1_registry = LLMProviderRegistry()
|
||||
agent2_registry = LLMProviderRegistry()
|
||||
|
||||
# Each agent registers its own "custom" provider
|
||||
agent1_registry.register("custom", Agent1Provider)
|
||||
agent2_registry.register("custom", Agent2Provider)
|
||||
|
||||
print("Agent 1 registry providers:", agent1_registry.list())
|
||||
print("Agent 2 registry providers:", agent2_registry.list())
|
||||
print()
|
||||
|
||||
# Create providers from each registry
|
||||
provider1 = create_llm_provider("custom/model-v1", registry=agent1_registry)
|
||||
provider2 = create_llm_provider("custom/model-v1", registry=agent2_registry)
|
||||
|
||||
# They have different implementations!
|
||||
print("Provider 1 response:", provider1.generate("Hello"))
|
||||
print("Provider 2 response:", provider2.generate("Hello"))
|
||||
print()
|
||||
|
||||
# Verify isolation
|
||||
print("Verification:")
|
||||
print(f" provider1.agent_name = {provider1.agent_name}")
|
||||
print(f" provider2.agent_name = {provider2.agent_name}")
|
||||
print(f" Same class? {type(provider1) == type(provider2)}")
|
||||
print()
|
||||
|
||||
|
||||
def demonstrate_parallel_agents():
|
||||
"""Simulate parallel agent execution with isolated registries."""
|
||||
print("=== Parallel Agent Simulation ===\n")
|
||||
|
||||
def run_agent(agent_id, registry):
|
||||
"""Simulate an agent run with its own registry."""
|
||||
provider = create_llm_provider("custom/model", registry=registry)
|
||||
result = provider.generate(f"Query from agent {agent_id}")
|
||||
return result
|
||||
|
||||
# Setup registries
|
||||
registries = {}
|
||||
for i in range(3):
|
||||
reg = LLMProviderRegistry()
|
||||
# Each agent gets a slightly different provider
|
||||
class DynamicProvider:
|
||||
provider_id = "custom"
|
||||
def __init__(self, model_id, config=None):
|
||||
self.model_id = model_id
|
||||
self.agent_id = i # Capture agent ID
|
||||
def generate(self, prompt):
|
||||
return f"Agent-{self.agent_id} processed: {prompt}"
|
||||
|
||||
reg.register("custom", DynamicProvider)
|
||||
registries[i] = reg
|
||||
|
||||
# Simulate parallel execution
|
||||
print("Simulating parallel agent runs:")
|
||||
for agent_id, registry in registries.items():
|
||||
result = run_agent(agent_id, registry)
|
||||
print(f" {result}")
|
||||
|
||||
print()
|
||||
print("Each agent used its own isolated registry - no collisions!")
|
||||
print()
|
||||
|
||||
|
||||
def demonstrate_shared_vs_isolated():
|
||||
"""Compare shared (global) vs isolated registry patterns."""
|
||||
print("=== Shared vs Isolated Patterns ===\n")
|
||||
|
||||
# Reset global registry
|
||||
_reset_default_registry()
|
||||
|
||||
# Pattern 1: Shared (global) registry
|
||||
print("Pattern 1: Shared Registry (use for app-wide providers)")
|
||||
register_llm_provider("shared-provider", Agent1Provider)
|
||||
|
||||
# All code sees the same provider
|
||||
global_reg = get_default_llm_registry()
|
||||
print(f" Global registry has 'shared-provider': {global_reg.has('shared-provider')}")
|
||||
|
||||
# Pattern 2: Isolated registry
|
||||
print("\nPattern 2: Isolated Registry (use for per-agent/per-run providers)")
|
||||
isolated = LLMProviderRegistry()
|
||||
isolated.register("isolated-provider", Agent2Provider)
|
||||
|
||||
# Only code with reference to this registry sees it
|
||||
print(f" Isolated registry has 'isolated-provider': {isolated.has('isolated-provider')}")
|
||||
print(f" Global registry has 'isolated-provider': {global_reg.has('isolated-provider')}")
|
||||
|
||||
print()
|
||||
print("Recommendation:")
|
||||
print(" - Use global registry for app-wide, stable providers")
|
||||
print(" - Use isolated registries for per-agent or per-run customization")
|
||||
print()
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("Isolated Registry Example - Multi-Agent Safety")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
demonstrate_collision_problem()
|
||||
demonstrate_isolated_solution()
|
||||
demonstrate_parallel_agents()
|
||||
demonstrate_shared_vs_isolated()
|
||||
|
||||
print("=" * 60)
|
||||
print("Example Complete")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,83 @@
|
||||
# PraisonAI Profiling Examples
|
||||
|
||||
This directory contains examples for profiling PraisonAI agent performance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install praisonai
|
||||
export OPENAI_API_KEY=your_key_here
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
### 1. Basic Profiling (`basic_profiling.py`)
|
||||
|
||||
Demonstrates programmatic profiling of a single query:
|
||||
|
||||
```bash
|
||||
python basic_profiling.py
|
||||
```
|
||||
|
||||
### 2. Suite Profiling (`suite_profiling.py`)
|
||||
|
||||
Runs a comprehensive profiling suite with multiple scenarios:
|
||||
|
||||
```bash
|
||||
python suite_profiling.py
|
||||
```
|
||||
|
||||
### 3. Optimization Demo (`optimization_example.py`)
|
||||
|
||||
Demonstrates Tier 0/1/2 performance optimizations:
|
||||
|
||||
```bash
|
||||
python optimization_example.py
|
||||
```
|
||||
|
||||
## CLI Commands
|
||||
|
||||
You can also profile directly from the command line:
|
||||
|
||||
```bash
|
||||
# Profile a query
|
||||
praisonai profile query "What is 2+2?"
|
||||
|
||||
# Profile with file grouping
|
||||
praisonai profile query "Hello" --show-files --limit 20
|
||||
|
||||
# Profile imports
|
||||
praisonai profile imports
|
||||
|
||||
# Profile startup time
|
||||
praisonai profile startup
|
||||
|
||||
# Run comprehensive suite
|
||||
praisonai profile suite --quick
|
||||
|
||||
# Create performance baseline
|
||||
praisonai profile snapshot --baseline
|
||||
|
||||
# Compare against baseline
|
||||
praisonai profile snapshot current --compare
|
||||
|
||||
# Show optimization status
|
||||
praisonai profile optimize --show
|
||||
```
|
||||
|
||||
## Output
|
||||
|
||||
- **Text output**: Human-readable timing breakdown
|
||||
- **JSON output**: Machine-readable for CI/CD (`--format json`)
|
||||
- **Artifacts**: Binary cProfile data and reports (`--save ./results`)
|
||||
- **Snapshots**: Performance baselines for regression detection
|
||||
|
||||
## Performance Optimizations
|
||||
|
||||
Enable opt-in optimizations via environment variables:
|
||||
|
||||
```bash
|
||||
export PRAISONAI_LITE_MODE=1
|
||||
export PRAISONAI_SKIP_TYPE_VALIDATION=1
|
||||
export PRAISONAI_MINIMAL_IMPORTS=1
|
||||
```
|
||||
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
Basic Profiling Example for PraisonAI
|
||||
|
||||
This example demonstrates how to use the unified profiler programmatically
|
||||
to measure agent performance.
|
||||
|
||||
Run with:
|
||||
python basic_profiling.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
from praisonai.cli.execution import (
|
||||
ExecutionRequest,
|
||||
Profiler,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Run a profiled query and display results."""
|
||||
|
||||
# Configure profiler (Layer 1 = basic profiling with function stats)
|
||||
config = ProfilerConfig(
|
||||
layer=1, # 0=minimal, 1=basic, 2=deep
|
||||
limit=15, # Show top 15 functions
|
||||
show_callers=False, # Set True for caller info (layer 2)
|
||||
show_callees=False, # Set True for callee info (layer 2)
|
||||
)
|
||||
|
||||
# Create execution request
|
||||
request = ExecutionRequest(
|
||||
prompt="What is 2+2?",
|
||||
agent_name="ProfiledAgent",
|
||||
)
|
||||
|
||||
# Run profiled query
|
||||
print("🔬 Running profiled query...")
|
||||
profiler = Profiler(config)
|
||||
result, report = profiler.profile_sync(request)
|
||||
|
||||
# Print formatted report
|
||||
print("\n" + report.to_text())
|
||||
|
||||
# Access timing data programmatically
|
||||
print("\n📊 Timing Summary:")
|
||||
print(f" Imports: {report.timing.imports_ms:>10.2f} ms")
|
||||
print(f" Agent Init: {report.timing.agent_init_ms:>10.2f} ms")
|
||||
print(f" Execution: {report.timing.execution_ms:>10.2f} ms")
|
||||
print(f" Total: {report.timing.total_ms:>10.2f} ms")
|
||||
|
||||
# Show top 5 functions
|
||||
if report.functions:
|
||||
print("\n🔥 Top 5 Functions by Cumulative Time:")
|
||||
for i, func in enumerate(report.functions[:5], 1):
|
||||
print(f" {i}. {func.name}: {func.cumulative_time_ms:.2f}ms")
|
||||
|
||||
# Show response
|
||||
print(f"\n💬 Response: {result.output}")
|
||||
|
||||
# Export to JSON
|
||||
print("\n📄 JSON Export available via: report.to_json()")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
Performance Optimization Example for PraisonAI
|
||||
|
||||
This example demonstrates how to use the Tier 0/1/2 optimization features
|
||||
programmatically.
|
||||
|
||||
Run with:
|
||||
python optimization_example.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
|
||||
def demo_provider_cache():
|
||||
"""Demonstrate provider caching (Tier 0)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 0: Provider Cache Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
get_provider_cache,
|
||||
resolve_model_provider,
|
||||
)
|
||||
|
||||
cache = get_provider_cache()
|
||||
cache.invalidate() # Start fresh
|
||||
|
||||
# Resolve some models
|
||||
models = ['gpt-4o', 'claude-3-opus', 'gemini-pro', 'gpt-4o-mini']
|
||||
|
||||
print("\nResolving models:")
|
||||
for model in models:
|
||||
provider, name = resolve_model_provider(model)
|
||||
print(f" {model} -> {provider}/{name}")
|
||||
|
||||
# Check cache stats
|
||||
stats = cache.get_stats()
|
||||
print("\nCache Stats:")
|
||||
print(f" Hits: {stats['hits']}")
|
||||
print(f" Misses: {stats['misses']}")
|
||||
print(f" Size: {stats['size']}")
|
||||
|
||||
|
||||
def demo_lazy_importer():
|
||||
"""Demonstrate lazy imports (Tier 0)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 0: Lazy Importer Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import LazyImporter
|
||||
import time
|
||||
|
||||
LazyImporter.clear()
|
||||
|
||||
# Time lazy import
|
||||
start = time.perf_counter()
|
||||
json_module = LazyImporter.get('json')
|
||||
first_import = (time.perf_counter() - start) * 1000
|
||||
|
||||
# Second import (cached)
|
||||
start = time.perf_counter()
|
||||
json_module2 = LazyImporter.get('json')
|
||||
cached_import = (time.perf_counter() - start) * 1000
|
||||
|
||||
print(f"\nFirst import: {first_import:.3f}ms")
|
||||
print(f"Cached import: {cached_import:.3f}ms")
|
||||
print(f"Same object: {json_module is json_module2}")
|
||||
|
||||
|
||||
def demo_client_pool():
|
||||
"""Demonstrate client pooling (Tier 1)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 1: Client Pool Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import ClientPool
|
||||
|
||||
pool = ClientPool(max_size=3)
|
||||
|
||||
# Create some mock clients
|
||||
def create_client(name):
|
||||
return {'name': name, 'created': True}
|
||||
|
||||
# Get or create clients
|
||||
client1 = pool.get_or_create('openai', lambda: create_client('openai'))
|
||||
_ = pool.get_or_create('anthropic', lambda: create_client('anthropic'))
|
||||
|
||||
# Reuse existing client
|
||||
client1_again = pool.get_or_create('openai', lambda: create_client('openai_new'))
|
||||
|
||||
print(f"\nPool size: {pool.size()}")
|
||||
print(f"Client reused: {client1 is client1_again}")
|
||||
print(f"Client1 name: {client1['name']}")
|
||||
|
||||
|
||||
def demo_prewarm():
|
||||
"""Demonstrate pre-warming (Tier 1)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 1: Pre-warming Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import PrewarmManager
|
||||
import time
|
||||
|
||||
PrewarmManager.reset()
|
||||
|
||||
print(f"\nPre-warming enabled: {PrewarmManager.is_enabled()}")
|
||||
|
||||
# Enable and prewarm
|
||||
PrewarmManager.enable()
|
||||
print(f"After enable: {PrewarmManager.is_enabled()}")
|
||||
|
||||
# Prewarm OpenAI (runs in background)
|
||||
PrewarmManager.prewarm_provider('openai')
|
||||
print("Pre-warming OpenAI in background...")
|
||||
|
||||
# Wait a bit for background thread
|
||||
time.sleep(0.5)
|
||||
|
||||
print(f"OpenAI prewarmed: {PrewarmManager.is_prewarmed('openai')}")
|
||||
|
||||
PrewarmManager.reset()
|
||||
|
||||
|
||||
def demo_lite_mode():
|
||||
"""Demonstrate lite mode (Tier 2)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 2: Lite Mode Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
LiteModeConfig,
|
||||
is_lite_mode,
|
||||
)
|
||||
|
||||
# Default config
|
||||
config = LiteModeConfig()
|
||||
print(f"\nDefault lite mode: {config.enabled}")
|
||||
|
||||
# Config from env
|
||||
print("\nTo enable lite mode, set environment variables:")
|
||||
print(" export PRAISONAI_LITE_MODE=1")
|
||||
print(" export PRAISONAI_SKIP_TYPE_VALIDATION=1")
|
||||
|
||||
print(f"\nCurrent is_lite_mode(): {is_lite_mode()}")
|
||||
|
||||
|
||||
def demo_perf_snapshot():
|
||||
"""Demonstrate performance snapshots (Tier 2)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Tier 2: Performance Snapshot Demo")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
PerfSnapshot,
|
||||
PerfComparison,
|
||||
PerfSnapshotManager,
|
||||
)
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
# Create manager with temp directory
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
manager = PerfSnapshotManager(snapshot_dir=Path(tmpdir))
|
||||
|
||||
# Create baseline
|
||||
baseline = PerfSnapshot(
|
||||
timestamp='2025-01-01T00:00:00Z',
|
||||
name='baseline',
|
||||
startup_cold_ms=100.0,
|
||||
startup_warm_ms=80.0,
|
||||
import_time_ms=500.0,
|
||||
query_time_ms=2000.0,
|
||||
first_token_ms=300.0,
|
||||
)
|
||||
manager.save_baseline(baseline)
|
||||
print("\nBaseline saved")
|
||||
|
||||
# Create current snapshot (simulated improvement)
|
||||
current = PerfSnapshot(
|
||||
timestamp='2025-01-02T00:00:00Z',
|
||||
name='current',
|
||||
startup_cold_ms=95.0, # 5% faster
|
||||
startup_warm_ms=78.0,
|
||||
import_time_ms=480.0, # 4% faster
|
||||
query_time_ms=1900.0, # 5% faster
|
||||
first_token_ms=280.0,
|
||||
)
|
||||
|
||||
# Compare
|
||||
comparison = PerfComparison(baseline=baseline, current=current)
|
||||
|
||||
print(f"\nStartup improvement: {-comparison.startup_cold_diff_pct:.1f}%")
|
||||
print(f"Import improvement: {-comparison.import_time_diff_pct:.1f}%")
|
||||
print(f"Query improvement: {-comparison.query_time_diff_pct:.1f}%")
|
||||
print(f"Is regression: {comparison.is_regression()}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all demos."""
|
||||
print("=" * 60)
|
||||
print("PraisonAI Performance Optimization Demo")
|
||||
print("=" * 60)
|
||||
|
||||
demo_provider_cache()
|
||||
demo_lazy_importer()
|
||||
demo_client_pool()
|
||||
demo_prewarm()
|
||||
demo_lite_mode()
|
||||
demo_perf_snapshot()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("All demos complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,227 @@
|
||||
{
|
||||
"timestamp": "2025-12-31T18:45:38.187539Z",
|
||||
"metadata": {
|
||||
"python_version": "3.12.11",
|
||||
"platform": "macOS-15.7.4-arm64-arm-64bit",
|
||||
"praisonai_version": "2.9.2",
|
||||
"timestamp": "2025-12-31T18:45:38.226558Z"
|
||||
},
|
||||
"startup": {
|
||||
"cold_ms": 80.9305000002496,
|
||||
"warm_ms": 55.266541996388696
|
||||
},
|
||||
"import_analysis": [
|
||||
{
|
||||
"module": "praisonaiagents",
|
||||
"self_us": 2419,
|
||||
"cumulative_us": 1645390,
|
||||
"self_ms": 2.419,
|
||||
"cumulative_ms": 1645.39
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.workflows",
|
||||
"self_us": 670,
|
||||
"cumulative_us": 1375919,
|
||||
"self_ms": 0.67,
|
||||
"cumulative_ms": 1375.919
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.workflows.yaml_parser",
|
||||
"self_us": 775,
|
||||
"cumulative_us": 1369435,
|
||||
"self_ms": 0.775,
|
||||
"cumulative_ms": 1369.435
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.agent.agent",
|
||||
"self_us": 19,
|
||||
"cumulative_us": 1349174,
|
||||
"self_ms": 0.019,
|
||||
"cumulative_ms": 1349.174
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.agent",
|
||||
"self_us": 537,
|
||||
"cumulative_us": 1349155,
|
||||
"self_ms": 0.537,
|
||||
"cumulative_ms": 1349.155
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.agent.agent",
|
||||
"self_us": 1453,
|
||||
"cumulative_us": 1285299,
|
||||
"self_ms": 1.453,
|
||||
"cumulative_ms": 1285.299
|
||||
},
|
||||
{
|
||||
"module": "praisonaiagents.llm.openai_client",
|
||||
"self_us": 7887,
|
||||
"cumulative_us": 1123578,
|
||||
"self_ms": 7.887,
|
||||
"cumulative_ms": 1123.578
|
||||
},
|
||||
{
|
||||
"module": "openai",
|
||||
"self_us": 1083,
|
||||
"cumulative_us": 1115692,
|
||||
"self_ms": 1.083,
|
||||
"cumulative_ms": 1115.692
|
||||
},
|
||||
{
|
||||
"module": "openai.types",
|
||||
"self_us": 2606,
|
||||
"cumulative_us": 957990,
|
||||
"self_ms": 2.606,
|
||||
"cumulative_ms": 957.99
|
||||
},
|
||||
{
|
||||
"module": "openai.types.batch",
|
||||
"self_us": 1787,
|
||||
"cumulative_us": 485667,
|
||||
"self_ms": 1.787,
|
||||
"cumulative_ms": 485.667
|
||||
},
|
||||
{
|
||||
"module": "openai._models",
|
||||
"self_us": 13423,
|
||||
"cumulative_us": 450798,
|
||||
"self_ms": 13.423,
|
||||
"cumulative_ms": 450.798
|
||||
},
|
||||
{
|
||||
"module": "openai.types.eval_create_params",
|
||||
"self_us": 644,
|
||||
"cumulative_us": 393275,
|
||||
"self_ms": 0.644,
|
||||
"cumulative_ms": 393.275
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.grader_inputs_param",
|
||||
"self_us": 17,
|
||||
"cumulative_us": 359515,
|
||||
"self_ms": 0.017,
|
||||
"cumulative_ms": 359.515
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders",
|
||||
"self_us": 392,
|
||||
"cumulative_us": 359498,
|
||||
"self_ms": 0.392,
|
||||
"cumulative_ms": 359.498
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.multi_grader",
|
||||
"self_us": 3847,
|
||||
"cumulative_us": 351251,
|
||||
"self_ms": 3.847,
|
||||
"cumulative_ms": 351.251
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.label_model_grader",
|
||||
"self_us": 5450,
|
||||
"cumulative_us": 337291,
|
||||
"self_ms": 5.45,
|
||||
"cumulative_ms": 337.291
|
||||
},
|
||||
{
|
||||
"module": "openai.types.graders.grader_inputs",
|
||||
"self_us": 545,
|
||||
"cumulative_us": 331841,
|
||||
"self_ms": 0.545,
|
||||
"cumulative_ms": 331.841
|
||||
},
|
||||
{
|
||||
"module": "openai.types.responses.response_input_text",
|
||||
"self_us": 21,
|
||||
"cumulative_us": 331297,
|
||||
"self_ms": 0.021,
|
||||
"cumulative_ms": 331.297
|
||||
},
|
||||
{
|
||||
"module": "openai.types.responses",
|
||||
"self_us": 1301,
|
||||
"cumulative_us": 331277,
|
||||
"self_ms": 1.301,
|
||||
"cumulative_ms": 331.277
|
||||
},
|
||||
{
|
||||
"module": "rich.logging",
|
||||
"self_us": 1066,
|
||||
"cumulative_us": 238256,
|
||||
"self_ms": 1.066,
|
||||
"cumulative_ms": 238.256
|
||||
}
|
||||
],
|
||||
"scenarios": [
|
||||
{
|
||||
"name": "simple_query",
|
||||
"config": {
|
||||
"prompt": "What is 2+2?",
|
||||
"model": "default",
|
||||
"stream": false,
|
||||
"iterations": 2
|
||||
},
|
||||
"import_time_stats": {
|
||||
"mean": 0.008833496394800022,
|
||||
"min": 0.005916997906751931,
|
||||
"max": 0.011749994882848114,
|
||||
"stdev": 0.0041245517164382375,
|
||||
"median": 0.008833496394800022
|
||||
},
|
||||
"total_time_stats": {
|
||||
"mean": 1727.4090415012324,
|
||||
"min": 1367.2264170018025,
|
||||
"max": 2087.5916660006624,
|
||||
"stdev": 509.3751524982296,
|
||||
"median": 1727.4090415012324
|
||||
},
|
||||
"first_token_stats": null,
|
||||
"raw_import_times_ms": [
|
||||
0.011749994882848114,
|
||||
0.005916997906751931
|
||||
],
|
||||
"raw_total_times_ms": [
|
||||
1367.2264170018025,
|
||||
2087.5916660006624
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "simple_stream",
|
||||
"config": {
|
||||
"prompt": "What is 2+2?",
|
||||
"model": "default",
|
||||
"stream": true,
|
||||
"iterations": 2
|
||||
},
|
||||
"import_time_stats": {
|
||||
"mean": 0.008792001608526334,
|
||||
"min": 0.006375004886649549,
|
||||
"max": 0.01120899833040312,
|
||||
"stdev": 0.0034181495442894613,
|
||||
"median": 0.008792001608526334
|
||||
},
|
||||
"total_time_stats": {
|
||||
"mean": 2321.0982290001994,
|
||||
"min": 1697.1089580038097,
|
||||
"max": 2945.087499996589,
|
||||
"stdev": 882.4540898183949,
|
||||
"median": 2321.0982290001994
|
||||
},
|
||||
"first_token_stats": {
|
||||
"mean": 2312.8511880022415,
|
||||
"min": 1681.2440420035273,
|
||||
"max": 2944.458334000956,
|
||||
"stdev": 893.2273919631452,
|
||||
"median": 2312.8511880022415
|
||||
},
|
||||
"raw_import_times_ms": [
|
||||
0.01120899833040312,
|
||||
0.006375004886649549
|
||||
],
|
||||
"raw_total_times_ms": [
|
||||
2945.087499996589,
|
||||
1697.1089580038097
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
"""
|
||||
Profile Suite Example for PraisonAI
|
||||
|
||||
This example demonstrates how to run a comprehensive profiling suite
|
||||
programmatically with multiple scenarios.
|
||||
|
||||
Run with:
|
||||
python suite_profiling.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
from praisonai.cli.features.profiler import (
|
||||
ScenarioConfig,
|
||||
run_profile_suite,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Run a profiling suite with custom scenarios."""
|
||||
|
||||
# Define custom scenarios
|
||||
scenarios = [
|
||||
ScenarioConfig(
|
||||
name="simple_query",
|
||||
prompt="What is 2+2?",
|
||||
stream=False,
|
||||
iterations=2,
|
||||
warmup=1,
|
||||
),
|
||||
ScenarioConfig(
|
||||
name="simple_stream",
|
||||
prompt="What is 2+2?",
|
||||
stream=True,
|
||||
iterations=2,
|
||||
warmup=1,
|
||||
),
|
||||
]
|
||||
|
||||
# Run suite with custom output directory
|
||||
print("🔬 Running Profile Suite...")
|
||||
result = run_profile_suite(
|
||||
output_dir="./profile_results",
|
||||
scenarios=scenarios,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Print summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Profile Suite Results")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n📊 Startup Times:")
|
||||
print(f" Cold Start: {result.startup_cold_ms:.2f}ms")
|
||||
print(f" Warm Start: {result.startup_warm_ms:.2f}ms")
|
||||
|
||||
if result.import_analysis:
|
||||
print("\n📦 Top Import:")
|
||||
print(f" {result.import_analysis[0]['module']}: {result.import_analysis[0]['cumulative_ms']:.2f}ms")
|
||||
|
||||
print("\n🎯 Scenario Results:")
|
||||
for scenario in result.scenarios:
|
||||
stats = scenario.get_stats(scenario.total_times)
|
||||
print(f" {scenario.name}:")
|
||||
print(f" Mean: {stats['mean']:.2f}ms")
|
||||
print(f" Min: {stats['min']:.2f}ms")
|
||||
print(f" Max: {stats['max']:.2f}ms")
|
||||
|
||||
print("\n✅ Results saved to: ./profile_results/")
|
||||
print(" - suite_results.json (machine-readable)")
|
||||
print(" - suite_report.txt (human-readable)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
Session Persistence Example
|
||||
|
||||
Demonstrates automatic session persistence with zero configuration.
|
||||
When you provide a session_id to an Agent, conversation history is
|
||||
automatically saved to disk and restored on subsequent runs.
|
||||
|
||||
Usage:
|
||||
python session_persistence_example.py
|
||||
|
||||
Requirements:
|
||||
- praisonaiagents package installed
|
||||
- OPENAI_API_KEY environment variable set
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
def example_basic_persistence():
|
||||
"""Basic session persistence example."""
|
||||
print("=" * 60)
|
||||
print("EXAMPLE 1: Basic Session Persistence")
|
||||
print("=" * 60)
|
||||
|
||||
session_id = "demo-session-basic"
|
||||
|
||||
# Create agent with session_id
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
session_id=session_id,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
# First message
|
||||
print("\n[User]: My favorite color is blue and my lucky number is 7.")
|
||||
response1 = agent.chat("My favorite color is blue and my lucky number is 7.")
|
||||
print(f"[Assistant]: {response1}")
|
||||
|
||||
# Second message - agent remembers
|
||||
print("\n[User]: What is my favorite color?")
|
||||
response2 = agent.chat("What is my favorite color?")
|
||||
print(f"[Assistant]: {response2}")
|
||||
|
||||
print(f"\n✅ Session saved to: ~/.praison/sessions/{session_id}.json")
|
||||
print(f" Chat history length: {len(agent.chat_history)} messages")
|
||||
|
||||
|
||||
def example_session_restoration():
|
||||
"""Demonstrates session restoration across Agent instances."""
|
||||
print("\n" + "=" * 60)
|
||||
print("EXAMPLE 2: Session Restoration")
|
||||
print("=" * 60)
|
||||
|
||||
session_id = "demo-session-restore"
|
||||
|
||||
# First agent instance
|
||||
print("\n--- First Agent Instance ---")
|
||||
agent1 = Agent(
|
||||
name="MemoryBot",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
session_id=session_id,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
print("[User]: Remember that my pet's name is Max.")
|
||||
response1 = agent1.chat("Remember that my pet's name is Max.")
|
||||
print(f"[Assistant]: {response1}")
|
||||
|
||||
# Simulate new process by creating new agent instance
|
||||
print("\n--- New Agent Instance (simulating restart) ---")
|
||||
agent2 = Agent(
|
||||
name="MemoryBot",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
session_id=session_id,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
print("[User]: What is my pet's name?")
|
||||
response2 = agent2.chat("What is my pet's name?")
|
||||
print(f"[Assistant]: {response2}")
|
||||
|
||||
if "max" in response2.lower():
|
||||
print("\n✅ Session restored successfully! Agent remembered 'Max'.")
|
||||
else:
|
||||
print("\n⚠️ Agent may not have remembered correctly.")
|
||||
|
||||
|
||||
def example_in_memory_only():
|
||||
"""Demonstrates in-memory memory without persistence."""
|
||||
print("\n" + "=" * 60)
|
||||
print("EXAMPLE 3: In-Memory Only (No session_id)")
|
||||
print("=" * 60)
|
||||
|
||||
# Agent without session_id - in-memory only
|
||||
agent = Agent(
|
||||
name="TempBot",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
print("\n[User]: My name is Alice.")
|
||||
response1 = agent.chat("My name is Alice.")
|
||||
print(f"[Assistant]: {response1}")
|
||||
|
||||
print("\n[User]: What is my name?")
|
||||
response2 = agent.chat("What is my name?")
|
||||
print(f"[Assistant]: {response2}")
|
||||
|
||||
print(f"\n✅ In-memory chat history: {len(agent.chat_history)} messages")
|
||||
print(" (This will be lost when the Agent instance is garbage collected)")
|
||||
|
||||
|
||||
def example_direct_store_access():
|
||||
"""Demonstrates direct access to the session store."""
|
||||
print("\n" + "=" * 60)
|
||||
print("EXAMPLE 4: Direct Session Store Access")
|
||||
print("=" * 60)
|
||||
|
||||
from praisonaiagents.session import get_default_session_store
|
||||
|
||||
store = get_default_session_store()
|
||||
session_id = "demo-direct-access"
|
||||
|
||||
# Add messages directly
|
||||
store.add_user_message(session_id, "Hello from direct access!")
|
||||
store.add_assistant_message(session_id, "Hi! I received your message.")
|
||||
|
||||
# Get history
|
||||
history = store.get_chat_history(session_id)
|
||||
print(f"\n✅ Added {len(history)} messages directly to session store")
|
||||
|
||||
for msg in history:
|
||||
print(f" [{msg['role']}]: {msg['content']}")
|
||||
|
||||
# List all sessions
|
||||
sessions = store.list_sessions()
|
||||
print(f"\n📋 Total sessions in store: {len(sessions)}")
|
||||
|
||||
# Clean up demo session
|
||||
store.delete_session(session_id)
|
||||
print(f" Deleted demo session: {session_id}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all examples."""
|
||||
print("\n" + "=" * 60)
|
||||
print("SESSION PERSISTENCE EXAMPLES")
|
||||
print("=" * 60)
|
||||
|
||||
# Run examples
|
||||
example_basic_persistence()
|
||||
example_session_restoration()
|
||||
example_in_memory_only()
|
||||
example_direct_store_access()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("ALL EXAMPLES COMPLETE")
|
||||
print("=" * 60)
|
||||
print("\nSession files are stored in: ~/.praison/sessions/")
|
||||
print("You can inspect them with: cat ~/.praison/sessions/<session_id>.json")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,82 @@
|
||||
# RAG Examples
|
||||
|
||||
This directory contains examples demonstrating RAG (Retrieval Augmented Generation) with PraisonAI.
|
||||
|
||||
## Examples
|
||||
|
||||
### 1. 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
|
||||
```
|
||||
|
||||
### 2. 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
|
||||
```
|
||||
|
||||
### 3. RAG Evaluation (`rag_evaluation.py`)
|
||||
|
||||
Evaluation example showing:
|
||||
- Golden query testing
|
||||
- Retrieval accuracy metrics
|
||||
- Answer relevance checking
|
||||
|
||||
```bash
|
||||
python rag_evaluation.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
|
||||
```
|
||||
@@ -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())
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
Basic TUI Usage Example for PraisonAI.
|
||||
|
||||
This example demonstrates how to launch the interactive TUI.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Ensure API key is set
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Please set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
|
||||
def main():
|
||||
"""Launch the TUI."""
|
||||
try:
|
||||
from praisonai.cli.features.tui import run_tui
|
||||
|
||||
# Launch with default settings
|
||||
run_tui()
|
||||
|
||||
except ImportError:
|
||||
print("Textual is required for TUI. Install with:")
|
||||
print(" pip install praisonai[tui]")
|
||||
exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
Headless TUI Simulation Example for PraisonAI.
|
||||
|
||||
This example demonstrates how to run TUI simulations without
|
||||
launching the interactive UI - useful for testing and CI.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonai.cli.features.tui import (
|
||||
TuiOrchestrator,
|
||||
SimulationRunner,
|
||||
MockProvider,
|
||||
)
|
||||
from praisonai.cli.features.tui.orchestrator import OutputMode
|
||||
from praisonai.cli.features.queue import QueueConfig
|
||||
|
||||
|
||||
async def run_simulation():
|
||||
"""Run a headless simulation."""
|
||||
|
||||
# Create orchestrator with mock provider (no real API calls)
|
||||
config = QueueConfig(enable_persistence=False)
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.PRETTY,
|
||||
debug=True,
|
||||
)
|
||||
|
||||
# Create simulation runner
|
||||
runner = SimulationRunner(orchestrator, assert_mode=True)
|
||||
|
||||
# Define simulation script
|
||||
script = {
|
||||
"session_id": "sim-example",
|
||||
"model": "gpt-4o-mini",
|
||||
"steps": [
|
||||
# Navigate to different screens
|
||||
{"action": "navigate", "args": {"screen": "main"}},
|
||||
{"action": "focus", "args": {"widget": "composer"}},
|
||||
|
||||
# Change model
|
||||
{
|
||||
"action": "model",
|
||||
"args": {"model": "gpt-4"},
|
||||
"expected": {"model": "gpt-4"}
|
||||
},
|
||||
|
||||
# Navigate to queue screen
|
||||
{"action": "navigate", "args": {"screen": "queue"}},
|
||||
|
||||
# Take a snapshot
|
||||
{"action": "snapshot"},
|
||||
|
||||
# Navigate back
|
||||
{"action": "navigate", "args": {"screen": "main"}},
|
||||
|
||||
# Small delay
|
||||
{"action": "sleep", "args": {"seconds": 0.5}},
|
||||
]
|
||||
}
|
||||
|
||||
# Run the simulation
|
||||
print("Running headless TUI simulation...")
|
||||
print("=" * 50)
|
||||
|
||||
success = await runner.run_script(script)
|
||||
|
||||
# Print results
|
||||
print("=" * 50)
|
||||
summary = runner.get_summary()
|
||||
|
||||
print(f"\nSimulation {'PASSED' if success else 'FAILED'}")
|
||||
print(f" Assertions passed: {summary['assertions_passed']}")
|
||||
print(f" Assertions failed: {summary['assertions_failed']}")
|
||||
|
||||
if summary['errors']:
|
||||
print("\nErrors:")
|
||||
for error in summary['errors']:
|
||||
print(f" - {error}")
|
||||
|
||||
return success
|
||||
|
||||
|
||||
async def run_with_events():
|
||||
"""Run simulation and capture all events."""
|
||||
|
||||
config = QueueConfig(enable_persistence=False)
|
||||
orchestrator = TuiOrchestrator(
|
||||
queue_config=config,
|
||||
output_mode=OutputMode.SILENT,
|
||||
)
|
||||
|
||||
# Capture events
|
||||
events = []
|
||||
orchestrator.add_event_callback(lambda e: events.append(e))
|
||||
|
||||
await orchestrator.start(session_id="event-capture")
|
||||
|
||||
# Perform some actions
|
||||
orchestrator.set_model("gpt-4")
|
||||
orchestrator.navigate_screen("queue")
|
||||
orchestrator.set_focus("queue-panel")
|
||||
orchestrator.navigate_screen("main")
|
||||
|
||||
await orchestrator.stop()
|
||||
|
||||
# Print captured events
|
||||
print("\nCaptured Events:")
|
||||
print("-" * 40)
|
||||
for event in events:
|
||||
print(f" {event.event_type.value}")
|
||||
|
||||
return events
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point."""
|
||||
print("PraisonAI TUI Headless Simulation Example")
|
||||
print("=" * 50)
|
||||
|
||||
# Run simulation
|
||||
success = asyncio.run(run_simulation())
|
||||
|
||||
# Run event capture
|
||||
print("\n" + "=" * 50)
|
||||
events = asyncio.run(run_with_events())
|
||||
|
||||
print(f"\nTotal events captured: {len(events)}")
|
||||
|
||||
return 0 if success else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
Queue Management Example for PraisonAI.
|
||||
|
||||
This example demonstrates how to use the queue system programmatically.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonai.cli.features.queue import (
|
||||
QueueManager,
|
||||
QueueConfig,
|
||||
QueuedRun,
|
||||
RunState,
|
||||
RunPriority,
|
||||
)
|
||||
|
||||
|
||||
async def main():
|
||||
"""Demonstrate queue management."""
|
||||
|
||||
print("PraisonAI Queue Management Example")
|
||||
print("=" * 50)
|
||||
|
||||
# Create queue configuration
|
||||
config = QueueConfig(
|
||||
max_concurrent_global=2,
|
||||
max_concurrent_per_agent=1,
|
||||
max_queue_size=10,
|
||||
enable_persistence=False, # Use in-memory for demo
|
||||
)
|
||||
|
||||
# Track outputs
|
||||
outputs = {}
|
||||
|
||||
async def on_output(run_id: str, chunk: str):
|
||||
if run_id not in outputs:
|
||||
outputs[run_id] = ""
|
||||
outputs[run_id] += chunk
|
||||
print(f" [{run_id[:8]}] {chunk}", end="", flush=True)
|
||||
|
||||
async def on_complete(run_id: str, run: QueuedRun):
|
||||
print(f"\n [{run_id[:8]}] Completed: {run.state.value}")
|
||||
|
||||
async def on_error(run_id: str, error: Exception):
|
||||
print(f"\n [{run_id[:8]}] Error: {error}")
|
||||
|
||||
# Create queue manager
|
||||
manager = QueueManager(
|
||||
config=config,
|
||||
on_output=on_output,
|
||||
on_complete=on_complete,
|
||||
on_error=on_error,
|
||||
)
|
||||
|
||||
# Start the manager
|
||||
await manager.start(recover=False)
|
||||
print("Queue manager started")
|
||||
|
||||
# Submit some runs with different priorities
|
||||
print("\nSubmitting runs...")
|
||||
|
||||
run1 = await manager.submit(
|
||||
input_content="Low priority task",
|
||||
agent_name="Agent1",
|
||||
priority=RunPriority.LOW,
|
||||
)
|
||||
print(f" Submitted run1: {run1} (LOW priority)")
|
||||
|
||||
run2 = await manager.submit(
|
||||
input_content="High priority task",
|
||||
agent_name="Agent2",
|
||||
priority=RunPriority.HIGH,
|
||||
)
|
||||
print(f" Submitted run2: {run2} (HIGH priority)")
|
||||
|
||||
run3 = await manager.submit(
|
||||
input_content="Normal priority task",
|
||||
agent_name="Agent1",
|
||||
priority=RunPriority.NORMAL,
|
||||
)
|
||||
print(f" Submitted run3: {run3} (NORMAL priority)")
|
||||
|
||||
# Check queue status
|
||||
print(f"\nQueue status:")
|
||||
print(f" Queued: {manager.queued_count}")
|
||||
print(f" Running: {manager.running_count}")
|
||||
|
||||
# List runs
|
||||
runs = manager.list_runs()
|
||||
print(f"\nAll runs ({len(runs)}):")
|
||||
for run in runs:
|
||||
print(f" {run.run_id[:8]}: {run.agent_name} - {run.state.value} ({run.priority.name})")
|
||||
|
||||
# Wait a bit for processing
|
||||
print("\nWaiting for processing...")
|
||||
await asyncio.sleep(2)
|
||||
|
||||
# Cancel a run
|
||||
if runs:
|
||||
cancel_id = runs[-1].run_id
|
||||
print(f"\nCancelling run: {cancel_id[:8]}")
|
||||
await manager.cancel(cancel_id)
|
||||
|
||||
# Get stats
|
||||
stats = manager.get_stats()
|
||||
print(f"\nQueue statistics:")
|
||||
print(f" Total runs: {stats.total_runs}")
|
||||
print(f" Queued: {stats.queued_count}")
|
||||
print(f" Running: {stats.running_count}")
|
||||
print(f" Succeeded: {stats.succeeded_count}")
|
||||
print(f" Failed: {stats.failed_count}")
|
||||
print(f" Cancelled: {stats.cancelled_count}")
|
||||
|
||||
# Stop the manager
|
||||
await manager.stop()
|
||||
print("\nQueue manager stopped")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("Example complete!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,84 @@
|
||||
/**
|
||||
* Exa Web Search Tool Example
|
||||
*
|
||||
* Demonstrates semantic web search with Exa AI.
|
||||
*
|
||||
* Prerequisites:
|
||||
* npm install praisonai-ts @exalabs/ai-sdk
|
||||
* export EXA_API_KEY=your-api-key
|
||||
* export OPENAI_API_KEY=your-api-key
|
||||
*/
|
||||
|
||||
// For local development, use relative imports
|
||||
// For production, use: import { Agent, exaSearch } from 'praisonai-ts'
|
||||
import { Agent } from '../../../src/praisonai-ts/src';
|
||||
import { exaSearch } from '../../../src/praisonai-ts/src/tools/builtins';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Exa Web Search Tool Example ===\n');
|
||||
|
||||
// Example 1: Basic Semantic Search
|
||||
console.log('1. Basic Semantic Search');
|
||||
console.log('------------------------');
|
||||
|
||||
const searchAgent = new Agent({
|
||||
name: 'SemanticSearcher',
|
||||
instructions: 'You search the web using semantic search and provide detailed summaries.',
|
||||
tools: [exaSearch({ numResults: 5, type: 'auto' })],
|
||||
});
|
||||
|
||||
const searchResult = await searchAgent.chat('Find the latest AI research breakthroughs');
|
||||
console.log('Search Result:', searchResult);
|
||||
console.log();
|
||||
|
||||
// Example 2: Company Research
|
||||
console.log('2. Company Research');
|
||||
console.log('-------------------');
|
||||
|
||||
const companyAgent = new Agent({
|
||||
name: 'CompanyResearcher',
|
||||
instructions: 'You research companies and provide analysis.',
|
||||
tools: [
|
||||
exaSearch({
|
||||
type: 'auto',
|
||||
numResults: 5,
|
||||
category: 'company',
|
||||
contents: {
|
||||
text: { maxCharacters: 1000 },
|
||||
summary: true,
|
||||
livecrawl: 'preferred',
|
||||
},
|
||||
}),
|
||||
],
|
||||
});
|
||||
|
||||
const companyResult = await companyAgent.chat(
|
||||
'Find AI startups in Europe founded after 2020'
|
||||
);
|
||||
console.log('Company Result:', companyResult);
|
||||
console.log();
|
||||
|
||||
// Example 3: Academic Paper Search
|
||||
console.log('3. Academic Paper Search');
|
||||
console.log('------------------------');
|
||||
|
||||
const academicAgent = new Agent({
|
||||
name: 'AcademicResearcher',
|
||||
instructions: 'You search for academic papers and summarize findings.',
|
||||
tools: [
|
||||
exaSearch({
|
||||
type: 'neural',
|
||||
category: 'research paper',
|
||||
numResults: 5,
|
||||
includeDomains: ['arxiv.org', 'nature.com', 'science.org'],
|
||||
}),
|
||||
],
|
||||
});
|
||||
|
||||
const academicResult = await academicAgent.chat(
|
||||
'Find recent papers on transformer architectures and attention mechanisms'
|
||||
);
|
||||
console.log('Academic Result:', academicResult);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,86 @@
|
||||
/**
|
||||
* Tavily Search Tool Example
|
||||
*
|
||||
* Demonstrates web search, content extraction, and crawling with Tavily.
|
||||
*
|
||||
* Prerequisites:
|
||||
* npm install praisonai @tavily/ai-sdk
|
||||
* export TAVILY_API_KEY=your-api-key
|
||||
* export OPENAI_API_KEY=your-api-key
|
||||
*/
|
||||
|
||||
// For local development, use relative imports
|
||||
// For production, use: import { Agent, tavilySearch } from 'praisonai-ts'
|
||||
import { Agent } from '../../../src/praisonai-ts/src';
|
||||
import { tavilySearch, tavilyExtract, tavilyCrawl } from '../../../src/praisonai-ts/src/tools/builtins';
|
||||
|
||||
async function main() {
|
||||
console.log('=== Tavily Search Tool Example ===\n');
|
||||
|
||||
// Example 1: Basic Web Search
|
||||
console.log('1. Basic Web Search');
|
||||
console.log('-------------------');
|
||||
|
||||
const searchAgent = new Agent({
|
||||
name: 'WebSearcher',
|
||||
instructions: 'You search the web for information and provide concise summaries with sources.',
|
||||
tools: [tavilySearch({ maxResults: 5, includeAnswer: true })],
|
||||
});
|
||||
|
||||
const searchResult = await searchAgent.chat('What are the latest developments in AI agents?');
|
||||
console.log('Search Result:', searchResult);
|
||||
console.log();
|
||||
|
||||
// Example 2: News Search with Time Filter
|
||||
console.log('2. News Search with Time Filter');
|
||||
console.log('-------------------------------');
|
||||
|
||||
const newsAgent = new Agent({
|
||||
name: 'NewsSearcher',
|
||||
instructions: 'You search for recent news and summarize key points.',
|
||||
tools: [
|
||||
tavilySearch({
|
||||
topic: 'news',
|
||||
timeRange: 'week',
|
||||
maxResults: 5,
|
||||
}),
|
||||
],
|
||||
});
|
||||
|
||||
const newsResult = await newsAgent.chat('Latest news about OpenAI');
|
||||
console.log('News Result:', newsResult);
|
||||
console.log();
|
||||
|
||||
// Example 3: Content Extraction
|
||||
console.log('3. Content Extraction');
|
||||
console.log('--------------------');
|
||||
|
||||
const extractAgent = new Agent({
|
||||
name: 'ContentExtractor',
|
||||
instructions: 'You extract and summarize content from web pages.',
|
||||
tools: [tavilyExtract()],
|
||||
});
|
||||
|
||||
const extractResult = await extractAgent.chat(
|
||||
'Extract the main content from https://docs.tavily.com and summarize it'
|
||||
);
|
||||
console.log('Extract Result:', extractResult);
|
||||
console.log();
|
||||
|
||||
// Example 4: Website Crawling
|
||||
console.log('4. Website Crawling');
|
||||
console.log('------------------');
|
||||
|
||||
const crawlAgent = new Agent({
|
||||
name: 'WebCrawler',
|
||||
instructions: 'You crawl websites and analyze their structure and content.',
|
||||
tools: [tavilyCrawl({ maxDepth: 2 })],
|
||||
});
|
||||
|
||||
const crawlResult = await crawlAgent.chat(
|
||||
'Crawl https://docs.tavily.com and list the main documentation sections'
|
||||
);
|
||||
console.log('Crawl Result:', crawlResult);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -189,7 +189,7 @@ def __getattr__(name):
|
||||
_lazy_cache[name] = AutoAgents
|
||||
return AutoAgents
|
||||
|
||||
# Session (imports requests)
|
||||
# Session (imports requests lazily via session/api.py)
|
||||
elif name == "Session":
|
||||
from .session import Session
|
||||
_lazy_cache[name] = Session
|
||||
@@ -234,6 +234,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
|
||||
@@ -463,4 +481,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
|
||||
|
||||
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,129 @@ 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 _process_knowledge(self, knowledge_item):
|
||||
"""Process and store knowledge from a file path, URL, or string."""
|
||||
try:
|
||||
@@ -2094,6 +2236,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 +2320,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 +2355,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 +3260,31 @@ Write the complete compiled report:"""
|
||||
# Return the single result if only one step
|
||||
return results[0]["result"] if results else None
|
||||
|
||||
def switch_model(self, new_model: str) -> None:
|
||||
"""
|
||||
Switch the agent's LLM model while preserving conversation history.
|
||||
|
||||
Args:
|
||||
new_model: The new model name to switch to (e.g., "gpt-4o", "claude-3-sonnet")
|
||||
"""
|
||||
# Store the new model name
|
||||
self.llm = new_model
|
||||
|
||||
# Recreate the LLM instance with the new model
|
||||
try:
|
||||
from ..llm.llm import LLM
|
||||
self._llm_instance = LLM(
|
||||
model=new_model,
|
||||
base_url=self._openai_base_url,
|
||||
api_key=self._openai_api_key,
|
||||
)
|
||||
self._using_custom_llm = True
|
||||
except ImportError:
|
||||
# If LLM class not available, just update the model string
|
||||
pass
|
||||
|
||||
# Chat history is preserved in self.chat_history (no action needed)
|
||||
|
||||
def start(self, prompt: str, **kwargs):
|
||||
"""Start the agent with a prompt. This is a convenience method that wraps chat()."""
|
||||
# Load history from past sessions if history_in_context is set
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
"""
|
||||
Message Queue for PraisonAI Agents.
|
||||
|
||||
Provides queuing for agent prompts when agent is busy.
|
||||
Thread-safe implementation for concurrent access.
|
||||
"""
|
||||
|
||||
import threading
|
||||
import heapq
|
||||
import time
|
||||
import logging
|
||||
from typing import Optional, Any, List, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MessagePriority(Enum):
|
||||
"""Message priority levels."""
|
||||
LOW = 1
|
||||
NORMAL = 5
|
||||
HIGH = 10
|
||||
URGENT = 20
|
||||
|
||||
|
||||
@dataclass(order=True)
|
||||
class QueuedMessage:
|
||||
"""A message in the queue with priority and timestamp."""
|
||||
priority: int
|
||||
timestamp: float = field(compare=False)
|
||||
content: Any = field(compare=False)
|
||||
metadata: dict = field(default_factory=dict, compare=False)
|
||||
|
||||
def __post_init__(self):
|
||||
# Negate priority for max-heap behavior (higher priority = lower number in heap)
|
||||
self.priority = -self.priority
|
||||
|
||||
|
||||
class AgentMessageQueue:
|
||||
"""
|
||||
Thread-safe message queue for agent prompts.
|
||||
|
||||
Supports priority-based ordering and FIFO within same priority.
|
||||
|
||||
Usage:
|
||||
queue = AgentMessageQueue()
|
||||
queue.enqueue("Hello", priority=5)
|
||||
queue.enqueue("Urgent!", priority=10)
|
||||
|
||||
msg = queue.dequeue() # Returns "Urgent!" first
|
||||
"""
|
||||
|
||||
def __init__(self, max_size: int = 100):
|
||||
"""
|
||||
Initialize message queue.
|
||||
|
||||
Args:
|
||||
max_size: Maximum number of messages in queue (0 = unlimited)
|
||||
"""
|
||||
self._queue: List[QueuedMessage] = []
|
||||
self._lock = threading.RLock()
|
||||
self._not_empty = threading.Condition(self._lock)
|
||||
self._max_size = max_size
|
||||
self._message_count = 0
|
||||
|
||||
def enqueue(
|
||||
self,
|
||||
content: Any,
|
||||
priority: int = MessagePriority.NORMAL.value,
|
||||
metadata: Optional[dict] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Add a message to the queue.
|
||||
|
||||
Args:
|
||||
content: Message content
|
||||
priority: Message priority (higher = more urgent)
|
||||
metadata: Optional metadata dict
|
||||
|
||||
Returns:
|
||||
True if message was added, False if queue is full
|
||||
"""
|
||||
with self._lock:
|
||||
if self._max_size > 0 and len(self._queue) >= self._max_size:
|
||||
logger.warning(f"Message queue full (max_size={self._max_size})")
|
||||
return False
|
||||
|
||||
message = QueuedMessage(
|
||||
priority=priority,
|
||||
timestamp=time.time(),
|
||||
content=content,
|
||||
metadata=metadata or {}
|
||||
)
|
||||
|
||||
heapq.heappush(self._queue, message)
|
||||
self._message_count += 1
|
||||
self._not_empty.notify()
|
||||
|
||||
logger.debug(f"Message enqueued (priority={priority}, queue_size={len(self._queue)})")
|
||||
return True
|
||||
|
||||
def dequeue(self, timeout: Optional[float] = None) -> Optional[Any]:
|
||||
"""
|
||||
Remove and return the highest priority message.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait for a message (None = don't wait)
|
||||
|
||||
Returns:
|
||||
Message content or None if queue is empty
|
||||
"""
|
||||
with self._not_empty:
|
||||
if not self._queue:
|
||||
if timeout is not None:
|
||||
self._not_empty.wait(timeout)
|
||||
|
||||
if not self._queue:
|
||||
return None
|
||||
|
||||
message = heapq.heappop(self._queue)
|
||||
logger.debug(f"Message dequeued (queue_size={len(self._queue)})")
|
||||
return message.content
|
||||
|
||||
def peek(self) -> Optional[Any]:
|
||||
"""Return the highest priority message without removing it."""
|
||||
with self._lock:
|
||||
if not self._queue:
|
||||
return None
|
||||
return self._queue[0].content
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
"""Check if queue is empty."""
|
||||
with self._lock:
|
||||
return len(self._queue) == 0
|
||||
|
||||
def size(self) -> int:
|
||||
"""Get current queue size."""
|
||||
with self._lock:
|
||||
return len(self._queue)
|
||||
|
||||
def clear(self) -> int:
|
||||
"""Clear all messages from queue. Returns number of cleared messages."""
|
||||
with self._lock:
|
||||
count = len(self._queue)
|
||||
self._queue.clear()
|
||||
return count
|
||||
|
||||
def get_all(self) -> List[Any]:
|
||||
"""Get all messages in priority order without removing them."""
|
||||
with self._lock:
|
||||
# Sort by priority (remember priority is negated)
|
||||
sorted_queue = sorted(self._queue)
|
||||
return [msg.content for msg in sorted_queue]
|
||||
|
||||
@property
|
||||
def total_messages_processed(self) -> int:
|
||||
"""Total number of messages ever enqueued."""
|
||||
return self._message_count
|
||||
|
||||
|
||||
class AsyncAgentMessageQueue:
|
||||
"""
|
||||
Async-compatible message queue for agent prompts.
|
||||
|
||||
Wraps AgentMessageQueue with async interface.
|
||||
"""
|
||||
|
||||
def __init__(self, max_size: int = 100):
|
||||
self._sync_queue = AgentMessageQueue(max_size=max_size)
|
||||
|
||||
async def enqueue(
|
||||
self,
|
||||
content: Any,
|
||||
priority: int = MessagePriority.NORMAL.value,
|
||||
metadata: Optional[dict] = None
|
||||
) -> bool:
|
||||
"""Async enqueue."""
|
||||
return self._sync_queue.enqueue(content, priority, metadata)
|
||||
|
||||
async def dequeue(self, timeout: Optional[float] = None) -> Optional[Any]:
|
||||
"""Async dequeue."""
|
||||
import asyncio
|
||||
|
||||
if timeout is None:
|
||||
return self._sync_queue.dequeue(timeout=0)
|
||||
|
||||
# Poll with async sleep for timeout
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
result = self._sync_queue.dequeue(timeout=0)
|
||||
if result is not None:
|
||||
return result
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
return None
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return self._sync_queue.is_empty()
|
||||
|
||||
def size(self) -> int:
|
||||
return self._sync_queue.size()
|
||||
@@ -0,0 +1,226 @@
|
||||
"""
|
||||
Automatic Conversation Summarization for PraisonAI Agents.
|
||||
|
||||
Provides automatic summarization when context window fills up.
|
||||
Reuses existing telemetry/token_collector for token tracking.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any, Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SummarizationConfig:
|
||||
"""Configuration for auto-summarization."""
|
||||
enabled: bool = False
|
||||
threshold: float = 0.8 # 80% of context window
|
||||
context_window: int = 128000 # Default for GPT-4
|
||||
summary_model: Optional[str] = None # Use smaller model for summaries
|
||||
preserve_system: bool = True # Keep system message
|
||||
preserve_recent: int = 2 # Keep last N message pairs
|
||||
|
||||
|
||||
class SummarizationManager:
|
||||
"""
|
||||
Manages automatic conversation summarization.
|
||||
|
||||
Tracks token usage and triggers summarization when threshold is reached.
|
||||
Thread-safe for multi-agent use.
|
||||
"""
|
||||
|
||||
# Context window sizes for common models
|
||||
MODEL_CONTEXT_WINDOWS = {
|
||||
"gpt-4o": 128000,
|
||||
"gpt-4o-mini": 128000,
|
||||
"gpt-4-turbo": 128000,
|
||||
"gpt-4": 8192,
|
||||
"gpt-3.5-turbo": 16385,
|
||||
"claude-3-opus": 200000,
|
||||
"claude-3-sonnet": 200000,
|
||||
"claude-3-haiku": 200000,
|
||||
"claude-3-5-sonnet": 200000,
|
||||
"gemini-pro": 32000,
|
||||
"gemini-1.5-pro": 1000000,
|
||||
"gemini-1.5-flash": 1000000,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
context_window: int = 128000,
|
||||
threshold: float = 0.8,
|
||||
summary_model: Optional[str] = None,
|
||||
preserve_system: bool = True,
|
||||
preserve_recent: int = 2,
|
||||
):
|
||||
"""
|
||||
Initialize summarization manager.
|
||||
|
||||
Args:
|
||||
context_window: Maximum context window size in tokens
|
||||
threshold: Percentage of context window to trigger summarization (0.0-1.0)
|
||||
summary_model: Model to use for generating summaries (None = use same model)
|
||||
preserve_system: Whether to preserve system message after summarization
|
||||
preserve_recent: Number of recent message pairs to preserve
|
||||
"""
|
||||
if not 0.0 < threshold <= 1.0:
|
||||
raise ValueError("threshold must be between 0.0 and 1.0")
|
||||
|
||||
self.context_window = context_window
|
||||
self.threshold = threshold
|
||||
self.summary_model = summary_model
|
||||
self.preserve_system = preserve_system
|
||||
self.preserve_recent = preserve_recent
|
||||
self.current_tokens = 0
|
||||
self._summarization_count = 0
|
||||
|
||||
@classmethod
|
||||
def for_model(cls, model: str, threshold: float = 0.8, **kwargs) -> "SummarizationManager":
|
||||
"""Create a SummarizationManager configured for a specific model."""
|
||||
# Normalize model name
|
||||
model_lower = model.lower()
|
||||
|
||||
# Find matching context window
|
||||
context_window = 128000 # Default
|
||||
for model_key, window in cls.MODEL_CONTEXT_WINDOWS.items():
|
||||
if model_key in model_lower:
|
||||
context_window = window
|
||||
break
|
||||
|
||||
return cls(context_window=context_window, threshold=threshold, **kwargs)
|
||||
|
||||
def add_tokens(self, tokens: int) -> None:
|
||||
"""Add tokens to the current count."""
|
||||
self.current_tokens += tokens
|
||||
|
||||
def set_tokens(self, tokens: int) -> None:
|
||||
"""Set the current token count."""
|
||||
self.current_tokens = tokens
|
||||
|
||||
def reset_tokens(self) -> None:
|
||||
"""Reset token count after summarization."""
|
||||
self.current_tokens = 0
|
||||
|
||||
def should_summarize(self) -> bool:
|
||||
"""Check if summarization should be triggered."""
|
||||
threshold_tokens = int(self.context_window * self.threshold)
|
||||
return self.current_tokens >= threshold_tokens
|
||||
|
||||
def get_usage_percentage(self) -> float:
|
||||
"""Get current context window usage as percentage."""
|
||||
return self.current_tokens / self.context_window
|
||||
|
||||
def generate_summary_prompt(self, messages: List[Dict[str, str]]) -> str:
|
||||
"""Generate prompt for summarizing conversation."""
|
||||
conversation = "\n".join([
|
||||
f"{msg['role'].upper()}: {msg['content']}"
|
||||
for msg in messages
|
||||
if msg['role'] != 'system'
|
||||
])
|
||||
|
||||
return f"""Summarize the following conversation concisely, preserving:
|
||||
1. Key decisions and conclusions
|
||||
2. Important context and facts
|
||||
3. Any pending tasks or questions
|
||||
4. Technical details that may be needed later
|
||||
|
||||
Keep the summary under 500 words.
|
||||
|
||||
CONVERSATION:
|
||||
{conversation}
|
||||
|
||||
SUMMARY:"""
|
||||
|
||||
def prepare_summarized_history(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
summary: str
|
||||
) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Prepare new message history after summarization.
|
||||
|
||||
Args:
|
||||
messages: Original message history
|
||||
summary: Generated summary
|
||||
|
||||
Returns:
|
||||
New message history with summary replacing old messages
|
||||
"""
|
||||
new_history = []
|
||||
|
||||
# Preserve system message if configured
|
||||
if self.preserve_system and messages and messages[0].get('role') == 'system':
|
||||
new_history.append(messages[0])
|
||||
|
||||
# Add summary as assistant message
|
||||
new_history.append({
|
||||
"role": "assistant",
|
||||
"content": f"[Previous conversation summary]\n{summary}"
|
||||
})
|
||||
|
||||
# Preserve recent messages
|
||||
if self.preserve_recent > 0:
|
||||
# Get last N pairs (user + assistant)
|
||||
recent_count = self.preserve_recent * 2
|
||||
recent_messages = [
|
||||
m for m in messages[-recent_count:]
|
||||
if m.get('role') != 'system'
|
||||
]
|
||||
new_history.extend(recent_messages)
|
||||
|
||||
self._summarization_count += 1
|
||||
logger.info(f"Conversation summarized (count: {self._summarization_count})")
|
||||
|
||||
return new_history
|
||||
|
||||
@property
|
||||
def summarization_count(self) -> int:
|
||||
"""Number of times summarization has been performed."""
|
||||
return self._summarization_count
|
||||
|
||||
|
||||
async def summarize_conversation(
|
||||
messages: List[Dict[str, str]],
|
||||
llm_call: Callable,
|
||||
manager: SummarizationManager,
|
||||
) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Summarize conversation using LLM.
|
||||
|
||||
Args:
|
||||
messages: Current message history
|
||||
llm_call: Async function to call LLM
|
||||
manager: SummarizationManager instance
|
||||
|
||||
Returns:
|
||||
New message history with summary
|
||||
"""
|
||||
prompt = manager.generate_summary_prompt(messages)
|
||||
|
||||
# Call LLM to generate summary
|
||||
summary = await llm_call(prompt)
|
||||
|
||||
# Prepare new history
|
||||
new_history = manager.prepare_summarized_history(messages, summary)
|
||||
|
||||
# Reset token count (will be recalculated on next message)
|
||||
manager.reset_tokens()
|
||||
|
||||
return new_history
|
||||
|
||||
|
||||
def summarize_conversation_sync(
|
||||
messages: List[Dict[str, str]],
|
||||
llm_call: Callable,
|
||||
manager: SummarizationManager,
|
||||
) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Synchronous version of summarize_conversation.
|
||||
"""
|
||||
prompt = manager.generate_summary_prompt(messages)
|
||||
summary = llm_call(prompt)
|
||||
new_history = manager.prepare_summarized_history(messages, summary)
|
||||
manager.reset_tokens()
|
||||
return new_history
|
||||
@@ -0,0 +1,238 @@
|
||||
"""
|
||||
Agent Profiles and Modes for PraisonAI Agents.
|
||||
|
||||
Provides built-in agent profiles and mode configurations
|
||||
similar to OpenCode's agent system.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
class AgentMode(str, Enum):
|
||||
"""Agent execution modes."""
|
||||
|
||||
PRIMARY = "primary" # Main agent, can spawn subagents
|
||||
SUBAGENT = "subagent" # Spawned by another agent
|
||||
ALL = "all" # Can be used in any context
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentProfile:
|
||||
"""
|
||||
A built-in agent profile configuration.
|
||||
|
||||
Attributes:
|
||||
name: Unique agent name
|
||||
description: What this agent does
|
||||
mode: Execution mode (primary/subagent/all)
|
||||
system_prompt: System prompt for the agent
|
||||
tools: List of tool names this agent can use
|
||||
temperature: LLM temperature setting
|
||||
top_p: LLM top_p setting
|
||||
max_steps: Maximum execution steps
|
||||
hidden: Whether to hide from agent list
|
||||
color: Display color (for UI)
|
||||
"""
|
||||
|
||||
name: str
|
||||
description: str = ""
|
||||
mode: AgentMode = AgentMode.ALL
|
||||
system_prompt: str = ""
|
||||
tools: List[str] = field(default_factory=list)
|
||||
temperature: float = 0.7
|
||||
top_p: float = 1.0
|
||||
max_steps: int = 50
|
||||
hidden: bool = False
|
||||
color: str = "#3B82F6"
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert profile to dictionary."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"mode": self.mode.value,
|
||||
"system_prompt": self.system_prompt,
|
||||
"tools": self.tools,
|
||||
"temperature": self.temperature,
|
||||
"top_p": self.top_p,
|
||||
"max_steps": self.max_steps,
|
||||
"hidden": self.hidden,
|
||||
"color": self.color,
|
||||
"metadata": self.metadata,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "AgentProfile":
|
||||
"""Create profile from dictionary."""
|
||||
return cls(
|
||||
name=data.get("name", ""),
|
||||
description=data.get("description", ""),
|
||||
mode=AgentMode(data.get("mode", "all")),
|
||||
system_prompt=data.get("system_prompt", ""),
|
||||
tools=data.get("tools", []),
|
||||
temperature=data.get("temperature", 0.7),
|
||||
top_p=data.get("top_p", 1.0),
|
||||
max_steps=data.get("max_steps", 50),
|
||||
hidden=data.get("hidden", False),
|
||||
color=data.get("color", "#3B82F6"),
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
|
||||
|
||||
# Built-in agent profiles
|
||||
BUILTIN_PROFILES: Dict[str, AgentProfile] = {
|
||||
"general": AgentProfile(
|
||||
name="general",
|
||||
description="General-purpose coding assistant",
|
||||
mode=AgentMode.PRIMARY,
|
||||
system_prompt=(
|
||||
"You are a helpful coding assistant. You can help with a wide variety "
|
||||
"of programming tasks including writing code, debugging, explaining "
|
||||
"concepts, and answering questions about software development."
|
||||
),
|
||||
tools=["read_file", "write_file", "bash", "search"],
|
||||
temperature=0.7,
|
||||
color="#3B82F6",
|
||||
),
|
||||
|
||||
"coder": AgentProfile(
|
||||
name="coder",
|
||||
description="Focused code implementation agent",
|
||||
mode=AgentMode.ALL,
|
||||
system_prompt=(
|
||||
"You are an expert code implementation agent. Your focus is on writing "
|
||||
"clean, efficient, and well-documented code. You follow best practices "
|
||||
"and design patterns. When implementing features, you consider edge cases "
|
||||
"and error handling."
|
||||
),
|
||||
tools=["read_file", "write_file", "bash"],
|
||||
temperature=0.3,
|
||||
color="#10B981",
|
||||
),
|
||||
|
||||
"planner": AgentProfile(
|
||||
name="planner",
|
||||
description="Task planning and decomposition agent",
|
||||
mode=AgentMode.SUBAGENT,
|
||||
system_prompt=(
|
||||
"You are a planning agent. Your role is to break down complex tasks "
|
||||
"into smaller, manageable steps. You create clear, actionable plans "
|
||||
"with specific milestones and success criteria. You do not execute "
|
||||
"tasks yourself, only plan them."
|
||||
),
|
||||
tools=["read_file", "search"],
|
||||
temperature=0.5,
|
||||
max_steps=20,
|
||||
color="#F59E0B",
|
||||
),
|
||||
|
||||
"reviewer": AgentProfile(
|
||||
name="reviewer",
|
||||
description="Code review and quality agent",
|
||||
mode=AgentMode.SUBAGENT,
|
||||
system_prompt=(
|
||||
"You are a code review agent. Your role is to review code for quality, "
|
||||
"correctness, security, and best practices. You provide constructive "
|
||||
"feedback and suggest improvements. You look for bugs, security issues, "
|
||||
"and opportunities for optimization."
|
||||
),
|
||||
tools=["read_file", "search"],
|
||||
temperature=0.3,
|
||||
max_steps=30,
|
||||
color="#EF4444",
|
||||
),
|
||||
|
||||
"explorer": AgentProfile(
|
||||
name="explorer",
|
||||
description="Codebase exploration and understanding agent",
|
||||
mode=AgentMode.SUBAGENT,
|
||||
system_prompt=(
|
||||
"You are a codebase exploration agent. Your role is to understand and "
|
||||
"navigate codebases. You can find relevant files, understand code "
|
||||
"structure, and identify patterns. You help locate where changes need "
|
||||
"to be made."
|
||||
),
|
||||
tools=["read_file", "search", "grep"],
|
||||
temperature=0.5,
|
||||
max_steps=40,
|
||||
color="#8B5CF6",
|
||||
),
|
||||
|
||||
"debugger": AgentProfile(
|
||||
name="debugger",
|
||||
description="Debugging and troubleshooting agent",
|
||||
mode=AgentMode.SUBAGENT,
|
||||
system_prompt=(
|
||||
"You are a debugging agent. Your role is to identify and fix bugs. "
|
||||
"You analyze error messages, trace code execution, and identify root "
|
||||
"causes. You suggest minimal fixes that address the issue without "
|
||||
"introducing new problems."
|
||||
),
|
||||
tools=["read_file", "bash", "search"],
|
||||
temperature=0.3,
|
||||
max_steps=30,
|
||||
color="#EC4899",
|
||||
),
|
||||
|
||||
"compaction": AgentProfile(
|
||||
name="compaction",
|
||||
description="Context compaction agent (internal)",
|
||||
mode=AgentMode.SUBAGENT,
|
||||
system_prompt=(
|
||||
"You are a context compaction agent. Your role is to summarize "
|
||||
"conversation history while preserving essential information. "
|
||||
"You create concise summaries that maintain context for continued "
|
||||
"work."
|
||||
),
|
||||
tools=[],
|
||||
temperature=0.3,
|
||||
max_steps=5,
|
||||
hidden=True,
|
||||
color="#6B7280",
|
||||
),
|
||||
|
||||
"title": AgentProfile(
|
||||
name="title",
|
||||
description="Session title generation agent (internal)",
|
||||
mode=AgentMode.SUBAGENT,
|
||||
system_prompt=(
|
||||
"You are a title generation agent. Your role is to create short, "
|
||||
"descriptive titles for conversations based on their content. "
|
||||
"Titles should be concise (3-7 words) and capture the main topic."
|
||||
),
|
||||
tools=[],
|
||||
temperature=0.7,
|
||||
max_steps=1,
|
||||
hidden=True,
|
||||
color="#6B7280",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_profile(name: str) -> Optional[AgentProfile]:
|
||||
"""Get a built-in agent profile by name."""
|
||||
return BUILTIN_PROFILES.get(name)
|
||||
|
||||
|
||||
def list_profiles(include_hidden: bool = False) -> List[AgentProfile]:
|
||||
"""List all built-in agent profiles."""
|
||||
profiles = list(BUILTIN_PROFILES.values())
|
||||
if not include_hidden:
|
||||
profiles = [p for p in profiles if not p.hidden]
|
||||
return profiles
|
||||
|
||||
|
||||
def register_profile(profile: AgentProfile) -> None:
|
||||
"""Register a custom agent profile."""
|
||||
BUILTIN_PROFILES[profile.name] = profile
|
||||
|
||||
|
||||
def get_profiles_by_mode(mode: AgentMode) -> List[AgentProfile]:
|
||||
"""Get all profiles that can be used in a specific mode."""
|
||||
return [
|
||||
p for p in BUILTIN_PROFILES.values()
|
||||
if p.mode == mode or p.mode == AgentMode.ALL
|
||||
]
|
||||
@@ -7,9 +7,12 @@ It extends the existing callback system to require human approval before executi
|
||||
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import Dict, Set, Optional, Callable, Any, Literal
|
||||
import json
|
||||
import os
|
||||
from typing import Dict, Set, Optional, Callable, Any, Literal, List
|
||||
from functools import wraps
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
@@ -28,6 +31,172 @@ approval_callback: Optional[Callable] = None
|
||||
# Context variable to track if we're in an approved execution context
|
||||
_approved_context: ContextVar[Set[str]] = ContextVar('approved_context', default=set())
|
||||
|
||||
# Global permission allowlist
|
||||
_permission_allowlist: Optional["PermissionAllowlist"] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolPermission:
|
||||
"""Permission entry for a tool."""
|
||||
tool_name: str
|
||||
allowed_paths: List[str] = field(default_factory=list)
|
||||
session_only: bool = False # If True, permission expires with session
|
||||
|
||||
|
||||
class PermissionAllowlist:
|
||||
"""
|
||||
Persistent permission allowlist for tools.
|
||||
|
||||
Allows pre-approving tools and paths to skip interactive approval.
|
||||
Can be saved/loaded from JSON for persistence across sessions.
|
||||
|
||||
Usage:
|
||||
allowlist = PermissionAllowlist()
|
||||
allowlist.add_tool("read_file")
|
||||
allowlist.add_tool("write_file", paths=["./src", "./tests"])
|
||||
|
||||
if allowlist.is_allowed("read_file"):
|
||||
# Skip approval prompt
|
||||
pass
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._tools: Dict[str, ToolPermission] = {}
|
||||
self._session_tools: Set[str] = set()
|
||||
|
||||
def add_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
paths: Optional[List[str]] = None,
|
||||
session_only: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Add a tool to the allowlist.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool to allow
|
||||
paths: Optional list of allowed paths (empty = all paths)
|
||||
session_only: If True, permission expires with session
|
||||
"""
|
||||
self._tools[tool_name] = ToolPermission(
|
||||
tool_name=tool_name,
|
||||
allowed_paths=paths or [],
|
||||
session_only=session_only
|
||||
)
|
||||
if session_only:
|
||||
self._session_tools.add(tool_name)
|
||||
|
||||
def remove_tool(self, tool_name: str) -> bool:
|
||||
"""Remove a tool from the allowlist."""
|
||||
if tool_name in self._tools:
|
||||
del self._tools[tool_name]
|
||||
self._session_tools.discard(tool_name)
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_allowed(self, tool_name: str, path: Optional[str] = None) -> bool:
|
||||
"""
|
||||
Check if a tool is allowed.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
path: Optional path to check against allowed paths
|
||||
|
||||
Returns:
|
||||
True if tool is allowed (optionally for the given path)
|
||||
"""
|
||||
if tool_name not in self._tools:
|
||||
return False
|
||||
|
||||
permission = self._tools[tool_name]
|
||||
|
||||
# If no paths specified, tool is allowed for all paths
|
||||
if not permission.allowed_paths:
|
||||
return True
|
||||
|
||||
# If path not provided, check if tool has any permissions
|
||||
if path is None:
|
||||
return True
|
||||
|
||||
# Check if path matches any allowed path
|
||||
for allowed_path in permission.allowed_paths:
|
||||
# Normalize paths for comparison
|
||||
norm_allowed = os.path.normpath(allowed_path)
|
||||
norm_path = os.path.normpath(path)
|
||||
|
||||
# Check if path starts with allowed path
|
||||
if norm_path.startswith(norm_allowed):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
"""Check if allowlist is empty."""
|
||||
return len(self._tools) == 0
|
||||
|
||||
def list_tools(self) -> List[str]:
|
||||
"""List all allowed tools."""
|
||||
return list(self._tools.keys())
|
||||
|
||||
def clear_session_permissions(self) -> None:
|
||||
"""Clear session-only permissions."""
|
||||
for tool_name in list(self._session_tools):
|
||||
if tool_name in self._tools:
|
||||
del self._tools[tool_name]
|
||||
self._session_tools.clear()
|
||||
|
||||
def save(self, filepath: str) -> None:
|
||||
"""Save allowlist to JSON file."""
|
||||
data = {
|
||||
"tools": {
|
||||
name: {
|
||||
"allowed_paths": perm.allowed_paths,
|
||||
"session_only": perm.session_only
|
||||
}
|
||||
for name, perm in self._tools.items()
|
||||
if not perm.session_only # Don't persist session-only permissions
|
||||
}
|
||||
}
|
||||
|
||||
os.makedirs(os.path.dirname(filepath) or ".", exist_ok=True)
|
||||
with open(filepath, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
@classmethod
|
||||
def load(cls, filepath: str) -> "PermissionAllowlist":
|
||||
"""Load allowlist from JSON file."""
|
||||
allowlist = cls()
|
||||
|
||||
if not os.path.exists(filepath):
|
||||
return allowlist
|
||||
|
||||
with open(filepath, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
for tool_name, tool_data in data.get("tools", {}).items():
|
||||
allowlist.add_tool(
|
||||
tool_name,
|
||||
paths=tool_data.get("allowed_paths", []),
|
||||
session_only=tool_data.get("session_only", False)
|
||||
)
|
||||
|
||||
return allowlist
|
||||
|
||||
|
||||
def get_permission_allowlist() -> PermissionAllowlist:
|
||||
"""Get the global permission allowlist."""
|
||||
global _permission_allowlist
|
||||
if _permission_allowlist is None:
|
||||
_permission_allowlist = PermissionAllowlist()
|
||||
return _permission_allowlist
|
||||
|
||||
|
||||
def set_permission_allowlist(allowlist: PermissionAllowlist) -> None:
|
||||
"""Set the global permission allowlist."""
|
||||
global _permission_allowlist
|
||||
_permission_allowlist = allowlist
|
||||
|
||||
|
||||
class ApprovalDecision:
|
||||
"""Result of an approval request"""
|
||||
def __init__(self, approved: bool, modified_args: Optional[Dict[str, Any]] = None, reason: str = ""):
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
Background Job Manager for PraisonAI Agents.
|
||||
|
||||
Provides job lifecycle tracking, auto-backgrounding based on threshold,
|
||||
and status inspection for long-running agent tasks.
|
||||
"""
|
||||
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
from dataclasses import dataclass, field
|
||||
from concurrent.futures import ThreadPoolExecutor, Future
|
||||
|
||||
|
||||
class JobStatus(str, Enum):
|
||||
"""Status of a background job."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobInfo:
|
||||
"""Information about a background job."""
|
||||
job_id: str
|
||||
status: JobStatus
|
||||
created_at: float = field(default_factory=time.time)
|
||||
started_at: Optional[float] = None
|
||||
completed_at: Optional[float] = None
|
||||
result: Any = None
|
||||
error: Optional[str] = None
|
||||
|
||||
@property
|
||||
def duration(self) -> Optional[float]:
|
||||
"""Get job duration in seconds."""
|
||||
if self.started_at is None:
|
||||
return None
|
||||
end_time = self.completed_at or time.time()
|
||||
return end_time - self.started_at
|
||||
|
||||
|
||||
class BackgroundJobManager:
|
||||
"""
|
||||
Manages background jobs for agent tasks.
|
||||
|
||||
Features:
|
||||
- Job lifecycle tracking (pending, running, completed, failed)
|
||||
- Auto-backgrounding based on execution time threshold
|
||||
- Thread-safe job status inspection
|
||||
- Configurable thread pool for concurrent execution
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_workers: int = 4,
|
||||
auto_background_threshold: float = 5.0,
|
||||
):
|
||||
"""
|
||||
Initialize the background job manager.
|
||||
|
||||
Args:
|
||||
max_workers: Maximum number of concurrent background jobs
|
||||
auto_background_threshold: Time in seconds after which a job
|
||||
should be automatically backgrounded
|
||||
"""
|
||||
self.max_workers = max_workers
|
||||
self.auto_background_threshold = auto_background_threshold
|
||||
self._executor = ThreadPoolExecutor(max_workers=max_workers)
|
||||
self._jobs: Dict[str, JobInfo] = {}
|
||||
self._futures: Dict[str, Future] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def active_jobs(self) -> int:
|
||||
"""Get the number of active (pending or running) jobs."""
|
||||
with self._lock:
|
||||
return sum(
|
||||
1 for job in self._jobs.values()
|
||||
if job.status in (JobStatus.PENDING, JobStatus.RUNNING)
|
||||
)
|
||||
|
||||
def start_job(
|
||||
self,
|
||||
func: Callable[[], Any],
|
||||
job_id: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Start a new background job.
|
||||
|
||||
Args:
|
||||
func: The function to execute in the background
|
||||
job_id: Optional custom job ID (auto-generated if not provided)
|
||||
|
||||
Returns:
|
||||
The job ID
|
||||
"""
|
||||
job_id = job_id or str(uuid.uuid4())[:8]
|
||||
|
||||
job_info = JobInfo(
|
||||
job_id=job_id,
|
||||
status=JobStatus.PENDING,
|
||||
)
|
||||
|
||||
with self._lock:
|
||||
self._jobs[job_id] = job_info
|
||||
|
||||
def _run_job():
|
||||
with self._lock:
|
||||
self._jobs[job_id].status = JobStatus.RUNNING
|
||||
self._jobs[job_id].started_at = time.time()
|
||||
|
||||
try:
|
||||
result = func()
|
||||
with self._lock:
|
||||
self._jobs[job_id].status = JobStatus.COMPLETED
|
||||
self._jobs[job_id].completed_at = time.time()
|
||||
self._jobs[job_id].result = result
|
||||
return result
|
||||
except Exception as e:
|
||||
with self._lock:
|
||||
self._jobs[job_id].status = JobStatus.FAILED
|
||||
self._jobs[job_id].completed_at = time.time()
|
||||
self._jobs[job_id].error = str(e)
|
||||
raise
|
||||
|
||||
future = self._executor.submit(_run_job)
|
||||
with self._lock:
|
||||
self._futures[job_id] = future
|
||||
|
||||
return job_id
|
||||
|
||||
def get_status(self, job_id: str) -> JobStatus:
|
||||
"""
|
||||
Get the status of a job.
|
||||
|
||||
Args:
|
||||
job_id: The job ID to check
|
||||
|
||||
Returns:
|
||||
The job status
|
||||
|
||||
Raises:
|
||||
KeyError: If job_id is not found
|
||||
"""
|
||||
with self._lock:
|
||||
if job_id not in self._jobs:
|
||||
raise KeyError(f"Job {job_id} not found")
|
||||
return self._jobs[job_id].status
|
||||
|
||||
def get_job_info(self, job_id: str) -> JobInfo:
|
||||
"""
|
||||
Get full information about a job.
|
||||
|
||||
Args:
|
||||
job_id: The job ID to check
|
||||
|
||||
Returns:
|
||||
The job information
|
||||
|
||||
Raises:
|
||||
KeyError: If job_id is not found
|
||||
"""
|
||||
with self._lock:
|
||||
if job_id not in self._jobs:
|
||||
raise KeyError(f"Job {job_id} not found")
|
||||
return self._jobs[job_id]
|
||||
|
||||
def get_result(self, job_id: str, timeout: Optional[float] = None) -> Any:
|
||||
"""
|
||||
Get the result of a completed job.
|
||||
|
||||
Args:
|
||||
job_id: The job ID
|
||||
timeout: Maximum time to wait for completion
|
||||
|
||||
Returns:
|
||||
The job result
|
||||
|
||||
Raises:
|
||||
KeyError: If job_id is not found
|
||||
TimeoutError: If timeout is reached before completion
|
||||
RuntimeError: If job failed
|
||||
"""
|
||||
with self._lock:
|
||||
if job_id not in self._futures:
|
||||
raise KeyError(f"Job {job_id} not found")
|
||||
future = self._futures[job_id]
|
||||
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except TimeoutError:
|
||||
raise TimeoutError(f"Job {job_id} did not complete within {timeout} seconds")
|
||||
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""
|
||||
Attempt to cancel a job.
|
||||
|
||||
Args:
|
||||
job_id: The job ID to cancel
|
||||
|
||||
Returns:
|
||||
True if cancellation was successful, False otherwise
|
||||
"""
|
||||
with self._lock:
|
||||
if job_id not in self._futures:
|
||||
return False
|
||||
future = self._futures[job_id]
|
||||
cancelled = future.cancel()
|
||||
if cancelled:
|
||||
self._jobs[job_id].status = JobStatus.CANCELLED
|
||||
self._jobs[job_id].completed_at = time.time()
|
||||
return cancelled
|
||||
|
||||
def list_jobs(self, status: Optional[JobStatus] = None) -> Dict[str, JobInfo]:
|
||||
"""
|
||||
List all jobs, optionally filtered by status.
|
||||
|
||||
Args:
|
||||
status: Optional status filter
|
||||
|
||||
Returns:
|
||||
Dictionary of job_id -> JobInfo
|
||||
"""
|
||||
with self._lock:
|
||||
if status is None:
|
||||
return dict(self._jobs)
|
||||
return {
|
||||
job_id: info
|
||||
for job_id, info in self._jobs.items()
|
||||
if info.status == status
|
||||
}
|
||||
|
||||
def cleanup_completed(self, max_age: float = 3600.0) -> int:
|
||||
"""
|
||||
Remove completed jobs older than max_age seconds.
|
||||
|
||||
Args:
|
||||
max_age: Maximum age in seconds for completed jobs
|
||||
|
||||
Returns:
|
||||
Number of jobs removed
|
||||
"""
|
||||
now = time.time()
|
||||
removed = 0
|
||||
|
||||
with self._lock:
|
||||
to_remove = []
|
||||
for job_id, info in self._jobs.items():
|
||||
if info.status in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
|
||||
if info.completed_at and (now - info.completed_at) > max_age:
|
||||
to_remove.append(job_id)
|
||||
|
||||
for job_id in to_remove:
|
||||
del self._jobs[job_id]
|
||||
if job_id in self._futures:
|
||||
del self._futures[job_id]
|
||||
removed += 1
|
||||
|
||||
return removed
|
||||
|
||||
def shutdown(self, wait: bool = True) -> None:
|
||||
"""
|
||||
Shutdown the job manager.
|
||||
|
||||
Args:
|
||||
wait: Whether to wait for pending jobs to complete
|
||||
"""
|
||||
self._executor.shutdown(wait=wait)
|
||||
|
||||
|
||||
# Global job manager instance (lazy initialized)
|
||||
_global_job_manager: Optional[BackgroundJobManager] = None
|
||||
_global_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_job_manager() -> BackgroundJobManager:
|
||||
"""Get the global job manager instance."""
|
||||
global _global_job_manager
|
||||
with _global_lock:
|
||||
if _global_job_manager is None:
|
||||
_global_job_manager = BackgroundJobManager()
|
||||
return _global_job_manager
|
||||
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
Event Bus Module for PraisonAI Agents.
|
||||
|
||||
Provides a typed publish/subscribe event system for real-time communication
|
||||
between components. Extends the existing hooks system with a more general
|
||||
event-driven architecture.
|
||||
|
||||
Features:
|
||||
- Typed event definitions with dataclass payloads
|
||||
- Sync and async subscribers
|
||||
- Event filtering by type
|
||||
- Global and scoped event buses
|
||||
- SSE-compatible event streaming
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded
|
||||
- No overhead when not subscribed
|
||||
- Optional dependency for server features
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.bus import EventBus, Event
|
||||
|
||||
# Create event bus
|
||||
bus = EventBus()
|
||||
|
||||
# Subscribe to events
|
||||
@bus.on("session.created")
|
||||
def handle_session(event):
|
||||
print(f"Session created: {event.data}")
|
||||
|
||||
# Publish events
|
||||
bus.publish("session.created", {"session_id": "abc123"})
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
"EventBus",
|
||||
"Event",
|
||||
"EventType",
|
||||
"Subscriber",
|
||||
"get_default_bus",
|
||||
"set_default_bus",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components."""
|
||||
if name == "EventBus":
|
||||
from .bus import EventBus
|
||||
return EventBus
|
||||
|
||||
if name == "Event":
|
||||
from .event import Event
|
||||
return Event
|
||||
|
||||
if name == "EventType":
|
||||
from .event import EventType
|
||||
return EventType
|
||||
|
||||
if name == "Subscriber":
|
||||
from .bus import Subscriber
|
||||
return Subscriber
|
||||
|
||||
if name == "get_default_bus":
|
||||
from .bus import get_default_bus
|
||||
return get_default_bus
|
||||
|
||||
if name == "set_default_bus":
|
||||
from .bus import set_default_bus
|
||||
return set_default_bus
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,346 @@
|
||||
"""
|
||||
Event Bus implementation for PraisonAI Agents.
|
||||
|
||||
Provides publish/subscribe functionality for typed events.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional, Set, Union
|
||||
from dataclasses import dataclass, field
|
||||
from .event import Event, EventType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Subscriber:
|
||||
"""
|
||||
A subscriber to the event bus.
|
||||
|
||||
Attributes:
|
||||
callback: Function to call when event is received
|
||||
event_types: Set of event types to subscribe to (None = all)
|
||||
is_async: Whether the callback is async
|
||||
id: Unique subscriber identifier
|
||||
"""
|
||||
callback: Callable[[Event], Any]
|
||||
event_types: Optional[Set[str]] = None
|
||||
is_async: bool = False
|
||||
id: str = field(default_factory=lambda: str(__import__("uuid").uuid4()))
|
||||
|
||||
def matches(self, event_type: str) -> bool:
|
||||
"""Check if subscriber should receive this event type."""
|
||||
if self.event_types is None:
|
||||
return True
|
||||
return event_type in self.event_types
|
||||
|
||||
|
||||
class EventBus:
|
||||
"""
|
||||
Event bus for publish/subscribe communication.
|
||||
|
||||
Thread-safe implementation supporting both sync and async subscribers.
|
||||
|
||||
Example:
|
||||
bus = EventBus()
|
||||
|
||||
# Subscribe with decorator
|
||||
@bus.on(EventType.SESSION_CREATED)
|
||||
def handle_session(event):
|
||||
print(f"Session: {event.data}")
|
||||
|
||||
# Subscribe to all events
|
||||
bus.subscribe(lambda e: print(e))
|
||||
|
||||
# Publish event
|
||||
bus.publish(EventType.SESSION_CREATED, {"session_id": "abc"})
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the event bus."""
|
||||
self._subscribers: List[Subscriber] = []
|
||||
self._lock = threading.RLock()
|
||||
self._event_history: List[Event] = []
|
||||
self._max_history = 1000
|
||||
|
||||
def subscribe(
|
||||
self,
|
||||
callback: Callable[[Event], Any],
|
||||
event_types: Optional[Union[str, List[str], Set[str]]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Subscribe to events.
|
||||
|
||||
Args:
|
||||
callback: Function to call when event is received
|
||||
event_types: Optional event type(s) to filter by
|
||||
|
||||
Returns:
|
||||
Subscriber ID for unsubscribing
|
||||
"""
|
||||
# Normalize event_types to a set
|
||||
if event_types is None:
|
||||
types_set = None
|
||||
elif isinstance(event_types, str):
|
||||
types_set = {event_types}
|
||||
else:
|
||||
types_set = set(event_types)
|
||||
|
||||
# Check if callback is async
|
||||
is_async = asyncio.iscoroutinefunction(callback)
|
||||
|
||||
subscriber = Subscriber(
|
||||
callback=callback,
|
||||
event_types=types_set,
|
||||
is_async=is_async,
|
||||
)
|
||||
|
||||
with self._lock:
|
||||
self._subscribers.append(subscriber)
|
||||
|
||||
logger.debug(f"Subscribed {subscriber.id} to {types_set or 'all events'}")
|
||||
return subscriber.id
|
||||
|
||||
def unsubscribe(self, subscriber_id: str) -> bool:
|
||||
"""
|
||||
Unsubscribe from events.
|
||||
|
||||
Args:
|
||||
subscriber_id: The subscriber ID returned from subscribe()
|
||||
|
||||
Returns:
|
||||
True if found and removed, False otherwise
|
||||
"""
|
||||
with self._lock:
|
||||
for i, sub in enumerate(self._subscribers):
|
||||
if sub.id == subscriber_id:
|
||||
self._subscribers.pop(i)
|
||||
logger.debug(f"Unsubscribed {subscriber_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
def on(
|
||||
self,
|
||||
*event_types: Union[str, EventType],
|
||||
) -> Callable:
|
||||
"""
|
||||
Decorator to subscribe a function to events.
|
||||
|
||||
Args:
|
||||
event_types: Event type(s) to subscribe to
|
||||
|
||||
Returns:
|
||||
Decorator function
|
||||
|
||||
Example:
|
||||
@bus.on(EventType.SESSION_CREATED)
|
||||
def handle_session(event):
|
||||
print(event.data)
|
||||
"""
|
||||
# Convert EventType enums to strings
|
||||
types = {
|
||||
t.value if isinstance(t, EventType) else t
|
||||
for t in event_types
|
||||
} if event_types else None
|
||||
|
||||
def decorator(func: Callable[[Event], Any]) -> Callable:
|
||||
self.subscribe(func, types)
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def publish(
|
||||
self,
|
||||
event_type: Union[str, EventType],
|
||||
data: Optional[Dict[str, Any]] = None,
|
||||
source: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Event:
|
||||
"""
|
||||
Publish an event to all matching subscribers.
|
||||
|
||||
Args:
|
||||
event_type: The event type
|
||||
data: Event payload data
|
||||
source: Optional source identifier
|
||||
metadata: Optional additional metadata
|
||||
|
||||
Returns:
|
||||
The published Event object
|
||||
"""
|
||||
# Convert EventType enum to string
|
||||
type_str = event_type.value if isinstance(event_type, EventType) else event_type
|
||||
|
||||
event = Event(
|
||||
type=type_str,
|
||||
data=data or {},
|
||||
source=source,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
return self.publish_event(event)
|
||||
|
||||
def publish_event(self, event: Event) -> Event:
|
||||
"""
|
||||
Publish a pre-constructed event.
|
||||
|
||||
Args:
|
||||
event: The event to publish
|
||||
|
||||
Returns:
|
||||
The published event
|
||||
"""
|
||||
logger.debug(f"Publishing event: {event.type}")
|
||||
|
||||
# Store in history
|
||||
with self._lock:
|
||||
self._event_history.append(event)
|
||||
if len(self._event_history) > self._max_history:
|
||||
self._event_history = self._event_history[-self._max_history:]
|
||||
|
||||
# Get matching subscribers
|
||||
subscribers = [
|
||||
sub for sub in self._subscribers
|
||||
if sub.matches(event.type)
|
||||
]
|
||||
|
||||
# Dispatch to subscribers
|
||||
for sub in subscribers:
|
||||
try:
|
||||
if sub.is_async:
|
||||
# Schedule async callback
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.create_task(sub.callback(event))
|
||||
except RuntimeError:
|
||||
# No running loop, run synchronously
|
||||
asyncio.run(sub.callback(event))
|
||||
else:
|
||||
sub.callback(event)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in subscriber {sub.id}: {e}")
|
||||
|
||||
return event
|
||||
|
||||
async def publish_async(
|
||||
self,
|
||||
event_type: Union[str, EventType],
|
||||
data: Optional[Dict[str, Any]] = None,
|
||||
source: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Event:
|
||||
"""
|
||||
Async version of publish that awaits all async subscribers.
|
||||
|
||||
Args:
|
||||
event_type: The event type
|
||||
data: Event payload data
|
||||
source: Optional source identifier
|
||||
metadata: Optional additional metadata
|
||||
|
||||
Returns:
|
||||
The published Event object
|
||||
"""
|
||||
type_str = event_type.value if isinstance(event_type, EventType) else event_type
|
||||
|
||||
event = Event(
|
||||
type=type_str,
|
||||
data=data or {},
|
||||
source=source,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
logger.debug(f"Publishing async event: {event.type}")
|
||||
|
||||
# Store in history
|
||||
with self._lock:
|
||||
self._event_history.append(event)
|
||||
if len(self._event_history) > self._max_history:
|
||||
self._event_history = self._event_history[-self._max_history:]
|
||||
|
||||
subscribers = [
|
||||
sub for sub in self._subscribers
|
||||
if sub.matches(event.type)
|
||||
]
|
||||
|
||||
# Dispatch to subscribers
|
||||
tasks = []
|
||||
for sub in subscribers:
|
||||
try:
|
||||
if sub.is_async:
|
||||
tasks.append(sub.callback(event))
|
||||
else:
|
||||
sub.callback(event)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in subscriber {sub.id}: {e}")
|
||||
|
||||
# Await all async tasks
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
return event
|
||||
|
||||
def get_history(
|
||||
self,
|
||||
event_type: Optional[str] = None,
|
||||
limit: int = 100,
|
||||
) -> List[Event]:
|
||||
"""
|
||||
Get recent event history.
|
||||
|
||||
Args:
|
||||
event_type: Optional filter by event type
|
||||
limit: Maximum number of events to return
|
||||
|
||||
Returns:
|
||||
List of recent events
|
||||
"""
|
||||
with self._lock:
|
||||
events = self._event_history.copy()
|
||||
|
||||
if event_type:
|
||||
events = [e for e in events if e.type == event_type]
|
||||
|
||||
return events[-limit:]
|
||||
|
||||
def clear_history(self):
|
||||
"""Clear the event history."""
|
||||
with self._lock:
|
||||
self._event_history.clear()
|
||||
|
||||
def clear_subscribers(self):
|
||||
"""Remove all subscribers."""
|
||||
with self._lock:
|
||||
self._subscribers.clear()
|
||||
|
||||
@property
|
||||
def subscriber_count(self) -> int:
|
||||
"""Get the number of subscribers."""
|
||||
with self._lock:
|
||||
return len(self._subscribers)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"EventBus(subscribers={self.subscriber_count})"
|
||||
|
||||
|
||||
# Global default bus
|
||||
_default_bus: Optional[EventBus] = None
|
||||
_bus_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_default_bus() -> EventBus:
|
||||
"""Get the default global event bus."""
|
||||
global _default_bus
|
||||
with _bus_lock:
|
||||
if _default_bus is None:
|
||||
_default_bus = EventBus()
|
||||
return _default_bus
|
||||
|
||||
|
||||
def set_default_bus(bus: EventBus):
|
||||
"""Set the default global event bus."""
|
||||
global _default_bus
|
||||
with _bus_lock:
|
||||
_default_bus = bus
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
Event definitions for the Event Bus.
|
||||
|
||||
Provides typed event structures for the publish/subscribe system.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Optional
|
||||
from enum import Enum
|
||||
import time
|
||||
import uuid
|
||||
|
||||
|
||||
class EventType(str, Enum):
|
||||
"""Standard event types for the agent system."""
|
||||
|
||||
# Session events
|
||||
SESSION_CREATED = "session.created"
|
||||
SESSION_UPDATED = "session.updated"
|
||||
SESSION_DELETED = "session.deleted"
|
||||
SESSION_FORKED = "session.forked"
|
||||
SESSION_REVERTED = "session.reverted"
|
||||
|
||||
# Message events
|
||||
MESSAGE_CREATED = "message.created"
|
||||
MESSAGE_UPDATED = "message.updated"
|
||||
MESSAGE_PART_CREATED = "message.part.created"
|
||||
MESSAGE_PART_UPDATED = "message.part.updated"
|
||||
|
||||
# Permission events
|
||||
PERMISSION_ASKED = "permission.asked"
|
||||
PERMISSION_REPLIED = "permission.replied"
|
||||
|
||||
# Agent events
|
||||
AGENT_STARTED = "agent.started"
|
||||
AGENT_COMPLETED = "agent.completed"
|
||||
AGENT_ERROR = "agent.error"
|
||||
|
||||
# Tool events
|
||||
TOOL_STARTED = "tool.started"
|
||||
TOOL_COMPLETED = "tool.completed"
|
||||
TOOL_ERROR = "tool.error"
|
||||
|
||||
# Snapshot events
|
||||
SNAPSHOT_CREATED = "snapshot.created"
|
||||
SNAPSHOT_RESTORED = "snapshot.restored"
|
||||
|
||||
# Server events
|
||||
SERVER_STARTED = "server.started"
|
||||
SERVER_STOPPED = "server.stopped"
|
||||
CLIENT_CONNECTED = "client.connected"
|
||||
CLIENT_DISCONNECTED = "client.disconnected"
|
||||
|
||||
# Compaction events
|
||||
COMPACTION_STARTED = "compaction.started"
|
||||
COMPACTION_COMPLETED = "compaction.completed"
|
||||
|
||||
# Custom event (for user-defined events)
|
||||
CUSTOM = "custom"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Event:
|
||||
"""
|
||||
A typed event for the event bus.
|
||||
|
||||
Attributes:
|
||||
type: The event type (from EventType enum or custom string)
|
||||
data: The event payload data
|
||||
id: Unique event identifier
|
||||
timestamp: Unix timestamp when event was created
|
||||
source: Optional source identifier (e.g., session_id, agent_name)
|
||||
metadata: Optional additional metadata
|
||||
"""
|
||||
type: str
|
||||
data: Dict[str, Any] = field(default_factory=dict)
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
source: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert event to dictionary for serialization."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"type": self.type,
|
||||
"data": self.data,
|
||||
"timestamp": self.timestamp,
|
||||
"source": self.source,
|
||||
"metadata": self.metadata,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Event":
|
||||
"""Create event from dictionary."""
|
||||
return cls(
|
||||
id=data.get("id", str(uuid.uuid4())),
|
||||
type=data.get("type", EventType.CUSTOM.value),
|
||||
data=data.get("data", {}),
|
||||
timestamp=data.get("timestamp", time.time()),
|
||||
source=data.get("source"),
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
|
||||
def to_sse(self) -> str:
|
||||
"""Format event for Server-Sent Events."""
|
||||
import json
|
||||
return f"id: {self.id}\nevent: {self.type}\ndata: {json.dumps(self.to_dict())}\n\n"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Event(type={self.type!r}, id={self.id[:8]}..., source={self.source!r})"
|
||||
@@ -109,6 +109,10 @@ class ContextCompactor:
|
||||
compacted = self._summarize(messages)
|
||||
elif self.strategy == CompactionStrategy.SMART:
|
||||
compacted = self._smart_compact(messages)
|
||||
elif self.strategy == CompactionStrategy.PRUNE:
|
||||
compacted = self._prune(messages)
|
||||
elif self.strategy == CompactionStrategy.LLM_SUMMARIZE:
|
||||
compacted = self._llm_summarize(messages)
|
||||
else:
|
||||
compacted = self._truncate(messages)
|
||||
|
||||
@@ -223,6 +227,89 @@ class ContextCompactor:
|
||||
# For now, use sliding window as base
|
||||
return self._sliding_window(messages)
|
||||
|
||||
def _prune(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Prune old tool outputs while keeping tool calls.
|
||||
|
||||
This reduces token usage by removing verbose tool outputs
|
||||
from older messages while preserving the context of what
|
||||
tools were called.
|
||||
"""
|
||||
result = []
|
||||
|
||||
# Separate system and other messages
|
||||
system_msgs = [m for m in messages if m.get("role") == "system"]
|
||||
other_msgs = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages intact
|
||||
recent = other_msgs[-self.preserve_recent:]
|
||||
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
|
||||
# Prune older messages
|
||||
for msg in older:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
# If this is a tool result, truncate it
|
||||
if role == "tool" or msg.get("tool_call_id"):
|
||||
if isinstance(content, str) and len(content) > 500:
|
||||
pruned_msg = msg.copy()
|
||||
pruned_msg["content"] = content[:200] + "\n...[output truncated]..."
|
||||
result.append(pruned_msg)
|
||||
else:
|
||||
result.append(msg)
|
||||
else:
|
||||
result.append(msg)
|
||||
|
||||
result.extend(recent)
|
||||
return result
|
||||
|
||||
def _llm_summarize(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Use LLM to summarize older messages.
|
||||
|
||||
Note: This is a placeholder that returns a structured summary.
|
||||
Actual LLM integration should be done at the agent level.
|
||||
"""
|
||||
result = []
|
||||
|
||||
# Keep system messages
|
||||
system_msgs = [m for m in messages if m.get("role") == "system"]
|
||||
other_msgs = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:]
|
||||
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
|
||||
if older:
|
||||
# Create structured summary for LLM to process
|
||||
summary_parts = []
|
||||
for msg in older:
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str) and content:
|
||||
# Extract key information
|
||||
summary_parts.append(f"[{role}]: {content[:150]}...")
|
||||
|
||||
if summary_parts:
|
||||
summary = (
|
||||
"[Compacted conversation history - summarize key points]\n"
|
||||
+ "\n".join(summary_parts[:10])
|
||||
)
|
||||
result.append({
|
||||
"role": "system",
|
||||
"content": summary,
|
||||
"_compacted": True,
|
||||
"_original_count": len(older),
|
||||
})
|
||||
|
||||
result.extend(recent)
|
||||
return result
|
||||
|
||||
def get_stats(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""Get statistics about messages."""
|
||||
total_tokens = self.count_total_tokens(messages)
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
"""
|
||||
Compaction Strategies for PraisonAI Agents.
|
||||
Compaction strategies for context management.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CompactionStrategy(str, Enum):
|
||||
"""Strategy for compacting context."""
|
||||
TRUNCATE = "truncate" # Remove oldest messages
|
||||
SUMMARIZE = "summarize" # Summarize old messages
|
||||
SLIDING = "sliding" # Sliding window of recent messages
|
||||
SMART = "smart" # Intelligent selection based on relevance
|
||||
"""Available compaction strategies."""
|
||||
|
||||
TRUNCATE = "truncate"
|
||||
SLIDING = "sliding"
|
||||
SUMMARIZE = "summarize"
|
||||
SMART = "smart"
|
||||
LLM_SUMMARIZE = "llm_summarize" # Use LLM for summarization
|
||||
PRUNE = "prune" # Remove old tool outputs
|
||||
@@ -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
|
||||
|
||||
@@ -150,7 +150,10 @@ def _clean_display_content(content: str, max_length: int = 20000) -> str:
|
||||
return content.strip()
|
||||
|
||||
def display_interaction(message, response, markdown=True, generation_time=None, console=None, agent_name=None, agent_role=None, agent_tools=None, task_name=None, task_description=None, task_id=None):
|
||||
"""Synchronous version of display_interaction."""
|
||||
"""Synchronous version of display_interaction.
|
||||
|
||||
Displays the task/message and response in clean panels with response time in title.
|
||||
"""
|
||||
if console is None:
|
||||
console = Console()
|
||||
|
||||
@@ -179,16 +182,18 @@ def display_interaction(message, response, markdown=True, generation_time=None,
|
||||
task_description=task_description,
|
||||
task_id=task_id
|
||||
)
|
||||
# Rest of the display logic...
|
||||
|
||||
# Build response title with time (Agno-style: "Response (1.2s)")
|
||||
response_title = "Response"
|
||||
if generation_time:
|
||||
console.print(Text(f"Response generated in {generation_time:.1f}s", style="dim"))
|
||||
response_title = f"Response ({generation_time:.1f}s)"
|
||||
|
||||
if markdown:
|
||||
console.print(Panel.fit(Markdown(message), title="Task", border_style="cyan"))
|
||||
console.print(Panel.fit(Markdown(response), title="Response", border_style="cyan"))
|
||||
console.print(Panel.fit(Markdown(response), title=response_title, border_style="cyan"))
|
||||
else:
|
||||
console.print(Panel.fit(Text(message, style="bold green"), title="Task", border_style="cyan"))
|
||||
console.print(Panel.fit(Text(response, style="bold blue"), title="Response", border_style="cyan"))
|
||||
console.print(Panel.fit(Text(response, style="bold blue"), title=response_title, border_style="cyan"))
|
||||
|
||||
def display_self_reflection(message: str, console=None):
|
||||
if not message or not message.strip():
|
||||
@@ -226,7 +231,16 @@ def display_instruction(message: str, console=None, agent_name: str = None, agen
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
console.print(Panel.fit(Text(message, style="bold blue"), title="Instruction", border_style="cyan"))
|
||||
|
||||
def display_tool_call(message: str, console=None):
|
||||
def display_tool_call(message: str, console=None, tool_name: str = None, tool_input: dict = None, tool_output: str = None):
|
||||
"""Display tool call information in a consolidated panel.
|
||||
|
||||
Args:
|
||||
message: The tool call message (legacy format)
|
||||
console: Rich console for output
|
||||
tool_name: Name of the tool being called
|
||||
tool_input: Input arguments to the tool
|
||||
tool_output: Output from the tool (if available)
|
||||
"""
|
||||
logging.debug(f"display_tool_call called with message: {repr(message)}")
|
||||
if not message or not message.strip():
|
||||
logging.debug("Empty message in display_tool_call, returning early")
|
||||
@@ -235,11 +249,21 @@ def display_tool_call(message: str, console=None):
|
||||
logging.debug(f"Cleaned message in display_tool_call: {repr(message)}")
|
||||
|
||||
# Execute synchronous callbacks (always, even when console is None)
|
||||
execute_sync_callback('tool_call', message=message)
|
||||
execute_sync_callback('tool_call', message=message, tool_name=tool_name, tool_input=tool_input, tool_output=tool_output)
|
||||
|
||||
# Only print panel if console is provided (verbose mode)
|
||||
if console is not None:
|
||||
console.print(Panel.fit(Text(message, style="bold cyan"), title="Tool Call", border_style="green"))
|
||||
# Build consolidated tool call display (Agno-style)
|
||||
if tool_name and tool_input is not None:
|
||||
# Format: tool_name(arg1=val1, arg2=val2)
|
||||
args_str = ", ".join(f"{k}={v}" for k, v in tool_input.items()) if tool_input else ""
|
||||
tool_display = f"• {tool_name}({args_str})"
|
||||
if tool_output:
|
||||
tool_display += f"\n → {tool_output}"
|
||||
console.print(Panel.fit(Text(tool_display, style="bold cyan"), title="Tool Call", border_style="green"))
|
||||
else:
|
||||
# Legacy format
|
||||
console.print(Panel.fit(Text(message, style="bold cyan"), title="Tool Call", border_style="green"))
|
||||
|
||||
def display_error(message: str, console=None):
|
||||
if not message or not message.strip():
|
||||
|
||||
@@ -11,8 +11,15 @@ import platform
|
||||
from typing import Any, List, Optional, Callable, Iterable, Union
|
||||
from functools import wraps, partial
|
||||
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
try:
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
MCP_AVAILABLE = True
|
||||
except ImportError:
|
||||
MCP_AVAILABLE = False
|
||||
ClientSession = None
|
||||
StdioServerParameters = None
|
||||
stdio_client = None
|
||||
|
||||
class MCPToolRunner(threading.Thread):
|
||||
"""A dedicated thread for running MCP operations."""
|
||||
@@ -196,10 +203,17 @@ class MCP:
|
||||
debug: Enable debug logging for MCP operations (default: False)
|
||||
**kwargs: Additional parameters for StdioServerParameters
|
||||
"""
|
||||
# Check if MCP is available
|
||||
if not MCP_AVAILABLE:
|
||||
raise ImportError(
|
||||
"MCP (Model Context Protocol) package is not installed. "
|
||||
"Install it with: pip install praisonaiagents[mcp]"
|
||||
)
|
||||
|
||||
# Handle backward compatibility with named parameter 'command'
|
||||
if command_or_string is None and command is not None:
|
||||
command_or_string = command
|
||||
|
||||
|
||||
# Set up logging - default to WARNING level to hide INFO messages
|
||||
if debug:
|
||||
logging.getLogger("mcp-wrapper").setLevel(logging.DEBUG)
|
||||
|
||||
@@ -16,11 +16,17 @@ import weakref
|
||||
from typing import List, Dict, Any, Optional, Callable, Iterable, Union
|
||||
from urllib.parse import urlparse, urljoin
|
||||
|
||||
from mcp import ClientSession
|
||||
try:
|
||||
from mcp import ClientSession
|
||||
MCP_AVAILABLE = True
|
||||
except ImportError:
|
||||
MCP_AVAILABLE = False
|
||||
ClientSession = None
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
except ImportError:
|
||||
raise ImportError("aiohttp is required for HTTP Stream transport. Install with: pip install praisonaiagents[mcp]")
|
||||
aiohttp = None
|
||||
|
||||
logger = logging.getLogger("mcp-http-stream")
|
||||
|
||||
@@ -507,9 +513,23 @@ class HTTPStreamMCPClient:
|
||||
timeout: Timeout in seconds for operations (default: 60)
|
||||
options: Additional configuration options for the transport
|
||||
"""
|
||||
# Check if MCP is available
|
||||
if not MCP_AVAILABLE:
|
||||
raise ImportError(
|
||||
"MCP (Model Context Protocol) package is not installed. "
|
||||
"Install it with: pip install praisonaiagents[mcp]"
|
||||
)
|
||||
|
||||
# Check if aiohttp is available
|
||||
if aiohttp is None:
|
||||
raise ImportError(
|
||||
"aiohttp is required for HTTP Stream transport. "
|
||||
"Install it with: pip install praisonaiagents[mcp]"
|
||||
)
|
||||
|
||||
# Parse URL to extract base URL and endpoint
|
||||
parsed = urlparse(server_url)
|
||||
|
||||
|
||||
# If the URL already has a path, use it; otherwise use default /mcp endpoint
|
||||
if parsed.path and parsed.path != '/':
|
||||
self.base_url = server_url
|
||||
|
||||
@@ -11,8 +11,14 @@ import inspect
|
||||
import json
|
||||
from typing import List, Dict, Any, Optional, Callable, Iterable
|
||||
|
||||
from mcp import ClientSession
|
||||
from mcp.client.sse import sse_client
|
||||
try:
|
||||
from mcp import ClientSession
|
||||
from mcp.client.sse import sse_client
|
||||
MCP_AVAILABLE = True
|
||||
except ImportError:
|
||||
MCP_AVAILABLE = False
|
||||
ClientSession = None
|
||||
sse_client = None
|
||||
|
||||
logger = logging.getLogger("mcp-sse")
|
||||
|
||||
@@ -169,12 +175,19 @@ class SSEMCPClient:
|
||||
debug: Whether to enable debug logging
|
||||
timeout: Timeout in seconds for operations (default: 60)
|
||||
"""
|
||||
# Check if MCP is available
|
||||
if not MCP_AVAILABLE:
|
||||
raise ImportError(
|
||||
"MCP (Model Context Protocol) package is not installed. "
|
||||
"Install it with: pip install praisonaiagents[mcp]"
|
||||
)
|
||||
|
||||
self.server_url = server_url
|
||||
self.debug = debug
|
||||
self.timeout = timeout
|
||||
self.session = None
|
||||
self.tools = []
|
||||
|
||||
|
||||
# Set up logging
|
||||
if debug:
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
@@ -236,3 +236,46 @@ def functions_to_mcp_tools(functions: List[Callable]) -> List[Dict[str, Any]]:
|
||||
List of MCP tool schemas
|
||||
"""
|
||||
return [function_to_mcp_schema(func) for func in functions]
|
||||
|
||||
|
||||
def filter_disabled_tools(
|
||||
tools: List[Dict[str, Any]],
|
||||
disabled_tools: List[str]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Filter out disabled tools from a tool list.
|
||||
|
||||
Args:
|
||||
tools: List of tool definitions (dicts with 'name' key)
|
||||
disabled_tools: List of tool names to disable
|
||||
|
||||
Returns:
|
||||
Filtered list of tools
|
||||
"""
|
||||
if not disabled_tools:
|
||||
return tools
|
||||
|
||||
disabled_set = set(disabled_tools)
|
||||
return [t for t in tools if t.get("name") not in disabled_set]
|
||||
|
||||
|
||||
def filter_tools_by_allowlist(
|
||||
tools: List[Dict[str, Any]],
|
||||
allowed_tools: List[str]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Filter tools to only include those in allowlist.
|
||||
|
||||
Args:
|
||||
tools: List of tool definitions (dicts with 'name' key)
|
||||
allowed_tools: List of tool names to allow (empty = allow all)
|
||||
|
||||
Returns:
|
||||
Filtered list of tools
|
||||
"""
|
||||
if not allowed_tools:
|
||||
return tools
|
||||
|
||||
allowed_set = set(allowed_tools)
|
||||
return [t for t in tools if t.get("name") in allowed_set]
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
"""
|
||||
Permissions Module for PraisonAI Agents.
|
||||
|
||||
Provides pattern-based permission rules, persistent approvals,
|
||||
and doom loop detection for safe agent execution.
|
||||
|
||||
Features:
|
||||
- Pattern-based permission rules (allow, deny, ask)
|
||||
- Persistent approval storage
|
||||
- Per-agent permission rulesets
|
||||
- Doom loop detection and prevention
|
||||
- Integration with existing approval system
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.permissions import PermissionManager, PermissionRule
|
||||
|
||||
# Create permission manager
|
||||
manager = PermissionManager()
|
||||
|
||||
# Add rules
|
||||
manager.add_rule(PermissionRule(
|
||||
pattern="bash:*",
|
||||
action="ask",
|
||||
description="Require approval for shell commands"
|
||||
))
|
||||
|
||||
# Check permission
|
||||
result = manager.check("bash:rm -rf /tmp/test")
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
"PermissionManager",
|
||||
"PermissionRule",
|
||||
"PermissionAction",
|
||||
"PermissionResult",
|
||||
"DoomLoopDetector",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components."""
|
||||
if name == "PermissionManager":
|
||||
from .manager import PermissionManager
|
||||
return PermissionManager
|
||||
|
||||
if name == "PermissionRule":
|
||||
from .rules import PermissionRule
|
||||
return PermissionRule
|
||||
|
||||
if name == "PermissionAction":
|
||||
from .rules import PermissionAction
|
||||
return PermissionAction
|
||||
|
||||
if name == "PermissionResult":
|
||||
from .rules import PermissionResult
|
||||
return PermissionResult
|
||||
|
||||
if name == "DoomLoopDetector":
|
||||
from .doom_loop import DoomLoopDetector
|
||||
return DoomLoopDetector
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,256 @@
|
||||
"""
|
||||
Doom Loop Detection for PraisonAI Agents.
|
||||
|
||||
Detects and prevents agents from getting stuck in repetitive loops.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default thresholds
|
||||
DEFAULT_LOOP_THRESHOLD = 3
|
||||
DEFAULT_WINDOW_SECONDS = 60
|
||||
DEFAULT_MAX_TOOL_CALLS = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCallRecord:
|
||||
"""Record of a tool call for loop detection."""
|
||||
|
||||
tool_name: str
|
||||
arguments_hash: str
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
session_id: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"tool_name": self.tool_name,
|
||||
"arguments_hash": self.arguments_hash,
|
||||
"timestamp": self.timestamp,
|
||||
"session_id": self.session_id,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class DoomLoopResult:
|
||||
"""Result of a doom loop check."""
|
||||
|
||||
is_loop: bool
|
||||
reason: str = ""
|
||||
loop_count: int = 0
|
||||
tool_name: Optional[str] = None
|
||||
recommendation: str = ""
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"is_loop": self.is_loop,
|
||||
"reason": self.reason,
|
||||
"loop_count": self.loop_count,
|
||||
"tool_name": self.tool_name,
|
||||
"recommendation": self.recommendation,
|
||||
}
|
||||
|
||||
|
||||
class DoomLoopDetector:
|
||||
"""
|
||||
Detects and prevents doom loops in agent execution.
|
||||
|
||||
A doom loop occurs when an agent repeatedly calls the same
|
||||
tool with the same arguments, indicating it's stuck.
|
||||
|
||||
Example:
|
||||
detector = DoomLoopDetector()
|
||||
|
||||
# Record tool calls
|
||||
detector.record("bash", {"command": "ls"})
|
||||
detector.record("bash", {"command": "ls"})
|
||||
detector.record("bash", {"command": "ls"})
|
||||
|
||||
# Check for loop
|
||||
result = detector.check("bash", {"command": "ls"})
|
||||
if result.is_loop:
|
||||
print(f"Doom loop detected: {result.reason}")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
loop_threshold: int = DEFAULT_LOOP_THRESHOLD,
|
||||
window_seconds: float = DEFAULT_WINDOW_SECONDS,
|
||||
max_tool_calls: int = DEFAULT_MAX_TOOL_CALLS,
|
||||
):
|
||||
"""
|
||||
Initialize the doom loop detector.
|
||||
|
||||
Args:
|
||||
loop_threshold: Number of identical calls to trigger loop detection
|
||||
window_seconds: Time window for counting calls
|
||||
max_tool_calls: Maximum total tool calls before warning
|
||||
"""
|
||||
self.loop_threshold = loop_threshold
|
||||
self.window_seconds = window_seconds
|
||||
self.max_tool_calls = max_tool_calls
|
||||
|
||||
self._records: List[ToolCallRecord] = []
|
||||
self._call_counts: Dict[str, int] = defaultdict(int)
|
||||
self._session_id: Optional[str] = None
|
||||
|
||||
def set_session(self, session_id: str):
|
||||
"""Set the current session ID."""
|
||||
self._session_id = session_id
|
||||
|
||||
def _hash_arguments(self, arguments: Dict[str, Any]) -> str:
|
||||
"""Create a hash of tool arguments."""
|
||||
import hashlib
|
||||
import json
|
||||
|
||||
try:
|
||||
arg_str = json.dumps(arguments, sort_keys=True, default=str)
|
||||
return hashlib.md5(arg_str.encode()).hexdigest()[:16]
|
||||
except (TypeError, ValueError):
|
||||
return "unhashable"
|
||||
|
||||
def _cleanup_old_records(self):
|
||||
"""Remove records outside the time window."""
|
||||
cutoff = time.time() - self.window_seconds
|
||||
self._records = [r for r in self._records if r.timestamp > cutoff]
|
||||
|
||||
def record(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: Optional[Dict[str, Any]] = None,
|
||||
) -> ToolCallRecord:
|
||||
"""
|
||||
Record a tool call.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
arguments: Tool arguments
|
||||
|
||||
Returns:
|
||||
The recorded ToolCallRecord
|
||||
"""
|
||||
self._cleanup_old_records()
|
||||
|
||||
args_hash = self._hash_arguments(arguments or {})
|
||||
|
||||
record = ToolCallRecord(
|
||||
tool_name=tool_name,
|
||||
arguments_hash=args_hash,
|
||||
session_id=self._session_id,
|
||||
)
|
||||
|
||||
self._records.append(record)
|
||||
self._call_counts[tool_name] += 1
|
||||
|
||||
return record
|
||||
|
||||
def check(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: Optional[Dict[str, Any]] = None,
|
||||
) -> DoomLoopResult:
|
||||
"""
|
||||
Check if calling this tool would indicate a doom loop.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
arguments: Tool arguments
|
||||
|
||||
Returns:
|
||||
DoomLoopResult indicating if a loop is detected
|
||||
"""
|
||||
self._cleanup_old_records()
|
||||
|
||||
args_hash = self._hash_arguments(arguments or {})
|
||||
|
||||
# Count identical calls in window
|
||||
identical_count = sum(
|
||||
1 for r in self._records
|
||||
if r.tool_name == tool_name and r.arguments_hash == args_hash
|
||||
)
|
||||
|
||||
# Check for loop
|
||||
if identical_count >= self.loop_threshold:
|
||||
return DoomLoopResult(
|
||||
is_loop=True,
|
||||
reason=f"Tool '{tool_name}' called {identical_count} times with same arguments",
|
||||
loop_count=identical_count,
|
||||
tool_name=tool_name,
|
||||
recommendation="Consider a different approach or ask for user guidance",
|
||||
)
|
||||
|
||||
# Check for excessive tool calls
|
||||
total_calls = len(self._records)
|
||||
if total_calls >= self.max_tool_calls:
|
||||
return DoomLoopResult(
|
||||
is_loop=True,
|
||||
reason=f"Excessive tool calls ({total_calls}) in time window",
|
||||
loop_count=total_calls,
|
||||
recommendation="Reduce tool usage or break task into smaller steps",
|
||||
)
|
||||
|
||||
# Check for tool-specific excessive calls
|
||||
tool_count = self._call_counts.get(tool_name, 0)
|
||||
if tool_count >= self.max_tool_calls // 2:
|
||||
return DoomLoopResult(
|
||||
is_loop=True,
|
||||
reason=f"Tool '{tool_name}' called excessively ({tool_count} times)",
|
||||
loop_count=tool_count,
|
||||
tool_name=tool_name,
|
||||
recommendation=f"Consider alternative to '{tool_name}'",
|
||||
)
|
||||
|
||||
return DoomLoopResult(
|
||||
is_loop=False,
|
||||
loop_count=identical_count,
|
||||
tool_name=tool_name,
|
||||
)
|
||||
|
||||
def record_and_check(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: Optional[Dict[str, Any]] = None,
|
||||
) -> DoomLoopResult:
|
||||
"""
|
||||
Record a tool call and check for doom loop.
|
||||
|
||||
Convenience method that combines record() and check().
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
arguments: Tool arguments
|
||||
|
||||
Returns:
|
||||
DoomLoopResult indicating if a loop is detected
|
||||
"""
|
||||
result = self.check(tool_name, arguments)
|
||||
if not result.is_loop:
|
||||
self.record(tool_name, arguments)
|
||||
return result
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get statistics about tool calls."""
|
||||
self._cleanup_old_records()
|
||||
|
||||
return {
|
||||
"total_calls": len(self._records),
|
||||
"unique_tools": len(set(r.tool_name for r in self._records)),
|
||||
"call_counts": dict(self._call_counts),
|
||||
"window_seconds": self.window_seconds,
|
||||
"loop_threshold": self.loop_threshold,
|
||||
}
|
||||
|
||||
def reset(self):
|
||||
"""Reset all records."""
|
||||
self._records.clear()
|
||||
self._call_counts.clear()
|
||||
|
||||
def reset_tool(self, tool_name: str):
|
||||
"""Reset records for a specific tool."""
|
||||
self._records = [r for r in self._records if r.tool_name != tool_name]
|
||||
self._call_counts.pop(tool_name, None)
|
||||
@@ -0,0 +1,359 @@
|
||||
"""
|
||||
Permission Manager for PraisonAI Agents.
|
||||
|
||||
Manages permission rules, persistent approvals, and integrates
|
||||
with doom loop detection.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from .rules import PermissionRule, PermissionAction, PermissionResult, PersistentApproval
|
||||
from .doom_loop import DoomLoopDetector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default storage path
|
||||
DEFAULT_PERMISSIONS_DIR = os.path.expanduser("~/.praison/permissions")
|
||||
|
||||
|
||||
class PermissionManager:
|
||||
"""
|
||||
Manages permission rules and approvals.
|
||||
|
||||
Provides pattern-based permission checking with persistent
|
||||
approval storage and doom loop detection.
|
||||
|
||||
Example:
|
||||
manager = PermissionManager()
|
||||
|
||||
# Add rules
|
||||
manager.add_rule(PermissionRule(
|
||||
pattern="bash:rm *",
|
||||
action=PermissionAction.ASK,
|
||||
description="Require approval for rm commands"
|
||||
))
|
||||
|
||||
manager.add_rule(PermissionRule(
|
||||
pattern="read:*",
|
||||
action=PermissionAction.ALLOW,
|
||||
description="Allow all read operations"
|
||||
))
|
||||
|
||||
# Check permission
|
||||
result = manager.check("bash:rm -rf /tmp/test")
|
||||
if result.needs_approval:
|
||||
# Ask user
|
||||
approved = ask_user(result.target)
|
||||
manager.approve(result.target, approved)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage_dir: Optional[str] = None,
|
||||
agent_name: Optional[str] = None,
|
||||
approval_callback: Optional[Callable[[str, str], bool]] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the permission manager.
|
||||
|
||||
Args:
|
||||
storage_dir: Directory for persistent storage
|
||||
agent_name: Optional agent name for filtering rules
|
||||
approval_callback: Optional callback for user approval
|
||||
"""
|
||||
self.storage_dir = storage_dir or DEFAULT_PERMISSIONS_DIR
|
||||
self.agent_name = agent_name
|
||||
self.approval_callback = approval_callback
|
||||
|
||||
self._rules: List[PermissionRule] = []
|
||||
self._approvals: List[PersistentApproval] = []
|
||||
self._doom_detector = DoomLoopDetector()
|
||||
self._lock = threading.RLock()
|
||||
|
||||
# Ensure storage directory exists
|
||||
os.makedirs(self.storage_dir, exist_ok=True)
|
||||
|
||||
# Load persistent data
|
||||
self._load_rules()
|
||||
self._load_approvals()
|
||||
|
||||
def _get_rules_path(self) -> str:
|
||||
"""Get path to rules file."""
|
||||
return os.path.join(self.storage_dir, "rules.json")
|
||||
|
||||
def _get_approvals_path(self) -> str:
|
||||
"""Get path to approvals file."""
|
||||
return os.path.join(self.storage_dir, "approvals.json")
|
||||
|
||||
def _load_rules(self):
|
||||
"""Load rules from disk."""
|
||||
path = self._get_rules_path()
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
with open(path, "r") as f:
|
||||
data = json.load(f)
|
||||
self._rules = [PermissionRule.from_dict(r) for r in data]
|
||||
except (json.JSONDecodeError, IOError) as e:
|
||||
logger.warning(f"Failed to load rules: {e}")
|
||||
|
||||
def _save_rules(self):
|
||||
"""Save rules to disk."""
|
||||
path = self._get_rules_path()
|
||||
try:
|
||||
with open(path, "w") as f:
|
||||
json.dump([r.to_dict() for r in self._rules], f, indent=2)
|
||||
except IOError as e:
|
||||
logger.error(f"Failed to save rules: {e}")
|
||||
|
||||
def _load_approvals(self):
|
||||
"""Load approvals from disk."""
|
||||
path = self._get_approvals_path()
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
with open(path, "r") as f:
|
||||
data = json.load(f)
|
||||
self._approvals = [PersistentApproval.from_dict(a) for a in data]
|
||||
# Filter out expired approvals
|
||||
self._approvals = [a for a in self._approvals if a.is_valid()]
|
||||
except (json.JSONDecodeError, IOError) as e:
|
||||
logger.warning(f"Failed to load approvals: {e}")
|
||||
|
||||
def _save_approvals(self):
|
||||
"""Save approvals to disk."""
|
||||
path = self._get_approvals_path()
|
||||
try:
|
||||
# Filter out expired and one-time approvals
|
||||
persistent = [
|
||||
a for a in self._approvals
|
||||
if a.is_valid() and a.scope in ("session", "always")
|
||||
]
|
||||
with open(path, "w") as f:
|
||||
json.dump([a.to_dict() for a in persistent], f, indent=2)
|
||||
except IOError as e:
|
||||
logger.error(f"Failed to save approvals: {e}")
|
||||
|
||||
def add_rule(self, rule: PermissionRule) -> str:
|
||||
"""
|
||||
Add a permission rule.
|
||||
|
||||
Args:
|
||||
rule: The rule to add
|
||||
|
||||
Returns:
|
||||
The rule ID
|
||||
"""
|
||||
with self._lock:
|
||||
self._rules.append(rule)
|
||||
# Sort by priority (higher first)
|
||||
self._rules.sort(key=lambda r: r.priority, reverse=True)
|
||||
self._save_rules()
|
||||
return rule.id
|
||||
|
||||
def remove_rule(self, rule_id: str) -> bool:
|
||||
"""
|
||||
Remove a permission rule.
|
||||
|
||||
Args:
|
||||
rule_id: The rule ID to remove
|
||||
|
||||
Returns:
|
||||
True if found and removed
|
||||
"""
|
||||
with self._lock:
|
||||
for i, rule in enumerate(self._rules):
|
||||
if rule.id == rule_id:
|
||||
self._rules.pop(i)
|
||||
self._save_rules()
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_rules(self, agent_name: Optional[str] = None) -> List[PermissionRule]:
|
||||
"""
|
||||
Get all rules, optionally filtered by agent.
|
||||
|
||||
Args:
|
||||
agent_name: Optional agent name to filter by
|
||||
|
||||
Returns:
|
||||
List of matching rules
|
||||
"""
|
||||
with self._lock:
|
||||
if agent_name is None:
|
||||
return self._rules.copy()
|
||||
return [
|
||||
r for r in self._rules
|
||||
if r.agent_name is None or r.agent_name == agent_name
|
||||
]
|
||||
|
||||
def check(self, target: str, agent_name: Optional[str] = None) -> PermissionResult:
|
||||
"""
|
||||
Check permission for a target.
|
||||
|
||||
Args:
|
||||
target: The target to check (e.g., "bash:rm -rf /tmp")
|
||||
agent_name: Optional agent name
|
||||
|
||||
Returns:
|
||||
PermissionResult with the decision
|
||||
"""
|
||||
agent = agent_name or self.agent_name
|
||||
|
||||
with self._lock:
|
||||
# Check persistent approvals first
|
||||
for approval in self._approvals:
|
||||
if approval.matches(target, agent):
|
||||
return PermissionResult(
|
||||
action=PermissionAction.ALLOW if approval.approved else PermissionAction.DENY,
|
||||
target=target,
|
||||
reason=f"Persistent approval: {'approved' if approval.approved else 'denied'}",
|
||||
approved=approval.approved,
|
||||
)
|
||||
|
||||
# Check rules
|
||||
for rule in self._rules:
|
||||
if rule.agent_name and agent and rule.agent_name != agent:
|
||||
continue
|
||||
|
||||
if rule.matches(target):
|
||||
return PermissionResult(
|
||||
action=rule.action,
|
||||
rule=rule,
|
||||
target=target,
|
||||
reason=rule.description or f"Matched rule: {rule.pattern}",
|
||||
)
|
||||
|
||||
# Default: ask
|
||||
return PermissionResult(
|
||||
action=PermissionAction.ASK,
|
||||
target=target,
|
||||
reason="No matching rule, requires approval",
|
||||
)
|
||||
|
||||
def approve(
|
||||
self,
|
||||
target: str,
|
||||
approved: bool,
|
||||
scope: str = "once",
|
||||
agent_name: Optional[str] = None,
|
||||
) -> PersistentApproval:
|
||||
"""
|
||||
Record an approval decision.
|
||||
|
||||
Args:
|
||||
target: The target that was approved/denied
|
||||
approved: Whether it was approved
|
||||
scope: Scope of approval (once, session, always)
|
||||
agent_name: Optional agent name
|
||||
|
||||
Returns:
|
||||
The created PersistentApproval
|
||||
"""
|
||||
approval = PersistentApproval(
|
||||
pattern=target,
|
||||
approved=approved,
|
||||
scope=scope,
|
||||
agent_name=agent_name or self.agent_name,
|
||||
)
|
||||
|
||||
with self._lock:
|
||||
self._approvals.append(approval)
|
||||
if scope in ("session", "always"):
|
||||
self._save_approvals()
|
||||
|
||||
return approval
|
||||
|
||||
def check_and_approve(
|
||||
self,
|
||||
target: str,
|
||||
agent_name: Optional[str] = None,
|
||||
) -> PermissionResult:
|
||||
"""
|
||||
Check permission and request approval if needed.
|
||||
|
||||
Uses the approval_callback if set.
|
||||
|
||||
Args:
|
||||
target: The target to check
|
||||
agent_name: Optional agent name
|
||||
|
||||
Returns:
|
||||
PermissionResult with final decision
|
||||
"""
|
||||
result = self.check(target, agent_name)
|
||||
|
||||
if result.needs_approval and self.approval_callback:
|
||||
approved = self.approval_callback(target, result.reason)
|
||||
result.approved = approved
|
||||
self.approve(target, approved, scope="once", agent_name=agent_name)
|
||||
|
||||
return result
|
||||
|
||||
def check_doom_loop(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
"""
|
||||
Check for doom loop and record tool call.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
arguments: Tool arguments
|
||||
|
||||
Returns:
|
||||
DoomLoopResult
|
||||
"""
|
||||
return self._doom_detector.record_and_check(tool_name, arguments)
|
||||
|
||||
def get_doom_stats(self) -> Dict[str, Any]:
|
||||
"""Get doom loop detector statistics."""
|
||||
return self._doom_detector.get_stats()
|
||||
|
||||
def reset_doom_detector(self):
|
||||
"""Reset the doom loop detector."""
|
||||
self._doom_detector.reset()
|
||||
|
||||
def clear_approvals(self, scope: Optional[str] = None):
|
||||
"""
|
||||
Clear approvals.
|
||||
|
||||
Args:
|
||||
scope: Optional scope to clear (None = all)
|
||||
"""
|
||||
with self._lock:
|
||||
if scope is None:
|
||||
self._approvals.clear()
|
||||
else:
|
||||
self._approvals = [a for a in self._approvals if a.scope != scope]
|
||||
self._save_approvals()
|
||||
|
||||
def clear_rules(self):
|
||||
"""Clear all rules."""
|
||||
with self._lock:
|
||||
self._rules.clear()
|
||||
self._save_rules()
|
||||
|
||||
def set_approval_callback(self, callback: Callable[[str, str], bool]):
|
||||
"""Set the approval callback."""
|
||||
self.approval_callback = callback
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Export manager state to dictionary."""
|
||||
with self._lock:
|
||||
return {
|
||||
"rules": [r.to_dict() for r in self._rules],
|
||||
"approvals": [a.to_dict() for a in self._approvals],
|
||||
"agent_name": self.agent_name,
|
||||
}
|
||||
|
||||
def from_dict(self, data: Dict[str, Any]):
|
||||
"""Import manager state from dictionary."""
|
||||
with self._lock:
|
||||
self._rules = [PermissionRule.from_dict(r) for r in data.get("rules", [])]
|
||||
self._approvals = [PersistentApproval.from_dict(a) for a in data.get("approvals", [])]
|
||||
if "agent_name" in data:
|
||||
self.agent_name = data["agent_name"]
|
||||
@@ -0,0 +1,211 @@
|
||||
"""
|
||||
Permission rules for PraisonAI Agents.
|
||||
|
||||
Provides pattern-based permission rules with allow/deny/ask actions.
|
||||
"""
|
||||
|
||||
import fnmatch
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
class PermissionAction(str, Enum):
|
||||
"""Actions for permission rules."""
|
||||
|
||||
ALLOW = "allow"
|
||||
DENY = "deny"
|
||||
ASK = "ask"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PermissionRule:
|
||||
"""
|
||||
A pattern-based permission rule.
|
||||
|
||||
Attributes:
|
||||
pattern: Glob or regex pattern to match against
|
||||
action: What to do when pattern matches (allow/deny/ask)
|
||||
description: Human-readable description
|
||||
is_regex: Whether pattern is a regex (default: glob)
|
||||
priority: Higher priority rules are checked first
|
||||
agent_name: Optional agent name this rule applies to
|
||||
enabled: Whether the rule is active
|
||||
"""
|
||||
|
||||
pattern: str
|
||||
action: PermissionAction = PermissionAction.ASK
|
||||
description: str = ""
|
||||
is_regex: bool = False
|
||||
priority: int = 0
|
||||
agent_name: Optional[str] = None
|
||||
enabled: bool = True
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
created_at: float = field(default_factory=time.time)
|
||||
|
||||
def matches(self, target: str) -> bool:
|
||||
"""
|
||||
Check if this rule matches the target.
|
||||
|
||||
Args:
|
||||
target: The string to match against (e.g., "bash:rm -rf /tmp")
|
||||
|
||||
Returns:
|
||||
True if the pattern matches
|
||||
"""
|
||||
if not self.enabled:
|
||||
return False
|
||||
|
||||
if self.is_regex:
|
||||
try:
|
||||
return bool(re.match(self.pattern, target))
|
||||
except re.error:
|
||||
return False
|
||||
else:
|
||||
return fnmatch.fnmatch(target, self.pattern)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert rule to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"pattern": self.pattern,
|
||||
"action": self.action.value,
|
||||
"description": self.description,
|
||||
"is_regex": self.is_regex,
|
||||
"priority": self.priority,
|
||||
"agent_name": self.agent_name,
|
||||
"enabled": self.enabled,
|
||||
"created_at": self.created_at,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "PermissionRule":
|
||||
"""Create rule from dictionary."""
|
||||
return cls(
|
||||
id=data.get("id", str(uuid.uuid4())),
|
||||
pattern=data.get("pattern", "*"),
|
||||
action=PermissionAction(data.get("action", "ask")),
|
||||
description=data.get("description", ""),
|
||||
is_regex=data.get("is_regex", False),
|
||||
priority=data.get("priority", 0),
|
||||
agent_name=data.get("agent_name"),
|
||||
enabled=data.get("enabled", True),
|
||||
created_at=data.get("created_at", time.time()),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PermissionResult:
|
||||
"""
|
||||
Result of a permission check.
|
||||
|
||||
Attributes:
|
||||
action: The resulting action (allow/deny/ask)
|
||||
rule: The rule that matched (if any)
|
||||
target: The target that was checked
|
||||
reason: Human-readable reason
|
||||
approved: Whether the action was approved (for ask actions)
|
||||
"""
|
||||
|
||||
action: PermissionAction
|
||||
rule: Optional[PermissionRule] = None
|
||||
target: str = ""
|
||||
reason: str = ""
|
||||
approved: Optional[bool] = None
|
||||
|
||||
@property
|
||||
def is_allowed(self) -> bool:
|
||||
"""Check if the action is allowed."""
|
||||
if self.action == PermissionAction.ALLOW:
|
||||
return True
|
||||
if self.action == PermissionAction.ASK:
|
||||
return self.approved is True
|
||||
return False
|
||||
|
||||
@property
|
||||
def is_denied(self) -> bool:
|
||||
"""Check if the action is denied."""
|
||||
if self.action == PermissionAction.DENY:
|
||||
return True
|
||||
if self.action == PermissionAction.ASK:
|
||||
return self.approved is False
|
||||
return False
|
||||
|
||||
@property
|
||||
def needs_approval(self) -> bool:
|
||||
"""Check if the action needs user approval."""
|
||||
return self.action == PermissionAction.ASK and self.approved is None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert result to dictionary."""
|
||||
return {
|
||||
"action": self.action.value,
|
||||
"rule_id": self.rule.id if self.rule else None,
|
||||
"target": self.target,
|
||||
"reason": self.reason,
|
||||
"approved": self.approved,
|
||||
"is_allowed": self.is_allowed,
|
||||
"is_denied": self.is_denied,
|
||||
"needs_approval": self.needs_approval,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class PersistentApproval:
|
||||
"""
|
||||
A persistent approval for a pattern.
|
||||
|
||||
Used to remember user decisions for future requests.
|
||||
"""
|
||||
|
||||
pattern: str
|
||||
approved: bool
|
||||
scope: str = "once" # once, session, always
|
||||
created_at: float = field(default_factory=time.time)
|
||||
expires_at: Optional[float] = None
|
||||
agent_name: Optional[str] = None
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
|
||||
def is_valid(self) -> bool:
|
||||
"""Check if the approval is still valid."""
|
||||
if self.expires_at and time.time() > self.expires_at:
|
||||
return False
|
||||
return True
|
||||
|
||||
def matches(self, target: str, agent_name: Optional[str] = None) -> bool:
|
||||
"""Check if this approval applies to the target."""
|
||||
if not self.is_valid():
|
||||
return False
|
||||
|
||||
if self.agent_name and agent_name and self.agent_name != agent_name:
|
||||
return False
|
||||
|
||||
return fnmatch.fnmatch(target, self.pattern)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert approval to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"pattern": self.pattern,
|
||||
"approved": self.approved,
|
||||
"scope": self.scope,
|
||||
"created_at": self.created_at,
|
||||
"expires_at": self.expires_at,
|
||||
"agent_name": self.agent_name,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "PersistentApproval":
|
||||
"""Create approval from dictionary."""
|
||||
return cls(
|
||||
id=data.get("id", str(uuid.uuid4())),
|
||||
pattern=data.get("pattern", "*"),
|
||||
approved=data.get("approved", False),
|
||||
scope=data.get("scope", "once"),
|
||||
created_at=data.get("created_at", time.time()),
|
||||
expires_at=data.get("expires_at"),
|
||||
agent_name=data.get("agent_name"),
|
||||
)
|
||||
@@ -0,0 +1,47 @@
|
||||
"""
|
||||
Plugin Module for PraisonAI Agents.
|
||||
|
||||
Provides dynamic plugin loading and hook-based extension system.
|
||||
|
||||
Features:
|
||||
- Dynamic plugin discovery and loading
|
||||
- Hook-based extension points
|
||||
- Tool plugins
|
||||
- Auth plugins
|
||||
- Configuration plugins
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.plugins import PluginManager
|
||||
|
||||
# Create plugin manager
|
||||
manager = PluginManager()
|
||||
|
||||
# Load plugins from directory
|
||||
manager.load_from_directory("./plugins")
|
||||
|
||||
# Get loaded plugins
|
||||
plugins = manager.list_plugins()
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
"PluginManager",
|
||||
"Plugin",
|
||||
"PluginHook",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components."""
|
||||
if name == "PluginManager":
|
||||
from .manager import PluginManager
|
||||
return PluginManager
|
||||
|
||||
if name == "Plugin":
|
||||
from .plugin import Plugin
|
||||
return Plugin
|
||||
|
||||
if name == "PluginHook":
|
||||
from .plugin import PluginHook
|
||||
return PluginHook
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,299 @@
|
||||
"""
|
||||
Plugin Manager for PraisonAI Agents.
|
||||
|
||||
Provides dynamic plugin discovery, loading, and lifecycle management.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Type
|
||||
|
||||
from .plugin import Plugin, PluginHook, PluginInfo, FunctionPlugin
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PluginManager:
|
||||
"""
|
||||
Manages plugin discovery, loading, and execution.
|
||||
|
||||
Example:
|
||||
manager = PluginManager()
|
||||
|
||||
# Load plugins from directory
|
||||
manager.load_from_directory("./plugins")
|
||||
|
||||
# Register a plugin directly
|
||||
manager.register(MyPlugin())
|
||||
|
||||
# Execute hooks
|
||||
args = manager.execute_hook(PluginHook.BEFORE_TOOL, "bash", {"cmd": "ls"})
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._plugins: Dict[str, Plugin] = {}
|
||||
self._enabled: Dict[str, bool] = {}
|
||||
|
||||
def register(self, plugin: Plugin) -> bool:
|
||||
"""
|
||||
Register a plugin.
|
||||
|
||||
Args:
|
||||
plugin: The plugin to register
|
||||
|
||||
Returns:
|
||||
True if registered successfully
|
||||
"""
|
||||
try:
|
||||
info = plugin.info
|
||||
if info.name in self._plugins:
|
||||
logger.warning(f"Plugin '{info.name}' already registered")
|
||||
return False
|
||||
|
||||
self._plugins[info.name] = plugin
|
||||
self._enabled[info.name] = True
|
||||
|
||||
# Initialize plugin
|
||||
plugin.on_init({})
|
||||
|
||||
logger.info(f"Registered plugin: {info.name} v{info.version}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to register plugin: {e}")
|
||||
return False
|
||||
|
||||
def unregister(self, name: str) -> bool:
|
||||
"""
|
||||
Unregister a plugin.
|
||||
|
||||
Args:
|
||||
name: The plugin name
|
||||
|
||||
Returns:
|
||||
True if unregistered successfully
|
||||
"""
|
||||
if name not in self._plugins:
|
||||
return False
|
||||
|
||||
try:
|
||||
plugin = self._plugins[name]
|
||||
plugin.on_shutdown()
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during plugin shutdown: {e}")
|
||||
|
||||
del self._plugins[name]
|
||||
del self._enabled[name]
|
||||
|
||||
logger.info(f"Unregistered plugin: {name}")
|
||||
return True
|
||||
|
||||
def enable(self, name: str) -> bool:
|
||||
"""Enable a plugin."""
|
||||
if name in self._enabled:
|
||||
self._enabled[name] = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable(self, name: str) -> bool:
|
||||
"""Disable a plugin."""
|
||||
if name in self._enabled:
|
||||
self._enabled[name] = False
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_enabled(self, name: str) -> bool:
|
||||
"""Check if a plugin is enabled."""
|
||||
return self._enabled.get(name, False)
|
||||
|
||||
def get_plugin(self, name: str) -> Optional[Plugin]:
|
||||
"""Get a plugin by name."""
|
||||
return self._plugins.get(name)
|
||||
|
||||
def list_plugins(self) -> List[PluginInfo]:
|
||||
"""List all registered plugins."""
|
||||
return [p.info for p in self._plugins.values()]
|
||||
|
||||
def load_from_directory(self, directory: str) -> int:
|
||||
"""
|
||||
Load plugins from a directory.
|
||||
|
||||
Looks for Python files with a `create_plugin` function
|
||||
or classes that inherit from Plugin.
|
||||
|
||||
Args:
|
||||
directory: Path to the plugins directory
|
||||
|
||||
Returns:
|
||||
Number of plugins loaded
|
||||
"""
|
||||
directory = Path(directory)
|
||||
if not directory.exists():
|
||||
logger.warning(f"Plugin directory does not exist: {directory}")
|
||||
return 0
|
||||
|
||||
loaded = 0
|
||||
|
||||
for file_path in directory.glob("*.py"):
|
||||
if file_path.name.startswith("_"):
|
||||
continue
|
||||
|
||||
try:
|
||||
plugin = self._load_plugin_file(file_path)
|
||||
if plugin and self.register(plugin):
|
||||
loaded += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load plugin from {file_path}: {e}")
|
||||
|
||||
return loaded
|
||||
|
||||
def _load_plugin_file(self, file_path: Path) -> Optional[Plugin]:
|
||||
"""Load a plugin from a Python file."""
|
||||
module_name = f"praison_plugin_{file_path.stem}"
|
||||
|
||||
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
||||
if spec is None or spec.loader is None:
|
||||
return None
|
||||
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading module {file_path}: {e}")
|
||||
del sys.modules[module_name]
|
||||
return None
|
||||
|
||||
# Look for create_plugin function
|
||||
if hasattr(module, "create_plugin"):
|
||||
return module.create_plugin()
|
||||
|
||||
# Look for Plugin subclass
|
||||
for name in dir(module):
|
||||
obj = getattr(module, name)
|
||||
if (
|
||||
isinstance(obj, type)
|
||||
and issubclass(obj, Plugin)
|
||||
and obj is not Plugin
|
||||
and obj is not FunctionPlugin
|
||||
):
|
||||
return obj()
|
||||
|
||||
return None
|
||||
|
||||
def load_from_module(self, module_path: str) -> bool:
|
||||
"""
|
||||
Load a plugin from a module path.
|
||||
|
||||
Args:
|
||||
module_path: Dotted module path (e.g., "mypackage.plugins.myplugin")
|
||||
|
||||
Returns:
|
||||
True if loaded successfully
|
||||
"""
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
|
||||
if hasattr(module, "create_plugin"):
|
||||
plugin = module.create_plugin()
|
||||
return self.register(plugin)
|
||||
|
||||
# Look for Plugin subclass
|
||||
for name in dir(module):
|
||||
obj = getattr(module, name)
|
||||
if (
|
||||
isinstance(obj, type)
|
||||
and issubclass(obj, Plugin)
|
||||
and obj is not Plugin
|
||||
and obj is not FunctionPlugin
|
||||
):
|
||||
return self.register(obj())
|
||||
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load plugin from module {module_path}: {e}")
|
||||
return False
|
||||
|
||||
def execute_hook(
|
||||
self,
|
||||
hook: PluginHook,
|
||||
*args,
|
||||
**kwargs
|
||||
) -> Any:
|
||||
"""
|
||||
Execute a hook across all enabled plugins.
|
||||
|
||||
Args:
|
||||
hook: The hook to execute
|
||||
*args: Arguments to pass to hook
|
||||
**kwargs: Keyword arguments to pass to hook
|
||||
|
||||
Returns:
|
||||
The result (may be modified by plugins)
|
||||
"""
|
||||
result = args[0] if args else None
|
||||
|
||||
for name, plugin in self._plugins.items():
|
||||
if not self._enabled.get(name, False):
|
||||
continue
|
||||
|
||||
if hook not in plugin.info.hooks:
|
||||
continue
|
||||
|
||||
try:
|
||||
method = getattr(plugin, hook.value, None)
|
||||
if method:
|
||||
if hook == PluginHook.BEFORE_TOOL:
|
||||
result = method(args[0], args[1] if len(args) > 1 else {})
|
||||
elif hook == PluginHook.AFTER_TOOL:
|
||||
result = method(args[0], args[1] if len(args) > 1 else None)
|
||||
elif hook == PluginHook.BEFORE_AGENT:
|
||||
result = method(args[0], kwargs.get("context", {}))
|
||||
elif hook == PluginHook.AFTER_AGENT:
|
||||
result = method(args[0], kwargs.get("context", {}))
|
||||
elif hook == PluginHook.BEFORE_LLM:
|
||||
result = method(args[0], args[1] if len(args) > 1 else {})
|
||||
elif hook == PluginHook.AFTER_LLM:
|
||||
result = method(args[0], kwargs.get("usage", {}))
|
||||
else:
|
||||
result = method(*args, **kwargs)
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing hook {hook.value} in plugin {name}: {e}")
|
||||
|
||||
return result
|
||||
|
||||
def get_all_tools(self) -> List[Dict[str, Any]]:
|
||||
"""Get all tools from all enabled plugins."""
|
||||
tools = []
|
||||
|
||||
for name, plugin in self._plugins.items():
|
||||
if not self._enabled.get(name, False):
|
||||
continue
|
||||
|
||||
try:
|
||||
plugin_tools = plugin.get_tools()
|
||||
tools.extend(plugin_tools)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting tools from plugin {name}: {e}")
|
||||
|
||||
return tools
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown all plugins."""
|
||||
for name in list(self._plugins.keys()):
|
||||
self.unregister(name)
|
||||
|
||||
|
||||
# Global plugin manager instance
|
||||
_default_manager: Optional[PluginManager] = None
|
||||
|
||||
|
||||
def get_plugin_manager() -> PluginManager:
|
||||
"""Get the global plugin manager."""
|
||||
global _default_manager
|
||||
if _default_manager is None:
|
||||
_default_manager = PluginManager()
|
||||
return _default_manager
|
||||
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
Plugin definitions for PraisonAI Agents.
|
||||
|
||||
Provides base plugin class and hook definitions.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PluginHook(str, Enum):
|
||||
"""Available plugin hooks."""
|
||||
|
||||
# Lifecycle hooks
|
||||
ON_INIT = "on_init"
|
||||
ON_SHUTDOWN = "on_shutdown"
|
||||
|
||||
# Agent hooks
|
||||
BEFORE_AGENT = "before_agent"
|
||||
AFTER_AGENT = "after_agent"
|
||||
|
||||
# Tool hooks
|
||||
BEFORE_TOOL = "before_tool"
|
||||
AFTER_TOOL = "after_tool"
|
||||
|
||||
# Message hooks
|
||||
BEFORE_MESSAGE = "before_message"
|
||||
AFTER_MESSAGE = "after_message"
|
||||
|
||||
# LLM hooks
|
||||
BEFORE_LLM = "before_llm"
|
||||
AFTER_LLM = "after_llm"
|
||||
|
||||
# Permission hooks
|
||||
ON_PERMISSION_ASK = "on_permission_ask"
|
||||
|
||||
# Config hooks
|
||||
ON_CONFIG = "on_config"
|
||||
|
||||
# Auth hooks
|
||||
ON_AUTH = "on_auth"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PluginInfo:
|
||||
"""Plugin metadata."""
|
||||
|
||||
name: str
|
||||
version: str = "1.0.0"
|
||||
description: str = ""
|
||||
author: str = ""
|
||||
hooks: List[PluginHook] = field(default_factory=list)
|
||||
dependencies: List[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"version": self.version,
|
||||
"description": self.description,
|
||||
"author": self.author,
|
||||
"hooks": [h.value for h in self.hooks],
|
||||
"dependencies": self.dependencies,
|
||||
}
|
||||
|
||||
|
||||
class Plugin(ABC):
|
||||
"""
|
||||
Base class for PraisonAI plugins.
|
||||
|
||||
Plugins can extend agent functionality by implementing
|
||||
hook methods that are called at various points in the
|
||||
agent lifecycle.
|
||||
|
||||
Example:
|
||||
class MyPlugin(Plugin):
|
||||
@property
|
||||
def info(self) -> PluginInfo:
|
||||
return PluginInfo(
|
||||
name="my_plugin",
|
||||
version="1.0.0",
|
||||
hooks=[PluginHook.BEFORE_TOOL]
|
||||
)
|
||||
|
||||
def before_tool(self, tool_name: str, args: dict) -> dict:
|
||||
# Modify args before tool execution
|
||||
return args
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def info(self) -> PluginInfo:
|
||||
"""Get plugin information."""
|
||||
pass
|
||||
|
||||
def on_init(self, context: Dict[str, Any]) -> None:
|
||||
"""Called when plugin is initialized."""
|
||||
pass
|
||||
|
||||
def on_shutdown(self) -> None:
|
||||
"""Called when plugin is shutting down."""
|
||||
pass
|
||||
|
||||
def before_agent(self, prompt: str, context: Dict[str, Any]) -> str:
|
||||
"""Called before agent execution. Can modify prompt."""
|
||||
return prompt
|
||||
|
||||
def after_agent(self, response: str, context: Dict[str, Any]) -> str:
|
||||
"""Called after agent execution. Can modify response."""
|
||||
return response
|
||||
|
||||
def before_tool(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Called before tool execution. Can modify args."""
|
||||
return args
|
||||
|
||||
def after_tool(self, tool_name: str, result: Any) -> Any:
|
||||
"""Called after tool execution. Can modify result."""
|
||||
return result
|
||||
|
||||
def before_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Called before message is processed. Can modify message."""
|
||||
return message
|
||||
|
||||
def after_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Called after message is processed. Can modify message."""
|
||||
return message
|
||||
|
||||
def before_llm(self, messages: List[Dict], params: Dict[str, Any]) -> tuple:
|
||||
"""Called before LLM call. Can modify messages and params."""
|
||||
return messages, params
|
||||
|
||||
def after_llm(self, response: str, usage: Dict[str, Any]) -> str:
|
||||
"""Called after LLM call. Can modify response."""
|
||||
return response
|
||||
|
||||
def on_permission_ask(self, target: str, reason: str) -> Optional[bool]:
|
||||
"""Called when permission is requested. Return True/False to auto-approve/deny."""
|
||||
return None
|
||||
|
||||
def on_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Called to modify configuration."""
|
||||
return config
|
||||
|
||||
def on_auth(self, auth_type: str, credentials: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
||||
"""Called for authentication. Return credentials or None."""
|
||||
return None
|
||||
|
||||
def get_tools(self) -> List[Dict[str, Any]]:
|
||||
"""Return additional tools provided by this plugin."""
|
||||
return []
|
||||
|
||||
|
||||
class FunctionPlugin(Plugin):
|
||||
"""
|
||||
A simple plugin that wraps functions for specific hooks.
|
||||
|
||||
Example:
|
||||
def my_before_tool(tool_name, args):
|
||||
return args
|
||||
|
||||
plugin = FunctionPlugin(
|
||||
name="my_plugin",
|
||||
hooks={PluginHook.BEFORE_TOOL: my_before_tool}
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
hooks: Optional[Dict[PluginHook, Callable]] = None,
|
||||
version: str = "1.0.0",
|
||||
description: str = "",
|
||||
):
|
||||
self._name = name
|
||||
self._hooks = hooks or {}
|
||||
self._version = version
|
||||
self._description = description
|
||||
|
||||
@property
|
||||
def info(self) -> PluginInfo:
|
||||
return PluginInfo(
|
||||
name=self._name,
|
||||
version=self._version,
|
||||
description=self._description,
|
||||
hooks=list(self._hooks.keys()),
|
||||
)
|
||||
|
||||
def before_tool(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
|
||||
if PluginHook.BEFORE_TOOL in self._hooks:
|
||||
return self._hooks[PluginHook.BEFORE_TOOL](tool_name, args)
|
||||
return args
|
||||
|
||||
def after_tool(self, tool_name: str, result: Any) -> Any:
|
||||
if PluginHook.AFTER_TOOL in self._hooks:
|
||||
return self._hooks[PluginHook.AFTER_TOOL](tool_name, result)
|
||||
return result
|
||||
|
||||
def before_agent(self, prompt: str, context: Dict[str, Any]) -> str:
|
||||
if PluginHook.BEFORE_AGENT in self._hooks:
|
||||
return self._hooks[PluginHook.BEFORE_AGENT](prompt, context)
|
||||
return prompt
|
||||
|
||||
def after_agent(self, response: str, context: Dict[str, Any]) -> str:
|
||||
if PluginHook.AFTER_AGENT in self._hooks:
|
||||
return self._hooks[PluginHook.AFTER_AGENT](response, context)
|
||||
return response
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
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"),
|
||||
"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, 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,203 @@
|
||||
"""
|
||||
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 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,534 @@
|
||||
"""
|
||||
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, 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 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}")
|
||||
Loaded 100 of 445 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user