mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 17:33:33 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ff3a4f2016 | ||
|
|
0d5444fc23 | ||
|
|
d7d5dab17b | ||
|
|
05678394e6 | ||
|
|
54cddc09b4 | ||
|
|
9cbab01785 | ||
|
|
de9d4bd3c6 | ||
|
|
aba2265e5e | ||
|
|
ebb9b24bff | ||
|
|
759ef2c655 | ||
|
|
cb4318f6e8 | ||
|
|
d9bb869d5e | ||
|
|
34b6aeb76b | ||
|
|
078b7801e1 | ||
|
|
a970e849b4 | ||
|
|
8e68f4e71c | ||
|
|
308758cdc2 | ||
|
|
f8fa61b082 | ||
|
|
955d236924 | ||
|
|
d9d63c3a13 | ||
|
|
5bee6eb470 | ||
|
|
7cb1676b7c | ||
|
|
e2e9c3f9ee | ||
|
|
cbe8beeca6 |
No files matched your search
+3
-1
@@ -80,4 +80,6 @@ CopilotKit*
|
||||
issue_details.json
|
||||
|
||||
# MCP Registry tokens
|
||||
.mcpregistry_*
|
||||
.mcpregistry_*
|
||||
|
||||
*.log
|
||||
@@ -204,6 +204,12 @@ npm install praisonai
|
||||
| ↳ Add Custom Knowledge | [Example](examples/python/concepts/knowledge-agents.py) | [📖](https://docs.praison.ai/features/knowledge) |
|
||||
| ↳ RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [📖](https://docs.praison.ai/features/rag) |
|
||||
| ↳ Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [📖](https://docs.praison.ai/features/chat-with-pdf) |
|
||||
| ↳ Data Readers (PDF, DOCX, etc.) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-readers-api) |
|
||||
| ↳ Vector Store Selection | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-vector-store-api) |
|
||||
| ↳ Retrieval Strategies | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-retrieval-api) |
|
||||
| ↳ Rerankers | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-reranker-api) |
|
||||
| ↳ Index Types (Vector/Keyword/Hybrid) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-index-api) |
|
||||
| ↳ Query Engines (Sub-Question, etc.) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-query-engine-api) |
|
||||
| **🔬 Research & Intelligence** | | |
|
||||
| ↳ Deep Research Agents | [Example](examples/python/agents/research-agent.py) | [📖](https://docs.praison.ai/agents/deep-research) |
|
||||
| ↳ Query Rewriter Agent | [Example](#5-query-rewriter-agent) | [📖](https://docs.praison.ai/agents/query-rewriter) |
|
||||
@@ -261,6 +267,16 @@ npm install praisonai
|
||||
| ↳ 100+ LLM Support | [Example](examples/python/providers/openai/openai_gpt4_example.py) | [📖](https://docs.praison.ai/models) |
|
||||
| ↳ Callback Agents | [Example](examples/python/general/advanced-callback-systems.py) | [📖](https://docs.praison.ai/features/callbacks) |
|
||||
| ↳ Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
|
||||
| ↳ Middleware System | [Example](examples/middleware/basic_middleware.py) | [📖](https://docs.praison.ai/features/middleware) |
|
||||
| ↳ Configurable Model | [Example](examples/middleware/configurable_model.py) | [📖](https://docs.praison.ai/features/configurable-model) |
|
||||
| ↳ Rate Limiter | [Example](examples/middleware/rate_limiter.py) | [📖](https://docs.praison.ai/features/rate-limiter) |
|
||||
| ↳ Injected Tool State | [Example](examples/middleware/injected_state.py) | [📖](https://docs.praison.ai/features/injected-state) |
|
||||
| ↳ Shadow Git Checkpoints | [Example](#10-shadow-git-checkpoints) | [📖](https://docs.praison.ai/features/checkpoints) |
|
||||
| ↳ Background Tasks | [Example](examples/background/basic_background.py) | [📖](https://docs.praison.ai/features/background-tasks) |
|
||||
| ↳ Policy Engine | [Example](examples/policy/basic_policy.py) | [📖](https://docs.praison.ai/features/policy-engine) |
|
||||
| ↳ Thinking Budgets | [Example](examples/thinking/basic_thinking.py) | [📖](https://docs.praison.ai/features/thinking-budgets) |
|
||||
| ↳ Output Styles | [Example](examples/output/basic_output.py) | [📖](https://docs.praison.ai/features/output-styles) |
|
||||
| ↳ Context Compaction | [Example](examples/compaction/basic_compaction.py) | [📖](https://docs.praison.ai/features/context-compaction) |
|
||||
| **📊 Monitoring & Management** | | |
|
||||
| ↳ Sessions Management | [Example](examples/python/sessions/comprehensive-session-management.py) | [📖](https://docs.praison.ai/features/sessions) |
|
||||
| ↳ Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
@@ -278,6 +294,12 @@ npm install praisonai
|
||||
| ↳ Git Integration | [Example](examples/python/cli/git_integration_example.py) | [📖](https://docs.praison.ai/docs/cli/git-integration) |
|
||||
| ↳ Sandbox Execution | [Example](examples/python/cli/sandbox_execution_example.py) | [📖](https://docs.praison.ai/docs/cli/sandbox-execution) |
|
||||
| ↳ CLI Compare | [Example](examples/compare/cli_compare_basic.py) | [📖](https://docs.praison.ai/docs/cli/compare) |
|
||||
| ↳ Profile/Benchmark | [Example](#profile-benchmark) | [📖](https://docs.praison.ai/docs/cli/profile) |
|
||||
| ↳ Auto Mode | [Example](#auto-mode) | [📖](https://docs.praison.ai/docs/cli/auto) |
|
||||
| ↳ Init | [Example](#init) | [📖](https://docs.praison.ai/docs/cli/init) |
|
||||
| ↳ File Input | [Example](#file-input) | [📖](https://docs.praison.ai/docs/cli/file-input) |
|
||||
| ↳ Final Agent | [Example](#final-agent) | [📖](https://docs.praison.ai/docs/cli/final-agent) |
|
||||
| ↳ Max Tokens | [Example](#max-tokens) | [📖](https://docs.praison.ai/docs/cli/max-tokens) |
|
||||
| **🧪 Evaluation** | | |
|
||||
| ↳ Accuracy Evaluation | [Example](examples/eval/accuracy_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
| ↳ Performance Evaluation | [Example](examples/eval/performance_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
@@ -857,21 +879,258 @@ callbacks:
|
||||
|
||||
### 9. Hooks
|
||||
|
||||
Configure in `.praison/hooks.json`:
|
||||
Intercept and modify agent behavior at various lifecycle points:
|
||||
|
||||
```python
|
||||
from praisonaiagents.memory import HooksManager
|
||||
from praisonaiagents.hooks import (
|
||||
HookRegistry, HookRunner, HookEvent, HookResult,
|
||||
BeforeToolInput
|
||||
)
|
||||
|
||||
hooks = HooksManager()
|
||||
# Create a hook registry
|
||||
registry = HookRegistry()
|
||||
|
||||
# Register Python hooks
|
||||
hooks.register("pre_write_code", lambda ctx: print(f"Writing {ctx['file']}"))
|
||||
# Log all tool calls
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def log_tools(event_data: BeforeToolInput) -> HookResult:
|
||||
print(f"Tool: {event_data.tool_name}")
|
||||
return HookResult.allow()
|
||||
|
||||
# Block dangerous operations
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def security_check(event_data: BeforeToolInput) -> HookResult:
|
||||
if "delete" in event_data.tool_name.lower():
|
||||
return HookResult.deny("Delete operations blocked")
|
||||
return HookResult.allow()
|
||||
|
||||
# Execute hooks
|
||||
result = hooks.execute("pre_write_code", {"file": "main.py"})
|
||||
runner = HookRunner(registry)
|
||||
```
|
||||
|
||||
### 10. Field Names Reference (A-I-G-S)
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai hooks list # List registered hooks
|
||||
praisonai hooks test before_tool # Test hooks for an event
|
||||
praisonai hooks run "echo test" # Run a command hook
|
||||
praisonai hooks validate hooks.json # Validate configuration
|
||||
```
|
||||
|
||||
|
||||
### 10. Shadow Git Checkpoints
|
||||
|
||||
File-level undo/restore using shadow git:
|
||||
|
||||
```python
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
|
||||
service = CheckpointService(workspace_dir="./my_project")
|
||||
await service.initialize()
|
||||
|
||||
# Save checkpoint before changes
|
||||
result = await service.save("Before refactoring")
|
||||
|
||||
# Make changes...
|
||||
|
||||
# Restore if needed
|
||||
await service.restore(result.checkpoint.id)
|
||||
|
||||
# View diff
|
||||
diff = await service.diff()
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai checkpoint save "Before changes" # Save checkpoint
|
||||
praisonai checkpoint list # List checkpoints
|
||||
praisonai checkpoint diff # Show changes
|
||||
praisonai checkpoint restore abc123 # Restore to checkpoint
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/checkpoints)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/checkpoint)
|
||||
- [💻 Example](examples/checkpoints/basic_checkpoints.py)
|
||||
|
||||
---
|
||||
|
||||
### 11. Background Tasks
|
||||
|
||||
Run agent tasks asynchronously without blocking:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from praisonaiagents.background import BackgroundRunner, BackgroundConfig
|
||||
|
||||
async def main():
|
||||
config = BackgroundConfig(max_concurrent_tasks=3)
|
||||
runner = BackgroundRunner(config=config)
|
||||
|
||||
async def my_task(name: str) -> str:
|
||||
await asyncio.sleep(2)
|
||||
return f"Task {name} completed"
|
||||
|
||||
task = await runner.submit(my_task, args=("example",), name="my_task")
|
||||
await task.wait(timeout=10.0)
|
||||
print(task.result)
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai background list # List running tasks
|
||||
praisonai background status <id> # Check task status
|
||||
praisonai background cancel <id> # Cancel a task
|
||||
praisonai background clear # Clear completed tasks
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/background-tasks)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/background)
|
||||
- [💻 Example](examples/background/basic_background.py)
|
||||
|
||||
---
|
||||
|
||||
### 12. Policy Engine
|
||||
|
||||
Control what agents can and cannot do with policy-based execution:
|
||||
|
||||
```python
|
||||
from praisonaiagents.policy import (
|
||||
PolicyEngine, Policy, PolicyRule, PolicyAction
|
||||
)
|
||||
|
||||
engine = PolicyEngine()
|
||||
|
||||
policy = Policy(
|
||||
name="no_delete",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="Delete operations blocked"
|
||||
)
|
||||
]
|
||||
)
|
||||
engine.add_policy(policy)
|
||||
|
||||
result = engine.check("tool:delete_file", {})
|
||||
print(f"Allowed: {result.allowed}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai policy list # List policies
|
||||
praisonai policy check "tool:name" # Check if allowed
|
||||
praisonai policy init # Create template
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/policy-engine)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/policy)
|
||||
- [💻 Example](examples/policy/basic_policy.py)
|
||||
|
||||
---
|
||||
|
||||
### 13. Thinking Budgets
|
||||
|
||||
Configure token budgets for extended thinking:
|
||||
|
||||
```python
|
||||
from praisonaiagents.thinking import ThinkingBudget, ThinkingTracker
|
||||
|
||||
# Use predefined levels
|
||||
budget = ThinkingBudget.high() # 16,000 tokens
|
||||
|
||||
# Track usage
|
||||
tracker = ThinkingTracker()
|
||||
session = tracker.start_session(budget_tokens=16000)
|
||||
tracker.end_session(session, tokens_used=12000)
|
||||
|
||||
summary = tracker.get_summary()
|
||||
print(f"Utilization: {summary['average_utilization']:.1%}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai thinking status # Show current budget
|
||||
praisonai thinking set high # Set budget level
|
||||
praisonai thinking stats # Show usage statistics
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/thinking-budgets)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/thinking)
|
||||
- [💻 Example](examples/thinking/basic_thinking.py)
|
||||
|
||||
---
|
||||
|
||||
### 14. Output Styles
|
||||
|
||||
Configure how agents format their responses:
|
||||
|
||||
```python
|
||||
from praisonaiagents.output import OutputStyle, OutputFormatter
|
||||
|
||||
# Use preset styles
|
||||
style = OutputStyle.concise()
|
||||
formatter = OutputFormatter(style)
|
||||
|
||||
# Format output
|
||||
text = "# Hello\n\nThis is **bold** text."
|
||||
plain = formatter.format(text)
|
||||
print(plain)
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai output status # Show current style
|
||||
praisonai output set concise # Set output style
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/output-styles)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/output-style)
|
||||
- [💻 Example](examples/output/basic_output.py)
|
||||
|
||||
---
|
||||
|
||||
### 15. Context Compaction
|
||||
|
||||
Automatically manage context window size:
|
||||
|
||||
```python
|
||||
from praisonaiagents.compaction import (
|
||||
ContextCompactor, CompactionStrategy
|
||||
)
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=4000,
|
||||
strategy=CompactionStrategy.SLIDING,
|
||||
preserve_recent=3
|
||||
)
|
||||
|
||||
messages = [...] # Your conversation history
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f"Compression: {result.compression_ratio:.1%}")
|
||||
```
|
||||
|
||||
**CLI Commands:**
|
||||
```bash
|
||||
praisonai compaction status # Show settings
|
||||
praisonai compaction set sliding # Set strategy
|
||||
praisonai compaction stats # Show statistics
|
||||
```
|
||||
|
||||
**Links:**
|
||||
- [📖 Coding Docs](https://docs.praison.ai/features/context-compaction)
|
||||
- [📖 CLI Docs](https://docs.praison.ai/docs/cli/compaction)
|
||||
- [💻 Example](examples/compaction/basic_compaction.py)
|
||||
|
||||
---
|
||||
|
||||
### 16. Field Names Reference (A-I-G-S)
|
||||
|
||||
PraisonAI accepts both old (agents.yaml) and new (workflow.yaml) field names. Use the **canonical names** for new projects:
|
||||
|
||||
@@ -905,7 +1164,7 @@ steps: # Define steps (not 'tasks')
|
||||
|
||||
> **Note:** The parser accepts both old and new names. Run `praisonai workflow validate <file.yaml>` to see suggestions for canonical names.
|
||||
|
||||
### 11. Extended agents.yaml with Workflow Patterns
|
||||
### 12. Extended agents.yaml with Workflow Patterns
|
||||
|
||||
**Feature Parity:** Both `agents.yaml` and `workflow.yaml` now support the same features:
|
||||
- All workflow patterns (route, parallel, loop, repeat)
|
||||
@@ -981,6 +1240,124 @@ Run with the same simple command:
|
||||
praisonai agents.yaml
|
||||
```
|
||||
|
||||
### 12. MCP (Model Context Protocol)
|
||||
|
||||
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
|
||||
|
||||
#### MCP Client (Consume MCP Servers)
|
||||
```python
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
# stdio - Local NPX/Python servers
|
||||
agent = Agent(tools=MCP("npx @modelcontextprotocol/server-memory"))
|
||||
|
||||
# Streamable HTTP - Production servers
|
||||
agent = Agent(tools=MCP("https://api.example.com/mcp"))
|
||||
|
||||
# WebSocket - Real-time bidirectional
|
||||
agent = Agent(tools=MCP("wss://api.example.com/mcp", auth_token="token"))
|
||||
|
||||
# SSE (Legacy) - Backward compatibility
|
||||
agent = Agent(tools=MCP("http://localhost:8080/sse"))
|
||||
|
||||
# With environment variables
|
||||
agent = Agent(
|
||||
tools=MCP(
|
||||
command="npx",
|
||||
args=["-y", "@modelcontextprotocol/server-brave-search"],
|
||||
env={"BRAVE_API_KEY": "your-key"}
|
||||
)
|
||||
)
|
||||
|
||||
# Multiple MCP servers + regular functions
|
||||
def my_custom_tool(query: str) -> str:
|
||||
"""Custom tool function."""
|
||||
return f"Result: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="MultiToolAgent",
|
||||
instructions="Agent with multiple MCP servers",
|
||||
tools=[
|
||||
MCP("uvx mcp-server-time"), # Time tools
|
||||
MCP("npx @modelcontextprotocol/server-memory"), # Memory tools
|
||||
my_custom_tool # Regular function
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
#### MCP Server (Expose Tools as MCP Server)
|
||||
|
||||
Expose your Python functions as MCP tools for Claude Desktop, Cursor, and other MCP clients:
|
||||
|
||||
```python
|
||||
from praisonaiagents.mcp import ToolsMCPServer
|
||||
|
||||
def search_web(query: str, max_results: int = 5) -> dict:
|
||||
"""Search the web for information."""
|
||||
return {"results": [f"Result for {query}"]}
|
||||
|
||||
def calculate(expression: str) -> dict:
|
||||
"""Evaluate a mathematical expression."""
|
||||
return {"result": eval(expression)}
|
||||
|
||||
# Create and run MCP server
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
server.register_tools([search_web, calculate])
|
||||
server.run() # stdio for Claude Desktop
|
||||
# server.run_sse(host="0.0.0.0", port=8080) # SSE for web clients
|
||||
```
|
||||
|
||||
#### MCP Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Session Management | Automatic Mcp-Session-Id handling |
|
||||
| Protocol Versioning | Mcp-Protocol-Version header |
|
||||
| Resumability | SSE stream recovery via Last-Event-ID |
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
### 13. A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
#### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
#### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
---
|
||||
|
||||
## 🎯 CLI / No-Code Interface
|
||||
@@ -1730,124 +2107,6 @@ agent = Agent(
|
||||
| Web Fetch | Anthropic |
|
||||
| Prompt Caching | OpenAI (auto), Anthropic, Bedrock, Deepseek |
|
||||
|
||||
## MCP (Model Context Protocol)
|
||||
|
||||
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
|
||||
|
||||
### MCP Client (Consume MCP Servers)
|
||||
```python
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
# stdio - Local NPX/Python servers
|
||||
agent = Agent(tools=MCP("npx @modelcontextprotocol/server-memory"))
|
||||
|
||||
# Streamable HTTP - Production servers
|
||||
agent = Agent(tools=MCP("https://api.example.com/mcp"))
|
||||
|
||||
# WebSocket - Real-time bidirectional
|
||||
agent = Agent(tools=MCP("wss://api.example.com/mcp", auth_token="token"))
|
||||
|
||||
# SSE (Legacy) - Backward compatibility
|
||||
agent = Agent(tools=MCP("http://localhost:8080/sse"))
|
||||
|
||||
# With environment variables
|
||||
agent = Agent(
|
||||
tools=MCP(
|
||||
command="npx",
|
||||
args=["-y", "@modelcontextprotocol/server-brave-search"],
|
||||
env={"BRAVE_API_KEY": "your-key"}
|
||||
)
|
||||
)
|
||||
|
||||
# Multiple MCP servers + regular functions
|
||||
def my_custom_tool(query: str) -> str:
|
||||
"""Custom tool function."""
|
||||
return f"Result: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="MultiToolAgent",
|
||||
instructions="Agent with multiple MCP servers",
|
||||
tools=[
|
||||
MCP("uvx mcp-server-time"), # Time tools
|
||||
MCP("npx @modelcontextprotocol/server-memory"), # Memory tools
|
||||
my_custom_tool # Regular function
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
### MCP Server (Expose Tools as MCP Server)
|
||||
|
||||
Expose your Python functions as MCP tools for Claude Desktop, Cursor, and other MCP clients:
|
||||
|
||||
```python
|
||||
from praisonaiagents.mcp import ToolsMCPServer
|
||||
|
||||
def search_web(query: str, max_results: int = 5) -> dict:
|
||||
"""Search the web for information."""
|
||||
return {"results": [f"Result for {query}"]}
|
||||
|
||||
def calculate(expression: str) -> dict:
|
||||
"""Evaluate a mathematical expression."""
|
||||
return {"result": eval(expression)}
|
||||
|
||||
# Create and run MCP server
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
server.register_tools([search_web, calculate])
|
||||
server.run() # stdio for Claude Desktop
|
||||
# server.run_sse(host="0.0.0.0", port=8080) # SSE for web clients
|
||||
```
|
||||
|
||||
### MCP Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Session Management | Automatic Mcp-Session-Id handling |
|
||||
| Protocol Versioning | Mcp-Protocol-Version header |
|
||||
| Resumability | SSE stream recovery via Last-Event-ID |
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
## A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
## CLI Features
|
||||
|
||||
| Feature | Docs |
|
||||
@@ -2388,6 +2647,84 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
|
||||
---
|
||||
|
||||
## 📚 Knowledge & Retrieval (RAG)
|
||||
|
||||
PraisonAI provides a complete knowledge stack for building RAG applications with multiple vector stores, retrieval strategies, rerankers, and query modes.
|
||||
|
||||
### Knowledge CLI Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai knowledge add <file\|dir\|url>` | Add documents to knowledge base |
|
||||
| `praisonai knowledge query <question>` | Query knowledge base with RAG |
|
||||
| `praisonai knowledge list` | List indexed documents |
|
||||
| `praisonai knowledge clear` | Clear knowledge base |
|
||||
| `praisonai knowledge stats` | Show knowledge base statistics |
|
||||
|
||||
### Knowledge CLI Options
|
||||
|
||||
| Option | Values | Description |
|
||||
|--------|--------|-------------|
|
||||
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | Vector store backend |
|
||||
| `--retrieval` | `basic`, `fusion`, `recursive`, `auto_merge` | Retrieval strategy |
|
||||
| `--reranker` | `simple`, `llm`, `cross_encoder`, `cohere` | Reranking method |
|
||||
| `--index-type` | `vector`, `keyword`, `hybrid` | Index type |
|
||||
| `--query-mode` | `default`, `sub_question`, `summarize` | Query mode |
|
||||
|
||||
### Knowledge CLI Examples
|
||||
|
||||
```bash
|
||||
# Add documents
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge add https://example.com/page.html
|
||||
praisonai knowledge add "*.pdf"
|
||||
|
||||
# Query with advanced options
|
||||
praisonai knowledge query "How to authenticate?" --retrieval fusion --reranker llm
|
||||
|
||||
# Full advanced query
|
||||
praisonai knowledge query "authentication flow" \
|
||||
--vector-store chroma \
|
||||
--retrieval fusion \
|
||||
--reranker llm \
|
||||
--index-type hybrid \
|
||||
--query-mode sub_question
|
||||
```
|
||||
|
||||
### Knowledge SDK Usage
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, Knowledge
|
||||
|
||||
# Simple usage with Agent
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
knowledge=["docs/manual.pdf", "data/faq.txt"],
|
||||
knowledge_config={
|
||||
"vector_store": {"provider": "chroma"}
|
||||
}
|
||||
)
|
||||
response = agent.chat("How do I authenticate?")
|
||||
|
||||
# Direct Knowledge usage
|
||||
knowledge = Knowledge()
|
||||
knowledge.add("document.pdf")
|
||||
results = knowledge.search("authentication", limit=5)
|
||||
```
|
||||
|
||||
### Knowledge Stack Features Table
|
||||
|
||||
| Feature | Description | SDK Docs | CLI Docs |
|
||||
|---------|-------------|----------|----------|
|
||||
| **Data Readers** | Load PDF, Markdown, Text, HTML, URLs | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Vector Stores** | ChromaDB, Pinecone, Qdrant, Weaviate, In-Memory | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Retrieval Strategies** | Basic, Fusion (RRF), Recursive, Auto-Merge | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Rerankers** | Simple, LLM, Cross-Encoder, Cohere | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Index Types** | Vector, Keyword (BM25), Hybrid | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Query Engines** | Default, Sub-Question, Summarize | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
|
||||
---
|
||||
|
||||
## 🔬 Advanced Features
|
||||
|
||||
### Research & Intelligence
|
||||
@@ -2475,6 +2812,141 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
|
||||
---
|
||||
|
||||
## 💾 Persistence (Databases)
|
||||
|
||||
Enable automatic conversation persistence with 2 lines of code:
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
db=db(database_url="postgresql://localhost/mydb"), # db(...) shortcut
|
||||
session_id="my-session" # Optional: defaults to per-hour ID (YYYYMMDDHH)
|
||||
)
|
||||
agent.chat("Hello!") # Auto-persists messages, runs, traces
|
||||
```
|
||||
|
||||
### Persistence CLI Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai persistence doctor` | Validate DB connectivity |
|
||||
| `praisonai persistence run` | Run agent with persistence |
|
||||
| `praisonai persistence resume` | Resume existing session |
|
||||
| `praisonai persistence export` | Export session to JSONL |
|
||||
| `praisonai persistence import` | Import session from JSONL |
|
||||
| `praisonai persistence migrate` | Apply schema migrations |
|
||||
| `praisonai persistence status` | Show schema status |
|
||||
|
||||
### 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
|
||||
```
|
||||
|
||||
### Databases Table
|
||||
|
||||
| Database | Store Type | Install | Example | Docs |
|
||||
|----------|------------|---------|---------|------|
|
||||
| PostgreSQL | Conversation | `pip install "praisonai[tools]"` | [simple_db_agent.py](examples/persistence/simple_db_agent.py) | [docs](https://docs.praison.ai/docs/databases/postgres) |
|
||||
| MySQL | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SQLite | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SingleStore | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Supabase | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| SurrealDB | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Qdrant | Knowledge | `pip install "praisonai[tools]"` | [knowledge_qdrant.py](examples/persistence/knowledge_qdrant.py) | [docs](https://docs.praison.ai/docs/databases/qdrant) |
|
||||
| ChromaDB | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Pinecone | Knowledge | `pip install pinecone` | [pinecone_wow.py](examples/vector/pinecone_wow.py) | [docs](https://docs.praison.ai/docs/databases/pinecone) |
|
||||
| Weaviate | Knowledge | `pip install weaviate-client` | [weaviate_wow.py](examples/vector/weaviate_wow.py) | [docs](https://docs.praison.ai/docs/databases/weaviate) |
|
||||
| LanceDB | Knowledge | `pip install lancedb` | [lancedb_real_wow.py](examples/vector/lancedb_real_wow.py) | [docs](https://docs.praison.ai/docs/databases/lancedb) |
|
||||
| Milvus | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| PGVector | Knowledge | `pip install psycopg2-binary` | [pgvector_real_wow.py](examples/vector/pgvector_real_wow.py) | [docs](https://docs.praison.ai/docs/databases/pgvector) |
|
||||
| Redis Vector | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Cassandra | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| ClickHouse | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Redis | State | `pip install "praisonai[tools]"` | [state_redis.py](examples/persistence/state_redis.py) | [docs](https://docs.praison.ai/docs/databases/redis) |
|
||||
| MongoDB | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| DynamoDB | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Firestore | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Upstash | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Memory | State | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Tools Table
|
||||
|
||||
Install all tools with: `pip install "praisonai[tools]"`
|
||||
|
||||
| Tool | Category | Import | Docs |
|
||||
|------|----------|--------|------|
|
||||
| Tavily | Web Search | `from praisonai_tools import TavilyTool` | [docs](https://docs.praison.ai/docs/tools/external/tavily) |
|
||||
| DuckDuckGo | Web Search | `from praisonai_tools import DuckDuckGoTool` | [docs](https://docs.praison.ai/docs/tools/external/duckduckgo) |
|
||||
| Exa | Web Search | `from praisonai_tools import ExaTool` | [docs](https://docs.praison.ai/docs/tools/external/exa) |
|
||||
| Serper | Web Search | `from praisonai_tools import SerperTool` | [docs](https://docs.praison.ai/docs/tools/external/serper) |
|
||||
| Jina | Web Reader | `from praisonai_tools import JinaTool` | [docs](https://docs.praison.ai/docs/tools/external/jina) |
|
||||
| Firecrawl | Web Scraping | `from praisonai_tools import FirecrawlTool` | [docs](https://docs.praison.ai/docs/tools/external/firecrawl) |
|
||||
| Crawl4AI | Web Scraping | `from praisonai_tools import Crawl4AITool` | [docs](https://docs.praison.ai/docs/tools/external/crawl4ai) |
|
||||
| Wikipedia | Knowledge | `from praisonai_tools import WikipediaTool` | [docs](https://docs.praison.ai/docs/tools/external/wikipedia) |
|
||||
| ArXiv | Research | `from praisonai_tools import ArxivTool` | [docs](https://docs.praison.ai/docs/tools/external/arxiv) |
|
||||
| HackerNews | News | `from praisonai_tools import HackerNewsTool` | [docs](https://docs.praison.ai/docs/tools/external/hackernews) |
|
||||
| YouTube | Media | `from praisonai_tools import YouTubeTool` | [docs](https://docs.praison.ai/docs/tools/external/youtube) |
|
||||
| Weather | Data | `from praisonai_tools import WeatherTool` | [docs](https://docs.praison.ai/docs/tools/external/weather) |
|
||||
| PostgreSQL | Database | `from praisonai_tools import PostgresTool` | [docs](https://docs.praison.ai/docs/tools/external/postgres) |
|
||||
| MySQL | Database | `from praisonai_tools import MySQLTool` | [docs](https://docs.praison.ai/docs/tools/external/mysql) |
|
||||
| SQLite | Database | `from praisonai_tools import SQLiteTool` | [docs](https://docs.praison.ai/docs/tools/external/sqlite) |
|
||||
| MongoDB | Database | `from praisonai_tools import MongoDBTool` | [docs](https://docs.praison.ai/docs/tools/external/mongodb) |
|
||||
| Redis | Database | `from praisonai_tools import RedisTool` | [docs](https://docs.praison.ai/docs/tools/external/redis) |
|
||||
| Qdrant | Vector DB | `from praisonai_tools import QdrantTool` | [docs](https://docs.praison.ai/docs/tools/external/qdrant) |
|
||||
| GitHub | DevOps | `from praisonai_tools import GitHubTool` | [docs](https://docs.praison.ai/docs/tools/external/github) |
|
||||
| Slack | Communication | `from praisonai_tools import SlackTool` | [docs](https://docs.praison.ai/docs/tools/external/slack) |
|
||||
| Discord | Communication | `from praisonai_tools import DiscordTool` | [docs](https://docs.praison.ai/docs/tools/external/discord) |
|
||||
| Telegram | Communication | `from praisonai_tools import TelegramTool` | [docs](https://docs.praison.ai/docs/tools/external/telegram) |
|
||||
| Email | Communication | `from praisonai_tools import EmailTool` | [docs](https://docs.praison.ai/docs/tools/external/email) |
|
||||
| Notion | Productivity | `from praisonai_tools import NotionTool` | [docs](https://docs.praison.ai/docs/tools/external/notion) |
|
||||
| File | File System | `from praisonai_tools import FileTool` | [docs](https://docs.praison.ai/docs/tools/external/file) |
|
||||
| Shell | System | `from praisonai_tools import ShellTool` | [docs](https://docs.praison.ai/docs/tools/external/shell) |
|
||||
| Python | Code | `from praisonai_tools import PythonTool` | [docs](https://docs.praison.ai/docs/tools/external/python) |
|
||||
| JSON | Data | `from praisonai_tools import JSONTool` | [docs](https://docs.praison.ai/docs/tools/external/json) |
|
||||
| CSV | Data | `from praisonai_tools import CSVTool` | [docs](https://docs.praison.ai/docs/tools/external/csv) |
|
||||
| Calculator | Math | `from praisonai_tools import CalculatorTool` | [docs](https://docs.praison.ai/docs/tools/external/calculator) |
|
||||
|
||||
> See [full tools documentation](https://docs.praison.ai/docs/tools/tools) for all 100+ available tools.
|
||||
|
||||
---
|
||||
|
||||
## 🎓 Video Tutorials
|
||||
|
||||
Learn PraisonAI through our comprehensive video series:
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.3.79" \
|
||||
"praisonai>=2.5.2" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.3.79" \
|
||||
"praisonai>=2.5.2" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.3.79" \
|
||||
"praisonai>=2.5.2" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
Async Runs Examples
|
||||
===================
|
||||
|
||||
This directory contains examples for the PraisonAI Async Jobs API.
|
||||
|
||||
Examples:
|
||||
---------
|
||||
1. sdk_submit_poll_result.py - Python SDK example for submit/poll/result
|
||||
2. cli_submit_poll_result.sh - Bash/curl example for submit/poll/result
|
||||
3. idempotency.sh - Demonstrates idempotency key usage
|
||||
4. cancel.sh - Demonstrates job cancellation
|
||||
5. sse_stream.sh - Demonstrates SSE streaming
|
||||
|
||||
Prerequisites:
|
||||
--------------
|
||||
1. Set OPENAI_API_KEY environment variable
|
||||
2. Start the jobs server:
|
||||
|
||||
cd /path/to/praisonai-package/src/praisonai
|
||||
python -m uvicorn praisonai.jobs.server:create_app --port 8005 --factory
|
||||
|
||||
3. Run examples:
|
||||
|
||||
python sdk_submit_poll_result.py
|
||||
bash cli_submit_poll_result.sh
|
||||
bash idempotency.sh
|
||||
bash cancel.sh
|
||||
bash sse_stream.sh
|
||||
|
||||
Environment Variables:
|
||||
----------------------
|
||||
- OPENAI_API_KEY: Required for agent execution
|
||||
- PRAISONAI_BASE_URL: API server URL (default: http://127.0.0.1:8005)
|
||||
"""
|
||||
@@ -0,0 +1,54 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: Cancel a Running Job
|
||||
#
|
||||
# This example demonstrates how to cancel a running job.
|
||||
#
|
||||
# Requirements:
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash cancel.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI Job Cancellation Example"
|
||||
echo "============================================================"
|
||||
|
||||
# Submit a long-running job
|
||||
echo ""
|
||||
echo "1. Submitting a job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompt": "Write a detailed essay about the history of computing."}')
|
||||
|
||||
JOB_ID=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID"
|
||||
|
||||
# Wait briefly
|
||||
echo ""
|
||||
echo "2. Waiting 1 second..."
|
||||
sleep 1
|
||||
|
||||
# Cancel the job
|
||||
echo ""
|
||||
echo "3. Cancelling job..."
|
||||
CANCEL_RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs/$JOB_ID/cancel")
|
||||
STATUS=$(echo "$CANCEL_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['status'])")
|
||||
echo " Status after cancel: $STATUS"
|
||||
|
||||
# Verify cancellation
|
||||
echo ""
|
||||
echo "4. Verification:"
|
||||
if [ "$STATUS" = "cancelled" ]; then
|
||||
echo " ✓ SUCCESS: Job was cancelled"
|
||||
else
|
||||
echo " ℹ Job status: $STATUS (may have completed before cancel)"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Cancellation example completed!"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,75 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: Submit, Poll, and Get Result for Async Jobs
|
||||
#
|
||||
# This example demonstrates how to use the PraisonAI CLI to:
|
||||
# 1. Submit an async job
|
||||
# 2. Check job status
|
||||
# 3. Get the result
|
||||
#
|
||||
# Requirements:
|
||||
# - OPENAI_API_KEY environment variable set
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash cli_submit_poll_result.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI Async Jobs CLI Example"
|
||||
echo "============================================================"
|
||||
|
||||
# Check if server is running
|
||||
echo ""
|
||||
echo "Checking server health..."
|
||||
if ! curl -s "$API_URL/health" > /dev/null 2>&1; then
|
||||
echo "✗ Server not available at $API_URL"
|
||||
echo " Start the server with: python -m uvicorn praisonai.jobs.server:create_app --port 8005 --factory"
|
||||
exit 1
|
||||
fi
|
||||
echo "✓ Server is healthy"
|
||||
|
||||
# 1. Submit a job
|
||||
echo ""
|
||||
echo "1. Submitting job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompt": "What is 2+2? Answer with just the number."}')
|
||||
|
||||
JOB_ID=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID"
|
||||
|
||||
# 2. Poll for status
|
||||
echo ""
|
||||
echo "2. Polling for status..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/api/v1/runs/$JOB_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['status'])")
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['progress']['percentage'])")
|
||||
|
||||
echo " Status: $STATUS | Progress: ${PROGRESS}%"
|
||||
|
||||
if [ "$STATUS" = "succeeded" ] || [ "$STATUS" = "failed" ] || [ "$STATUS" = "cancelled" ]; then
|
||||
break
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
|
||||
# 3. Get result
|
||||
echo ""
|
||||
echo "3. Getting result..."
|
||||
if [ "$STATUS" = "succeeded" ]; then
|
||||
RESULT=$(curl -s "$API_URL/api/v1/runs/$JOB_ID/result")
|
||||
echo " Result: $(echo "$RESULT" | python3 -c "import sys,json; print(json.load(sys.stdin)['result'])")"
|
||||
echo " Duration: $(echo "$RESULT" | python3 -c "import sys,json; print(f\"{json.load(sys.stdin)['duration_seconds']:.2f}s\")")"
|
||||
else
|
||||
echo " Job did not succeed: $STATUS"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Example completed!"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,57 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: Idempotency Key Usage
|
||||
#
|
||||
# This example demonstrates how idempotency keys prevent duplicate job submissions.
|
||||
# Submitting the same request with the same idempotency key returns the existing job.
|
||||
#
|
||||
# Requirements:
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash idempotency.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
IDEM_KEY="example-idem-key-$(date +%s)"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI Idempotency Key Example"
|
||||
echo "============================================================"
|
||||
|
||||
# First submission
|
||||
echo ""
|
||||
echo "1. First submission with Idempotency-Key: $IDEM_KEY"
|
||||
RESPONSE1=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Idempotency-Key: $IDEM_KEY" \
|
||||
-d '{"prompt": "What is 1+1?"}')
|
||||
|
||||
JOB_ID1=$(echo "$RESPONSE1" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID1"
|
||||
|
||||
# Second submission with same key
|
||||
echo ""
|
||||
echo "2. Second submission with SAME Idempotency-Key: $IDEM_KEY"
|
||||
RESPONSE2=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Idempotency-Key: $IDEM_KEY" \
|
||||
-d '{"prompt": "Different prompt but same key"}')
|
||||
|
||||
JOB_ID2=$(echo "$RESPONSE2" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID2"
|
||||
|
||||
# Verify same job ID
|
||||
echo ""
|
||||
echo "3. Verification:"
|
||||
if [ "$JOB_ID1" = "$JOB_ID2" ]; then
|
||||
echo " ✓ SUCCESS: Same job ID returned (no duplicate created)"
|
||||
else
|
||||
echo " ✗ FAILURE: Different job IDs returned"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Idempotency example completed!"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,121 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SDK Example: Submit, Poll, and Get Result for Async Jobs
|
||||
|
||||
This example demonstrates how to use the PraisonAI Jobs SDK to:
|
||||
1. Submit an async job
|
||||
2. Poll for status
|
||||
3. Get the result
|
||||
|
||||
Requirements:
|
||||
- OPENAI_API_KEY environment variable set
|
||||
- Jobs server running on http://127.0.0.1:8005
|
||||
|
||||
Usage:
|
||||
python sdk_submit_poll_result.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import httpx
|
||||
|
||||
# Configuration
|
||||
API_URL = os.getenv("PRAISONAI_BASE_URL", "http://127.0.0.1:8005")
|
||||
POLL_INTERVAL = 2 # seconds
|
||||
|
||||
|
||||
def submit_job(prompt: str, idempotency_key: str = None) -> dict:
|
||||
"""Submit a job and return the response."""
|
||||
headers = {}
|
||||
if idempotency_key:
|
||||
headers["Idempotency-Key"] = idempotency_key
|
||||
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.post(
|
||||
f"{API_URL}/api/v1/runs",
|
||||
json={"prompt": prompt},
|
||||
headers=headers if headers else None
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_status(job_id: str) -> dict:
|
||||
"""Get job status."""
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.get(f"{API_URL}/api/v1/runs/{job_id}")
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_result(job_id: str) -> dict:
|
||||
"""Get job result."""
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.get(f"{API_URL}/api/v1/runs/{job_id}/result")
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def wait_for_completion(job_id: str, max_wait: int = 120) -> dict:
|
||||
"""Poll until job completes."""
|
||||
start = time.time()
|
||||
while time.time() - start < max_wait:
|
||||
status = get_status(job_id)
|
||||
print(f" Status: {status['status']} | Progress: {status['progress']['percentage']:.0f}%")
|
||||
|
||||
if status["status"] in ("succeeded", "failed", "cancelled"):
|
||||
return status
|
||||
|
||||
# Honor retry_after if present
|
||||
wait_time = status.get("retry_after") or POLL_INTERVAL
|
||||
time.sleep(wait_time)
|
||||
|
||||
raise TimeoutError(f"Job {job_id} did not complete within {max_wait}s")
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("PraisonAI Async Jobs SDK Example")
|
||||
print("=" * 60)
|
||||
|
||||
# Check if server is running
|
||||
try:
|
||||
with httpx.Client(timeout=5.0) as client:
|
||||
health = client.get(f"{API_URL}/health")
|
||||
health.raise_for_status()
|
||||
print(f"✓ Server healthy: {health.json()['status']}")
|
||||
except Exception as e:
|
||||
print(f"✗ Server not available at {API_URL}: {e}")
|
||||
print(" Start the server with: python -m uvicorn praisonai.jobs.server:create_app --port 8005 --factory")
|
||||
return
|
||||
|
||||
# 1. Submit a job
|
||||
print("\n1. Submitting job...")
|
||||
prompt = "What is the capital of France? Answer in one word."
|
||||
result = submit_job(prompt)
|
||||
job_id = result["job_id"]
|
||||
print(f" Job ID: {job_id}")
|
||||
print(f" Status: {result['status']}")
|
||||
print(f" Poll URL: {result['poll_url']}")
|
||||
|
||||
# 2. Wait for completion
|
||||
print("\n2. Waiting for completion...")
|
||||
final_status = wait_for_completion(job_id)
|
||||
print(f" Final status: {final_status['status']}")
|
||||
|
||||
# 3. Get result
|
||||
if final_status["status"] == "succeeded":
|
||||
print("\n3. Getting result...")
|
||||
result = get_result(job_id)
|
||||
print(f" Result: {result['result']}")
|
||||
print(f" Duration: {result['duration_seconds']:.2f}s")
|
||||
else:
|
||||
print(f"\n3. Job did not succeed: {final_status.get('error', 'Unknown error')}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Example completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,51 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: SSE Streaming
|
||||
#
|
||||
# This example demonstrates how to stream job progress via Server-Sent Events.
|
||||
#
|
||||
# Requirements:
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash sse_stream.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI SSE Streaming Example"
|
||||
echo "============================================================"
|
||||
|
||||
# Submit a job
|
||||
echo ""
|
||||
echo "1. Submitting job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompt": "Count from 1 to 5."}')
|
||||
|
||||
JOB_ID=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID"
|
||||
|
||||
# Stream progress
|
||||
echo ""
|
||||
echo "2. Streaming progress (Ctrl+C to stop)..."
|
||||
echo " Connecting to: $API_URL/api/v1/runs/$JOB_ID/stream"
|
||||
echo ""
|
||||
|
||||
# Use curl to stream SSE events (timeout after 60s)
|
||||
curl -s -N --max-time 60 "$API_URL/api/v1/runs/$JOB_ID/stream" | while read -r line; do
|
||||
if [[ "$line" == data:* ]]; then
|
||||
data="${line#data: }"
|
||||
if [ "$data" = "[DONE]" ]; then
|
||||
echo " Stream completed"
|
||||
break
|
||||
fi
|
||||
echo " Event: $data"
|
||||
fi
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Streaming example completed!"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,92 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Background Tasks Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the background task system to:
|
||||
1. Run tasks asynchronously
|
||||
2. Track task progress
|
||||
3. Cancel running tasks
|
||||
4. Handle task completion
|
||||
|
||||
Usage:
|
||||
python basic_background.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
from praisonaiagents.background import BackgroundRunner, BackgroundConfig, TaskStatus
|
||||
|
||||
|
||||
async def main():
|
||||
print("=" * 60)
|
||||
print("Background Tasks Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a background runner with config
|
||||
config = BackgroundConfig(max_concurrent_tasks=3)
|
||||
runner = BackgroundRunner(config=config)
|
||||
|
||||
# Define some example tasks
|
||||
def slow_task(duration: float, name: str) -> str:
|
||||
"""A slow task that simulates work."""
|
||||
print(f" 🔄 Task '{name}' starting (will take {duration}s)")
|
||||
time.sleep(duration)
|
||||
return f"Task '{name}' completed after {duration}s"
|
||||
|
||||
async def async_task(duration: float, name: str) -> str:
|
||||
"""An async task that simulates work."""
|
||||
print(f" 🔄 Async task '{name}' starting (will take {duration}s)")
|
||||
await asyncio.sleep(duration)
|
||||
return f"Async task '{name}' completed after {duration}s"
|
||||
|
||||
# Submit tasks
|
||||
print("\n--- Submitting Tasks ---")
|
||||
|
||||
task1 = await runner.submit(slow_task, args=(1.0, "quick-job"), name="task1")
|
||||
print(f"✅ Submitted task 1: {task1.id[:8]}")
|
||||
|
||||
task2 = await runner.submit(slow_task, args=(2.0, "medium-job"), name="task2")
|
||||
print(f"✅ Submitted task 2: {task2.id[:8]}")
|
||||
|
||||
task3 = await runner.submit(async_task, args=(1.5, "async-job"), name="task3")
|
||||
print(f"✅ Submitted async task 3: {task3.id[:8]}")
|
||||
|
||||
# List all tasks
|
||||
print("\n--- Task Status ---")
|
||||
for task in runner.tasks:
|
||||
print(f" [{task.id[:8]}] {task.status.value}")
|
||||
|
||||
# Wait for task 1 to complete
|
||||
print("\n--- Waiting for Task 1 ---")
|
||||
await task1.wait(timeout=5.0)
|
||||
print(f"✅ Task 1 result: {task1.result}")
|
||||
|
||||
# Check status of all tasks
|
||||
print("\n--- Updated Task Status ---")
|
||||
for task in [task1, task2, task3]:
|
||||
status = "✅" if task.status == TaskStatus.COMPLETED else "🔄"
|
||||
print(f" {status} [{task.id[:8]}] {task.status.value}")
|
||||
|
||||
# Wait for all remaining tasks
|
||||
print("\n--- Waiting for All Tasks ---")
|
||||
await task2.wait(timeout=10.0)
|
||||
await task3.wait(timeout=10.0)
|
||||
|
||||
# Final status
|
||||
print("\n--- Final Results ---")
|
||||
for task in [task1, task2, task3]:
|
||||
if task.result:
|
||||
print(f" ✅ {task.result}")
|
||||
|
||||
# Cleanup
|
||||
runner.clear_completed()
|
||||
print("\n✅ Cleared completed tasks")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,110 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Checkpoints Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the shadow git checkpointing system to:
|
||||
1. Save checkpoints before making changes
|
||||
2. List available checkpoints
|
||||
3. View diffs between checkpoints
|
||||
4. Restore to a previous checkpoint
|
||||
|
||||
Usage:
|
||||
python basic_checkpoints.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
|
||||
|
||||
async def main():
|
||||
# Create a temporary workspace for demonstration
|
||||
with tempfile.TemporaryDirectory() as workspace:
|
||||
print("=" * 60)
|
||||
print("Shadow Git Checkpointing Demo")
|
||||
print("=" * 60)
|
||||
print(f"\nWorkspace: {workspace}")
|
||||
|
||||
# Create some initial files
|
||||
test_file = Path(workspace) / "example.py"
|
||||
test_file.write_text("# Initial content\nprint('Hello, World!')\n")
|
||||
print(f"\n✅ Created initial file: {test_file.name}")
|
||||
|
||||
# Initialize the checkpoint service
|
||||
service = CheckpointService(workspace_dir=workspace)
|
||||
|
||||
try:
|
||||
# Initialize the shadow git repository
|
||||
print("\n--- Initializing Checkpoint Service ---")
|
||||
await service.initialize()
|
||||
print("✅ Shadow git repository initialized")
|
||||
|
||||
# Save first checkpoint
|
||||
print("\n--- Saving Checkpoint 1: Initial State ---")
|
||||
result1 = await service.save("Initial state")
|
||||
print(f"✅ Checkpoint saved: {result1.checkpoint.id[:8]}")
|
||||
print(f" Message: {result1.checkpoint.message}")
|
||||
|
||||
# Make some changes
|
||||
print("\n--- Making Changes ---")
|
||||
test_file.write_text("# Modified content\nprint('Hello, PraisonAI!')\nprint('Checkpoints are awesome!')\n")
|
||||
print("✅ Modified example.py")
|
||||
|
||||
# Create a new file
|
||||
new_file = Path(workspace) / "utils.py"
|
||||
new_file.write_text("def helper():\n return 'I help!'\n")
|
||||
print("✅ Created utils.py")
|
||||
|
||||
# Save second checkpoint
|
||||
print("\n--- Saving Checkpoint 2: After Changes ---")
|
||||
result2 = await service.save("Added features")
|
||||
print(f"✅ Checkpoint saved: {result2.checkpoint.id[:8]}")
|
||||
|
||||
# List all checkpoints
|
||||
print("\n--- Listing All Checkpoints ---")
|
||||
checkpoints = await service.list_checkpoints()
|
||||
for i, cp in enumerate(checkpoints, 1):
|
||||
print(f" {i}. [{cp.id[:8]}] {cp.message} ({cp.timestamp})")
|
||||
|
||||
# Show diff
|
||||
print("\n--- Current Diff (uncommitted changes) ---")
|
||||
diff = await service.diff()
|
||||
if diff.files:
|
||||
for f in diff.files:
|
||||
print(f" {f.status}: {f.path}")
|
||||
else:
|
||||
print(" No uncommitted changes")
|
||||
|
||||
# Make more changes to show diff
|
||||
test_file.write_text("# Even more changes\nprint('This will be reverted')\n")
|
||||
print("\n✅ Made more changes to example.py")
|
||||
|
||||
diff = await service.diff()
|
||||
print("\n--- New Diff ---")
|
||||
if diff.files:
|
||||
for f in diff.files:
|
||||
print(f" {f.status}: {f.path}")
|
||||
|
||||
# Restore to checkpoint 2
|
||||
print(f"\n--- Restoring to Checkpoint: {result2.checkpoint.id[:8]} ---")
|
||||
await service.restore(result2.checkpoint.id)
|
||||
print("✅ Restored successfully!")
|
||||
|
||||
# Verify restoration
|
||||
content = test_file.read_text()
|
||||
print(f"\n--- Verified Content After Restore ---")
|
||||
print(content)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n❌ Error: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,93 @@
|
||||
# Persistence CLI Commands
|
||||
|
||||
## Quick Reference
|
||||
|
||||
```bash
|
||||
# Check database connectivity
|
||||
praisonai persistence doctor \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--state-url redis://localhost:6379
|
||||
|
||||
# Run agent with persistence
|
||||
praisonai persistence run \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--session-id my-session
|
||||
|
||||
# Resume a session
|
||||
praisonai persistence resume \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--session-id my-session
|
||||
|
||||
# Export session to file
|
||||
praisonai persistence export \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--session-id my-session \
|
||||
--output session_backup.jsonl
|
||||
|
||||
# Import session from file
|
||||
praisonai persistence import \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--file session_backup.jsonl
|
||||
|
||||
# Check schema status
|
||||
praisonai persistence status \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--state-url redis://localhost:6379
|
||||
|
||||
# Apply migrations
|
||||
praisonai persistence migrate \
|
||||
--conversation-url postgresql://localhost:5432/praisonai \
|
||||
--state-url redis://localhost:6379
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
Instead of CLI flags, you can use environment variables:
|
||||
|
||||
```bash
|
||||
export PRAISON_CONVERSATION_URL=postgresql://localhost:5432/praisonai
|
||||
export PRAISON_STATE_URL=redis://localhost:6379
|
||||
export PRAISON_KNOWLEDGE_URL=http://localhost:6333
|
||||
|
||||
# Then run without URL flags
|
||||
praisonai persistence doctor
|
||||
praisonai persistence status
|
||||
```
|
||||
|
||||
## Docker Setup
|
||||
|
||||
```bash
|
||||
# Start all services
|
||||
docker run -d --name praison-postgres \
|
||||
-e POSTGRES_PASSWORD=praison123 \
|
||||
-e POSTGRES_DB=praisonai \
|
||||
-p 5432:5432 postgres:16
|
||||
|
||||
docker run -d --name praison-redis \
|
||||
-p 6379:6379 redis:7
|
||||
|
||||
docker run -d --name praison-qdrant \
|
||||
-p 6333:6333 -p 6334:6334 qdrant/qdrant
|
||||
|
||||
# Verify
|
||||
docker ps
|
||||
```
|
||||
|
||||
## Example Workflow
|
||||
|
||||
```bash
|
||||
# 1. Start services
|
||||
docker-compose up -d
|
||||
|
||||
# 2. Check connectivity
|
||||
praisonai persistence doctor
|
||||
|
||||
# 3. Run agent
|
||||
praisonai persistence run --session-id demo-session
|
||||
|
||||
# 4. Export for backup
|
||||
praisonai persistence export --session-id demo-session --output backup.jsonl
|
||||
|
||||
# 5. Import to another environment
|
||||
praisonai persistence import --file backup.jsonl
|
||||
```
|
||||
@@ -0,0 +1,124 @@
|
||||
# PraisonAI Persistence CLI - End-to-End Guide
|
||||
|
||||
This guide demonstrates the CLI commands for database persistence.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
# Start Docker containers
|
||||
docker run -d --name praison-postgres -p 5432:5432 \
|
||||
-e POSTGRES_PASSWORD=praison123 \
|
||||
-e POSTGRES_DB=praisonai \
|
||||
postgres:16
|
||||
|
||||
docker run -d --name praison-qdrant -p 6333:6333 qdrant/qdrant
|
||||
docker run -d --name praison-redis -p 6379:6379 redis:7
|
||||
```
|
||||
|
||||
## 1. Doctor Command - Validate Connectivity
|
||||
|
||||
```bash
|
||||
# Check all stores
|
||||
praisonai persistence doctor \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--knowledge-url "http://localhost:6333" \
|
||||
--state-url "redis://localhost:6379"
|
||||
```
|
||||
|
||||
**Expected Output:**
|
||||
```
|
||||
==================================================
|
||||
PraisonAI Persistence Doctor
|
||||
==================================================
|
||||
|
||||
[Conversation Store] Testing: postgresql://postgres:****@localhost:5432/praisonai
|
||||
✅ Connected (postgres)
|
||||
|
||||
[Knowledge Store] Testing: http://localhost:6333
|
||||
✅ Connected (qdrant)
|
||||
|
||||
[State Store] Testing: redis://localhost:6379
|
||||
✅ Connected (redis)
|
||||
|
||||
==================================================
|
||||
Results: 3/3 stores connected successfully
|
||||
==================================================
|
||||
```
|
||||
|
||||
## 2. Run Command - Execute Agent with Persistence
|
||||
|
||||
```bash
|
||||
# Run agent with persistence (dry-run first)
|
||||
praisonai persistence run \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--dry-run \
|
||||
"Hello, my name is Bob"
|
||||
|
||||
# Actually run it
|
||||
praisonai persistence run \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
"Hello, my name is Bob"
|
||||
```
|
||||
|
||||
**Expected Output:**
|
||||
```
|
||||
[Session: cli-demo-001]
|
||||
Agent: Hello Bob! Nice to meet you. How can I help you today?
|
||||
```
|
||||
|
||||
## 3. Resume Command - Continue a Session
|
||||
|
||||
```bash
|
||||
# Show history
|
||||
praisonai persistence resume \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--show-history
|
||||
|
||||
# Continue conversation
|
||||
praisonai persistence resume \
|
||||
--session-id "cli-demo-001" \
|
||||
--conversation-url "postgresql://postgres:praison123@localhost:5432/praisonai" \
|
||||
--continue "What's my name?"
|
||||
```
|
||||
|
||||
**Expected Output:**
|
||||
```
|
||||
[Session: cli-demo-001]
|
||||
Messages in history: 2
|
||||
|
||||
--- Conversation History ---
|
||||
[USER] Hello, my name is Bob
|
||||
[ASSISTANT] Hello Bob! Nice to meet you...
|
||||
--- End History ---
|
||||
|
||||
Agent: Your name is Bob!
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
Instead of passing URLs on command line, use environment variables:
|
||||
|
||||
```bash
|
||||
export PRAISON_CONVERSATION_URL="postgresql://postgres:praison123@localhost:5432/praisonai"
|
||||
export PRAISON_KNOWLEDGE_URL="http://localhost:6333"
|
||||
export PRAISON_STATE_URL="redis://localhost:6379"
|
||||
export OPENAI_API_KEY="your-key"
|
||||
|
||||
# Now commands are simpler
|
||||
praisonai persistence doctor --all
|
||||
praisonai persistence run --session-id "my-session" "Hello!"
|
||||
```
|
||||
|
||||
## Full CLI Reference
|
||||
|
||||
```
|
||||
praisonai persistence --help
|
||||
praisonai persistence doctor --help
|
||||
praisonai persistence run --help
|
||||
praisonai persistence resume --help
|
||||
```
|
||||
@@ -0,0 +1,146 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Context Compaction Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use context compaction to:
|
||||
1. Manage conversation context length
|
||||
2. Use different compaction strategies
|
||||
3. Track compaction results
|
||||
4. Preserve important messages
|
||||
|
||||
Usage:
|
||||
python basic_compaction.py
|
||||
"""
|
||||
|
||||
from praisonaiagents.compaction import (
|
||||
ContextCompactor, CompactionConfig, CompactionStrategy, CompactionResult
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("Context Compaction Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# ==========================================================================
|
||||
# Sample Messages
|
||||
# ==========================================================================
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful AI assistant."},
|
||||
{"role": "user", "content": "Hello! Can you help me with Python?"},
|
||||
{"role": "assistant", "content": "Of course! I'd be happy to help you with Python. What would you like to know?"},
|
||||
{"role": "user", "content": "How do I read a file in Python?"},
|
||||
{"role": "assistant", "content": "To read a file in Python, you can use the built-in open() function. Here's an example:\n\nwith open('filename.txt', 'r') as f:\n content = f.read()\n\nThis opens the file, reads its content, and automatically closes it."},
|
||||
{"role": "user", "content": "What about writing to a file?"},
|
||||
{"role": "assistant", "content": "To write to a file, use the 'w' mode:\n\nwith open('filename.txt', 'w') as f:\n f.write('Hello, World!')\n\nUse 'a' mode to append instead of overwriting."},
|
||||
{"role": "user", "content": "Can you explain list comprehensions?"},
|
||||
{"role": "assistant", "content": "List comprehensions are a concise way to create lists. Instead of:\n\nresult = []\nfor x in range(10):\n result.append(x * 2)\n\nYou can write:\n\nresult = [x * 2 for x in range(10)]"},
|
||||
]
|
||||
|
||||
print(f"\n--- Original Messages ---")
|
||||
print(f" Message count: {len(messages)}")
|
||||
|
||||
# ==========================================================================
|
||||
# Basic Compaction
|
||||
# ==========================================================================
|
||||
print("\n--- Basic Compaction (Truncate Strategy) ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=50, # Low limit to trigger compaction
|
||||
strategy=CompactionStrategy.TRUNCATE,
|
||||
preserve_system=True,
|
||||
preserve_recent=2
|
||||
)
|
||||
|
||||
# Check stats before
|
||||
stats = compactor.get_stats(messages)
|
||||
print(f" Total tokens: {stats['total_tokens']}")
|
||||
print(f" Max tokens: {stats['max_tokens']}")
|
||||
print(f" Needs compaction: {stats['needs_compaction']}")
|
||||
|
||||
# Compact
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f"\n After compaction:")
|
||||
print(f" Messages: {len(messages)} -> {len(compacted)}")
|
||||
print(f" Tokens: {result.original_tokens} -> {result.compacted_tokens}")
|
||||
print(f" Saved: {result.tokens_saved} tokens")
|
||||
print(f" Compression: {result.compression_ratio:.1%}")
|
||||
|
||||
# ==========================================================================
|
||||
# Sliding Window Strategy
|
||||
# ==========================================================================
|
||||
print("\n--- Sliding Window Strategy ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=50,
|
||||
strategy=CompactionStrategy.SLIDING,
|
||||
preserve_system=True,
|
||||
preserve_recent=3
|
||||
)
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f" Messages kept: {result.messages_kept}")
|
||||
print(f" Messages removed: {result.messages_removed}")
|
||||
print(f" Strategy: {result.strategy_used.value}")
|
||||
|
||||
# ==========================================================================
|
||||
# Summarize Strategy
|
||||
# ==========================================================================
|
||||
print("\n--- Summarize Strategy ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=50,
|
||||
strategy=CompactionStrategy.SUMMARIZE,
|
||||
preserve_system=True,
|
||||
preserve_recent=2
|
||||
)
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f" Messages kept: {result.messages_kept}")
|
||||
print(f" Strategy: {result.strategy_used.value}")
|
||||
|
||||
# Show what messages remain
|
||||
print(f"\n Remaining message roles:")
|
||||
for msg in compacted:
|
||||
role = msg.get("role", "unknown")
|
||||
content_preview = msg.get("content", "")[:50]
|
||||
print(f" [{role}] {content_preview}...")
|
||||
|
||||
# ==========================================================================
|
||||
# No Compaction Needed
|
||||
# ==========================================================================
|
||||
print("\n--- No Compaction Needed ---")
|
||||
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=10000, # High limit
|
||||
strategy=CompactionStrategy.TRUNCATE
|
||||
)
|
||||
|
||||
compacted, result = compactor.compact(messages)
|
||||
|
||||
print(f" Was compacted: {result.was_compacted}")
|
||||
print(f" Messages unchanged: {len(compacted) == len(messages)}")
|
||||
|
||||
# ==========================================================================
|
||||
# Result Serialization
|
||||
# ==========================================================================
|
||||
print("\n--- Result Serialization ---")
|
||||
|
||||
compactor = ContextCompactor(max_tokens=50)
|
||||
_, result = compactor.compact(messages)
|
||||
|
||||
data = result.to_dict()
|
||||
print(f" Original tokens: {data['original_tokens']}")
|
||||
print(f" Compacted tokens: {data['compacted_tokens']}")
|
||||
print(f" Compression ratio: {data['compression_ratio']:.2f}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Hooks Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the hooks system to:
|
||||
1. Log tool calls before execution
|
||||
2. Block dangerous operations
|
||||
3. Add context to agent responses
|
||||
4. Monitor session lifecycle
|
||||
|
||||
Usage:
|
||||
python basic_hooks.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.hooks import (
|
||||
HookRegistry, HookRunner, HookEvent, HookResult,
|
||||
BeforeToolInput, AfterToolInput
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
# Create a hook registry
|
||||
registry = HookRegistry()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 1: Log all tool calls
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def log_tool_calls(event_data: BeforeToolInput) -> HookResult:
|
||||
"""Log every tool call before execution."""
|
||||
print(f"\n📝 [LOG] Tool: {event_data.tool_name}")
|
||||
print(f" Input: {event_data.tool_input}")
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 2: Block dangerous file operations (using function logic)
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def block_delete_operations(event_data: BeforeToolInput) -> HookResult:
|
||||
"""Block any delete or remove operations."""
|
||||
dangerous_keywords = ['delete', 'remove', 'rm', 'unlink']
|
||||
tool_lower = event_data.tool_name.lower()
|
||||
|
||||
if any(kw in tool_lower for kw in dangerous_keywords):
|
||||
print(f"\n🚫 [BLOCKED] Dangerous operation: {event_data.tool_name}")
|
||||
return HookResult.deny(
|
||||
f"Operation '{event_data.tool_name}' is blocked by security policy"
|
||||
)
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 3: Require confirmation for write operations
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def confirm_write_operations(event_data: BeforeToolInput) -> HookResult:
|
||||
"""Request confirmation for write operations."""
|
||||
write_keywords = ['write', 'save', 'create', 'update']
|
||||
tool_lower = event_data.tool_name.lower()
|
||||
|
||||
if any(kw in tool_lower for kw in write_keywords):
|
||||
print(f"\n⚠️ [CONFIRM] Write operation: {event_data.tool_name}")
|
||||
# In a real scenario, you might prompt the user here
|
||||
# For this example, we'll auto-approve
|
||||
return HookResult.allow("Auto-approved for demo")
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Hook 4: Add timing information after tool execution
|
||||
# ==========================================================================
|
||||
@registry.on(HookEvent.AFTER_TOOL)
|
||||
def log_tool_completion(event_data: AfterToolInput) -> HookResult:
|
||||
"""Log tool completion with timing."""
|
||||
print(f"\n✅ [DONE] Tool: {event_data.tool_name}")
|
||||
print(f" Duration: {event_data.execution_time_ms:.2f}ms")
|
||||
if event_data.tool_error:
|
||||
print(f" Error: {event_data.tool_error}")
|
||||
return HookResult.allow()
|
||||
|
||||
# ==========================================================================
|
||||
# Test the hooks
|
||||
# ==========================================================================
|
||||
print("=" * 60)
|
||||
print("Testing Hooks System")
|
||||
print("=" * 60)
|
||||
|
||||
runner = HookRunner(registry)
|
||||
|
||||
# Test 1: Normal tool call (should be logged and allowed)
|
||||
print("\n--- Test 1: Normal read operation ---")
|
||||
input1 = BeforeToolInput(
|
||||
session_id="test",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="read_file",
|
||||
tool_input={"path": "/tmp/test.txt"}
|
||||
)
|
||||
results = asyncio.run(runner.execute(HookEvent.BEFORE_TOOL, input1))
|
||||
print(f"Result: {'ALLOWED' if not HookRunner.is_blocked(results) else 'BLOCKED'}")
|
||||
|
||||
# Test 2: Delete operation (should be blocked)
|
||||
print("\n--- Test 2: Delete operation ---")
|
||||
input2 = BeforeToolInput(
|
||||
session_id="test",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="delete_file",
|
||||
tool_input={"path": "/important/data.txt"}
|
||||
)
|
||||
results = asyncio.run(runner.execute(HookEvent.BEFORE_TOOL, input2))
|
||||
print(f"Result: {'ALLOWED' if not HookRunner.is_blocked(results) else 'BLOCKED'}")
|
||||
if HookRunner.is_blocked(results):
|
||||
print(f"Reason: {HookRunner.get_blocking_reason(results)}")
|
||||
|
||||
# Test 3: Write operation (should be confirmed)
|
||||
print("\n--- Test 3: Write operation ---")
|
||||
input3 = BeforeToolInput(
|
||||
session_id="test",
|
||||
cwd="/tmp",
|
||||
event_name="before_tool",
|
||||
timestamp="2024-01-01T00:00:00",
|
||||
tool_name="write_file",
|
||||
tool_input={"path": "/tmp/output.txt", "content": "Hello"}
|
||||
)
|
||||
results = asyncio.run(runner.execute(HookEvent.BEFORE_TOOL, input3))
|
||||
print(f"Result: {'ALLOWED' if not HookRunner.is_blocked(results) else 'BLOCKED'}")
|
||||
|
||||
# Show registered hooks
|
||||
print("\n" + "=" * 60)
|
||||
print("Registered Hooks:")
|
||||
print("=" * 60)
|
||||
for event, hooks in registry.list_hooks().items():
|
||||
print(f"\n{event}:")
|
||||
for hook in hooks:
|
||||
print(f" - {hook['name']} (matcher: {hook['matcher'] or '*'})")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,55 @@
|
||||
"""
|
||||
Basic Middleware Example - PraisonAI Agents
|
||||
|
||||
Demonstrates before/after hooks and wrap decorators for model and tool calls.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, tool
|
||||
from praisonaiagents.hooks import (
|
||||
before_model, after_model, wrap_model_call,
|
||||
before_tool, after_tool, wrap_tool_call,
|
||||
InvocationContext, ModelRequest, ModelResponse
|
||||
)
|
||||
|
||||
# Simple tool for demonstration
|
||||
@tool
|
||||
def get_weather(city: str) -> str:
|
||||
"""Get weather for a city."""
|
||||
return f"Sunny, 22°C in {city}"
|
||||
|
||||
# Before model hook - adds context
|
||||
@before_model
|
||||
def add_context(request: ModelRequest) -> ModelRequest:
|
||||
print(f"[before_model] Adding context to request")
|
||||
return request
|
||||
|
||||
# After model hook - logs response
|
||||
@after_model
|
||||
def log_response(response: ModelResponse) -> ModelResponse:
|
||||
print(f"[after_model] Response received")
|
||||
return response
|
||||
|
||||
# Wrap tool call - retry on error
|
||||
@wrap_tool_call
|
||||
def retry_on_error(tool_call, call_next):
|
||||
print(f"[wrap_tool_call] Executing tool")
|
||||
try:
|
||||
return call_next(tool_call)
|
||||
except Exception as e:
|
||||
print(f"[wrap_tool_call] Retrying after error: {e}")
|
||||
return call_next(tool_call)
|
||||
|
||||
# Create agent with hooks
|
||||
agent = Agent(
|
||||
name="WeatherBot",
|
||||
instructions="You help with weather queries.",
|
||||
tools=[get_weather],
|
||||
hooks=[add_context, log_response, retry_on_error]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test the tool directly
|
||||
result = get_weather("London")
|
||||
print(f"Weather: {result}")
|
||||
|
||||
print("\n✓ Middleware example complete")
|
||||
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
Configurable Model Example - PraisonAI Agents
|
||||
|
||||
Demonstrates runtime model switching without recreating the agent.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Create agent with configurable model enabled
|
||||
agent = Agent(
|
||||
name="FlexBot",
|
||||
instructions="You are a helpful assistant.",
|
||||
llm="gpt-4o-mini",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Default model call
|
||||
print("Using default model (gpt-4o-mini)...")
|
||||
# response = agent.chat("Say hello in 5 words")
|
||||
|
||||
# Override model per-call
|
||||
print("\nOverriding to different model...")
|
||||
# response = agent.chat("Say hello in 5 words", config={"model": "gpt-4o"})
|
||||
|
||||
# Override temperature
|
||||
print("\nOverriding temperature...")
|
||||
# response = agent.chat("Say hello creatively", config={"temperature": 0.9})
|
||||
|
||||
print("\n✓ Configurable model example complete")
|
||||
print("Note: Uncomment agent.chat() calls to run with API key")
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
Injected State Example - PraisonAI Agents
|
||||
|
||||
Demonstrates injecting agent state into tools without exposing it in the schema.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, tool
|
||||
from praisonaiagents.tools import Injected
|
||||
from praisonaiagents.tools.injected import AgentState, with_injection_context
|
||||
|
||||
# Tool with injected state - state param is NOT in the public schema
|
||||
@tool
|
||||
def show_context(query: str, state: Injected[dict]) -> str:
|
||||
"""Show the current agent context."""
|
||||
session_id = state.get('session_id', 'unknown')
|
||||
agent_id = state.get('agent_id', 'unknown')
|
||||
return f"Query: {query}, Session: {session_id}, Agent: {agent_id}"
|
||||
|
||||
# Create agent with the tool
|
||||
agent = Agent(
|
||||
name="ContextBot",
|
||||
instructions="You help show context information.",
|
||||
tools=[show_context],
|
||||
session_id="my-session-123"
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Verify injected param is not in schema
|
||||
schema = show_context.get_schema()
|
||||
params = schema['function']['parameters']['properties']
|
||||
print(f"Schema params: {list(params.keys())}")
|
||||
assert 'state' not in params, "state should NOT be in schema"
|
||||
print("✓ 'state' correctly excluded from schema")
|
||||
|
||||
# Test with manual injection context
|
||||
mock_state = AgentState(
|
||||
agent_id="test-agent",
|
||||
run_id="run-1",
|
||||
session_id="session-abc"
|
||||
)
|
||||
|
||||
with with_injection_context(mock_state):
|
||||
result = show_context(query="hello")
|
||||
print(f"Result: {result}")
|
||||
|
||||
# Test via agent.execute_tool
|
||||
result = agent.execute_tool("show_context", {"query": "test"})
|
||||
print(f"Agent result: {result}")
|
||||
|
||||
print("\n✓ Injected state example complete")
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
Rate Limiter Example - PraisonAI Agents
|
||||
|
||||
Demonstrates token bucket rate limiting for LLM API calls.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
# Create rate limiter: 60 requests per minute
|
||||
limiter = RateLimiter(requests_per_minute=60, burst=5)
|
||||
|
||||
# Create agent with rate limiter
|
||||
agent = Agent(
|
||||
name="RateLimitedBot",
|
||||
instructions="You are a helpful assistant.",
|
||||
rate_limiter=limiter
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(f"Rate limiter: {limiter}")
|
||||
print(f"Available tokens: {limiter.available_tokens}")
|
||||
|
||||
# Demonstrate rate limiting
|
||||
for i in range(3):
|
||||
limiter.acquire()
|
||||
print(f"Request {i+1} acquired, tokens left: {limiter.available_tokens:.1f}")
|
||||
|
||||
print("\n✓ Rate limiter example complete")
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
Multi-Agent Persistence Example
|
||||
|
||||
Demonstrates how multiple agents can share or isolate conversation
|
||||
history using the persistence layer.
|
||||
"""
|
||||
|
||||
from praisonai.persistence.factory import create_conversation_store, create_state_store
|
||||
from praisonai.persistence.conversation.base import ConversationSession, ConversationMessage
|
||||
import uuid
|
||||
|
||||
|
||||
def main():
|
||||
# Create stores
|
||||
conv_store = create_conversation_store("sqlite", path=":memory:")
|
||||
state_store = create_state_store("memory")
|
||||
|
||||
print("=== Multi-Agent Persistence Demo ===\n")
|
||||
|
||||
# Scenario 1: Shared session (agents collaborate)
|
||||
print("--- Scenario 1: Shared Session ---")
|
||||
shared_session_id = f"shared-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
session = ConversationSession(
|
||||
session_id=shared_session_id,
|
||||
user_id="user-123",
|
||||
name="Shared Conversation"
|
||||
)
|
||||
conv_store.create_session(session)
|
||||
|
||||
# Agent 1 (Researcher) adds context
|
||||
conv_store.add_message(shared_session_id, ConversationMessage(
|
||||
role="assistant",
|
||||
content="[Researcher] I found relevant information about the topic.",
|
||||
metadata={"agent_id": "researcher", "agent_role": "research"}
|
||||
))
|
||||
|
||||
# Agent 2 (Writer) uses the context
|
||||
conv_store.add_message(shared_session_id, ConversationMessage(
|
||||
role="assistant",
|
||||
content="[Writer] Based on the research, here's a summary...",
|
||||
metadata={"agent_id": "writer", "agent_role": "writing"}
|
||||
))
|
||||
|
||||
# Agent 3 (Editor) reviews
|
||||
conv_store.add_message(shared_session_id, ConversationMessage(
|
||||
role="assistant",
|
||||
content="[Editor] I've reviewed and polished the content.",
|
||||
metadata={"agent_id": "editor", "agent_role": "editing"}
|
||||
))
|
||||
|
||||
messages = conv_store.get_messages(shared_session_id)
|
||||
print(f"Shared session has {len(messages)} messages from different agents")
|
||||
for msg in messages:
|
||||
agent = msg.metadata.get("agent_id", "unknown") if msg.metadata else "unknown"
|
||||
print(f" - {agent}: {msg.content[:50]}...")
|
||||
|
||||
# Scenario 2: Isolated sessions (agents work independently)
|
||||
print("\n--- Scenario 2: Isolated Sessions ---")
|
||||
|
||||
agent_sessions = {}
|
||||
for agent_name in ["planner", "executor", "validator"]:
|
||||
session_id = f"{agent_name}-{uuid.uuid4().hex[:8]}"
|
||||
agent_sessions[agent_name] = session_id
|
||||
|
||||
conv_store.create_session(ConversationSession(
|
||||
session_id=session_id,
|
||||
agent_id=agent_name,
|
||||
name=f"{agent_name.title()} Session"
|
||||
))
|
||||
|
||||
# Each agent has its own conversation
|
||||
conv_store.add_message(session_id, ConversationMessage(
|
||||
role="user",
|
||||
content=f"Task for {agent_name}"
|
||||
))
|
||||
conv_store.add_message(session_id, ConversationMessage(
|
||||
role="assistant",
|
||||
content=f"[{agent_name.title()}] Processing task..."
|
||||
))
|
||||
|
||||
print("Created isolated sessions for each agent:")
|
||||
for agent, sid in agent_sessions.items():
|
||||
msgs = conv_store.get_messages(sid)
|
||||
print(f" - {agent}: {len(msgs)} messages")
|
||||
|
||||
# Scenario 3: Shared state across agents
|
||||
print("\n--- Scenario 3: Shared State ---")
|
||||
|
||||
# Agents share state for coordination
|
||||
state_store.set("workflow:status", {"phase": "research", "progress": 0})
|
||||
state_store.set("workflow:context", {"topic": "AI Agents", "deadline": "2024-12-31"})
|
||||
|
||||
# Agent 1 updates progress
|
||||
state = state_store.get("workflow:status")
|
||||
state["progress"] = 33
|
||||
state["phase"] = "writing"
|
||||
state_store.set("workflow:status", state)
|
||||
print(f"Agent 1 updated state: {state}")
|
||||
|
||||
# Agent 2 reads and updates
|
||||
state = state_store.get("workflow:status")
|
||||
state["progress"] = 66
|
||||
state["phase"] = "editing"
|
||||
state_store.set("workflow:status", state)
|
||||
print(f"Agent 2 updated state: {state}")
|
||||
|
||||
# Agent 3 completes
|
||||
state = state_store.get("workflow:status")
|
||||
state["progress"] = 100
|
||||
state["phase"] = "complete"
|
||||
state_store.set("workflow:status", state)
|
||||
print(f"Agent 3 completed: {state}")
|
||||
|
||||
# Cleanup
|
||||
conv_store.close()
|
||||
state_store.close()
|
||||
print("\n✅ Multi-agent persistence demo complete!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Multi-Agent Shared Session - DB Persistence Test"""
|
||||
import sys
|
||||
sys.path.insert(0, 'src/praisonai')
|
||||
from praisonai.persistence import create_conversation_store
|
||||
from praisonai.persistence.conversation.base import ConversationSession, ConversationMessage
|
||||
|
||||
# Create shared session for multiple agents
|
||||
store = create_conversation_store("sqlite", path="/tmp/multi_agent.db")
|
||||
session_id = "shared-session-123"
|
||||
|
||||
# Create session
|
||||
session = ConversationSession(session_id=session_id, agent_id="multi")
|
||||
try:
|
||||
store.create_session(session)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Agent 1 writes
|
||||
msg1 = ConversationMessage(session_id=session_id, role="user", content="Agent1: Hello")
|
||||
store.add_message(session_id, msg1)
|
||||
msg2 = ConversationMessage(session_id=session_id, role="assistant", content="Agent1: Hi there!")
|
||||
store.add_message(session_id, msg2)
|
||||
|
||||
# Agent 2 writes to same session
|
||||
msg3 = ConversationMessage(session_id=session_id, role="user", content="Agent2: What did Agent1 say?")
|
||||
store.add_message(session_id, msg3)
|
||||
msg4 = ConversationMessage(session_id=session_id, role="assistant", content="Agent2: Agent1 said Hello")
|
||||
store.add_message(session_id, msg4)
|
||||
|
||||
# Verify both agents' messages are in session
|
||||
messages = store.get_messages(session_id)
|
||||
print(f"Total messages in shared session: {len(messages)}")
|
||||
assert len(messages) >= 4, "Should have messages from both agents"
|
||||
|
||||
# Verify content
|
||||
contents = [m.content for m in messages]
|
||||
assert any("Agent1" in c for c in contents), "Should have Agent1 messages"
|
||||
assert any("Agent2" in c for c in contents), "Should have Agent2 messages"
|
||||
|
||||
print("PASSED: Multi-agent shared session persistence")
|
||||
@@ -0,0 +1,16 @@
|
||||
"""Multi-Agent Shared Session Demo"""
|
||||
from praisonaiagents import Agent
|
||||
from praisonai.persistence import create_conversation_store
|
||||
|
||||
store = create_conversation_store("sqlite", path="/tmp/multi_agent.db")
|
||||
|
||||
agent1 = Agent(name="Researcher", llm="gpt-4o-mini", db=store, session_id="shared-session")
|
||||
agent2 = Agent(name="Writer", llm="gpt-4o-mini", db=store, session_id="shared-session")
|
||||
|
||||
r1 = agent1.chat("Research topic: AI in healthcare. Key finding: reduces diagnosis time by 50%")
|
||||
print(f"Researcher: {r1[:50]}...")
|
||||
|
||||
r2 = agent2.chat("Based on the research, what was the key finding?")
|
||||
print(f"Writer: {r2[:50]}...")
|
||||
assert "50" in r2 or "diagnosis" in r2.lower(), "Writer should see Researcher's context"
|
||||
print("PASSED: Multi-agent shared session works")
|
||||
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
Basic Observability Usage Example
|
||||
|
||||
This example demonstrates how to use the observability module
|
||||
with PraisonAI agents.
|
||||
|
||||
Usage:
|
||||
# Set your provider's API key first
|
||||
export LANGFUSE_PUBLIC_KEY=pk-lf-xxx
|
||||
export LANGFUSE_SECRET_KEY=sk-lf-xxx
|
||||
|
||||
python basic_usage.py
|
||||
"""
|
||||
|
||||
from praisonai_tools.observability import obs
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Initialize observability (auto-detects provider from env vars)
|
||||
obs.init()
|
||||
|
||||
# Or specify a provider explicitly
|
||||
# obs.init(provider="langfuse", project_name="my-project")
|
||||
|
||||
# Create an agent
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant.",
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
|
||||
# Use trace context manager for complete workflow tracing
|
||||
with obs.trace("chat-session", session_id="user-123"):
|
||||
# Use span context manager for individual operations
|
||||
with obs.span("user-query", kind=obs.SpanKind.AGENT):
|
||||
response = agent.chat("What is the capital of France?")
|
||||
print(response)
|
||||
|
||||
# Log LLM calls manually if needed
|
||||
obs.log_llm_call(
|
||||
model="gpt-4o-mini",
|
||||
input_messages="What is 2+2?",
|
||||
output="4",
|
||||
input_tokens=10,
|
||||
output_tokens=1,
|
||||
)
|
||||
|
||||
# Check diagnostics
|
||||
print("\nDiagnostics:")
|
||||
print(obs.doctor())
|
||||
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
Langfuse Integration Example
|
||||
|
||||
This example shows how to use Langfuse for LLM observability.
|
||||
|
||||
Setup:
|
||||
1. Sign up at https://langfuse.com/
|
||||
2. Get your API keys from the project settings
|
||||
3. Set environment variables:
|
||||
export LANGFUSE_PUBLIC_KEY=pk-lf-xxx
|
||||
export LANGFUSE_SECRET_KEY=sk-lf-xxx
|
||||
4. Install dependencies:
|
||||
pip install opentelemetry-sdk opentelemetry-exporter-otlp
|
||||
|
||||
Usage:
|
||||
python langfuse_example.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools.observability import obs
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Initialize Langfuse
|
||||
success = obs.init(
|
||||
provider="langfuse",
|
||||
project_name="praisonai-demo",
|
||||
)
|
||||
|
||||
if not success:
|
||||
print("Failed to initialize Langfuse. Check your API keys.")
|
||||
print("Required: LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY")
|
||||
exit(1)
|
||||
|
||||
print("Langfuse initialized successfully!")
|
||||
print(f"View traces at: https://cloud.langfuse.com/")
|
||||
|
||||
# Create agent
|
||||
agent = Agent(
|
||||
instructions="You are a helpful coding assistant.",
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
|
||||
# Run with tracing
|
||||
with obs.trace("coding-session", user_id="developer-1"):
|
||||
response = agent.chat("Write a Python function to calculate fibonacci numbers")
|
||||
print(response)
|
||||
|
||||
print("\nCheck Langfuse dashboard for traces!")
|
||||
@@ -0,0 +1,22 @@
|
||||
"""MLflow Local Observability Demo (no API key needed)"""
|
||||
import mlflow
|
||||
import time
|
||||
|
||||
mlflow.set_tracking_uri("file:///tmp/mlflow_demo")
|
||||
mlflow.set_experiment("praisonai-demo")
|
||||
|
||||
with mlflow.start_run(run_name="agent-trace"):
|
||||
mlflow.log_param("agent_name", "Assistant")
|
||||
mlflow.log_param("model", "gpt-4o-mini")
|
||||
|
||||
# Simulate agent work
|
||||
start = time.time()
|
||||
time.sleep(0.1)
|
||||
latency = time.time() - start
|
||||
|
||||
mlflow.log_metric("latency_ms", latency * 1000)
|
||||
mlflow.log_metric("tokens_used", 150)
|
||||
|
||||
print(f"Logged run to MLflow: latency={latency*1000:.1f}ms")
|
||||
|
||||
print("PASSED: MLflow local observability works")
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
Multi-Agent Tracing Example
|
||||
|
||||
This example shows how to trace multi-agent workflows with
|
||||
parent/child span relationships.
|
||||
|
||||
Usage:
|
||||
# Set your provider's API key
|
||||
export LANGFUSE_PUBLIC_KEY=pk-lf-xxx
|
||||
export LANGFUSE_SECRET_KEY=sk-lf-xxx
|
||||
|
||||
python multi_agent_tracing.py
|
||||
"""
|
||||
|
||||
from praisonai_tools.observability import obs
|
||||
from praisonai_tools.observability.base import SpanKind
|
||||
from praisonaiagents import Agent, PraisonAIAgents
|
||||
|
||||
# Initialize observability
|
||||
obs.init()
|
||||
|
||||
# Create agents
|
||||
researcher = Agent(
|
||||
name="Researcher",
|
||||
instructions="You research topics and provide detailed information.",
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
|
||||
writer = Agent(
|
||||
name="Writer",
|
||||
instructions="You write clear, engaging content based on research.",
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
|
||||
# Trace the multi-agent workflow
|
||||
with obs.trace("content-creation", metadata={"workflow": "research-write"}):
|
||||
|
||||
# Research phase
|
||||
with obs.span("research-phase", kind=SpanKind.AGENT) as research_span:
|
||||
research_span.attributes["agent"] = "Researcher"
|
||||
research_result = researcher.chat("Research the benefits of AI in healthcare")
|
||||
|
||||
# Writing phase
|
||||
with obs.span("writing-phase", kind=SpanKind.AGENT) as write_span:
|
||||
write_span.attributes["agent"] = "Writer"
|
||||
write_span.attributes["input_length"] = len(str(research_result))
|
||||
final_content = writer.chat(f"Write a blog post based on: {research_result}")
|
||||
|
||||
print("Final Content:")
|
||||
print(final_content)
|
||||
|
||||
# Check trace info
|
||||
print("\nDiagnostics:", obs.doctor())
|
||||
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
Qdrant Knowledge Store - Agent-First Example
|
||||
|
||||
Demonstrates using Qdrant for knowledge-based RAG with an Agent.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name praison-qdrant -p 6333:6333 qdrant/qdrant
|
||||
|
||||
Run:
|
||||
python knowledge_qdrant.py
|
||||
|
||||
Expected Output:
|
||||
Agent responds using knowledge from documents
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
import tempfile
|
||||
import os
|
||||
import shutil
|
||||
|
||||
# Create sample knowledge document
|
||||
sample_doc = """
|
||||
# AI and Programming Guide
|
||||
|
||||
## Python Programming
|
||||
Python is a versatile programming language used for web development, data science, and AI.
|
||||
It has a simple syntax and extensive libraries.
|
||||
|
||||
## Machine Learning
|
||||
Machine learning enables computers to learn patterns from data without explicit programming.
|
||||
Common frameworks include TensorFlow, PyTorch, and scikit-learn.
|
||||
|
||||
## Databases
|
||||
PostgreSQL is a powerful open-source relational database system.
|
||||
It supports both SQL and JSON data types.
|
||||
"""
|
||||
|
||||
# Save to temp file
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
doc_path = os.path.join(temp_dir, "guide.txt")
|
||||
with open(doc_path, "w") as f:
|
||||
f.write(sample_doc)
|
||||
|
||||
print("=== Qdrant Knowledge Store Demo (Agent-First) ===")
|
||||
|
||||
# Agent-first approach: use knowledge parameter with Qdrant
|
||||
agent = Agent(
|
||||
name="KnowledgeAssistant",
|
||||
instructions="You are a helpful assistant with access to technical documentation.",
|
||||
knowledge=[doc_path],
|
||||
knowledge_config={
|
||||
"vector_store": "qdrant",
|
||||
"url": "http://localhost:6333"
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What programming language is good for AI?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Cleanup
|
||||
shutil.rmtree(temp_dir)
|
||||
|
||||
print("\n=== Demo Complete ===")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonai.persistence.factory import create_knowledge_store
|
||||
# store = create_knowledge_store("qdrant", url="http://localhost:6333")
|
||||
@@ -0,0 +1,43 @@
|
||||
"""
|
||||
Minimal Agent with Database Persistence
|
||||
|
||||
This is the simplest way to add persistence to your agent.
|
||||
Messages, runs, and traces are automatically saved.
|
||||
|
||||
Run:
|
||||
python minimal_agent_db.py
|
||||
|
||||
Expected output:
|
||||
- Agent responds to your question
|
||||
- Data is persisted to PostgreSQL
|
||||
- Session can be resumed later with same session_id
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create database connection (PostgreSQL + Redis)
|
||||
my_db = db(
|
||||
database_url="postgresql://postgres:praison123@localhost:5432/praisonai",
|
||||
state_url="redis://localhost:6379" # Optional: for state/tracing
|
||||
)
|
||||
|
||||
# Create agent with persistence
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant. Be concise.",
|
||||
db=my_db,
|
||||
# session_id is optional - defaults to per-hour ID (YYYYMMDDHH-hash)
|
||||
# Set explicitly for session continuity across runs:
|
||||
session_id="my-persistent-session"
|
||||
)
|
||||
|
||||
# Chat - messages are automatically persisted
|
||||
response = agent.chat("What is the capital of France?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Verify persistence
|
||||
data = my_db.export_session("my-persistent-session")
|
||||
print(f"\nMessages stored: {len(data.get('messages', []))}")
|
||||
|
||||
# Clean up
|
||||
my_db.close()
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
PostgreSQL with Runs and Traces
|
||||
|
||||
Demonstrates full persistence with run tracking and tracing.
|
||||
|
||||
Run:
|
||||
python postgres_runs_traces.py
|
||||
|
||||
Expected output:
|
||||
- Agent responds with tool usage
|
||||
- Runs and traces are persisted
|
||||
- Export shows full session data
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
def get_weather(city: str) -> str:
|
||||
"""Get weather for a city."""
|
||||
return f"Weather in {city}: Sunny, 72°F"
|
||||
|
||||
# Create database with PostgreSQL + Redis
|
||||
my_db = db(
|
||||
database_url="postgresql://postgres:praison123@localhost:5432/praisonai",
|
||||
state_url="redis://localhost:6379"
|
||||
)
|
||||
|
||||
# Create agent with tools
|
||||
agent = Agent(
|
||||
name="WeatherBot",
|
||||
instructions="You help with weather. Use the get_weather tool when asked.",
|
||||
tools=[get_weather],
|
||||
db=my_db,
|
||||
session_id="weather-session-001"
|
||||
)
|
||||
|
||||
# First chat - triggers run tracking
|
||||
print("=== First Question ===")
|
||||
response1 = agent.chat("What's the weather in Tokyo?")
|
||||
print(f"Response: {response1}\n")
|
||||
|
||||
# Second chat - same session, new run
|
||||
print("=== Second Question ===")
|
||||
response2 = agent.chat("How about New York?")
|
||||
print(f"Response: {response2}\n")
|
||||
|
||||
# Export session data
|
||||
print("=== Session Data ===")
|
||||
data = my_db.export_session("weather-session-001")
|
||||
print(f"Session ID: {data.get('session_id')}")
|
||||
print(f"Messages: {len(data.get('messages', []))}")
|
||||
print(f"Agent: {data.get('agent_name')}")
|
||||
|
||||
# Show messages
|
||||
print("\n=== Messages ===")
|
||||
for msg in data.get('messages', []):
|
||||
role = msg.get('role', 'unknown')
|
||||
content = msg.get('content', '')[:60]
|
||||
print(f" [{role}] {content}...")
|
||||
|
||||
my_db.close()
|
||||
print("\n✅ Done")
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Redis State Store - Agent-First Example
|
||||
|
||||
Use Redis for state management with an Agent.
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name redis -p 6379:6379 redis:7
|
||||
|
||||
Run:
|
||||
python redis_state.py
|
||||
|
||||
Expected output:
|
||||
Agent responds with state persisted to Redis
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
print("=== Redis State Store (Agent-First) ===")
|
||||
|
||||
# Agent-first approach: use db parameter with Redis state store
|
||||
my_db = db(
|
||||
database_url="sqlite:///conversations.db", # Conversations
|
||||
state_url="redis://localhost:6379" # State/tracing
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant.",
|
||||
db=my_db,
|
||||
session_id="redis-state-session"
|
||||
)
|
||||
|
||||
# Chat - state is automatically managed via Redis
|
||||
response = agent.chat("Hello! Remember my name is Alice.")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Second message to test state
|
||||
response2 = agent.chat("What is my name?")
|
||||
print(f"Response: {response2}")
|
||||
|
||||
my_db.close()
|
||||
print("\n✅ Done")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonaiagents import db
|
||||
# redis_db = db.RedisDB(host="localhost", port=6379)
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
Session Resume Example
|
||||
|
||||
Demonstrates how to resume a conversation from a previous session.
|
||||
|
||||
Run:
|
||||
python session_resume.py
|
||||
|
||||
Expected output:
|
||||
- First run: Creates new session
|
||||
- Second run: Resumes with history
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create database
|
||||
my_db = db(
|
||||
database_url="postgresql://postgres:praison123@localhost:5432/praisonai",
|
||||
state_url="redis://localhost:6379"
|
||||
)
|
||||
|
||||
SESSION_ID = "resume-demo-session"
|
||||
|
||||
# Create agent with explicit session_id
|
||||
agent = Agent(
|
||||
name="MemoryBot",
|
||||
instructions="You remember our conversation. Be helpful and concise.",
|
||||
db=my_db,
|
||||
session_id=SESSION_ID
|
||||
)
|
||||
|
||||
# Check if this is a new or resumed session
|
||||
data = my_db.export_session(SESSION_ID)
|
||||
existing_messages = len(data.get('messages', []))
|
||||
|
||||
if existing_messages == 0:
|
||||
print("=== NEW SESSION ===")
|
||||
print("This is a new conversation.\n")
|
||||
|
||||
response = agent.chat("My name is Alice and I love Python programming.")
|
||||
print(f"Response: {response}\n")
|
||||
|
||||
response = agent.chat("What's my favorite programming language?")
|
||||
print(f"Response: {response}\n")
|
||||
else:
|
||||
print(f"=== RESUMED SESSION ({existing_messages} previous messages) ===")
|
||||
print("Continuing previous conversation.\n")
|
||||
|
||||
# Show previous messages
|
||||
print("Previous messages:")
|
||||
for msg in data.get('messages', [])[:4]:
|
||||
role = msg.get('role', 'unknown')
|
||||
content = msg.get('content', '')[:60]
|
||||
print(f" [{role}] {content}...")
|
||||
print()
|
||||
|
||||
# Continue conversation
|
||||
response = agent.chat("What's my name and what do I like?")
|
||||
print(f"Response: {response}\n")
|
||||
|
||||
# Show final state
|
||||
final_data = my_db.export_session(SESSION_ID)
|
||||
print(f"Total messages now: {len(final_data.get('messages', []))}")
|
||||
|
||||
my_db.close()
|
||||
print("✅ Done")
|
||||
@@ -0,0 +1,58 @@
|
||||
"""
|
||||
Simple Database Persistence Example for PraisonAI Agents.
|
||||
|
||||
This example shows how to enable automatic conversation persistence
|
||||
with just 2 lines of code.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name praison-postgres -p 5432:5432 \
|
||||
-e POSTGRES_PASSWORD=praison123 \
|
||||
-e POSTGRES_DB=praisonai \
|
||||
postgres:16
|
||||
|
||||
Environment:
|
||||
export OPENAI_API_KEY=your-key
|
||||
|
||||
Run:
|
||||
python simple_agent_postgres.py
|
||||
|
||||
Expected Output:
|
||||
=== First Run ===
|
||||
Agent: Nice to meet you, Alice! Python is a great choice...
|
||||
Session ID: demo-session-001
|
||||
Messages persisted: 2
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create database adapter (simplified import)
|
||||
db_instance = db(
|
||||
database_url="postgresql://postgres:praison123@localhost:5432/praisonai"
|
||||
)
|
||||
|
||||
# Create agent with persistence enabled
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant. Keep responses brief.",
|
||||
db=db_instance,
|
||||
session_id="demo-session-001", # Same session_id = resume conversation
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Chat with persistence
|
||||
print("=== First Run ===")
|
||||
response = agent.chat("My name is Alice and I love Python programming.")
|
||||
print(f"Agent: {response}\n")
|
||||
|
||||
# Show session info
|
||||
print(f"Session ID: {agent.session_id}")
|
||||
print(f"Messages in history: {len(agent.chat_history)}")
|
||||
|
||||
# Close DB connection
|
||||
db_instance.close()
|
||||
|
||||
print("\n=== Session persisted! ===")
|
||||
print("Run again with same session_id to resume.")
|
||||
@@ -0,0 +1,41 @@
|
||||
"""
|
||||
SQLite Local Database
|
||||
|
||||
Use SQLite for simple local persistence without external services.
|
||||
|
||||
Run:
|
||||
python sqlite_local.py
|
||||
|
||||
Expected output:
|
||||
- Agent responds
|
||||
- Data persisted to local SQLite file
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
# Create SQLite database (local file)
|
||||
db_path = os.path.join(tempfile.gettempdir(), "my_agent.db")
|
||||
my_db = db.SQLiteDB(path=db_path)
|
||||
|
||||
print(f"Database: {db_path}")
|
||||
|
||||
# Create agent
|
||||
agent = Agent(
|
||||
name="LocalBot",
|
||||
instructions="You are a helpful assistant.",
|
||||
db=my_db,
|
||||
session_id="local-session"
|
||||
)
|
||||
|
||||
# Chat
|
||||
response = agent.chat("Hello! What can you help me with?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Verify
|
||||
data = my_db.export_session("local-session")
|
||||
print(f"\nMessages stored: {len(data.get('messages', []))}")
|
||||
|
||||
my_db.close()
|
||||
print("✅ Done")
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Redis State Store - Agent-First Example
|
||||
|
||||
Use Redis for state management with an Agent.
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name praison-redis -p 6379:6379 redis:7
|
||||
|
||||
Run:
|
||||
python state_redis.py
|
||||
|
||||
Expected Output:
|
||||
Agent responds with state persisted to Redis
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
print("=== Redis State Store (Agent-First) ===")
|
||||
|
||||
# Agent-first approach: use db parameter with Redis state store
|
||||
my_db = db(
|
||||
database_url="sqlite:///conversations.db", # Conversations
|
||||
state_url="redis://localhost:6379" # State/tracing
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant.",
|
||||
db=my_db,
|
||||
session_id="redis-state-example"
|
||||
)
|
||||
|
||||
# Chat - state is automatically managed via Redis
|
||||
response = agent.chat("Hello! Remember my favorite color is blue.")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Second message to test state
|
||||
response2 = agent.chat("What is my favorite color?")
|
||||
print(f"Response: {response2}")
|
||||
|
||||
my_db.close()
|
||||
print("\n=== Demo Complete ===")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonai.persistence.factory import create_state_store
|
||||
# store = create_state_store("redis", url="redis://localhost:6379")
|
||||
@@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Policy Engine Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use the policy engine to:
|
||||
1. Create policies with rules
|
||||
2. Check if actions are allowed
|
||||
3. Block dangerous operations
|
||||
4. Use convenience policy creators
|
||||
|
||||
Usage:
|
||||
python basic_policy.py
|
||||
"""
|
||||
|
||||
from praisonaiagents.policy import (
|
||||
PolicyEngine, Policy, PolicyRule, PolicyAction, PolicyConfig,
|
||||
create_deny_tools_policy, create_read_only_policy
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("Policy Engine Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a policy engine
|
||||
engine = PolicyEngine()
|
||||
|
||||
# ==========================================================================
|
||||
# Policy 1: Block dangerous tools
|
||||
# ==========================================================================
|
||||
print("\n--- Creating Policies ---")
|
||||
|
||||
dangerous_tools_policy = Policy(
|
||||
name="block_dangerous",
|
||||
description="Block dangerous file operations",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:delete_*",
|
||||
reason="Delete operations are blocked",
|
||||
name="deny_delete"
|
||||
),
|
||||
PolicyRule(
|
||||
action=PolicyAction.DENY,
|
||||
resource="tool:rm_*",
|
||||
reason="Remove operations are blocked",
|
||||
name="deny_rm"
|
||||
),
|
||||
],
|
||||
priority=100
|
||||
)
|
||||
engine.add_policy(dangerous_tools_policy)
|
||||
print(f"✅ Added policy: {dangerous_tools_policy.name}")
|
||||
|
||||
# ==========================================================================
|
||||
# Policy 2: Allow read operations
|
||||
# ==========================================================================
|
||||
read_policy = Policy(
|
||||
name="allow_read",
|
||||
description="Allow read operations",
|
||||
rules=[
|
||||
PolicyRule(
|
||||
action=PolicyAction.ALLOW,
|
||||
resource="tool:read_*",
|
||||
name="allow_read"
|
||||
),
|
||||
PolicyRule(
|
||||
action=PolicyAction.ALLOW,
|
||||
resource="tool:list_*",
|
||||
name="allow_list"
|
||||
),
|
||||
]
|
||||
)
|
||||
engine.add_policy(read_policy)
|
||||
print(f"✅ Added policy: {read_policy.name}")
|
||||
|
||||
# ==========================================================================
|
||||
# Test policy checks
|
||||
# ==========================================================================
|
||||
print("\n--- Testing Policy Checks ---")
|
||||
|
||||
test_cases = [
|
||||
("tool:read_file", "Reading a file"),
|
||||
("tool:list_directory", "Listing directory"),
|
||||
("tool:delete_file", "Deleting a file"),
|
||||
("tool:rm_directory", "Removing directory"),
|
||||
("tool:write_file", "Writing a file"),
|
||||
]
|
||||
|
||||
for resource, description in test_cases:
|
||||
result = engine.check(resource, {})
|
||||
status = "✅ ALLOWED" if result.allowed else "🚫 DENIED"
|
||||
reason = f" ({result.reason})" if result.reason else ""
|
||||
print(f" {status}: {description}{reason}")
|
||||
|
||||
# ==========================================================================
|
||||
# Using convenience functions
|
||||
# ==========================================================================
|
||||
print("\n--- Using Convenience Functions ---")
|
||||
|
||||
# Create a read-only policy
|
||||
read_only = create_read_only_policy()
|
||||
print(f"✅ Created read-only policy with {len(read_only.rules)} rules")
|
||||
|
||||
# Create a deny tools policy
|
||||
deny_tools = create_deny_tools_policy(
|
||||
["execute_*", "shell_*", "system_*"],
|
||||
reason="System commands are blocked"
|
||||
)
|
||||
print(f"✅ Created deny-tools policy with {len(deny_tools.rules)} rules")
|
||||
|
||||
# ==========================================================================
|
||||
# Strict mode example
|
||||
# ==========================================================================
|
||||
print("\n--- Strict Mode Example ---")
|
||||
|
||||
strict_engine = PolicyEngine(PolicyConfig(strict_mode=True))
|
||||
|
||||
# In strict mode, unknown resources are denied
|
||||
result = strict_engine.check("tool:unknown_operation", {})
|
||||
print(f" Unknown operation in strict mode: {'DENIED' if not result.allowed else 'ALLOWED'}")
|
||||
|
||||
# ==========================================================================
|
||||
# Policy serialization
|
||||
# ==========================================================================
|
||||
print("\n--- Policy Serialization ---")
|
||||
|
||||
policy_dict = dangerous_tools_policy.to_dict()
|
||||
print(f" Serialized policy: {policy_dict['name']} with {len(policy_dict['rules'])} rules")
|
||||
|
||||
# Recreate from dict
|
||||
restored = Policy.from_dict(policy_dict)
|
||||
print(f" Restored policy: {restored.name}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -15,6 +15,7 @@ Features demonstrated:
|
||||
from praisonaiagents import Agent
|
||||
import tempfile
|
||||
import os
|
||||
import shutil
|
||||
|
||||
# Create a sample knowledge document for demonstration
|
||||
sample_document = """
|
||||
@@ -91,38 +92,15 @@ doc_path = os.path.join(temp_dir, "ai_healthcare_guide.txt")
|
||||
with open(doc_path, "w") as f:
|
||||
f.write(sample_document)
|
||||
|
||||
# Create an agent with advanced knowledge management configuration
|
||||
# Create an agent with knowledge management
|
||||
knowledge_agent = Agent(
|
||||
name="AdvancedKnowledgeAgent",
|
||||
role="Healthcare AI Knowledge Expert",
|
||||
goal="Provide comprehensive answers about AI in healthcare using advanced knowledge retrieval",
|
||||
backstory="You are an expert in healthcare AI with access to comprehensive knowledge bases and advanced retrieval capabilities.",
|
||||
|
||||
# Advanced knowledge configuration
|
||||
# Knowledge configuration - simple and working
|
||||
knowledge=[doc_path],
|
||||
knowledge_config={
|
||||
"vector_store": {
|
||||
"provider": "chroma",
|
||||
"collection_name": "healthcare_ai_knowledge"
|
||||
},
|
||||
"chunking": {
|
||||
"strategy": "semantic", # Advanced semantic chunking
|
||||
"chunk_size": 500, # Optimal chunk size for retrieval
|
||||
"chunk_overlap": 50, # Overlap for context preservation
|
||||
"separators": ["\n\n", "\n", ".", "!", "?"] # Smart separators
|
||||
},
|
||||
"retrieval": {
|
||||
"search_type": "similarity",
|
||||
"k": 5, # Retrieve top 5 most relevant chunks
|
||||
"score_threshold": 0.7, # Minimum relevance score
|
||||
"rerank": True # Enable reranking for better results
|
||||
},
|
||||
"embedding": {
|
||||
"provider": "openai",
|
||||
"model": "text-embedding-3-small", # Efficient embedding model
|
||||
"dimensions": 1536
|
||||
}
|
||||
},
|
||||
|
||||
instructions="""You are an expert in healthcare AI. Use the knowledge base to provide
|
||||
comprehensive, accurate answers. Always cite specific information from the documents
|
||||
@@ -180,7 +158,6 @@ relationship_result = knowledge_agent.start(
|
||||
print(f"Relationship analysis result:\n{relationship_result}\n")
|
||||
|
||||
# Clean up temporary files
|
||||
import shutil
|
||||
shutil.rmtree(temp_dir)
|
||||
|
||||
print("="*80)
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run all examples and collect PASS/SKIP/FAIL results."""
|
||||
import subprocess
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
EXAMPLES_DIR = Path(__file__).parent
|
||||
RESULTS = []
|
||||
|
||||
def run_example(path: Path) -> tuple:
|
||||
"""Run an example and return (name, status, reason)."""
|
||||
name = f"{path.parent.name}/{path.name}"
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(path)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
env={**os.environ, "PYTHONPATH": str(EXAMPLES_DIR.parent / "src")}
|
||||
)
|
||||
output = result.stdout + result.stderr
|
||||
if "PASSED" in output:
|
||||
return (name, "PASS", "")
|
||||
elif "SKIPPED" in output:
|
||||
return (name, "SKIP", output.split("SKIPPED:")[-1].strip()[:50])
|
||||
elif result.returncode != 0:
|
||||
error = output.strip().split("\n")[-1][:80]
|
||||
return (name, "FAIL", error)
|
||||
else:
|
||||
return (name, "PASS", "")
|
||||
except subprocess.TimeoutExpired:
|
||||
return (name, "FAIL", "Timeout")
|
||||
except Exception as e:
|
||||
return (name, "FAIL", str(e)[:50])
|
||||
|
||||
def main():
|
||||
print("=" * 70)
|
||||
print("EXAMPLE RUNNER - PASS/SKIP/FAIL SUMMARY")
|
||||
print("=" * 70)
|
||||
|
||||
# Find all example files
|
||||
examples = []
|
||||
for subdir in ["db", "vector", "observability", "multi_agent", "workflows"]:
|
||||
subpath = EXAMPLES_DIR / subdir
|
||||
if subpath.exists():
|
||||
examples.extend(sorted(subpath.glob("*_wow.py")))
|
||||
|
||||
print(f"\nFound {len(examples)} examples to run\n")
|
||||
|
||||
for ex in examples:
|
||||
name, status, reason = run_example(ex)
|
||||
RESULTS.append((name, status, reason))
|
||||
icon = {"PASS": "✅", "SKIP": "⏭️", "FAIL": "❌"}.get(status, "?")
|
||||
reason_str = f" ({reason})" if reason else ""
|
||||
print(f"{icon} {status:4} | {name:45} {reason_str}")
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 70)
|
||||
passed = sum(1 for _, s, _ in RESULTS if s == "PASS")
|
||||
skipped = sum(1 for _, s, _ in RESULTS if s == "SKIP")
|
||||
failed = sum(1 for _, s, _ in RESULTS if s == "FAIL")
|
||||
print(f"SUMMARY: {passed} PASS | {skipped} SKIP | {failed} FAIL | {len(RESULTS)} TOTAL")
|
||||
print("=" * 70)
|
||||
|
||||
return 0 if failed == 0 else 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,94 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example: Creating Skills with PraisonAI CLI
|
||||
|
||||
This example demonstrates how to create Agent Skills using the CLI.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import tempfile
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def example_create_skill_with_template():
|
||||
"""Create a skill using template mode (no API key required)."""
|
||||
print("=" * 60)
|
||||
print("Example 1: Create skill with template")
|
||||
print("=" * 60)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
result = subprocess.run([
|
||||
"praisonai", "skills", "create",
|
||||
"--name", "my-template-skill",
|
||||
"--description", "A skill for processing text files",
|
||||
"--author", "example-team",
|
||||
"--license", "MIT",
|
||||
"--template",
|
||||
"--output-dir", tmpdir
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode == 0:
|
||||
skill_md = Path(tmpdir) / "my-template-skill" / "SKILL.md"
|
||||
print(f"\nGenerated SKILL.md:\n{skill_md.read_text()[:500]}...")
|
||||
else:
|
||||
print(f"Error: {result.stderr}")
|
||||
|
||||
|
||||
def example_create_skill_with_ai():
|
||||
"""Create a skill using AI generation (requires API key)."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Create skill with AI generation")
|
||||
print("=" * 60)
|
||||
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Skipping: OPENAI_API_KEY not set")
|
||||
return
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
result = subprocess.run([
|
||||
"praisonai", "skills", "create",
|
||||
"--name", "csv-analyzer",
|
||||
"--description", "Analyze CSV files and generate statistical reports",
|
||||
"--output-dir", tmpdir
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode == 0:
|
||||
skill_md = Path(tmpdir) / "csv-analyzer" / "SKILL.md"
|
||||
print(f"\nGenerated SKILL.md:\n{skill_md.read_text()[:800]}...")
|
||||
|
||||
# Check if script was generated
|
||||
script_path = Path(tmpdir) / "csv-analyzer" / "scripts" / "skill.py"
|
||||
if script_path.exists():
|
||||
print(f"\nGenerated skill.py:\n{script_path.read_text()[:500]}...")
|
||||
|
||||
|
||||
def example_create_skill_with_script():
|
||||
"""Create a skill with auto-generated Python script."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 3: Create skill with script")
|
||||
print("=" * 60)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
result = subprocess.run([
|
||||
"praisonai", "skills", "create",
|
||||
"--name", "data-processor",
|
||||
"--description", "Process and transform data files",
|
||||
"--template",
|
||||
"--script",
|
||||
"--output-dir", tmpdir
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode == 0:
|
||||
script_path = Path(tmpdir) / "data-processor" / "scripts" / "skill.py"
|
||||
if script_path.exists():
|
||||
print(f"\nGenerated skill.py:\n{script_path.read_text()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
example_create_skill_with_template()
|
||||
example_create_skill_with_ai()
|
||||
example_create_skill_with_script()
|
||||
@@ -0,0 +1,137 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example: Multi-Agent CLI Usage
|
||||
|
||||
This example demonstrates how to use the praisonai agents command
|
||||
to define and run multiple agents from the command line.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
|
||||
|
||||
def example_single_agent():
|
||||
"""Run a single agent with a simple task."""
|
||||
print("=" * 60)
|
||||
print("Example 1: Single Agent")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "assistant:Helper",
|
||||
"--task", "What is the capital of France?"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
if result.returncode != 0:
|
||||
print(f"Error: {result.stderr}")
|
||||
|
||||
|
||||
def example_agent_with_tools():
|
||||
"""Run an agent with tools."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Agent with Tools")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "researcher:Research Analyst:internet_search",
|
||||
"--task", "Find the latest news about renewable energy"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_multiple_agents_sequential():
|
||||
"""Run multiple agents in sequential mode."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 3: Multiple Agents (Sequential)")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "researcher:Research Analyst:internet_search",
|
||||
"--agent", "writer:Content Writer:write_file",
|
||||
"--task", "Research AI trends and write a brief summary"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_multiple_agents_parallel():
|
||||
"""Run multiple agents in parallel mode."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 4: Multiple Agents (Parallel)")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "analyst1:Market Analyst",
|
||||
"--agent", "analyst2:Tech Analyst",
|
||||
"--process", "parallel",
|
||||
"--task", "Analyze the current state of AI industry"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_with_custom_llm():
|
||||
"""Run agents with a custom LLM model."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 5: Custom LLM Model")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "coder:Developer:execute_command",
|
||||
"--llm", "gpt-4o-mini",
|
||||
"--task", "Write a simple hello world Python script"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_with_instructions():
|
||||
"""Run agents with additional instructions."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 6: With Additional Instructions")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "run",
|
||||
"--agent", "analyst:Data Analyst",
|
||||
"--instructions", "Be concise and provide only key insights",
|
||||
"--task", "Explain the benefits of using AI in healthcare"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
def example_list_tools():
|
||||
"""List available tools for agents."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 7: List Available Tools")
|
||||
print("=" * 60)
|
||||
|
||||
result = subprocess.run([
|
||||
"praisonai", "agents", "list"
|
||||
], capture_output=True, text=True)
|
||||
|
||||
print(result.stdout)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Check for API key
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
print("Warning: OPENAI_API_KEY not set. Some examples may fail.")
|
||||
print()
|
||||
|
||||
# Run examples
|
||||
example_list_tools()
|
||||
example_single_agent()
|
||||
# Uncomment to run more examples:
|
||||
# example_agent_with_tools()
|
||||
# example_multiple_agents_sequential()
|
||||
# example_multiple_agents_parallel()
|
||||
# example_with_custom_llm()
|
||||
# example_with_instructions()
|
||||
@@ -0,0 +1,137 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic Thinking Budgets Example for PraisonAI Agents.
|
||||
|
||||
This example demonstrates how to use thinking budgets to:
|
||||
1. Configure token budgets for extended thinking
|
||||
2. Use predefined budget levels
|
||||
3. Track thinking usage
|
||||
4. Adapt budgets based on complexity
|
||||
|
||||
Usage:
|
||||
python basic_thinking.py
|
||||
"""
|
||||
|
||||
from praisonaiagents.thinking import (
|
||||
ThinkingBudget, ThinkingConfig, ThinkingUsage, ThinkingTracker
|
||||
)
|
||||
from praisonaiagents.thinking.budget import BudgetLevel
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("Thinking Budgets Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# ==========================================================================
|
||||
# Budget Levels
|
||||
# ==========================================================================
|
||||
print("\n--- Budget Levels ---")
|
||||
|
||||
levels = [
|
||||
("Minimal", ThinkingBudget.minimal()),
|
||||
("Low", ThinkingBudget.low()),
|
||||
("Medium", ThinkingBudget.medium()),
|
||||
("High", ThinkingBudget.high()),
|
||||
("Maximum", ThinkingBudget.maximum()),
|
||||
]
|
||||
|
||||
for name, budget in levels:
|
||||
print(f" {name}: {budget.max_tokens:,} tokens")
|
||||
|
||||
# ==========================================================================
|
||||
# Custom Budget
|
||||
# ==========================================================================
|
||||
print("\n--- Custom Budget ---")
|
||||
|
||||
custom_budget = ThinkingBudget(
|
||||
max_tokens=12000,
|
||||
min_tokens=2000,
|
||||
adaptive=True,
|
||||
max_time_seconds=120.0
|
||||
)
|
||||
print(f" Max tokens: {custom_budget.max_tokens:,}")
|
||||
print(f" Min tokens: {custom_budget.min_tokens:,}")
|
||||
print(f" Adaptive: {custom_budget.adaptive}")
|
||||
print(f" Time limit: {custom_budget.max_time_seconds}s")
|
||||
|
||||
# ==========================================================================
|
||||
# Complexity-Based Scaling
|
||||
# ==========================================================================
|
||||
print("\n--- Complexity-Based Token Allocation ---")
|
||||
|
||||
budget = ThinkingBudget(
|
||||
max_tokens=16000,
|
||||
min_tokens=2000,
|
||||
adaptive=True
|
||||
)
|
||||
|
||||
complexities = [0.0, 0.25, 0.5, 0.75, 1.0]
|
||||
for complexity in complexities:
|
||||
tokens = budget.get_tokens_for_complexity(complexity)
|
||||
bar = "█" * int(tokens / 1000)
|
||||
print(f" Complexity {complexity:.2f}: {tokens:,} tokens {bar}")
|
||||
|
||||
# ==========================================================================
|
||||
# Usage Tracking
|
||||
# ==========================================================================
|
||||
print("\n--- Usage Tracking ---")
|
||||
|
||||
usage = ThinkingUsage(
|
||||
budget_tokens=8000,
|
||||
tokens_used=5000,
|
||||
time_seconds=30.0,
|
||||
budget_time=60.0
|
||||
)
|
||||
|
||||
print(f" Tokens used: {usage.tokens_used:,} / {usage.budget_tokens:,}")
|
||||
print(f" Tokens remaining: {usage.tokens_remaining:,}")
|
||||
print(f" Utilization: {usage.token_utilization:.1%}")
|
||||
print(f" Over budget: {usage.is_over_budget}")
|
||||
print(f" Time remaining: {usage.time_remaining}s")
|
||||
|
||||
# ==========================================================================
|
||||
# Session Tracking
|
||||
# ==========================================================================
|
||||
print("\n--- Session Tracking ---")
|
||||
|
||||
tracker = ThinkingTracker()
|
||||
|
||||
# Simulate multiple thinking sessions
|
||||
sessions = [
|
||||
(8000, 4000, 20.0),
|
||||
(8000, 6000, 35.0),
|
||||
(16000, 12000, 60.0),
|
||||
(8000, 9000, 45.0), # Over budget
|
||||
]
|
||||
|
||||
for budget_tokens, tokens_used, time_seconds in sessions:
|
||||
session = tracker.start_session(budget_tokens=budget_tokens)
|
||||
tracker.end_session(session, tokens_used=tokens_used, time_seconds=time_seconds)
|
||||
|
||||
summary = tracker.get_summary()
|
||||
print(f" Total sessions: {summary['session_count']}")
|
||||
print(f" Total tokens used: {summary['total_tokens_used']:,}")
|
||||
print(f" Average tokens/session: {summary['average_tokens_per_session']:,.0f}")
|
||||
print(f" Average utilization: {summary['average_utilization']:.1%}")
|
||||
print(f" Over-budget sessions: {summary['over_budget_count']}")
|
||||
|
||||
# ==========================================================================
|
||||
# Serialization
|
||||
# ==========================================================================
|
||||
print("\n--- Serialization ---")
|
||||
|
||||
budget = ThinkingBudget.from_level(BudgetLevel.HIGH)
|
||||
data = budget.to_dict()
|
||||
print(f" Serialized: level={data['level']}, max_tokens={data['max_tokens']}")
|
||||
|
||||
restored = ThinkingBudget.from_dict(data)
|
||||
print(f" Restored: {restored.max_tokens:,} tokens")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Demo Complete!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
"""ArXiv Tool Example.
|
||||
|
||||
This example demonstrates how to use the ArXiv tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import ArxivTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize ArXiv tool (no API key needed)
|
||||
arxiv = ArxivTool()
|
||||
|
||||
# Example 1: Search ArXiv
|
||||
print("=" * 60)
|
||||
print("Example 1: Search ArXiv")
|
||||
print("=" * 60)
|
||||
|
||||
results = arxiv.search("transformer neural networks", max_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} papers:")
|
||||
for paper in results:
|
||||
print(f"\n Title: {paper.get('title', 'N/A')[:60]}...")
|
||||
authors = paper.get('authors', [])
|
||||
if authors:
|
||||
print(f" Authors: {', '.join(authors[:2])}...")
|
||||
print(f" Published: {paper.get('published', 'N/A')}")
|
||||
print(f" PDF: {paper.get('pdf_url', 'N/A')}")
|
||||
|
||||
print("\n✅ ArXiv tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
"""Calculator Tool Example.
|
||||
|
||||
This example demonstrates how to use the Calculator tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import CalculatorTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Calculator tool
|
||||
calc = CalculatorTool()
|
||||
|
||||
# Example 1: Basic arithmetic
|
||||
print("=" * 60)
|
||||
print("Example 1: Basic Arithmetic")
|
||||
print("=" * 60)
|
||||
|
||||
expressions = [
|
||||
"2 + 2",
|
||||
"10 * 5",
|
||||
"100 / 4",
|
||||
"2 ** 10",
|
||||
"(10 + 5) * 2 - 3"
|
||||
]
|
||||
|
||||
for expr in expressions:
|
||||
result = calc.evaluate(expr)
|
||||
print(f" {expr} = {result.get('result', result)}")
|
||||
|
||||
# Example 2: Percentage calculation
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Percentage Calculation")
|
||||
print("=" * 60)
|
||||
|
||||
result = calc.evaluate("250 * 0.15")
|
||||
print(f" 15% of 250 = {result.get('result', result)}")
|
||||
|
||||
result = calc.evaluate("1000 * (1 + 0.05) ** 3")
|
||||
print(f" $1000 at 5% for 3 years = ${result.get('result', 0):.2f}")
|
||||
|
||||
print("\n✅ Calculator tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+75
@@ -0,0 +1,75 @@
|
||||
"""DuckDuckGo Search Tool Example.
|
||||
|
||||
This example demonstrates how to use the DuckDuckGo search tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import DuckDuckGoTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize DuckDuckGo tool (no API key needed)
|
||||
ddg = DuckDuckGoTool()
|
||||
|
||||
# Example 1: Web search
|
||||
print("=" * 60)
|
||||
print("Example 1: Web Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = ddg.search("Python programming best practices", max_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r['title']}")
|
||||
print(f" URL: {r['url']}")
|
||||
print(f" Snippet: {r['snippet'][:100]}...")
|
||||
print()
|
||||
|
||||
# Example 2: News search
|
||||
print("=" * 60)
|
||||
print("Example 2: News Search")
|
||||
print("=" * 60)
|
||||
|
||||
news = ddg.news("artificial intelligence", max_results=3)
|
||||
|
||||
if news and "error" in news[0]:
|
||||
print(f"Error: {news[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(news)} news articles:")
|
||||
for n in news:
|
||||
print(f" - {n['title']}")
|
||||
print(f" Source: {n.get('source', 'N/A')}")
|
||||
print(f" Date: {n.get('date', 'N/A')}")
|
||||
print()
|
||||
|
||||
# Example 3: Image search
|
||||
print("=" * 60)
|
||||
print("Example 3: Image Search")
|
||||
print("=" * 60)
|
||||
|
||||
images = ddg.images("sunset landscape", max_results=3)
|
||||
|
||||
if images and "error" in images[0]:
|
||||
print(f"Error: {images[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(images)} images:")
|
||||
for img in images:
|
||||
print(f" - {img['title']}")
|
||||
print(f" Image URL: {img['image_url'][:60]}...")
|
||||
print()
|
||||
|
||||
print("✅ DuckDuckGo tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Vendored
+48
@@ -0,0 +1,48 @@
|
||||
"""Exa Search Tool Example.
|
||||
|
||||
This example demonstrates how to use the Exa search tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export EXA_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import ExaTool
|
||||
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("EXA_API_KEY"):
|
||||
print("Error: EXA_API_KEY environment variable not set")
|
||||
print("Get your API key from https://exa.ai/")
|
||||
return
|
||||
|
||||
# Initialize Exa tool
|
||||
exa = ExaTool()
|
||||
|
||||
# Example 1: Neural search
|
||||
print("=" * 60)
|
||||
print("Example 1: Neural Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = exa.search("best practices for Python development", num_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r.get('title', 'N/A')[:50]}...")
|
||||
print(f" URL: {r.get('url', 'N/A')[:60]}...")
|
||||
print(f" Score: {r.get('score', 'N/A')}")
|
||||
print()
|
||||
|
||||
print("✅ Exa tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
"""Firecrawl Tool Example.
|
||||
|
||||
This example demonstrates how to use the Firecrawl tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export FIRECRAWL_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import FirecrawlTool
|
||||
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("FIRECRAWL_API_KEY"):
|
||||
print("Error: FIRECRAWL_API_KEY environment variable not set")
|
||||
print("Get your API key from https://firecrawl.dev/")
|
||||
return
|
||||
|
||||
# Initialize Firecrawl tool
|
||||
firecrawl = FirecrawlTool()
|
||||
|
||||
# Example 1: Scrape a page
|
||||
print("=" * 60)
|
||||
print("Example 1: Scrape a Page")
|
||||
print("=" * 60)
|
||||
|
||||
result = firecrawl.scrape("https://example.com")
|
||||
|
||||
if "error" in result:
|
||||
print(f"Error: {result['error']}")
|
||||
else:
|
||||
print(f"URL: {result.get('url', 'N/A')}")
|
||||
markdown = result.get('markdown', '')
|
||||
print(f"Content preview: {markdown[:300]}...")
|
||||
|
||||
print("\n✅ Firecrawl tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
"""Hacker News Tool Example.
|
||||
|
||||
This example demonstrates how to use the Hacker News tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import HackerNewsTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Hacker News tool
|
||||
hn = HackerNewsTool()
|
||||
|
||||
# Example 1: Get top stories
|
||||
print("=" * 60)
|
||||
print("Example 1: Top Stories")
|
||||
print("=" * 60)
|
||||
|
||||
stories = hn.get_top_stories(limit=5)
|
||||
|
||||
if stories and "error" in stories[0]:
|
||||
print(f"Error: {stories[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(stories)} top stories:")
|
||||
for story in stories:
|
||||
print(f" - {story.get('title', 'N/A')[:50]}...")
|
||||
print(f" Score: {story.get('score', 'N/A')} | Comments: {story.get('descendants', 'N/A')}")
|
||||
print(f" URL: {story.get('url', 'N/A')[:60]}...")
|
||||
print()
|
||||
|
||||
print("✅ Hacker News tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Vendored
+40
@@ -0,0 +1,40 @@
|
||||
"""Jina Reader Tool Example.
|
||||
|
||||
This example demonstrates how to use the Jina Reader tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: Works without API key (with rate limits)
|
||||
"""
|
||||
|
||||
from praisonai_tools import JinaTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Jina tool (works without API key)
|
||||
jina = JinaTool()
|
||||
|
||||
# Example 1: Read a URL
|
||||
print("=" * 60)
|
||||
print("Example 1: Read URL Content")
|
||||
print("=" * 60)
|
||||
|
||||
result = jina.read("https://example.com")
|
||||
|
||||
if "error" in result:
|
||||
print(f"Error: {result['error']}")
|
||||
else:
|
||||
print(f"URL: {result.get('url', 'N/A')}")
|
||||
print(f"Title: {result.get('title', 'N/A')}")
|
||||
content = result.get('content', '')
|
||||
print(f"Content preview: {content[:300]}...")
|
||||
|
||||
print("\n✅ Jina tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+56
@@ -0,0 +1,56 @@
|
||||
"""PostgreSQL Tool Example.
|
||||
|
||||
This example demonstrates how to use the PostgreSQL tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
Docker: docker run -d --name postgres -e POSTGRES_PASSWORD=praison123 -e POSTGRES_DB=praisonai -p 5432:5432 postgres:16
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
from praisonai_tools import PostgresTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize PostgreSQL tool
|
||||
pg = PostgresTool(
|
||||
host="localhost",
|
||||
port=5432,
|
||||
database="praisonai",
|
||||
user="postgres",
|
||||
password="praison123"
|
||||
)
|
||||
|
||||
# Example 1: List tables
|
||||
print("=" * 60)
|
||||
print("Example 1: List Tables")
|
||||
print("=" * 60)
|
||||
|
||||
tables = pg.list_tables()
|
||||
|
||||
if isinstance(tables, dict) and "error" in tables:
|
||||
print(f"Error: {tables['error']}")
|
||||
else:
|
||||
print(f"Found {len(tables)} tables:")
|
||||
for t in tables[:10]:
|
||||
print(f" - {t.get('table_name', t)}")
|
||||
|
||||
# Example 2: Query
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Execute Query")
|
||||
print("=" * 60)
|
||||
|
||||
result = pg.query("SELECT version()")
|
||||
|
||||
if isinstance(result, dict) and "error" in result:
|
||||
print(f"Error: {result['error']}")
|
||||
else:
|
||||
print(f"PostgreSQL version: {result}")
|
||||
|
||||
print("\n✅ PostgreSQL tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
"""Serper Tool Example.
|
||||
|
||||
This example demonstrates how to use the Serper (Google Search) tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export SERPER_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import SerperTool
|
||||
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("SERPER_API_KEY"):
|
||||
print("Error: SERPER_API_KEY environment variable not set")
|
||||
print("Get your API key from https://serper.dev/")
|
||||
return
|
||||
|
||||
# Initialize Serper tool
|
||||
serper = SerperTool()
|
||||
|
||||
# Example 1: Web search
|
||||
print("=" * 60)
|
||||
print("Example 1: Google Web Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = serper.search("Python programming tutorials", max_results=3)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r.get('title', 'N/A')[:50]}...")
|
||||
print(f" URL: {r.get('link', 'N/A')[:60]}...")
|
||||
print()
|
||||
|
||||
print("✅ Serper tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
"""Tavily Search Tool Example.
|
||||
|
||||
This example demonstrates how to use the Tavily search tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
export TAVILY_API_KEY=your_api_key
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonai_tools import TavilyTool
|
||||
|
||||
def main():
|
||||
# Check for API key
|
||||
if not os.getenv("TAVILY_API_KEY"):
|
||||
print("Error: TAVILY_API_KEY environment variable not set")
|
||||
print("Get your API key from https://tavily.com/")
|
||||
return
|
||||
|
||||
# Initialize Tavily tool
|
||||
tavily = TavilyTool()
|
||||
|
||||
# Example 1: Basic search
|
||||
print("=" * 60)
|
||||
print("Example 1: Basic Search")
|
||||
print("=" * 60)
|
||||
|
||||
results = tavily.search("What is quantum computing?", max_results=3)
|
||||
|
||||
if "error" in results:
|
||||
print(f"Error: {results['error']}")
|
||||
else:
|
||||
print(f"Query: {results.get('query')}")
|
||||
if "answer" in results:
|
||||
print(f"\nAI Answer: {results['answer'][:200]}...")
|
||||
print(f"\nResults ({len(results.get('results', []))}):")
|
||||
for r in results.get("results", []):
|
||||
print(f" - {r['title']}")
|
||||
print(f" URL: {r['url']}")
|
||||
print(f" Score: {r.get('score', 'N/A')}")
|
||||
|
||||
# Example 2: Search context for RAG
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Search Context (for RAG)")
|
||||
print("=" * 60)
|
||||
|
||||
context = tavily.search_context("Python best practices 2024")
|
||||
print(f"Context (first 300 chars): {context[:300]}...")
|
||||
|
||||
print("\n✅ Tavily tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
"""Wikipedia Tool Example.
|
||||
|
||||
This example demonstrates how to use the Wikipedia tool with PraisonAI agents.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
|
||||
Usage:
|
||||
python tool.py
|
||||
|
||||
Note: No API key required!
|
||||
"""
|
||||
|
||||
from praisonai_tools import WikipediaTool
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize Wikipedia tool (no API key needed)
|
||||
wiki = WikipediaTool()
|
||||
|
||||
# Example 1: Search Wikipedia
|
||||
print("=" * 60)
|
||||
print("Example 1: Search Wikipedia")
|
||||
print("=" * 60)
|
||||
|
||||
results = wiki.search("Python programming", max_results=5)
|
||||
|
||||
if results and "error" in results[0]:
|
||||
print(f"Error: {results[0]['error']}")
|
||||
else:
|
||||
print(f"Found {len(results)} results:")
|
||||
for r in results:
|
||||
print(f" - {r['title']}")
|
||||
|
||||
# Example 2: Get page summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Get Page Summary")
|
||||
print("=" * 60)
|
||||
|
||||
summary = wiki.summary("Artificial intelligence", sentences=3)
|
||||
|
||||
if "error" in summary:
|
||||
print(f"Error: {summary['error']}")
|
||||
else:
|
||||
print(f"Title: {summary['title']}")
|
||||
print(f"Summary: {summary['summary'][:300]}...")
|
||||
|
||||
# Example 3: Get full page
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 3: Get Full Page")
|
||||
print("=" * 60)
|
||||
|
||||
page = wiki.get_page("Machine learning")
|
||||
|
||||
if "error" in page:
|
||||
print(f"Error: {page['error']}")
|
||||
if "options" in page:
|
||||
print(f"Options: {page['options'][:5]}")
|
||||
else:
|
||||
print(f"Title: {page['title']}")
|
||||
print(f"URL: {page['url']}")
|
||||
print(f"Categories: {page['categories'][:3]}")
|
||||
print(f"Content preview: {page['content'][:200]}...")
|
||||
|
||||
print("\n✅ Wikipedia tool working correctly!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,23 @@
|
||||
"""ChromaDB Vector Store - Agent-First Example"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Agent-first approach: use knowledge parameter with ChromaDB
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "chroma",
|
||||
"path": "/tmp/chroma_demo"
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: ChromaDB with Agent")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonai.persistence import create_knowledge_store
|
||||
# store = create_knowledge_store("chroma", path="/tmp/chroma_demo")
|
||||
@@ -0,0 +1,25 @@
|
||||
"""LanceDB Vector Store - Agent-First Example"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Agent-first approach: use knowledge parameter with LanceDB
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "lancedb",
|
||||
"path": "/tmp/lancedb_test"
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: LanceDB with Agent")
|
||||
|
||||
# --- Advanced: Direct LanceDB Usage ---
|
||||
# import lancedb
|
||||
# db = lancedb.connect("/tmp/lancedb_test")
|
||||
# data = [{"id": "1", "text": "ML is AI", "vector": [0.1] * 128}]
|
||||
# table = db.create_table("demo", data)
|
||||
@@ -0,0 +1,23 @@
|
||||
"""LanceDB Vector Store - Agent-First Example"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Agent-first approach: use knowledge parameter with LanceDB
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "lancedb",
|
||||
"path": "/tmp/lancedb_demo"
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: LanceDB with Agent")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonai.persistence import create_knowledge_store
|
||||
# store = create_knowledge_store("lancedb", path="/tmp/lancedb_demo")
|
||||
@@ -0,0 +1,31 @@
|
||||
"""PGVector Vector Store - Agent-First Example
|
||||
|
||||
Docker: docker run -d --name pgvector -e POSTGRES_PASSWORD=postgres -p 5433:5432 pgvector/pgvector:pg16
|
||||
"""
|
||||
import os
|
||||
from praisonaiagents import Agent
|
||||
|
||||
url = os.getenv("PGVECTOR_URL", "postgresql://postgres:postgres@localhost:5433/postgres")
|
||||
|
||||
# Agent-first approach: use knowledge parameter with PGVector
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "pgvector",
|
||||
"url": url
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: PGVector with Agent")
|
||||
|
||||
# --- Advanced: Direct PGVector Usage ---
|
||||
# import psycopg2
|
||||
# conn = psycopg2.connect(url)
|
||||
# cur = conn.cursor()
|
||||
# cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
@@ -0,0 +1,29 @@
|
||||
"""PGVector Knowledge Store - Agent-First Example (requires Docker)
|
||||
|
||||
Docker: docker run -d --name pgvector -e POSTGRES_PASSWORD=postgres -p 5433:5432 pgvector/pgvector:pg16
|
||||
"""
|
||||
import os
|
||||
from praisonaiagents import Agent
|
||||
|
||||
url = os.getenv("PGVECTOR_URL", "postgresql://postgres:postgres@localhost:5433/postgres")
|
||||
|
||||
# Agent-first approach: use knowledge parameter with PGVector
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "pgvector",
|
||||
"url": url
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: PGVector with Agent")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonai.persistence import create_knowledge_store
|
||||
# store = create_knowledge_store("pgvector", url=url)
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Pinecone Vector Store - Agent-First Example
|
||||
|
||||
Requires: export PINECONE_API_KEY=...
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
from praisonaiagents import Agent
|
||||
|
||||
api_key = os.getenv("PINECONE_API_KEY")
|
||||
if not api_key:
|
||||
print("SKIPPED: Pinecone - PINECONE_API_KEY not set")
|
||||
sys.exit(0)
|
||||
|
||||
# Agent-first approach: use knowledge parameter with Pinecone
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "pinecone",
|
||||
"api_key": api_key,
|
||||
"environment": "us-east-1"
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: Pinecone with Agent")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from pinecone import Pinecone
|
||||
# pc = Pinecone(api_key=api_key)
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Qdrant Vector Store - Agent-First Example (requires Docker)
|
||||
|
||||
Docker: docker run -d --name qdrant -p 6333:6333 qdrant/qdrant
|
||||
"""
|
||||
import os
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Agent-first approach: use knowledge parameter with Qdrant
|
||||
url = os.getenv("QDRANT_URL", "http://localhost:6333")
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "qdrant",
|
||||
"url": url
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: Qdrant with Agent")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonai.persistence import create_knowledge_store
|
||||
# store = create_knowledge_store("qdrant", url=url)
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Weaviate Cloud Vector Store - Agent-First Example
|
||||
|
||||
Requires: export WEAVIATE_URL=... and WEAVIATE_API_KEY=...
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
from praisonaiagents import Agent
|
||||
|
||||
weaviate_url = os.getenv("WEAVIATE_URL", "https://your-cluster.weaviate.cloud")
|
||||
weaviate_key = os.getenv("WEAVIATE_API_KEY")
|
||||
if not weaviate_key:
|
||||
print("SKIPPED: Weaviate - WEAVIATE_API_KEY not set")
|
||||
sys.exit(0)
|
||||
|
||||
# Agent-first approach: use knowledge parameter with Weaviate
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant with access to documents.",
|
||||
knowledge=["./docs/guide.pdf"], # Add your documents here
|
||||
knowledge_config={
|
||||
"vector_store": "weaviate",
|
||||
"url": weaviate_url,
|
||||
"api_key": weaviate_key
|
||||
}
|
||||
)
|
||||
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What information do you have?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
print("PASSED: Weaviate with Agent")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# import weaviate
|
||||
# from weaviate.classes.init import Auth
|
||||
# client = weaviate.connect_to_weaviate_cloud(cluster_url=weaviate_url, auth_credentials=Auth.api_key(weaviate_key))
|
||||
@@ -62,6 +62,8 @@ docs/generated/
|
||||
*.log
|
||||
*.sqlite
|
||||
*.db
|
||||
# But keep the db module directory
|
||||
!praisonaiagents/db/
|
||||
|
||||
# OS generated files
|
||||
.DS_Store
|
||||
|
||||
@@ -43,8 +43,7 @@ from .tools.base import BaseTool, ToolResult, ToolValidationError, validate_tool
|
||||
from .tools.decorator import tool, FunctionTool
|
||||
from .tools.registry import get_registry, register_tool, get_tool, ToolRegistry
|
||||
from .agents.autoagents import AutoAgents
|
||||
from .knowledge.knowledge import Knowledge
|
||||
from .knowledge.chunking import Chunking
|
||||
# Knowledge is lazy-loaded to avoid importing heavy deps (chromadb, mem0) at startup
|
||||
# MCP support (optional)
|
||||
try:
|
||||
from .mcp.mcp import MCP
|
||||
@@ -53,7 +52,9 @@ except ImportError:
|
||||
_mcp_available = False
|
||||
MCP = None
|
||||
from .session import Session
|
||||
from .memory.memory import Memory
|
||||
# Memory is lazy-loaded to avoid importing chromadb at startup
|
||||
from .db import db
|
||||
from .obs import obs
|
||||
# Workflows - import from dedicated workflows module
|
||||
from .workflows import (
|
||||
Workflow, WorkflowStep, WorkflowContext, StepResult,
|
||||
@@ -65,8 +66,16 @@ from .guardrails import GuardrailResult, LLMGuardrail
|
||||
|
||||
# Fast Context support (lazy loaded to avoid performance impact)
|
||||
def __getattr__(name):
|
||||
"""Lazy load FastContext to avoid impacting package load time."""
|
||||
if name == "FastContext":
|
||||
"""Lazy load heavy modules to avoid impacting package load time."""
|
||||
# Knowledge module (imports chromadb, mem0)
|
||||
if name == "Knowledge":
|
||||
from praisonaiagents.knowledge.knowledge import Knowledge
|
||||
return Knowledge
|
||||
elif name == "Chunking":
|
||||
from praisonaiagents.knowledge.chunking import Chunking
|
||||
return Chunking
|
||||
# FastContext support
|
||||
elif name == "FastContext":
|
||||
from praisonaiagents.context.fast import FastContext
|
||||
return FastContext
|
||||
elif name == "FastContextResult":
|
||||
@@ -91,6 +100,10 @@ def __getattr__(name):
|
||||
elif name == "SkillLoader":
|
||||
from praisonaiagents.skills import SkillLoader
|
||||
return SkillLoader
|
||||
# Memory module (imports chromadb)
|
||||
elif name == "Memory":
|
||||
from praisonaiagents.memory.memory import Memory
|
||||
return Memory
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
# Planning mode support
|
||||
@@ -257,6 +270,8 @@ __all__ = [
|
||||
'AutoAgents',
|
||||
'Session',
|
||||
'Memory',
|
||||
'db',
|
||||
'obs',
|
||||
'display_interaction',
|
||||
'display_self_reflection',
|
||||
'display_instruction',
|
||||
|
||||
@@ -252,7 +252,12 @@ class Agent:
|
||||
history_in_context: Optional[int] = None,
|
||||
auto_save: Optional[str] = None,
|
||||
skills: Optional[List[str]] = None,
|
||||
skills_dirs: Optional[List[str]] = None
|
||||
skills_dirs: Optional[List[str]] = None,
|
||||
db: Optional[Any] = None,
|
||||
session_id: Optional[str] = None,
|
||||
hooks: Optional[List[Any]] = None,
|
||||
llm_config: Optional[Dict[str, Any]] = None,
|
||||
rate_limiter: Optional[Any] = None
|
||||
):
|
||||
"""Initialize an Agent instance.
|
||||
|
||||
@@ -393,6 +398,17 @@ class Agent:
|
||||
# Flag to track if final result has been displayed to prevent duplicates
|
||||
self._final_display_shown = False
|
||||
|
||||
# Store hooks for middleware system (zero overhead when empty)
|
||||
self._hooks = hooks or []
|
||||
self._middleware_manager = None # Lazy init
|
||||
|
||||
# Store llm_config for configurable model switching
|
||||
self._llm_config = llm_config or {}
|
||||
self._llm_configurable = self._llm_config.get('configurable', False)
|
||||
|
||||
# Store rate limiter (optional, zero overhead when None)
|
||||
self._rate_limiter = rate_limiter
|
||||
|
||||
# Store OpenAI client parameters for lazy initialization
|
||||
self._openai_api_key = api_key
|
||||
self._openai_base_url = base_url
|
||||
@@ -615,6 +631,11 @@ Your Goal: {self.goal}
|
||||
self._skill_manager = None # Lazy loaded
|
||||
self._skills_initialized = False
|
||||
|
||||
# Database persistence (lazy - no imports until used)
|
||||
self._db = db
|
||||
self._session_id = session_id
|
||||
self._db_initialized = False
|
||||
|
||||
@property
|
||||
def console(self):
|
||||
"""Lazily initialize Rich Console only when needed."""
|
||||
@@ -651,8 +672,45 @@ Your Goal: {self.goal}
|
||||
self._skill_manager.discover(self._skills_dirs, include_defaults=False)
|
||||
|
||||
self._skills_initialized = True
|
||||
|
||||
# Auto-add skill execution tools if not already present
|
||||
self._add_skill_tools()
|
||||
return self._skill_manager
|
||||
|
||||
def _add_skill_tools(self):
|
||||
"""Add tools required for skill execution (read_file, run_skill_script).
|
||||
|
||||
Uses lazy imports from praisonaiagents.tools to avoid performance impact
|
||||
when skills are not used.
|
||||
"""
|
||||
# Check if tools already include required capabilities
|
||||
tool_names = set()
|
||||
for tool in self.tools:
|
||||
if callable(tool) and hasattr(tool, '__name__'):
|
||||
tool_names.add(tool.__name__)
|
||||
elif hasattr(tool, 'name'):
|
||||
tool_names.add(tool.name)
|
||||
|
||||
# Add read_file if not present
|
||||
if 'read_file' not in tool_names:
|
||||
try:
|
||||
from ..tools import read_file
|
||||
self.tools.append(read_file)
|
||||
logging.debug("Added read_file tool for skill support")
|
||||
except ImportError:
|
||||
logging.warning("Could not import read_file tool for skills")
|
||||
|
||||
# Add run_skill_script from skill_tools module
|
||||
if 'run_skill_script' not in tool_names:
|
||||
try:
|
||||
from ..tools.skill_tools import create_skill_tools
|
||||
# Create skill tools with current working directory
|
||||
skill_tools = create_skill_tools()
|
||||
self.tools.append(skill_tools.run_skill_script)
|
||||
logging.debug("Added run_skill_script tool for skill support")
|
||||
except ImportError as e:
|
||||
logging.warning(f"Could not import skill_tools: {e}")
|
||||
|
||||
def get_skills_prompt(self) -> str:
|
||||
"""Get the XML prompt for available skills.
|
||||
|
||||
@@ -1321,6 +1379,13 @@ Your Goal: {self.goal}"""
|
||||
if self.verbose:
|
||||
self._display_memory_info()
|
||||
|
||||
# Add skills prompt if skills are configured
|
||||
if self._skills or self._skills_dirs:
|
||||
skills_prompt = self.get_skills_prompt()
|
||||
if skills_prompt:
|
||||
system_prompt += f"\n\n## Available Skills\n{skills_prompt}"
|
||||
system_prompt += "\n\nWhen a skill is relevant to the task, read its SKILL.md file to get detailed instructions. If the skill has scripts in its scripts/ directory, you can execute them using the execute_code or run_script tool."
|
||||
|
||||
# Add tool usage instructions if tools are available
|
||||
# Use provided tools or fall back to self.tools
|
||||
tools_to_use = tools if tools is not None else self.tools
|
||||
@@ -1571,8 +1636,32 @@ Your Goal: {self.goal}"""
|
||||
def execute_tool(self, function_name, arguments):
|
||||
"""
|
||||
Execute a tool dynamically based on the function name and arguments.
|
||||
Injects agent state for tools with Injected[T] parameters.
|
||||
"""
|
||||
logging.debug(f"{self.name} executing tool {function_name} with arguments: {arguments}")
|
||||
|
||||
# Set up injection context for tools with Injected parameters
|
||||
from ..tools.injected import AgentState, with_injection_context
|
||||
state = AgentState(
|
||||
agent_id=self.name,
|
||||
run_id=getattr(self, '_current_run_id', 'unknown'),
|
||||
session_id=getattr(self, '_session_id', None) or 'default',
|
||||
last_user_message=self.chat_history[-1].get('content') if self.chat_history else None,
|
||||
metadata={'agent_name': self.name}
|
||||
)
|
||||
|
||||
# Execute within injection context
|
||||
return self._execute_tool_with_context(function_name, arguments, state)
|
||||
|
||||
def _execute_tool_with_context(self, function_name, arguments, state):
|
||||
"""Execute tool within injection context."""
|
||||
from ..tools.injected import with_injection_context
|
||||
|
||||
with with_injection_context(state):
|
||||
return self._execute_tool_impl(function_name, arguments)
|
||||
|
||||
def _execute_tool_impl(self, function_name, arguments):
|
||||
"""Internal tool execution implementation."""
|
||||
|
||||
# Check if approval is required for this tool
|
||||
from ..approval import is_approval_required, console_approval_callback, get_risk_level, mark_approved, ApprovalDecision, get_approval_callback
|
||||
@@ -1913,7 +2002,112 @@ Your Goal: {self.goal}"""
|
||||
# expand=False
|
||||
# )
|
||||
|
||||
def chat(self, prompt, temperature=1.0, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=None, task_name=None, task_description=None, task_id=None):
|
||||
def _init_db_session(self):
|
||||
"""Initialize DB session if db adapter is provided (lazy, first chat only)."""
|
||||
if self._db is None or self._db_initialized:
|
||||
return
|
||||
|
||||
# Generate session_id if not provided: default to per-hour ID (YYYYMMDDHH-agentname)
|
||||
if self._session_id is None:
|
||||
import hashlib
|
||||
from datetime import datetime, timezone
|
||||
# Per-hour session ID: YYYYMMDDHH (UTC) + agent name hash for uniqueness
|
||||
hour_str = datetime.now(timezone.utc).strftime("%Y%m%d%H")
|
||||
agent_hash = hashlib.md5(self.name.encode()).hexdigest()[:6]
|
||||
self._session_id = f"{hour_str}-{agent_hash}"
|
||||
|
||||
# Call db adapter's on_agent_start to get previous messages
|
||||
try:
|
||||
history = self._db.on_agent_start(
|
||||
agent_name=self.name,
|
||||
session_id=self._session_id,
|
||||
user_id=self.user_id,
|
||||
metadata={"role": self.role, "goal": self.goal}
|
||||
)
|
||||
|
||||
# Restore chat history from previous session
|
||||
if history:
|
||||
for msg in history:
|
||||
self.chat_history.append({
|
||||
"role": msg.role,
|
||||
"content": msg.content
|
||||
})
|
||||
logging.info(f"Resumed session {self._session_id} with {len(history)} messages")
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to initialize DB session: {e}")
|
||||
|
||||
self._db_initialized = True
|
||||
self._current_run_id = None # Track current run
|
||||
|
||||
def _start_run(self, input_content: str):
|
||||
"""Start a new run (turn) for persistence tracking."""
|
||||
if self._db is None:
|
||||
return
|
||||
|
||||
import uuid
|
||||
self._current_run_id = f"run-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
try:
|
||||
if hasattr(self._db, 'on_run_start'):
|
||||
self._db.on_run_start(
|
||||
session_id=self._session_id,
|
||||
run_id=self._current_run_id,
|
||||
input_content=input_content,
|
||||
metadata={"agent_name": self.name}
|
||||
)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to start run: {e}")
|
||||
|
||||
def _end_run(self, output_content: str, status: str = "completed", metrics: dict = None):
|
||||
"""End the current run (turn)."""
|
||||
if self._db is None or self._current_run_id is None:
|
||||
return
|
||||
|
||||
try:
|
||||
if hasattr(self._db, 'on_run_end'):
|
||||
self._db.on_run_end(
|
||||
session_id=self._session_id,
|
||||
run_id=self._current_run_id,
|
||||
output_content=output_content,
|
||||
status=status,
|
||||
metrics=metrics or {},
|
||||
metadata={"agent_name": self.name}
|
||||
)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to end run: {e}")
|
||||
|
||||
self._current_run_id = None
|
||||
|
||||
def _persist_message(self, role: str, content: str):
|
||||
"""Persist a message to the DB if adapter is provided."""
|
||||
if self._db is None:
|
||||
return
|
||||
|
||||
try:
|
||||
if role == "user":
|
||||
self._db.on_user_message(self._session_id, content)
|
||||
elif role == "assistant":
|
||||
self._db.on_agent_message(self._session_id, content)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to persist message: {e}")
|
||||
|
||||
@property
|
||||
def session_id(self) -> Optional[str]:
|
||||
"""Get the current session ID."""
|
||||
return self._session_id
|
||||
|
||||
def chat(self, prompt, temperature=1.0, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=None, task_name=None, task_description=None, task_id=None, config=None):
|
||||
# Apply rate limiter if configured (before any LLM call)
|
||||
if self._rate_limiter is not None:
|
||||
self._rate_limiter.acquire()
|
||||
|
||||
# Initialize DB session on first chat (lazy)
|
||||
self._init_db_session()
|
||||
|
||||
# Start a new run for this chat turn
|
||||
prompt_str = prompt if isinstance(prompt, str) else str(prompt)
|
||||
self._start_run(prompt_str)
|
||||
|
||||
# Reset the final display flag for each new conversation
|
||||
self._final_display_shown = False
|
||||
|
||||
@@ -2010,6 +2204,8 @@ Your Goal: {self.goal}"""
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
# Persist user message to DB
|
||||
self._persist_message("user", normalized_content)
|
||||
|
||||
try:
|
||||
# Pass everything to LLM class
|
||||
@@ -2039,6 +2235,8 @@ Your Goal: {self.goal}"""
|
||||
)
|
||||
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
|
||||
# Log completion time if in debug mode
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
@@ -2083,6 +2281,8 @@ Your Goal: {self.goal}"""
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
# Persist user message to DB (OpenAI path)
|
||||
self._persist_message("user", normalized_content)
|
||||
|
||||
reflection_count = 0
|
||||
start_time = time.time()
|
||||
@@ -2122,6 +2322,8 @@ Your Goal: {self.goal}"""
|
||||
# Add to chat history and return raw response
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
# Apply guardrail validation even for JSON output
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
@@ -2137,6 +2339,8 @@ Your Goal: {self.goal}"""
|
||||
if not self.self_reflect:
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB (non-reflect path)
|
||||
self._persist_message("assistant", response_text)
|
||||
if self.verbose:
|
||||
logging.debug(f"Agent {self.name} final response: {response_text}")
|
||||
# Return only reasoning content if reasoning_steps is True
|
||||
@@ -2225,11 +2429,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
self._end_run(validated_response, "completed", {"duration_ms": (time.time() - start_time) * 1000})
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
self._end_run(None, "error", {"error": str(e)})
|
||||
return None
|
||||
|
||||
# Check if we've hit max reflections
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Background Agents Module for PraisonAI Agents.
|
||||
|
||||
Provides the ability to run agents in the background, allowing:
|
||||
- Long-running tasks without blocking
|
||||
- Task queuing and management
|
||||
- Progress monitoring and notifications
|
||||
- Graceful cancellation
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Background processing only when explicitly started
|
||||
- No overhead when not in use
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.background import BackgroundRunner, BackgroundTask
|
||||
|
||||
# Create a background runner
|
||||
runner = BackgroundRunner()
|
||||
|
||||
# Submit a task
|
||||
task = runner.submit(
|
||||
agent=my_agent,
|
||||
prompt="Research AI trends",
|
||||
callback=on_complete
|
||||
)
|
||||
|
||||
# Check status
|
||||
print(task.status) # "running", "completed", "failed"
|
||||
|
||||
# Wait for completion
|
||||
result = await task.wait()
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core classes
|
||||
"BackgroundRunner",
|
||||
"BackgroundTask",
|
||||
# Status enum
|
||||
"TaskStatus",
|
||||
# Configuration
|
||||
"BackgroundConfig",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "BackgroundRunner":
|
||||
from .runner import BackgroundRunner
|
||||
return BackgroundRunner
|
||||
|
||||
if name == "BackgroundTask":
|
||||
from .task import BackgroundTask
|
||||
return BackgroundTask
|
||||
|
||||
if name == "TaskStatus":
|
||||
from .task import TaskStatus
|
||||
return TaskStatus
|
||||
|
||||
if name == "BackgroundConfig":
|
||||
from .config import BackgroundConfig
|
||||
return BackgroundConfig
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
Background Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class BackgroundConfig:
|
||||
"""Configuration for background task execution."""
|
||||
max_concurrent_tasks: int = 5
|
||||
default_timeout: Optional[float] = None # None = no timeout
|
||||
auto_cleanup: bool = True # Auto-remove completed tasks
|
||||
cleanup_delay: float = 300.0 # Seconds to keep completed tasks
|
||||
enable_notifications: bool = True
|
||||
log_progress: bool = True
|
||||
@@ -0,0 +1,314 @@
|
||||
"""
|
||||
Background Runner for PraisonAI Agents.
|
||||
|
||||
Manages background task execution.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any, Callable, Union
|
||||
from datetime import datetime
|
||||
|
||||
from .task import BackgroundTask, TaskStatus
|
||||
from .config import BackgroundConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BackgroundRunner:
|
||||
"""
|
||||
Manages background task execution.
|
||||
|
||||
Provides:
|
||||
- Task submission and queuing
|
||||
- Concurrent execution management
|
||||
- Progress monitoring
|
||||
- Task cancellation
|
||||
|
||||
Example:
|
||||
runner = BackgroundRunner()
|
||||
|
||||
# Submit a task
|
||||
task = await runner.submit(
|
||||
func=my_agent.start,
|
||||
args=("Research AI trends",),
|
||||
name="research_task"
|
||||
)
|
||||
|
||||
# Check status
|
||||
print(task.status)
|
||||
|
||||
# Wait for result
|
||||
result = await task.wait()
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[BackgroundConfig] = None):
|
||||
"""
|
||||
Initialize the background runner.
|
||||
|
||||
Args:
|
||||
config: Configuration for background execution
|
||||
"""
|
||||
self.config = config or BackgroundConfig()
|
||||
self._tasks: Dict[str, BackgroundTask] = {}
|
||||
self._semaphore = asyncio.Semaphore(self.config.max_concurrent_tasks)
|
||||
self._running = False
|
||||
self._cleanup_task: Optional[asyncio.Task] = None
|
||||
|
||||
@property
|
||||
def tasks(self) -> List[BackgroundTask]:
|
||||
"""Get all tasks."""
|
||||
return list(self._tasks.values())
|
||||
|
||||
@property
|
||||
def running_tasks(self) -> List[BackgroundTask]:
|
||||
"""Get currently running tasks."""
|
||||
return [t for t in self._tasks.values() if t.is_running]
|
||||
|
||||
@property
|
||||
def pending_tasks(self) -> List[BackgroundTask]:
|
||||
"""Get pending tasks."""
|
||||
return [t for t in self._tasks.values() if t.status == TaskStatus.PENDING]
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
func: Callable,
|
||||
args: tuple = (),
|
||||
kwargs: Optional[Dict[str, Any]] = None,
|
||||
name: Optional[str] = None,
|
||||
timeout: Optional[float] = None,
|
||||
on_complete: Optional[Callable[[BackgroundTask], None]] = None,
|
||||
on_progress: Optional[Callable[[BackgroundTask], None]] = None
|
||||
) -> BackgroundTask:
|
||||
"""
|
||||
Submit a task for background execution.
|
||||
|
||||
Args:
|
||||
func: Function to execute (can be sync or async)
|
||||
args: Positional arguments for the function
|
||||
kwargs: Keyword arguments for the function
|
||||
name: Optional task name
|
||||
timeout: Optional timeout in seconds
|
||||
on_complete: Callback when task completes
|
||||
on_progress: Callback for progress updates
|
||||
|
||||
Returns:
|
||||
BackgroundTask object for tracking
|
||||
"""
|
||||
kwargs = kwargs or {}
|
||||
|
||||
task = BackgroundTask(
|
||||
name=name or func.__name__,
|
||||
metadata={
|
||||
"func_name": func.__name__,
|
||||
"timeout": timeout or self.config.default_timeout
|
||||
}
|
||||
)
|
||||
|
||||
self._tasks[task.id] = task
|
||||
|
||||
# Create the execution coroutine
|
||||
async def execute():
|
||||
async with self._semaphore:
|
||||
task.start()
|
||||
logger.info(f"Background task started: {task.name} ({task.id})")
|
||||
|
||||
try:
|
||||
# Check if function is async
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
result = await asyncio.wait_for(
|
||||
func(*args, **kwargs),
|
||||
timeout=timeout or self.config.default_timeout
|
||||
)
|
||||
else:
|
||||
# Run sync function in thread pool
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await asyncio.wait_for(
|
||||
loop.run_in_executor(None, lambda: func(*args, **kwargs)),
|
||||
timeout=timeout or self.config.default_timeout
|
||||
)
|
||||
|
||||
task.complete(result)
|
||||
logger.info(f"Background task completed: {task.name} ({task.id})")
|
||||
|
||||
if on_complete:
|
||||
try:
|
||||
on_complete(task)
|
||||
except Exception as e:
|
||||
logger.warning(f"on_complete callback error: {e}")
|
||||
|
||||
return result
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
task.fail(f"Task timed out after {timeout}s")
|
||||
logger.warning(f"Background task timed out: {task.name}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
task.cancel()
|
||||
logger.info(f"Background task cancelled: {task.name}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
task.fail(str(e))
|
||||
logger.error(f"Background task failed: {task.name} - {e}")
|
||||
|
||||
if on_complete:
|
||||
try:
|
||||
on_complete(task)
|
||||
except Exception as cb_error:
|
||||
logger.warning(f"on_complete callback error: {cb_error}")
|
||||
|
||||
# Create and store the future
|
||||
task._future = asyncio.create_task(execute())
|
||||
|
||||
return task
|
||||
|
||||
async def submit_agent(
|
||||
self,
|
||||
agent: Any,
|
||||
prompt: str,
|
||||
name: Optional[str] = None,
|
||||
timeout: Optional[float] = None,
|
||||
on_complete: Optional[Callable[[BackgroundTask], None]] = None
|
||||
) -> BackgroundTask:
|
||||
"""
|
||||
Submit an agent task for background execution.
|
||||
|
||||
Args:
|
||||
agent: Agent instance to run
|
||||
prompt: Prompt to send to the agent
|
||||
name: Optional task name
|
||||
timeout: Optional timeout in seconds
|
||||
on_complete: Callback when task completes
|
||||
|
||||
Returns:
|
||||
BackgroundTask object for tracking
|
||||
"""
|
||||
task_name = name or f"agent_{getattr(agent, 'name', 'unknown')}"
|
||||
|
||||
# Determine the method to call
|
||||
if hasattr(agent, 'start'):
|
||||
func = agent.start
|
||||
elif hasattr(agent, 'chat'):
|
||||
func = agent.chat
|
||||
elif hasattr(agent, 'run'):
|
||||
func = agent.run
|
||||
else:
|
||||
raise ValueError("Agent must have start(), chat(), or run() method")
|
||||
|
||||
return await self.submit(
|
||||
func=func,
|
||||
args=(prompt,),
|
||||
name=task_name,
|
||||
timeout=timeout,
|
||||
on_complete=on_complete
|
||||
)
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[BackgroundTask]:
|
||||
"""Get a task by ID."""
|
||||
return self._tasks.get(task_id)
|
||||
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
"""
|
||||
Cancel a task.
|
||||
|
||||
Args:
|
||||
task_id: ID of task to cancel
|
||||
|
||||
Returns:
|
||||
True if cancelled, False if not found or already completed
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if task is None:
|
||||
return False
|
||||
|
||||
if task.is_completed:
|
||||
return False
|
||||
|
||||
task.cancel()
|
||||
|
||||
if task._future and not task._future.done():
|
||||
task._future.cancel()
|
||||
|
||||
return True
|
||||
|
||||
async def cancel_all(self):
|
||||
"""Cancel all running tasks."""
|
||||
for task in self.running_tasks:
|
||||
await self.cancel_task(task.id)
|
||||
|
||||
async def wait_all(self, timeout: Optional[float] = None) -> List[BackgroundTask]:
|
||||
"""
|
||||
Wait for all tasks to complete.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait
|
||||
|
||||
Returns:
|
||||
List of completed tasks
|
||||
"""
|
||||
futures = [t._future for t in self._tasks.values() if t._future]
|
||||
|
||||
if not futures:
|
||||
return list(self._tasks.values())
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(*futures, return_exceptions=True),
|
||||
timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
return list(self._tasks.values())
|
||||
|
||||
def clear_completed(self):
|
||||
"""Remove completed tasks from tracking."""
|
||||
completed_ids = [
|
||||
task_id for task_id, task in self._tasks.items()
|
||||
if task.is_completed
|
||||
]
|
||||
for task_id in completed_ids:
|
||||
del self._tasks[task_id]
|
||||
|
||||
def list_tasks(self, status: Optional[TaskStatus] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List tasks with optional status filter.
|
||||
|
||||
Args:
|
||||
status: Optional status to filter by
|
||||
|
||||
Returns:
|
||||
List of task dictionaries
|
||||
"""
|
||||
tasks = self._tasks.values()
|
||||
|
||||
if status:
|
||||
tasks = [t for t in tasks if t.status == status]
|
||||
|
||||
return [t.to_dict() for t in tasks]
|
||||
|
||||
async def start_cleanup_loop(self):
|
||||
"""Start the automatic cleanup loop."""
|
||||
if self._cleanup_task is not None:
|
||||
return
|
||||
|
||||
async def cleanup_loop():
|
||||
while True:
|
||||
await asyncio.sleep(self.config.cleanup_delay)
|
||||
if self.config.auto_cleanup:
|
||||
self.clear_completed()
|
||||
|
||||
self._cleanup_task = asyncio.create_task(cleanup_loop())
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the runner and cancel all tasks."""
|
||||
await self.cancel_all()
|
||||
|
||||
if self._cleanup_task:
|
||||
self._cleanup_task.cancel()
|
||||
try:
|
||||
await self._cleanup_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._cleanup_task = None
|
||||
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
Background Task for PraisonAI Agents.
|
||||
|
||||
Represents a task running in the background.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
import asyncio
|
||||
import logging
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Any, Callable, Dict
|
||||
from datetime import datetime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TaskStatus(str, Enum):
|
||||
"""Status of a background task."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BackgroundTask:
|
||||
"""
|
||||
Represents a background task.
|
||||
|
||||
Tracks the status, progress, and result of a task running
|
||||
in the background.
|
||||
"""
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
|
||||
name: str = ""
|
||||
status: TaskStatus = TaskStatus.PENDING
|
||||
progress: float = 0.0 # 0.0 to 1.0
|
||||
result: Any = None
|
||||
error: Optional[str] = None
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
started_at: Optional[datetime] = None
|
||||
completed_at: Optional[datetime] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# Internal
|
||||
_future: Optional[asyncio.Future] = field(default=None, repr=False)
|
||||
_cancel_event: Optional[asyncio.Event] = field(default=None, repr=False)
|
||||
|
||||
def __post_init__(self):
|
||||
self._cancel_event = asyncio.Event()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if task is currently running."""
|
||||
return self.status == TaskStatus.RUNNING
|
||||
|
||||
@property
|
||||
def is_completed(self) -> bool:
|
||||
"""Check if task has completed (success or failure)."""
|
||||
return self.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED)
|
||||
|
||||
@property
|
||||
def is_successful(self) -> bool:
|
||||
"""Check if task completed successfully."""
|
||||
return self.status == TaskStatus.COMPLETED
|
||||
|
||||
@property
|
||||
def duration_seconds(self) -> Optional[float]:
|
||||
"""Get task duration in seconds."""
|
||||
if self.started_at is None:
|
||||
return None
|
||||
end_time = self.completed_at or datetime.now()
|
||||
return (end_time - self.started_at).total_seconds()
|
||||
|
||||
def start(self):
|
||||
"""Mark task as started."""
|
||||
self.status = TaskStatus.RUNNING
|
||||
self.started_at = datetime.now()
|
||||
|
||||
def complete(self, result: Any = None):
|
||||
"""Mark task as completed."""
|
||||
self.status = TaskStatus.COMPLETED
|
||||
self.result = result
|
||||
self.progress = 1.0
|
||||
self.completed_at = datetime.now()
|
||||
|
||||
def fail(self, error: str):
|
||||
"""Mark task as failed."""
|
||||
self.status = TaskStatus.FAILED
|
||||
self.error = error
|
||||
self.completed_at = datetime.now()
|
||||
|
||||
def cancel(self):
|
||||
"""Request task cancellation."""
|
||||
if self._cancel_event:
|
||||
self._cancel_event.set()
|
||||
self.status = TaskStatus.CANCELLED
|
||||
self.completed_at = datetime.now()
|
||||
|
||||
def update_progress(self, progress: float, message: Optional[str] = None):
|
||||
"""Update task progress."""
|
||||
self.progress = min(max(progress, 0.0), 1.0)
|
||||
if message:
|
||||
self.metadata["last_message"] = message
|
||||
|
||||
async def wait(self, timeout: Optional[float] = None) -> Any:
|
||||
"""
|
||||
Wait for task completion.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait in seconds
|
||||
|
||||
Returns:
|
||||
Task result
|
||||
|
||||
Raises:
|
||||
asyncio.TimeoutError: If timeout exceeded
|
||||
RuntimeError: If task failed
|
||||
"""
|
||||
if self._future is None:
|
||||
if self.is_completed:
|
||||
if self.status == TaskStatus.FAILED:
|
||||
raise RuntimeError(self.error)
|
||||
return self.result
|
||||
raise RuntimeError("Task has no associated future")
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(self._future, timeout=timeout)
|
||||
return result
|
||||
except asyncio.TimeoutError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(str(e))
|
||||
|
||||
def should_cancel(self) -> bool:
|
||||
"""Check if cancellation was requested."""
|
||||
return self._cancel_event is not None and self._cancel_event.is_set()
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"name": self.name,
|
||||
"status": self.status.value,
|
||||
"progress": self.progress,
|
||||
"result": str(self.result) if self.result else None,
|
||||
"error": self.error,
|
||||
"created_at": self.created_at.isoformat(),
|
||||
"started_at": self.started_at.isoformat() if self.started_at else None,
|
||||
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
|
||||
"duration_seconds": self.duration_seconds,
|
||||
"metadata": self.metadata
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
"""
|
||||
Checkpoints Module for PraisonAI Agents.
|
||||
|
||||
Provides file-level checkpointing using a shadow git repository to track
|
||||
and restore file changes made by agents. This enables:
|
||||
|
||||
- Automatic checkpointing before file modifications
|
||||
- Rewind to any previous checkpoint
|
||||
- Diff between checkpoints
|
||||
- Restore files and conversation state together
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Checkpoints only created when enabled
|
||||
- No overhead when checkpoints are disabled
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.checkpoints import CheckpointService
|
||||
|
||||
# Create checkpoint service
|
||||
service = CheckpointService(
|
||||
workspace_dir="/path/to/project",
|
||||
storage_dir="~/.praison/checkpoints"
|
||||
)
|
||||
|
||||
# Initialize shadow git
|
||||
await service.initialize()
|
||||
|
||||
# Save a checkpoint
|
||||
checkpoint_id = await service.save("Before refactoring")
|
||||
|
||||
# Restore to a checkpoint
|
||||
await service.restore(checkpoint_id)
|
||||
|
||||
# Get diff between checkpoints
|
||||
diff = await service.diff(from_id, to_id)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core service
|
||||
"CheckpointService",
|
||||
# Data types
|
||||
"Checkpoint",
|
||||
"CheckpointDiff",
|
||||
"CheckpointConfig",
|
||||
# Events
|
||||
"CheckpointEvent",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "CheckpointService":
|
||||
from .service import CheckpointService
|
||||
return CheckpointService
|
||||
|
||||
if name in ("Checkpoint", "CheckpointDiff", "CheckpointConfig"):
|
||||
from .types import Checkpoint, CheckpointDiff, CheckpointConfig
|
||||
return locals()[name]
|
||||
|
||||
if name == "CheckpointEvent":
|
||||
from .types import CheckpointEvent
|
||||
return CheckpointEvent
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,497 @@
|
||||
"""
|
||||
Checkpoint Service for PraisonAI Agents.
|
||||
|
||||
Implements shadow git repository for file-level checkpointing.
|
||||
"""
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
import logging
|
||||
import shutil
|
||||
from typing import Optional, List, Dict, Any, Callable
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from .types import (
|
||||
CheckpointConfig, Checkpoint, CheckpointDiff, FileDiff,
|
||||
CheckpointResult, CheckpointEvent
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Protected paths that should never be checkpointed
|
||||
PROTECTED_PATHS = [
|
||||
os.path.expanduser("~"),
|
||||
os.path.expanduser("~/Desktop"),
|
||||
os.path.expanduser("~/Documents"),
|
||||
os.path.expanduser("~/Downloads"),
|
||||
"/",
|
||||
"/tmp",
|
||||
"/var",
|
||||
"/etc",
|
||||
]
|
||||
|
||||
|
||||
class CheckpointService:
|
||||
"""
|
||||
Shadow Git Checkpoint Service.
|
||||
|
||||
Uses a separate git repository to track file changes in the workspace,
|
||||
enabling checkpoint creation, restoration, and diffing without
|
||||
interfering with the workspace's own git repository.
|
||||
|
||||
Example:
|
||||
service = CheckpointService(
|
||||
workspace_dir="/path/to/project",
|
||||
storage_dir="~/.praison/checkpoints"
|
||||
)
|
||||
|
||||
await service.initialize()
|
||||
|
||||
# Save checkpoint before making changes
|
||||
result = await service.save("Before refactoring")
|
||||
|
||||
# Make changes...
|
||||
|
||||
# Restore if needed
|
||||
await service.restore(result.checkpoint.id)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace_dir: str,
|
||||
storage_dir: Optional[str] = None,
|
||||
enabled: bool = True,
|
||||
auto_checkpoint: bool = True,
|
||||
max_checkpoints: int = 100
|
||||
):
|
||||
"""
|
||||
Initialize the checkpoint service.
|
||||
|
||||
Args:
|
||||
workspace_dir: Directory to track
|
||||
storage_dir: Where to store checkpoint data
|
||||
enabled: Whether checkpoints are enabled
|
||||
auto_checkpoint: Auto-checkpoint before file modifications
|
||||
max_checkpoints: Maximum checkpoints to keep
|
||||
"""
|
||||
self.config = CheckpointConfig(
|
||||
workspace_dir=workspace_dir,
|
||||
storage_dir=storage_dir,
|
||||
enabled=enabled,
|
||||
auto_checkpoint=auto_checkpoint,
|
||||
max_checkpoints=max_checkpoints
|
||||
)
|
||||
|
||||
self._initialized = False
|
||||
self._checkpoints: List[Checkpoint] = []
|
||||
self._event_handlers: Dict[CheckpointEvent, List[Callable]] = {
|
||||
event: [] for event in CheckpointEvent
|
||||
}
|
||||
|
||||
@property
|
||||
def workspace_dir(self) -> str:
|
||||
"""Get the workspace directory."""
|
||||
return self.config.workspace_dir
|
||||
|
||||
@property
|
||||
def checkpoint_dir(self) -> str:
|
||||
"""Get the checkpoint storage directory."""
|
||||
return self.config.get_checkpoint_dir()
|
||||
|
||||
@property
|
||||
def git_dir(self) -> str:
|
||||
"""Get the shadow git directory."""
|
||||
return os.path.join(self.checkpoint_dir, ".git")
|
||||
|
||||
@property
|
||||
def is_initialized(self) -> bool:
|
||||
"""Check if the service is initialized."""
|
||||
return self._initialized
|
||||
|
||||
def on(self, event: CheckpointEvent, handler: Callable):
|
||||
"""Register an event handler."""
|
||||
self._event_handlers[event].append(handler)
|
||||
|
||||
def _emit(self, event: CheckpointEvent, data: Any = None):
|
||||
"""Emit an event to all handlers."""
|
||||
for handler in self._event_handlers[event]:
|
||||
try:
|
||||
handler(data)
|
||||
except Exception as e:
|
||||
logger.warning(f"Event handler error: {e}")
|
||||
|
||||
async def initialize(self) -> bool:
|
||||
"""
|
||||
Initialize the shadow git repository.
|
||||
|
||||
Returns:
|
||||
True if initialization succeeded
|
||||
"""
|
||||
if not self.config.enabled:
|
||||
return False
|
||||
|
||||
# Check for protected paths
|
||||
workspace_real = os.path.realpath(self.config.workspace_dir)
|
||||
for protected in PROTECTED_PATHS:
|
||||
if workspace_real == os.path.realpath(protected):
|
||||
logger.warning(f"Cannot checkpoint protected path: {workspace_real}")
|
||||
return False
|
||||
|
||||
# Check if workspace exists
|
||||
if not os.path.isdir(self.config.workspace_dir):
|
||||
logger.error(f"Workspace does not exist: {self.config.workspace_dir}")
|
||||
return False
|
||||
|
||||
# Create checkpoint directory
|
||||
os.makedirs(self.checkpoint_dir, exist_ok=True)
|
||||
|
||||
# Initialize shadow git if needed
|
||||
if not os.path.exists(self.git_dir):
|
||||
success = await self._init_shadow_git()
|
||||
if not success:
|
||||
return False
|
||||
|
||||
# Load existing checkpoints
|
||||
await self._load_checkpoints()
|
||||
|
||||
self._initialized = True
|
||||
self._emit(CheckpointEvent.INITIALIZED, {"workspace": self.config.workspace_dir})
|
||||
|
||||
return True
|
||||
|
||||
async def _init_shadow_git(self) -> bool:
|
||||
"""Initialize the shadow git repository."""
|
||||
try:
|
||||
# Initialize git repo
|
||||
await self._run_git("init")
|
||||
|
||||
# Configure git
|
||||
await self._run_git("config", "user.name", "PraisonAI Checkpoints")
|
||||
await self._run_git("config", "user.email", "[email protected]")
|
||||
|
||||
# Set worktree to workspace
|
||||
await self._run_git("config", "core.worktree", self.config.workspace_dir)
|
||||
|
||||
# Create exclude file
|
||||
exclude_dir = os.path.join(self.git_dir, "info")
|
||||
os.makedirs(exclude_dir, exist_ok=True)
|
||||
|
||||
exclude_file = os.path.join(exclude_dir, "exclude")
|
||||
with open(exclude_file, "w") as f:
|
||||
for pattern in self.config.exclude_patterns:
|
||||
f.write(f"{pattern}\n")
|
||||
|
||||
# Also exclude the checkpoint directory itself
|
||||
f.write(f"{self.checkpoint_dir}\n")
|
||||
|
||||
# Create initial empty commit
|
||||
await self._run_git("commit", "--allow-empty", "-m", "Initial checkpoint")
|
||||
|
||||
logger.info(f"Initialized shadow git at {self.checkpoint_dir}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize shadow git: {e}")
|
||||
return False
|
||||
|
||||
async def _run_git(self, *args: str) -> str:
|
||||
"""Run a git command in the shadow repository."""
|
||||
env = self._get_sanitized_env()
|
||||
|
||||
cmd = ["git", f"--git-dir={self.git_dir}", f"--work-tree={self.config.workspace_dir}"]
|
||||
cmd.extend(args)
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
cwd=self.config.workspace_dir
|
||||
)
|
||||
|
||||
stdout, stderr = await process.communicate()
|
||||
|
||||
if process.returncode != 0:
|
||||
error_msg = stderr.decode().strip()
|
||||
raise RuntimeError(f"Git command failed: {' '.join(args)}\n{error_msg}")
|
||||
|
||||
return stdout.decode().strip()
|
||||
|
||||
def _get_sanitized_env(self) -> Dict[str, str]:
|
||||
"""Get environment with git variables sanitized."""
|
||||
env = os.environ.copy()
|
||||
|
||||
# Remove git environment variables that could interfere
|
||||
git_vars = [
|
||||
"GIT_DIR", "GIT_WORK_TREE", "GIT_INDEX_FILE",
|
||||
"GIT_OBJECT_DIRECTORY", "GIT_ALTERNATE_OBJECT_DIRECTORIES"
|
||||
]
|
||||
for var in git_vars:
|
||||
env.pop(var, None)
|
||||
|
||||
return env
|
||||
|
||||
async def save(self, message: str, allow_empty: bool = False) -> CheckpointResult:
|
||||
"""
|
||||
Save a checkpoint.
|
||||
|
||||
Args:
|
||||
message: Checkpoint message
|
||||
allow_empty: Allow checkpoint even if no changes
|
||||
|
||||
Returns:
|
||||
CheckpointResult with the created checkpoint
|
||||
"""
|
||||
if not self._initialized:
|
||||
return CheckpointResult.fail("Service not initialized")
|
||||
|
||||
try:
|
||||
# Stage all changes
|
||||
await self._run_git("add", "-A")
|
||||
|
||||
# Check if there are changes
|
||||
try:
|
||||
await self._run_git("diff", "--cached", "--quiet")
|
||||
has_changes = False
|
||||
except RuntimeError:
|
||||
has_changes = True
|
||||
|
||||
if not has_changes and not allow_empty:
|
||||
return CheckpointResult.fail("No changes to checkpoint")
|
||||
|
||||
# Create commit
|
||||
commit_args = ["commit", "-m", message]
|
||||
if allow_empty:
|
||||
commit_args.append("--allow-empty")
|
||||
|
||||
await self._run_git(*commit_args)
|
||||
|
||||
# Get commit info
|
||||
commit_hash = await self._run_git("rev-parse", "HEAD")
|
||||
timestamp = await self._run_git("log", "-1", "--format=%aI")
|
||||
|
||||
checkpoint = Checkpoint(
|
||||
id=commit_hash,
|
||||
short_id=commit_hash[:8],
|
||||
message=message,
|
||||
timestamp=datetime.fromisoformat(timestamp)
|
||||
)
|
||||
|
||||
self._checkpoints.append(checkpoint)
|
||||
|
||||
# Prune old checkpoints if needed
|
||||
await self._prune_checkpoints()
|
||||
|
||||
self._emit(CheckpointEvent.CHECKPOINT_CREATED, checkpoint)
|
||||
logger.info(f"Created checkpoint: {checkpoint.short_id} - {message}")
|
||||
|
||||
return CheckpointResult.ok(checkpoint)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._emit(CheckpointEvent.ERROR, {"error": error_msg})
|
||||
return CheckpointResult.fail(error_msg)
|
||||
|
||||
async def restore(self, checkpoint_id: str) -> CheckpointResult:
|
||||
"""
|
||||
Restore workspace to a checkpoint.
|
||||
|
||||
Args:
|
||||
checkpoint_id: Checkpoint ID (commit hash) to restore
|
||||
|
||||
Returns:
|
||||
CheckpointResult indicating success/failure
|
||||
"""
|
||||
if not self._initialized:
|
||||
return CheckpointResult.fail("Service not initialized")
|
||||
|
||||
try:
|
||||
# Clean untracked files
|
||||
await self._run_git("clean", "-f", "-d")
|
||||
|
||||
# Reset to checkpoint
|
||||
await self._run_git("reset", "--hard", checkpoint_id)
|
||||
|
||||
# Get checkpoint info
|
||||
message = await self._run_git("log", "-1", "--format=%s", checkpoint_id)
|
||||
timestamp = await self._run_git("log", "-1", "--format=%aI", checkpoint_id)
|
||||
|
||||
checkpoint = Checkpoint(
|
||||
id=checkpoint_id,
|
||||
short_id=checkpoint_id[:8],
|
||||
message=message,
|
||||
timestamp=datetime.fromisoformat(timestamp)
|
||||
)
|
||||
|
||||
self._emit(CheckpointEvent.CHECKPOINT_RESTORED, checkpoint)
|
||||
logger.info(f"Restored to checkpoint: {checkpoint.short_id}")
|
||||
|
||||
return CheckpointResult.ok(checkpoint)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._emit(CheckpointEvent.ERROR, {"error": error_msg})
|
||||
return CheckpointResult.fail(error_msg)
|
||||
|
||||
async def diff(
|
||||
self,
|
||||
from_id: Optional[str] = None,
|
||||
to_id: Optional[str] = None
|
||||
) -> CheckpointDiff:
|
||||
"""
|
||||
Get diff between checkpoints.
|
||||
|
||||
Args:
|
||||
from_id: Starting checkpoint (default: previous checkpoint)
|
||||
to_id: Ending checkpoint (default: current working directory)
|
||||
|
||||
Returns:
|
||||
CheckpointDiff with file changes
|
||||
"""
|
||||
if not self._initialized:
|
||||
return CheckpointDiff(from_checkpoint="", to_checkpoint="", files=[])
|
||||
|
||||
try:
|
||||
# Default from_id to HEAD~1
|
||||
if from_id is None:
|
||||
try:
|
||||
from_id = await self._run_git("rev-parse", "HEAD~1")
|
||||
except RuntimeError:
|
||||
from_id = await self._run_git("rev-parse", "HEAD")
|
||||
|
||||
# Build diff command
|
||||
if to_id:
|
||||
diff_output = await self._run_git(
|
||||
"diff", "--stat", "--numstat", from_id, to_id
|
||||
)
|
||||
else:
|
||||
# Diff against working directory
|
||||
await self._run_git("add", "-A") # Stage to include untracked
|
||||
diff_output = await self._run_git(
|
||||
"diff", "--stat", "--numstat", "--cached", from_id
|
||||
)
|
||||
|
||||
# Parse diff output
|
||||
files = self._parse_diff_output(diff_output)
|
||||
|
||||
return CheckpointDiff(
|
||||
from_checkpoint=from_id[:8] if from_id else "",
|
||||
to_checkpoint=to_id[:8] if to_id else None,
|
||||
files=files
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get diff: {e}")
|
||||
return CheckpointDiff(from_checkpoint="", to_checkpoint="", files=[])
|
||||
|
||||
def _parse_diff_output(self, output: str) -> List[FileDiff]:
|
||||
"""Parse git diff --numstat output."""
|
||||
files = []
|
||||
|
||||
for line in output.strip().split("\n"):
|
||||
if not line or "\t" not in line:
|
||||
continue
|
||||
|
||||
parts = line.split("\t")
|
||||
if len(parts) >= 3:
|
||||
try:
|
||||
additions = int(parts[0]) if parts[0] != "-" else 0
|
||||
deletions = int(parts[1]) if parts[1] != "-" else 0
|
||||
path = parts[2]
|
||||
|
||||
# Determine status
|
||||
if additions > 0 and deletions == 0:
|
||||
status = "added"
|
||||
elif additions == 0 and deletions > 0:
|
||||
status = "deleted"
|
||||
else:
|
||||
status = "modified"
|
||||
|
||||
files.append(FileDiff(
|
||||
path=path,
|
||||
absolute_path=os.path.join(self.config.workspace_dir, path),
|
||||
status=status,
|
||||
additions=additions,
|
||||
deletions=deletions
|
||||
))
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return files
|
||||
|
||||
async def list_checkpoints(self, limit: int = 50) -> List[Checkpoint]:
|
||||
"""
|
||||
List all checkpoints.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of checkpoints to return
|
||||
|
||||
Returns:
|
||||
List of checkpoints, newest first
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
try:
|
||||
# Get commit log
|
||||
log_output = await self._run_git(
|
||||
"log", f"-{limit}", "--format=%H|%s|%aI"
|
||||
)
|
||||
|
||||
checkpoints = []
|
||||
for line in log_output.strip().split("\n"):
|
||||
if not line or "|" not in line:
|
||||
continue
|
||||
|
||||
parts = line.split("|", 2)
|
||||
if len(parts) >= 3:
|
||||
checkpoints.append(Checkpoint(
|
||||
id=parts[0],
|
||||
short_id=parts[0][:8],
|
||||
message=parts[1],
|
||||
timestamp=datetime.fromisoformat(parts[2])
|
||||
))
|
||||
|
||||
return checkpoints
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list checkpoints: {e}")
|
||||
return []
|
||||
|
||||
async def _load_checkpoints(self):
|
||||
"""Load existing checkpoints from git log."""
|
||||
self._checkpoints = await self.list_checkpoints()
|
||||
|
||||
async def _prune_checkpoints(self):
|
||||
"""Prune old checkpoints if over limit."""
|
||||
if len(self._checkpoints) <= self.config.max_checkpoints:
|
||||
return
|
||||
|
||||
# Keep only the most recent checkpoints
|
||||
# Note: This doesn't actually delete git history, just our tracking
|
||||
self._checkpoints = self._checkpoints[:self.config.max_checkpoints]
|
||||
|
||||
async def get_checkpoint(self, checkpoint_id: str) -> Optional[Checkpoint]:
|
||||
"""Get a specific checkpoint by ID."""
|
||||
for cp in self._checkpoints:
|
||||
if cp.id.startswith(checkpoint_id) or cp.short_id == checkpoint_id:
|
||||
return cp
|
||||
return None
|
||||
|
||||
async def cleanup(self):
|
||||
"""Clean up the checkpoint service."""
|
||||
# Nothing to clean up currently
|
||||
self._initialized = False
|
||||
|
||||
async def delete_all(self) -> bool:
|
||||
"""Delete all checkpoint data for this workspace."""
|
||||
try:
|
||||
if os.path.exists(self.checkpoint_dir):
|
||||
shutil.rmtree(self.checkpoint_dir)
|
||||
self._checkpoints = []
|
||||
self._initialized = False
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete checkpoints: {e}")
|
||||
return False
|
||||
@@ -0,0 +1,158 @@
|
||||
"""
|
||||
Checkpoint Types for PraisonAI Agents.
|
||||
|
||||
Defines the core types, enums, and dataclasses for the checkpoint system.
|
||||
"""
|
||||
|
||||
import os
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class CheckpointEvent(str, Enum):
|
||||
"""Events emitted by the checkpoint service."""
|
||||
INITIALIZED = "initialized"
|
||||
CHECKPOINT_CREATED = "checkpoint_created"
|
||||
CHECKPOINT_RESTORED = "checkpoint_restored"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointConfig:
|
||||
"""Configuration for the checkpoint service."""
|
||||
workspace_dir: str
|
||||
storage_dir: Optional[str] = None
|
||||
enabled: bool = True
|
||||
auto_checkpoint: bool = True # Auto-checkpoint before file modifications
|
||||
max_checkpoints: int = 100 # Maximum checkpoints to keep
|
||||
exclude_patterns: List[str] = field(default_factory=lambda: [
|
||||
".git",
|
||||
".praison",
|
||||
"__pycache__",
|
||||
"*.pyc",
|
||||
".env",
|
||||
"node_modules",
|
||||
".venv",
|
||||
"venv",
|
||||
"*.log"
|
||||
])
|
||||
|
||||
def __post_init__(self):
|
||||
# Expand paths
|
||||
self.workspace_dir = os.path.expanduser(self.workspace_dir)
|
||||
if self.storage_dir:
|
||||
self.storage_dir = os.path.expanduser(self.storage_dir)
|
||||
else:
|
||||
# Default storage in user's home directory
|
||||
self.storage_dir = os.path.expanduser("~/.praison/checkpoints")
|
||||
|
||||
def get_checkpoint_dir(self) -> str:
|
||||
"""Get the checkpoint directory for this workspace."""
|
||||
# Hash the workspace path to create a unique directory
|
||||
import hashlib
|
||||
workspace_hash = hashlib.sha256(
|
||||
self.workspace_dir.encode()
|
||||
).hexdigest()[:12]
|
||||
return os.path.join(self.storage_dir, workspace_hash)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Checkpoint:
|
||||
"""Represents a single checkpoint."""
|
||||
id: str # Git commit hash
|
||||
short_id: str # Short hash (8 chars)
|
||||
message: str
|
||||
timestamp: datetime
|
||||
files_changed: int = 0
|
||||
insertions: int = 0
|
||||
deletions: int = 0
|
||||
|
||||
@classmethod
|
||||
def from_git_commit(cls, commit_hash: str, message: str, timestamp: str) -> "Checkpoint":
|
||||
"""Create a Checkpoint from git commit info."""
|
||||
return cls(
|
||||
id=commit_hash,
|
||||
short_id=commit_hash[:8],
|
||||
message=message,
|
||||
timestamp=datetime.fromisoformat(timestamp) if isinstance(timestamp, str) else timestamp
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"short_id": self.short_id,
|
||||
"message": self.message,
|
||||
"timestamp": self.timestamp.isoformat(),
|
||||
"files_changed": self.files_changed,
|
||||
"insertions": self.insertions,
|
||||
"deletions": self.deletions
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileDiff:
|
||||
"""Represents a diff for a single file."""
|
||||
path: str
|
||||
absolute_path: str
|
||||
status: str # added, modified, deleted
|
||||
additions: int = 0
|
||||
deletions: int = 0
|
||||
before_content: Optional[str] = None
|
||||
after_content: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"path": self.path,
|
||||
"absolute_path": self.absolute_path,
|
||||
"status": self.status,
|
||||
"additions": self.additions,
|
||||
"deletions": self.deletions
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointDiff:
|
||||
"""Represents a diff between two checkpoints."""
|
||||
from_checkpoint: str
|
||||
to_checkpoint: Optional[str] # None means current working directory
|
||||
files: List[FileDiff] = field(default_factory=list)
|
||||
total_additions: int = 0
|
||||
total_deletions: int = 0
|
||||
|
||||
def __post_init__(self):
|
||||
# Calculate totals
|
||||
self.total_additions = sum(f.additions for f in self.files)
|
||||
self.total_deletions = sum(f.deletions for f in self.files)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"from_checkpoint": self.from_checkpoint,
|
||||
"to_checkpoint": self.to_checkpoint,
|
||||
"files": [f.to_dict() for f in self.files],
|
||||
"total_additions": self.total_additions,
|
||||
"total_deletions": self.total_deletions
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointResult:
|
||||
"""Result of a checkpoint operation."""
|
||||
success: bool
|
||||
checkpoint: Optional[Checkpoint] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def ok(cls, checkpoint: Checkpoint) -> "CheckpointResult":
|
||||
"""Create a successful result."""
|
||||
return cls(success=True, checkpoint=checkpoint)
|
||||
|
||||
@classmethod
|
||||
def fail(cls, error: str) -> "CheckpointResult":
|
||||
"""Create a failed result."""
|
||||
return cls(success=False, error=error)
|
||||
@@ -0,0 +1,57 @@
|
||||
"""
|
||||
Auto Context Compaction Module for PraisonAI Agents.
|
||||
|
||||
Provides automatic context management for long conversations:
|
||||
- Message summarization
|
||||
- Context window management
|
||||
- Token counting and limits
|
||||
- Intelligent message pruning
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Compaction only runs when needed
|
||||
- No overhead when context is within limits
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.compaction import ContextCompactor, CompactionConfig
|
||||
|
||||
# Create a compactor
|
||||
compactor = ContextCompactor(
|
||||
max_tokens=8000,
|
||||
strategy="summarize"
|
||||
)
|
||||
|
||||
# Compact messages when needed
|
||||
compacted = compactor.compact(messages)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core classes
|
||||
"ContextCompactor",
|
||||
"CompactionConfig",
|
||||
# Strategies
|
||||
"CompactionStrategy",
|
||||
# Results
|
||||
"CompactionResult",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "ContextCompactor":
|
||||
from .compactor import ContextCompactor
|
||||
return ContextCompactor
|
||||
|
||||
if name == "CompactionConfig":
|
||||
from .config import CompactionConfig
|
||||
return CompactionConfig
|
||||
|
||||
if name == "CompactionStrategy":
|
||||
from .strategy import CompactionStrategy
|
||||
return CompactionStrategy
|
||||
|
||||
if name == "CompactionResult":
|
||||
from .result import CompactionResult
|
||||
return CompactionResult
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
Context Compactor for PraisonAI Agents.
|
||||
|
||||
Manages context window by compacting messages when needed.
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
from .config import CompactionConfig
|
||||
from .strategy import CompactionStrategy
|
||||
from .result import CompactionResult
|
||||
|
||||
|
||||
class ContextCompactor:
|
||||
"""
|
||||
Compacts conversation context to fit within token limits.
|
||||
|
||||
Example:
|
||||
compactor = ContextCompactor(max_tokens=8000)
|
||||
|
||||
# Check if compaction needed
|
||||
if compactor.needs_compaction(messages):
|
||||
result = compactor.compact(messages)
|
||||
messages = result.messages
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_tokens: int = 8000,
|
||||
target_tokens: Optional[int] = None,
|
||||
strategy: CompactionStrategy = CompactionStrategy.TRUNCATE,
|
||||
preserve_system: bool = True,
|
||||
preserve_recent: int = 5
|
||||
):
|
||||
"""
|
||||
Initialize the compactor.
|
||||
|
||||
Args:
|
||||
max_tokens: Maximum tokens before compaction
|
||||
target_tokens: Target tokens after compaction
|
||||
strategy: Compaction strategy to use
|
||||
preserve_system: Keep system messages
|
||||
preserve_recent: Number of recent messages to preserve
|
||||
"""
|
||||
self.max_tokens = max_tokens
|
||||
self.target_tokens = target_tokens or int(max_tokens * 0.75)
|
||||
self.strategy = strategy
|
||||
self.preserve_system = preserve_system
|
||||
self.preserve_recent = preserve_recent
|
||||
|
||||
def estimate_tokens(self, text: str) -> int:
|
||||
"""Estimate token count for text."""
|
||||
# Rough estimate: ~4 chars per token
|
||||
return len(text) // 4
|
||||
|
||||
def count_message_tokens(self, message: Dict[str, Any]) -> int:
|
||||
"""Count tokens in a message."""
|
||||
content = message.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return self.estimate_tokens(content)
|
||||
elif isinstance(content, list):
|
||||
# Handle multi-part content
|
||||
total = 0
|
||||
for part in content:
|
||||
if isinstance(part, dict):
|
||||
total += self.estimate_tokens(str(part.get("text", "")))
|
||||
else:
|
||||
total += self.estimate_tokens(str(part))
|
||||
return total
|
||||
return 0
|
||||
|
||||
def count_total_tokens(self, messages: List[Dict[str, Any]]) -> int:
|
||||
"""Count total tokens in messages."""
|
||||
return sum(self.count_message_tokens(m) for m in messages)
|
||||
|
||||
def needs_compaction(self, messages: List[Dict[str, Any]]) -> bool:
|
||||
"""Check if messages need compaction."""
|
||||
return self.count_total_tokens(messages) > self.max_tokens
|
||||
|
||||
def compact(
|
||||
self,
|
||||
messages: List[Dict[str, Any]]
|
||||
) -> tuple[List[Dict[str, Any]], CompactionResult]:
|
||||
"""
|
||||
Compact messages to fit within token limit.
|
||||
|
||||
Args:
|
||||
messages: List of messages to compact
|
||||
|
||||
Returns:
|
||||
Tuple of (compacted messages, result)
|
||||
"""
|
||||
original_tokens = self.count_total_tokens(messages)
|
||||
|
||||
if original_tokens <= self.max_tokens:
|
||||
return messages, CompactionResult(
|
||||
original_tokens=original_tokens,
|
||||
compacted_tokens=original_tokens,
|
||||
messages_removed=0,
|
||||
messages_kept=len(messages),
|
||||
strategy_used=self.strategy
|
||||
)
|
||||
|
||||
if self.strategy == CompactionStrategy.TRUNCATE:
|
||||
compacted = self._truncate(messages)
|
||||
elif self.strategy == CompactionStrategy.SLIDING:
|
||||
compacted = self._sliding_window(messages)
|
||||
elif self.strategy == CompactionStrategy.SUMMARIZE:
|
||||
compacted = self._summarize(messages)
|
||||
elif self.strategy == CompactionStrategy.SMART:
|
||||
compacted = self._smart_compact(messages)
|
||||
else:
|
||||
compacted = self._truncate(messages)
|
||||
|
||||
compacted_tokens = self.count_total_tokens(compacted)
|
||||
|
||||
result = CompactionResult(
|
||||
original_tokens=original_tokens,
|
||||
compacted_tokens=compacted_tokens,
|
||||
messages_removed=len(messages) - len(compacted),
|
||||
messages_kept=len(compacted),
|
||||
strategy_used=self.strategy
|
||||
)
|
||||
|
||||
return compacted, result
|
||||
|
||||
def _truncate(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Truncate oldest messages."""
|
||||
result = []
|
||||
|
||||
# Separate system and non-system messages
|
||||
system_msgs = []
|
||||
other_msgs = []
|
||||
|
||||
for msg in messages:
|
||||
if self.preserve_system and msg.get("role") == "system":
|
||||
system_msgs.append(msg)
|
||||
else:
|
||||
other_msgs.append(msg)
|
||||
|
||||
# Always keep system messages
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:] if other_msgs else []
|
||||
|
||||
# Add recent messages
|
||||
result.extend(recent)
|
||||
|
||||
# If still over limit, truncate more
|
||||
while self.count_total_tokens(result) > self.target_tokens and len(result) > 1:
|
||||
# Remove oldest non-system message
|
||||
for i, msg in enumerate(result):
|
||||
if msg.get("role") != "system":
|
||||
result.pop(i)
|
||||
break
|
||||
|
||||
return result
|
||||
|
||||
def _sliding_window(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Keep a sliding window of recent messages."""
|
||||
result = []
|
||||
|
||||
# Keep system messages
|
||||
for msg in messages:
|
||||
if self.preserve_system and msg.get("role") == "system":
|
||||
result.append(msg)
|
||||
|
||||
# Add messages from end until we hit target
|
||||
non_system = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
for msg in reversed(non_system):
|
||||
if self.count_total_tokens(result + [msg]) <= self.target_tokens:
|
||||
result.insert(len([m for m in result if m.get("role") == "system"]), msg)
|
||||
else:
|
||||
break
|
||||
|
||||
# Ensure messages are in order
|
||||
system_msgs = [m for m in result if m.get("role") == "system"]
|
||||
other_msgs = [m for m in result if m.get("role") != "system"]
|
||||
|
||||
return system_msgs + other_msgs
|
||||
|
||||
def _summarize(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Summarize old messages (simplified version)."""
|
||||
result = []
|
||||
|
||||
# Keep system messages
|
||||
system_msgs = [m for m in messages if m.get("role") == "system"]
|
||||
other_msgs = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
result.extend(system_msgs)
|
||||
|
||||
# Keep recent messages
|
||||
recent = other_msgs[-self.preserve_recent:]
|
||||
|
||||
# Summarize older messages
|
||||
older = other_msgs[:-self.preserve_recent] if len(other_msgs) > self.preserve_recent else []
|
||||
|
||||
if older:
|
||||
# Create a simple summary
|
||||
summary_parts = []
|
||||
for msg in older:
|
||||
role = msg.get("role", "unknown")
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str) and content:
|
||||
# Take first 100 chars
|
||||
summary_parts.append(f"{role}: {content[:100]}...")
|
||||
|
||||
if summary_parts:
|
||||
summary = "[Previous conversation summary]\n" + "\n".join(summary_parts[:5])
|
||||
result.append({
|
||||
"role": "system",
|
||||
"content": summary
|
||||
})
|
||||
|
||||
result.extend(recent)
|
||||
|
||||
return result
|
||||
|
||||
def _smart_compact(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Smart compaction based on message importance."""
|
||||
# For now, use sliding window as base
|
||||
return self._sliding_window(messages)
|
||||
|
||||
def get_stats(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""Get statistics about messages."""
|
||||
total_tokens = self.count_total_tokens(messages)
|
||||
|
||||
return {
|
||||
"message_count": len(messages),
|
||||
"total_tokens": total_tokens,
|
||||
"max_tokens": self.max_tokens,
|
||||
"target_tokens": self.target_tokens,
|
||||
"needs_compaction": total_tokens > self.max_tokens,
|
||||
"utilization": total_tokens / self.max_tokens if self.max_tokens > 0 else 0
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
Compaction Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionConfig:
|
||||
"""Configuration for context compaction."""
|
||||
enabled: bool = True
|
||||
max_tokens: int = 8000
|
||||
target_tokens: int = 6000 # Target after compaction
|
||||
preserve_system: bool = True # Always keep system messages
|
||||
preserve_recent: int = 5 # Keep last N messages
|
||||
auto_compact: bool = True # Automatically compact when needed
|
||||
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
Compaction Result for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from .strategy import CompactionStrategy
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionResult:
|
||||
"""Result of a compaction operation."""
|
||||
original_tokens: int
|
||||
compacted_tokens: int
|
||||
messages_removed: int
|
||||
messages_kept: int
|
||||
strategy_used: CompactionStrategy
|
||||
summary: str = ""
|
||||
|
||||
@property
|
||||
def tokens_saved(self) -> int:
|
||||
"""Get number of tokens saved."""
|
||||
return self.original_tokens - self.compacted_tokens
|
||||
|
||||
@property
|
||||
def compression_ratio(self) -> float:
|
||||
"""Get compression ratio."""
|
||||
if self.original_tokens == 0:
|
||||
return 0.0
|
||||
return self.compacted_tokens / self.original_tokens
|
||||
|
||||
@property
|
||||
def was_compacted(self) -> bool:
|
||||
"""Check if compaction occurred."""
|
||||
return self.messages_removed > 0 or self.tokens_saved > 0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"original_tokens": self.original_tokens,
|
||||
"compacted_tokens": self.compacted_tokens,
|
||||
"tokens_saved": self.tokens_saved,
|
||||
"messages_removed": self.messages_removed,
|
||||
"messages_kept": self.messages_kept,
|
||||
"strategy_used": self.strategy_used.value,
|
||||
"compression_ratio": self.compression_ratio,
|
||||
"summary": self.summary
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
Compaction Strategies for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CompactionStrategy(str, Enum):
|
||||
"""Strategy for compacting context."""
|
||||
TRUNCATE = "truncate" # Remove oldest messages
|
||||
SUMMARIZE = "summarize" # Summarize old messages
|
||||
SLIDING = "sliding" # Sliding window of recent messages
|
||||
SMART = "smart" # Intelligent selection based on relevance
|
||||
@@ -0,0 +1,85 @@
|
||||
"""
|
||||
Database adapter interface for PraisonAI Agents.
|
||||
|
||||
This module provides the protocol and types for database persistence.
|
||||
Implementations are provided by the wrapper layer (praisonai.db).
|
||||
|
||||
Usage (simplest - recommended):
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
db=db(database_url="postgresql://localhost/mydb"), # db(...) shortcut
|
||||
session_id="my-session" # optional: defaults to per-hour ID
|
||||
)
|
||||
agent.chat("Hello!") # auto-persists messages, runs, traces, tool calls
|
||||
|
||||
Alternative (explicit backend):
|
||||
db=db.PraisonDB(database_url="...") # Auto-detect backend
|
||||
db=db.PostgresDB(host="localhost") # PostgreSQL
|
||||
db=db.SQLiteDB(path="data.db") # SQLite
|
||||
db=db.RedisDB(host="localhost") # Redis (state only)
|
||||
"""
|
||||
|
||||
from .protocol import (
|
||||
DbAdapter,
|
||||
AsyncDbAdapter,
|
||||
DbMessage,
|
||||
DbToolCall,
|
||||
DbRun,
|
||||
DbSpan,
|
||||
DbTrace,
|
||||
)
|
||||
|
||||
|
||||
class _LazyDbModule:
|
||||
"""
|
||||
Lazy proxy for database backends.
|
||||
|
||||
Allows `from praisonaiagents import db` without importing heavy deps.
|
||||
Actual backend classes are loaded only when accessed.
|
||||
"""
|
||||
|
||||
_BACKENDS = {
|
||||
"PraisonDB": "praisonai.db.adapter",
|
||||
"PostgresDB": "praisonai.db.adapter",
|
||||
"SQLiteDB": "praisonai.db.adapter",
|
||||
"RedisDB": "praisonai.db.adapter",
|
||||
}
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name in self._BACKENDS:
|
||||
try:
|
||||
import importlib
|
||||
module = importlib.import_module(self._BACKENDS[name])
|
||||
return getattr(module, name)
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
f"Database backend '{name}' requires the praisonai package. "
|
||||
f"Install with: pip install praisonai\n"
|
||||
f"Original error: {e}"
|
||||
) from e
|
||||
raise AttributeError(f"module 'db' has no attribute {name!r}")
|
||||
|
||||
def __repr__(self):
|
||||
return "<module 'praisonaiagents.db' (lazy database backends)>"
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
"""Shortcut: db(...) is equivalent to db.PraisonDB(...)"""
|
||||
return self.PraisonDB(**kwargs)
|
||||
|
||||
|
||||
# Singleton lazy module instance
|
||||
db = _LazyDbModule()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DbAdapter",
|
||||
"AsyncDbAdapter",
|
||||
"DbMessage",
|
||||
"DbToolCall",
|
||||
"DbRun",
|
||||
"DbSpan",
|
||||
"DbTrace",
|
||||
"db",
|
||||
]
|
||||
@@ -0,0 +1,464 @@
|
||||
"""
|
||||
Database adapter protocol for PraisonAI Agents.
|
||||
|
||||
This module uses only stdlib (typing, dataclasses) to define the interface.
|
||||
No external dependencies are imported here.
|
||||
|
||||
Implementations are provided by the wrapper layer (praisonai.db).
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
|
||||
import time
|
||||
|
||||
|
||||
@dataclass
|
||||
class DbToolCall:
|
||||
"""A tool call to be persisted."""
|
||||
tool_name: str
|
||||
args: Dict[str, Any]
|
||||
result: Any = None
|
||||
call_id: Optional[str] = None
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
duration_ms: Optional[float] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DbMessage:
|
||||
"""A message to be persisted."""
|
||||
role: str # "user", "assistant", "system", "tool"
|
||||
content: str
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
id: Optional[str] = None
|
||||
tool_calls: Optional[List["DbToolCall"]] = None # For assistant messages with tool calls
|
||||
run_id: Optional[str] = None # Group messages by run
|
||||
|
||||
|
||||
@dataclass
|
||||
class DbRun:
|
||||
"""A single agent run (turn) to be persisted."""
|
||||
run_id: str
|
||||
session_id: str
|
||||
agent_name: str
|
||||
started_at: float = field(default_factory=time.time)
|
||||
ended_at: Optional[float] = None
|
||||
status: str = "running" # running, completed, error, cancelled
|
||||
input_content: Optional[str] = None
|
||||
output_content: Optional[str] = None
|
||||
messages: List[DbMessage] = field(default_factory=list)
|
||||
tool_calls: List[DbToolCall] = field(default_factory=list)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
metrics: Dict[str, Any] = field(default_factory=dict) # tokens, latency, etc.
|
||||
|
||||
|
||||
@dataclass
|
||||
class DbSpan:
|
||||
"""A span for tracing/observability."""
|
||||
span_id: str
|
||||
trace_id: str
|
||||
name: str
|
||||
started_at: float = field(default_factory=time.time)
|
||||
ended_at: Optional[float] = None
|
||||
parent_span_id: Optional[str] = None
|
||||
status: str = "running" # running, ok, error
|
||||
attributes: Dict[str, Any] = field(default_factory=dict)
|
||||
events: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DbTrace:
|
||||
"""A trace for observability (groups spans)."""
|
||||
trace_id: str
|
||||
session_id: Optional[str] = None
|
||||
run_id: Optional[str] = None
|
||||
agent_name: Optional[str] = None
|
||||
user_id: Optional[str] = None
|
||||
started_at: float = field(default_factory=time.time)
|
||||
ended_at: Optional[float] = None
|
||||
status: str = "running" # running, ok, error
|
||||
spans: List[DbSpan] = field(default_factory=list)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class DbAdapter(Protocol):
|
||||
"""
|
||||
Protocol for database adapters that provide persistence for agents.
|
||||
|
||||
Implementations must provide these methods. The agent will call them
|
||||
at appropriate lifecycle points if a db adapter is provided.
|
||||
|
||||
All methods are synchronous. Async support can be added via wrapper.
|
||||
"""
|
||||
|
||||
def on_agent_start(
|
||||
self,
|
||||
agent_name: str,
|
||||
session_id: str,
|
||||
user_id: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> List[DbMessage]:
|
||||
"""
|
||||
Called when agent starts a session.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent
|
||||
session_id: Session identifier
|
||||
user_id: Optional user identifier
|
||||
metadata: Optional session metadata
|
||||
|
||||
Returns:
|
||||
List of previous messages for this session (for resume)
|
||||
"""
|
||||
...
|
||||
|
||||
def on_user_message(
|
||||
self,
|
||||
session_id: str,
|
||||
content: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when user sends a message.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
content: Message content
|
||||
metadata: Optional message metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def on_agent_message(
|
||||
self,
|
||||
session_id: str,
|
||||
content: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when agent produces a response.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
content: Response content
|
||||
metadata: Optional message metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def on_tool_call(
|
||||
self,
|
||||
session_id: str,
|
||||
tool_name: str,
|
||||
args: Dict[str, Any],
|
||||
result: Any,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when a tool is executed.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
tool_name: Name of the tool
|
||||
args: Tool arguments
|
||||
result: Tool result
|
||||
metadata: Optional call metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def on_agent_end(
|
||||
self,
|
||||
session_id: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when agent session ends.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
metadata: Optional end metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def on_run_start(
|
||||
self,
|
||||
session_id: str,
|
||||
run_id: str,
|
||||
input_content: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when a new run (turn) starts.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
run_id: Run identifier
|
||||
input_content: User input for this run
|
||||
metadata: Optional run metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def on_run_end(
|
||||
self,
|
||||
session_id: str,
|
||||
run_id: str,
|
||||
output_content: Optional[str] = None,
|
||||
status: str = "completed",
|
||||
metrics: Optional[Dict[str, Any]] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when a run (turn) ends.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
run_id: Run identifier
|
||||
output_content: Agent output for this run
|
||||
status: Run status (completed, error, cancelled)
|
||||
metrics: Optional metrics (tokens, latency)
|
||||
metadata: Optional run metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def get_runs(
|
||||
self,
|
||||
session_id: str,
|
||||
limit: Optional[int] = None,
|
||||
) -> List["DbRun"]:
|
||||
"""
|
||||
Get runs for a session.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
limit: Maximum number of runs to return
|
||||
|
||||
Returns:
|
||||
List of runs for this session
|
||||
"""
|
||||
...
|
||||
|
||||
def export_session(
|
||||
self,
|
||||
session_id: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Export a session as a dictionary (for JSONL export).
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with session data, messages, runs
|
||||
"""
|
||||
...
|
||||
|
||||
def import_session(
|
||||
self,
|
||||
data: Dict[str, Any],
|
||||
) -> str:
|
||||
"""
|
||||
Import a session from a dictionary.
|
||||
|
||||
Args:
|
||||
data: Session data dictionary
|
||||
|
||||
Returns:
|
||||
Session ID of imported session
|
||||
"""
|
||||
...
|
||||
|
||||
# --- Tracing/Observability ---
|
||||
|
||||
def on_trace_start(
|
||||
self,
|
||||
trace_id: str,
|
||||
session_id: Optional[str] = None,
|
||||
run_id: Optional[str] = None,
|
||||
agent_name: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when a new trace starts.
|
||||
|
||||
Args:
|
||||
trace_id: Unique trace identifier
|
||||
session_id: Optional session identifier
|
||||
run_id: Optional run identifier
|
||||
agent_name: Optional agent name
|
||||
user_id: Optional user identifier
|
||||
metadata: Optional trace metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def on_trace_end(
|
||||
self,
|
||||
trace_id: str,
|
||||
status: str = "ok",
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when a trace ends.
|
||||
|
||||
Args:
|
||||
trace_id: Trace identifier
|
||||
status: Trace status (ok, error)
|
||||
metadata: Optional end metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def on_span_start(
|
||||
self,
|
||||
span_id: str,
|
||||
trace_id: str,
|
||||
name: str,
|
||||
parent_span_id: Optional[str] = None,
|
||||
attributes: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when a new span starts.
|
||||
|
||||
Args:
|
||||
span_id: Unique span identifier
|
||||
trace_id: Parent trace identifier
|
||||
name: Span name (e.g., "llm_call", "tool_execution")
|
||||
parent_span_id: Optional parent span for nesting
|
||||
attributes: Optional span attributes
|
||||
"""
|
||||
...
|
||||
|
||||
def on_span_end(
|
||||
self,
|
||||
span_id: str,
|
||||
status: str = "ok",
|
||||
attributes: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called when a span ends.
|
||||
|
||||
Args:
|
||||
span_id: Span identifier
|
||||
status: Span status (ok, error)
|
||||
attributes: Optional end attributes
|
||||
"""
|
||||
...
|
||||
|
||||
def get_traces(
|
||||
self,
|
||||
session_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> List["DbTrace"]:
|
||||
"""
|
||||
Get traces with optional filters.
|
||||
|
||||
Args:
|
||||
session_id: Filter by session
|
||||
user_id: Filter by user
|
||||
limit: Maximum number of traces
|
||||
|
||||
Returns:
|
||||
List of traces
|
||||
"""
|
||||
...
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the database connection."""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class AsyncDbAdapter(Protocol):
|
||||
"""
|
||||
Async protocol for database adapters.
|
||||
|
||||
Provides async variants of all DbAdapter methods for use with
|
||||
async/await patterns and async backends like asyncpg.
|
||||
"""
|
||||
|
||||
async def on_agent_start(
|
||||
self,
|
||||
agent_name: str,
|
||||
session_id: str,
|
||||
user_id: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> List["DbMessage"]:
|
||||
"""Async version of on_agent_start."""
|
||||
...
|
||||
|
||||
async def on_user_message(
|
||||
self,
|
||||
session_id: str,
|
||||
content: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Async version of on_user_message."""
|
||||
...
|
||||
|
||||
async def on_agent_message(
|
||||
self,
|
||||
session_id: str,
|
||||
content: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Async version of on_agent_message."""
|
||||
...
|
||||
|
||||
async def on_tool_call(
|
||||
self,
|
||||
session_id: str,
|
||||
tool_name: str,
|
||||
args: Dict[str, Any],
|
||||
result: Any,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Async version of on_tool_call."""
|
||||
...
|
||||
|
||||
async def on_agent_end(
|
||||
self,
|
||||
session_id: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Async version of on_agent_end."""
|
||||
...
|
||||
|
||||
async def on_run_start(
|
||||
self,
|
||||
session_id: str,
|
||||
run_id: str,
|
||||
input_content: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Async version of on_run_start."""
|
||||
...
|
||||
|
||||
async def on_run_end(
|
||||
self,
|
||||
session_id: str,
|
||||
run_id: str,
|
||||
output_content: Optional[str] = None,
|
||||
status: str = "completed",
|
||||
metrics: Optional[Dict[str, Any]] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Async version of on_run_end."""
|
||||
...
|
||||
|
||||
async def export_session(
|
||||
self,
|
||||
session_id: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""Async version of export_session."""
|
||||
...
|
||||
|
||||
async def import_session(
|
||||
self,
|
||||
data: Dict[str, Any],
|
||||
) -> str:
|
||||
"""Async version of import_session."""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Async close the database connection."""
|
||||
...
|
||||
@@ -0,0 +1 @@
|
||||
"""TDD tests for db integration."""
|
||||
@@ -0,0 +1,246 @@
|
||||
"""
|
||||
TDD tests for PraisonAI Agents DB integration.
|
||||
|
||||
Tests:
|
||||
1. Lazy import verification - no heavy deps on import
|
||||
2. Agent(db=MockDb) calls hook lifecycle in correct order
|
||||
3. Session resume injects history
|
||||
4. db=None behavior unchanged
|
||||
"""
|
||||
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
def test_lazy_import_no_heavy_deps():
|
||||
"""Test that importing praisonaiagents.db does not import DB libs."""
|
||||
print("=== Test: Lazy Import Verification ===")
|
||||
|
||||
# Record modules before import
|
||||
before_modules = set(sys.modules.keys())
|
||||
|
||||
# Import only the db module (not full praisonaiagents which loads Knowledge)
|
||||
from praisonaiagents.db import DbAdapter, DbMessage
|
||||
|
||||
# Check that heavy DB libs are NOT loaded by db module itself
|
||||
# Note: praisonaiagents main package loads qdrant/chroma via Knowledge
|
||||
# but the db module should not add any additional heavy deps
|
||||
heavy_libs = ['psycopg2', 'pymysql', 'redis'] # DB-specific libs
|
||||
loaded_heavy = [lib for lib in heavy_libs if lib in sys.modules]
|
||||
|
||||
if loaded_heavy:
|
||||
print(f" ❌ Heavy DB libs loaded by db module: {loaded_heavy}")
|
||||
return False
|
||||
|
||||
print(" ✅ No heavy DB libs loaded by praisonaiagents.db module")
|
||||
return True
|
||||
|
||||
|
||||
def test_agent_accepts_db_parameter():
|
||||
"""Test that Agent accepts db and session_id parameters."""
|
||||
print("\n=== Test: Agent Accepts db Parameter ===")
|
||||
|
||||
import inspect
|
||||
from praisonaiagents import Agent
|
||||
|
||||
sig = inspect.signature(Agent.__init__)
|
||||
params = list(sig.parameters.keys())
|
||||
|
||||
if 'db' not in params:
|
||||
print(" ❌ Agent does not accept 'db' parameter")
|
||||
return False
|
||||
|
||||
if 'session_id' not in params:
|
||||
print(" ❌ Agent does not accept 'session_id' parameter")
|
||||
return False
|
||||
|
||||
print(" ✅ Agent accepts db and session_id parameters")
|
||||
return True
|
||||
|
||||
|
||||
class MockDbAdapter:
|
||||
"""Mock DB adapter for testing hook lifecycle."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
self.history_to_return = []
|
||||
|
||||
def on_agent_start(self, agent_name, session_id, user_id=None, metadata=None):
|
||||
self.calls.append(('on_agent_start', agent_name, session_id))
|
||||
return self.history_to_return
|
||||
|
||||
def on_user_message(self, session_id, content, metadata=None):
|
||||
self.calls.append(('on_user_message', session_id, content))
|
||||
|
||||
def on_agent_message(self, session_id, content, metadata=None):
|
||||
self.calls.append(('on_agent_message', session_id, content))
|
||||
|
||||
def on_tool_call(self, session_id, tool_name, args, result, metadata=None):
|
||||
self.calls.append(('on_tool_call', session_id, tool_name))
|
||||
|
||||
def on_agent_end(self, session_id, metadata=None):
|
||||
self.calls.append(('on_agent_end', session_id))
|
||||
|
||||
def close(self):
|
||||
self.calls.append(('close',))
|
||||
|
||||
|
||||
def test_db_hook_lifecycle():
|
||||
"""Test that Agent calls db hooks in correct order."""
|
||||
print("\n=== Test: DB Hook Lifecycle ===")
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
mock_db = MockDbAdapter()
|
||||
|
||||
# Create agent with mock db
|
||||
agent = Agent(
|
||||
name="TestAgent",
|
||||
instructions="Test agent",
|
||||
db=mock_db,
|
||||
session_id="test-session-001",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Verify db is stored
|
||||
if agent._db is not mock_db:
|
||||
print(" ❌ Agent did not store db adapter")
|
||||
return False
|
||||
|
||||
if agent._session_id != "test-session-001":
|
||||
print(" ❌ Agent did not store session_id")
|
||||
return False
|
||||
|
||||
print(" ✅ Agent stores db and session_id correctly")
|
||||
return True
|
||||
|
||||
|
||||
def test_db_none_unchanged_behavior():
|
||||
"""Test that db=None behavior is unchanged."""
|
||||
print("\n=== Test: db=None Unchanged Behavior ===")
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Create agent without db
|
||||
agent = Agent(
|
||||
name="TestAgent",
|
||||
instructions="Test agent",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Verify db is None
|
||||
if agent._db is not None:
|
||||
print(" ❌ Agent._db should be None when not provided")
|
||||
return False
|
||||
|
||||
# Verify _init_db_session does nothing when db is None
|
||||
agent._init_db_session()
|
||||
|
||||
if agent._db_initialized:
|
||||
print(" ❌ _db_initialized should remain False when db is None")
|
||||
return False
|
||||
|
||||
print(" ✅ db=None behavior unchanged")
|
||||
return True
|
||||
|
||||
|
||||
def test_session_id_auto_generation():
|
||||
"""Test that session_id is auto-generated if not provided."""
|
||||
print("\n=== Test: Session ID Auto-Generation ===")
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
mock_db = MockDbAdapter()
|
||||
|
||||
agent = Agent(
|
||||
name="TestAgent",
|
||||
instructions="Test agent",
|
||||
db=mock_db,
|
||||
# No session_id provided
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Trigger db initialization
|
||||
agent._init_db_session()
|
||||
|
||||
if agent._session_id is None:
|
||||
print(" ❌ session_id should be auto-generated")
|
||||
return False
|
||||
|
||||
if not agent._session_id.startswith("session-"):
|
||||
print(f" ❌ session_id format unexpected: {agent._session_id}")
|
||||
return False
|
||||
|
||||
print(f" ✅ session_id auto-generated: {agent._session_id}")
|
||||
return True
|
||||
|
||||
|
||||
def test_protocol_compliance():
|
||||
"""Test that DbAdapter protocol is properly defined."""
|
||||
print("\n=== Test: DbAdapter Protocol Compliance ===")
|
||||
|
||||
from praisonaiagents.db.protocol import DbAdapter, DbMessage
|
||||
|
||||
# Check DbMessage fields
|
||||
msg = DbMessage(role="user", content="test")
|
||||
if msg.role != "user" or msg.content != "test":
|
||||
print(" ❌ DbMessage fields not working")
|
||||
return False
|
||||
|
||||
# Check protocol methods
|
||||
required_methods = [
|
||||
'on_agent_start',
|
||||
'on_user_message',
|
||||
'on_agent_message',
|
||||
'on_tool_call',
|
||||
'on_agent_end',
|
||||
'close'
|
||||
]
|
||||
|
||||
mock = MockDbAdapter()
|
||||
for method in required_methods:
|
||||
if not hasattr(mock, method):
|
||||
print(f" ❌ MockDbAdapter missing {method}")
|
||||
return False
|
||||
|
||||
# Check protocol isinstance
|
||||
if not isinstance(mock, DbAdapter):
|
||||
print(" ❌ MockDbAdapter should implement DbAdapter protocol")
|
||||
return False
|
||||
|
||||
print(" ✅ DbAdapter protocol properly defined")
|
||||
return True
|
||||
|
||||
|
||||
def run_all_tests():
|
||||
"""Run all TDD tests."""
|
||||
print("=" * 60)
|
||||
print("PraisonAI Agents DB Integration - TDD Tests")
|
||||
print("=" * 60)
|
||||
|
||||
results = []
|
||||
|
||||
results.append(("Lazy Import", test_lazy_import_no_heavy_deps()))
|
||||
results.append(("Agent db Parameter", test_agent_accepts_db_parameter()))
|
||||
results.append(("DB Hook Lifecycle", test_db_hook_lifecycle()))
|
||||
results.append(("db=None Unchanged", test_db_none_unchanged_behavior()))
|
||||
results.append(("Session ID Auto-Gen", test_session_id_auto_generation()))
|
||||
results.append(("Protocol Compliance", test_protocol_compliance()))
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
passed = sum(1 for _, r in results if r)
|
||||
total = len(results)
|
||||
print(f"Results: {passed}/{total} tests passed")
|
||||
|
||||
for name, result in results:
|
||||
status = "✅" if result else "❌"
|
||||
print(f" {status} {name}")
|
||||
|
||||
print("=" * 60)
|
||||
|
||||
return passed == total
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
success = run_all_tests()
|
||||
sys.exit(0 if success else 1)
|
||||
@@ -0,0 +1,137 @@
|
||||
"""
|
||||
Hooks Module for PraisonAI Agents.
|
||||
|
||||
Provides a powerful hook system for intercepting and modifying agent behavior
|
||||
at various lifecycle points. Unlike callbacks (which are for UI events),
|
||||
hooks can intercept, modify, or block tool execution.
|
||||
|
||||
Features:
|
||||
- Event-based hook system (BeforeTool, AfterTool, BeforeAgent, etc.)
|
||||
- Shell command hooks for external integrations
|
||||
- Python function hooks for in-process customization
|
||||
- Matcher patterns for selective hook execution
|
||||
- Sequential and parallel hook execution
|
||||
- Decision outcomes (allow, deny, block, ask)
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Hooks only execute when registered
|
||||
- No overhead when hooks are disabled
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.hooks import HookRegistry, HookEvent
|
||||
|
||||
# Register a Python function hook
|
||||
registry = HookRegistry()
|
||||
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def my_hook(event_data):
|
||||
if event_data.tool_name == "dangerous_tool":
|
||||
return HookResult(decision="deny", reason="Tool blocked by policy")
|
||||
return HookResult(decision="allow")
|
||||
|
||||
# Register a shell command hook
|
||||
registry.register_command_hook(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="python /path/to/validator.py",
|
||||
matcher="write_*" # Only match tools starting with write_
|
||||
)
|
||||
|
||||
# Use with Agent
|
||||
agent = Agent(
|
||||
name="MyAgent",
|
||||
hooks=registry
|
||||
)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core types
|
||||
"HookEvent",
|
||||
"HookDecision",
|
||||
"HookResult",
|
||||
"HookInput",
|
||||
"HookOutput",
|
||||
# Hook definitions
|
||||
"HookDefinition",
|
||||
"CommandHook",
|
||||
"FunctionHook",
|
||||
# Registry and runner
|
||||
"HookRegistry",
|
||||
"HookRunner",
|
||||
# Event-specific inputs
|
||||
"BeforeToolInput",
|
||||
"AfterToolInput",
|
||||
"BeforeAgentInput",
|
||||
"AfterAgentInput",
|
||||
"SessionStartInput",
|
||||
"SessionEndInput",
|
||||
# Middleware types
|
||||
"InvocationContext",
|
||||
"ModelRequest",
|
||||
"ModelResponse",
|
||||
"ToolRequest",
|
||||
"ToolResponse",
|
||||
# Middleware decorators
|
||||
"before_model",
|
||||
"after_model",
|
||||
"wrap_model_call",
|
||||
"before_tool",
|
||||
"after_tool",
|
||||
"wrap_tool_call",
|
||||
# Middleware utilities
|
||||
"MiddlewareChain",
|
||||
"AsyncMiddlewareChain",
|
||||
"MiddlewareManager",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name in ("HookEvent", "HookDecision", "HookResult", "HookInput", "HookOutput"):
|
||||
from .types import HookEvent, HookDecision, HookResult, HookInput, HookOutput
|
||||
return locals()[name]
|
||||
|
||||
if name in ("HookDefinition", "CommandHook", "FunctionHook"):
|
||||
from .types import HookDefinition, CommandHook, FunctionHook
|
||||
return locals()[name]
|
||||
|
||||
if name == "HookRegistry":
|
||||
from .registry import HookRegistry
|
||||
return HookRegistry
|
||||
|
||||
if name == "HookRunner":
|
||||
from .runner import HookRunner
|
||||
return HookRunner
|
||||
|
||||
if name in ("BeforeToolInput", "AfterToolInput", "BeforeAgentInput",
|
||||
"AfterAgentInput", "SessionStartInput", "SessionEndInput"):
|
||||
from .events import (
|
||||
BeforeToolInput, AfterToolInput, BeforeAgentInput,
|
||||
AfterAgentInput, SessionStartInput, SessionEndInput
|
||||
)
|
||||
return locals()[name]
|
||||
|
||||
# Middleware types
|
||||
if name in ("InvocationContext", "ModelRequest", "ModelResponse",
|
||||
"ToolRequest", "ToolResponse"):
|
||||
from .middleware import (
|
||||
InvocationContext, ModelRequest, ModelResponse,
|
||||
ToolRequest, ToolResponse
|
||||
)
|
||||
return locals()[name]
|
||||
|
||||
# Middleware decorators
|
||||
if name in ("before_model", "after_model", "wrap_model_call",
|
||||
"before_tool", "after_tool", "wrap_tool_call"):
|
||||
from .middleware import (
|
||||
before_model, after_model, wrap_model_call,
|
||||
before_tool, after_tool, wrap_tool_call
|
||||
)
|
||||
return locals()[name]
|
||||
|
||||
# Middleware utilities
|
||||
if name in ("MiddlewareChain", "AsyncMiddlewareChain", "MiddlewareManager"):
|
||||
from .middleware import MiddlewareChain, AsyncMiddlewareChain, MiddlewareManager
|
||||
return locals()[name]
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,195 @@
|
||||
"""
|
||||
Event-specific input types for the hook system.
|
||||
|
||||
Each event type has its own input class with relevant fields.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List
|
||||
from .types import HookInput
|
||||
|
||||
|
||||
@dataclass
|
||||
class BeforeToolInput(HookInput):
|
||||
"""Input for BeforeTool hooks."""
|
||||
tool_name: str = ""
|
||||
tool_input: Dict[str, Any] = field(default_factory=dict)
|
||||
tool_description: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"tool_name": self.tool_name,
|
||||
"tool_input": self.tool_input,
|
||||
"tool_description": self.tool_description
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class AfterToolInput(HookInput):
|
||||
"""Input for AfterTool hooks."""
|
||||
tool_name: str = ""
|
||||
tool_input: Dict[str, Any] = field(default_factory=dict)
|
||||
tool_output: Any = None
|
||||
tool_error: Optional[str] = None
|
||||
execution_time_ms: float = 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"tool_name": self.tool_name,
|
||||
"tool_input": self.tool_input,
|
||||
"tool_output": str(self.tool_output) if self.tool_output else None,
|
||||
"tool_error": self.tool_error,
|
||||
"execution_time_ms": self.execution_time_ms
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class BeforeAgentInput(HookInput):
|
||||
"""Input for BeforeAgent hooks."""
|
||||
prompt: str = ""
|
||||
conversation_history: List[Dict[str, Any]] = field(default_factory=list)
|
||||
tools_available: List[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"prompt": self.prompt,
|
||||
"conversation_history_length": len(self.conversation_history),
|
||||
"tools_available": self.tools_available
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class AfterAgentInput(HookInput):
|
||||
"""Input for AfterAgent hooks."""
|
||||
prompt: str = ""
|
||||
response: str = ""
|
||||
tools_used: List[str] = field(default_factory=list)
|
||||
total_tokens: int = 0
|
||||
execution_time_ms: float = 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"prompt": self.prompt,
|
||||
"response": self.response[:500] if self.response else "", # Truncate for JSON
|
||||
"tools_used": self.tools_used,
|
||||
"total_tokens": self.total_tokens,
|
||||
"execution_time_ms": self.execution_time_ms
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionStartInput(HookInput):
|
||||
"""Input for SessionStart hooks."""
|
||||
source: str = "startup" # startup, resume, clear
|
||||
session_name: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"source": self.source,
|
||||
"session_name": self.session_name
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionEndInput(HookInput):
|
||||
"""Input for SessionEnd hooks."""
|
||||
reason: str = "exit" # exit, clear, logout, error
|
||||
total_turns: int = 0
|
||||
total_tokens: int = 0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"reason": self.reason,
|
||||
"total_turns": self.total_turns,
|
||||
"total_tokens": self.total_tokens
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class BeforeLLMInput(HookInput):
|
||||
"""Input for BeforeLLM hooks."""
|
||||
messages: List[Dict[str, Any]] = field(default_factory=list)
|
||||
model: str = ""
|
||||
temperature: Optional[float] = None
|
||||
max_tokens: Optional[int] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"messages_count": len(self.messages),
|
||||
"model": self.model,
|
||||
"temperature": self.temperature,
|
||||
"max_tokens": self.max_tokens
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class AfterLLMInput(HookInput):
|
||||
"""Input for AfterLLM hooks."""
|
||||
messages: List[Dict[str, Any]] = field(default_factory=list)
|
||||
response: str = ""
|
||||
model: str = ""
|
||||
tokens_used: int = 0
|
||||
latency_ms: float = 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"messages_count": len(self.messages),
|
||||
"response_length": len(self.response),
|
||||
"model": self.model,
|
||||
"tokens_used": self.tokens_used,
|
||||
"latency_ms": self.latency_ms
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class OnErrorInput(HookInput):
|
||||
"""Input for OnError hooks."""
|
||||
error_type: str = ""
|
||||
error_message: str = ""
|
||||
stack_trace: Optional[str] = None
|
||||
context: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"error_type": self.error_type,
|
||||
"error_message": self.error_message,
|
||||
"stack_trace": self.stack_trace,
|
||||
"context": self.context
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
@dataclass
|
||||
class OnRetryInput(HookInput):
|
||||
"""Input for OnRetry hooks."""
|
||||
retry_count: int = 0
|
||||
max_retries: int = 3
|
||||
error_message: str = ""
|
||||
operation: str = "" # tool_call, llm_request, etc.
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"retry_count": self.retry_count,
|
||||
"max_retries": self.max_retries,
|
||||
"error_message": self.error_message,
|
||||
"operation": self.operation
|
||||
})
|
||||
return base
|
||||
@@ -0,0 +1,439 @@
|
||||
"""
|
||||
Middleware System for PraisonAI Agents.
|
||||
|
||||
Provides wrap_model_call and wrap_tool_call patterns for intercepting
|
||||
and modifying agent behavior at model/tool call boundaries.
|
||||
|
||||
Zero Performance Impact:
|
||||
- All middleware is optional
|
||||
- Empty middleware chain is O(1) fast path
|
||||
- No overhead when middleware not registered
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.hooks import before_model, wrap_model_call, wrap_tool_call
|
||||
|
||||
@before_model
|
||||
def add_context(request):
|
||||
request.messages.append({"role": "system", "content": "Extra"})
|
||||
return request
|
||||
|
||||
@wrap_model_call
|
||||
def retry_on_error(request, call_next):
|
||||
for _ in range(3):
|
||||
try:
|
||||
return call_next(request)
|
||||
except Exception:
|
||||
pass
|
||||
raise RuntimeError("All retries failed")
|
||||
|
||||
agent = Agent(name="Test", hooks=[add_context, retry_on_error])
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, TypeVar, Generic
|
||||
from functools import wraps
|
||||
|
||||
|
||||
@dataclass
|
||||
class InvocationContext:
|
||||
"""Context passed through middleware chain.
|
||||
|
||||
Contains identifiers for multi-agent safety and optional metadata.
|
||||
"""
|
||||
agent_id: str
|
||||
run_id: str
|
||||
session_id: str
|
||||
tool_name: Optional[str] = None
|
||||
model_name: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelRequest:
|
||||
"""Request data for model calls."""
|
||||
messages: List[Dict[str, Any]]
|
||||
model: str
|
||||
temperature: float = 1.0
|
||||
context: Optional[InvocationContext] = None
|
||||
tools: Optional[List[Any]] = None
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelResponse:
|
||||
"""Response data from model calls."""
|
||||
content: str
|
||||
model: str
|
||||
usage: Optional[Dict[str, int]] = None
|
||||
context: Optional[InvocationContext] = None
|
||||
tool_calls: Optional[List[Any]] = None
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolRequest:
|
||||
"""Request data for tool calls."""
|
||||
tool_name: str
|
||||
arguments: Dict[str, Any]
|
||||
context: Optional[InvocationContext] = None
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResponse:
|
||||
"""Response data from tool calls."""
|
||||
tool_name: str
|
||||
result: Any
|
||||
error: Optional[str] = None
|
||||
context: Optional[InvocationContext] = None
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
# Type aliases for middleware functions
|
||||
T = TypeVar('T')
|
||||
BeforeModelFn = Callable[[ModelRequest], ModelRequest]
|
||||
AfterModelFn = Callable[[ModelResponse], ModelResponse]
|
||||
WrapModelCallFn = Callable[[ModelRequest, Callable[[ModelRequest], ModelResponse]], ModelResponse]
|
||||
BeforeToolFn = Callable[[ToolRequest], ToolRequest]
|
||||
AfterToolFn = Callable[[ToolResponse], ToolResponse]
|
||||
WrapToolCallFn = Callable[[ToolRequest, Callable[[ToolRequest], ToolResponse]], ToolResponse]
|
||||
|
||||
|
||||
class MiddlewareChain:
|
||||
"""Synchronous middleware chain executor.
|
||||
|
||||
Composes middleware functions in registration order.
|
||||
First registered middleware wraps all subsequent ones.
|
||||
"""
|
||||
|
||||
def __init__(self, middlewares: Optional[List[Callable]] = None):
|
||||
self._middlewares = middlewares or []
|
||||
|
||||
def execute(self, request: Any, final_handler: Callable[[Any], Any]) -> Any:
|
||||
"""Execute the middleware chain with a final handler.
|
||||
|
||||
Args:
|
||||
request: The request object to pass through
|
||||
final_handler: The actual operation to perform
|
||||
|
||||
Returns:
|
||||
The result from the chain
|
||||
"""
|
||||
if not self._middlewares:
|
||||
return final_handler(request)
|
||||
|
||||
# Build the chain from inside out
|
||||
# Last middleware wraps final_handler
|
||||
# First middleware is outermost
|
||||
def build_chain(index: int) -> Callable[[Any], Any]:
|
||||
if index >= len(self._middlewares):
|
||||
return final_handler
|
||||
|
||||
middleware = self._middlewares[index]
|
||||
next_handler = build_chain(index + 1)
|
||||
|
||||
def handler(req: Any) -> Any:
|
||||
return middleware(req, next_handler)
|
||||
|
||||
return handler
|
||||
|
||||
chain = build_chain(0)
|
||||
return chain(request)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._middlewares)
|
||||
|
||||
|
||||
class AsyncMiddlewareChain:
|
||||
"""Asynchronous middleware chain executor.
|
||||
|
||||
Same as MiddlewareChain but for async middleware and handlers.
|
||||
"""
|
||||
|
||||
def __init__(self, middlewares: Optional[List[Callable]] = None):
|
||||
self._middlewares = middlewares or []
|
||||
|
||||
async def execute(self, request: Any, final_handler: Callable) -> Any:
|
||||
"""Execute the async middleware chain.
|
||||
|
||||
Args:
|
||||
request: The request object to pass through
|
||||
final_handler: The actual async operation to perform
|
||||
|
||||
Returns:
|
||||
The result from the chain
|
||||
"""
|
||||
if not self._middlewares:
|
||||
return await final_handler(request)
|
||||
|
||||
# Build the chain from inside out
|
||||
def build_chain(index: int) -> Callable:
|
||||
if index >= len(self._middlewares):
|
||||
return final_handler
|
||||
|
||||
middleware = self._middlewares[index]
|
||||
next_handler = build_chain(index + 1)
|
||||
|
||||
async def handler(req: Any) -> Any:
|
||||
return await middleware(req, next_handler)
|
||||
|
||||
return handler
|
||||
|
||||
chain = build_chain(0)
|
||||
return await chain(request)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._middlewares)
|
||||
|
||||
|
||||
# Decorator functions for tagging middleware
|
||||
|
||||
def before_model(func: BeforeModelFn) -> BeforeModelFn:
|
||||
"""Decorator to mark a function as a before_model hook.
|
||||
|
||||
The function receives a ModelRequest and should return a (possibly modified) ModelRequest.
|
||||
|
||||
Example:
|
||||
@before_model
|
||||
def add_context(request):
|
||||
request.messages.append({"role": "system", "content": "Extra"})
|
||||
return request
|
||||
"""
|
||||
func._hook_type = 'before_model'
|
||||
return func
|
||||
|
||||
|
||||
def after_model(func: AfterModelFn) -> AfterModelFn:
|
||||
"""Decorator to mark a function as an after_model hook.
|
||||
|
||||
The function receives a ModelResponse and should return a (possibly modified) ModelResponse.
|
||||
|
||||
Example:
|
||||
@after_model
|
||||
def log_response(response):
|
||||
print(f"Model returned: {response.content[:50]}")
|
||||
return response
|
||||
"""
|
||||
func._hook_type = 'after_model'
|
||||
return func
|
||||
|
||||
|
||||
def wrap_model_call(func: WrapModelCallFn) -> WrapModelCallFn:
|
||||
"""Decorator to mark a function as a wrap_model_call middleware.
|
||||
|
||||
The function receives (request, call_next) and should call call_next(request)
|
||||
to continue the chain, or return early to short-circuit.
|
||||
|
||||
Example:
|
||||
@wrap_model_call
|
||||
def retry_on_error(request, call_next):
|
||||
for _ in range(3):
|
||||
try:
|
||||
return call_next(request)
|
||||
except Exception:
|
||||
pass
|
||||
raise RuntimeError("All retries failed")
|
||||
"""
|
||||
func._hook_type = 'wrap_model_call'
|
||||
return func
|
||||
|
||||
|
||||
def before_tool(func: BeforeToolFn) -> BeforeToolFn:
|
||||
"""Decorator to mark a function as a before_tool hook.
|
||||
|
||||
The function receives a ToolRequest and should return a (possibly modified) ToolRequest.
|
||||
|
||||
Example:
|
||||
@before_tool
|
||||
def validate_args(request):
|
||||
if 'dangerous' in request.arguments:
|
||||
raise ValueError("Dangerous argument detected")
|
||||
return request
|
||||
"""
|
||||
func._hook_type = 'before_tool'
|
||||
return func
|
||||
|
||||
|
||||
def after_tool(func: AfterToolFn) -> AfterToolFn:
|
||||
"""Decorator to mark a function as an after_tool hook.
|
||||
|
||||
The function receives a ToolResponse and should return a (possibly modified) ToolResponse.
|
||||
|
||||
Example:
|
||||
@after_tool
|
||||
def log_result(response):
|
||||
print(f"Tool {response.tool_name} returned: {response.result}")
|
||||
return response
|
||||
"""
|
||||
func._hook_type = 'after_tool'
|
||||
return func
|
||||
|
||||
|
||||
def wrap_tool_call(func: WrapToolCallFn) -> WrapToolCallFn:
|
||||
"""Decorator to mark a function as a wrap_tool_call middleware.
|
||||
|
||||
The function receives (request, call_next) and should call call_next(request)
|
||||
to continue the chain, or return early to short-circuit.
|
||||
|
||||
Example:
|
||||
@wrap_tool_call
|
||||
def retry_flaky_tool(request, call_next):
|
||||
last_error = None
|
||||
for _ in range(3):
|
||||
try:
|
||||
return call_next(request)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
raise last_error
|
||||
"""
|
||||
func._hook_type = 'wrap_tool_call'
|
||||
return func
|
||||
|
||||
|
||||
def get_hook_type(func: Callable) -> Optional[str]:
|
||||
"""Get the hook type of a decorated function."""
|
||||
return getattr(func, '_hook_type', None)
|
||||
|
||||
|
||||
def is_middleware(func: Callable) -> bool:
|
||||
"""Check if a function is a middleware (wrap_* type)."""
|
||||
hook_type = get_hook_type(func)
|
||||
return hook_type in ('wrap_model_call', 'wrap_tool_call')
|
||||
|
||||
|
||||
def categorize_hooks(hooks: List[Callable]) -> Dict[str, List[Callable]]:
|
||||
"""Categorize hooks by their type.
|
||||
|
||||
Returns:
|
||||
Dict with keys: before_model, after_model, wrap_model_call,
|
||||
before_tool, after_tool, wrap_tool_call
|
||||
"""
|
||||
categories = {
|
||||
'before_model': [],
|
||||
'after_model': [],
|
||||
'wrap_model_call': [],
|
||||
'before_tool': [],
|
||||
'after_tool': [],
|
||||
'wrap_tool_call': [],
|
||||
}
|
||||
|
||||
for hook in hooks or []:
|
||||
hook_type = get_hook_type(hook)
|
||||
if hook_type and hook_type in categories:
|
||||
categories[hook_type].append(hook)
|
||||
|
||||
return categories
|
||||
|
||||
|
||||
class MiddlewareManager:
|
||||
"""Manages middleware for an agent.
|
||||
|
||||
Provides methods to execute before/after hooks and middleware chains
|
||||
for both model and tool calls.
|
||||
"""
|
||||
|
||||
def __init__(self, hooks: Optional[List[Callable]] = None):
|
||||
self._hooks = hooks or []
|
||||
self._categorized = categorize_hooks(self._hooks)
|
||||
|
||||
# Pre-build chains for efficiency
|
||||
self._model_chain = MiddlewareChain(self._categorized['wrap_model_call'])
|
||||
self._tool_chain = MiddlewareChain(self._categorized['wrap_tool_call'])
|
||||
|
||||
@property
|
||||
def has_model_hooks(self) -> bool:
|
||||
"""Check if any model hooks are registered."""
|
||||
return bool(
|
||||
self._categorized['before_model'] or
|
||||
self._categorized['after_model'] or
|
||||
self._categorized['wrap_model_call']
|
||||
)
|
||||
|
||||
@property
|
||||
def has_tool_hooks(self) -> bool:
|
||||
"""Check if any tool hooks are registered."""
|
||||
return bool(
|
||||
self._categorized['before_tool'] or
|
||||
self._categorized['after_tool'] or
|
||||
self._categorized['wrap_tool_call']
|
||||
)
|
||||
|
||||
def run_before_model(self, request: ModelRequest) -> ModelRequest:
|
||||
"""Run all before_model hooks."""
|
||||
for hook in self._categorized['before_model']:
|
||||
request = hook(request)
|
||||
return request
|
||||
|
||||
def run_after_model(self, response: ModelResponse) -> ModelResponse:
|
||||
"""Run all after_model hooks (in reverse order)."""
|
||||
for hook in reversed(self._categorized['after_model']):
|
||||
response = hook(response)
|
||||
return response
|
||||
|
||||
def run_before_tool(self, request: ToolRequest) -> ToolRequest:
|
||||
"""Run all before_tool hooks."""
|
||||
for hook in self._categorized['before_tool']:
|
||||
request = hook(request)
|
||||
return request
|
||||
|
||||
def run_after_tool(self, response: ToolResponse) -> ToolResponse:
|
||||
"""Run all after_tool hooks (in reverse order)."""
|
||||
for hook in reversed(self._categorized['after_tool']):
|
||||
response = hook(response)
|
||||
return response
|
||||
|
||||
def execute_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
model_fn: Callable[[ModelRequest], ModelResponse]
|
||||
) -> ModelResponse:
|
||||
"""Execute a model call with all hooks and middleware.
|
||||
|
||||
Order:
|
||||
1. before_model hooks (in order)
|
||||
2. wrap_model_call middleware chain
|
||||
3. actual model call
|
||||
4. after_model hooks (in reverse order)
|
||||
"""
|
||||
# Fast path: no hooks
|
||||
if not self.has_model_hooks:
|
||||
return model_fn(request)
|
||||
|
||||
# Run before hooks
|
||||
request = self.run_before_model(request)
|
||||
|
||||
# Build wrapped handler that includes after hooks
|
||||
def wrapped_model_fn(req: ModelRequest) -> ModelResponse:
|
||||
response = model_fn(req)
|
||||
return self.run_after_model(response)
|
||||
|
||||
# Execute through middleware chain
|
||||
return self._model_chain.execute(request, wrapped_model_fn)
|
||||
|
||||
def execute_tool_call(
|
||||
self,
|
||||
request: ToolRequest,
|
||||
tool_fn: Callable[[ToolRequest], ToolResponse]
|
||||
) -> ToolResponse:
|
||||
"""Execute a tool call with all hooks and middleware.
|
||||
|
||||
Order:
|
||||
1. before_tool hooks (in order)
|
||||
2. wrap_tool_call middleware chain
|
||||
3. actual tool call
|
||||
4. after_tool hooks (in reverse order)
|
||||
"""
|
||||
# Fast path: no hooks
|
||||
if not self.has_tool_hooks:
|
||||
return tool_fn(request)
|
||||
|
||||
# Run before hooks
|
||||
request = self.run_before_tool(request)
|
||||
|
||||
# Build wrapped handler that includes after hooks
|
||||
def wrapped_tool_fn(req: ToolRequest) -> ToolResponse:
|
||||
response = tool_fn(req)
|
||||
return self.run_after_tool(response)
|
||||
|
||||
# Execute through middleware chain
|
||||
return self._tool_chain.execute(request, wrapped_tool_fn)
|
||||
@@ -0,0 +1,336 @@
|
||||
"""
|
||||
Hook Registry for PraisonAI Agents.
|
||||
|
||||
Manages registration and lookup of hooks for different events.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Callable, Union
|
||||
from functools import wraps
|
||||
|
||||
from .types import (
|
||||
HookEvent, HookDefinition, CommandHook, FunctionHook,
|
||||
HookInput, HookResult
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HookRegistry:
|
||||
"""
|
||||
Registry for managing hooks.
|
||||
|
||||
Provides methods to register, unregister, and lookup hooks
|
||||
for different events.
|
||||
|
||||
Example:
|
||||
registry = HookRegistry()
|
||||
|
||||
# Register a function hook using decorator
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def validate_tool(event_data):
|
||||
if event_data.tool_name == "dangerous":
|
||||
return HookResult.deny("Tool is dangerous")
|
||||
return HookResult.allow()
|
||||
|
||||
# Register a command hook
|
||||
registry.register_command(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
command="python /path/to/validator.py",
|
||||
matcher="write_*"
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(self, enabled: bool = True):
|
||||
"""
|
||||
Initialize the hook registry.
|
||||
|
||||
Args:
|
||||
enabled: Whether hooks are enabled globally
|
||||
"""
|
||||
self._hooks: Dict[HookEvent, List[HookDefinition]] = {
|
||||
event: [] for event in HookEvent
|
||||
}
|
||||
self._enabled = enabled
|
||||
self._global_timeout = 60.0
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Check if hooks are enabled."""
|
||||
return self._enabled
|
||||
|
||||
@enabled.setter
|
||||
def enabled(self, value: bool):
|
||||
"""Enable or disable hooks."""
|
||||
self._enabled = value
|
||||
|
||||
def register(self, hook: HookDefinition) -> str:
|
||||
"""
|
||||
Register a hook definition.
|
||||
|
||||
Args:
|
||||
hook: The hook definition to register
|
||||
|
||||
Returns:
|
||||
The hook ID
|
||||
"""
|
||||
self._hooks[hook.event].append(hook)
|
||||
logger.debug(f"Registered hook '{hook.name}' for event '{hook.event.value}'")
|
||||
return hook.id
|
||||
|
||||
def register_function(
|
||||
self,
|
||||
event: HookEvent,
|
||||
func: Callable[[HookInput], HookResult],
|
||||
matcher: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0,
|
||||
is_async: bool = False
|
||||
) -> str:
|
||||
"""
|
||||
Register a Python function as a hook.
|
||||
|
||||
Args:
|
||||
event: The event to hook
|
||||
func: The function to call
|
||||
matcher: Optional regex pattern to match targets
|
||||
name: Optional name for the hook
|
||||
description: Optional description
|
||||
sequential: Whether to run sequentially
|
||||
timeout: Timeout in seconds
|
||||
is_async: Whether the function is async
|
||||
|
||||
Returns:
|
||||
The hook ID
|
||||
"""
|
||||
hook = FunctionHook(
|
||||
event=event,
|
||||
func=func,
|
||||
matcher=matcher,
|
||||
name=name or func.__name__,
|
||||
description=description,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
is_async=is_async
|
||||
)
|
||||
return self.register(hook)
|
||||
|
||||
def register_command(
|
||||
self,
|
||||
event: HookEvent,
|
||||
command: str,
|
||||
matcher: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
shell: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Register a shell command as a hook.
|
||||
|
||||
Args:
|
||||
event: The event to hook
|
||||
command: The shell command to execute
|
||||
matcher: Optional regex pattern to match targets
|
||||
name: Optional name for the hook
|
||||
description: Optional description
|
||||
sequential: Whether to run sequentially
|
||||
timeout: Timeout in seconds
|
||||
env: Additional environment variables
|
||||
shell: Whether to run in shell mode
|
||||
|
||||
Returns:
|
||||
The hook ID
|
||||
"""
|
||||
hook = CommandHook(
|
||||
event=event,
|
||||
command=command,
|
||||
matcher=matcher,
|
||||
name=name,
|
||||
description=description,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
env=env or {},
|
||||
shell=shell
|
||||
)
|
||||
return self.register(hook)
|
||||
|
||||
def on(
|
||||
self,
|
||||
event: HookEvent,
|
||||
matcher: Optional[str] = None,
|
||||
sequential: bool = False,
|
||||
timeout: float = 60.0
|
||||
) -> Callable:
|
||||
"""
|
||||
Decorator to register a function as a hook.
|
||||
|
||||
Args:
|
||||
event: The event to hook
|
||||
matcher: Optional regex pattern to match targets
|
||||
sequential: Whether to run sequentially
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
Decorator function
|
||||
|
||||
Example:
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def my_hook(event_data):
|
||||
return HookResult.allow()
|
||||
"""
|
||||
def decorator(func: Callable[[HookInput], HookResult]) -> Callable:
|
||||
import asyncio
|
||||
is_async = asyncio.iscoroutinefunction(func)
|
||||
|
||||
self.register_function(
|
||||
event=event,
|
||||
func=func,
|
||||
matcher=matcher,
|
||||
sequential=sequential,
|
||||
timeout=timeout,
|
||||
is_async=is_async
|
||||
)
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
def unregister(self, hook_id: str) -> bool:
|
||||
"""
|
||||
Unregister a hook by ID.
|
||||
|
||||
Args:
|
||||
hook_id: The hook ID to unregister
|
||||
|
||||
Returns:
|
||||
True if found and removed, False otherwise
|
||||
"""
|
||||
for event in HookEvent:
|
||||
for i, hook in enumerate(self._hooks[event]):
|
||||
if hook.id == hook_id:
|
||||
self._hooks[event].pop(i)
|
||||
logger.debug(f"Unregistered hook '{hook_id}'")
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_hooks(
|
||||
self,
|
||||
event: HookEvent,
|
||||
target: Optional[str] = None
|
||||
) -> List[HookDefinition]:
|
||||
"""
|
||||
Get all hooks for an event, optionally filtered by target.
|
||||
|
||||
Args:
|
||||
event: The event to get hooks for
|
||||
target: Optional target to filter by (e.g., tool name)
|
||||
|
||||
Returns:
|
||||
List of matching hook definitions
|
||||
"""
|
||||
if not self._enabled:
|
||||
return []
|
||||
|
||||
hooks = self._hooks.get(event, [])
|
||||
|
||||
if target is None:
|
||||
return [h for h in hooks if h.enabled]
|
||||
|
||||
return [h for h in hooks if h.enabled and h.matches(target)]
|
||||
|
||||
def has_hooks(self, event: HookEvent) -> bool:
|
||||
"""Check if there are any hooks registered for an event."""
|
||||
return bool(self.get_hooks(event))
|
||||
|
||||
def clear(self, event: Optional[HookEvent] = None):
|
||||
"""
|
||||
Clear all hooks or hooks for a specific event.
|
||||
|
||||
Args:
|
||||
event: Optional event to clear hooks for
|
||||
"""
|
||||
if event is None:
|
||||
for e in HookEvent:
|
||||
self._hooks[e] = []
|
||||
else:
|
||||
self._hooks[event] = []
|
||||
|
||||
def list_hooks(self) -> Dict[str, List[Dict]]:
|
||||
"""
|
||||
List all registered hooks.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping event names to hook info
|
||||
"""
|
||||
result = {}
|
||||
for event in HookEvent:
|
||||
hooks = self._hooks[event]
|
||||
if hooks:
|
||||
result[event.value] = [
|
||||
{
|
||||
"id": h.id,
|
||||
"name": h.name,
|
||||
"type": "command" if isinstance(h, CommandHook) else "function",
|
||||
"matcher": h.matcher,
|
||||
"enabled": h.enabled,
|
||||
"sequential": h.sequential
|
||||
}
|
||||
for h in hooks
|
||||
]
|
||||
return result
|
||||
|
||||
def enable_hook(self, hook_id: str) -> bool:
|
||||
"""Enable a specific hook."""
|
||||
for event in HookEvent:
|
||||
for hook in self._hooks[event]:
|
||||
if hook.id == hook_id:
|
||||
hook.enabled = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable_hook(self, hook_id: str) -> bool:
|
||||
"""Disable a specific hook."""
|
||||
for event in HookEvent:
|
||||
for hook in self._hooks[event]:
|
||||
if hook.id == hook_id:
|
||||
hook.enabled = False
|
||||
return True
|
||||
return False
|
||||
|
||||
def set_global_timeout(self, timeout: float):
|
||||
"""Set the global timeout for all hooks."""
|
||||
self._global_timeout = timeout
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Get total number of registered hooks."""
|
||||
return sum(len(hooks) for hooks in self._hooks.values())
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"HookRegistry(enabled={self._enabled}, hooks={len(self)})"
|
||||
|
||||
|
||||
# Global default registry
|
||||
_default_registry: Optional[HookRegistry] = None
|
||||
|
||||
|
||||
def get_default_registry() -> HookRegistry:
|
||||
"""Get the default global hook registry."""
|
||||
global _default_registry
|
||||
if _default_registry is None:
|
||||
_default_registry = HookRegistry()
|
||||
return _default_registry
|
||||
|
||||
|
||||
def set_default_registry(registry: HookRegistry):
|
||||
"""Set the default global hook registry."""
|
||||
global _default_registry
|
||||
_default_registry = registry
|
||||
@@ -0,0 +1,450 @@
|
||||
"""
|
||||
Hook Runner for PraisonAI Agents.
|
||||
|
||||
Executes hooks registered in the registry, supporting both
|
||||
Python functions and shell commands.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import asyncio
|
||||
import logging
|
||||
import subprocess
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
from .types import (
|
||||
HookEvent, HookDefinition, CommandHook, FunctionHook,
|
||||
HookInput, HookResult, HookExecutionResult
|
||||
)
|
||||
from .registry import HookRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Exit codes for command hooks
|
||||
EXIT_CODE_SUCCESS = 0
|
||||
EXIT_CODE_BLOCKING_ERROR = 2
|
||||
EXIT_CODE_NON_BLOCKING_ERROR = 1
|
||||
|
||||
|
||||
class HookRunner:
|
||||
"""
|
||||
Executes hooks from a registry.
|
||||
|
||||
Supports:
|
||||
- Python function hooks (sync and async)
|
||||
- Shell command hooks
|
||||
- Sequential and parallel execution
|
||||
- Timeout handling
|
||||
- Result aggregation
|
||||
|
||||
Example:
|
||||
registry = HookRegistry()
|
||||
runner = HookRunner(registry)
|
||||
|
||||
# Execute hooks for an event
|
||||
results = await runner.execute(
|
||||
event=HookEvent.BEFORE_TOOL,
|
||||
input_data=BeforeToolInput(
|
||||
tool_name="write_file",
|
||||
tool_input={"path": "/tmp/test.txt"}
|
||||
)
|
||||
)
|
||||
|
||||
# Check if any hook blocked
|
||||
if runner.is_blocked(results):
|
||||
print("Tool execution blocked by hook")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
registry: Optional[HookRegistry] = None,
|
||||
default_timeout: float = 60.0,
|
||||
cwd: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
Initialize the hook runner.
|
||||
|
||||
Args:
|
||||
registry: Hook registry to use
|
||||
default_timeout: Default timeout for hooks
|
||||
cwd: Working directory for command hooks
|
||||
"""
|
||||
self._registry = registry or HookRegistry()
|
||||
self._default_timeout = default_timeout
|
||||
self._cwd = cwd or os.getcwd()
|
||||
|
||||
@property
|
||||
def registry(self) -> HookRegistry:
|
||||
"""Get the hook registry."""
|
||||
return self._registry
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
target: Optional[str] = None
|
||||
) -> List[HookExecutionResult]:
|
||||
"""
|
||||
Execute all hooks for an event.
|
||||
|
||||
Args:
|
||||
event: The event to execute hooks for
|
||||
input_data: Input data for the hooks
|
||||
target: Optional target to filter hooks (e.g., tool name)
|
||||
|
||||
Returns:
|
||||
List of execution results
|
||||
"""
|
||||
hooks = self._registry.get_hooks(event, target)
|
||||
|
||||
if not hooks:
|
||||
return []
|
||||
|
||||
# Separate sequential and parallel hooks
|
||||
sequential_hooks = [h for h in hooks if h.sequential]
|
||||
parallel_hooks = [h for h in hooks if not h.sequential]
|
||||
|
||||
results = []
|
||||
|
||||
# Execute parallel hooks first
|
||||
if parallel_hooks:
|
||||
parallel_results = await self._execute_parallel(
|
||||
parallel_hooks, event, input_data
|
||||
)
|
||||
results.extend(parallel_results)
|
||||
|
||||
# Execute sequential hooks
|
||||
if sequential_hooks:
|
||||
sequential_results = await self._execute_sequential(
|
||||
sequential_hooks, event, input_data
|
||||
)
|
||||
results.extend(sequential_results)
|
||||
|
||||
return results
|
||||
|
||||
def execute_sync(
|
||||
self,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
target: Optional[str] = None
|
||||
) -> List[HookExecutionResult]:
|
||||
"""
|
||||
Synchronous version of execute.
|
||||
|
||||
Args:
|
||||
event: The event to execute hooks for
|
||||
input_data: Input data for the hooks
|
||||
target: Optional target to filter hooks
|
||||
|
||||
Returns:
|
||||
List of execution results
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# If we're in an async context, create a task
|
||||
future = asyncio.ensure_future(self.execute(event, input_data, target))
|
||||
return loop.run_until_complete(future)
|
||||
except RuntimeError:
|
||||
# No event loop running, create one
|
||||
return asyncio.run(self.execute(event, input_data, target))
|
||||
|
||||
async def _execute_parallel(
|
||||
self,
|
||||
hooks: List[HookDefinition],
|
||||
event: HookEvent,
|
||||
input_data: HookInput
|
||||
) -> List[HookExecutionResult]:
|
||||
"""Execute hooks in parallel."""
|
||||
tasks = [
|
||||
self._execute_single(hook, event, input_data)
|
||||
for hook in hooks
|
||||
]
|
||||
return await asyncio.gather(*tasks)
|
||||
|
||||
async def _execute_sequential(
|
||||
self,
|
||||
hooks: List[HookDefinition],
|
||||
event: HookEvent,
|
||||
input_data: HookInput
|
||||
) -> List[HookExecutionResult]:
|
||||
"""Execute hooks sequentially, passing modified input to next hook."""
|
||||
results = []
|
||||
current_input = input_data
|
||||
|
||||
for hook in hooks:
|
||||
result = await self._execute_single(hook, event, current_input)
|
||||
results.append(result)
|
||||
|
||||
# If hook modified the input, use it for next hook
|
||||
if result.success and result.output and result.output.modified_input:
|
||||
# Update input with modifications
|
||||
for key, value in result.output.modified_input.items():
|
||||
if hasattr(current_input, key):
|
||||
setattr(current_input, key, value)
|
||||
|
||||
# Stop if hook blocked
|
||||
if result.output and result.output.is_denied():
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
async def _execute_single(
|
||||
self,
|
||||
hook: HookDefinition,
|
||||
event: HookEvent,
|
||||
input_data: HookInput
|
||||
) -> HookExecutionResult:
|
||||
"""Execute a single hook."""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
if isinstance(hook, CommandHook):
|
||||
return await self._execute_command_hook(hook, event, input_data, start_time)
|
||||
elif isinstance(hook, FunctionHook):
|
||||
return await self._execute_function_hook(hook, event, input_data, start_time)
|
||||
else:
|
||||
raise ValueError(f"Unknown hook type: {type(hook)}")
|
||||
except Exception as e:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
logger.warning(f"Hook '{hook.name}' failed: {e}")
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=str(e),
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
async def _execute_command_hook(
|
||||
self,
|
||||
hook: CommandHook,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
start_time: float
|
||||
) -> HookExecutionResult:
|
||||
"""Execute a command hook."""
|
||||
timeout = hook.timeout or self._default_timeout
|
||||
|
||||
# Prepare environment
|
||||
env = os.environ.copy()
|
||||
env.update(hook.env)
|
||||
env["PRAISON_PROJECT_DIR"] = self._cwd
|
||||
env["PRAISON_EVENT"] = event.value
|
||||
|
||||
# Expand command variables
|
||||
command = self._expand_command(hook.command, input_data)
|
||||
|
||||
try:
|
||||
# Create subprocess
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
cwd=self._cwd
|
||||
)
|
||||
|
||||
# Send input as JSON
|
||||
input_json = json.dumps(input_data.to_dict()).encode()
|
||||
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
process.communicate(input=input_json),
|
||||
timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=f"Hook timed out after {timeout}s",
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
duration = (time.time() - start_time) * 1000
|
||||
exit_code = process.returncode
|
||||
stdout_str = stdout.decode() if stdout else ""
|
||||
stderr_str = stderr.decode() if stderr else ""
|
||||
|
||||
# Parse output
|
||||
output = self._parse_command_output(stdout_str, stderr_str, exit_code)
|
||||
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=exit_code == EXIT_CODE_SUCCESS,
|
||||
output=output,
|
||||
stdout=stdout_str,
|
||||
stderr=stderr_str,
|
||||
exit_code=exit_code,
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=str(e),
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
async def _execute_function_hook(
|
||||
self,
|
||||
hook: FunctionHook,
|
||||
event: HookEvent,
|
||||
input_data: HookInput,
|
||||
start_time: float
|
||||
) -> HookExecutionResult:
|
||||
"""Execute a function hook."""
|
||||
if hook.func is None:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error="Hook function is None",
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
timeout = hook.timeout or self._default_timeout
|
||||
|
||||
try:
|
||||
if hook.is_async:
|
||||
# Execute async function with timeout
|
||||
result = await asyncio.wait_for(
|
||||
hook.func(input_data),
|
||||
timeout=timeout
|
||||
)
|
||||
else:
|
||||
# Execute sync function in thread pool with timeout
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await asyncio.wait_for(
|
||||
loop.run_in_executor(None, hook.func, input_data),
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
duration = (time.time() - start_time) * 1000
|
||||
|
||||
# Normalize result
|
||||
if result is None:
|
||||
result = HookResult.allow()
|
||||
elif not isinstance(result, HookResult):
|
||||
result = HookResult.allow()
|
||||
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=True,
|
||||
output=result,
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=f"Hook timed out after {timeout}s",
|
||||
duration_ms=duration
|
||||
)
|
||||
except Exception as e:
|
||||
duration = (time.time() - start_time) * 1000
|
||||
return HookExecutionResult(
|
||||
hook_id=hook.id,
|
||||
hook_name=hook.name or "unknown",
|
||||
event=event,
|
||||
success=False,
|
||||
error=str(e),
|
||||
duration_ms=duration
|
||||
)
|
||||
|
||||
def _expand_command(self, command: str, input_data: HookInput) -> str:
|
||||
"""Expand variables in command string."""
|
||||
return command.replace(
|
||||
"$PRAISON_PROJECT_DIR", self._cwd
|
||||
).replace(
|
||||
"$PRAISON_CWD", input_data.cwd
|
||||
)
|
||||
|
||||
def _parse_command_output(
|
||||
self,
|
||||
stdout: str,
|
||||
stderr: str,
|
||||
exit_code: int
|
||||
) -> HookResult:
|
||||
"""Parse command output into HookResult."""
|
||||
# Try to parse JSON from stdout
|
||||
if stdout.strip():
|
||||
try:
|
||||
data = json.loads(stdout.strip())
|
||||
|
||||
# Handle double-encoded JSON
|
||||
if isinstance(data, str):
|
||||
data = json.loads(data)
|
||||
|
||||
return HookResult(
|
||||
decision=data.get("decision", "allow"),
|
||||
reason=data.get("reason"),
|
||||
modified_input=data.get("modified_input"),
|
||||
additional_context=data.get("additional_context"),
|
||||
suppress_output=data.get("suppress_output", False)
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Convert exit code to decision
|
||||
if exit_code == EXIT_CODE_SUCCESS:
|
||||
return HookResult(
|
||||
decision="allow",
|
||||
additional_context=stdout.strip() if stdout.strip() else None
|
||||
)
|
||||
elif exit_code == EXIT_CODE_BLOCKING_ERROR:
|
||||
return HookResult(
|
||||
decision="deny",
|
||||
reason=stderr.strip() or stdout.strip() or "Blocked by hook"
|
||||
)
|
||||
else:
|
||||
return HookResult(
|
||||
decision="allow",
|
||||
additional_context=f"Warning: {stderr.strip() or stdout.strip()}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_blocked(results: List[HookExecutionResult]) -> bool:
|
||||
"""Check if any hook blocked execution."""
|
||||
for result in results:
|
||||
if result.output and result.output.is_denied():
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def get_blocking_reason(results: List[HookExecutionResult]) -> Optional[str]:
|
||||
"""Get the reason for blocking from results."""
|
||||
for result in results:
|
||||
if result.output and result.output.is_denied():
|
||||
return result.output.reason
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def aggregate_context(results: List[HookExecutionResult]) -> Optional[str]:
|
||||
"""Aggregate additional context from all results."""
|
||||
contexts = []
|
||||
for result in results:
|
||||
if result.output and result.output.additional_context:
|
||||
contexts.append(result.output.additional_context)
|
||||
return "\n".join(contexts) if contexts else None
|
||||
@@ -0,0 +1,197 @@
|
||||
"""
|
||||
Hook Types for PraisonAI Agents.
|
||||
|
||||
Defines the core types, enums, and dataclasses for the hook system.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List, Callable, Union, Literal
|
||||
|
||||
|
||||
class HookEvent(str, Enum):
|
||||
"""Event names for the hook system."""
|
||||
BEFORE_TOOL = "before_tool"
|
||||
AFTER_TOOL = "after_tool"
|
||||
BEFORE_AGENT = "before_agent"
|
||||
AFTER_AGENT = "after_agent"
|
||||
SESSION_START = "session_start"
|
||||
SESSION_END = "session_end"
|
||||
BEFORE_LLM = "before_llm"
|
||||
AFTER_LLM = "after_llm"
|
||||
ON_ERROR = "on_error"
|
||||
ON_RETRY = "on_retry"
|
||||
|
||||
|
||||
# Decision types for hook outputs
|
||||
HookDecision = Literal["allow", "deny", "block", "ask", None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookInput:
|
||||
"""Base hook input - common fields for all events."""
|
||||
session_id: str
|
||||
cwd: str
|
||||
event_name: str
|
||||
timestamp: str
|
||||
agent_name: Optional[str] = None
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"cwd": self.cwd,
|
||||
"event_name": self.event_name,
|
||||
"timestamp": self.timestamp,
|
||||
"agent_name": self.agent_name,
|
||||
**self.extra
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookOutput:
|
||||
"""Base hook output - common fields for all events."""
|
||||
proceed: bool = True
|
||||
stop_reason: Optional[str] = None
|
||||
suppress_output: bool = False
|
||||
system_message: Optional[str] = None
|
||||
decision: HookDecision = None
|
||||
reason: Optional[str] = None
|
||||
modified_data: Optional[Dict[str, Any]] = None
|
||||
|
||||
def is_blocking(self) -> bool:
|
||||
"""Check if this output represents a blocking decision."""
|
||||
return self.decision in ("block", "deny")
|
||||
|
||||
def should_stop(self) -> bool:
|
||||
"""Check if execution should stop."""
|
||||
return not self.proceed
|
||||
|
||||
def get_reason(self) -> str:
|
||||
"""Get the effective reason for blocking or stopping."""
|
||||
return self.reason or self.stop_reason or "No reason provided"
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookResult:
|
||||
"""Result from a hook execution."""
|
||||
decision: HookDecision = "allow"
|
||||
reason: Optional[str] = None
|
||||
modified_input: Optional[Dict[str, Any]] = None
|
||||
additional_context: Optional[str] = None
|
||||
suppress_output: bool = False
|
||||
|
||||
@classmethod
|
||||
def allow(cls, reason: Optional[str] = None) -> "HookResult":
|
||||
"""Create an allow result."""
|
||||
return cls(decision="allow", reason=reason)
|
||||
|
||||
@classmethod
|
||||
def deny(cls, reason: str) -> "HookResult":
|
||||
"""Create a deny result."""
|
||||
return cls(decision="deny", reason=reason)
|
||||
|
||||
@classmethod
|
||||
def block(cls, reason: str) -> "HookResult":
|
||||
"""Create a block result."""
|
||||
return cls(decision="block", reason=reason)
|
||||
|
||||
@classmethod
|
||||
def ask(cls, reason: str) -> "HookResult":
|
||||
"""Create an ask result (requires user confirmation)."""
|
||||
return cls(decision="ask", reason=reason)
|
||||
|
||||
def is_allowed(self) -> bool:
|
||||
"""Check if the result allows execution."""
|
||||
return self.decision in ("allow", None)
|
||||
|
||||
def is_denied(self) -> bool:
|
||||
"""Check if the result denies execution."""
|
||||
return self.decision in ("deny", "block")
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookDefinition:
|
||||
"""Hook definition with matcher and configuration."""
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
|
||||
event: HookEvent = HookEvent.BEFORE_TOOL
|
||||
matcher: Optional[str] = None # Regex pattern to match tool names
|
||||
sequential: bool = False # Execute hooks sequentially vs parallel
|
||||
enabled: bool = True
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
timeout: float = 60.0 # Timeout in seconds
|
||||
|
||||
def matches(self, target: str) -> bool:
|
||||
"""Check if this hook matches the target (tool name, etc.)."""
|
||||
if self.matcher is None:
|
||||
return True
|
||||
|
||||
import re
|
||||
try:
|
||||
return bool(re.match(self.matcher, target))
|
||||
except re.error:
|
||||
# Invalid regex, fall back to simple string match
|
||||
return self.matcher in target
|
||||
|
||||
|
||||
@dataclass
|
||||
class CommandHook(HookDefinition):
|
||||
"""Hook that executes a shell command."""
|
||||
command: str = ""
|
||||
shell: bool = True
|
||||
env: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.name:
|
||||
self.name = f"command_hook_{self.id}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionHook(HookDefinition):
|
||||
"""Hook that executes a Python function."""
|
||||
func: Optional[Callable[[HookInput], HookResult]] = None
|
||||
is_async: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.name and self.func:
|
||||
self.name = self.func.__name__
|
||||
elif not self.name:
|
||||
self.name = f"function_hook_{self.id}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookExecutionResult:
|
||||
"""Result of executing a single hook."""
|
||||
hook_id: str
|
||||
hook_name: str
|
||||
event: HookEvent
|
||||
success: bool
|
||||
output: Optional[HookResult] = None
|
||||
stdout: Optional[str] = None
|
||||
stderr: Optional[str] = None
|
||||
exit_code: Optional[int] = None
|
||||
duration_ms: float = 0.0
|
||||
error: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"hook_id": self.hook_id,
|
||||
"hook_name": self.hook_name,
|
||||
"event": self.event.value,
|
||||
"success": self.success,
|
||||
"output": self.output.__dict__ if self.output else None,
|
||||
"stdout": self.stdout,
|
||||
"stderr": self.stderr,
|
||||
"exit_code": self.exit_code,
|
||||
"duration_ms": self.duration_ms,
|
||||
"error": self.error
|
||||
}
|
||||
|
||||
|
||||
# Type alias for hook functions
|
||||
HookFunction = Callable[[HookInput], Union[HookResult, None]]
|
||||
AsyncHookFunction = Callable[[HookInput], Any] # Returns Awaitable[HookResult]
|
||||
@@ -1,8 +1,130 @@
|
||||
"""
|
||||
PraisonAI Knowledge - Advanced knowledge management system with configurable features
|
||||
PraisonAI Knowledge - Advanced knowledge management system with configurable features.
|
||||
|
||||
This module provides:
|
||||
- Document readers (ReaderProtocol, ReaderRegistry)
|
||||
- Vector stores (VectorStoreProtocol, VectorStoreRegistry)
|
||||
- Retrieval strategies (RetrieverProtocol, RetrievalStrategy)
|
||||
- Rerankers (RerankerProtocol, RerankerRegistry)
|
||||
- Index types (IndexProtocol, IndexType)
|
||||
- Query engines (QueryEngineProtocol, QueryMode)
|
||||
"""
|
||||
|
||||
# Core Knowledge class (always available)
|
||||
from praisonaiagents.knowledge.knowledge import Knowledge
|
||||
from praisonaiagents.knowledge.chunking import Chunking
|
||||
|
||||
__all__ = ["Knowledge", "Chunking"]
|
||||
# Lazy loading for protocols and registries to avoid import overhead
|
||||
_LAZY_IMPORTS = {
|
||||
# Readers
|
||||
"Document": ("praisonaiagents.knowledge.readers", "Document"),
|
||||
"ReaderProtocol": ("praisonaiagents.knowledge.readers", "ReaderProtocol"),
|
||||
"ReaderRegistry": ("praisonaiagents.knowledge.readers", "ReaderRegistry"),
|
||||
"get_reader_registry": ("praisonaiagents.knowledge.readers", "get_reader_registry"),
|
||||
"detect_source_kind": ("praisonaiagents.knowledge.readers", "detect_source_kind"),
|
||||
|
||||
# Vector stores
|
||||
"VectorRecord": ("praisonaiagents.knowledge.vector_store", "VectorRecord"),
|
||||
"VectorStoreProtocol": ("praisonaiagents.knowledge.vector_store", "VectorStoreProtocol"),
|
||||
"VectorStoreRegistry": ("praisonaiagents.knowledge.vector_store", "VectorStoreRegistry"),
|
||||
"get_vector_store_registry": ("praisonaiagents.knowledge.vector_store", "get_vector_store_registry"),
|
||||
"InMemoryVectorStore": ("praisonaiagents.knowledge.vector_store", "InMemoryVectorStore"),
|
||||
|
||||
# Retrieval
|
||||
"RetrievalResult": ("praisonaiagents.knowledge.retrieval", "RetrievalResult"),
|
||||
"RetrievalStrategy": ("praisonaiagents.knowledge.retrieval", "RetrievalStrategy"),
|
||||
"RetrieverProtocol": ("praisonaiagents.knowledge.retrieval", "RetrieverProtocol"),
|
||||
"RetrieverRegistry": ("praisonaiagents.knowledge.retrieval", "RetrieverRegistry"),
|
||||
"get_retriever_registry": ("praisonaiagents.knowledge.retrieval", "get_retriever_registry"),
|
||||
"reciprocal_rank_fusion": ("praisonaiagents.knowledge.retrieval", "reciprocal_rank_fusion"),
|
||||
"merge_adjacent_chunks": ("praisonaiagents.knowledge.retrieval", "merge_adjacent_chunks"),
|
||||
|
||||
# Rerankers
|
||||
"RerankResult": ("praisonaiagents.knowledge.rerankers", "RerankResult"),
|
||||
"RerankerProtocol": ("praisonaiagents.knowledge.rerankers", "RerankerProtocol"),
|
||||
"RerankerRegistry": ("praisonaiagents.knowledge.rerankers", "RerankerRegistry"),
|
||||
"get_reranker_registry": ("praisonaiagents.knowledge.rerankers", "get_reranker_registry"),
|
||||
"SimpleReranker": ("praisonaiagents.knowledge.rerankers", "SimpleReranker"),
|
||||
|
||||
# Index
|
||||
"IndexType": ("praisonaiagents.knowledge.index", "IndexType"),
|
||||
"IndexStats": ("praisonaiagents.knowledge.index", "IndexStats"),
|
||||
"IndexProtocol": ("praisonaiagents.knowledge.index", "IndexProtocol"),
|
||||
"IndexRegistry": ("praisonaiagents.knowledge.index", "IndexRegistry"),
|
||||
"get_index_registry": ("praisonaiagents.knowledge.index", "get_index_registry"),
|
||||
"KeywordIndex": ("praisonaiagents.knowledge.index", "KeywordIndex"),
|
||||
|
||||
# Query engine
|
||||
"QueryMode": ("praisonaiagents.knowledge.query_engine", "QueryMode"),
|
||||
"QueryResult": ("praisonaiagents.knowledge.query_engine", "QueryResult"),
|
||||
"QueryEngineProtocol": ("praisonaiagents.knowledge.query_engine", "QueryEngineProtocol"),
|
||||
"QueryEngineRegistry": ("praisonaiagents.knowledge.query_engine", "QueryEngineRegistry"),
|
||||
"get_query_engine_registry": ("praisonaiagents.knowledge.query_engine", "get_query_engine_registry"),
|
||||
"decompose_question": ("praisonaiagents.knowledge.query_engine", "decompose_question"),
|
||||
"SimpleQueryEngine": ("praisonaiagents.knowledge.query_engine", "SimpleQueryEngine"),
|
||||
"SubQuestionEngine": ("praisonaiagents.knowledge.query_engine", "SubQuestionEngine"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load protocols and registries."""
|
||||
if name in _LAZY_IMPORTS:
|
||||
module_path, attr_name = _LAZY_IMPORTS[name]
|
||||
import importlib
|
||||
module = importlib.import_module(module_path)
|
||||
return getattr(module, attr_name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__():
|
||||
"""List available attributes."""
|
||||
return list(_LAZY_IMPORTS.keys()) + ["Knowledge", "Chunking"]
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Core
|
||||
"Knowledge",
|
||||
"Chunking",
|
||||
# Readers
|
||||
"Document",
|
||||
"ReaderProtocol",
|
||||
"ReaderRegistry",
|
||||
"get_reader_registry",
|
||||
"detect_source_kind",
|
||||
# Vector stores
|
||||
"VectorRecord",
|
||||
"VectorStoreProtocol",
|
||||
"VectorStoreRegistry",
|
||||
"get_vector_store_registry",
|
||||
"InMemoryVectorStore",
|
||||
# Retrieval
|
||||
"RetrievalResult",
|
||||
"RetrievalStrategy",
|
||||
"RetrieverProtocol",
|
||||
"RetrieverRegistry",
|
||||
"get_retriever_registry",
|
||||
"reciprocal_rank_fusion",
|
||||
"merge_adjacent_chunks",
|
||||
# Rerankers
|
||||
"RerankResult",
|
||||
"RerankerProtocol",
|
||||
"RerankerRegistry",
|
||||
"get_reranker_registry",
|
||||
"SimpleReranker",
|
||||
# Index
|
||||
"IndexType",
|
||||
"IndexStats",
|
||||
"IndexProtocol",
|
||||
"IndexRegistry",
|
||||
"get_index_registry",
|
||||
"KeywordIndex",
|
||||
# Query engine
|
||||
"QueryMode",
|
||||
"QueryResult",
|
||||
"QueryEngineProtocol",
|
||||
"QueryEngineRegistry",
|
||||
"get_query_engine_registry",
|
||||
"decompose_question",
|
||||
"SimpleQueryEngine",
|
||||
"SubQuestionEngine",
|
||||
]
|
||||
@@ -0,0 +1,332 @@
|
||||
"""
|
||||
Index Type Selection for PraisonAI Agents.
|
||||
|
||||
This module defines index types and protocols.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from enum import Enum
|
||||
import logging
|
||||
import math
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IndexType(str, Enum):
|
||||
"""Available index types."""
|
||||
VECTOR = "vector"
|
||||
KEYWORD = "keyword"
|
||||
HYBRID = "hybrid"
|
||||
GRAPH = "graph"
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexStats:
|
||||
"""Statistics about an index."""
|
||||
index_type: IndexType
|
||||
document_count: int
|
||||
chunk_count: int
|
||||
total_tokens: int = 0
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class IndexProtocol(Protocol):
|
||||
"""
|
||||
Protocol for document indices.
|
||||
|
||||
Implementations must provide methods for indexing and querying.
|
||||
"""
|
||||
|
||||
name: str
|
||||
index_type: IndexType
|
||||
|
||||
def add_documents(
|
||||
self,
|
||||
texts: List[str],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
embeddings: Optional[List[List[float]]] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
Add documents to the index.
|
||||
|
||||
Args:
|
||||
texts: List of text content
|
||||
metadatas: Optional list of metadata dicts
|
||||
ids: Optional list of IDs
|
||||
embeddings: Optional pre-computed embeddings (for vector index)
|
||||
|
||||
Returns:
|
||||
List of document IDs
|
||||
"""
|
||||
...
|
||||
|
||||
def query(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
query_embedding: Optional[List[float]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query the index.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
top_k: Number of results
|
||||
filter: Optional metadata filter
|
||||
query_embedding: Optional pre-computed query embedding
|
||||
|
||||
Returns:
|
||||
List of results with text, score, metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete documents from the index."""
|
||||
...
|
||||
|
||||
def stats(self) -> IndexStats:
|
||||
"""Get index statistics."""
|
||||
...
|
||||
|
||||
|
||||
class IndexRegistry:
|
||||
"""Registry for index types."""
|
||||
|
||||
_instance: Optional["IndexRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "IndexRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._indices: Dict[str, Callable[..., IndexProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(self, name: str, factory: Callable[..., IndexProtocol]) -> None:
|
||||
"""Register an index factory."""
|
||||
self._indices[name] = factory
|
||||
|
||||
def get(self, name: str, **kwargs) -> Optional[IndexProtocol]:
|
||||
"""Get an index by name."""
|
||||
if name in self._indices:
|
||||
try:
|
||||
return self._indices[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize index '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_indices(self) -> List[str]:
|
||||
"""List all registered index names."""
|
||||
return list(self._indices.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered indices."""
|
||||
self._indices.clear()
|
||||
|
||||
|
||||
def get_index_registry() -> IndexRegistry:
|
||||
"""Get the global index registry instance."""
|
||||
return IndexRegistry()
|
||||
|
||||
|
||||
class KeywordIndex:
|
||||
"""
|
||||
Simple BM25-like keyword index.
|
||||
|
||||
Pure Python implementation with no external dependencies.
|
||||
"""
|
||||
|
||||
name: str = "keyword"
|
||||
index_type: IndexType = IndexType.KEYWORD
|
||||
|
||||
def __init__(self, k1: float = 1.5, b: float = 0.75, **kwargs):
|
||||
self.k1 = k1
|
||||
self.b = b
|
||||
self._documents: Dict[str, Dict[str, Any]] = {} # id -> {text, metadata, terms}
|
||||
self._term_doc_freq: Dict[str, int] = {} # term -> doc count
|
||||
self._avg_doc_len: float = 0.0
|
||||
self._total_docs: int = 0
|
||||
|
||||
def add_documents(
|
||||
self,
|
||||
texts: List[str],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
embeddings: Optional[List[List[float]]] = None
|
||||
) -> List[str]:
|
||||
"""Add documents to the keyword index."""
|
||||
import uuid
|
||||
|
||||
metadatas = metadatas or [{} for _ in texts]
|
||||
ids = ids or [str(uuid.uuid4()) for _ in texts]
|
||||
|
||||
total_len = sum(len(doc["text"].split()) for doc in self._documents.values())
|
||||
|
||||
for text, metadata, doc_id in zip(texts, metadatas, ids):
|
||||
terms = self._tokenize(text)
|
||||
term_freq = {}
|
||||
for term in terms:
|
||||
term_freq[term] = term_freq.get(term, 0) + 1
|
||||
|
||||
self._documents[doc_id] = {
|
||||
"text": text,
|
||||
"metadata": metadata,
|
||||
"terms": term_freq,
|
||||
"length": len(terms)
|
||||
}
|
||||
|
||||
# Update document frequency
|
||||
for term in set(terms):
|
||||
self._term_doc_freq[term] = self._term_doc_freq.get(term, 0) + 1
|
||||
|
||||
total_len += len(terms)
|
||||
|
||||
self._total_docs = len(self._documents)
|
||||
self._avg_doc_len = total_len / self._total_docs if self._total_docs > 0 else 0
|
||||
|
||||
return ids
|
||||
|
||||
def query(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
query_embedding: Optional[List[float]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Query using BM25 scoring."""
|
||||
query_terms = self._tokenize(query)
|
||||
|
||||
scores = []
|
||||
for doc_id, doc in self._documents.items():
|
||||
# Apply filter
|
||||
if filter:
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if doc["metadata"].get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if not match:
|
||||
continue
|
||||
|
||||
score = self._bm25_score(query_terms, doc)
|
||||
scores.append({
|
||||
"id": doc_id,
|
||||
"text": doc["text"],
|
||||
"score": score,
|
||||
"metadata": doc["metadata"]
|
||||
})
|
||||
|
||||
# Sort by score
|
||||
scores.sort(key=lambda x: x["score"], reverse=True)
|
||||
return scores[:top_k]
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete documents from the index."""
|
||||
if delete_all:
|
||||
count = len(self._documents)
|
||||
self._documents.clear()
|
||||
self._term_doc_freq.clear()
|
||||
self._total_docs = 0
|
||||
self._avg_doc_len = 0
|
||||
return count
|
||||
|
||||
deleted = 0
|
||||
to_delete = []
|
||||
|
||||
if ids:
|
||||
to_delete = [id_ for id_ in ids if id_ in self._documents]
|
||||
elif filter:
|
||||
for doc_id, doc in self._documents.items():
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if doc["metadata"].get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if match:
|
||||
to_delete.append(doc_id)
|
||||
|
||||
for doc_id in to_delete:
|
||||
doc = self._documents[doc_id]
|
||||
# Update term frequencies
|
||||
for term in doc["terms"]:
|
||||
if term in self._term_doc_freq:
|
||||
self._term_doc_freq[term] -= 1
|
||||
if self._term_doc_freq[term] <= 0:
|
||||
del self._term_doc_freq[term]
|
||||
del self._documents[doc_id]
|
||||
deleted += 1
|
||||
|
||||
self._total_docs = len(self._documents)
|
||||
if self._total_docs > 0:
|
||||
self._avg_doc_len = sum(d["length"] for d in self._documents.values()) / self._total_docs
|
||||
else:
|
||||
self._avg_doc_len = 0
|
||||
|
||||
return deleted
|
||||
|
||||
def stats(self) -> IndexStats:
|
||||
"""Get index statistics."""
|
||||
return IndexStats(
|
||||
index_type=self.index_type,
|
||||
document_count=self._total_docs,
|
||||
chunk_count=self._total_docs,
|
||||
total_tokens=sum(d["length"] for d in self._documents.values()),
|
||||
metadata={"unique_terms": len(self._term_doc_freq)}
|
||||
)
|
||||
|
||||
def _tokenize(self, text: str) -> List[str]:
|
||||
"""Simple tokenization."""
|
||||
import re
|
||||
# Lowercase and split on non-alphanumeric
|
||||
tokens = re.findall(r'\b\w+\b', text.lower())
|
||||
# Remove very short tokens
|
||||
return [t for t in tokens if len(t) > 1]
|
||||
|
||||
def _bm25_score(self, query_terms: List[str], doc: Dict[str, Any]) -> float:
|
||||
"""Calculate BM25 score."""
|
||||
score = 0.0
|
||||
doc_len = doc["length"]
|
||||
term_freq = doc["terms"]
|
||||
|
||||
for term in query_terms:
|
||||
if term not in term_freq:
|
||||
continue
|
||||
|
||||
tf = term_freq[term]
|
||||
df = self._term_doc_freq.get(term, 0)
|
||||
|
||||
# IDF component
|
||||
idf = math.log((self._total_docs - df + 0.5) / (df + 0.5) + 1)
|
||||
|
||||
# TF component with length normalization
|
||||
tf_norm = (tf * (self.k1 + 1)) / (tf + self.k1 * (1 - self.b + self.b * doc_len / self._avg_doc_len))
|
||||
|
||||
score += idf * tf_norm
|
||||
|
||||
return score
|
||||
|
||||
|
||||
# Register the keyword index by default
|
||||
def _register_default_indices():
|
||||
"""Register default indices."""
|
||||
registry = get_index_registry()
|
||||
registry.register("keyword", KeywordIndex)
|
||||
|
||||
|
||||
_register_default_indices()
|
||||
@@ -0,0 +1,351 @@
|
||||
"""
|
||||
Query Engine Patterns for PraisonAI Agents.
|
||||
|
||||
This module defines query modes and protocols.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from enum import Enum
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QueryMode(str, Enum):
|
||||
"""Available query modes."""
|
||||
DEFAULT = "default"
|
||||
SUB_QUESTION = "sub_question"
|
||||
SQL = "sql"
|
||||
ROUTER = "router"
|
||||
SUMMARIZE = "summarize"
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
"""
|
||||
Result from a query operation.
|
||||
|
||||
Attributes:
|
||||
answer: The synthesized answer
|
||||
sources: List of source documents/chunks used
|
||||
sub_questions: Optional list of sub-questions (for sub_question mode)
|
||||
metadata: Additional metadata
|
||||
"""
|
||||
answer: str
|
||||
sources: List[Dict[str, Any]] = field(default_factory=list)
|
||||
sub_questions: Optional[List[str]] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"answer": self.answer,
|
||||
"sources": self.sources,
|
||||
"sub_questions": self.sub_questions,
|
||||
"metadata": self.metadata
|
||||
}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class QueryEngineProtocol(Protocol):
|
||||
"""
|
||||
Protocol for query engines.
|
||||
|
||||
Implementations must provide a query method that takes a question
|
||||
and returns an answer with sources.
|
||||
"""
|
||||
|
||||
name: str
|
||||
mode: QueryMode
|
||||
|
||||
def query(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""
|
||||
Query the knowledge base.
|
||||
|
||||
Args:
|
||||
question: The user's question
|
||||
context: Optional pre-retrieved context
|
||||
**kwargs: Additional engine-specific options
|
||||
|
||||
Returns:
|
||||
QueryResult with answer and sources
|
||||
"""
|
||||
...
|
||||
|
||||
async def aquery(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Async version of query."""
|
||||
...
|
||||
|
||||
|
||||
class QueryEngineRegistry:
|
||||
"""Registry for query engines."""
|
||||
|
||||
_instance: Optional["QueryEngineRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "QueryEngineRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._engines: Dict[str, Callable[..., QueryEngineProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(self, name: str, factory: Callable[..., QueryEngineProtocol]) -> None:
|
||||
"""Register a query engine factory."""
|
||||
self._engines[name] = factory
|
||||
|
||||
def get(self, name: str, **kwargs) -> Optional[QueryEngineProtocol]:
|
||||
"""Get a query engine by name."""
|
||||
if name in self._engines:
|
||||
try:
|
||||
return self._engines[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize query engine '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_engines(self) -> List[str]:
|
||||
"""List all registered engine names."""
|
||||
return list(self._engines.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered engines."""
|
||||
self._engines.clear()
|
||||
|
||||
|
||||
def get_query_engine_registry() -> QueryEngineRegistry:
|
||||
"""Get the global query engine registry instance."""
|
||||
return QueryEngineRegistry()
|
||||
|
||||
|
||||
def decompose_question(question: str) -> List[str]:
|
||||
"""
|
||||
Simple deterministic question decomposition.
|
||||
|
||||
Splits complex questions into sub-questions without LLM.
|
||||
Used as fallback when LLM is not available.
|
||||
|
||||
Args:
|
||||
question: The complex question
|
||||
|
||||
Returns:
|
||||
List of sub-questions
|
||||
"""
|
||||
sub_questions = []
|
||||
|
||||
# Split on common conjunctions
|
||||
conjunctions = [" and ", " also ", " additionally ", " furthermore ", " moreover "]
|
||||
parts = [question]
|
||||
|
||||
for conj in conjunctions:
|
||||
new_parts = []
|
||||
for part in parts:
|
||||
new_parts.extend(part.split(conj))
|
||||
parts = new_parts
|
||||
|
||||
# Clean up parts
|
||||
for part in parts:
|
||||
part = part.strip()
|
||||
if part:
|
||||
# Ensure it ends with a question mark
|
||||
if not part.endswith("?"):
|
||||
part = part + "?"
|
||||
# Capitalize first letter
|
||||
part = part[0].upper() + part[1:] if len(part) > 1 else part.upper()
|
||||
sub_questions.append(part)
|
||||
|
||||
# If no decomposition happened, return original
|
||||
if len(sub_questions) <= 1:
|
||||
return [question]
|
||||
|
||||
return sub_questions
|
||||
|
||||
|
||||
def synthesize_answer(
|
||||
question: str,
|
||||
contexts: List[str],
|
||||
max_context_length: int = 4000
|
||||
) -> str:
|
||||
"""
|
||||
Simple answer synthesis without LLM.
|
||||
|
||||
Concatenates relevant context with the question.
|
||||
Used as fallback when LLM is not available.
|
||||
|
||||
Args:
|
||||
question: The original question
|
||||
contexts: List of context strings
|
||||
max_context_length: Maximum total context length
|
||||
|
||||
Returns:
|
||||
Synthesized answer string
|
||||
"""
|
||||
if not contexts:
|
||||
return "No relevant information found."
|
||||
|
||||
# Truncate contexts to fit
|
||||
total_length = 0
|
||||
selected_contexts = []
|
||||
|
||||
for ctx in contexts:
|
||||
if total_length + len(ctx) > max_context_length:
|
||||
remaining = max_context_length - total_length
|
||||
if remaining > 100:
|
||||
selected_contexts.append(ctx[:remaining] + "...")
|
||||
break
|
||||
selected_contexts.append(ctx)
|
||||
total_length += len(ctx)
|
||||
|
||||
# Format as a simple answer
|
||||
context_text = "\n\n".join(selected_contexts)
|
||||
|
||||
return f"Based on the available information:\n\n{context_text}"
|
||||
|
||||
|
||||
class SimpleQueryEngine:
|
||||
"""
|
||||
Simple query engine for testing and fallback.
|
||||
|
||||
Uses basic context concatenation without LLM synthesis.
|
||||
"""
|
||||
|
||||
name: str = "simple"
|
||||
mode: QueryMode = QueryMode.DEFAULT
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def query(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Simple query that returns context as answer."""
|
||||
contexts = context or []
|
||||
answer = synthesize_answer(question, contexts)
|
||||
|
||||
return QueryResult(
|
||||
answer=answer,
|
||||
sources=[{"text": ctx} for ctx in contexts],
|
||||
metadata={"mode": self.mode.value}
|
||||
)
|
||||
|
||||
async def aquery(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.query(question, context, **kwargs)
|
||||
|
||||
|
||||
class SubQuestionEngine:
|
||||
"""
|
||||
Sub-question decomposition query engine.
|
||||
|
||||
Decomposes complex questions into simpler sub-questions.
|
||||
Uses deterministic fallback when LLM is not available.
|
||||
"""
|
||||
|
||||
name: str = "sub_question"
|
||||
mode: QueryMode = QueryMode.SUB_QUESTION
|
||||
|
||||
def __init__(self, llm: Optional[Any] = None, **kwargs):
|
||||
self.llm = llm
|
||||
|
||||
def query(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
retriever: Optional[Any] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Query with sub-question decomposition."""
|
||||
# Decompose question
|
||||
if self.llm:
|
||||
sub_questions = self._decompose_with_llm(question)
|
||||
else:
|
||||
sub_questions = decompose_question(question)
|
||||
|
||||
# Gather context for each sub-question
|
||||
all_contexts = list(context or [])
|
||||
sub_answers = []
|
||||
|
||||
for sub_q in sub_questions:
|
||||
if retriever:
|
||||
# Retrieve for each sub-question
|
||||
try:
|
||||
results = retriever.retrieve(sub_q, top_k=3)
|
||||
for r in results:
|
||||
if hasattr(r, 'text'):
|
||||
all_contexts.append(r.text)
|
||||
elif isinstance(r, dict):
|
||||
all_contexts.append(r.get('text', str(r)))
|
||||
except Exception as e:
|
||||
logger.warning(f"Retrieval failed for sub-question: {e}")
|
||||
|
||||
# Simple sub-answer (would use LLM in full implementation)
|
||||
sub_answers.append(f"Sub-question: {sub_q}")
|
||||
|
||||
# Synthesize final answer
|
||||
answer = synthesize_answer(question, all_contexts)
|
||||
|
||||
return QueryResult(
|
||||
answer=answer,
|
||||
sources=[{"text": ctx} for ctx in all_contexts[:10]],
|
||||
sub_questions=sub_questions,
|
||||
metadata={"mode": self.mode.value, "sub_question_count": len(sub_questions)}
|
||||
)
|
||||
|
||||
async def aquery(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.query(question, context, **kwargs)
|
||||
|
||||
def _decompose_with_llm(self, question: str) -> List[str]:
|
||||
"""Decompose question using LLM."""
|
||||
try:
|
||||
prompt = f"""Break down this complex question into simpler sub-questions that can be answered independently.
|
||||
|
||||
Question: {question}
|
||||
|
||||
Return only the sub-questions, one per line, without numbering or bullets."""
|
||||
|
||||
response = self.llm.chat(prompt)
|
||||
if response:
|
||||
lines = response.strip().split("\n")
|
||||
return [line.strip() for line in lines if line.strip()]
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM decomposition failed: {e}")
|
||||
|
||||
return decompose_question(question)
|
||||
|
||||
|
||||
# Register default engines
|
||||
def _register_default_engines():
|
||||
"""Register default query engines."""
|
||||
registry = get_query_engine_registry()
|
||||
registry.register("default", SimpleQueryEngine)
|
||||
registry.register("simple", SimpleQueryEngine)
|
||||
registry.register("sub_question", SubQuestionEngine)
|
||||
|
||||
|
||||
_register_default_engines()
|
||||
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
Data Reader Adapter Protocol for PraisonAI Agents.
|
||||
|
||||
This module defines the protocol for document readers and a registry system.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer (praisonai.adapters.readers).
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
import os
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""
|
||||
Lightweight document representation.
|
||||
|
||||
Attributes:
|
||||
content: The text content of the document
|
||||
metadata: Optional metadata dict (source, filename, page, etc.)
|
||||
doc_id: Optional unique identifier
|
||||
embedding: Optional pre-computed embedding
|
||||
"""
|
||||
content: str
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
doc_id: Optional[str] = None
|
||||
embedding: Optional[List[float]] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.doc_id is None:
|
||||
import uuid
|
||||
self.doc_id = str(uuid.uuid4())
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"content": self.content,
|
||||
"metadata": self.metadata,
|
||||
"doc_id": self.doc_id,
|
||||
"embedding": self.embedding
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Document":
|
||||
"""Create from dictionary."""
|
||||
return cls(
|
||||
content=data.get("content", ""),
|
||||
metadata=data.get("metadata", {}),
|
||||
doc_id=data.get("doc_id"),
|
||||
embedding=data.get("embedding")
|
||||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ReaderProtocol(Protocol):
|
||||
"""
|
||||
Protocol for document readers.
|
||||
|
||||
Implementations must provide:
|
||||
- name: Identifier for the reader
|
||||
- supported_extensions: List of file extensions this reader handles
|
||||
- load(): Method to load documents from a source
|
||||
"""
|
||||
|
||||
name: str
|
||||
supported_extensions: List[str]
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Document]:
|
||||
"""
|
||||
Load documents from a source.
|
||||
|
||||
Args:
|
||||
source: File path, URL, or other source identifier
|
||||
metadata: Optional metadata to attach to all documents
|
||||
|
||||
Returns:
|
||||
List of Document objects
|
||||
"""
|
||||
...
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""
|
||||
Check if this reader can handle the given source.
|
||||
|
||||
Args:
|
||||
source: File path, URL, or other source identifier
|
||||
|
||||
Returns:
|
||||
True if this reader can handle the source
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class ReaderRegistry:
|
||||
"""
|
||||
Registry for document readers.
|
||||
|
||||
Provides lazy loading and automatic reader selection based on source type.
|
||||
"""
|
||||
|
||||
_instance: Optional["ReaderRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "ReaderRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._readers: Dict[str, Callable[[], ReaderProtocol]] = {}
|
||||
cls._instance._extension_map: Dict[str, str] = {}
|
||||
cls._instance._initialized_readers: Dict[str, ReaderProtocol] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[[], ReaderProtocol],
|
||||
extensions: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""
|
||||
Register a reader factory.
|
||||
|
||||
Args:
|
||||
name: Unique name for the reader
|
||||
factory: Callable that returns a ReaderProtocol instance
|
||||
extensions: List of file extensions this reader handles
|
||||
"""
|
||||
self._readers[name] = factory
|
||||
if extensions:
|
||||
for ext in extensions:
|
||||
ext_lower = ext.lower().lstrip(".")
|
||||
self._extension_map[ext_lower] = name
|
||||
|
||||
def get(self, name: str) -> Optional[ReaderProtocol]:
|
||||
"""
|
||||
Get a reader by name (lazy initialization).
|
||||
|
||||
Args:
|
||||
name: Reader name
|
||||
|
||||
Returns:
|
||||
ReaderProtocol instance or None
|
||||
"""
|
||||
if name in self._initialized_readers:
|
||||
return self._initialized_readers[name]
|
||||
|
||||
if name in self._readers:
|
||||
try:
|
||||
reader = self._readers[name]()
|
||||
self._initialized_readers[name] = reader
|
||||
return reader
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize reader '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def get_for_source(self, source: str) -> Optional[ReaderProtocol]:
|
||||
"""
|
||||
Get the best reader for a given source.
|
||||
|
||||
Args:
|
||||
source: File path, URL, or other source identifier
|
||||
|
||||
Returns:
|
||||
Best matching ReaderProtocol or None
|
||||
"""
|
||||
source_kind = detect_source_kind(source)
|
||||
|
||||
if source_kind == "file":
|
||||
ext = os.path.splitext(source)[1].lower().lstrip(".")
|
||||
if ext in self._extension_map:
|
||||
return self.get(self._extension_map[ext])
|
||||
elif source_kind == "url":
|
||||
# Try URL reader
|
||||
if "url" in self._readers:
|
||||
return self.get("url")
|
||||
elif source_kind == "directory":
|
||||
if "directory" in self._readers:
|
||||
return self.get("directory")
|
||||
|
||||
# Fallback to auto reader if available
|
||||
if "auto" in self._readers:
|
||||
return self.get("auto")
|
||||
|
||||
return None
|
||||
|
||||
def list_readers(self) -> List[str]:
|
||||
"""List all registered reader names."""
|
||||
return list(self._readers.keys())
|
||||
|
||||
def list_extensions(self) -> Dict[str, str]:
|
||||
"""List all registered extensions and their readers."""
|
||||
return dict(self._extension_map)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered readers (mainly for testing)."""
|
||||
self._readers.clear()
|
||||
self._extension_map.clear()
|
||||
self._initialized_readers.clear()
|
||||
|
||||
|
||||
def get_reader_registry() -> ReaderRegistry:
|
||||
"""Get the global reader registry instance."""
|
||||
return ReaderRegistry()
|
||||
|
||||
|
||||
def detect_source_kind(source: str) -> str:
|
||||
"""
|
||||
Detect the kind of source without importing heavy libraries.
|
||||
|
||||
Args:
|
||||
source: File path, URL, directory, or glob pattern
|
||||
|
||||
Returns:
|
||||
One of: "file", "url", "directory", "glob", "unknown"
|
||||
"""
|
||||
if not isinstance(source, str):
|
||||
return "unknown"
|
||||
|
||||
source = source.strip()
|
||||
|
||||
# Check for URL
|
||||
if source.startswith(("http://", "https://", "ftp://", "s3://", "gs://")):
|
||||
return "url"
|
||||
|
||||
# Check for glob pattern
|
||||
if "*" in source or "?" in source or "[" in source:
|
||||
return "glob"
|
||||
|
||||
# Check if it's a directory
|
||||
if os.path.isdir(source):
|
||||
return "directory"
|
||||
|
||||
# Check if it's a file
|
||||
if os.path.isfile(source):
|
||||
return "file"
|
||||
|
||||
# Check if it looks like a file path (has extension)
|
||||
if os.path.splitext(source)[1]:
|
||||
return "file"
|
||||
|
||||
return "unknown"
|
||||
|
||||
|
||||
def get_file_extension(source: str) -> str:
|
||||
"""
|
||||
Get the file extension from a source path.
|
||||
|
||||
Args:
|
||||
source: File path or URL
|
||||
|
||||
Returns:
|
||||
Lowercase extension without dot, or empty string
|
||||
"""
|
||||
# Handle URLs by extracting path
|
||||
if source.startswith(("http://", "https://")):
|
||||
from urllib.parse import urlparse
|
||||
path = urlparse(source).path
|
||||
return os.path.splitext(path)[1].lower().lstrip(".")
|
||||
|
||||
return os.path.splitext(source)[1].lower().lstrip(".")
|
||||
|
||||
|
||||
# Default extension mappings (reader name -> extensions)
|
||||
DEFAULT_EXTENSION_MAPPINGS = {
|
||||
"text": ["txt", "text"],
|
||||
"markdown": ["md", "markdown"],
|
||||
"json": ["json", "jsonl"],
|
||||
"csv": ["csv", "tsv"],
|
||||
"html": ["html", "htm"],
|
||||
"xml": ["xml"],
|
||||
"pdf": ["pdf"],
|
||||
"docx": ["docx", "doc"],
|
||||
"xlsx": ["xlsx", "xls"],
|
||||
"pptx": ["pptx", "ppt"],
|
||||
"image": ["jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp"],
|
||||
"audio": ["mp3", "wav", "ogg", "m4a", "flac"],
|
||||
"video": ["mp4", "avi", "mov", "mkv", "webm"],
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
Reranker Protocol for PraisonAI Agents.
|
||||
|
||||
This module defines the protocol for rerankers and a registry system.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by praisonai-tools or wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankResult:
|
||||
"""
|
||||
Result from a reranking operation.
|
||||
|
||||
Attributes:
|
||||
text: The text content
|
||||
score: Reranking score (higher is more relevant)
|
||||
original_index: Original position before reranking
|
||||
metadata: Optional metadata
|
||||
"""
|
||||
text: str
|
||||
score: float
|
||||
original_index: int
|
||||
metadata: Dict[str, Any] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.metadata is None:
|
||||
self.metadata = {}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RerankerProtocol(Protocol):
|
||||
"""
|
||||
Protocol for rerankers.
|
||||
|
||||
Implementations must provide a rerank method that takes a query
|
||||
and a list of documents and returns reranked results.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""
|
||||
Rerank documents based on relevance to query.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of document texts to rerank
|
||||
top_k: Optional limit on results (None = return all)
|
||||
**kwargs: Additional reranker-specific options
|
||||
|
||||
Returns:
|
||||
List of RerankResult objects sorted by relevance
|
||||
"""
|
||||
...
|
||||
|
||||
async def arerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""Async version of rerank."""
|
||||
...
|
||||
|
||||
|
||||
class RerankerRegistry:
|
||||
"""
|
||||
Registry for rerankers.
|
||||
"""
|
||||
|
||||
_instance: Optional["RerankerRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "RerankerRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._rerankers: Dict[str, Callable[..., RerankerProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[..., RerankerProtocol]
|
||||
) -> None:
|
||||
"""Register a reranker factory."""
|
||||
self._rerankers[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
**kwargs
|
||||
) -> Optional[RerankerProtocol]:
|
||||
"""Get a reranker by name."""
|
||||
if name in self._rerankers:
|
||||
try:
|
||||
return self._rerankers[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize reranker '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_rerankers(self) -> List[str]:
|
||||
"""List all registered reranker names."""
|
||||
return list(self._rerankers.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered rerankers."""
|
||||
self._rerankers.clear()
|
||||
|
||||
|
||||
def get_reranker_registry() -> RerankerRegistry:
|
||||
"""Get the global reranker registry instance."""
|
||||
return RerankerRegistry()
|
||||
|
||||
|
||||
class SimpleReranker:
|
||||
"""
|
||||
Simple keyword-based reranker for testing and fallback.
|
||||
|
||||
Uses term frequency overlap for scoring. No external dependencies.
|
||||
"""
|
||||
|
||||
name: str = "simple"
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""Rerank using simple term overlap."""
|
||||
query_terms = set(query.lower().split())
|
||||
|
||||
results = []
|
||||
for i, doc in enumerate(documents):
|
||||
doc_terms = set(doc.lower().split())
|
||||
|
||||
# Calculate Jaccard similarity
|
||||
if not query_terms or not doc_terms:
|
||||
score = 0.0
|
||||
else:
|
||||
intersection = len(query_terms & doc_terms)
|
||||
union = len(query_terms | doc_terms)
|
||||
score = intersection / union if union > 0 else 0.0
|
||||
|
||||
results.append(RerankResult(
|
||||
text=doc,
|
||||
score=score,
|
||||
original_index=i
|
||||
))
|
||||
|
||||
# Sort by score descending
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
|
||||
if top_k is not None:
|
||||
results = results[:top_k]
|
||||
|
||||
return results
|
||||
|
||||
async def arerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.rerank(query, documents, top_k, **kwargs)
|
||||
|
||||
|
||||
# Register the simple reranker by default
|
||||
def _register_default_rerankers():
|
||||
"""Register default rerankers."""
|
||||
registry = get_reranker_registry()
|
||||
registry.register("simple", SimpleReranker)
|
||||
|
||||
|
||||
_register_default_rerankers()
|
||||
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
Retrieval Strategy Patterns for PraisonAI Agents.
|
||||
|
||||
This module defines protocols and strategies for document retrieval.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from enum import Enum
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RetrievalStrategy(str, Enum):
|
||||
"""Available retrieval strategies."""
|
||||
BASIC = "basic"
|
||||
FUSION = "fusion"
|
||||
RECURSIVE = "recursive"
|
||||
AUTO_MERGE = "auto_merge"
|
||||
HYBRID = "hybrid"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetrievalResult:
|
||||
"""
|
||||
Result from a retrieval operation.
|
||||
|
||||
Attributes:
|
||||
text: The retrieved text content
|
||||
score: Relevance score (0-1)
|
||||
metadata: Source metadata
|
||||
doc_id: Document identifier
|
||||
chunk_index: Index of chunk within document
|
||||
"""
|
||||
text: str
|
||||
score: float
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
doc_id: Optional[str] = None
|
||||
chunk_index: Optional[int] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"text": self.text,
|
||||
"score": self.score,
|
||||
"metadata": self.metadata,
|
||||
"doc_id": self.doc_id,
|
||||
"chunk_index": self.chunk_index
|
||||
}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RetrieverProtocol(Protocol):
|
||||
"""
|
||||
Protocol for document retrievers.
|
||||
|
||||
Implementations must provide a retrieve method that takes a query
|
||||
and returns relevant documents.
|
||||
"""
|
||||
|
||||
name: str
|
||||
strategy: RetrievalStrategy
|
||||
|
||||
def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
Retrieve relevant documents for a query.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
top_k: Maximum number of results
|
||||
filter: Optional metadata filter
|
||||
**kwargs: Additional retriever-specific options
|
||||
|
||||
Returns:
|
||||
List of RetrievalResult objects sorted by relevance
|
||||
"""
|
||||
...
|
||||
|
||||
async def aretrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[RetrievalResult]:
|
||||
"""Async version of retrieve."""
|
||||
...
|
||||
|
||||
|
||||
class RetrieverRegistry:
|
||||
"""
|
||||
Registry for retrieval strategies.
|
||||
"""
|
||||
|
||||
_instance: Optional["RetrieverRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "RetrieverRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._retrievers: Dict[str, Callable[..., RetrieverProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[..., RetrieverProtocol]
|
||||
) -> None:
|
||||
"""Register a retriever factory."""
|
||||
self._retrievers[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
**kwargs
|
||||
) -> Optional[RetrieverProtocol]:
|
||||
"""Get a retriever by name."""
|
||||
if name in self._retrievers:
|
||||
try:
|
||||
return self._retrievers[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize retriever '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_retrievers(self) -> List[str]:
|
||||
"""List all registered retriever names."""
|
||||
return list(self._retrievers.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered retrievers."""
|
||||
self._retrievers.clear()
|
||||
|
||||
|
||||
def get_retriever_registry() -> RetrieverRegistry:
|
||||
"""Get the global retriever registry instance."""
|
||||
return RetrieverRegistry()
|
||||
|
||||
|
||||
def reciprocal_rank_fusion(
|
||||
result_lists: List[List[RetrievalResult]],
|
||||
k: int = 60
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
Reciprocal Rank Fusion (RRF) algorithm for combining multiple result lists.
|
||||
|
||||
Pure Python implementation with no external dependencies.
|
||||
|
||||
Args:
|
||||
result_lists: List of result lists from different retrievers
|
||||
k: RRF constant (default 60)
|
||||
|
||||
Returns:
|
||||
Fused and re-ranked list of results
|
||||
"""
|
||||
# Track scores by document ID
|
||||
fused_scores: Dict[str, float] = {}
|
||||
doc_map: Dict[str, RetrievalResult] = {}
|
||||
|
||||
for results in result_lists:
|
||||
for rank, result in enumerate(results):
|
||||
doc_key = result.doc_id or result.text[:100] # Use text prefix as fallback key
|
||||
|
||||
# RRF score: 1 / (k + rank)
|
||||
rrf_score = 1.0 / (k + rank + 1)
|
||||
|
||||
if doc_key in fused_scores:
|
||||
fused_scores[doc_key] += rrf_score
|
||||
else:
|
||||
fused_scores[doc_key] = rrf_score
|
||||
doc_map[doc_key] = result
|
||||
|
||||
# Sort by fused score
|
||||
sorted_keys = sorted(fused_scores.keys(), key=lambda x: fused_scores[x], reverse=True)
|
||||
|
||||
# Build result list with updated scores
|
||||
fused_results = []
|
||||
for key in sorted_keys:
|
||||
result = doc_map[key]
|
||||
fused_results.append(RetrievalResult(
|
||||
text=result.text,
|
||||
score=fused_scores[key],
|
||||
metadata=result.metadata,
|
||||
doc_id=result.doc_id,
|
||||
chunk_index=result.chunk_index
|
||||
))
|
||||
|
||||
return fused_results
|
||||
|
||||
|
||||
def merge_adjacent_chunks(
|
||||
results: List[RetrievalResult],
|
||||
max_gap: int = 1
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
Merge adjacent chunks from the same document.
|
||||
|
||||
Args:
|
||||
results: List of retrieval results
|
||||
max_gap: Maximum gap between chunk indices to merge
|
||||
|
||||
Returns:
|
||||
Merged results
|
||||
"""
|
||||
if not results:
|
||||
return results
|
||||
|
||||
# Group by document ID
|
||||
doc_groups: Dict[str, List[RetrievalResult]] = {}
|
||||
for result in results:
|
||||
doc_id = result.doc_id or "unknown"
|
||||
if doc_id not in doc_groups:
|
||||
doc_groups[doc_id] = []
|
||||
doc_groups[doc_id].append(result)
|
||||
|
||||
merged_results = []
|
||||
|
||||
for doc_id, chunks in doc_groups.items():
|
||||
# Sort by chunk index
|
||||
chunks_with_index = [c for c in chunks if c.chunk_index is not None]
|
||||
chunks_without_index = [c for c in chunks if c.chunk_index is None]
|
||||
|
||||
if not chunks_with_index:
|
||||
merged_results.extend(chunks_without_index)
|
||||
continue
|
||||
|
||||
chunks_with_index.sort(key=lambda x: x.chunk_index or 0)
|
||||
|
||||
# Merge adjacent chunks
|
||||
current_group = [chunks_with_index[0]]
|
||||
|
||||
for chunk in chunks_with_index[1:]:
|
||||
last_chunk = current_group[-1]
|
||||
last_idx = last_chunk.chunk_index or 0
|
||||
curr_idx = chunk.chunk_index or 0
|
||||
|
||||
if curr_idx - last_idx <= max_gap:
|
||||
current_group.append(chunk)
|
||||
else:
|
||||
# Merge current group and start new one
|
||||
merged_results.append(_merge_chunk_group(current_group))
|
||||
current_group = [chunk]
|
||||
|
||||
# Don't forget the last group
|
||||
if current_group:
|
||||
merged_results.append(_merge_chunk_group(current_group))
|
||||
|
||||
merged_results.extend(chunks_without_index)
|
||||
|
||||
# Sort by score
|
||||
merged_results.sort(key=lambda x: x.score, reverse=True)
|
||||
return merged_results
|
||||
|
||||
|
||||
def _merge_chunk_group(chunks: List[RetrievalResult]) -> RetrievalResult:
|
||||
"""Merge a group of chunks into a single result."""
|
||||
if len(chunks) == 1:
|
||||
return chunks[0]
|
||||
|
||||
# Combine text
|
||||
merged_text = "\n".join(c.text for c in chunks)
|
||||
|
||||
# Average score
|
||||
avg_score = sum(c.score for c in chunks) / len(chunks)
|
||||
|
||||
# Use first chunk's metadata
|
||||
metadata = chunks[0].metadata.copy()
|
||||
metadata["merged_chunks"] = len(chunks)
|
||||
|
||||
return RetrievalResult(
|
||||
text=merged_text,
|
||||
score=avg_score,
|
||||
metadata=metadata,
|
||||
doc_id=chunks[0].doc_id,
|
||||
chunk_index=chunks[0].chunk_index
|
||||
)
|
||||
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
Vector Store Adapter Protocol for PraisonAI Agents.
|
||||
|
||||
This module defines the protocol for vector stores and a registry system.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by praisonai-tools or wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorRecord:
|
||||
"""
|
||||
A record in a vector store.
|
||||
|
||||
Attributes:
|
||||
id: Unique identifier
|
||||
text: The text content
|
||||
embedding: The vector embedding
|
||||
metadata: Optional metadata dict
|
||||
score: Optional similarity score (for query results)
|
||||
"""
|
||||
id: str
|
||||
text: str
|
||||
embedding: List[float]
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
score: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"text": self.text,
|
||||
"embedding": self.embedding,
|
||||
"metadata": self.metadata,
|
||||
"score": self.score
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "VectorRecord":
|
||||
"""Create from dictionary."""
|
||||
return cls(
|
||||
id=data["id"],
|
||||
text=data["text"],
|
||||
embedding=data["embedding"],
|
||||
metadata=data.get("metadata", {}),
|
||||
score=data.get("score")
|
||||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VectorStoreProtocol(Protocol):
|
||||
"""
|
||||
Protocol for vector stores.
|
||||
|
||||
Implementations must provide methods for:
|
||||
- Adding vectors
|
||||
- Querying by vector similarity
|
||||
- Deleting vectors
|
||||
- Getting count
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def add(
|
||||
self,
|
||||
texts: List[str],
|
||||
embeddings: List[List[float]],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
Add vectors to the store.
|
||||
|
||||
Args:
|
||||
texts: List of text content
|
||||
embeddings: List of embedding vectors
|
||||
metadatas: Optional list of metadata dicts
|
||||
ids: Optional list of IDs (auto-generated if not provided)
|
||||
namespace: Optional namespace for multi-tenant isolation
|
||||
|
||||
Returns:
|
||||
List of IDs for the added vectors
|
||||
"""
|
||||
...
|
||||
|
||||
def query(
|
||||
self,
|
||||
embedding: List[float],
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""
|
||||
Query vectors by similarity.
|
||||
|
||||
Args:
|
||||
embedding: Query embedding vector
|
||||
top_k: Number of results to return
|
||||
namespace: Optional namespace filter
|
||||
filter: Optional metadata filter
|
||||
|
||||
Returns:
|
||||
List of VectorRecord objects sorted by similarity
|
||||
"""
|
||||
...
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""
|
||||
Delete vectors from the store.
|
||||
|
||||
Args:
|
||||
ids: Optional list of IDs to delete
|
||||
namespace: Optional namespace filter
|
||||
filter: Optional metadata filter
|
||||
delete_all: If True, delete all vectors
|
||||
|
||||
Returns:
|
||||
Number of vectors deleted
|
||||
"""
|
||||
...
|
||||
|
||||
def count(self, namespace: Optional[str] = None) -> int:
|
||||
"""
|
||||
Get the number of vectors in the store.
|
||||
|
||||
Args:
|
||||
namespace: Optional namespace filter
|
||||
|
||||
Returns:
|
||||
Number of vectors
|
||||
"""
|
||||
...
|
||||
|
||||
def get(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""
|
||||
Get vectors by ID.
|
||||
|
||||
Args:
|
||||
ids: List of IDs to retrieve
|
||||
namespace: Optional namespace filter
|
||||
|
||||
Returns:
|
||||
List of VectorRecord objects
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class VectorStoreRegistry:
|
||||
"""
|
||||
Registry for vector stores.
|
||||
|
||||
Provides lazy loading and configuration-based initialization.
|
||||
"""
|
||||
|
||||
_instance: Optional["VectorStoreRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "VectorStoreRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._stores: Dict[str, Callable[..., VectorStoreProtocol]] = {}
|
||||
cls._instance._initialized_stores: Dict[str, VectorStoreProtocol] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[..., VectorStoreProtocol]
|
||||
) -> None:
|
||||
"""
|
||||
Register a vector store factory.
|
||||
|
||||
Args:
|
||||
name: Unique name for the store (e.g., "chroma", "pinecone")
|
||||
factory: Callable that returns a VectorStoreProtocol instance
|
||||
"""
|
||||
self._stores[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> Optional[VectorStoreProtocol]:
|
||||
"""
|
||||
Get a vector store by name (lazy initialization).
|
||||
|
||||
Args:
|
||||
name: Store name
|
||||
config: Optional configuration dict
|
||||
namespace: Optional namespace for multi-tenant isolation
|
||||
|
||||
Returns:
|
||||
VectorStoreProtocol instance or None
|
||||
"""
|
||||
cache_key = f"{name}:{namespace or 'default'}"
|
||||
|
||||
if cache_key in self._initialized_stores:
|
||||
return self._initialized_stores[cache_key]
|
||||
|
||||
if name in self._stores:
|
||||
try:
|
||||
store = self._stores[name](config=config, namespace=namespace)
|
||||
self._initialized_stores[cache_key] = store
|
||||
return store
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize vector store '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_stores(self) -> List[str]:
|
||||
"""List all registered store names."""
|
||||
return list(self._stores.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered stores (mainly for testing)."""
|
||||
self._stores.clear()
|
||||
self._initialized_stores.clear()
|
||||
|
||||
|
||||
def get_vector_store_registry() -> VectorStoreRegistry:
|
||||
"""Get the global vector store registry instance."""
|
||||
return VectorStoreRegistry()
|
||||
|
||||
|
||||
class InMemoryVectorStore:
|
||||
"""
|
||||
Pure Python in-memory vector store for testing and default usage.
|
||||
|
||||
Uses cosine similarity for queries. No external dependencies.
|
||||
"""
|
||||
|
||||
name: str = "memory"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
namespace: Optional[str] = None
|
||||
):
|
||||
self.config = config or {}
|
||||
self.default_namespace = namespace or "default"
|
||||
self._data: Dict[str, Dict[str, VectorRecord]] = {} # namespace -> id -> record
|
||||
|
||||
def _get_namespace_data(self, namespace: Optional[str] = None) -> Dict[str, VectorRecord]:
|
||||
"""Get or create namespace data."""
|
||||
ns = namespace or self.default_namespace
|
||||
if ns not in self._data:
|
||||
self._data[ns] = {}
|
||||
return self._data[ns]
|
||||
|
||||
def add(
|
||||
self,
|
||||
texts: List[str],
|
||||
embeddings: List[List[float]],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""Add vectors to the store."""
|
||||
import uuid
|
||||
|
||||
ns_data = self._get_namespace_data(namespace)
|
||||
metadatas = metadatas or [{} for _ in texts]
|
||||
ids = ids or [str(uuid.uuid4()) for _ in texts]
|
||||
|
||||
result_ids = []
|
||||
for i, (text, embedding, metadata, id_) in enumerate(zip(texts, embeddings, metadatas, ids)):
|
||||
record = VectorRecord(
|
||||
id=id_,
|
||||
text=text,
|
||||
embedding=embedding,
|
||||
metadata=metadata
|
||||
)
|
||||
ns_data[id_] = record
|
||||
result_ids.append(id_)
|
||||
|
||||
return result_ids
|
||||
|
||||
def query(
|
||||
self,
|
||||
embedding: List[float],
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""Query vectors by cosine similarity."""
|
||||
ns_data = self._get_namespace_data(namespace)
|
||||
|
||||
if not ns_data:
|
||||
return []
|
||||
|
||||
# Calculate cosine similarity for all records
|
||||
results = []
|
||||
for record in ns_data.values():
|
||||
# Apply metadata filter if provided
|
||||
if filter:
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if record.metadata.get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if not match:
|
||||
continue
|
||||
|
||||
score = self._cosine_similarity(embedding, record.embedding)
|
||||
results.append(VectorRecord(
|
||||
id=record.id,
|
||||
text=record.text,
|
||||
embedding=record.embedding,
|
||||
metadata=record.metadata,
|
||||
score=score
|
||||
))
|
||||
|
||||
# Sort by score descending and return top_k
|
||||
results.sort(key=lambda x: x.score or 0, reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete vectors from the store."""
|
||||
ns = namespace or self.default_namespace
|
||||
|
||||
if delete_all:
|
||||
if ns in self._data:
|
||||
count = len(self._data[ns])
|
||||
self._data[ns] = {}
|
||||
return count
|
||||
return 0
|
||||
|
||||
if ns not in self._data:
|
||||
return 0
|
||||
|
||||
ns_data = self._data[ns]
|
||||
deleted = 0
|
||||
|
||||
if ids:
|
||||
for id_ in ids:
|
||||
if id_ in ns_data:
|
||||
del ns_data[id_]
|
||||
deleted += 1
|
||||
elif filter:
|
||||
to_delete = []
|
||||
for id_, record in ns_data.items():
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if record.metadata.get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if match:
|
||||
to_delete.append(id_)
|
||||
for id_ in to_delete:
|
||||
del ns_data[id_]
|
||||
deleted += 1
|
||||
|
||||
return deleted
|
||||
|
||||
def count(self, namespace: Optional[str] = None) -> int:
|
||||
"""Get the number of vectors in the store."""
|
||||
ns = namespace or self.default_namespace
|
||||
if ns in self._data:
|
||||
return len(self._data[ns])
|
||||
return 0
|
||||
|
||||
def get(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""Get vectors by ID."""
|
||||
ns_data = self._get_namespace_data(namespace)
|
||||
results = []
|
||||
for id_ in ids:
|
||||
if id_ in ns_data:
|
||||
results.append(ns_data[id_])
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def _cosine_similarity(a: List[float], b: List[float]) -> float:
|
||||
"""Calculate cosine similarity between two vectors."""
|
||||
if len(a) != len(b):
|
||||
return 0.0
|
||||
|
||||
dot_product = sum(x * y for x, y in zip(a, b))
|
||||
norm_a = sum(x * x for x in a) ** 0.5
|
||||
norm_b = sum(x * x for x in b) ** 0.5
|
||||
|
||||
if norm_a == 0 or norm_b == 0:
|
||||
return 0.0
|
||||
|
||||
return dot_product / (norm_a * norm_b)
|
||||
|
||||
|
||||
# Register the in-memory store by default
|
||||
def _register_default_stores():
|
||||
"""Register default vector stores."""
|
||||
registry = get_vector_store_registry()
|
||||
registry.register("memory", InMemoryVectorStore)
|
||||
|
||||
|
||||
_register_default_stores()
|
||||
@@ -27,6 +27,7 @@ from .model_router import (
|
||||
TaskComplexity,
|
||||
create_routing_agent
|
||||
)
|
||||
from .rate_limiter import RateLimiter
|
||||
|
||||
__all__ = [
|
||||
"LLM",
|
||||
@@ -46,5 +47,6 @@ __all__ = [
|
||||
"ModelRouter",
|
||||
"ModelProfile",
|
||||
"TaskComplexity",
|
||||
"create_routing_agent"
|
||||
"create_routing_agent",
|
||||
"RateLimiter"
|
||||
]
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
Rate Limiter for PraisonAI Agents.
|
||||
|
||||
Provides token bucket rate limiting for LLM API calls.
|
||||
|
||||
Zero Performance Impact:
|
||||
- When rate_limiter is None, no overhead
|
||||
- Uses monotonic clock for accuracy
|
||||
- Injectable clock/sleep for deterministic testing
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
from praisonaiagents import Agent
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
agent = Agent(name="Test", rate_limiter=limiter)
|
||||
"""
|
||||
|
||||
import time
|
||||
import asyncio
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Callable
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimiter:
|
||||
"""Token bucket rate limiter for API calls.
|
||||
|
||||
Uses the token bucket algorithm:
|
||||
- Tokens are added at a fixed rate (requests_per_minute / 60 per second)
|
||||
- Each request consumes one token
|
||||
- If no tokens available, wait until one is available
|
||||
- Burst allows multiple requests in quick succession
|
||||
|
||||
Args:
|
||||
requests_per_minute: Maximum requests per minute
|
||||
tokens_per_minute: Optional token-based rate limiting (for future use)
|
||||
burst: Maximum burst size (default: 1)
|
||||
|
||||
Example:
|
||||
limiter = RateLimiter(requests_per_minute=60) # 1 req/sec
|
||||
limiter.acquire() # Blocks if rate exceeded
|
||||
|
||||
# Async usage
|
||||
await limiter.acquire_async()
|
||||
"""
|
||||
requests_per_minute: int
|
||||
tokens_per_minute: Optional[int] = None
|
||||
burst: int = 1
|
||||
|
||||
# Internal state
|
||||
_tokens: float = field(default=None, init=False, repr=False)
|
||||
_last_update: float = field(default=None, init=False, repr=False)
|
||||
_lock: asyncio.Lock = field(default=None, init=False, repr=False)
|
||||
|
||||
# Injectable functions for testing
|
||||
_get_time: Callable[[], float] = field(default=None, init=False, repr=False)
|
||||
_sleep: Callable[[float], None] = field(default=None, init=False, repr=False)
|
||||
_async_sleep: Callable[[float], None] = field(default=None, init=False, repr=False)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize internal state."""
|
||||
self._tokens = float(self.burst)
|
||||
self._last_update = time.monotonic()
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# Default implementations (can be overridden for testing)
|
||||
if self._get_time is None:
|
||||
self._get_time = time.monotonic
|
||||
if self._sleep is None:
|
||||
self._sleep = time.sleep
|
||||
if self._async_sleep is None:
|
||||
self._async_sleep = asyncio.sleep
|
||||
|
||||
@property
|
||||
def _rate(self) -> float:
|
||||
"""Tokens per second."""
|
||||
return self.requests_per_minute / 60.0
|
||||
|
||||
def _refill(self) -> None:
|
||||
"""Refill tokens based on elapsed time."""
|
||||
now = self._get_time()
|
||||
elapsed = now - self._last_update
|
||||
self._tokens = min(self.burst, self._tokens + elapsed * self._rate)
|
||||
self._last_update = now
|
||||
|
||||
def _wait_time(self) -> float:
|
||||
"""Calculate time to wait for next token."""
|
||||
if self._tokens >= 1.0:
|
||||
return 0.0
|
||||
return (1.0 - self._tokens) / self._rate
|
||||
|
||||
def acquire(self) -> None:
|
||||
"""Acquire a token, blocking if necessary.
|
||||
|
||||
This is the synchronous version. Use acquire_async() for async code.
|
||||
"""
|
||||
self._refill()
|
||||
|
||||
wait = self._wait_time()
|
||||
if wait > 0:
|
||||
self._sleep(wait)
|
||||
self._refill()
|
||||
|
||||
self._tokens -= 1.0
|
||||
|
||||
async def acquire_async(self) -> None:
|
||||
"""Acquire a token asynchronously, waiting if necessary.
|
||||
|
||||
This is the async version. Use acquire() for sync code.
|
||||
"""
|
||||
async with self._lock:
|
||||
self._refill()
|
||||
|
||||
wait = self._wait_time()
|
||||
if wait > 0:
|
||||
await self._async_sleep(wait)
|
||||
self._refill()
|
||||
|
||||
self._tokens -= 1.0
|
||||
|
||||
def try_acquire(self) -> bool:
|
||||
"""Try to acquire a token without blocking.
|
||||
|
||||
Returns:
|
||||
True if token acquired, False if rate limit exceeded
|
||||
"""
|
||||
self._refill()
|
||||
|
||||
if self._tokens >= 1.0:
|
||||
self._tokens -= 1.0
|
||||
return True
|
||||
return False
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the rate limiter to initial state."""
|
||||
self._tokens = float(self.burst)
|
||||
self._last_update = self._get_time()
|
||||
|
||||
@property
|
||||
def available_tokens(self) -> float:
|
||||
"""Get current available tokens (for monitoring)."""
|
||||
self._refill()
|
||||
return self._tokens
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"RateLimiter(rpm={self.requests_per_minute}, burst={self.burst})"
|
||||
@@ -0,0 +1,65 @@
|
||||
"""
|
||||
LSP Integration Module for PraisonAI Agents.
|
||||
|
||||
Provides Language Server Protocol integration for code intelligence:
|
||||
- Diagnostics (errors, warnings)
|
||||
- Code completion suggestions
|
||||
- Go to definition
|
||||
- Find references
|
||||
- Hover information
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- LSP client only initialized when needed
|
||||
- No overhead when not in use
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.lsp import LSPClient
|
||||
|
||||
# Create client for a language
|
||||
client = LSPClient(language="python")
|
||||
|
||||
# Start the language server
|
||||
await client.start()
|
||||
|
||||
# Get diagnostics for a file
|
||||
diagnostics = await client.get_diagnostics("/path/to/file.py")
|
||||
|
||||
# Get completions at a position
|
||||
completions = await client.get_completions("/path/to/file.py", line=10, character=5)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Core client
|
||||
"LSPClient",
|
||||
# Data types
|
||||
"Diagnostic",
|
||||
"DiagnosticSeverity",
|
||||
"CompletionItem",
|
||||
"Location",
|
||||
"Position",
|
||||
"Range",
|
||||
# Configuration
|
||||
"LSPConfig",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name == "LSPClient":
|
||||
from .client import LSPClient
|
||||
return LSPClient
|
||||
|
||||
if name in ("Diagnostic", "DiagnosticSeverity", "CompletionItem",
|
||||
"Location", "Position", "Range"):
|
||||
from .types import (
|
||||
Diagnostic, DiagnosticSeverity, CompletionItem,
|
||||
Location, Position, Range
|
||||
)
|
||||
return locals()[name]
|
||||
|
||||
if name == "LSPConfig":
|
||||
from .config import LSPConfig
|
||||
return LSPConfig
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,489 @@
|
||||
"""
|
||||
LSP Client for PraisonAI Agents.
|
||||
|
||||
Provides Language Server Protocol client implementation.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any
|
||||
from pathlib import Path
|
||||
|
||||
from .types import (
|
||||
Diagnostic, DiagnosticSeverity, CompletionItem, Location,
|
||||
Position, Range, TextDocumentItem, TextDocumentIdentifier,
|
||||
TextDocumentPositionParams
|
||||
)
|
||||
from .config import LSPConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LSPClient:
|
||||
"""
|
||||
Language Server Protocol client.
|
||||
|
||||
Communicates with language servers to provide code intelligence
|
||||
features like diagnostics, completions, and navigation.
|
||||
|
||||
Example:
|
||||
client = LSPClient(language="python")
|
||||
await client.start()
|
||||
|
||||
# Get diagnostics
|
||||
diagnostics = await client.get_diagnostics("/path/to/file.py")
|
||||
|
||||
# Get completions
|
||||
completions = await client.get_completions(
|
||||
"/path/to/file.py", line=10, character=5
|
||||
)
|
||||
|
||||
await client.stop()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
language: str,
|
||||
command: Optional[str] = None,
|
||||
args: Optional[List[str]] = None,
|
||||
root_uri: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
Initialize the LSP client.
|
||||
|
||||
Args:
|
||||
language: Programming language (python, javascript, etc.)
|
||||
command: Language server command (uses default if not specified)
|
||||
args: Command arguments
|
||||
root_uri: Workspace root URI
|
||||
"""
|
||||
self.config = LSPConfig(
|
||||
language=language,
|
||||
command=command,
|
||||
args=args or [],
|
||||
root_uri=root_uri
|
||||
)
|
||||
|
||||
self._process: Optional[asyncio.subprocess.Process] = None
|
||||
self._request_id = 0
|
||||
self._pending_requests: Dict[int, asyncio.Future] = {}
|
||||
self._diagnostics: Dict[str, List[Diagnostic]] = {}
|
||||
self._initialized = False
|
||||
self._reader_task: Optional[asyncio.Task] = None
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if the language server is running."""
|
||||
return self._process is not None and self._process.returncode is None
|
||||
|
||||
async def start(self) -> bool:
|
||||
"""
|
||||
Start the language server.
|
||||
|
||||
Returns:
|
||||
True if started successfully
|
||||
"""
|
||||
if self.is_running:
|
||||
return True
|
||||
|
||||
try:
|
||||
cmd = [self.config.command] + self.config.args
|
||||
|
||||
self._process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
# Start reader task
|
||||
self._reader_task = asyncio.create_task(self._read_responses())
|
||||
|
||||
# Initialize the server
|
||||
root_uri = self.config.root_uri or f"file://{os.getcwd()}"
|
||||
|
||||
result = await self._send_request("initialize", {
|
||||
"processId": os.getpid(),
|
||||
"rootUri": root_uri,
|
||||
"capabilities": {
|
||||
"textDocument": {
|
||||
"completion": {"completionItem": {"snippetSupport": True}},
|
||||
"hover": {},
|
||||
"definition": {},
|
||||
"references": {},
|
||||
"publishDiagnostics": {}
|
||||
}
|
||||
},
|
||||
"initializationOptions": self.config.initialization_options
|
||||
})
|
||||
|
||||
if result:
|
||||
await self._send_notification("initialized", {})
|
||||
self._initialized = True
|
||||
logger.info(f"LSP server started: {self.config.command}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except FileNotFoundError:
|
||||
logger.error(f"Language server not found: {self.config.command}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start LSP server: {e}")
|
||||
return False
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the language server."""
|
||||
if not self.is_running:
|
||||
return
|
||||
|
||||
try:
|
||||
await self._send_request("shutdown", None)
|
||||
await self._send_notification("exit", None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if self._reader_task:
|
||||
self._reader_task.cancel()
|
||||
try:
|
||||
await self._reader_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self._process:
|
||||
self._process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(self._process.wait(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
self._process.kill()
|
||||
|
||||
self._process = None
|
||||
self._initialized = False
|
||||
|
||||
async def open_document(self, file_path: str, text: Optional[str] = None) -> bool:
|
||||
"""
|
||||
Open a document in the language server.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
text: File content (reads from disk if not provided)
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if not self._initialized:
|
||||
return False
|
||||
|
||||
if text is None:
|
||||
try:
|
||||
with open(file_path, "r") as f:
|
||||
text = f.read()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to read file: {e}")
|
||||
return False
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
await self._send_notification("textDocument/didOpen", {
|
||||
"textDocument": TextDocumentItem(
|
||||
uri=uri,
|
||||
language_id=self.config.language,
|
||||
version=1,
|
||||
text=text
|
||||
).to_dict()
|
||||
})
|
||||
|
||||
return True
|
||||
|
||||
async def close_document(self, file_path: str):
|
||||
"""Close a document in the language server."""
|
||||
if not self._initialized:
|
||||
return
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
await self._send_notification("textDocument/didClose", {
|
||||
"textDocument": TextDocumentIdentifier(uri=uri).to_dict()
|
||||
})
|
||||
|
||||
async def get_diagnostics(self, file_path: str) -> List[Diagnostic]:
|
||||
"""
|
||||
Get diagnostics for a file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
List of diagnostics
|
||||
"""
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
return self._diagnostics.get(uri, [])
|
||||
|
||||
async def get_completions(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int
|
||||
) -> List[CompletionItem]:
|
||||
"""
|
||||
Get completions at a position.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
|
||||
Returns:
|
||||
List of completion items
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/completion", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
items = result.get("items", result) if isinstance(result, dict) else result
|
||||
|
||||
return [
|
||||
CompletionItem.from_dict(item)
|
||||
for item in (items if isinstance(items, list) else [])
|
||||
]
|
||||
|
||||
async def get_definition(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int
|
||||
) -> List[Location]:
|
||||
"""
|
||||
Get definition location(s) for a symbol.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
|
||||
Returns:
|
||||
List of locations
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/definition", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
if isinstance(result, dict):
|
||||
return [Location.from_dict(result)]
|
||||
elif isinstance(result, list):
|
||||
return [Location.from_dict(loc) for loc in result]
|
||||
|
||||
return []
|
||||
|
||||
async def get_references(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int,
|
||||
include_declaration: bool = True
|
||||
) -> List[Location]:
|
||||
"""
|
||||
Get references to a symbol.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
include_declaration: Include the declaration in results
|
||||
|
||||
Returns:
|
||||
List of locations
|
||||
"""
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/references", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character},
|
||||
"context": {"includeDeclaration": include_declaration}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
return [Location.from_dict(loc) for loc in result]
|
||||
|
||||
async def get_hover(
|
||||
self,
|
||||
file_path: str,
|
||||
line: int,
|
||||
character: int
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Get hover information for a position.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
line: Line number (0-indexed)
|
||||
character: Character position (0-indexed)
|
||||
|
||||
Returns:
|
||||
Hover content as string
|
||||
"""
|
||||
if not self._initialized:
|
||||
return None
|
||||
|
||||
uri = f"file://{os.path.abspath(file_path)}"
|
||||
|
||||
result = await self._send_request("textDocument/hover", {
|
||||
"textDocument": {"uri": uri},
|
||||
"position": {"line": line, "character": character}
|
||||
})
|
||||
|
||||
if result is None:
|
||||
return None
|
||||
|
||||
contents = result.get("contents")
|
||||
if isinstance(contents, str):
|
||||
return contents
|
||||
elif isinstance(contents, dict):
|
||||
return contents.get("value", str(contents))
|
||||
elif isinstance(contents, list):
|
||||
return "\n".join(
|
||||
c.get("value", str(c)) if isinstance(c, dict) else str(c)
|
||||
for c in contents
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
async def _send_request(self, method: str, params: Any) -> Any:
|
||||
"""Send a request and wait for response."""
|
||||
if not self.is_running:
|
||||
return None
|
||||
|
||||
self._request_id += 1
|
||||
request_id = self._request_id
|
||||
|
||||
message = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"method": method,
|
||||
"params": params
|
||||
}
|
||||
|
||||
future = asyncio.get_event_loop().create_future()
|
||||
self._pending_requests[request_id] = future
|
||||
|
||||
await self._send_message(message)
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(future, timeout=self.config.timeout)
|
||||
return result
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(f"Request timed out: {method}")
|
||||
self._pending_requests.pop(request_id, None)
|
||||
return None
|
||||
|
||||
async def _send_notification(self, method: str, params: Any):
|
||||
"""Send a notification (no response expected)."""
|
||||
if not self.is_running:
|
||||
return
|
||||
|
||||
message = {
|
||||
"jsonrpc": "2.0",
|
||||
"method": method,
|
||||
"params": params
|
||||
}
|
||||
|
||||
await self._send_message(message)
|
||||
|
||||
async def _send_message(self, message: Dict[str, Any]):
|
||||
"""Send a JSON-RPC message."""
|
||||
if self._process is None or self._process.stdin is None:
|
||||
return
|
||||
|
||||
content = json.dumps(message)
|
||||
header = f"Content-Length: {len(content)}\r\n\r\n"
|
||||
|
||||
self._process.stdin.write(header.encode() + content.encode())
|
||||
await self._process.stdin.drain()
|
||||
|
||||
async def _read_responses(self):
|
||||
"""Read responses from the language server."""
|
||||
if self._process is None or self._process.stdout is None:
|
||||
return
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Read header
|
||||
header = b""
|
||||
while b"\r\n\r\n" not in header:
|
||||
chunk = await self._process.stdout.read(1)
|
||||
if not chunk:
|
||||
return
|
||||
header += chunk
|
||||
|
||||
# Parse content length
|
||||
content_length = 0
|
||||
for line in header.decode().split("\r\n"):
|
||||
if line.startswith("Content-Length:"):
|
||||
content_length = int(line.split(":")[1].strip())
|
||||
break
|
||||
|
||||
if content_length == 0:
|
||||
continue
|
||||
|
||||
# Read content
|
||||
content = await self._process.stdout.read(content_length)
|
||||
message = json.loads(content.decode())
|
||||
|
||||
await self._handle_message(message)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error reading LSP response: {e}")
|
||||
|
||||
async def _handle_message(self, message: Dict[str, Any]):
|
||||
"""Handle an incoming message."""
|
||||
if "id" in message:
|
||||
# Response to a request
|
||||
request_id = message["id"]
|
||||
if request_id in self._pending_requests:
|
||||
future = self._pending_requests.pop(request_id)
|
||||
if "error" in message:
|
||||
future.set_exception(
|
||||
RuntimeError(message["error"].get("message", "Unknown error"))
|
||||
)
|
||||
else:
|
||||
future.set_result(message.get("result"))
|
||||
|
||||
elif "method" in message:
|
||||
# Notification from server
|
||||
method = message["method"]
|
||||
params = message.get("params", {})
|
||||
|
||||
if method == "textDocument/publishDiagnostics":
|
||||
uri = params.get("uri", "")
|
||||
diagnostics = [
|
||||
Diagnostic.from_dict(d)
|
||||
for d in params.get("diagnostics", [])
|
||||
]
|
||||
self._diagnostics[uri] = diagnostics
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
LSP Configuration for PraisonAI Agents.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, List
|
||||
|
||||
|
||||
# Default language server commands
|
||||
DEFAULT_SERVERS = {
|
||||
"python": {
|
||||
"command": "pylsp",
|
||||
"args": [],
|
||||
},
|
||||
"javascript": {
|
||||
"command": "typescript-language-server",
|
||||
"args": ["--stdio"],
|
||||
},
|
||||
"typescript": {
|
||||
"command": "typescript-language-server",
|
||||
"args": ["--stdio"],
|
||||
},
|
||||
"rust": {
|
||||
"command": "rust-analyzer",
|
||||
"args": [],
|
||||
},
|
||||
"go": {
|
||||
"command": "gopls",
|
||||
"args": [],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class LSPConfig:
|
||||
"""Configuration for LSP client."""
|
||||
language: str
|
||||
command: Optional[str] = None
|
||||
args: List[str] = field(default_factory=list)
|
||||
root_uri: Optional[str] = None
|
||||
initialization_options: Dict = field(default_factory=dict)
|
||||
timeout: float = 30.0
|
||||
|
||||
def __post_init__(self):
|
||||
# Use default server if not specified
|
||||
if self.command is None:
|
||||
if self.language in DEFAULT_SERVERS:
|
||||
server = DEFAULT_SERVERS[self.language]
|
||||
self.command = server["command"]
|
||||
self.args = server["args"]
|
||||
else:
|
||||
raise ValueError(f"No default server for language: {self.language}")
|
||||
@@ -0,0 +1,162 @@
|
||||
"""
|
||||
LSP Types for PraisonAI Agents.
|
||||
|
||||
Defines the core types for Language Server Protocol integration.
|
||||
"""
|
||||
|
||||
from enum import IntEnum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
|
||||
class DiagnosticSeverity(IntEnum):
|
||||
"""Severity of a diagnostic."""
|
||||
ERROR = 1
|
||||
WARNING = 2
|
||||
INFORMATION = 3
|
||||
HINT = 4
|
||||
|
||||
|
||||
@dataclass
|
||||
class Position:
|
||||
"""Position in a text document (0-indexed)."""
|
||||
line: int
|
||||
character: int
|
||||
|
||||
def to_dict(self) -> Dict[str, int]:
|
||||
return {"line": self.line, "character": self.character}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, int]) -> "Position":
|
||||
return cls(line=data["line"], character=data["character"])
|
||||
|
||||
|
||||
@dataclass
|
||||
class Range:
|
||||
"""Range in a text document."""
|
||||
start: Position
|
||||
end: Position
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {"start": self.start.to_dict(), "end": self.end.to_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Range":
|
||||
return cls(
|
||||
start=Position.from_dict(data["start"]),
|
||||
end=Position.from_dict(data["end"])
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Location:
|
||||
"""Location in a document."""
|
||||
uri: str
|
||||
range: Range
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {"uri": self.uri, "range": self.range.to_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Location":
|
||||
return cls(uri=data["uri"], range=Range.from_dict(data["range"]))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Diagnostic:
|
||||
"""Represents a diagnostic (error, warning, etc.)."""
|
||||
range: Range
|
||||
message: str
|
||||
severity: DiagnosticSeverity = DiagnosticSeverity.ERROR
|
||||
code: Optional[str] = None
|
||||
source: Optional[str] = None
|
||||
related_information: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"range": self.range.to_dict(),
|
||||
"message": self.message,
|
||||
"severity": self.severity.value,
|
||||
"code": self.code,
|
||||
"source": self.source
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Diagnostic":
|
||||
return cls(
|
||||
range=Range.from_dict(data["range"]),
|
||||
message=data["message"],
|
||||
severity=DiagnosticSeverity(data.get("severity", 1)),
|
||||
code=data.get("code"),
|
||||
source=data.get("source")
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompletionItem:
|
||||
"""Represents a completion item."""
|
||||
label: str
|
||||
kind: int = 1 # Text
|
||||
detail: Optional[str] = None
|
||||
documentation: Optional[str] = None
|
||||
insert_text: Optional[str] = None
|
||||
sort_text: Optional[str] = None
|
||||
filter_text: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"label": self.label,
|
||||
"kind": self.kind,
|
||||
"detail": self.detail,
|
||||
"documentation": self.documentation,
|
||||
"insertText": self.insert_text
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "CompletionItem":
|
||||
return cls(
|
||||
label=data["label"],
|
||||
kind=data.get("kind", 1),
|
||||
detail=data.get("detail"),
|
||||
documentation=data.get("documentation"),
|
||||
insert_text=data.get("insertText")
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextDocumentItem:
|
||||
"""Represents a text document."""
|
||||
uri: str
|
||||
language_id: str
|
||||
version: int
|
||||
text: str
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"uri": self.uri,
|
||||
"languageId": self.language_id,
|
||||
"version": self.version,
|
||||
"text": self.text
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextDocumentIdentifier:
|
||||
"""Identifies a text document."""
|
||||
uri: str
|
||||
|
||||
def to_dict(self) -> Dict[str, str]:
|
||||
return {"uri": self.uri}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextDocumentPositionParams:
|
||||
"""Parameters for text document position requests."""
|
||||
text_document: TextDocumentIdentifier
|
||||
position: Position
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"textDocument": self.text_document.to_dict(),
|
||||
"position": self.position.to_dict()
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
"""
|
||||
Observability interface for PraisonAI Agents.
|
||||
|
||||
This module provides the protocol and types for observability/tracing.
|
||||
Implementations are provided by the wrapper layer (praisonai_tools.observability).
|
||||
|
||||
Usage (simplest - recommended):
|
||||
from praisonaiagents import Agent, obs
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
observability=obs.auto(), # Auto-detect from env vars
|
||||
)
|
||||
agent.chat("Hello!") # Auto-traces to configured provider
|
||||
|
||||
Alternative (explicit provider):
|
||||
observability=obs.langfuse() # Langfuse
|
||||
observability=obs.langsmith() # LangSmith
|
||||
observability=obs.agentops() # AgentOps
|
||||
observability=obs.arize() # Arize Phoenix
|
||||
observability=obs.datadog() # Datadog
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from dataclasses import dataclass, field
|
||||
import time
|
||||
|
||||
|
||||
@dataclass
|
||||
class Span:
|
||||
"""Represents a span in a trace."""
|
||||
span_id: str
|
||||
trace_id: str
|
||||
name: str
|
||||
parent_span_id: Optional[str] = None
|
||||
start_time: float = field(default_factory=time.time)
|
||||
end_time: Optional[float] = None
|
||||
attributes: Dict[str, Any] = field(default_factory=dict)
|
||||
status: str = "ok"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Trace:
|
||||
"""Represents a complete trace."""
|
||||
trace_id: str
|
||||
session_id: Optional[str] = None
|
||||
agent_name: Optional[str] = None
|
||||
user_id: Optional[str] = None
|
||||
start_time: float = field(default_factory=time.time)
|
||||
end_time: Optional[float] = None
|
||||
spans: List[Span] = field(default_factory=list)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
status: str = "ok"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ObservabilityAdapter(Protocol):
|
||||
"""Protocol for observability adapters."""
|
||||
|
||||
def on_trace_start(
|
||||
self,
|
||||
trace_id: str,
|
||||
session_id: Optional[str] = None,
|
||||
agent_name: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Called when a trace starts."""
|
||||
...
|
||||
|
||||
def on_trace_end(
|
||||
self,
|
||||
trace_id: str,
|
||||
status: str = "ok",
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Called when a trace ends."""
|
||||
...
|
||||
|
||||
def on_span_start(
|
||||
self,
|
||||
span_id: str,
|
||||
trace_id: str,
|
||||
name: str,
|
||||
parent_span_id: Optional[str] = None,
|
||||
attributes: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Called when a span starts."""
|
||||
...
|
||||
|
||||
def on_span_end(
|
||||
self,
|
||||
span_id: str,
|
||||
status: str = "ok",
|
||||
attributes: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Called when a span ends."""
|
||||
...
|
||||
|
||||
def on_llm_call(
|
||||
self,
|
||||
span_id: str,
|
||||
model: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
response: Optional[str] = None,
|
||||
tokens: Optional[Dict[str, int]] = None,
|
||||
latency_ms: Optional[float] = None,
|
||||
) -> None:
|
||||
"""Called when an LLM call is made."""
|
||||
...
|
||||
|
||||
def on_tool_call(
|
||||
self,
|
||||
span_id: str,
|
||||
tool_name: str,
|
||||
args: Dict[str, Any],
|
||||
result: Any,
|
||||
latency_ms: Optional[float] = None,
|
||||
) -> None:
|
||||
"""Called when a tool is invoked."""
|
||||
...
|
||||
|
||||
def flush(self) -> None:
|
||||
"""Flush any pending data."""
|
||||
...
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the adapter and release resources."""
|
||||
...
|
||||
|
||||
|
||||
class _LazyObsModule:
|
||||
"""
|
||||
Lazy proxy for observability backends.
|
||||
|
||||
Allows `from praisonaiagents import obs` without importing heavy deps.
|
||||
Actual provider classes are loaded only when accessed.
|
||||
"""
|
||||
|
||||
_PROVIDERS = {
|
||||
"langfuse": "praisonai_tools.observability.providers.langfuse_provider",
|
||||
"langsmith": "praisonai_tools.observability.providers.langsmith_provider",
|
||||
"agentops": "praisonai_tools.observability.providers.agentops_provider",
|
||||
"arize": "praisonai_tools.observability.providers.arize_phoenix_provider",
|
||||
"phoenix": "praisonai_tools.observability.providers.arize_phoenix_provider",
|
||||
"datadog": "praisonai_tools.observability.providers.datadog_provider",
|
||||
"mlflow": "praisonai_tools.observability.providers.mlflow_provider",
|
||||
"openlit": "praisonai_tools.observability.providers.openlit_provider",
|
||||
"opik": "praisonai_tools.observability.providers.opik_provider",
|
||||
"portkey": "praisonai_tools.observability.providers.portkey_provider",
|
||||
"traceloop": "praisonai_tools.observability.providers.traceloop_provider",
|
||||
"weave": "praisonai_tools.observability.providers.weave_provider",
|
||||
"braintrust": "praisonai_tools.observability.providers.braintrust_provider",
|
||||
"langtrace": "praisonai_tools.observability.providers.langtrace_provider",
|
||||
"langwatch": "praisonai_tools.observability.providers.langwatch_provider",
|
||||
"maxim": "praisonai_tools.observability.providers.maxim_provider",
|
||||
}
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name in self._PROVIDERS:
|
||||
return self._create_provider_factory(name)
|
||||
raise AttributeError(f"module 'obs' has no attribute {name!r}")
|
||||
|
||||
def _create_provider_factory(self, provider_name: str):
|
||||
"""Create a factory function for a provider."""
|
||||
def factory(**kwargs):
|
||||
try:
|
||||
import importlib
|
||||
module = importlib.import_module(self._PROVIDERS[provider_name])
|
||||
# Get the provider class (naming convention: XxxProvider)
|
||||
class_name = f"{provider_name.title().replace('_', '')}Provider"
|
||||
provider_cls = getattr(module, class_name, None)
|
||||
if provider_cls:
|
||||
return provider_cls(**kwargs)
|
||||
# Fallback: try to find any class ending with Provider
|
||||
for attr_name in dir(module):
|
||||
if attr_name.endswith("Provider"):
|
||||
return getattr(module, attr_name)(**kwargs)
|
||||
raise AttributeError(f"No provider class found in {module}")
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
f"Observability provider '{provider_name}' requires praisonai-tools. "
|
||||
f"Install with: pip install praisonai-tools\n"
|
||||
f"Original error: {e}"
|
||||
) from e
|
||||
return factory
|
||||
|
||||
def auto(self, **kwargs):
|
||||
"""
|
||||
Auto-detect observability provider from environment variables.
|
||||
|
||||
Checks for common env vars like LANGFUSE_PUBLIC_KEY, LANGSMITH_API_KEY, etc.
|
||||
Returns None if no provider is detected.
|
||||
"""
|
||||
import os
|
||||
|
||||
# Check for common provider env vars
|
||||
provider_env_map = {
|
||||
"LANGFUSE_PUBLIC_KEY": "langfuse",
|
||||
"LANGSMITH_API_KEY": "langsmith",
|
||||
"AGENTOPS_API_KEY": "agentops",
|
||||
"ARIZE_API_KEY": "arize",
|
||||
"DATADOG_API_KEY": "datadog",
|
||||
"MLFLOW_TRACKING_URI": "mlflow",
|
||||
"OPENLIT_API_KEY": "openlit",
|
||||
"OPIK_API_KEY": "opik",
|
||||
"PORTKEY_API_KEY": "portkey",
|
||||
"TRACELOOP_API_KEY": "traceloop",
|
||||
"WANDB_API_KEY": "weave",
|
||||
"BRAINTRUST_API_KEY": "braintrust",
|
||||
}
|
||||
|
||||
for env_var, provider in provider_env_map.items():
|
||||
if os.getenv(env_var):
|
||||
return self._create_provider_factory(provider)(**kwargs)
|
||||
|
||||
return None
|
||||
|
||||
def __repr__(self):
|
||||
return "<module 'praisonaiagents.obs' (lazy observability backends)>"
|
||||
|
||||
def __call__(self, provider: str = "auto", **kwargs):
|
||||
"""
|
||||
Shortcut: obs("langfuse") or obs() for auto-detection.
|
||||
"""
|
||||
if provider == "auto":
|
||||
return self.auto(**kwargs)
|
||||
return self._create_provider_factory(provider)(**kwargs)
|
||||
|
||||
|
||||
# Singleton lazy module instance
|
||||
obs = _LazyObsModule()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ObservabilityAdapter",
|
||||
"Span",
|
||||
"Trace",
|
||||
"obs",
|
||||
]
|
||||
Loaded 100 of 251 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user