mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 17:33:33 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
08b21dabe3 | ||
|
|
8c70d840df | ||
|
|
d8e944c1d3 | ||
|
|
79b2b88f72 | ||
|
|
f5dfbd5670 | ||
|
|
aa6aa5933d | ||
|
|
ff3a4f2016 | ||
|
|
0d5444fc23 | ||
|
|
d7d5dab17b | ||
|
|
05678394e6 | ||
|
|
54cddc09b4 | ||
|
|
9cbab01785 | ||
|
|
de9d4bd3c6 | ||
|
|
aba2265e5e |
No files matched your search
+3
-1
@@ -80,4 +80,6 @@ CopilotKit*
|
||||
issue_details.json
|
||||
|
||||
# MCP Registry tokens
|
||||
.mcpregistry_*
|
||||
.mcpregistry_*
|
||||
|
||||
*.log
|
||||
@@ -35,46 +35,108 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
|
||||
|
||||
---
|
||||
|
||||
> **Quick Paths:**
|
||||
> - 🚀 **New here?** → [Quick Start](#-quick-start)
|
||||
> - 📦 **Installing?** → [Installation](#-installation)
|
||||
> - 💻 **Python SDK?** → [Python Examples](#-using-python-code)
|
||||
> - 🎯 **CLI user?** → [CLI Reference](#-cli--no-code-interface)
|
||||
> - 🔧 **Need config?** → [Configuration](#-configuration--integration)
|
||||
> - 🤝 **Contributing?** → [Development](#-development)
|
||||
|
||||
---
|
||||
|
||||
## 📑 Table of Contents
|
||||
|
||||
### Getting Started
|
||||
- [⚡ Performance](#-performance)
|
||||
- [🚀 Quick Start](#-quick-start)
|
||||
- [📦 Installation](#-installation)
|
||||
- [Python SDK](#python-sdk)
|
||||
- [JavaScript SDK](#javascript-sdk)
|
||||
- [💻 Usage](#-usage)
|
||||
- [Python Code Examples](#python-code-examples)
|
||||
- [CLI / No-Code Interface](#cli--no-code-interface)
|
||||
- [JavaScript Code Examples](#javascript-code-examples)
|
||||
- [Environment Variables](#environment-variables)
|
||||
|
||||
### Core Concepts
|
||||
- [✨ Key Features](#-key-features)
|
||||
- [🌐 Supported Providers](#-supported-providers)
|
||||
|
||||
### Python SDK
|
||||
- [📘 Python Examples](#-using-python-code)
|
||||
- [1. Single Agent](#1-single-agent)
|
||||
- [2. Multi Agents](#2-multi-agents)
|
||||
- [3. Planning Mode](#3-agent-with-planning-mode)
|
||||
- [4. Deep Research](#4-deep-research-agent)
|
||||
- [5. Query Rewriter](#5-query-rewriter-agent)
|
||||
- [6. Agent Memory](#6-agent-memory-zero-dependencies)
|
||||
- [7. Rules & Instructions](#7-rules--instructions)
|
||||
- [8. Auto-Generated Memories](#8-auto-generated-memories)
|
||||
- [9. Agentic Workflows](#9-agentic-workflows)
|
||||
- [10. Hooks](#10-hooks)
|
||||
- [11. Shadow Git Checkpoints](#11-shadow-git-checkpoints)
|
||||
- [12. Background Tasks](#12-background-tasks)
|
||||
- [13. Policy Engine](#13-policy-engine)
|
||||
- [14. Thinking Budgets](#14-thinking-budgets)
|
||||
- [15. Output Styles](#15-output-styles)
|
||||
- [16. Context Compaction](#16-context-compaction)
|
||||
- [17. Field Names Reference](#17-field-names-reference-a-i-g-s)
|
||||
- [18. Extended agents.yaml](#18-extended-agentsyaml-with-workflow-patterns)
|
||||
- [19. MCP Protocol](#19-mcp-model-context-protocol)
|
||||
- [20. A2A Protocol](#20-a2a-agent2agent-protocol)
|
||||
- [🛠️ Custom Tools](#️-custom-tools)
|
||||
|
||||
### JavaScript SDK
|
||||
- [💻 JavaScript Examples](#-using-javascript-code)
|
||||
|
||||
### CLI Reference
|
||||
- [🎯 CLI Overview](#-cli--no-code-interface)
|
||||
- [Auto Mode](#auto-mode)
|
||||
- [Interactive Mode](#interactive-mode-cli)
|
||||
- [Deep Research CLI](#deep-research-cli)
|
||||
- [Planning Mode CLI](#planning-mode-cli)
|
||||
- [Memory CLI](#memory-cli)
|
||||
- [Workflow CLI](#workflow-cli)
|
||||
- [Knowledge CLI](#knowledge-cli)
|
||||
- [Session CLI](#session-cli)
|
||||
- [Tools CLI](#tools-cli)
|
||||
- [MCP Config CLI](#mcp-config-cli)
|
||||
- [External Agents CLI](#external-agents-cli)
|
||||
- [CLI Features Summary](#cli-features)
|
||||
|
||||
### Configuration
|
||||
- [🔧 Configuration & Integration](#-configuration--integration)
|
||||
- [Ollama Integration](#ollama-integration)
|
||||
- [Groq Integration](#groq-integration)
|
||||
- [100+ Models Support](#100-models-support)
|
||||
- [📋 Workflows](#-workflows)
|
||||
- [Workflow Patterns](#workflow-patterns)
|
||||
- [YAML Configuration](#yaml-configuration)
|
||||
- [Agents Playbook](#agents-playbook)
|
||||
- [🧠 Memory & Context](#-memory--context)
|
||||
- [📋 Agents Playbook](#-agents-playbook)
|
||||
|
||||
### Advanced Features
|
||||
- [🔬 Advanced Features](#-advanced-features)
|
||||
- [Research & Intelligence](#research--intelligence)
|
||||
- [Memory & Caching](#memory--caching)
|
||||
- [Planning & Workflows](#planning--workflows)
|
||||
- [Specialised Agents](#specialised-agents)
|
||||
- [MCP Protocol](#mcp-protocol)
|
||||
- [A2A & A2UI Protocols](#a2a--a2ui-protocols)
|
||||
- [Safety & Control](#safety--control)
|
||||
- [Developer Tools](#developer-tools)
|
||||
|
||||
### Architecture
|
||||
- [📊 Process Types & Patterns](#-process-types--patterns)
|
||||
- [Sequential Process](#sequential-process)
|
||||
- [Hierarchical Process](#hierarchical-process)
|
||||
- [Workflow Process](#workflow-process)
|
||||
- [Agentic Patterns](#agentic-patterns)
|
||||
|
||||
### Data & Persistence
|
||||
- [💾 Persistence (Databases)](#-persistence-databases)
|
||||
- [📚 Knowledge & Retrieval (RAG)](#-knowledge--retrieval-rag)
|
||||
- [🔧 Tools Table](#-tools-table)
|
||||
|
||||
### Learning & Community
|
||||
- [🎓 Video Tutorials](#-video-tutorials)
|
||||
- [⭐ Star History](#-star-history)
|
||||
- [👥 Contributing](#-contributing)
|
||||
- [🔧 Development](#-development)
|
||||
- [❓ FAQ & Troubleshooting](#-faq--troubleshooting)
|
||||
|
||||
---
|
||||
|
||||
@@ -143,6 +205,30 @@ pip install praisonai
|
||||
npm install praisonai
|
||||
```
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Required | Description |
|
||||
|----------|----------|-------------|
|
||||
| `OPENAI_API_KEY` | Yes* | OpenAI API key |
|
||||
| `ANTHROPIC_API_KEY` | No | Anthropic Claude API key |
|
||||
| `GOOGLE_API_KEY` | No | Google Gemini API key |
|
||||
| `GROQ_API_KEY` | No | Groq API key |
|
||||
| `OPENAI_BASE_URL` | No | Custom API endpoint (for Ollama, Groq, etc.) |
|
||||
|
||||
> *At least one LLM provider API key is required.
|
||||
|
||||
```bash
|
||||
# Set your API key
|
||||
export OPENAI_API_KEY=your_key_here
|
||||
|
||||
# For Ollama (local models)
|
||||
export OPENAI_BASE_URL=http://localhost:11434/v1
|
||||
|
||||
# For Groq
|
||||
export OPENAI_API_KEY=your_groq_key
|
||||
export OPENAI_BASE_URL=https://api.groq.com/openai/v1
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 💻 Usage
|
||||
@@ -169,137 +255,263 @@ npm install praisonai
|
||||
|
||||
## ✨ Key Features
|
||||
|
||||
<details open>
|
||||
<summary><strong>🤖 Core Agents</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| **🤖 Core Agents** | | |
|
||||
| ↳ Single Agent | [Example](examples/python/agents/single-agent.py) | [📖](https://docs.praison.ai/agents/single) |
|
||||
| ↳ Multi Agents | [Example](examples/python/general/mini_agents_example.py) | [📖](https://docs.praison.ai/concepts/agents) |
|
||||
| ↳ Auto Agents | [Example](examples/python/general/auto_agents_example.py) | [📖](https://docs.praison.ai/features/autoagents) |
|
||||
| ↳ Self Reflection AI Agents | [Example](examples/python/concepts/self-reflection-details.py) | [📖](https://docs.praison.ai/features/selfreflection) |
|
||||
| ↳ Reasoning AI Agents | [Example](examples/python/concepts/reasoning-extraction.py) | [📖](https://docs.praison.ai/features/reasoning) |
|
||||
| ↳ Multi Modal AI Agents | [Example](examples/python/general/multimodal.py) | [📖](https://docs.praison.ai/features/multimodal) |
|
||||
| **🔄 Workflows** | | |
|
||||
| ↳ Simple Workflow | [Example](examples/python/workflows/simple_workflow.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| ↳ Workflow with Agents | [Example](examples/python/workflows/workflow_with_agents.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| ↳ Agentic Routing (`route()`) | [Example](examples/python/workflows/workflow_routing.py) | [📖](https://docs.praison.ai/features/routing) |
|
||||
| ↳ Parallel Execution (`parallel()`) | [Example](examples/python/workflows/workflow_parallel.py) | [📖](https://docs.praison.ai/features/parallelisation) |
|
||||
| ↳ Loop over List/CSV (`loop()`) | [Example](examples/python/workflows/workflow_loop_csv.py) | [📖](https://docs.praison.ai/features/repetitive) |
|
||||
| ↳ Evaluator-Optimizer (`repeat()`) | [Example](examples/python/workflows/workflow_repeat.py) | [📖](https://docs.praison.ai/features/evaluator-optimiser) |
|
||||
| ↳ Conditional Steps | [Example](examples/python/workflows/workflow_conditional.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| ↳ Workflow Branching | [Example](examples/python/workflows/workflow_branching.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| ↳ Workflow Early Stop | [Example](examples/python/workflows/workflow_early_stop.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| ↳ Workflow Checkpoints | [Example](examples/python/workflows/workflow_checkpoints.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| **💻 Code & Development** | | |
|
||||
| ↳ Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [📖](https://docs.praison.ai/features/codeagent) |
|
||||
| ↳ AI Code Editing Tools | [Example](examples/python/code/code_editing_example.py) | [📖](https://docs.praison.ai/code/editing) |
|
||||
| ↳ External Agents (All) | [Example](examples/python/code/external_agents_example.py) | [📖](https://docs.praison.ai/code/external-agents) |
|
||||
| ↳ Claude Code CLI | [Example](examples/python/code/claude_code_example.py) | [📖](https://docs.praison.ai/code/claude-code) |
|
||||
| ↳ Gemini CLI | [Example](examples/python/code/gemini_cli_example.py) | [📖](https://docs.praison.ai/code/gemini-cli) |
|
||||
| ↳ Codex CLI | [Example](examples/python/code/codex_cli_example.py) | [📖](https://docs.praison.ai/code/codex-cli) |
|
||||
| ↳ Cursor CLI | [Example](examples/python/code/cursor_cli_example.py) | [📖](https://docs.praison.ai/code/cursor-cli) |
|
||||
| **🧠 Memory & Knowledge** | | |
|
||||
| ↳ Memory (Short & Long Term) | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
|
||||
| ↳ File-Based Memory | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
|
||||
| ↳ Claude Memory Tool | [Example](#claude-memory-tool-cli) | [📖](https://docs.praison.ai/features/claude-memory-tool) |
|
||||
| ↳ 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) |
|
||||
| **🔬 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) |
|
||||
| ↳ Native Web Search | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/agents/websearch) |
|
||||
| ↳ Built-in Search Tools | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/tools/tavily) |
|
||||
| ↳ Unified Web Search | [Example](src/praisonai-agents/examples/web_search_example.py) | [📖](https://docs.praison.ai/tools/web-search) |
|
||||
| ↳ Web Fetch (Anthropic) | [Example](#web-search-web-fetch--prompt-caching) | [📖](https://docs.praison.ai/features/model-capabilities) |
|
||||
| **📋 Planning & Execution** | | |
|
||||
| ↳ Planning Mode | [Example](examples/python/agents/planning-agent.py) | [📖](https://docs.praison.ai/features/planning-mode) |
|
||||
| ↳ Planning Tools | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
|
||||
| ↳ Planning Reasoning | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
|
||||
| ↳ Prompt Chaining | [Example](examples/python/general/prompt_chaining.py) | [📖](https://docs.praison.ai/features/promptchaining) |
|
||||
| ↳ Evaluator Optimiser | [Example](examples/python/general/evaluator-optimiser.py) | [📖](https://docs.praison.ai/features/evaluator-optimiser) |
|
||||
| ↳ Orchestrator Workers | [Example](examples/python/general/orchestrator-workers.py) | [📖](https://docs.praison.ai/features/orchestrator-worker) |
|
||||
| **👥 Specialized Agents** | | |
|
||||
| ↳ Data Analyst Agent | [Example](examples/python/agents/data-analyst-agent.py) | [📖](https://docs.praison.ai/agents/data-analyst) |
|
||||
| ↳ Finance Agent | [Example](examples/python/agents/finance-agent.py) | [📖](https://docs.praison.ai/agents/finance) |
|
||||
| ↳ Shopping Agent | [Example](examples/python/agents/shopping-agent.py) | [📖](https://docs.praison.ai/agents/shopping) |
|
||||
| ↳ Recommendation Agent | [Example](examples/python/agents/recommendation-agent.py) | [📖](https://docs.praison.ai/agents/recommendation) |
|
||||
| ↳ Wikipedia Agent | [Example](examples/python/agents/wikipedia-agent.py) | [📖](https://docs.praison.ai/agents/wikipedia) |
|
||||
| ↳ Programming Agent | [Example](examples/python/agents/programming-agent.py) | [📖](https://docs.praison.ai/agents/programming) |
|
||||
| ↳ Math Agents | [Example](examples/python/agents/math-agent.py) | [📖](https://docs.praison.ai/features/mathagent) |
|
||||
| ↳ Markdown Agent | [Example](examples/python/agents/markdown-agent.py) | [📖](https://docs.praison.ai/agents/markdown) |
|
||||
| ↳ Prompt Expander Agent | [Example](#prompt-expansion) | [📖](https://docs.praison.ai/agents/prompt-expander) |
|
||||
| **🎨 Media & Multimodal** | | |
|
||||
| ↳ Image Generation Agent | [Example](examples/python/image/image-agent.py) | [📖](https://docs.praison.ai/features/image-generation) |
|
||||
| ↳ Image to Text Agent | [Example](examples/python/agents/image-to-text-agent.py) | [📖](https://docs.praison.ai/agents/image-to-text) |
|
||||
| ↳ Video Agent | [Example](examples/python/agents/video-agent.py) | [📖](https://docs.praison.ai/agents/video) |
|
||||
| ↳ Camera Integration | [Example](examples/python/camera/) | [📖](https://docs.praison.ai/features/camera-integration) |
|
||||
| **🔌 Protocols & Integration** | | |
|
||||
| ↳ MCP Transports | [Example](examples/python/mcp/mcp-transports-overview.py) | [📖](https://docs.praison.ai/mcp/transports) |
|
||||
| ↳ WebSocket MCP | [Example](examples/python/mcp/websocket-mcp.py) | [📖](https://docs.praison.ai/mcp/sse-transport) |
|
||||
| ↳ MCP Security | [Example](examples/python/mcp/mcp-security.py) | [📖](https://docs.praison.ai/mcp/transports) |
|
||||
| ↳ MCP Resumability | [Example](examples/python/mcp/mcp-resumability.py) | [📖](https://docs.praison.ai/mcp/sse-transport) |
|
||||
| ↳ MCP Config Management | [Example](#mcp-config-cli) | [📖](https://docs.praison.ai/docs/cli/mcp) |
|
||||
| ↳ LangChain Integrated Agents | [Example](examples/python/general/langchain_example.py) | [📖](https://docs.praison.ai/features/langchain) |
|
||||
| **🛡️ Safety & Control** | | |
|
||||
| ↳ Guardrails | [Example](examples/python/guardrails/comprehensive-guardrails-example.py) | [📖](https://docs.praison.ai/features/guardrails) |
|
||||
| ↳ Human Approval | [Example](examples/python/general/human_approval_example.py) | [📖](https://docs.praison.ai/features/approval) |
|
||||
| ↳ Rules & Instructions | [Example](#6-rules--instructions) | [📖](https://docs.praison.ai/features/rules) |
|
||||
| **⚙️ Advanced Features** | | |
|
||||
| ↳ Async & Parallel Processing | [Example](examples/python/general/async_example.py) | [📖](https://docs.praison.ai/features/async) |
|
||||
| ↳ Parallelisation | [Example](examples/python/general/parallelisation.py) | [📖](https://docs.praison.ai/features/parallelisation) |
|
||||
| ↳ Repetitive Agents | [Example](examples/python/concepts/repetitive-agents.py) | [📖](https://docs.praison.ai/features/repetitive) |
|
||||
| ↳ Agent Handoffs | [Example](examples/python/handoff/handoff_basic.py) | [📖](https://docs.praison.ai/features/handoffs) |
|
||||
| ↳ Stateful Agents | [Example](examples/python/stateful/workflow-state-example.py) | [📖](https://docs.praison.ai/features/stateful-agents) |
|
||||
| ↳ Autonomous Workflow | [Example](examples/python/general/autonomous-agent.py) | [📖](https://docs.praison.ai/features/autonomous-workflow) |
|
||||
| ↳ Structured Output Agents | [Example](examples/python/general/structured_agents_example.py) | [📖](https://docs.praison.ai/features/structured) |
|
||||
| ↳ Model Router | [Example](examples/python/agents/router-agent-cost-optimization.py) | [📖](https://docs.praison.ai/features/model-router) |
|
||||
| ↳ Prompt Caching | [Example](#web-search-web-fetch--prompt-caching) | [📖](https://docs.praison.ai/features/model-capabilities) |
|
||||
| ↳ Fast Context | [Example](examples/context/00_agent_fast_context_basic.py) | [📖](https://docs.praison.ai/features/fast-context) |
|
||||
| **🛠️ Tools & Configuration** | | |
|
||||
| ↳ 100+ Custom Tools | [Example](examples/python/general/tools_example.py) | [📖](https://docs.praison.ai/tools/tools) |
|
||||
| ↳ YAML Configuration | [Example](examples/cookbooks/yaml/secondary_market_research_agents.yaml) | [📖](https://docs.praison.ai/developers/agents-playbook) |
|
||||
| ↳ 100+ LLM Support | [Example](examples/python/providers/openai/openai_gpt4_example.py) | [📖](https://docs.praison.ai/models) |
|
||||
| ↳ Callback Agents | [Example](examples/python/general/advanced-callback-systems.py) | [📖](https://docs.praison.ai/features/callbacks) |
|
||||
| ↳ Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
|
||||
| ↳ Shadow Git Checkpoints | [Example](#10-shadow-git-checkpoints) | [📖](https://docs.praison.ai/features/checkpoints) |
|
||||
| ↳ Background Tasks | [Example](examples/background/basic_background.py) | [📖](https://docs.praison.ai/features/background-tasks) |
|
||||
| ↳ Policy Engine | [Example](examples/policy/basic_policy.py) | [📖](https://docs.praison.ai/features/policy-engine) |
|
||||
| ↳ Thinking Budgets | [Example](examples/thinking/basic_thinking.py) | [📖](https://docs.praison.ai/features/thinking-budgets) |
|
||||
| ↳ Output Styles | [Example](examples/output/basic_output.py) | [📖](https://docs.praison.ai/features/output-styles) |
|
||||
| ↳ Context Compaction | [Example](examples/compaction/basic_compaction.py) | [📖](https://docs.praison.ai/features/context-compaction) |
|
||||
| **📊 Monitoring & Management** | | |
|
||||
| ↳ Sessions Management | [Example](examples/python/sessions/comprehensive-session-management.py) | [📖](https://docs.praison.ai/features/sessions) |
|
||||
| ↳ Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| ↳ History in Context | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| ↳ Telemetry | [Example](examples/python/telemetry/production-telemetry-example.py) | [📖](https://docs.praison.ai/features/telemetry) |
|
||||
| ↳ Project Docs (.praison/docs/) | [Example](#docs-cli) | [📖](https://docs.praison.ai/docs/cli/docs) |
|
||||
| ↳ AI Commit Messages | [Example](#ai-commit-cli) | [📖](https://docs.praison.ai/docs/cli/commit) |
|
||||
| ↳ @Mentions in Prompts | [Example](#mentions-in-prompts) | [📖](https://docs.praison.ai/docs/cli/mentions) |
|
||||
| **🖥️ CLI Features** | | |
|
||||
| ↳ Slash Commands | [Example](examples/python/cli/slash_commands_example.py) | [📖](https://docs.praison.ai/docs/cli/slash-commands) |
|
||||
| ↳ Autonomy Modes | [Example](examples/python/cli/autonomy_modes_example.py) | [📖](https://docs.praison.ai/docs/cli/autonomy-modes) |
|
||||
| ↳ Cost Tracking | [Example](examples/python/cli/cost_tracking_example.py) | [📖](https://docs.praison.ai/docs/cli/cost-tracking) |
|
||||
| ↳ Repository Map | [Example](examples/python/cli/repo_map_example.py) | [📖](https://docs.praison.ai/docs/cli/repo-map) |
|
||||
| ↳ Interactive TUI | [Example](examples/python/cli/interactive_tui_example.py) | [📖](https://docs.praison.ai/docs/cli/interactive-tui) |
|
||||
| ↳ 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) |
|
||||
| ↳ Reliability Evaluation | [Example](examples/eval/reliability_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
| ↳ Criteria Evaluation | [Example](examples/eval/criteria_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
| **🎯 Agent Skills** | | |
|
||||
| ↳ Skills Management | [Example](examples/skills/basic_skill_usage.py) | [📖](https://docs.praison.ai/features/skills) |
|
||||
| ↳ Custom Skills | [Example](examples/skills/custom_skill_example.py) | [📖](https://docs.praison.ai/features/skills) |
|
||||
| **⏰ 24/7 Scheduling** | | |
|
||||
| ↳ Agent Scheduler | [Example](examples/python/scheduled_agents/news_checker_live.py) | [📖](https://docs.praison.ai/docs/cli/scheduler) |
|
||||
| Single Agent | [Example](examples/python/agents/single-agent.py) | [📖](https://docs.praison.ai/agents/single) |
|
||||
| Multi Agents | [Example](examples/python/general/mini_agents_example.py) | [📖](https://docs.praison.ai/concepts/agents) |
|
||||
| Auto Agents | [Example](examples/python/general/auto_agents_example.py) | [📖](https://docs.praison.ai/features/autoagents) |
|
||||
| Self Reflection AI Agents | [Example](examples/python/concepts/self-reflection-details.py) | [📖](https://docs.praison.ai/features/selfreflection) |
|
||||
| Reasoning AI Agents | [Example](examples/python/concepts/reasoning-extraction.py) | [📖](https://docs.praison.ai/features/reasoning) |
|
||||
| Multi Modal AI Agents | [Example](examples/python/general/multimodal.py) | [📖](https://docs.praison.ai/features/multimodal) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🔄 Workflows</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Simple Workflow | [Example](examples/python/workflows/simple_workflow.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| Workflow with Agents | [Example](examples/python/workflows/workflow_with_agents.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| Agentic Routing (`route()`) | [Example](examples/python/workflows/workflow_routing.py) | [📖](https://docs.praison.ai/features/routing) |
|
||||
| Parallel Execution (`parallel()`) | [Example](examples/python/workflows/workflow_parallel.py) | [📖](https://docs.praison.ai/features/parallelisation) |
|
||||
| Loop over List/CSV (`loop()`) | [Example](examples/python/workflows/workflow_loop_csv.py) | [📖](https://docs.praison.ai/features/repetitive) |
|
||||
| Evaluator-Optimizer (`repeat()`) | [Example](examples/python/workflows/workflow_repeat.py) | [📖](https://docs.praison.ai/features/evaluator-optimiser) |
|
||||
| Conditional Steps | [Example](examples/python/workflows/workflow_conditional.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| Workflow Branching | [Example](examples/python/workflows/workflow_branching.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| Workflow Early Stop | [Example](examples/python/workflows/workflow_early_stop.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
| Workflow Checkpoints | [Example](examples/python/workflows/workflow_checkpoints.py) | [📖](https://docs.praison.ai/features/workflows) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>💻 Code & Development</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [📖](https://docs.praison.ai/features/codeagent) |
|
||||
| AI Code Editing Tools | [Example](examples/python/code/code_editing_example.py) | [📖](https://docs.praison.ai/code/editing) |
|
||||
| External Agents (All) | [Example](examples/python/code/external_agents_example.py) | [📖](https://docs.praison.ai/code/external-agents) |
|
||||
| Claude Code CLI | [Example](examples/python/code/claude_code_example.py) | [📖](https://docs.praison.ai/code/claude-code) |
|
||||
| Gemini CLI | [Example](examples/python/code/gemini_cli_example.py) | [📖](https://docs.praison.ai/code/gemini-cli) |
|
||||
| Codex CLI | [Example](examples/python/code/codex_cli_example.py) | [📖](https://docs.praison.ai/code/codex-cli) |
|
||||
| Cursor CLI | [Example](examples/python/code/cursor_cli_example.py) | [📖](https://docs.praison.ai/code/cursor-cli) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🧠 Memory & Knowledge</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Memory (Short & Long Term) | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
|
||||
| File-Based Memory | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
|
||||
| Claude Memory Tool | [Example](#claude-memory-tool-cli) | [📖](https://docs.praison.ai/features/claude-memory-tool) |
|
||||
| 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) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🔬 Research & Intelligence</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| 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) |
|
||||
| Native Web Search | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/agents/websearch) |
|
||||
| Built-in Search Tools | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/tools/tavily) |
|
||||
| Unified Web Search | [Example](src/praisonai-agents/examples/web_search_example.py) | [📖](https://docs.praison.ai/tools/web-search) |
|
||||
| Web Fetch (Anthropic) | [Example](#web-search-web-fetch--prompt-caching) | [📖](https://docs.praison.ai/features/model-capabilities) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>📋 Planning & Execution</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Planning Mode | [Example](examples/python/agents/planning-agent.py) | [📖](https://docs.praison.ai/features/planning-mode) |
|
||||
| Planning Tools | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
|
||||
| Planning Reasoning | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
|
||||
| Prompt Chaining | [Example](examples/python/general/prompt_chaining.py) | [📖](https://docs.praison.ai/features/promptchaining) |
|
||||
| Evaluator Optimiser | [Example](examples/python/general/evaluator-optimiser.py) | [📖](https://docs.praison.ai/features/evaluator-optimiser) |
|
||||
| Orchestrator Workers | [Example](examples/python/general/orchestrator-workers.py) | [📖](https://docs.praison.ai/features/orchestrator-worker) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>👥 Specialized Agents</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Data Analyst Agent | [Example](examples/python/agents/data-analyst-agent.py) | [📖](https://docs.praison.ai/agents/data-analyst) |
|
||||
| Finance Agent | [Example](examples/python/agents/finance-agent.py) | [📖](https://docs.praison.ai/agents/finance) |
|
||||
| Shopping Agent | [Example](examples/python/agents/shopping-agent.py) | [📖](https://docs.praison.ai/agents/shopping) |
|
||||
| Recommendation Agent | [Example](examples/python/agents/recommendation-agent.py) | [📖](https://docs.praison.ai/agents/recommendation) |
|
||||
| Wikipedia Agent | [Example](examples/python/agents/wikipedia-agent.py) | [📖](https://docs.praison.ai/agents/wikipedia) |
|
||||
| Programming Agent | [Example](examples/python/agents/programming-agent.py) | [📖](https://docs.praison.ai/agents/programming) |
|
||||
| Math Agents | [Example](examples/python/agents/math-agent.py) | [📖](https://docs.praison.ai/features/mathagent) |
|
||||
| Markdown Agent | [Example](examples/python/agents/markdown-agent.py) | [📖](https://docs.praison.ai/agents/markdown) |
|
||||
| Prompt Expander Agent | [Example](#prompt-expansion) | [📖](https://docs.praison.ai/agents/prompt-expander) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🎨 Media & Multimodal</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Image Generation Agent | [Example](examples/python/image/image-agent.py) | [📖](https://docs.praison.ai/features/image-generation) |
|
||||
| Image to Text Agent | [Example](examples/python/agents/image-to-text-agent.py) | [📖](https://docs.praison.ai/agents/image-to-text) |
|
||||
| Video Agent | [Example](examples/python/agents/video-agent.py) | [📖](https://docs.praison.ai/agents/video) |
|
||||
| Camera Integration | [Example](examples/python/camera/) | [📖](https://docs.praison.ai/features/camera-integration) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🔌 Protocols & Integration</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| MCP Transports | [Example](examples/python/mcp/mcp-transports-overview.py) | [📖](https://docs.praison.ai/mcp/transports) |
|
||||
| WebSocket MCP | [Example](examples/python/mcp/websocket-mcp.py) | [📖](https://docs.praison.ai/mcp/sse-transport) |
|
||||
| MCP Security | [Example](examples/python/mcp/mcp-security.py) | [📖](https://docs.praison.ai/mcp/transports) |
|
||||
| MCP Resumability | [Example](examples/python/mcp/mcp-resumability.py) | [📖](https://docs.praison.ai/mcp/sse-transport) |
|
||||
| MCP Config Management | [Example](#mcp-config-cli) | [📖](https://docs.praison.ai/docs/cli/mcp) |
|
||||
| LangChain Integrated Agents | [Example](examples/python/general/langchain_example.py) | [📖](https://docs.praison.ai/features/langchain) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🛡️ Safety & Control</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Guardrails | [Example](examples/python/guardrails/comprehensive-guardrails-example.py) | [📖](https://docs.praison.ai/features/guardrails) |
|
||||
| Human Approval | [Example](examples/python/general/human_approval_example.py) | [📖](https://docs.praison.ai/features/approval) |
|
||||
| Rules & Instructions | [Example](#7-rules--instructions) | [📖](https://docs.praison.ai/features/rules) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>⚙️ Advanced Features</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Async & Parallel Processing | [Example](examples/python/general/async_example.py) | [📖](https://docs.praison.ai/features/async) |
|
||||
| Parallelisation | [Example](examples/python/general/parallelisation.py) | [📖](https://docs.praison.ai/features/parallelisation) |
|
||||
| Repetitive Agents | [Example](examples/python/concepts/repetitive-agents.py) | [📖](https://docs.praison.ai/features/repetitive) |
|
||||
| Agent Handoffs | [Example](examples/python/handoff/handoff_basic.py) | [📖](https://docs.praison.ai/features/handoffs) |
|
||||
| Stateful Agents | [Example](examples/python/stateful/workflow-state-example.py) | [📖](https://docs.praison.ai/features/stateful-agents) |
|
||||
| Autonomous Workflow | [Example](examples/python/general/autonomous-agent.py) | [📖](https://docs.praison.ai/features/autonomous-workflow) |
|
||||
| Structured Output Agents | [Example](examples/python/general/structured_agents_example.py) | [📖](https://docs.praison.ai/features/structured) |
|
||||
| Model Router | [Example](examples/python/agents/router-agent-cost-optimization.py) | [📖](https://docs.praison.ai/features/model-router) |
|
||||
| Prompt Caching | [Example](#web-search-web-fetch--prompt-caching) | [📖](https://docs.praison.ai/features/model-capabilities) |
|
||||
| Fast Context | [Example](examples/context/00_agent_fast_context_basic.py) | [📖](https://docs.praison.ai/features/fast-context) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🛠️ Tools & Configuration</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| 100+ Custom Tools | [Example](examples/python/general/tools_example.py) | [📖](https://docs.praison.ai/tools/tools) |
|
||||
| YAML Configuration | [Example](examples/cookbooks/yaml/secondary_market_research_agents.yaml) | [📖](https://docs.praison.ai/developers/agents-playbook) |
|
||||
| 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](#10-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](#11-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) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>📊 Monitoring & Management</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| 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) |
|
||||
| History in Context | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| Telemetry | [Example](examples/python/telemetry/production-telemetry-example.py) | [📖](https://docs.praison.ai/features/telemetry) |
|
||||
| Project Docs (.praison/docs/) | [Example](#docs-cli) | [📖](https://docs.praison.ai/docs/cli/docs) |
|
||||
| AI Commit Messages | [Example](#ai-commit-cli) | [📖](https://docs.praison.ai/docs/cli/commit) |
|
||||
| @Mentions in Prompts | [Example](#mentions-in-prompts) | [📖](https://docs.praison.ai/docs/cli/mentions) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🖥️ CLI Features</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Slash Commands | [Example](examples/python/cli/slash_commands_example.py) | [📖](https://docs.praison.ai/docs/cli/slash-commands) |
|
||||
| Autonomy Modes | [Example](examples/python/cli/autonomy_modes_example.py) | [📖](https://docs.praison.ai/docs/cli/autonomy-modes) |
|
||||
| Cost Tracking | [Example](examples/python/cli/cost_tracking_example.py) | [📖](https://docs.praison.ai/docs/cli/cost-tracking) |
|
||||
| Repository Map | [Example](examples/python/cli/repo_map_example.py) | [📖](https://docs.praison.ai/docs/cli/repo-map) |
|
||||
| Interactive TUI | [Example](examples/python/cli/interactive_tui_example.py) | [📖](https://docs.praison.ai/docs/cli/interactive-tui) |
|
||||
| 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) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🧪 Evaluation</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| 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) |
|
||||
| Reliability Evaluation | [Example](examples/eval/reliability_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
| Criteria Evaluation | [Example](examples/eval/criteria_example.py) | [📖](https://docs.praison.ai/docs/cli/eval) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🎯 Agent Skills</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Skills Management | [Example](examples/skills/basic_skill_usage.py) | [📖](https://docs.praison.ai/features/skills) |
|
||||
| Custom Skills | [Example](examples/skills/custom_skill_example.py) | [📖](https://docs.praison.ai/features/skills) |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>⏰ 24/7 Scheduling</strong></summary>
|
||||
|
||||
| Feature | Code | Docs |
|
||||
|---------|:----:|:----:|
|
||||
| Agent Scheduler | [Example](examples/python/scheduled_agents/news_checker_live.py) | [📖](https://docs.praison.ai/docs/cli/scheduler) |
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
@@ -533,7 +745,7 @@ memory.handle_command("/memory save my_session")
|
||||
| `memory="sqlite"` | Built-in | SQLite with indexing |
|
||||
| `memory="chromadb"` | chromadb | Vector/semantic search |
|
||||
|
||||
### 6. Rules & Instructions
|
||||
### 7. Rules & Instructions
|
||||
|
||||
PraisonAI auto-discovers instruction files from your project root and git root:
|
||||
|
||||
@@ -582,7 +794,7 @@ activation: always # always, glob, manual, ai_decision
|
||||
- Follow PEP 8
|
||||
```
|
||||
|
||||
### 7. Auto-Generated Memories
|
||||
### 8. Auto-Generated Memories
|
||||
|
||||
```python
|
||||
from praisonaiagents.memory import FileMemory, AutoMemory
|
||||
@@ -597,7 +809,7 @@ memories = auto.process_interaction(
|
||||
# Extracts: name="John", preference="Python for backend"
|
||||
```
|
||||
|
||||
### 8. Agentic Workflows
|
||||
### 9. Agentic Workflows
|
||||
|
||||
Create powerful multi-agent workflows with the `Workflow` class:
|
||||
|
||||
@@ -867,7 +1079,7 @@ callbacks:
|
||||
on_workflow_complete: log_complete
|
||||
```
|
||||
|
||||
### 9. Hooks
|
||||
### 10. Hooks
|
||||
|
||||
Intercept and modify agent behavior at various lifecycle points:
|
||||
|
||||
@@ -906,7 +1118,7 @@ praisonai hooks validate hooks.json # Validate configuration
|
||||
```
|
||||
|
||||
|
||||
### 10. Shadow Git Checkpoints
|
||||
### 11. Shadow Git Checkpoints
|
||||
|
||||
File-level undo/restore using shadow git:
|
||||
|
||||
@@ -943,7 +1155,7 @@ praisonai checkpoint restore abc123 # Restore to checkpoint
|
||||
|
||||
---
|
||||
|
||||
### 11. Background Tasks
|
||||
### 12. Background Tasks
|
||||
|
||||
Run agent tasks asynchronously without blocking:
|
||||
|
||||
@@ -981,7 +1193,7 @@ praisonai background clear # Clear completed tasks
|
||||
|
||||
---
|
||||
|
||||
### 12. Policy Engine
|
||||
### 13. Policy Engine
|
||||
|
||||
Control what agents can and cannot do with policy-based execution:
|
||||
|
||||
@@ -1022,7 +1234,7 @@ praisonai policy init # Create template
|
||||
|
||||
---
|
||||
|
||||
### 13. Thinking Budgets
|
||||
### 14. Thinking Budgets
|
||||
|
||||
Configure token budgets for extended thinking:
|
||||
|
||||
@@ -1055,7 +1267,7 @@ praisonai thinking stats # Show usage statistics
|
||||
|
||||
---
|
||||
|
||||
### 14. Output Styles
|
||||
### 15. Output Styles
|
||||
|
||||
Configure how agents format their responses:
|
||||
|
||||
@@ -1085,7 +1297,7 @@ praisonai output set concise # Set output style
|
||||
|
||||
---
|
||||
|
||||
### 15. Context Compaction
|
||||
### 16. Context Compaction
|
||||
|
||||
Automatically manage context window size:
|
||||
|
||||
@@ -1120,7 +1332,7 @@ praisonai compaction stats # Show statistics
|
||||
|
||||
---
|
||||
|
||||
### 16. Field Names Reference (A-I-G-S)
|
||||
### 17. 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:
|
||||
|
||||
@@ -1154,7 +1366,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.
|
||||
|
||||
### 12. Extended agents.yaml with Workflow Patterns
|
||||
### 18. 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)
|
||||
@@ -1230,7 +1442,7 @@ Run with the same simple command:
|
||||
praisonai agents.yaml
|
||||
```
|
||||
|
||||
### 12. MCP (Model Context Protocol)
|
||||
### 19. MCP (Model Context Protocol)
|
||||
|
||||
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
|
||||
|
||||
@@ -1306,7 +1518,7 @@ server.run() # stdio for Claude Desktop
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
### 13. A2A (Agent2Agent Protocol)
|
||||
### 20. 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.
|
||||
|
||||
@@ -1354,6 +1566,22 @@ app.include_router(a2a.get_router())
|
||||
|
||||
PraisonAI provides a powerful CLI for no-code automation and quick prototyping.
|
||||
|
||||
### CLI Quick Reference
|
||||
|
||||
| Category | Commands |
|
||||
|----------|----------|
|
||||
| **Execution** | `praisonai`, `--auto`, `--interactive`, `--chat` |
|
||||
| **Research** | `research`, `--query-rewrite`, `--deep-research` |
|
||||
| **Planning** | `--planning`, `--planning-tools`, `--planning-reasoning` |
|
||||
| **Workflows** | `workflow run`, `workflow list`, `workflow auto` |
|
||||
| **Memory** | `memory show`, `memory add`, `memory search`, `memory clear` |
|
||||
| **Knowledge** | `knowledge add`, `knowledge query`, `knowledge list` |
|
||||
| **Sessions** | `session list`, `session resume`, `session delete` |
|
||||
| **Tools** | `tools list`, `tools info`, `tools search` |
|
||||
| **MCP** | `mcp list`, `mcp create`, `mcp enable` |
|
||||
| **Development** | `commit`, `docs`, `checkpoint`, `hooks` |
|
||||
| **Scheduling** | `schedule start`, `schedule list`, `schedule stop` |
|
||||
|
||||
### Auto Mode
|
||||
```bash
|
||||
pip install praisonai
|
||||
@@ -2637,6 +2865,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
|
||||
@@ -2751,6 +3057,44 @@ agent.chat("Hello!") # Auto-persists messages, runs, traces
|
||||
| `praisonai persistence migrate` | Apply schema migrations |
|
||||
| `praisonai persistence status` | Show schema status |
|
||||
|
||||
### Knowledge CLI Commands {#knowledge-cli}
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai knowledge add <source>` | Add file, directory, URL, or glob pattern |
|
||||
| `praisonai knowledge query "<question>"` | Query knowledge base with RAG |
|
||||
| `praisonai knowledge list` | List indexed documents |
|
||||
| `praisonai knowledge clear` | Clear knowledge base |
|
||||
| `praisonai knowledge stats` | Show knowledge base statistics |
|
||||
|
||||
**Knowledge Query Flags:**
|
||||
|
||||
| Flag | Values | Default |
|
||||
|------|--------|---------|
|
||||
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | `chroma` |
|
||||
| `--retrieval-strategy` | `basic`, `fusion`, `recursive`, `auto_merge` | `basic` |
|
||||
| `--reranker` | `none`, `simple`, `llm`, `cross_encoder`, `cohere` | `none` |
|
||||
| `--index-type` | `vector`, `keyword`, `hybrid` | `vector` |
|
||||
| `--query-mode` | `default`, `sub_question`, `summarize` | `default` |
|
||||
| `--workspace` | Path to workspace directory | Current dir |
|
||||
| `--session` | Session ID for persistence | - |
|
||||
|
||||
**Examples:**
|
||||
|
||||
```bash
|
||||
# Add documents
|
||||
praisonai knowledge add document.pdf
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge add "*.md"
|
||||
|
||||
# Query with options
|
||||
praisonai knowledge query "How to authenticate?" \
|
||||
--vector-store chroma \
|
||||
--retrieval-strategy fusion \
|
||||
--reranker simple \
|
||||
--query-mode sub_question
|
||||
```
|
||||
|
||||
### Databases Table
|
||||
|
||||
| Database | Store Type | Install | Example | Docs |
|
||||
@@ -2763,11 +3107,11 @@ agent.chat("Hello!") # Auto-persists messages, runs, traces
|
||||
| SurrealDB | Conversation | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Qdrant | Knowledge | `pip install "praisonai[tools]"` | [knowledge_qdrant.py](examples/persistence/knowledge_qdrant.py) | [docs](https://docs.praison.ai/docs/databases/qdrant) |
|
||||
| ChromaDB | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Pinecone | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| Weaviate | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| LanceDB | Knowledge | `pip install "praisonai[tools]"` | - | [docs](https://docs.praison.ai/docs/databases/overview) |
|
||||
| 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 "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) |
|
||||
@@ -2897,6 +3241,119 @@ cd src/praisonai && uv publish
|
||||
|
||||
---
|
||||
|
||||
## ❓ FAQ & Troubleshooting
|
||||
|
||||
<details>
|
||||
<summary><strong>ModuleNotFoundError: No module named 'praisonaiagents'</strong></summary>
|
||||
|
||||
Install the package:
|
||||
```bash
|
||||
pip install praisonaiagents
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>API key not found / Authentication error</strong></summary>
|
||||
|
||||
Ensure your API key is set:
|
||||
```bash
|
||||
export OPENAI_API_KEY=your_key_here
|
||||
```
|
||||
|
||||
For other providers, see [Environment Variables](#environment-variables).
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>How do I use a local model (Ollama)?</strong></summary>
|
||||
|
||||
```bash
|
||||
# Start Ollama server first
|
||||
ollama serve
|
||||
|
||||
# Set environment variable
|
||||
export OPENAI_BASE_URL=http://localhost:11434/v1
|
||||
```
|
||||
|
||||
See [Ollama Integration](#ollama-integration) for more details.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>How do I persist conversations to a database?</strong></summary>
|
||||
|
||||
Use the `db` parameter:
|
||||
```python
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
db=db(database_url="postgresql://localhost/mydb"),
|
||||
session_id="my-session"
|
||||
)
|
||||
```
|
||||
|
||||
See [Persistence (Databases)](#-persistence-databases) for supported databases.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>How do I enable agent memory?</strong></summary>
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
memory=True, # Enables file-based memory (no extra deps!)
|
||||
user_id="user123"
|
||||
)
|
||||
```
|
||||
|
||||
See [Agent Memory](#6-agent-memory-zero-dependencies) for more options.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>How do I run multiple agents together?</strong></summary>
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, PraisonAIAgents
|
||||
|
||||
agent1 = Agent(instructions="Research topics")
|
||||
agent2 = Agent(instructions="Summarize findings")
|
||||
agents = PraisonAIAgents(agents=[agent1, agent2])
|
||||
agents.start()
|
||||
```
|
||||
|
||||
See [Multi Agents](#2-multi-agents) for more examples.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>How do I use MCP tools?</strong></summary>
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
agent = Agent(
|
||||
tools=MCP("npx @modelcontextprotocol/server-memory")
|
||||
)
|
||||
```
|
||||
|
||||
See [MCP Protocol](#19-mcp-model-context-protocol) for all transport options.
|
||||
|
||||
</details>
|
||||
|
||||
### Getting Help
|
||||
|
||||
- 📚 [Full Documentation](https://docs.praison.ai)
|
||||
- 🐛 [Report Issues](https://github.com/MervinPraison/PraisonAI/issues)
|
||||
- 💬 [Discussions](https://github.com/MervinPraison/PraisonAI/discussions)
|
||||
|
||||
---
|
||||
|
||||
<div align="center">
|
||||
<p><strong>Made with ❤️ by the PraisonAI Team</strong></p>
|
||||
<p>
|
||||
|
||||
@@ -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.4.2" \
|
||||
"praisonai>=2.5.5" \
|
||||
"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.4.2" \
|
||||
"praisonai>=2.5.5" \
|
||||
"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.4.2" \
|
||||
"praisonai>=2.5.5" \
|
||||
"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,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())
|
||||
@@ -1,7 +1,7 @@
|
||||
"""
|
||||
Qdrant Knowledge Store Example for PraisonAI.
|
||||
Qdrant Knowledge Store - Agent-First Example
|
||||
|
||||
Demonstrates vector insert + search through the knowledge store.
|
||||
Demonstrates using Qdrant for knowledge-based RAG with an Agent.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
@@ -13,67 +13,59 @@ Run:
|
||||
python knowledge_qdrant.py
|
||||
|
||||
Expected Output:
|
||||
=== Qdrant Knowledge Store Demo ===
|
||||
Created collection: demo_knowledge
|
||||
Inserted 3 documents
|
||||
Search results: 2 documents found
|
||||
Agent responds using knowledge from documents
|
||||
"""
|
||||
|
||||
from praisonai.persistence.factory import create_knowledge_store
|
||||
from praisonai.persistence.knowledge.base import KnowledgeDocument
|
||||
import random
|
||||
from praisonaiagents import Agent
|
||||
import tempfile
|
||||
import os
|
||||
import shutil
|
||||
|
||||
# Create Qdrant knowledge store
|
||||
store = create_knowledge_store("qdrant", url="http://localhost:6333")
|
||||
# Create sample knowledge document
|
||||
sample_doc = """
|
||||
# AI and Programming Guide
|
||||
|
||||
collection_name = "demo_knowledge"
|
||||
dimension = 384 # Embedding dimension
|
||||
## Python Programming
|
||||
Python is a versatile programming language used for web development, data science, and AI.
|
||||
It has a simple syntax and extensive libraries.
|
||||
|
||||
# Create collection
|
||||
print("=== Qdrant Knowledge Store Demo ===")
|
||||
if store.collection_exists(collection_name):
|
||||
store.delete_collection(collection_name)
|
||||
## Machine Learning
|
||||
Machine learning enables computers to learn patterns from data without explicit programming.
|
||||
Common frameworks include TensorFlow, PyTorch, and scikit-learn.
|
||||
|
||||
store.create_collection(collection_name, dimension=dimension)
|
||||
print(f"Created collection: {collection_name}")
|
||||
## Databases
|
||||
PostgreSQL is a powerful open-source relational database system.
|
||||
It supports both SQL and JSON data types.
|
||||
"""
|
||||
|
||||
# Create sample documents with embeddings
|
||||
random.seed(42)
|
||||
documents = [
|
||||
KnowledgeDocument(
|
||||
id="doc-1",
|
||||
content="Python is a versatile programming language used for web development, data science, and AI.",
|
||||
embedding=[random.random() for _ in range(dimension)],
|
||||
metadata={"category": "programming", "language": "python"}
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc-2",
|
||||
content="Machine learning enables computers to learn patterns from data without explicit programming.",
|
||||
embedding=[random.random() for _ in range(dimension)],
|
||||
metadata={"category": "ai", "topic": "ml"}
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc-3",
|
||||
content="PostgreSQL is a powerful open-source relational database system.",
|
||||
embedding=[random.random() for _ in range(dimension)],
|
||||
metadata={"category": "database", "type": "sql"}
|
||||
),
|
||||
]
|
||||
# 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)
|
||||
|
||||
# Insert documents
|
||||
ids = store.insert(collection_name, documents)
|
||||
print(f"Inserted {len(ids)} documents")
|
||||
print("=== Qdrant Knowledge Store Demo (Agent-First) ===")
|
||||
|
||||
# Search for similar documents
|
||||
query_embedding = [random.random() for _ in range(dimension)]
|
||||
results = store.search(collection_name, query_embedding, limit=2)
|
||||
print(f"Search results: {len(results)} documents found")
|
||||
# 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"
|
||||
}
|
||||
)
|
||||
|
||||
for doc in results:
|
||||
print(f" - {doc.id}: {doc.content[:50]}...")
|
||||
# Chat - agent uses knowledge for RAG
|
||||
response = agent.chat("What programming language is good for AI?")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# Cleanup
|
||||
store.delete_collection(collection_name)
|
||||
store.close()
|
||||
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")
|
||||
@@ -1,83 +1,46 @@
|
||||
"""
|
||||
Redis State Store
|
||||
Redis State Store - Agent-First Example
|
||||
|
||||
Use Redis for state management, tracing, and caching.
|
||||
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:
|
||||
- State operations (set/get/delete)
|
||||
- Tracing operations
|
||||
Agent responds with state persisted to Redis
|
||||
"""
|
||||
|
||||
from praisonaiagents import db
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create Redis-only database (for state/tracing)
|
||||
redis_db = db.RedisDB(host="localhost", port=6379)
|
||||
print("=== Redis State Store (Agent-First) ===")
|
||||
|
||||
print("=== Redis State Store ===")
|
||||
|
||||
# Test state operations
|
||||
print("\n1. State Operations:")
|
||||
redis_db._init_stores() # Initialize stores
|
||||
|
||||
# Set a value
|
||||
redis_db._state_store.set("user:123:preferences", {
|
||||
"theme": "dark",
|
||||
"language": "en"
|
||||
})
|
||||
print(" Set user preferences")
|
||||
|
||||
# Get the value
|
||||
prefs = redis_db._state_store.get("user:123:preferences")
|
||||
print(f" Got preferences: {prefs}")
|
||||
|
||||
# Set with TTL (expires in 60 seconds)
|
||||
redis_db._state_store.set("temp:token", "abc123", ttl=60)
|
||||
print(" Set temporary token (60s TTL)")
|
||||
|
||||
# Delete
|
||||
redis_db._state_store.delete("user:123:preferences")
|
||||
print(" Deleted preferences")
|
||||
|
||||
# Test tracing
|
||||
print("\n2. Tracing Operations:")
|
||||
|
||||
# Start a trace
|
||||
trace_id = "trace-example-001"
|
||||
redis_db.on_trace_start(
|
||||
trace_id=trace_id,
|
||||
session_id="example-session",
|
||||
agent_name="ExampleAgent",
|
||||
metadata={"purpose": "demo"}
|
||||
# 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
|
||||
)
|
||||
print(f" Started trace: {trace_id}")
|
||||
|
||||
# Add spans
|
||||
span_id = "span-llm-001"
|
||||
redis_db.on_span_start(
|
||||
span_id=span_id,
|
||||
trace_id=trace_id,
|
||||
name="llm_call",
|
||||
attributes={"model": "gpt-4", "temperature": 0.7}
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant.",
|
||||
db=my_db,
|
||||
session_id="redis-state-session"
|
||||
)
|
||||
print(f" Started span: {span_id}")
|
||||
|
||||
redis_db.on_span_end(
|
||||
span_id=span_id,
|
||||
status="ok",
|
||||
attributes={"tokens": 150, "latency_ms": 450}
|
||||
)
|
||||
print(f" Ended span: {span_id}")
|
||||
# Chat - state is automatically managed via Redis
|
||||
response = agent.chat("Hello! Remember my name is Alice.")
|
||||
print(f"Response: {response}")
|
||||
|
||||
# End trace
|
||||
redis_db.on_trace_end(
|
||||
trace_id=trace_id,
|
||||
status="ok",
|
||||
metadata={"total_spans": 1}
|
||||
)
|
||||
print(f" Ended trace: {trace_id}")
|
||||
# Second message to test state
|
||||
response2 = agent.chat("What is my name?")
|
||||
print(f"Response: {response2}")
|
||||
|
||||
redis_db.close()
|
||||
my_db.close()
|
||||
print("\n✅ Done")
|
||||
|
||||
# --- Advanced: Direct Store Usage ---
|
||||
# from praisonaiagents import db
|
||||
# redis_db = db.RedisDB(host="localhost", port=6379)
|
||||
@@ -26,23 +26,18 @@ Expected Output:
|
||||
Messages persisted: 2
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonai.db import PostgresDB
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create database adapter
|
||||
db = PostgresDB(
|
||||
host="localhost",
|
||||
port=5432,
|
||||
database="praisonai",
|
||||
user="postgres",
|
||||
password="praison123"
|
||||
# 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,
|
||||
db=db_instance,
|
||||
session_id="demo-session-001", # Same session_id = resume conversation
|
||||
verbose=False
|
||||
)
|
||||
@@ -57,7 +52,7 @@ print(f"Session ID: {agent.session_id}")
|
||||
print(f"Messages in history: {len(agent.chat_history)}")
|
||||
|
||||
# Close DB connection
|
||||
db.close()
|
||||
db_instance.close()
|
||||
|
||||
print("\n=== Session persisted! ===")
|
||||
print("Run again with same session_id to resume.")
|
||||
@@ -1,10 +1,7 @@
|
||||
"""
|
||||
Redis State Store Example for PraisonAI.
|
||||
Redis State Store - Agent-First Example
|
||||
|
||||
Demonstrates set/get/delete + TTL operations.
|
||||
|
||||
Requirements:
|
||||
pip install "praisonai[tools]"
|
||||
Use Redis for state management with an Agent.
|
||||
|
||||
Docker Setup:
|
||||
docker run -d --name praison-redis -p 6379:6379 redis:7
|
||||
@@ -13,55 +10,37 @@ Run:
|
||||
python state_redis.py
|
||||
|
||||
Expected Output:
|
||||
=== Redis State Store Demo ===
|
||||
Set key: session:user123:preferences
|
||||
Get value: {'theme': 'dark', 'language': 'en'}
|
||||
TTL remaining: 3600 seconds
|
||||
Hash operations: {'name': 'Alice', 'role': 'admin'}
|
||||
Agent responds with state persisted to Redis
|
||||
"""
|
||||
|
||||
from praisonai.persistence.factory import create_state_store
|
||||
from praisonaiagents import Agent, db
|
||||
|
||||
# Create Redis state store
|
||||
store = create_state_store("redis", url="redis://localhost:6379")
|
||||
print("=== Redis State Store (Agent-First) ===")
|
||||
|
||||
print("=== Redis State Store Demo ===")
|
||||
# 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
|
||||
)
|
||||
|
||||
# Basic set/get
|
||||
key = "session:user123:preferences"
|
||||
value = {"theme": "dark", "language": "en", "notifications": True}
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
instructions="You are a helpful assistant.",
|
||||
db=my_db,
|
||||
session_id="redis-state-example"
|
||||
)
|
||||
|
||||
store.set(key, value)
|
||||
print(f"Set key: {key}")
|
||||
# Chat - state is automatically managed via Redis
|
||||
response = agent.chat("Hello! Remember my favorite color is blue.")
|
||||
print(f"Response: {response}")
|
||||
|
||||
retrieved = store.get(key)
|
||||
print(f"Get value: {retrieved}")
|
||||
|
||||
# Set with TTL (expires in 1 hour)
|
||||
ttl_key = "session:user123:token"
|
||||
store.set(ttl_key, "abc123xyz", ttl=3600)
|
||||
remaining = store.ttl(ttl_key)
|
||||
print(f"TTL remaining: {remaining} seconds")
|
||||
|
||||
# Hash operations (for structured data)
|
||||
hash_key = "user:123:profile"
|
||||
store.hset(hash_key, "name", "Alice")
|
||||
store.hset(hash_key, "role", "admin")
|
||||
store.hset(hash_key, "last_login", "2024-12-24")
|
||||
|
||||
profile = store.hgetall(hash_key)
|
||||
print(f"Hash operations: {profile}")
|
||||
|
||||
# Check existence
|
||||
exists = store.exists(key)
|
||||
print(f"Key exists: {exists}")
|
||||
|
||||
# Delete keys
|
||||
store.delete(key)
|
||||
store.delete(ttl_key)
|
||||
store.delete(hash_key)
|
||||
print("Cleaned up test keys")
|
||||
|
||||
store.close()
|
||||
# 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")
|
||||
@@ -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,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,8 +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,
|
||||
@@ -66,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":
|
||||
@@ -92,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
|
||||
@@ -259,6 +271,7 @@ __all__ = [
|
||||
'Session',
|
||||
'Memory',
|
||||
'db',
|
||||
'obs',
|
||||
'display_interaction',
|
||||
'display_self_reflection',
|
||||
'display_instruction',
|
||||
|
||||
@@ -254,7 +254,10 @@ class Agent:
|
||||
skills: Optional[List[str]] = None,
|
||||
skills_dirs: Optional[List[str]] = None,
|
||||
db: Optional[Any] = None,
|
||||
session_id: Optional[str] = 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.
|
||||
|
||||
@@ -395,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
|
||||
@@ -1622,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
|
||||
@@ -2058,7 +2096,11 @@ Your Goal: {self.goal}"""
|
||||
"""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):
|
||||
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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -89,11 +89,9 @@ class AccuracyEvaluator(BaseEvaluator):
|
||||
EvaluationScore with score and reasoning
|
||||
"""
|
||||
try:
|
||||
from openai import OpenAI
|
||||
import litellm
|
||||
except ImportError:
|
||||
raise ImportError("OpenAI package required for accuracy evaluation. Install with: pip install openai")
|
||||
|
||||
client = OpenAI()
|
||||
raise ImportError("litellm package required for accuracy evaluation. Install with: pip install litellm")
|
||||
|
||||
prompt = f"""You are an expert evaluator. Compare the actual output against the expected output and provide a score from 1-10.
|
||||
|
||||
@@ -115,7 +113,8 @@ Respond in this exact format:
|
||||
SCORE: [number 1-10]
|
||||
REASONING: [brief explanation]"""
|
||||
|
||||
response = client.chat.completions.create(
|
||||
# Use litellm for unified multi-provider support
|
||||
response = litellm.completion(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=0.1,
|
||||
|
||||
@@ -99,11 +99,9 @@ class CriteriaEvaluator(BaseEvaluator):
|
||||
CriteriaScore with score, passed status, and reasoning
|
||||
"""
|
||||
try:
|
||||
from openai import OpenAI
|
||||
import litellm
|
||||
except ImportError:
|
||||
raise ImportError("OpenAI package required for criteria evaluation. Install with: pip install openai")
|
||||
|
||||
client = OpenAI()
|
||||
raise ImportError("litellm package required for criteria evaluation. Install with: pip install litellm")
|
||||
|
||||
if self.scoring_type == "binary":
|
||||
prompt = f"""You are an expert evaluator. Evaluate the following output against the given criteria.
|
||||
@@ -138,7 +136,8 @@ Respond in this exact format:
|
||||
SCORE: [number 1-10]
|
||||
REASONING: [brief explanation]"""
|
||||
|
||||
response = client.chat.completions.create(
|
||||
# Use litellm for unified multi-provider support
|
||||
response = litellm.completion(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=0.1,
|
||||
|
||||
@@ -65,6 +65,23 @@ __all__ = [
|
||||
"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",
|
||||
]
|
||||
|
||||
|
||||
@@ -94,4 +111,27 @@ def __getattr__(name: str):
|
||||
)
|
||||
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,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)
|
||||
@@ -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()
|
||||
@@ -91,19 +91,20 @@ class MongoDBMemory:
|
||||
raise Exception(f"Failed to initialize MongoDB: {e}")
|
||||
|
||||
def _init_embedding_model(self):
|
||||
"""Initialize embedding model from config."""
|
||||
"""Initialize embedding model from config using litellm for unified provider support."""
|
||||
try:
|
||||
# Set up embedding model based on config
|
||||
embedder_config = self.config.get("embedder", {})
|
||||
if embedder_config.get("provider") == "openai":
|
||||
import openai
|
||||
self.embedding_model = openai.OpenAI()
|
||||
self.embedding_model_name = embedder_config.get("config", {}).get("model", "text-embedding-3-small")
|
||||
else:
|
||||
# Default to OpenAI
|
||||
import openai
|
||||
self.embedding_model = openai.OpenAI()
|
||||
self.embedding_model_name = "text-embedding-3-small"
|
||||
provider = embedder_config.get("provider", "openai")
|
||||
model_name = embedder_config.get("config", {}).get("model", "text-embedding-3-small")
|
||||
|
||||
# Store model name for later use
|
||||
self.embedding_model_name = model_name
|
||||
|
||||
# Use litellm for embeddings - it handles all providers uniformly
|
||||
# We'll use litellm.embedding() in _get_embeddings() instead of storing a client
|
||||
self.embedding_provider = provider
|
||||
self.embedding_model = None # Will use litellm.embedding() directly
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to initialize embedding model: {e}")
|
||||
|
||||
@@ -172,13 +173,16 @@ class MongoDBMemory:
|
||||
logging.warning(f"Could not create vector search index: {e}")
|
||||
|
||||
def _get_embedding(self, text):
|
||||
"""Get embedding for text."""
|
||||
"""Get embedding for text using litellm for unified provider support."""
|
||||
try:
|
||||
response = self.embedding_model.embeddings.create(
|
||||
input=text,
|
||||
model=self.embedding_model_name
|
||||
import litellm
|
||||
|
||||
# Use litellm.embedding() for unified multi-provider support
|
||||
response = litellm.embedding(
|
||||
model=self.embedding_model_name,
|
||||
input=text
|
||||
)
|
||||
return response.data[0].embedding
|
||||
return response.data[0]["embedding"]
|
||||
except Exception as e:
|
||||
logging.error(f"Error getting embedding: {e}")
|
||||
return None
|
||||
@@ -500,6 +504,33 @@ class Knowledge:
|
||||
base_config["graph_store"] = self._config["graph_store"]
|
||||
return base_config
|
||||
|
||||
def _prepare_mem0_config(self, config):
|
||||
"""Prepare config for mem0 by removing PraisonAI-specific fields.
|
||||
|
||||
mem0's RerankerConfig only accepts 'provider' and 'config' fields.
|
||||
PraisonAI adds 'enabled' and 'default_rerank' for internal use.
|
||||
"""
|
||||
mem0_config = config.copy()
|
||||
|
||||
# Strip PraisonAI-specific reranker fields that mem0 doesn't accept
|
||||
if "reranker" in mem0_config:
|
||||
reranker = mem0_config["reranker"]
|
||||
if isinstance(reranker, dict):
|
||||
# Keep only mem0-compatible fields: provider, config
|
||||
mem0_reranker = {}
|
||||
if "provider" in reranker:
|
||||
mem0_reranker["provider"] = reranker["provider"]
|
||||
if "config" in reranker:
|
||||
mem0_reranker["config"] = reranker["config"]
|
||||
|
||||
# If no valid mem0 fields, remove reranker entirely
|
||||
if mem0_reranker:
|
||||
mem0_config["reranker"] = mem0_reranker
|
||||
else:
|
||||
del mem0_config["reranker"]
|
||||
|
||||
return mem0_config
|
||||
|
||||
@cached_property
|
||||
def memory(self):
|
||||
# Check if MongoDB provider is specified
|
||||
@@ -511,19 +542,22 @@ class Knowledge:
|
||||
# Fall back to default memory
|
||||
pass
|
||||
|
||||
# Prepare config for mem0 (strip PraisonAI-specific fields)
|
||||
mem0_config = self._prepare_mem0_config(self.config)
|
||||
|
||||
# Default Mem0 memory
|
||||
try:
|
||||
return CustomMemory.from_config(self.config)
|
||||
return CustomMemory.from_config(mem0_config)
|
||||
except (NotImplementedError, ValueError) as e:
|
||||
if "list_collections" in str(e) or "Extra fields not allowed" in str(e):
|
||||
# Keep only allowed fields
|
||||
vector_store_config = {
|
||||
"collection_name": self.config["vector_store"]["config"]["collection_name"],
|
||||
"path": self.config["vector_store"]["config"]["path"]
|
||||
"collection_name": mem0_config["vector_store"]["config"]["collection_name"],
|
||||
"path": mem0_config["vector_store"]["config"]["path"]
|
||||
}
|
||||
self.config["vector_store"]["config"] = vector_store_config
|
||||
mem0_config["vector_store"]["config"] = vector_store_config
|
||||
from mem0 import Memory
|
||||
return Memory.from_config(self.config)
|
||||
return Memory.from_config(mem0_config)
|
||||
raise
|
||||
|
||||
@cached_property
|
||||
|
||||
@@ -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,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",
|
||||
]
|
||||
@@ -4,16 +4,20 @@ Session Management for PraisonAI Agents
|
||||
A simple wrapper around existing stateful capabilities to provide a unified
|
||||
session API for developers building stateful agent applications.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import uuid
|
||||
import requests
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory.memory import Memory
|
||||
from .knowledge.knowledge import Knowledge
|
||||
|
||||
from .agent import Agent
|
||||
from .memory import Memory
|
||||
from .knowledge import Knowledge
|
||||
|
||||
|
||||
class Session:
|
||||
@@ -112,6 +116,7 @@ class Session:
|
||||
if self.is_remote:
|
||||
raise ValueError("Memory operations are not available for remote agent sessions")
|
||||
if self._memory is None:
|
||||
from .memory.memory import Memory
|
||||
self._memory = Memory(config=self.memory_config)
|
||||
return self._memory
|
||||
|
||||
@@ -121,6 +126,7 @@ class Session:
|
||||
if self.is_remote:
|
||||
raise ValueError("Knowledge operations are not available for remote agent sessions")
|
||||
if self._knowledge is None:
|
||||
from .knowledge.knowledge import Knowledge
|
||||
self._knowledge = Knowledge(config=self.knowledge_config)
|
||||
return self._knowledge
|
||||
|
||||
|
||||
@@ -2,6 +2,9 @@
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
# Export Injected type for tool state injection
|
||||
from .injected import Injected
|
||||
|
||||
# Map of function names to their module and class (if any)
|
||||
TOOL_MAPPINGS = {
|
||||
# Direct functions
|
||||
@@ -286,4 +289,4 @@ def __getattr__(name: str) -> Any:
|
||||
method = getattr(_instances[class_name], name)
|
||||
return method
|
||||
|
||||
__all__ = list(TOOL_MAPPINGS.keys())
|
||||
__all__ = list(TOOL_MAPPINGS.keys()) + ['Injected']
|
||||
@@ -14,6 +14,14 @@ Usage:
|
||||
@tool(name="web_search", description="Search the internet")
|
||||
def search(query: str, max_results: int = 5) -> list:
|
||||
return [...]
|
||||
|
||||
# With injected state:
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
'''Tool with injected state.'''
|
||||
return f"session={state.get('session_id')}"
|
||||
"""
|
||||
|
||||
import inspect
|
||||
@@ -22,12 +30,19 @@ import logging
|
||||
from typing import Any, Callable, Dict, Optional, Union, get_type_hints
|
||||
|
||||
from .base import BaseTool
|
||||
from .injected import (
|
||||
is_injected_type,
|
||||
get_injected_params,
|
||||
inject_state_into_kwargs,
|
||||
filter_injected_from_schema
|
||||
)
|
||||
|
||||
|
||||
class FunctionTool(BaseTool):
|
||||
"""A BaseTool wrapper for plain functions.
|
||||
|
||||
Created automatically by the @tool decorator.
|
||||
Supports Injected[T] parameters for state injection.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -42,6 +57,9 @@ class FunctionTool(BaseTool):
|
||||
self.description = description or func.__doc__ or f"Tool: {self.name}"
|
||||
self.version = version
|
||||
|
||||
# Detect injected parameters
|
||||
self._injected_params = get_injected_params(func)
|
||||
|
||||
# Generate schema from the original function, not run()
|
||||
self.parameters = self._generate_schema_from_func(func)
|
||||
|
||||
@@ -55,8 +73,16 @@ class FunctionTool(BaseTool):
|
||||
if not self.description:
|
||||
self.description = self.__class__.__doc__ or f"Tool: {self.name}"
|
||||
|
||||
@property
|
||||
def injected_params(self) -> Dict[str, Any]:
|
||||
"""Get the injected parameters for this tool."""
|
||||
return self._injected_params
|
||||
|
||||
def _generate_schema_from_func(self, func: Callable) -> Dict[str, Any]:
|
||||
"""Generate JSON Schema from the wrapped function's signature."""
|
||||
"""Generate JSON Schema from the wrapped function's signature.
|
||||
|
||||
Injected parameters are excluded from the schema.
|
||||
"""
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
@@ -71,8 +97,17 @@ class FunctionTool(BaseTool):
|
||||
if param_name in ('self', 'cls'):
|
||||
continue
|
||||
|
||||
# Skip injected parameters - they don't go in schema
|
||||
if param_name in self._injected_params:
|
||||
continue
|
||||
|
||||
# Get type hint
|
||||
param_type = hints.get(param_name, Any)
|
||||
|
||||
# Double-check it's not an Injected type
|
||||
if is_injected_type(param_type):
|
||||
continue
|
||||
|
||||
json_type = BaseTool._python_type_to_json(param_type)
|
||||
|
||||
schema["properties"][param_name] = {"type": json_type}
|
||||
@@ -86,11 +121,18 @@ class FunctionTool(BaseTool):
|
||||
return schema
|
||||
|
||||
def run(self, **kwargs) -> Any:
|
||||
"""Execute the wrapped function."""
|
||||
"""Execute the wrapped function with injected state."""
|
||||
# Inject state for any Injected parameters
|
||||
kwargs = inject_state_into_kwargs(kwargs, self._injected_params)
|
||||
return self._func(**kwargs)
|
||||
|
||||
def __call__(self, *args, **kwargs) -> Any:
|
||||
"""Allow calling with positional args like the original function."""
|
||||
"""Allow calling with positional args like the original function.
|
||||
|
||||
Injects state for Injected parameters.
|
||||
"""
|
||||
# Inject state for any Injected parameters
|
||||
kwargs = inject_state_into_kwargs(kwargs, self._injected_params)
|
||||
return self._func(*args, **kwargs)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
"""
|
||||
Injected State for PraisonAI Tools.
|
||||
|
||||
Provides a mechanism for tools to receive injected state (agent context)
|
||||
without exposing it in the tool's public schema.
|
||||
|
||||
Zero Performance Impact:
|
||||
- When no Injected params, no overhead
|
||||
- Context is thread-local for safety
|
||||
- Lazy evaluation of state
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
'''Tool with injected state.'''
|
||||
session_id = state.get('session_id')
|
||||
return f"Query: {query}, Session: {session_id}"
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Generic, List, Optional, TypeVar, get_type_hints, get_origin, get_args
|
||||
from contextvars import ContextVar
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
T = TypeVar('T')
|
||||
|
||||
|
||||
class Injected(Generic[T]):
|
||||
"""Type marker for injected parameters.
|
||||
|
||||
Parameters annotated with Injected[T] will:
|
||||
- Not appear in the tool's public schema
|
||||
- Be automatically injected at runtime with the current agent state
|
||||
|
||||
Supported types:
|
||||
- Injected[dict] - Receives state as a dictionary
|
||||
- Injected[AgentState] - Receives state as AgentState dataclass
|
||||
|
||||
Example:
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
return f"session={state.get('session_id')}"
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentState:
|
||||
"""State object injected into tools.
|
||||
|
||||
Contains context about the current agent execution.
|
||||
"""
|
||||
agent_id: str
|
||||
run_id: str
|
||||
session_id: str
|
||||
last_user_message: Optional[str] = None
|
||||
last_agent_message: Optional[str] = None
|
||||
memory: Optional[Any] = None
|
||||
previous_tool_results: List[Any] = field(default_factory=list)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for Injected[dict] usage."""
|
||||
return {
|
||||
'agent_id': self.agent_id,
|
||||
'run_id': self.run_id,
|
||||
'session_id': self.session_id,
|
||||
'last_user_message': self.last_user_message,
|
||||
'last_agent_message': self.last_agent_message,
|
||||
'previous_tool_results': self.previous_tool_results,
|
||||
'metadata': self.metadata,
|
||||
}
|
||||
|
||||
|
||||
# Thread-local context for current state
|
||||
_current_state: ContextVar[Optional[AgentState]] = ContextVar('current_state', default=None)
|
||||
|
||||
|
||||
def get_current_state() -> Optional[AgentState]:
|
||||
"""Get the current injection state."""
|
||||
return _current_state.get()
|
||||
|
||||
|
||||
def set_current_state(state: Optional[AgentState]) -> None:
|
||||
"""Set the current injection state."""
|
||||
_current_state.set(state)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def with_injection_context(state: AgentState):
|
||||
"""Context manager for setting injection state.
|
||||
|
||||
Usage:
|
||||
with with_injection_context(state):
|
||||
result = my_tool(query="test")
|
||||
"""
|
||||
token = _current_state.set(state)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_current_state.reset(token)
|
||||
|
||||
|
||||
def is_injected_type(annotation: Any) -> bool:
|
||||
"""Check if a type annotation is Injected[T]."""
|
||||
origin = get_origin(annotation)
|
||||
if origin is Injected:
|
||||
return True
|
||||
# Handle string annotations
|
||||
if isinstance(annotation, str) and 'Injected' in annotation:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_injected_type(annotation: Any) -> Optional[type]:
|
||||
"""Get the inner type T from Injected[T]."""
|
||||
origin = get_origin(annotation)
|
||||
if origin is Injected:
|
||||
args = get_args(annotation)
|
||||
if args:
|
||||
return args[0]
|
||||
return None
|
||||
|
||||
|
||||
def resolve_injected_value(annotation: Any, state: Optional[AgentState]) -> Any:
|
||||
"""Resolve the value to inject based on annotation type.
|
||||
|
||||
Args:
|
||||
annotation: The Injected[T] annotation
|
||||
state: The current agent state
|
||||
|
||||
Returns:
|
||||
The resolved value (dict or AgentState)
|
||||
"""
|
||||
if state is None:
|
||||
# Return empty dict if no state available
|
||||
return {}
|
||||
|
||||
inner_type = get_injected_type(annotation)
|
||||
|
||||
if inner_type is dict or inner_type is Dict:
|
||||
return state.to_dict()
|
||||
elif inner_type is AgentState:
|
||||
return state
|
||||
else:
|
||||
# Default to dict
|
||||
return state.to_dict()
|
||||
|
||||
|
||||
def get_injected_params(func) -> Dict[str, Any]:
|
||||
"""Get all injected parameters from a function.
|
||||
|
||||
Returns:
|
||||
Dict mapping param name to its Injected[T] annotation
|
||||
"""
|
||||
try:
|
||||
hints = get_type_hints(func)
|
||||
except Exception:
|
||||
# If we can't get hints, return empty
|
||||
return {}
|
||||
|
||||
injected = {}
|
||||
for name, hint in hints.items():
|
||||
if name == 'return':
|
||||
continue
|
||||
if is_injected_type(hint):
|
||||
injected[name] = hint
|
||||
|
||||
return injected
|
||||
|
||||
|
||||
def filter_injected_from_schema(properties: Dict[str, Any], required: List[str],
|
||||
injected_params: Dict[str, Any]) -> tuple:
|
||||
"""Remove injected parameters from schema.
|
||||
|
||||
Args:
|
||||
properties: Schema properties dict
|
||||
required: List of required param names
|
||||
injected_params: Dict of injected param names
|
||||
|
||||
Returns:
|
||||
Tuple of (filtered_properties, filtered_required)
|
||||
"""
|
||||
filtered_props = {k: v for k, v in properties.items() if k not in injected_params}
|
||||
filtered_required = [r for r in required if r not in injected_params]
|
||||
return filtered_props, filtered_required
|
||||
|
||||
|
||||
def inject_state_into_kwargs(kwargs: Dict[str, Any], injected_params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Inject state values into kwargs for injected parameters.
|
||||
|
||||
Args:
|
||||
kwargs: The original kwargs
|
||||
injected_params: Dict of injected param names to annotations
|
||||
|
||||
Returns:
|
||||
Updated kwargs with injected values
|
||||
"""
|
||||
if not injected_params:
|
||||
return kwargs
|
||||
|
||||
state = get_current_state()
|
||||
result = kwargs.copy()
|
||||
|
||||
for param_name, annotation in injected_params.items():
|
||||
if param_name not in result:
|
||||
result[param_name] = resolve_injected_value(annotation, state)
|
||||
|
||||
return result
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.2.1"
|
||||
version = "0.3.4"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -672,16 +672,16 @@ class TestPlanningAgent:
|
||||
|
||||
planner = PlanningAgent()
|
||||
|
||||
assert planner.llm_model == "gpt-4o-mini"
|
||||
assert planner.llm_model in ["gpt-4o-mini", "gpt-5-nano"]
|
||||
assert planner.read_only is True
|
||||
|
||||
def test_planning_agent_custom_llm(self):
|
||||
"""Test PlanningAgent with custom LLM."""
|
||||
from praisonaiagents.planning import PlanningAgent
|
||||
|
||||
planner = PlanningAgent(llm="gpt-4o")
|
||||
planner = PlanningAgent(llm="gpt-5-nano")
|
||||
|
||||
assert planner.llm_model == "gpt-4o"
|
||||
assert planner.llm_model == "gpt-5-nano"
|
||||
|
||||
def test_planning_agent_allowed_tools(self):
|
||||
"""Test PlanningAgent has only read-only tools."""
|
||||
@@ -840,10 +840,10 @@ class TestPraisonAIAgentsPlanningIntegration:
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
planning=True,
|
||||
planning_llm="gpt-4o"
|
||||
planning_llm="gpt-5-nano"
|
||||
)
|
||||
|
||||
assert agents.planning_llm == "gpt-4o"
|
||||
assert agents.planning_llm == "gpt-5-nano"
|
||||
|
||||
def test_agents_auto_approve_plan(self):
|
||||
"""Test auto-approve plan option."""
|
||||
@@ -867,7 +867,7 @@ class TestPraisonAIAgentsPlanningIntegration:
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
from praisonaiagents.planning import Plan, PlanStep
|
||||
|
||||
agent = Agent(name="Test", role="Tester", llm="gpt-4o-mini")
|
||||
agent = Agent(name="Test", role="Tester", llm="gpt-5-nano")
|
||||
task = Task(description="Test task", agent=agent)
|
||||
|
||||
agents = PraisonAIAgents(
|
||||
@@ -880,7 +880,7 @@ class TestPraisonAIAgentsPlanningIntegration:
|
||||
# Test that planning is enabled and configured correctly
|
||||
assert agents.planning is True
|
||||
assert agents.auto_approve_plan is True
|
||||
assert agents.planning_llm == "gpt-4o-mini"
|
||||
assert agents.planning_llm == "gpt-5-nano"
|
||||
assert agents._current_plan is None # No plan created yet
|
||||
assert agents._todo_list is None
|
||||
|
||||
|
||||
@@ -1189,16 +1189,16 @@ class TestWorkflowConfiguration:
|
||||
"""Test that workflow planning setting is stored."""
|
||||
def step1(ctx): return StepResult(output="Done")
|
||||
|
||||
workflow = Workflow(steps=[step1], planning=True, planning_llm="gpt-4o")
|
||||
workflow = Workflow(steps=[step1], planning=True, planning_llm="gpt-5-nano")
|
||||
assert workflow.planning == True
|
||||
assert workflow.planning_llm == "gpt-4o"
|
||||
assert workflow.planning_llm == "gpt-5-nano"
|
||||
|
||||
def test_workflow_default_llm(self):
|
||||
"""Test default LLM setting."""
|
||||
def step1(ctx): return StepResult(output="Done")
|
||||
|
||||
workflow = Workflow(steps=[step1], default_llm="gpt-4o-mini")
|
||||
assert workflow.default_llm == "gpt-4o-mini"
|
||||
workflow = Workflow(steps=[step1], default_llm="gpt-5-nano")
|
||||
assert workflow.default_llm == "gpt-5-nano"
|
||||
|
||||
def test_workflow_default_agent_config(self):
|
||||
"""Test default agent config."""
|
||||
@@ -1393,10 +1393,10 @@ class TestPlanningAndReasoning:
|
||||
workflow = Workflow(
|
||||
steps=[lambda ctx: StepResult(output="test")],
|
||||
planning=True,
|
||||
planning_llm="gpt-4o"
|
||||
planning_llm="gpt-5-nano"
|
||||
)
|
||||
assert workflow.planning == True
|
||||
assert workflow.planning_llm == "gpt-4o"
|
||||
assert workflow.planning_llm == "gpt-5-nano"
|
||||
|
||||
def test_workflow_reasoning_parameter_exists(self):
|
||||
"""Test that Workflow accepts reasoning parameter."""
|
||||
|
||||
@@ -91,14 +91,11 @@ class TestAccuracyEvaluator:
|
||||
|
||||
def test_judge_output(self):
|
||||
"""Test LLM judging of output."""
|
||||
with patch('openai.OpenAI') as mock_openai_class:
|
||||
mock_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_client
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = "SCORE: 8\nREASONING: Good match"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
def test_func(x):
|
||||
return "4"
|
||||
@@ -115,14 +112,11 @@ class TestAccuracyEvaluator:
|
||||
|
||||
def test_run_returns_accuracy_result(self):
|
||||
"""Test that run returns AccuracyResult."""
|
||||
with patch('openai.OpenAI') as mock_openai_class:
|
||||
mock_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_client
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = "SCORE: 9\nREASONING: Excellent"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
def test_func(x):
|
||||
return "4"
|
||||
@@ -140,14 +134,11 @@ class TestAccuracyEvaluator:
|
||||
|
||||
def test_evaluate_output(self):
|
||||
"""Test evaluate_output method for pre-generated outputs."""
|
||||
with patch('openai.OpenAI') as mock_openai_class:
|
||||
mock_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_client
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = "SCORE: 10\nREASONING: Perfect"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
def test_func(x):
|
||||
return "unused"
|
||||
@@ -325,14 +316,11 @@ class TestCriteriaEvaluator:
|
||||
|
||||
def test_judge_output_numeric(self):
|
||||
"""Test numeric scoring."""
|
||||
with patch('openai.OpenAI') as mock_openai_class:
|
||||
mock_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_client
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = "SCORE: 8\nREASONING: Good"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
def test_func(x):
|
||||
return "test"
|
||||
@@ -350,14 +338,11 @@ class TestCriteriaEvaluator:
|
||||
|
||||
def test_judge_output_binary(self):
|
||||
"""Test binary scoring."""
|
||||
with patch('openai.OpenAI') as mock_openai_class:
|
||||
mock_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_client
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = "RESULT: PASS\nREASONING: Meets criteria"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
def test_func(x):
|
||||
return "test"
|
||||
@@ -374,14 +359,11 @@ class TestCriteriaEvaluator:
|
||||
|
||||
def test_on_fail_callback(self):
|
||||
"""Test on_fail callback is called."""
|
||||
with patch('openai.OpenAI') as mock_openai_class:
|
||||
mock_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_client
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = "SCORE: 3\nREASONING: Poor"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
on_fail_called = []
|
||||
def on_fail(score):
|
||||
@@ -403,14 +385,11 @@ class TestCriteriaEvaluator:
|
||||
|
||||
def test_evaluate_output_criteria(self):
|
||||
"""Test evaluate_output method."""
|
||||
with patch('openai.OpenAI') as mock_openai_class:
|
||||
mock_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_client
|
||||
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = "SCORE: 9\nREASONING: Great"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
mock_completion.return_value = mock_response
|
||||
|
||||
def test_func(x):
|
||||
return "unused"
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
"""
|
||||
TDD Tests for Configurable Model (runtime model switching).
|
||||
|
||||
These tests are written FIRST before implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
|
||||
class TestConfigurableModelBasic:
|
||||
"""Test basic configurable model functionality."""
|
||||
|
||||
def test_agent_accepts_llm_config_configurable(self):
|
||||
"""Agent should accept llm_config with configurable=True."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
assert agent is not None
|
||||
assert getattr(agent, '_llm_configurable', False) is True
|
||||
|
||||
def test_agent_chat_accepts_config_parameter(self):
|
||||
"""Agent.chat() should accept config parameter for model override."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
# Should not raise - config parameter should be accepted
|
||||
# We'll mock the actual LLM call
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked"):
|
||||
result = agent.chat(
|
||||
"Hello",
|
||||
config={"model": "claude-3-5-sonnet"}
|
||||
)
|
||||
|
||||
def test_config_model_override(self):
|
||||
"""config.model should override default model for that call."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
# Track which model was used
|
||||
used_model = [None]
|
||||
|
||||
original_chat = agent._chat_completion
|
||||
def mock_chat(*args, **kwargs):
|
||||
# The model should be accessible somehow
|
||||
used_model[0] = getattr(agent, '_current_model', agent.llm)
|
||||
return "response"
|
||||
|
||||
with patch.object(agent, '_chat_completion', side_effect=mock_chat):
|
||||
agent.chat("Hello", config={"model": "claude-3-5-sonnet"})
|
||||
|
||||
# Model should have been switched for this call
|
||||
# Implementation will set _current_model or similar
|
||||
|
||||
def test_config_does_not_mutate_default(self):
|
||||
"""Per-call config should not change agent's default model."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
original_llm = agent.llm
|
||||
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked"):
|
||||
agent.chat("Hello", config={"model": "claude-3-5-sonnet"})
|
||||
|
||||
# Default should be unchanged
|
||||
assert agent.llm == original_llm
|
||||
|
||||
def test_config_temperature_override(self):
|
||||
"""config.temperature should override for that call."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
# Should accept temperature in config
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked") as mock:
|
||||
agent.chat("Hello", config={"temperature": 0.2})
|
||||
# Temperature should be passed through
|
||||
|
||||
|
||||
class TestConfigurableModelCaching:
|
||||
"""Test model client caching for configurable models."""
|
||||
|
||||
def test_model_client_cached(self):
|
||||
"""Model clients should be cached to avoid recreation."""
|
||||
from praisonaiagents.llm import LLM
|
||||
|
||||
llm = LLM(model="gpt-4o", configurable=True)
|
||||
|
||||
# First call with model A
|
||||
# Second call with model A should reuse client
|
||||
# This is internal behavior - we test via call count or similar
|
||||
|
||||
def test_cache_size_limited(self):
|
||||
"""Model client cache should have a size limit."""
|
||||
from praisonaiagents.llm import LLM
|
||||
|
||||
llm = LLM(model="gpt-4o", configurable=True, cache_size=3)
|
||||
|
||||
# Should not grow unbounded
|
||||
|
||||
|
||||
class TestConfigurableModelThreadSafety:
|
||||
"""Test thread safety of configurable models."""
|
||||
|
||||
def test_concurrent_calls_different_models(self):
|
||||
"""Concurrent calls with different models should not interfere."""
|
||||
import threading
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
results = {}
|
||||
errors = []
|
||||
|
||||
def call_with_model(model_name, thread_id):
|
||||
try:
|
||||
with patch.object(agent, '_chat_completion', return_value=f"response_{model_name}"):
|
||||
result = agent.chat("Hello", config={"model": model_name})
|
||||
results[thread_id] = result
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=call_with_model, args=("model_a", 1)),
|
||||
threading.Thread(target=call_with_model, args=("model_b", 2)),
|
||||
]
|
||||
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert len(errors) == 0
|
||||
|
||||
|
||||
class TestConfigurableModelValidation:
|
||||
"""Test validation of configurable model settings."""
|
||||
|
||||
def test_config_without_configurable_flag_ignored(self):
|
||||
"""If llm_config.configurable is False, config param should be ignored."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o"
|
||||
# No llm_config or configurable=False
|
||||
)
|
||||
|
||||
# Should still work, just ignore the config
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked"):
|
||||
result = agent.chat("Hello", config={"model": "other"})
|
||||
# Should use default model, not "other"
|
||||
|
||||
def test_invalid_config_keys_ignored(self):
|
||||
"""Unknown keys in config should be ignored, not error."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
# Should not raise for unknown keys
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked"):
|
||||
result = agent.chat("Hello", config={"unknown_key": "value"})
|
||||
|
||||
|
||||
class TestConfigurableModelProvider:
|
||||
"""Test provider switching in configurable models."""
|
||||
|
||||
def test_config_provider_override(self):
|
||||
"""config.provider should allow switching providers."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
# Should accept provider in config
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked"):
|
||||
agent.chat("Hello", config={"provider": "anthropic", "model": "claude-3-5-sonnet"})
|
||||
|
||||
def test_model_string_with_provider_prefix(self):
|
||||
"""Model string like 'anthropic/claude-3' should work."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
llm="gpt-4o",
|
||||
llm_config={"configurable": True}
|
||||
)
|
||||
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked"):
|
||||
agent.chat("Hello", config={"model": "anthropic/claude-3-5-sonnet"})
|
||||
@@ -0,0 +1,367 @@
|
||||
"""
|
||||
TDD Tests for InjectedState for Tools.
|
||||
|
||||
These tests are written FIRST before implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
|
||||
class TestInjectedTypeExists:
|
||||
"""Test that Injected type marker exists."""
|
||||
|
||||
def test_injected_type_exists(self):
|
||||
"""Injected type should exist in praisonaiagents.tools."""
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
# Should be a generic type marker
|
||||
assert Injected is not None
|
||||
|
||||
def test_injected_generic_usage(self):
|
||||
"""Injected should work as a generic type annotation."""
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
# Should be usable as Injected[dict], Injected[SomeType], etc.
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
return query
|
||||
|
||||
# Should have annotations
|
||||
annotations = my_tool.__annotations__
|
||||
assert 'state' in annotations
|
||||
|
||||
def test_injected_with_agent_state(self):
|
||||
"""Injected[AgentState] should work."""
|
||||
from praisonaiagents.tools import Injected
|
||||
from praisonaiagents.tools.injected import AgentState
|
||||
|
||||
def my_tool(query: str, state: Injected[AgentState]) -> str:
|
||||
return query
|
||||
|
||||
annotations = my_tool.__annotations__
|
||||
assert 'state' in annotations
|
||||
|
||||
|
||||
class TestAgentStateType:
|
||||
"""Test AgentState type for injection."""
|
||||
|
||||
def test_agent_state_exists(self):
|
||||
"""AgentState type should exist."""
|
||||
from praisonaiagents.tools.injected import AgentState
|
||||
|
||||
assert AgentState is not None
|
||||
|
||||
def test_agent_state_has_required_fields(self):
|
||||
"""AgentState should have agent_id, run_id, session_id."""
|
||||
from praisonaiagents.tools.injected import AgentState
|
||||
|
||||
state = AgentState(
|
||||
agent_id="agent-123",
|
||||
run_id="run-456",
|
||||
session_id="session-789"
|
||||
)
|
||||
|
||||
assert state.agent_id == "agent-123"
|
||||
assert state.run_id == "run-456"
|
||||
assert state.session_id == "session-789"
|
||||
|
||||
def test_agent_state_optional_fields(self):
|
||||
"""AgentState should have optional fields."""
|
||||
from praisonaiagents.tools.injected import AgentState
|
||||
|
||||
state = AgentState(
|
||||
agent_id="a",
|
||||
run_id="r",
|
||||
session_id="s",
|
||||
last_user_message="Hello",
|
||||
last_agent_message="Hi there",
|
||||
memory=Mock(),
|
||||
previous_tool_results=["result1", "result2"]
|
||||
)
|
||||
|
||||
assert state.last_user_message == "Hello"
|
||||
assert state.last_agent_message == "Hi there"
|
||||
assert state.memory is not None
|
||||
assert len(state.previous_tool_results) == 2
|
||||
|
||||
|
||||
class TestToolDecoratorInjection:
|
||||
"""Test @tool decorator handles Injected parameters."""
|
||||
|
||||
def test_tool_decorator_detects_injected(self):
|
||||
"""@tool should detect Injected parameters."""
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
"""A tool with injected state."""
|
||||
return f"query={query}, session={state.get('session_id')}"
|
||||
|
||||
# Should have metadata about injected params
|
||||
assert hasattr(my_tool, '_injected_params') or hasattr(my_tool, 'injected_params')
|
||||
|
||||
def test_injected_params_not_in_schema(self):
|
||||
"""Injected parameters should not appear in tool schema."""
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
"""A tool with injected state."""
|
||||
return query
|
||||
|
||||
schema = my_tool.get_schema()
|
||||
params = schema['function']['parameters']['properties']
|
||||
|
||||
# 'query' should be in schema
|
||||
assert 'query' in params
|
||||
# 'state' should NOT be in schema (it's injected)
|
||||
assert 'state' not in params
|
||||
|
||||
def test_injected_params_not_required(self):
|
||||
"""Injected parameters should not be in required list."""
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
"""A tool with injected state."""
|
||||
return query
|
||||
|
||||
schema = my_tool.get_schema()
|
||||
required = schema['function']['parameters'].get('required', [])
|
||||
|
||||
assert 'state' not in required
|
||||
assert 'query' in required
|
||||
|
||||
|
||||
class TestToolExecutionInjection:
|
||||
"""Test that injected state is provided at runtime."""
|
||||
|
||||
def test_injected_state_provided_at_runtime(self):
|
||||
"""When tool is called, injected state should be provided."""
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
from praisonaiagents.tools.injected import AgentState
|
||||
|
||||
received_state = [None]
|
||||
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[AgentState]) -> str:
|
||||
"""A tool with injected state."""
|
||||
received_state[0] = state
|
||||
return f"session={state.session_id}"
|
||||
|
||||
# Simulate execution with injection context
|
||||
from praisonaiagents.tools.injected import with_injection_context
|
||||
|
||||
mock_state = AgentState(
|
||||
agent_id="a",
|
||||
run_id="r",
|
||||
session_id="test-session"
|
||||
)
|
||||
|
||||
with with_injection_context(mock_state):
|
||||
result = my_tool(query="hello")
|
||||
|
||||
assert received_state[0] is not None
|
||||
assert received_state[0].session_id == "test-session"
|
||||
|
||||
def test_injected_dict_state(self):
|
||||
"""Injected[dict] should provide dict with state info."""
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
from praisonaiagents.tools.injected import with_injection_context, AgentState
|
||||
|
||||
received_state = [None]
|
||||
|
||||
@tool
|
||||
def my_tool(query: str, state: Injected[dict]) -> str:
|
||||
"""A tool with injected dict state."""
|
||||
received_state[0] = state
|
||||
return "ok"
|
||||
|
||||
mock_state = AgentState(
|
||||
agent_id="a",
|
||||
run_id="r",
|
||||
session_id="s"
|
||||
)
|
||||
|
||||
with with_injection_context(mock_state):
|
||||
my_tool(query="test")
|
||||
|
||||
assert isinstance(received_state[0], dict)
|
||||
assert 'session_id' in received_state[0]
|
||||
|
||||
|
||||
class TestAgentToolExecution:
|
||||
"""Test Agent executes tools with injection."""
|
||||
|
||||
def test_agent_injects_state_to_tools(self):
|
||||
"""Agent should inject state when executing tools."""
|
||||
from praisonaiagents import Agent, tool
|
||||
from praisonaiagents.tools import Injected
|
||||
|
||||
received_session = [None]
|
||||
|
||||
@tool
|
||||
def show_state(query: str, state: Injected[dict]) -> str:
|
||||
"""Show the injected state."""
|
||||
received_session[0] = state.get('session_id')
|
||||
return f"session={state.get('session_id')}"
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
tools=[show_state],
|
||||
session_id="my-session-123"
|
||||
)
|
||||
|
||||
# Execute the tool directly through agent
|
||||
result = agent.execute_tool("show_state", {"query": "test"})
|
||||
|
||||
# State should have been injected
|
||||
assert received_session[0] == "my-session-123"
|
||||
|
||||
|
||||
class TestInjectionContextManager:
|
||||
"""Test injection context manager."""
|
||||
|
||||
def test_with_injection_context_exists(self):
|
||||
"""with_injection_context should exist."""
|
||||
from praisonaiagents.tools.injected import with_injection_context
|
||||
|
||||
assert with_injection_context is not None
|
||||
|
||||
def test_context_manager_sets_state(self):
|
||||
"""Context manager should set current state."""
|
||||
from praisonaiagents.tools.injected import (
|
||||
with_injection_context,
|
||||
get_current_state,
|
||||
AgentState
|
||||
)
|
||||
|
||||
state = AgentState(agent_id="a", run_id="r", session_id="s")
|
||||
|
||||
with with_injection_context(state):
|
||||
current = get_current_state()
|
||||
assert current is state
|
||||
|
||||
# After context, should be None
|
||||
assert get_current_state() is None
|
||||
|
||||
def test_nested_contexts(self):
|
||||
"""Nested contexts should work correctly."""
|
||||
from praisonaiagents.tools.injected import (
|
||||
with_injection_context,
|
||||
get_current_state,
|
||||
AgentState
|
||||
)
|
||||
|
||||
state1 = AgentState(agent_id="a1", run_id="r1", session_id="s1")
|
||||
state2 = AgentState(agent_id="a2", run_id="r2", session_id="s2")
|
||||
|
||||
with with_injection_context(state1):
|
||||
assert get_current_state().agent_id == "a1"
|
||||
|
||||
with with_injection_context(state2):
|
||||
assert get_current_state().agent_id == "a2"
|
||||
|
||||
# Back to state1
|
||||
assert get_current_state().agent_id == "a1"
|
||||
|
||||
|
||||
class TestAsyncInjection:
|
||||
"""Test async tool injection."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_tool_injection(self):
|
||||
"""Async tools should also receive injected state."""
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
from praisonaiagents.tools.injected import with_injection_context, AgentState
|
||||
|
||||
received_state = [None]
|
||||
|
||||
@tool
|
||||
async def async_tool(query: str, state: Injected[dict]) -> str:
|
||||
"""Async tool with injected state."""
|
||||
received_state[0] = state
|
||||
return "ok"
|
||||
|
||||
state = AgentState(agent_id="a", run_id="r", session_id="async-session")
|
||||
|
||||
with with_injection_context(state):
|
||||
await async_tool(query="test")
|
||||
|
||||
assert received_state[0] is not None
|
||||
assert received_state[0].get('session_id') == "async-session"
|
||||
|
||||
|
||||
class TestInjectionZeroOverhead:
|
||||
"""Test that injection has zero overhead when not used."""
|
||||
|
||||
def test_no_injected_params_no_overhead(self):
|
||||
"""Tools without Injected params should have no overhead."""
|
||||
from praisonaiagents import tool
|
||||
|
||||
@tool
|
||||
def simple_tool(query: str) -> str:
|
||||
"""Simple tool without injection."""
|
||||
return query
|
||||
|
||||
# Should work normally without any injection context
|
||||
result = simple_tool(query="hello")
|
||||
assert result == "hello"
|
||||
|
||||
def test_schema_generation_fast(self):
|
||||
"""Schema generation should be fast even with Injected params."""
|
||||
from praisonaiagents import tool
|
||||
from praisonaiagents.tools import Injected
|
||||
import time
|
||||
|
||||
@tool
|
||||
def my_tool(a: str, b: int, c: Injected[dict], d: float) -> str:
|
||||
"""Tool with mixed params."""
|
||||
return a
|
||||
|
||||
start = time.monotonic()
|
||||
for _ in range(100):
|
||||
my_tool.get_schema()
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
# Should be very fast
|
||||
assert elapsed < 0.1 # 100 calls in < 100ms
|
||||
|
||||
|
||||
class TestBackwardCompatibility:
|
||||
"""Test backward compatibility with existing tools."""
|
||||
|
||||
def test_existing_tools_still_work(self):
|
||||
"""Existing tools without Injected should continue to work."""
|
||||
from praisonaiagents import tool
|
||||
|
||||
@tool
|
||||
def legacy_tool(query: str, count: int = 5) -> str:
|
||||
"""Legacy tool."""
|
||||
return f"{query} x {count}"
|
||||
|
||||
result = legacy_tool(query="test", count=3)
|
||||
assert result == "test x 3"
|
||||
|
||||
def test_existing_tool_schema_unchanged(self):
|
||||
"""Existing tool schemas should be unchanged."""
|
||||
from praisonaiagents import tool
|
||||
|
||||
@tool
|
||||
def legacy_tool(query: str, count: int = 5) -> str:
|
||||
"""Legacy tool."""
|
||||
return f"{query} x {count}"
|
||||
|
||||
schema = legacy_tool.get_schema()
|
||||
params = schema['function']['parameters']['properties']
|
||||
|
||||
assert 'query' in params
|
||||
assert 'count' in params
|
||||
assert len(params) == 2
|
||||
@@ -0,0 +1,403 @@
|
||||
"""
|
||||
TDD Tests for Middleware System (wrap_model_call, wrap_tool_call).
|
||||
|
||||
These tests are written FIRST before implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, MagicMock, patch, AsyncMock
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
class TestMiddlewareTypes:
|
||||
"""Test middleware type definitions."""
|
||||
|
||||
def test_invocation_context_exists(self):
|
||||
"""InvocationContext dataclass should exist with required fields."""
|
||||
from praisonaiagents.hooks.middleware import InvocationContext
|
||||
|
||||
ctx = InvocationContext(
|
||||
agent_id="agent-123",
|
||||
run_id="run-456",
|
||||
session_id="session-789"
|
||||
)
|
||||
assert ctx.agent_id == "agent-123"
|
||||
assert ctx.run_id == "run-456"
|
||||
assert ctx.session_id == "session-789"
|
||||
|
||||
def test_invocation_context_optional_fields(self):
|
||||
"""InvocationContext should have optional tool_name, model_name, metadata."""
|
||||
from praisonaiagents.hooks.middleware import InvocationContext
|
||||
|
||||
ctx = InvocationContext(
|
||||
agent_id="a",
|
||||
run_id="r",
|
||||
session_id="s",
|
||||
tool_name="my_tool",
|
||||
model_name="gpt-4o",
|
||||
metadata={"key": "value"}
|
||||
)
|
||||
assert ctx.tool_name == "my_tool"
|
||||
assert ctx.model_name == "gpt-4o"
|
||||
assert ctx.metadata == {"key": "value"}
|
||||
|
||||
def test_model_request_exists(self):
|
||||
"""ModelRequest dataclass for before_model hooks."""
|
||||
from praisonaiagents.hooks.middleware import ModelRequest
|
||||
|
||||
req = ModelRequest(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
model="gpt-4o",
|
||||
temperature=0.7,
|
||||
context=Mock()
|
||||
)
|
||||
assert req.messages == [{"role": "user", "content": "Hello"}]
|
||||
assert req.model == "gpt-4o"
|
||||
assert req.temperature == 0.7
|
||||
|
||||
def test_model_response_exists(self):
|
||||
"""ModelResponse dataclass for after_model hooks."""
|
||||
from praisonaiagents.hooks.middleware import ModelResponse
|
||||
|
||||
resp = ModelResponse(
|
||||
content="Hello back!",
|
||||
model="gpt-4o",
|
||||
usage={"prompt_tokens": 10, "completion_tokens": 5},
|
||||
context=Mock()
|
||||
)
|
||||
assert resp.content == "Hello back!"
|
||||
assert resp.model == "gpt-4o"
|
||||
|
||||
def test_tool_request_exists(self):
|
||||
"""ToolRequest dataclass for before_tool hooks."""
|
||||
from praisonaiagents.hooks.middleware import ToolRequest
|
||||
|
||||
req = ToolRequest(
|
||||
tool_name="get_weather",
|
||||
arguments={"city": "Paris"},
|
||||
context=Mock()
|
||||
)
|
||||
assert req.tool_name == "get_weather"
|
||||
assert req.arguments == {"city": "Paris"}
|
||||
|
||||
def test_tool_response_exists(self):
|
||||
"""ToolResponse dataclass for after_tool hooks."""
|
||||
from praisonaiagents.hooks.middleware import ToolResponse
|
||||
|
||||
resp = ToolResponse(
|
||||
tool_name="get_weather",
|
||||
result="Sunny, 22°C",
|
||||
error=None,
|
||||
context=Mock()
|
||||
)
|
||||
assert resp.result == "Sunny, 22°C"
|
||||
assert resp.error is None
|
||||
|
||||
|
||||
class TestMiddlewareDecorators:
|
||||
"""Test middleware decorator functions."""
|
||||
|
||||
def test_before_model_decorator_exists(self):
|
||||
"""before_model decorator should exist and tag functions."""
|
||||
from praisonaiagents.hooks import before_model
|
||||
|
||||
@before_model
|
||||
def my_hook(request):
|
||||
return request
|
||||
|
||||
assert hasattr(my_hook, '_hook_type')
|
||||
assert my_hook._hook_type == 'before_model'
|
||||
|
||||
def test_after_model_decorator_exists(self):
|
||||
"""after_model decorator should exist and tag functions."""
|
||||
from praisonaiagents.hooks import after_model
|
||||
|
||||
@after_model
|
||||
def my_hook(response):
|
||||
return response
|
||||
|
||||
assert hasattr(my_hook, '_hook_type')
|
||||
assert my_hook._hook_type == 'after_model'
|
||||
|
||||
def test_wrap_model_call_decorator_exists(self):
|
||||
"""wrap_model_call decorator should exist and tag functions."""
|
||||
from praisonaiagents.hooks import wrap_model_call
|
||||
|
||||
@wrap_model_call
|
||||
def my_middleware(request, call_next):
|
||||
return call_next(request)
|
||||
|
||||
assert hasattr(my_middleware, '_hook_type')
|
||||
assert my_middleware._hook_type == 'wrap_model_call'
|
||||
|
||||
def test_before_tool_decorator_exists(self):
|
||||
"""before_tool decorator should exist and tag functions."""
|
||||
from praisonaiagents.hooks import before_tool
|
||||
|
||||
@before_tool
|
||||
def my_hook(request):
|
||||
return request
|
||||
|
||||
assert hasattr(my_hook, '_hook_type')
|
||||
assert my_hook._hook_type == 'before_tool'
|
||||
|
||||
def test_after_tool_decorator_exists(self):
|
||||
"""after_tool decorator should exist and tag functions."""
|
||||
from praisonaiagents.hooks import after_tool
|
||||
|
||||
@after_tool
|
||||
def my_hook(response):
|
||||
return response
|
||||
|
||||
assert hasattr(my_hook, '_hook_type')
|
||||
assert my_hook._hook_type == 'after_tool'
|
||||
|
||||
def test_wrap_tool_call_decorator_exists(self):
|
||||
"""wrap_tool_call decorator should exist and tag functions."""
|
||||
from praisonaiagents.hooks import wrap_tool_call
|
||||
|
||||
@wrap_tool_call
|
||||
def my_middleware(request, call_next):
|
||||
return call_next(request)
|
||||
|
||||
assert hasattr(my_middleware, '_hook_type')
|
||||
assert my_middleware._hook_type == 'wrap_tool_call'
|
||||
|
||||
|
||||
class TestMiddlewareChaining:
|
||||
"""Test middleware chain composition."""
|
||||
|
||||
def test_wrap_model_call_chain_order(self):
|
||||
"""wrap_model_call hooks should compose in registration order."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
calls = []
|
||||
|
||||
def mw1(req, call_next):
|
||||
calls.append("mw1_before")
|
||||
result = call_next(req)
|
||||
calls.append("mw1_after")
|
||||
return result
|
||||
|
||||
def mw2(req, call_next):
|
||||
calls.append("mw2_before")
|
||||
result = call_next(req)
|
||||
calls.append("mw2_after")
|
||||
return result
|
||||
|
||||
def final_handler(req):
|
||||
calls.append("final")
|
||||
return "response"
|
||||
|
||||
chain = MiddlewareChain([mw1, mw2])
|
||||
result = chain.execute({"test": True}, final_handler)
|
||||
|
||||
assert calls == ["mw1_before", "mw2_before", "final", "mw2_after", "mw1_after"]
|
||||
assert result == "response"
|
||||
|
||||
def test_empty_middleware_chain(self):
|
||||
"""Empty middleware chain should just call final handler."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
def final_handler(req):
|
||||
return "direct"
|
||||
|
||||
chain = MiddlewareChain([])
|
||||
result = chain.execute({"test": True}, final_handler)
|
||||
assert result == "direct"
|
||||
|
||||
def test_middleware_can_modify_request(self):
|
||||
"""Middleware should be able to modify request before passing on."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
def add_context(req, call_next):
|
||||
req["extra"] = "added"
|
||||
return call_next(req)
|
||||
|
||||
def final_handler(req):
|
||||
return req.get("extra")
|
||||
|
||||
chain = MiddlewareChain([add_context])
|
||||
result = chain.execute({"original": True}, final_handler)
|
||||
assert result == "added"
|
||||
|
||||
def test_middleware_can_modify_response(self):
|
||||
"""Middleware should be able to modify response after call_next."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
def uppercase_response(req, call_next):
|
||||
result = call_next(req)
|
||||
return result.upper()
|
||||
|
||||
def final_handler(req):
|
||||
return "hello"
|
||||
|
||||
chain = MiddlewareChain([uppercase_response])
|
||||
result = chain.execute({}, final_handler)
|
||||
assert result == "HELLO"
|
||||
|
||||
def test_middleware_can_short_circuit(self):
|
||||
"""Middleware can return early without calling call_next."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
def block_middleware(req, call_next):
|
||||
if req.get("block"):
|
||||
return "blocked"
|
||||
return call_next(req)
|
||||
|
||||
def final_handler(req):
|
||||
return "reached"
|
||||
|
||||
chain = MiddlewareChain([block_middleware])
|
||||
|
||||
# Should be blocked
|
||||
result = chain.execute({"block": True}, final_handler)
|
||||
assert result == "blocked"
|
||||
|
||||
# Should pass through
|
||||
result = chain.execute({"block": False}, final_handler)
|
||||
assert result == "reached"
|
||||
|
||||
|
||||
class TestMiddlewareRetry:
|
||||
"""Test retry logic in middleware."""
|
||||
|
||||
def test_wrap_tool_call_retry_on_error(self):
|
||||
"""wrap_tool_call should support retry patterns."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
attempt_count = [0]
|
||||
|
||||
def retry_middleware(req, call_next):
|
||||
last_error = None
|
||||
for i in range(3):
|
||||
try:
|
||||
return call_next(req)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
raise last_error
|
||||
|
||||
def flaky_handler(req):
|
||||
attempt_count[0] += 1
|
||||
if attempt_count[0] < 3:
|
||||
raise RuntimeError("flaky")
|
||||
return "success"
|
||||
|
||||
chain = MiddlewareChain([retry_middleware])
|
||||
result = chain.execute({}, flaky_handler)
|
||||
assert result == "success"
|
||||
assert attempt_count[0] == 3
|
||||
|
||||
|
||||
class TestAsyncMiddleware:
|
||||
"""Test async middleware support."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_wrap_model_call(self):
|
||||
"""Async middleware should work with async handlers."""
|
||||
from praisonaiagents.hooks.middleware import AsyncMiddlewareChain
|
||||
|
||||
calls = []
|
||||
|
||||
async def async_mw(req, call_next):
|
||||
calls.append("before")
|
||||
result = await call_next(req)
|
||||
calls.append("after")
|
||||
return result
|
||||
|
||||
async def async_handler(req):
|
||||
calls.append("handler")
|
||||
return "async_result"
|
||||
|
||||
chain = AsyncMiddlewareChain([async_mw])
|
||||
result = await chain.execute({}, async_handler)
|
||||
|
||||
assert calls == ["before", "handler", "after"]
|
||||
assert result == "async_result"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_middleware_chain_order(self):
|
||||
"""Async middleware should maintain order."""
|
||||
from praisonaiagents.hooks.middleware import AsyncMiddlewareChain
|
||||
|
||||
order = []
|
||||
|
||||
async def mw1(req, call_next):
|
||||
order.append(1)
|
||||
result = await call_next(req)
|
||||
order.append(-1)
|
||||
return result
|
||||
|
||||
async def mw2(req, call_next):
|
||||
order.append(2)
|
||||
result = await call_next(req)
|
||||
order.append(-2)
|
||||
return result
|
||||
|
||||
async def handler(req):
|
||||
order.append(0)
|
||||
return "done"
|
||||
|
||||
chain = AsyncMiddlewareChain([mw1, mw2])
|
||||
await chain.execute({}, handler)
|
||||
|
||||
assert order == [1, 2, 0, -2, -1]
|
||||
|
||||
|
||||
class TestAgentMiddlewareIntegration:
|
||||
"""Test middleware integration with Agent class."""
|
||||
|
||||
def test_agent_accepts_hooks_parameter(self):
|
||||
"""Agent should accept hooks parameter."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.hooks import before_model
|
||||
|
||||
@before_model
|
||||
def my_hook(req):
|
||||
return req
|
||||
|
||||
# Should not raise
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
hooks=[my_hook]
|
||||
)
|
||||
assert agent is not None
|
||||
|
||||
def test_agent_hooks_empty_by_default(self):
|
||||
"""Agent should have empty hooks by default."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", instructions="Test")
|
||||
# Should have no hooks or empty hooks list
|
||||
hooks = getattr(agent, '_hooks', None) or getattr(agent, 'hooks', [])
|
||||
assert hooks is None or len(hooks) == 0
|
||||
|
||||
|
||||
class TestMiddlewareZeroOverhead:
|
||||
"""Test that middleware has zero overhead when not used."""
|
||||
|
||||
def test_no_hooks_fast_path(self):
|
||||
"""When no hooks registered, should take fast path."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
# Empty chain should be essentially a no-op
|
||||
chain = MiddlewareChain([])
|
||||
|
||||
call_count = [0]
|
||||
def handler(req):
|
||||
call_count[0] += 1
|
||||
return "result"
|
||||
|
||||
result = chain.execute({}, handler)
|
||||
assert result == "result"
|
||||
assert call_count[0] == 1
|
||||
|
||||
def test_middleware_chain_is_none_check(self):
|
||||
"""Should handle None middleware list gracefully."""
|
||||
from praisonaiagents.hooks.middleware import MiddlewareChain
|
||||
|
||||
chain = MiddlewareChain(None)
|
||||
result = chain.execute({}, lambda r: "ok")
|
||||
assert result == "ok"
|
||||
@@ -0,0 +1,260 @@
|
||||
"""
|
||||
TDD Tests for Rate Limiter.
|
||||
|
||||
These tests are written FIRST before implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import time
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
|
||||
class TestRateLimiterBasic:
|
||||
"""Test basic rate limiter functionality."""
|
||||
|
||||
def test_rate_limiter_exists(self):
|
||||
"""RateLimiter class should exist in praisonaiagents.llm."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
assert limiter is not None
|
||||
|
||||
def test_rate_limiter_rpm_parameter(self):
|
||||
"""RateLimiter should accept requests_per_minute."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
assert limiter.requests_per_minute == 60
|
||||
|
||||
def test_rate_limiter_acquire_sync(self):
|
||||
"""RateLimiter.acquire() should block if rate exceeded."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
# 1 request per second = 60 per minute
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
|
||||
# First acquire should be instant
|
||||
start = time.monotonic()
|
||||
limiter.acquire()
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 0.1 # Should be nearly instant
|
||||
|
||||
def test_rate_limiter_blocks_when_exceeded(self):
|
||||
"""RateLimiter should block when rate is exceeded."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
# Very low rate: 6 per minute = 1 per 10 seconds
|
||||
# But we'll use a mock clock to test without waiting
|
||||
limiter = RateLimiter(requests_per_minute=6)
|
||||
|
||||
# Inject a mock clock for testing
|
||||
mock_time = [0.0]
|
||||
limiter._get_time = lambda: mock_time[0]
|
||||
limiter._sleep = lambda s: mock_time.__setitem__(0, mock_time[0] + s)
|
||||
|
||||
# First request should pass
|
||||
limiter.acquire()
|
||||
|
||||
# Second request should need to wait ~10 seconds
|
||||
# (but our mock sleep just advances time)
|
||||
limiter.acquire()
|
||||
|
||||
# Time should have advanced
|
||||
assert mock_time[0] >= 10.0
|
||||
|
||||
def test_rate_limiter_token_bucket(self):
|
||||
"""RateLimiter should use token bucket algorithm."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60, burst=5)
|
||||
|
||||
# Reset internal state with mock time
|
||||
mock_time = [0.0]
|
||||
limiter._get_time = lambda: mock_time[0]
|
||||
limiter._sleep = lambda s: mock_time.__setitem__(0, mock_time[0] + s)
|
||||
limiter._tokens = 5.0 # Reset tokens to burst size
|
||||
limiter._last_update = 0.0 # Reset last update
|
||||
|
||||
for _ in range(5):
|
||||
limiter.acquire()
|
||||
|
||||
# Should not have waited much (burst allows 5 instant requests)
|
||||
assert mock_time[0] < 1.0
|
||||
|
||||
|
||||
class TestRateLimiterAsync:
|
||||
"""Test async rate limiter functionality."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limiter_acquire_async(self):
|
||||
"""RateLimiter.acquire_async() should work with async."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
|
||||
# Should have async acquire method
|
||||
await limiter.acquire_async()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limiter_async_blocks(self):
|
||||
"""Async acquire should also respect rate limits."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
import asyncio
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=6)
|
||||
|
||||
# Mock async sleep
|
||||
sleep_times = []
|
||||
original_sleep = asyncio.sleep
|
||||
|
||||
async def mock_sleep(s):
|
||||
sleep_times.append(s)
|
||||
# Don't actually sleep in tests
|
||||
|
||||
with patch('asyncio.sleep', mock_sleep):
|
||||
await limiter.acquire_async()
|
||||
await limiter.acquire_async()
|
||||
|
||||
# Should have tried to sleep
|
||||
|
||||
|
||||
class TestRateLimiterAgentIntegration:
|
||||
"""Test rate limiter integration with Agent."""
|
||||
|
||||
def test_agent_accepts_rate_limiter(self):
|
||||
"""Agent should accept rate_limiter parameter."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
rate_limiter=limiter
|
||||
)
|
||||
assert agent is not None
|
||||
assert getattr(agent, '_rate_limiter', None) is limiter or \
|
||||
getattr(agent, 'rate_limiter', None) is limiter
|
||||
|
||||
def test_agent_rate_limiter_applied_to_llm_calls(self):
|
||||
"""Rate limiter should be applied before LLM calls."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
acquire_count = [0]
|
||||
original_acquire = limiter.acquire
|
||||
|
||||
def counting_acquire():
|
||||
acquire_count[0] += 1
|
||||
return original_acquire()
|
||||
|
||||
limiter.acquire = counting_acquire
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent",
|
||||
rate_limiter=limiter
|
||||
)
|
||||
|
||||
with patch.object(agent, '_chat_completion', return_value="mocked"):
|
||||
agent.chat("Hello")
|
||||
|
||||
# Rate limiter should have been called
|
||||
assert acquire_count[0] >= 1
|
||||
|
||||
|
||||
class TestRateLimiterZeroOverhead:
|
||||
"""Test that rate limiter has zero overhead when not used."""
|
||||
|
||||
def test_no_rate_limiter_no_overhead(self):
|
||||
"""When no rate limiter, should not add any overhead."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
name="Test",
|
||||
instructions="Test agent"
|
||||
# No rate_limiter
|
||||
)
|
||||
|
||||
# Should not have rate limiter
|
||||
assert getattr(agent, '_rate_limiter', None) is None
|
||||
|
||||
def test_rate_limiter_none_check_fast(self):
|
||||
"""Checking for None rate limiter should be O(1)."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
# This is more of a design requirement than a test
|
||||
# The implementation should do: if self._rate_limiter: self._rate_limiter.acquire()
|
||||
|
||||
|
||||
class TestRateLimiterConfiguration:
|
||||
"""Test rate limiter configuration options."""
|
||||
|
||||
def test_rate_limiter_default_burst(self):
|
||||
"""RateLimiter should have sensible default burst."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
# Default burst should be 1 or based on rpm
|
||||
assert hasattr(limiter, 'burst') or hasattr(limiter, '_burst')
|
||||
|
||||
def test_rate_limiter_custom_burst(self):
|
||||
"""RateLimiter should accept custom burst parameter."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60, burst=10)
|
||||
burst = getattr(limiter, 'burst', None) or getattr(limiter, '_burst', None)
|
||||
assert burst == 10
|
||||
|
||||
def test_rate_limiter_tpm_optional(self):
|
||||
"""RateLimiter should optionally accept tokens_per_minute."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
# Should not raise
|
||||
limiter = RateLimiter(
|
||||
requests_per_minute=60,
|
||||
tokens_per_minute=100000
|
||||
)
|
||||
assert limiter is not None
|
||||
|
||||
|
||||
class TestRateLimiterDeterministic:
|
||||
"""Test that rate limiter is deterministic for testing."""
|
||||
|
||||
def test_injectable_clock(self):
|
||||
"""RateLimiter should allow injecting clock for testing."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
mock_time = [0.0]
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
limiter._get_time = lambda: mock_time[0]
|
||||
|
||||
# Should use our mock time
|
||||
t1 = limiter._get_time()
|
||||
mock_time[0] = 100.0
|
||||
t2 = limiter._get_time()
|
||||
|
||||
assert t1 == 0.0
|
||||
assert t2 == 100.0
|
||||
|
||||
def test_injectable_sleep(self):
|
||||
"""RateLimiter should allow injecting sleep for testing."""
|
||||
from praisonaiagents.llm import RateLimiter
|
||||
|
||||
sleep_calls = []
|
||||
|
||||
limiter = RateLimiter(requests_per_minute=60)
|
||||
limiter._sleep = lambda s: sleep_calls.append(s)
|
||||
|
||||
# Force a wait scenario
|
||||
limiter._get_time = lambda: 0.0
|
||||
limiter._tokens = 0 # No tokens available
|
||||
limiter._last_update = 0.0
|
||||
|
||||
limiter.acquire()
|
||||
|
||||
# Should have called our mock sleep
|
||||
# (actual behavior depends on implementation)
|
||||
Generated
+1
-1
@@ -2422,7 +2422,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.2.1"
|
||||
version = "0.3.4"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
{
|
||||
"version": "1.0.0",
|
||||
"lastUpdated": "2024-12-26",
|
||||
"implementationStatus": "Phase 3 Complete - Full Documentation",
|
||||
"summary": {
|
||||
"pythonCoreFeatures": 85,
|
||||
"pythonWrapperFeatures": 45,
|
||||
"typescriptFeatures": 48,
|
||||
"gapCount": 25,
|
||||
"priorityP0": 12,
|
||||
"priorityP1": 15,
|
||||
"priorityP2": 20,
|
||||
"priorityP3": 20
|
||||
},
|
||||
"pythonCoreSDK": {
|
||||
"path": "/Users/praison/praisonai-package/src/praisonai-agents/praisonaiagents",
|
||||
"exports": [
|
||||
"Agent", "ImageAgent", "ContextAgent", "DeepResearchAgent", "QueryRewriterAgent", "PromptExpanderAgent",
|
||||
"PraisonAIAgents", "Agents", "Task", "Tools", "BaseTool", "ToolResult", "tool", "FunctionTool",
|
||||
"ToolRegistry", "get_registry", "register_tool", "get_tool", "AutoAgents", "Session", "Memory",
|
||||
"Knowledge", "Chunking", "db", "obs", "MCP", "Workflow", "Pipeline", "WorkflowStep", "WorkflowContext",
|
||||
"StepResult", "Route", "Parallel", "Loop", "Repeat", "GuardrailResult", "LLMGuardrail",
|
||||
"Handoff", "handoff", "FlowDisplay", "FastContext", "SkillManager", "SkillProperties",
|
||||
"get_telemetry", "enable_telemetry", "MinimalTelemetry", "TelemetryCollector", "AGUI", "A2A"
|
||||
],
|
||||
"modules": {
|
||||
"agent": ["Agent", "ImageAgent", "ContextAgent", "DeepResearchAgent", "QueryRewriterAgent", "PromptExpanderAgent", "RouterAgent", "Handoff"],
|
||||
"agents": ["PraisonAIAgents", "AutoAgents"],
|
||||
"task": ["Task"],
|
||||
"tools": ["Tools", "BaseTool", "ToolResult", "tool", "FunctionTool", "ToolRegistry", "40+ built-in tools"],
|
||||
"llm": ["LLM", "OpenAIClient", "ModelRouter", "RateLimiter", "ModelCapabilities"],
|
||||
"memory": ["Memory", "AutoMemory", "FileMemory", "RulesManager", "DocsManager", "Hooks"],
|
||||
"knowledge": ["Knowledge", "Chunking", "VectorStore", "QueryEngine", "Readers", "Rerankers"],
|
||||
"db": ["DbAdapter", "AsyncDbAdapter", "DbMessage", "DbToolCall", "DbRun", "DbSpan", "DbTrace"],
|
||||
"obs": ["ObservabilityAdapter", "Span", "Trace", "16 provider integrations"],
|
||||
"mcp": ["MCP", "MCPServer", "MCPSession", "MCPTransport", "MCPWebSocket", "MCPSSE", "MCPHTTPStream"],
|
||||
"workflows": ["Workflow", "Pipeline", "WorkflowStep", "WorkflowContext", "StepResult", "Route", "Parallel", "Loop", "Repeat", "YAMLWorkflowParser"],
|
||||
"guardrails": ["GuardrailResult", "LLMGuardrail"],
|
||||
"eval": ["AccuracyEval", "PerformanceEval", "ReliabilityEval", "CriteriaEval"],
|
||||
"skills": ["SkillManager", "SkillProperties", "SkillLoader", "SkillDiscovery", "SkillValidator"],
|
||||
"telemetry": ["MinimalTelemetry", "TelemetryCollector", "PerformanceMonitor", "TokenCollector"],
|
||||
"context": ["FastContext", "FastContextResult"],
|
||||
"session": ["Session"],
|
||||
"ui": ["AGUI", "A2A"]
|
||||
}
|
||||
},
|
||||
"pythonWrapper": {
|
||||
"path": "/Users/praison/praisonai-package/src/praisonai/praisonai",
|
||||
"exports": ["PraisonAI"],
|
||||
"cliFeatures": [
|
||||
"slash_commands", "autonomy_mode", "cost_tracker", "repo_map", "interactive_tui",
|
||||
"git_integration", "sandbox_executor", "agent_scheduler", "agents", "at_mentions",
|
||||
"auto_memory", "background", "checkpoints", "compare", "eval", "external_agents",
|
||||
"fast_context", "flow_display", "guardrail", "handoff", "hooks", "image", "jobs",
|
||||
"knowledge", "mcp", "message_queue", "metrics", "n8n", "observability", "persistence",
|
||||
"profile", "profiling", "router", "session", "skills", "telemetry", "todo", "tools", "workflow"
|
||||
]
|
||||
},
|
||||
"typescriptSDK": {
|
||||
"path": "/Users/praison/praisonai-package/src/praisonai-ts/src",
|
||||
"exports": ["Agent", "PraisonAIAgents", "Task", "Tool", "BaseTool", "Memory", "BaseMemoryStore", "Process", "BaseProcess", "OpenAIService", "Logger"],
|
||||
"modules": {
|
||||
"agent": ["Agent (simple)", "Agent (types/proxy)", "PraisonAIAgents"],
|
||||
"task": ["Task"],
|
||||
"tools": ["Tool interface", "BaseTool", "arxivTools", "mcpSse"],
|
||||
"llm": ["OpenAIService (OpenAI only)"],
|
||||
"memory": ["Memory interface", "BaseMemoryStore (in-memory only)"],
|
||||
"process": ["Process interface", "BaseProcess"],
|
||||
"utils": ["Logger"]
|
||||
}
|
||||
},
|
||||
"gapMatrix": {
|
||||
"P0_CoreParity": [
|
||||
{"feature": "Multi-provider LLM support (LiteLLM equivalent)", "python": true, "typescript": false, "priority": "P0", "effort": "high"},
|
||||
{"feature": "Streaming responses", "python": true, "typescript": "partial", "priority": "P0", "effort": "medium"},
|
||||
{"feature": "Structured outputs (JSON schema)", "python": true, "typescript": false, "priority": "P0", "effort": "medium"},
|
||||
{"feature": "Tool/function calling with validation", "python": true, "typescript": "partial", "priority": "P0", "effort": "medium"},
|
||||
{"feature": "Session management (session_id, run_id)", "python": true, "typescript": false, "priority": "P0", "effort": "medium"},
|
||||
{"feature": "Multi-agent orchestration patterns", "python": true, "typescript": "partial", "priority": "P0", "effort": "high"},
|
||||
{"feature": "Agent handoff", "python": true, "typescript": false, "priority": "P0", "effort": "medium"},
|
||||
{"feature": "Tool decorator (@tool)", "python": true, "typescript": false, "priority": "P0", "effort": "low"},
|
||||
{"feature": "Tool registry", "python": true, "typescript": false, "priority": "P0", "effort": "medium"},
|
||||
{"feature": "Async agent execution", "python": true, "typescript": "partial", "priority": "P0", "effort": "low"},
|
||||
{"feature": "Error handling and retries", "python": true, "typescript": "partial", "priority": "P0", "effort": "medium"},
|
||||
{"feature": "Verbose/debug logging", "python": true, "typescript": "partial", "priority": "P0", "effort": "low"}
|
||||
],
|
||||
"P1_Persistence": [
|
||||
{"feature": "Database adapter protocol", "python": true, "typescript": false, "priority": "P1", "effort": "high"},
|
||||
{"feature": "SQLite adapter", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "PostgreSQL adapter", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "Redis adapter", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "Message persistence", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "Run/trace persistence", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "Memory persistence (vector)", "python": true, "typescript": false, "priority": "P1", "effort": "high"},
|
||||
{"feature": "Knowledge base (RAG)", "python": true, "typescript": false, "priority": "P1", "effort": "high"},
|
||||
{"feature": "Chunking strategies", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "Vector store integration", "python": true, "typescript": false, "priority": "P1", "effort": "high"},
|
||||
{"feature": "Semantic search/recall", "python": true, "typescript": false, "priority": "P1", "effort": "high"},
|
||||
{"feature": "Working memory", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "File-based memory", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "Auto-memory", "python": true, "typescript": false, "priority": "P1", "effort": "medium"},
|
||||
{"feature": "Session export/import", "python": true, "typescript": false, "priority": "P1", "effort": "medium"}
|
||||
],
|
||||
"P2_CLI": [
|
||||
{"feature": "CLI entrypoint", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Slash commands", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Interactive TUI", "python": true, "typescript": false, "priority": "P2", "effort": "high"},
|
||||
{"feature": "Cost tracking", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Repo map", "python": true, "typescript": false, "priority": "P2", "effort": "high"},
|
||||
{"feature": "Git integration", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Sandbox executor", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Autonomy modes", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Agent scheduler", "python": true, "typescript": false, "priority": "P2", "effort": "high"},
|
||||
{"feature": "Background jobs", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Checkpoints", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Flow display", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Profiling", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Metrics", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "External agent integrations", "python": true, "typescript": false, "priority": "P2", "effort": "high"},
|
||||
{"feature": "N8N integration", "python": true, "typescript": false, "priority": "P2", "effort": "high"},
|
||||
{"feature": "MCP CLI support", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Skills CLI", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Eval CLI", "python": true, "typescript": false, "priority": "P2", "effort": "medium"},
|
||||
{"feature": "Workflow CLI", "python": true, "typescript": false, "priority": "P2", "effort": "medium"}
|
||||
],
|
||||
"P3_Advanced": [
|
||||
{"feature": "Workflows/Pipelines", "python": true, "typescript": false, "priority": "P3", "effort": "high"},
|
||||
{"feature": "Route pattern", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Parallel pattern", "python": true, "typescript": "partial", "priority": "P3", "effort": "low"},
|
||||
{"feature": "Loop pattern", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Repeat pattern", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "YAML workflow parser", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Guardrails", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "LLM Guardrail", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Evaluation framework", "python": true, "typescript": false, "priority": "P3", "effort": "high"},
|
||||
{"feature": "Accuracy eval", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Performance eval", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Reliability eval", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Criteria eval", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Observability protocol", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Telemetry", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Performance monitoring", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Token tracking", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Skills system", "python": true, "typescript": false, "priority": "P3", "effort": "high"},
|
||||
{"feature": "Fast context", "python": true, "typescript": false, "priority": "P3", "effort": "medium"},
|
||||
{"feature": "Image agent", "python": true, "typescript": false, "priority": "P3", "effort": "medium"}
|
||||
]
|
||||
},
|
||||
"mastraCapabilities": {
|
||||
"implemented": [
|
||||
"Agent core", "Tool calling", "Streaming (partial)", "Multi-agent (partial)"
|
||||
],
|
||||
"missing": [
|
||||
"Memory with semantic recall", "Working memory", "Thread management", "Resource scoping",
|
||||
"Storage abstraction", "Vector store integration", "Embeddings", "Evaluation scorers",
|
||||
"Workflow engine", "Step execution", "Branching logic", "Event-driven workflows",
|
||||
"MCP integration", "Voice/TTS", "Observability hooks", "Request context",
|
||||
"Tool validation", "Tool suspend/resume", "Provider options per tool"
|
||||
]
|
||||
},
|
||||
"implementationPlan": {
|
||||
"phase1_P0": {
|
||||
"description": "Core agent parity - streaming, structured outputs, sessions, tool system",
|
||||
"tasks": [
|
||||
"Implement AI SDK provider abstraction (OpenAI, Anthropic, Google)",
|
||||
"Add streaming response support with proper typing",
|
||||
"Implement structured output with JSON schema/Zod",
|
||||
"Add session management (session_id, run_id, trace_id)",
|
||||
"Implement tool decorator and registry",
|
||||
"Add agent handoff support",
|
||||
"Improve error handling and retries"
|
||||
]
|
||||
},
|
||||
"phase2_P1": {
|
||||
"description": "Persistence layer - db adapters, memory, knowledge",
|
||||
"tasks": [
|
||||
"Implement DbAdapter protocol",
|
||||
"Add SQLite adapter",
|
||||
"Add message/run/trace persistence",
|
||||
"Implement Memory with persistence",
|
||||
"Add Knowledge base (RAG) support",
|
||||
"Implement vector store integration"
|
||||
]
|
||||
},
|
||||
"phase3_P2": {
|
||||
"description": "CLI parity",
|
||||
"tasks": [
|
||||
"Create CLI entrypoint",
|
||||
"Implement slash commands",
|
||||
"Add cost tracking",
|
||||
"Implement interactive mode"
|
||||
]
|
||||
},
|
||||
"phase4_P3": {
|
||||
"description": "Advanced features",
|
||||
"tasks": [
|
||||
"Implement Workflow engine",
|
||||
"Add workflow patterns (Route, Parallel, Loop, Repeat)",
|
||||
"Implement Guardrails",
|
||||
"Add Evaluation framework",
|
||||
"Implement Observability hooks",
|
||||
"Add Skills system"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,7 @@ const config: Config.InitialOptions = {
|
||||
preset: 'ts-jest',
|
||||
testEnvironment: 'node',
|
||||
roots: ['<rootDir>/tests'],
|
||||
testMatch: ['**/*.main.test.ts', '**/*.spec.ts'],
|
||||
testMatch: ['**/*.test.ts', '**/*.spec.ts'],
|
||||
moduleFileExtensions: ['ts', 'js', 'json'],
|
||||
coverageDirectory: 'coverage',
|
||||
collectCoverageFrom: [
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
{
|
||||
"name": "praisonai",
|
||||
"version": "1.0.19",
|
||||
"version": "1.1.0",
|
||||
"description": "PraisonAI TypeScript AI Agents Framework - Node.js, npm, and Javascript AI Agents Framework",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
"bin": {
|
||||
"praisonai-ts": "./dist/cli/index.js"
|
||||
},
|
||||
"scripts": {
|
||||
"build": "tsc",
|
||||
"test": "jest",
|
||||
@@ -62,7 +65,6 @@
|
||||
"fast-xml-parser": "^4.5.1",
|
||||
"node-fetch": "^2.6.9",
|
||||
"openai": "^4.81.0",
|
||||
"praisonai": "^1.0.19",
|
||||
"@modelcontextprotocol/sdk": "^1.12.1"
|
||||
},
|
||||
"optionalDependencies": {
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
/**
|
||||
* Context Agent - Agent with enhanced context management
|
||||
*/
|
||||
|
||||
import { createProvider, type LLMProvider, type Message } from '../llm/providers';
|
||||
import { Session } from '../session';
|
||||
import { KnowledgeBase, type SearchResult } from '../knowledge/rag';
|
||||
|
||||
export interface ContextAgentConfig {
|
||||
name?: string;
|
||||
instructions: string;
|
||||
llm?: string;
|
||||
knowledgeBase?: KnowledgeBase;
|
||||
contextWindow?: number;
|
||||
maxContextTokens?: number;
|
||||
verbose?: boolean;
|
||||
}
|
||||
|
||||
export interface ContextMessage {
|
||||
role: 'system' | 'user' | 'assistant';
|
||||
content: string;
|
||||
timestamp: number;
|
||||
metadata?: Record<string, any>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Context Agent - Agent with RAG and context management
|
||||
*/
|
||||
export class ContextAgent {
|
||||
readonly name: string;
|
||||
readonly instructions: string;
|
||||
|
||||
private provider: LLMProvider;
|
||||
private knowledgeBase?: KnowledgeBase;
|
||||
private contextWindow: number;
|
||||
private maxContextTokens: number;
|
||||
private messages: ContextMessage[] = [];
|
||||
private verbose: boolean;
|
||||
|
||||
constructor(config: ContextAgentConfig) {
|
||||
this.name = config.name || `ContextAgent_${Math.random().toString(36).substr(2, 9)}`;
|
||||
this.instructions = config.instructions;
|
||||
this.provider = createProvider(config.llm || 'gpt-4o-mini');
|
||||
this.knowledgeBase = config.knowledgeBase;
|
||||
this.contextWindow = config.contextWindow ?? 10;
|
||||
this.maxContextTokens = config.maxContextTokens ?? 4000;
|
||||
this.verbose = config.verbose ?? false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Chat with context awareness
|
||||
*/
|
||||
async chat(prompt: string): Promise<{ text: string; context?: SearchResult[] }> {
|
||||
// Add user message
|
||||
this.messages.push({
|
||||
role: 'user',
|
||||
content: prompt,
|
||||
timestamp: Date.now()
|
||||
});
|
||||
|
||||
// Build context from knowledge base
|
||||
let ragContext: SearchResult[] = [];
|
||||
let contextString = '';
|
||||
|
||||
if (this.knowledgeBase) {
|
||||
ragContext = await this.knowledgeBase.search(prompt);
|
||||
if (ragContext.length > 0) {
|
||||
contextString = this.knowledgeBase.buildContext(ragContext);
|
||||
if (this.verbose) {
|
||||
console.log(`[ContextAgent] Found ${ragContext.length} relevant documents`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build messages array with context
|
||||
const systemPrompt = contextString
|
||||
? `${this.instructions}\n\nRelevant context:\n${contextString}`
|
||||
: this.instructions;
|
||||
|
||||
const messages: Message[] = [
|
||||
{ role: 'system', content: systemPrompt },
|
||||
...this.getRecentMessages()
|
||||
];
|
||||
|
||||
// Generate response
|
||||
const result = await this.provider.generateText({ messages });
|
||||
|
||||
// Add assistant message
|
||||
this.messages.push({
|
||||
role: 'assistant',
|
||||
content: result.text,
|
||||
timestamp: Date.now()
|
||||
});
|
||||
|
||||
return {
|
||||
text: result.text,
|
||||
context: ragContext.length > 0 ? ragContext : undefined
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get recent messages within context window
|
||||
*/
|
||||
private getRecentMessages(): Message[] {
|
||||
const recent = this.messages.slice(-this.contextWindow);
|
||||
return recent.map(m => ({
|
||||
role: m.role,
|
||||
content: m.content
|
||||
}));
|
||||
}
|
||||
|
||||
/**
|
||||
* Add document to knowledge base
|
||||
*/
|
||||
async addDocument(id: string, content: string, metadata?: Record<string, any>): Promise<void> {
|
||||
if (!this.knowledgeBase) {
|
||||
throw new Error('No knowledge base configured');
|
||||
}
|
||||
await this.knowledgeBase.add({ id, content, metadata });
|
||||
}
|
||||
|
||||
/**
|
||||
* Search knowledge base
|
||||
*/
|
||||
async searchKnowledge(query: string, limit?: number): Promise<SearchResult[]> {
|
||||
if (!this.knowledgeBase) {
|
||||
return [];
|
||||
}
|
||||
return this.knowledgeBase.search(query, limit);
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear conversation history
|
||||
*/
|
||||
clearHistory(): void {
|
||||
this.messages = [];
|
||||
}
|
||||
|
||||
/**
|
||||
* Get conversation history
|
||||
*/
|
||||
getHistory(): ContextMessage[] {
|
||||
return [...this.messages];
|
||||
}
|
||||
|
||||
/**
|
||||
* Set knowledge base
|
||||
*/
|
||||
setKnowledgeBase(kb: KnowledgeBase): void {
|
||||
this.knowledgeBase = kb;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a context agent
|
||||
*/
|
||||
export function createContextAgent(config: ContextAgentConfig): ContextAgent {
|
||||
return new ContextAgent(config);
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
/**
|
||||
* Enhanced Agent - Agent with session management, provider abstraction, and tool support
|
||||
*/
|
||||
|
||||
import { createProvider, type LLMProvider, type Message, type ToolCall, type GenerateTextResult } from '../llm/providers';
|
||||
import { Session, Run, Trace, getSessionManager } from '../session';
|
||||
import { FunctionTool, ToolRegistry, tool } from '../tools/decorator';
|
||||
import { Logger } from '../utils/logger';
|
||||
|
||||
export interface EnhancedAgentConfig {
|
||||
name?: string;
|
||||
instructions: string;
|
||||
llm?: string;
|
||||
tools?: Array<FunctionTool | Function | { name: string; description?: string; parameters?: any; execute: Function }>;
|
||||
session?: Session;
|
||||
sessionId?: string;
|
||||
verbose?: boolean;
|
||||
stream?: boolean;
|
||||
maxToolCalls?: number;
|
||||
temperature?: number;
|
||||
maxTokens?: number;
|
||||
outputSchema?: any;
|
||||
}
|
||||
|
||||
export interface ChatOptions {
|
||||
stream?: boolean;
|
||||
temperature?: number;
|
||||
maxTokens?: number;
|
||||
outputSchema?: any;
|
||||
onToken?: (token: string) => void;
|
||||
}
|
||||
|
||||
export interface ChatResult {
|
||||
text: string;
|
||||
structured?: any;
|
||||
toolCalls?: ToolCall[];
|
||||
usage: {
|
||||
promptTokens: number;
|
||||
completionTokens: number;
|
||||
totalTokens: number;
|
||||
};
|
||||
runId: string;
|
||||
sessionId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Enhanced Agent with full feature support
|
||||
*/
|
||||
export class EnhancedAgent {
|
||||
readonly name: string;
|
||||
readonly instructions: string;
|
||||
readonly sessionId: string;
|
||||
|
||||
private provider: LLMProvider;
|
||||
private session: Session;
|
||||
private toolRegistry: ToolRegistry;
|
||||
private verbose: boolean;
|
||||
private stream: boolean;
|
||||
private maxToolCalls: number;
|
||||
private temperature: number;
|
||||
private maxTokens?: number;
|
||||
private outputSchema?: any;
|
||||
|
||||
constructor(config: EnhancedAgentConfig) {
|
||||
this.name = config.name || `Agent_${Math.random().toString(36).substr(2, 9)}`;
|
||||
this.instructions = config.instructions;
|
||||
this.verbose = config.verbose ?? true;
|
||||
this.stream = config.stream ?? true;
|
||||
this.maxToolCalls = config.maxToolCalls ?? 10;
|
||||
this.temperature = config.temperature ?? 0.7;
|
||||
this.maxTokens = config.maxTokens;
|
||||
this.outputSchema = config.outputSchema;
|
||||
|
||||
// Initialize provider
|
||||
this.provider = createProvider(config.llm || 'gpt-4o-mini');
|
||||
|
||||
// Initialize session
|
||||
if (config.session) {
|
||||
this.session = config.session;
|
||||
} else if (config.sessionId) {
|
||||
this.session = getSessionManager().getOrCreate(config.sessionId);
|
||||
} else {
|
||||
this.session = getSessionManager().create();
|
||||
}
|
||||
this.sessionId = this.session.id;
|
||||
|
||||
// Initialize tool registry
|
||||
this.toolRegistry = new ToolRegistry();
|
||||
if (config.tools) {
|
||||
this.registerTools(config.tools);
|
||||
}
|
||||
|
||||
Logger.setVerbose(this.verbose);
|
||||
}
|
||||
|
||||
private registerTools(tools: EnhancedAgentConfig['tools']): void {
|
||||
if (!tools) return;
|
||||
|
||||
for (const t of tools) {
|
||||
if (t instanceof FunctionTool) {
|
||||
this.toolRegistry.register(t);
|
||||
} else if (typeof t === 'function') {
|
||||
// Convert function to tool
|
||||
const funcName = t.name || `function_${Math.random().toString(36).substr(2, 9)}`;
|
||||
this.toolRegistry.register(tool({
|
||||
name: funcName,
|
||||
description: `Function ${funcName}`,
|
||||
execute: t as any,
|
||||
}));
|
||||
} else if (typeof t === 'object' && 'execute' in t) {
|
||||
// Tool config object
|
||||
this.toolRegistry.register(tool({
|
||||
name: t.name,
|
||||
description: t.description,
|
||||
parameters: t.parameters,
|
||||
execute: t.execute as any,
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a tool to the agent
|
||||
*/
|
||||
addTool(t: FunctionTool | { name: string; description?: string; parameters?: any; execute: Function }): this {
|
||||
if (t instanceof FunctionTool) {
|
||||
this.toolRegistry.register(t);
|
||||
} else {
|
||||
this.toolRegistry.register(tool({
|
||||
name: t.name,
|
||||
description: t.description,
|
||||
parameters: t.parameters,
|
||||
execute: t.execute as any,
|
||||
}));
|
||||
}
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all registered tools
|
||||
*/
|
||||
getTools(): FunctionTool[] {
|
||||
return this.toolRegistry.list();
|
||||
}
|
||||
|
||||
/**
|
||||
* Chat with the agent
|
||||
*/
|
||||
async chat(prompt: string, options: ChatOptions = {}): Promise<ChatResult> {
|
||||
const run = this.session.createRun().start();
|
||||
const trace = run.createTrace({ name: 'chat' }).start();
|
||||
|
||||
try {
|
||||
// Add user message to session
|
||||
this.session.addMessage({ role: 'user', content: prompt });
|
||||
|
||||
// Build messages array
|
||||
const messages: Message[] = [
|
||||
{ role: 'system', content: this.instructions },
|
||||
...this.session.getMessagesForLLM() as Message[],
|
||||
];
|
||||
|
||||
// Get tool definitions
|
||||
const tools = this.toolRegistry.list().length > 0
|
||||
? this.toolRegistry.getDefinitions()
|
||||
: undefined;
|
||||
|
||||
let result: GenerateTextResult;
|
||||
let fullText = '';
|
||||
let allToolCalls: ToolCall[] = [];
|
||||
let iterations = 0;
|
||||
|
||||
// Tool calling loop
|
||||
while (iterations < this.maxToolCalls) {
|
||||
iterations++;
|
||||
|
||||
if (options.stream && !tools && options.onToken) {
|
||||
// Streaming without tools
|
||||
const stream = await this.provider.streamText({
|
||||
messages,
|
||||
temperature: options.temperature ?? this.temperature,
|
||||
maxTokens: options.maxTokens ?? this.maxTokens,
|
||||
onToken: options.onToken,
|
||||
});
|
||||
|
||||
let usage = { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
|
||||
for await (const chunk of stream) {
|
||||
if (chunk.text) fullText += chunk.text;
|
||||
if (chunk.usage) usage = { ...usage, ...chunk.usage };
|
||||
}
|
||||
|
||||
result = {
|
||||
text: fullText,
|
||||
usage,
|
||||
finishReason: 'stop',
|
||||
};
|
||||
break;
|
||||
} else if (options.outputSchema) {
|
||||
// Structured output
|
||||
const objResult = await this.provider.generateObject({
|
||||
messages,
|
||||
schema: options.outputSchema ?? this.outputSchema,
|
||||
temperature: options.temperature ?? this.temperature,
|
||||
maxTokens: options.maxTokens ?? this.maxTokens,
|
||||
});
|
||||
|
||||
result = {
|
||||
text: JSON.stringify(objResult.object),
|
||||
usage: objResult.usage,
|
||||
finishReason: 'stop',
|
||||
};
|
||||
fullText = result.text;
|
||||
break;
|
||||
} else {
|
||||
// Regular generation with potential tool calls
|
||||
result = await this.provider.generateText({
|
||||
messages,
|
||||
temperature: options.temperature ?? this.temperature,
|
||||
maxTokens: options.maxTokens ?? this.maxTokens,
|
||||
tools,
|
||||
});
|
||||
|
||||
if (result.toolCalls && result.toolCalls.length > 0) {
|
||||
allToolCalls.push(...result.toolCalls);
|
||||
|
||||
// Add assistant message with tool calls
|
||||
messages.push({
|
||||
role: 'assistant',
|
||||
content: result.text || null,
|
||||
tool_calls: result.toolCalls,
|
||||
});
|
||||
|
||||
// Execute tool calls
|
||||
for (const tc of result.toolCalls) {
|
||||
const toolTrace = trace.createChild({ name: `tool:${tc.function.name}` }).start();
|
||||
|
||||
try {
|
||||
const toolFn = this.toolRegistry.get(tc.function.name);
|
||||
if (!toolFn) {
|
||||
throw new Error(`Tool '${tc.function.name}' not found`);
|
||||
}
|
||||
|
||||
const args = JSON.parse(tc.function.arguments);
|
||||
const toolResult = await toolFn.execute(args, {
|
||||
agentName: this.name,
|
||||
sessionId: this.sessionId,
|
||||
runId: run.id,
|
||||
});
|
||||
|
||||
const resultStr = typeof toolResult === 'string'
|
||||
? toolResult
|
||||
: JSON.stringify(toolResult);
|
||||
|
||||
messages.push({
|
||||
role: 'tool',
|
||||
content: resultStr,
|
||||
tool_call_id: tc.id,
|
||||
});
|
||||
|
||||
toolTrace.complete({ result: resultStr });
|
||||
} catch (error: any) {
|
||||
messages.push({
|
||||
role: 'tool',
|
||||
content: `Error: ${error.message}`,
|
||||
tool_call_id: tc.id,
|
||||
});
|
||||
toolTrace.fail(error);
|
||||
}
|
||||
}
|
||||
|
||||
// Continue loop to get final response
|
||||
continue;
|
||||
}
|
||||
|
||||
// No tool calls, we have our final response
|
||||
fullText = result.text;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Add assistant message to session
|
||||
this.session.addMessage({
|
||||
role: 'assistant',
|
||||
content: fullText,
|
||||
tool_calls: allToolCalls.length > 0 ? allToolCalls : undefined,
|
||||
});
|
||||
|
||||
trace.complete({ responseLength: fullText.length });
|
||||
run.complete();
|
||||
|
||||
// Parse structured output if schema was provided
|
||||
let structured: any;
|
||||
if (options.outputSchema) {
|
||||
try {
|
||||
structured = JSON.parse(fullText);
|
||||
} catch {
|
||||
// Not valid JSON
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
text: fullText,
|
||||
structured,
|
||||
toolCalls: allToolCalls.length > 0 ? allToolCalls : undefined,
|
||||
usage: result!.usage,
|
||||
runId: run.id,
|
||||
sessionId: this.sessionId,
|
||||
};
|
||||
} catch (error: any) {
|
||||
trace.fail(error);
|
||||
run.fail(error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Start method for compatibility
|
||||
*/
|
||||
async start(prompt: string, options?: ChatOptions): Promise<string> {
|
||||
const result = await this.chat(prompt, options);
|
||||
return result.text;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the session
|
||||
*/
|
||||
getSession(): Session {
|
||||
return this.session;
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear conversation history
|
||||
*/
|
||||
clearHistory(): void {
|
||||
this.session.clearMessages();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
/**
|
||||
* Agent Handoff - Transfer conversations between agents
|
||||
*/
|
||||
|
||||
import type { EnhancedAgent } from './enhanced';
|
||||
|
||||
export interface HandoffConfig {
|
||||
agent: EnhancedAgent;
|
||||
name?: string;
|
||||
description?: string;
|
||||
condition?: (context: HandoffContext) => boolean;
|
||||
transformContext?: (messages: any[]) => any[];
|
||||
}
|
||||
|
||||
export interface HandoffContext {
|
||||
messages: any[];
|
||||
lastMessage: string;
|
||||
topic?: string;
|
||||
metadata?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface HandoffResult {
|
||||
handedOffTo: string;
|
||||
response: string;
|
||||
context: HandoffContext;
|
||||
}
|
||||
|
||||
/**
|
||||
* Handoff class - Represents a handoff target
|
||||
*/
|
||||
export class Handoff {
|
||||
readonly targetAgent: EnhancedAgent;
|
||||
readonly name: string;
|
||||
readonly description: string;
|
||||
readonly condition?: (context: HandoffContext) => boolean;
|
||||
readonly transformContext?: (messages: any[]) => any[];
|
||||
|
||||
constructor(config: HandoffConfig) {
|
||||
this.targetAgent = config.agent;
|
||||
this.name = config.name || `handoff_to_${config.agent.name}`;
|
||||
this.description = config.description || `Transfer conversation to ${config.agent.name}`;
|
||||
this.condition = config.condition;
|
||||
this.transformContext = config.transformContext;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if handoff should be triggered
|
||||
*/
|
||||
shouldTrigger(context: HandoffContext): boolean {
|
||||
if (this.condition) {
|
||||
return this.condition(context);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute the handoff
|
||||
*/
|
||||
async execute(context: HandoffContext): Promise<HandoffResult> {
|
||||
let messages = context.messages;
|
||||
|
||||
if (this.transformContext) {
|
||||
messages = this.transformContext(messages);
|
||||
}
|
||||
|
||||
// Transfer context to target agent
|
||||
const response = await this.targetAgent.chat(context.lastMessage);
|
||||
|
||||
return {
|
||||
handedOffTo: this.targetAgent.name,
|
||||
response: response.text,
|
||||
context: {
|
||||
...context,
|
||||
messages,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get tool definition for LLM
|
||||
*/
|
||||
getToolDefinition(): { name: string; description: string; parameters: any } {
|
||||
return {
|
||||
name: this.name,
|
||||
description: this.description,
|
||||
parameters: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
reason: {
|
||||
type: 'string',
|
||||
description: 'Reason for the handoff',
|
||||
},
|
||||
},
|
||||
required: [],
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a handoff configuration
|
||||
*/
|
||||
export function handoff(config: HandoffConfig): Handoff {
|
||||
return new Handoff(config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Handoff filters - Common condition functions
|
||||
*/
|
||||
export const handoffFilters = {
|
||||
/**
|
||||
* Trigger handoff based on topic keywords
|
||||
*/
|
||||
topic: (keywords: string | string[]) => {
|
||||
const keywordList = Array.isArray(keywords) ? keywords : [keywords];
|
||||
return (context: HandoffContext): boolean => {
|
||||
const lastMessage = context.lastMessage.toLowerCase();
|
||||
return keywordList.some(kw => lastMessage.includes(kw.toLowerCase()));
|
||||
};
|
||||
},
|
||||
|
||||
/**
|
||||
* Trigger handoff based on metadata
|
||||
*/
|
||||
metadata: (key: string, value: any) => {
|
||||
return (context: HandoffContext): boolean => {
|
||||
return context.metadata?.[key] === value;
|
||||
};
|
||||
},
|
||||
|
||||
/**
|
||||
* Always trigger
|
||||
*/
|
||||
always: () => {
|
||||
return (): boolean => true;
|
||||
},
|
||||
|
||||
/**
|
||||
* Never trigger (manual only)
|
||||
*/
|
||||
never: () => {
|
||||
return (): boolean => false;
|
||||
},
|
||||
|
||||
/**
|
||||
* Combine multiple conditions with AND
|
||||
*/
|
||||
and: (...conditions: Array<(context: HandoffContext) => boolean>) => {
|
||||
return (context: HandoffContext): boolean => {
|
||||
return conditions.every(cond => cond(context));
|
||||
};
|
||||
},
|
||||
|
||||
/**
|
||||
* Combine multiple conditions with OR
|
||||
*/
|
||||
or: (...conditions: Array<(context: HandoffContext) => boolean>) => {
|
||||
return (context: HandoffContext): boolean => {
|
||||
return conditions.some(cond => cond(context));
|
||||
};
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,146 @@
|
||||
/**
|
||||
* Router Agent - Route requests to specialized agents
|
||||
*/
|
||||
|
||||
import type { EnhancedAgent } from './enhanced';
|
||||
|
||||
export interface RouteConfig {
|
||||
agent: EnhancedAgent;
|
||||
condition: (input: string, context?: RouteContext) => boolean | Promise<boolean>;
|
||||
priority?: number;
|
||||
}
|
||||
|
||||
export interface RouteContext {
|
||||
history?: string[];
|
||||
metadata?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface RouterConfig {
|
||||
name?: string;
|
||||
routes: RouteConfig[];
|
||||
defaultAgent?: EnhancedAgent;
|
||||
verbose?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Router Agent - Routes requests to the most appropriate agent
|
||||
*/
|
||||
export class RouterAgent {
|
||||
readonly name: string;
|
||||
private routes: RouteConfig[];
|
||||
private defaultAgent?: EnhancedAgent;
|
||||
private verbose: boolean;
|
||||
|
||||
constructor(config: RouterConfig) {
|
||||
this.name = config.name || 'Router';
|
||||
this.routes = config.routes.sort((a, b) => (b.priority || 0) - (a.priority || 0));
|
||||
this.defaultAgent = config.defaultAgent;
|
||||
this.verbose = config.verbose ?? false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Route a request to the appropriate agent
|
||||
*/
|
||||
async route(input: string, context?: RouteContext): Promise<{ agent: EnhancedAgent; response: string } | null> {
|
||||
for (const route of this.routes) {
|
||||
const matches = await route.condition(input, context);
|
||||
if (matches) {
|
||||
if (this.verbose) {
|
||||
console.log(`[Router] Routing to: ${route.agent.name}`);
|
||||
}
|
||||
const response = await route.agent.chat(input);
|
||||
return { agent: route.agent, response: response.text };
|
||||
}
|
||||
}
|
||||
|
||||
if (this.defaultAgent) {
|
||||
if (this.verbose) {
|
||||
console.log(`[Router] Using default agent: ${this.defaultAgent.name}`);
|
||||
}
|
||||
const response = await this.defaultAgent.chat(input);
|
||||
return { agent: this.defaultAgent, response: response.text };
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a route
|
||||
*/
|
||||
addRoute(config: RouteConfig): this {
|
||||
this.routes.push(config);
|
||||
this.routes.sort((a, b) => (b.priority || 0) - (a.priority || 0));
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all routes
|
||||
*/
|
||||
getRoutes(): RouteConfig[] {
|
||||
return [...this.routes];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Route condition helpers
|
||||
*/
|
||||
export const routeConditions = {
|
||||
/**
|
||||
* Match by keywords
|
||||
*/
|
||||
keywords: (keywords: string | string[]) => {
|
||||
const keywordList = Array.isArray(keywords) ? keywords : [keywords];
|
||||
return (input: string): boolean => {
|
||||
const lower = input.toLowerCase();
|
||||
return keywordList.some(kw => lower.includes(kw.toLowerCase()));
|
||||
};
|
||||
},
|
||||
|
||||
/**
|
||||
* Match by regex
|
||||
*/
|
||||
pattern: (regex: RegExp) => {
|
||||
return (input: string): boolean => regex.test(input);
|
||||
},
|
||||
|
||||
/**
|
||||
* Match by metadata
|
||||
*/
|
||||
metadata: (key: string, value: any) => {
|
||||
return (_input: string, context?: RouteContext): boolean => {
|
||||
return context?.metadata?.[key] === value;
|
||||
};
|
||||
},
|
||||
|
||||
/**
|
||||
* Always match (for default routes)
|
||||
*/
|
||||
always: () => {
|
||||
return (): boolean => true;
|
||||
},
|
||||
|
||||
/**
|
||||
* Combine conditions with AND
|
||||
*/
|
||||
and: (...conditions: Array<(input: string, context?: RouteContext) => boolean>) => {
|
||||
return (input: string, context?: RouteContext): boolean => {
|
||||
return conditions.every(c => c(input, context));
|
||||
};
|
||||
},
|
||||
|
||||
/**
|
||||
* Combine conditions with OR
|
||||
*/
|
||||
or: (...conditions: Array<(input: string, context?: RouteContext) => boolean>) => {
|
||||
return (input: string, context?: RouteContext): boolean => {
|
||||
return conditions.some(c => c(input, context));
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Create a router agent
|
||||
*/
|
||||
export function createRouter(config: RouterConfig): RouterAgent {
|
||||
return new RouterAgent(config);
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env node
|
||||
/**
|
||||
* PraisonAI TypeScript CLI
|
||||
*/
|
||||
|
||||
import { createProvider, isProviderAvailable, getAvailableProviders } from '../llm/providers';
|
||||
import { Session } from '../session';
|
||||
import { tool, ToolRegistry } from '../tools/decorator';
|
||||
|
||||
interface CLIOptions {
|
||||
model?: string;
|
||||
stream?: boolean;
|
||||
verbose?: boolean;
|
||||
sessionId?: string;
|
||||
}
|
||||
|
||||
async function chat(prompt: string, options: CLIOptions = {}): Promise<string> {
|
||||
const model = options.model || process.env.PRAISONAI_MODEL || 'openai/gpt-4o-mini';
|
||||
const provider = createProvider(model);
|
||||
|
||||
const session = new Session({ id: options.sessionId });
|
||||
session.addMessage({ role: 'user', content: prompt });
|
||||
|
||||
if (options.stream) {
|
||||
let fullText = '';
|
||||
const stream = await provider.streamText({
|
||||
messages: [{ role: 'user', content: prompt }],
|
||||
onToken: (token) => {
|
||||
process.stdout.write(token);
|
||||
fullText += token;
|
||||
}
|
||||
});
|
||||
|
||||
for await (const chunk of stream) {
|
||||
// Stream is consumed by onToken
|
||||
}
|
||||
console.log(); // New line after streaming
|
||||
return fullText;
|
||||
} else {
|
||||
const result = await provider.generateText({
|
||||
messages: [{ role: 'user', content: prompt }]
|
||||
});
|
||||
console.log(result.text);
|
||||
return result.text;
|
||||
}
|
||||
}
|
||||
|
||||
async function listProviders(): Promise<void> {
|
||||
console.log('Available Providers:');
|
||||
const available = getAvailableProviders();
|
||||
|
||||
const providers = ['openai', 'anthropic', 'google'];
|
||||
for (const p of providers) {
|
||||
const status = available.includes(p) ? '✅' : '❌';
|
||||
console.log(` ${status} ${p}`);
|
||||
}
|
||||
}
|
||||
|
||||
async function version(): Promise<void> {
|
||||
const pkg = require('../../package.json');
|
||||
console.log(`praisonai v${pkg.version}`);
|
||||
}
|
||||
|
||||
async function help(): Promise<void> {
|
||||
console.log(`
|
||||
PraisonAI TypeScript CLI
|
||||
|
||||
Usage:
|
||||
praisonai-ts chat <prompt> Chat with an AI agent
|
||||
praisonai-ts providers List available providers
|
||||
praisonai-ts version Show version
|
||||
praisonai-ts help Show this help
|
||||
|
||||
Options:
|
||||
--model, -m <model> Model to use (e.g., openai/gpt-4o-mini)
|
||||
--stream, -s Enable streaming output
|
||||
--verbose, -v Verbose output
|
||||
--session <id> Session ID for conversation continuity
|
||||
|
||||
Environment Variables:
|
||||
OPENAI_API_KEY OpenAI API key
|
||||
ANTHROPIC_API_KEY Anthropic API key
|
||||
GOOGLE_API_KEY Google API key
|
||||
PRAISONAI_MODEL Default model
|
||||
|
||||
Examples:
|
||||
praisonai-ts chat "Hello, how are you?"
|
||||
praisonai-ts chat "Write a poem" --stream
|
||||
praisonai-ts chat "Explain AI" -m anthropic/claude-sonnet-4-20250514
|
||||
praisonai-ts providers
|
||||
`);
|
||||
}
|
||||
|
||||
function parseArgs(args: string[]): { command: string; prompt: string; options: CLIOptions } {
|
||||
const options: CLIOptions = {};
|
||||
let command = 'help';
|
||||
let prompt = '';
|
||||
|
||||
let i = 0;
|
||||
while (i < args.length) {
|
||||
const arg = args[i];
|
||||
|
||||
if (arg === '--model' || arg === '-m') {
|
||||
options.model = args[++i];
|
||||
} else if (arg === '--stream' || arg === '-s') {
|
||||
options.stream = true;
|
||||
} else if (arg === '--verbose' || arg === '-v') {
|
||||
options.verbose = true;
|
||||
} else if (arg === '--session') {
|
||||
options.sessionId = args[++i];
|
||||
} else if (!arg.startsWith('-')) {
|
||||
if (!command || command === 'help') {
|
||||
command = arg;
|
||||
} else {
|
||||
prompt = arg;
|
||||
}
|
||||
}
|
||||
i++;
|
||||
}
|
||||
|
||||
return { command, prompt, options };
|
||||
}
|
||||
|
||||
async function main(): Promise<void> {
|
||||
const args = process.argv.slice(2);
|
||||
const { command, prompt, options } = parseArgs(args);
|
||||
|
||||
try {
|
||||
switch (command) {
|
||||
case 'chat':
|
||||
if (!prompt) {
|
||||
console.error('Error: Please provide a prompt');
|
||||
process.exit(1);
|
||||
}
|
||||
await chat(prompt, options);
|
||||
break;
|
||||
case 'providers':
|
||||
await listProviders();
|
||||
break;
|
||||
case 'version':
|
||||
await version();
|
||||
break;
|
||||
case 'help':
|
||||
default:
|
||||
await help();
|
||||
break;
|
||||
}
|
||||
} catch (error: any) {
|
||||
console.error('Error:', error.message);
|
||||
if (options.verbose) {
|
||||
console.error(error.stack);
|
||||
}
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
// Export for programmatic use
|
||||
export { chat, listProviders, version, help, parseArgs };
|
||||
|
||||
// Run CLI if executed directly
|
||||
if (require.main === module) {
|
||||
main();
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
/**
|
||||
* Evaluation Framework - Accuracy, Performance, and Reliability evaluation
|
||||
*/
|
||||
|
||||
export interface EvalResult {
|
||||
passed: boolean;
|
||||
score: number;
|
||||
message?: string;
|
||||
details?: Record<string, any>;
|
||||
duration: number;
|
||||
}
|
||||
|
||||
export interface AccuracyEvalConfig {
|
||||
input: string;
|
||||
expectedOutput: string;
|
||||
actualOutput: string;
|
||||
threshold?: number;
|
||||
}
|
||||
|
||||
export interface PerformanceEvalConfig {
|
||||
func: () => Promise<any>;
|
||||
iterations?: number;
|
||||
warmupRuns?: number;
|
||||
}
|
||||
|
||||
export interface PerformanceResult extends EvalResult {
|
||||
avgTime: number;
|
||||
minTime: number;
|
||||
maxTime: number;
|
||||
p95Time: number;
|
||||
times: number[];
|
||||
}
|
||||
|
||||
export interface ReliabilityEvalConfig {
|
||||
expectedToolCalls: string[];
|
||||
actualToolCalls: string[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Accuracy Evaluation - Compare actual output to expected
|
||||
*/
|
||||
export async function accuracyEval(config: AccuracyEvalConfig): Promise<EvalResult> {
|
||||
const start = Date.now();
|
||||
const threshold = config.threshold ?? 0.8;
|
||||
|
||||
const similarity = calculateSimilarity(config.expectedOutput, config.actualOutput);
|
||||
const passed = similarity >= threshold;
|
||||
|
||||
return {
|
||||
passed,
|
||||
score: similarity,
|
||||
message: passed ? 'Output matches expected' : 'Output does not match expected',
|
||||
details: {
|
||||
expected: config.expectedOutput,
|
||||
actual: config.actualOutput,
|
||||
threshold
|
||||
},
|
||||
duration: Date.now() - start
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Performance Evaluation - Measure execution time
|
||||
*/
|
||||
export async function performanceEval(config: PerformanceEvalConfig): Promise<PerformanceResult> {
|
||||
const iterations = config.iterations ?? 10;
|
||||
const warmupRuns = config.warmupRuns ?? 2;
|
||||
const times: number[] = [];
|
||||
|
||||
// Warmup runs
|
||||
for (let i = 0; i < warmupRuns; i++) {
|
||||
await config.func();
|
||||
}
|
||||
|
||||
// Actual runs
|
||||
const start = Date.now();
|
||||
for (let i = 0; i < iterations; i++) {
|
||||
const runStart = Date.now();
|
||||
await config.func();
|
||||
times.push(Date.now() - runStart);
|
||||
}
|
||||
|
||||
const sortedTimes = [...times].sort((a, b) => a - b);
|
||||
const avgTime = times.reduce((a, b) => a + b, 0) / times.length;
|
||||
const minTime = sortedTimes[0];
|
||||
const maxTime = sortedTimes[sortedTimes.length - 1];
|
||||
const p95Index = Math.floor(sortedTimes.length * 0.95);
|
||||
const p95Time = sortedTimes[p95Index] || maxTime;
|
||||
|
||||
return {
|
||||
passed: true,
|
||||
score: 1,
|
||||
avgTime,
|
||||
minTime,
|
||||
maxTime,
|
||||
p95Time,
|
||||
times,
|
||||
duration: Date.now() - start,
|
||||
details: {
|
||||
iterations,
|
||||
warmupRuns
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Reliability Evaluation - Check tool call accuracy
|
||||
*/
|
||||
export async function reliabilityEval(config: ReliabilityEvalConfig): Promise<EvalResult> {
|
||||
const start = Date.now();
|
||||
|
||||
const expected = new Set(config.expectedToolCalls);
|
||||
const actual = new Set(config.actualToolCalls);
|
||||
|
||||
const matched = config.expectedToolCalls.filter(t => actual.has(t));
|
||||
const missing = config.expectedToolCalls.filter(t => !actual.has(t));
|
||||
const extra = config.actualToolCalls.filter(t => !expected.has(t));
|
||||
|
||||
const score = expected.size > 0 ? matched.length / expected.size : 1;
|
||||
const passed = missing.length === 0;
|
||||
|
||||
return {
|
||||
passed,
|
||||
score,
|
||||
message: passed ? 'All expected tool calls made' : `Missing tool calls: ${missing.join(', ')}`,
|
||||
details: {
|
||||
matched,
|
||||
missing,
|
||||
extra,
|
||||
expected: config.expectedToolCalls,
|
||||
actual: config.actualToolCalls
|
||||
},
|
||||
duration: Date.now() - start
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculate text similarity (simple Jaccard similarity)
|
||||
*/
|
||||
function calculateSimilarity(a: string, b: string): number {
|
||||
const wordsA = new Set(a.toLowerCase().split(/\s+/));
|
||||
const wordsB = new Set(b.toLowerCase().split(/\s+/));
|
||||
|
||||
const intersection = new Set([...wordsA].filter(x => wordsB.has(x)));
|
||||
const union = new Set([...wordsA, ...wordsB]);
|
||||
|
||||
return union.size > 0 ? intersection.size / union.size : 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Eval Suite - Run multiple evaluations
|
||||
*/
|
||||
export class EvalSuite {
|
||||
private results: Map<string, EvalResult> = new Map();
|
||||
|
||||
async runAccuracy(name: string, config: AccuracyEvalConfig): Promise<EvalResult> {
|
||||
const result = await accuracyEval(config);
|
||||
this.results.set(name, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
async runPerformance(name: string, config: PerformanceEvalConfig): Promise<PerformanceResult> {
|
||||
const result = await performanceEval(config);
|
||||
this.results.set(name, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
async runReliability(name: string, config: ReliabilityEvalConfig): Promise<EvalResult> {
|
||||
const result = await reliabilityEval(config);
|
||||
this.results.set(name, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
getResults(): Map<string, EvalResult> {
|
||||
return new Map(this.results);
|
||||
}
|
||||
|
||||
getSummary(): { total: number; passed: number; failed: number; avgScore: number } {
|
||||
const results = Array.from(this.results.values());
|
||||
const passed = results.filter(r => r.passed).length;
|
||||
const avgScore = results.length > 0
|
||||
? results.reduce((a, b) => a + b.score, 0) / results.length
|
||||
: 0;
|
||||
|
||||
return {
|
||||
total: results.length,
|
||||
passed,
|
||||
failed: results.length - passed,
|
||||
avgScore
|
||||
};
|
||||
}
|
||||
|
||||
printSummary(): void {
|
||||
const summary = this.getSummary();
|
||||
console.log('\n=== Evaluation Summary ===');
|
||||
console.log(`Total: ${summary.total}`);
|
||||
console.log(`Passed: ${summary.passed}`);
|
||||
console.log(`Failed: ${summary.failed}`);
|
||||
console.log(`Avg Score: ${(summary.avgScore * 100).toFixed(1)}%`);
|
||||
|
||||
console.log('\nResults:');
|
||||
for (const [name, result] of this.results) {
|
||||
const status = result.passed ? '✅' : '❌';
|
||||
console.log(` ${status} ${name}: ${(result.score * 100).toFixed(1)}%`);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
/**
|
||||
* Guardrails - Input/output validation and safety checks
|
||||
*/
|
||||
|
||||
export type GuardrailStatus = 'passed' | 'failed' | 'warning';
|
||||
|
||||
export interface GuardrailResult {
|
||||
status: GuardrailStatus;
|
||||
message?: string;
|
||||
details?: Record<string, any>;
|
||||
modifiedContent?: string;
|
||||
}
|
||||
|
||||
export type GuardrailFunction = (
|
||||
content: string,
|
||||
context?: GuardrailContext
|
||||
) => Promise<GuardrailResult> | GuardrailResult;
|
||||
|
||||
export interface GuardrailContext {
|
||||
role: 'input' | 'output';
|
||||
agentName?: string;
|
||||
sessionId?: string;
|
||||
metadata?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface GuardrailConfig {
|
||||
name: string;
|
||||
description?: string;
|
||||
check: GuardrailFunction;
|
||||
onFail?: 'block' | 'warn' | 'modify';
|
||||
}
|
||||
|
||||
/**
|
||||
* Guardrail class
|
||||
*/
|
||||
export class Guardrail {
|
||||
readonly name: string;
|
||||
readonly description: string;
|
||||
readonly check: GuardrailFunction;
|
||||
readonly onFail: 'block' | 'warn' | 'modify';
|
||||
|
||||
constructor(config: GuardrailConfig) {
|
||||
this.name = config.name;
|
||||
this.description = config.description || `Guardrail: ${config.name}`;
|
||||
this.check = config.check;
|
||||
this.onFail = config.onFail || 'block';
|
||||
}
|
||||
|
||||
async run(content: string, context?: GuardrailContext): Promise<GuardrailResult> {
|
||||
try {
|
||||
return await this.check(content, context);
|
||||
} catch (error: any) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: `Guardrail error: ${error.message}`,
|
||||
details: { error: error.message },
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a guardrail
|
||||
*/
|
||||
export function guardrail(config: GuardrailConfig): Guardrail {
|
||||
return new Guardrail(config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Guardrail Manager - Run multiple guardrails
|
||||
*/
|
||||
export class GuardrailManager {
|
||||
private guardrails: Guardrail[] = [];
|
||||
|
||||
add(g: Guardrail): this {
|
||||
this.guardrails.push(g);
|
||||
return this;
|
||||
}
|
||||
|
||||
async runAll(
|
||||
content: string,
|
||||
context?: GuardrailContext
|
||||
): Promise<{ passed: boolean; results: Array<{ name: string; result: GuardrailResult }> }> {
|
||||
const results: Array<{ name: string; result: GuardrailResult }> = [];
|
||||
let passed = true;
|
||||
|
||||
for (const g of this.guardrails) {
|
||||
const result = await g.run(content, context);
|
||||
results.push({ name: g.name, result });
|
||||
|
||||
if (result.status === 'failed') {
|
||||
passed = false;
|
||||
if (g.onFail === 'block') {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { passed, results };
|
||||
}
|
||||
|
||||
get count(): number {
|
||||
return this.guardrails.length;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Built-in guardrails
|
||||
*/
|
||||
export const builtinGuardrails = {
|
||||
/**
|
||||
* Check for maximum length
|
||||
*/
|
||||
maxLength: (maxChars: number): Guardrail => {
|
||||
return guardrail({
|
||||
name: 'max_length',
|
||||
description: `Ensure content is under ${maxChars} characters`,
|
||||
check: (content) => {
|
||||
if (content.length > maxChars) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: `Content exceeds maximum length of ${maxChars} characters`,
|
||||
details: { length: content.length, max: maxChars },
|
||||
};
|
||||
}
|
||||
return { status: 'passed' };
|
||||
},
|
||||
});
|
||||
},
|
||||
|
||||
/**
|
||||
* Check for minimum length
|
||||
*/
|
||||
minLength: (minChars: number): Guardrail => {
|
||||
return guardrail({
|
||||
name: 'min_length',
|
||||
description: `Ensure content is at least ${minChars} characters`,
|
||||
check: (content) => {
|
||||
if (content.length < minChars) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: `Content is below minimum length of ${minChars} characters`,
|
||||
details: { length: content.length, min: minChars },
|
||||
};
|
||||
}
|
||||
return { status: 'passed' };
|
||||
},
|
||||
});
|
||||
},
|
||||
|
||||
/**
|
||||
* Check for blocked words
|
||||
*/
|
||||
blockedWords: (words: string[]): Guardrail => {
|
||||
return guardrail({
|
||||
name: 'blocked_words',
|
||||
description: 'Check for blocked words',
|
||||
check: (content) => {
|
||||
const lowerContent = content.toLowerCase();
|
||||
const found = words.filter(w => lowerContent.includes(w.toLowerCase()));
|
||||
if (found.length > 0) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: `Content contains blocked words`,
|
||||
details: { blockedWords: found },
|
||||
};
|
||||
}
|
||||
return { status: 'passed' };
|
||||
},
|
||||
});
|
||||
},
|
||||
|
||||
/**
|
||||
* Check for required words
|
||||
*/
|
||||
requiredWords: (words: string[]): Guardrail => {
|
||||
return guardrail({
|
||||
name: 'required_words',
|
||||
description: 'Check for required words',
|
||||
check: (content) => {
|
||||
const lowerContent = content.toLowerCase();
|
||||
const missing = words.filter(w => !lowerContent.includes(w.toLowerCase()));
|
||||
if (missing.length > 0) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: `Content missing required words`,
|
||||
details: { missingWords: missing },
|
||||
};
|
||||
}
|
||||
return { status: 'passed' };
|
||||
},
|
||||
});
|
||||
},
|
||||
|
||||
/**
|
||||
* Regex pattern check
|
||||
*/
|
||||
pattern: (regex: RegExp, mustMatch: boolean = true): Guardrail => {
|
||||
return guardrail({
|
||||
name: 'pattern',
|
||||
description: `Check content against pattern: ${regex}`,
|
||||
check: (content) => {
|
||||
const matches = regex.test(content);
|
||||
if (mustMatch && !matches) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: `Content does not match required pattern`,
|
||||
details: { pattern: regex.toString() },
|
||||
};
|
||||
}
|
||||
if (!mustMatch && matches) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: `Content matches forbidden pattern`,
|
||||
details: { pattern: regex.toString() },
|
||||
};
|
||||
}
|
||||
return { status: 'passed' };
|
||||
},
|
||||
});
|
||||
},
|
||||
|
||||
/**
|
||||
* JSON validity check
|
||||
*/
|
||||
validJson: (): Guardrail => {
|
||||
return guardrail({
|
||||
name: 'valid_json',
|
||||
description: 'Ensure content is valid JSON',
|
||||
check: (content) => {
|
||||
try {
|
||||
JSON.parse(content);
|
||||
return { status: 'passed' };
|
||||
} catch (e: any) {
|
||||
return {
|
||||
status: 'failed',
|
||||
message: 'Content is not valid JSON',
|
||||
details: { error: e.message },
|
||||
};
|
||||
}
|
||||
},
|
||||
});
|
||||
},
|
||||
};
|
||||
@@ -4,4 +4,74 @@ export * from './knowledge';
|
||||
export * from './llm';
|
||||
export * from './memory';
|
||||
export * from './process';
|
||||
export * from './tools';
|
||||
|
||||
// Export tools (excluding conflicting types)
|
||||
export {
|
||||
Tool, BaseTool,
|
||||
FunctionTool, tool, ToolRegistry, getRegistry, registerTool, getTool,
|
||||
type ToolConfig, type ToolContext, type ToolParameters
|
||||
} from './tools';
|
||||
export * from './tools/arxivTools';
|
||||
export * from './tools/mcpSse';
|
||||
|
||||
// Export session management
|
||||
export * from './session';
|
||||
|
||||
// Export database adapters
|
||||
export * from './db';
|
||||
|
||||
// Export workflows
|
||||
export * from './workflows';
|
||||
|
||||
// Export guardrails
|
||||
export * from './guardrails';
|
||||
|
||||
// Export handoff
|
||||
export { Handoff, handoff, handoffFilters, type HandoffConfig, type HandoffContext, type HandoffResult } from './agent/handoff';
|
||||
|
||||
// Export router agent
|
||||
export { RouterAgent, createRouter, routeConditions, type RouterConfig, type RouteConfig, type RouteContext } from './agent/router';
|
||||
|
||||
// Export context agent
|
||||
export { ContextAgent, createContextAgent, type ContextAgentConfig, type ContextMessage } from './agent/context';
|
||||
|
||||
// Export knowledge base (RAG)
|
||||
export { KnowledgeBase, createKnowledgeBase, type Document, type SearchResult, type EmbeddingProvider, type KnowledgeBaseConfig } from './knowledge/rag';
|
||||
|
||||
// Export evaluation framework
|
||||
export { accuracyEval, performanceEval, reliabilityEval, EvalSuite, type EvalResult, type PerformanceResult, type AccuracyEvalConfig, type PerformanceEvalConfig, type ReliabilityEvalConfig } from './eval';
|
||||
|
||||
// Export observability
|
||||
export { MemoryObservabilityAdapter, ConsoleObservabilityAdapter, setObservabilityAdapter, getObservabilityAdapter, type ObservabilityAdapter, type TraceContext, type SpanContext, type SpanData, type TraceData } from './observability';
|
||||
|
||||
// Export skills
|
||||
export { SkillManager, createSkillManager, parseSkillFile, type Skill, type SkillMetadata, type SkillDiscoveryOptions } from './skills';
|
||||
|
||||
// Export CLI
|
||||
export { chat, listProviders, version, help } from './cli';
|
||||
|
||||
// Export providers with explicit names to avoid conflicts
|
||||
export {
|
||||
createProvider,
|
||||
getDefaultProvider,
|
||||
parseModelString,
|
||||
isProviderAvailable,
|
||||
getAvailableProviders,
|
||||
OpenAIProvider,
|
||||
AnthropicProvider,
|
||||
GoogleProvider,
|
||||
BaseProvider,
|
||||
type LLMProvider,
|
||||
type ProviderConfig,
|
||||
type ProviderFactory,
|
||||
type GenerateTextOptions,
|
||||
type GenerateTextResult,
|
||||
type StreamTextOptions,
|
||||
type StreamChunk,
|
||||
type GenerateObjectOptions,
|
||||
type GenerateObjectResult,
|
||||
type TokenUsage,
|
||||
type Message as ProviderMessage,
|
||||
type ToolCall,
|
||||
type ToolDefinition as ProviderToolDefinition,
|
||||
} from './llm/providers';
|
||||
@@ -0,0 +1,193 @@
|
||||
/**
|
||||
* Knowledge Base (RAG) - Retrieval Augmented Generation
|
||||
*/
|
||||
|
||||
export interface Document {
|
||||
id: string;
|
||||
content: string;
|
||||
metadata?: Record<string, any>;
|
||||
embedding?: number[];
|
||||
}
|
||||
|
||||
export interface SearchResult {
|
||||
document: Document;
|
||||
score: number;
|
||||
}
|
||||
|
||||
export interface EmbeddingProvider {
|
||||
embed(text: string): Promise<number[]>;
|
||||
embedBatch(texts: string[]): Promise<number[][]>;
|
||||
}
|
||||
|
||||
export interface KnowledgeBaseConfig {
|
||||
embeddingProvider?: EmbeddingProvider;
|
||||
similarityThreshold?: number;
|
||||
maxResults?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Simple in-memory vector store for RAG
|
||||
*/
|
||||
export class KnowledgeBase {
|
||||
private documents: Map<string, Document> = new Map();
|
||||
private embeddingProvider?: EmbeddingProvider;
|
||||
private similarityThreshold: number;
|
||||
private maxResults: number;
|
||||
|
||||
constructor(config: KnowledgeBaseConfig = {}) {
|
||||
this.embeddingProvider = config.embeddingProvider;
|
||||
this.similarityThreshold = config.similarityThreshold ?? 0.7;
|
||||
this.maxResults = config.maxResults ?? 5;
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a document to the knowledge base
|
||||
*/
|
||||
async add(doc: Omit<Document, 'embedding'>): Promise<Document> {
|
||||
const document: Document = { ...doc };
|
||||
|
||||
if (this.embeddingProvider) {
|
||||
document.embedding = await this.embeddingProvider.embed(doc.content);
|
||||
}
|
||||
|
||||
this.documents.set(doc.id, document);
|
||||
return document;
|
||||
}
|
||||
|
||||
/**
|
||||
* Add multiple documents
|
||||
*/
|
||||
async addBatch(docs: Array<Omit<Document, 'embedding'>>): Promise<Document[]> {
|
||||
if (this.embeddingProvider && docs.length > 0) {
|
||||
const embeddings = await this.embeddingProvider.embedBatch(docs.map(d => d.content));
|
||||
return Promise.all(docs.map(async (doc, i) => {
|
||||
const document: Document = { ...doc, embedding: embeddings[i] };
|
||||
this.documents.set(doc.id, document);
|
||||
return document;
|
||||
}));
|
||||
}
|
||||
|
||||
return Promise.all(docs.map(doc => this.add(doc)));
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a document by ID
|
||||
*/
|
||||
get(id: string): Document | undefined {
|
||||
return this.documents.get(id);
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete a document
|
||||
*/
|
||||
delete(id: string): boolean {
|
||||
return this.documents.delete(id);
|
||||
}
|
||||
|
||||
/**
|
||||
* Search for similar documents
|
||||
*/
|
||||
async search(query: string, limit?: number): Promise<SearchResult[]> {
|
||||
const maxResults = limit ?? this.maxResults;
|
||||
|
||||
if (!this.embeddingProvider) {
|
||||
// Fallback to simple text matching
|
||||
return this.textSearch(query, maxResults);
|
||||
}
|
||||
|
||||
const queryEmbedding = await this.embeddingProvider.embed(query);
|
||||
const results: SearchResult[] = [];
|
||||
|
||||
for (const doc of this.documents.values()) {
|
||||
if (doc.embedding) {
|
||||
const score = this.cosineSimilarity(queryEmbedding, doc.embedding);
|
||||
if (score >= this.similarityThreshold) {
|
||||
results.push({ document: doc, score });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return results
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, maxResults);
|
||||
}
|
||||
|
||||
/**
|
||||
* Simple text-based search fallback
|
||||
*/
|
||||
private textSearch(query: string, limit: number): SearchResult[] {
|
||||
const queryLower = query.toLowerCase();
|
||||
const results: SearchResult[] = [];
|
||||
|
||||
for (const doc of this.documents.values()) {
|
||||
const contentLower = doc.content.toLowerCase();
|
||||
if (contentLower.includes(queryLower)) {
|
||||
const score = queryLower.length / contentLower.length;
|
||||
results.push({ document: doc, score: Math.min(score * 10, 1) });
|
||||
}
|
||||
}
|
||||
|
||||
return results
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, limit);
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculate cosine similarity between two vectors
|
||||
*/
|
||||
private cosineSimilarity(a: number[], b: number[]): number {
|
||||
if (a.length !== b.length) return 0;
|
||||
|
||||
let dotProduct = 0;
|
||||
let normA = 0;
|
||||
let normB = 0;
|
||||
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
dotProduct += a[i] * b[i];
|
||||
normA += a[i] * a[i];
|
||||
normB += b[i] * b[i];
|
||||
}
|
||||
|
||||
const denominator = Math.sqrt(normA) * Math.sqrt(normB);
|
||||
return denominator === 0 ? 0 : dotProduct / denominator;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all documents
|
||||
*/
|
||||
list(): Document[] {
|
||||
return Array.from(this.documents.values());
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear all documents
|
||||
*/
|
||||
clear(): void {
|
||||
this.documents.clear();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get document count
|
||||
*/
|
||||
get size(): number {
|
||||
return this.documents.size;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build context from search results for RAG
|
||||
*/
|
||||
buildContext(results: SearchResult[]): string {
|
||||
if (results.length === 0) return '';
|
||||
|
||||
return results
|
||||
.map((r, i) => `[${i + 1}] ${r.document.content}`)
|
||||
.join('\n\n');
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a knowledge base
|
||||
*/
|
||||
export function createKnowledgeBase(config?: KnowledgeBaseConfig): KnowledgeBase {
|
||||
return new KnowledgeBase(config);
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
/**
|
||||
* Anthropic Provider - Implementation for Anthropic Claude API
|
||||
*/
|
||||
|
||||
import { BaseProvider } from './base';
|
||||
import type {
|
||||
ProviderConfig,
|
||||
GenerateTextOptions,
|
||||
GenerateTextResult,
|
||||
StreamTextOptions,
|
||||
StreamChunk,
|
||||
GenerateObjectOptions,
|
||||
GenerateObjectResult,
|
||||
Message,
|
||||
ToolDefinition,
|
||||
ToolCall,
|
||||
} from './types';
|
||||
|
||||
interface AnthropicMessage {
|
||||
role: 'user' | 'assistant';
|
||||
content: string | Array<{ type: string; text?: string; tool_use_id?: string; content?: string }>;
|
||||
}
|
||||
|
||||
interface AnthropicTool {
|
||||
name: string;
|
||||
description: string;
|
||||
input_schema: Record<string, any>;
|
||||
}
|
||||
|
||||
export class AnthropicProvider extends BaseProvider {
|
||||
readonly providerId = 'anthropic';
|
||||
private apiKey: string;
|
||||
private baseUrl: string;
|
||||
|
||||
constructor(modelId: string, config: ProviderConfig = {}) {
|
||||
super(modelId, config);
|
||||
this.apiKey = config.apiKey || process.env.ANTHROPIC_API_KEY || '';
|
||||
this.baseUrl = config.baseUrl || 'https://api.anthropic.com';
|
||||
|
||||
if (!this.apiKey) {
|
||||
throw new Error('ANTHROPIC_API_KEY is required');
|
||||
}
|
||||
}
|
||||
|
||||
async generateText(options: GenerateTextOptions): Promise<GenerateTextResult> {
|
||||
return this.withRetry(async () => {
|
||||
const { systemPrompt, messages } = this.extractSystemPrompt(options.messages);
|
||||
|
||||
const response = await fetch(`${this.baseUrl}/v1/messages`, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'x-api-key': this.apiKey,
|
||||
'anthropic-version': '2023-06-01',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: this.modelId,
|
||||
max_tokens: options.maxTokens || 4096,
|
||||
system: systemPrompt,
|
||||
messages: this.formatMessages(messages),
|
||||
temperature: options.temperature ?? 0.7,
|
||||
tools: options.tools ? this.formatTools(options.tools) : undefined,
|
||||
stop_sequences: options.stop,
|
||||
top_p: options.topP,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
throw Object.assign(new Error(error.error?.message || `Anthropic API error: ${response.status}`), {
|
||||
status: response.status,
|
||||
});
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
let text = '';
|
||||
const toolCalls: ToolCall[] = [];
|
||||
|
||||
for (const block of data.content || []) {
|
||||
if (block.type === 'text') {
|
||||
text += block.text;
|
||||
} else if (block.type === 'tool_use') {
|
||||
toolCalls.push({
|
||||
id: block.id,
|
||||
type: 'function',
|
||||
function: {
|
||||
name: block.name,
|
||||
arguments: JSON.stringify(block.input),
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
text,
|
||||
toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
|
||||
usage: {
|
||||
promptTokens: data.usage?.input_tokens || 0,
|
||||
completionTokens: data.usage?.output_tokens || 0,
|
||||
totalTokens: (data.usage?.input_tokens || 0) + (data.usage?.output_tokens || 0),
|
||||
},
|
||||
finishReason: this.mapStopReason(data.stop_reason),
|
||||
raw: data,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
async streamText(options: StreamTextOptions): Promise<AsyncIterable<StreamChunk>> {
|
||||
const self = this;
|
||||
const { systemPrompt, messages } = this.extractSystemPrompt(options.messages);
|
||||
|
||||
return {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
const response = await fetch(`${self.baseUrl}/v1/messages`, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'x-api-key': self.apiKey,
|
||||
'anthropic-version': '2023-06-01',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: self.modelId,
|
||||
max_tokens: options.maxTokens || 4096,
|
||||
system: systemPrompt,
|
||||
messages: self.formatMessages(messages),
|
||||
temperature: options.temperature ?? 0.7,
|
||||
tools: options.tools ? self.formatTools(options.tools) : undefined,
|
||||
stream: true,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
throw new Error(error.error?.message || `Anthropic API error: ${response.status}`);
|
||||
}
|
||||
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) throw new Error('No response body');
|
||||
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
const toolCalls: ToolCall[] = [];
|
||||
let currentToolCall: Partial<ToolCall> | null = null;
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop() || '';
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith('data: ')) continue;
|
||||
const data = line.slice(6);
|
||||
if (data === '[DONE]') continue;
|
||||
|
||||
try {
|
||||
const event = JSON.parse(data);
|
||||
|
||||
if (event.type === 'content_block_delta') {
|
||||
if (event.delta?.type === 'text_delta') {
|
||||
const text = event.delta.text;
|
||||
if (options.onToken) options.onToken(text);
|
||||
yield { text };
|
||||
} else if (event.delta?.type === 'input_json_delta') {
|
||||
if (currentToolCall) {
|
||||
currentToolCall.function!.arguments += event.delta.partial_json;
|
||||
}
|
||||
}
|
||||
} else if (event.type === 'content_block_start') {
|
||||
if (event.content_block?.type === 'tool_use') {
|
||||
currentToolCall = {
|
||||
id: event.content_block.id,
|
||||
type: 'function',
|
||||
function: {
|
||||
name: event.content_block.name,
|
||||
arguments: '',
|
||||
},
|
||||
};
|
||||
}
|
||||
} else if (event.type === 'content_block_stop') {
|
||||
if (currentToolCall) {
|
||||
toolCalls.push(currentToolCall as ToolCall);
|
||||
currentToolCall = null;
|
||||
}
|
||||
} else if (event.type === 'message_delta') {
|
||||
yield {
|
||||
finishReason: self.mapStopReason(event.delta?.stop_reason),
|
||||
toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
|
||||
usage: event.usage ? {
|
||||
promptTokens: 0,
|
||||
completionTokens: event.usage.output_tokens,
|
||||
totalTokens: event.usage.output_tokens,
|
||||
} : undefined,
|
||||
};
|
||||
}
|
||||
} catch (e) {
|
||||
// Skip malformed JSON
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async generateObject<T = any>(options: GenerateObjectOptions<T>): Promise<GenerateObjectResult<T>> {
|
||||
const { systemPrompt, messages } = this.extractSystemPrompt(options.messages);
|
||||
|
||||
// Add JSON instruction to system prompt
|
||||
const jsonSystemPrompt = `${systemPrompt}\n\nYou must respond with valid JSON matching this schema:\n${JSON.stringify(options.schema, null, 2)}`;
|
||||
|
||||
return this.withRetry(async () => {
|
||||
const response = await fetch(`${this.baseUrl}/v1/messages`, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'x-api-key': this.apiKey,
|
||||
'anthropic-version': '2023-06-01',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: this.modelId,
|
||||
max_tokens: options.maxTokens || 4096,
|
||||
system: jsonSystemPrompt,
|
||||
messages: this.formatMessages(messages),
|
||||
temperature: options.temperature ?? 0.7,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
throw new Error(error.error?.message || `Anthropic API error: ${response.status}`);
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
let text = '';
|
||||
|
||||
for (const block of data.content || []) {
|
||||
if (block.type === 'text') {
|
||||
text += block.text;
|
||||
}
|
||||
}
|
||||
|
||||
// Extract JSON from response
|
||||
const jsonMatch = text.match(/\{[\s\S]*\}/);
|
||||
if (!jsonMatch) {
|
||||
throw new Error(`No JSON found in response: ${text}`);
|
||||
}
|
||||
|
||||
let parsed: T;
|
||||
try {
|
||||
parsed = JSON.parse(jsonMatch[0]);
|
||||
} catch (e) {
|
||||
throw new Error(`Failed to parse JSON: ${jsonMatch[0]}`);
|
||||
}
|
||||
|
||||
return {
|
||||
object: parsed,
|
||||
usage: {
|
||||
promptTokens: data.usage?.input_tokens || 0,
|
||||
completionTokens: data.usage?.output_tokens || 0,
|
||||
totalTokens: (data.usage?.input_tokens || 0) + (data.usage?.output_tokens || 0),
|
||||
},
|
||||
raw: data,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
private extractSystemPrompt(messages: Message[]): { systemPrompt: string; messages: Message[] } {
|
||||
const systemMessages = messages.filter(m => m.role === 'system');
|
||||
const otherMessages = messages.filter(m => m.role !== 'system');
|
||||
const systemPrompt = systemMessages.map(m => m.content).join('\n');
|
||||
return { systemPrompt, messages: otherMessages };
|
||||
}
|
||||
|
||||
protected formatMessages(messages: Message[]): AnthropicMessage[] {
|
||||
const result: AnthropicMessage[] = [];
|
||||
|
||||
for (const msg of messages) {
|
||||
if (msg.role === 'system') continue; // Handled separately
|
||||
|
||||
if (msg.role === 'tool') {
|
||||
// Tool results need to be part of user message in Anthropic
|
||||
result.push({
|
||||
role: 'user',
|
||||
content: [{
|
||||
type: 'tool_result',
|
||||
tool_use_id: msg.tool_call_id || '',
|
||||
content: msg.content || '',
|
||||
}],
|
||||
});
|
||||
} else if (msg.role === 'assistant' && msg.tool_calls) {
|
||||
const content: any[] = [];
|
||||
if (msg.content) {
|
||||
content.push({ type: 'text', text: msg.content });
|
||||
}
|
||||
for (const tc of msg.tool_calls) {
|
||||
content.push({
|
||||
type: 'tool_use',
|
||||
id: tc.id,
|
||||
name: tc.function.name,
|
||||
input: JSON.parse(tc.function.arguments),
|
||||
});
|
||||
}
|
||||
result.push({ role: 'assistant', content });
|
||||
} else {
|
||||
result.push({
|
||||
role: msg.role as 'user' | 'assistant',
|
||||
content: msg.content || '',
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
protected formatTools(tools: ToolDefinition[]): AnthropicTool[] {
|
||||
return tools.map(tool => ({
|
||||
name: tool.name,
|
||||
description: tool.description || `Function ${tool.name}`,
|
||||
input_schema: tool.parameters || { type: 'object', properties: {} },
|
||||
}));
|
||||
}
|
||||
|
||||
private mapStopReason(reason: string | null): GenerateTextResult['finishReason'] {
|
||||
switch (reason) {
|
||||
case 'end_turn': return 'stop';
|
||||
case 'max_tokens': return 'length';
|
||||
case 'tool_use': return 'tool_calls';
|
||||
default: return 'stop';
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
/**
|
||||
* Base Provider - Abstract base class for LLM providers
|
||||
*/
|
||||
|
||||
import type {
|
||||
LLMProvider,
|
||||
ProviderConfig,
|
||||
GenerateTextOptions,
|
||||
GenerateTextResult,
|
||||
StreamTextOptions,
|
||||
StreamChunk,
|
||||
GenerateObjectOptions,
|
||||
GenerateObjectResult,
|
||||
Message,
|
||||
ToolDefinition,
|
||||
} from './types';
|
||||
|
||||
export abstract class BaseProvider implements LLMProvider {
|
||||
abstract readonly providerId: string;
|
||||
readonly modelId: string;
|
||||
protected config: ProviderConfig;
|
||||
|
||||
constructor(modelId: string, config: ProviderConfig = {}) {
|
||||
this.modelId = modelId;
|
||||
this.config = {
|
||||
maxRetries: 3,
|
||||
timeout: 60000,
|
||||
...config,
|
||||
};
|
||||
}
|
||||
|
||||
abstract generateText(options: GenerateTextOptions): Promise<GenerateTextResult>;
|
||||
abstract streamText(options: StreamTextOptions): Promise<AsyncIterable<StreamChunk>>;
|
||||
abstract generateObject<T = any>(options: GenerateObjectOptions<T>): Promise<GenerateObjectResult<T>>;
|
||||
|
||||
/**
|
||||
* Convert tool definitions to provider-specific format
|
||||
*/
|
||||
protected abstract formatTools(tools: ToolDefinition[]): any[];
|
||||
|
||||
/**
|
||||
* Convert messages to provider-specific format
|
||||
*/
|
||||
protected abstract formatMessages(messages: Message[]): any[];
|
||||
|
||||
/**
|
||||
* Retry logic with exponential backoff
|
||||
*/
|
||||
protected async withRetry<T>(
|
||||
fn: () => Promise<T>,
|
||||
maxRetries: number = this.config.maxRetries || 3
|
||||
): Promise<T> {
|
||||
let lastError: Error | undefined;
|
||||
for (let attempt = 0; attempt <= maxRetries; attempt++) {
|
||||
try {
|
||||
return await fn();
|
||||
} catch (error: any) {
|
||||
lastError = error;
|
||||
// Check if error is retryable (rate limit, server error)
|
||||
const isRetryable =
|
||||
error.status === 429 ||
|
||||
error.status >= 500 ||
|
||||
error.code === 'ECONNRESET' ||
|
||||
error.code === 'ETIMEDOUT';
|
||||
|
||||
if (!isRetryable || attempt === maxRetries) {
|
||||
throw error;
|
||||
}
|
||||
|
||||
// Exponential backoff
|
||||
const delay = Math.min(1000 * Math.pow(2, attempt), 30000);
|
||||
await new Promise(resolve => setTimeout(resolve, delay));
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
/**
|
||||
* Google Provider - Implementation for Google Gemini API
|
||||
*/
|
||||
|
||||
import { BaseProvider } from './base';
|
||||
import type {
|
||||
ProviderConfig,
|
||||
GenerateTextOptions,
|
||||
GenerateTextResult,
|
||||
StreamTextOptions,
|
||||
StreamChunk,
|
||||
GenerateObjectOptions,
|
||||
GenerateObjectResult,
|
||||
Message,
|
||||
ToolDefinition,
|
||||
ToolCall,
|
||||
} from './types';
|
||||
|
||||
interface GeminiContent {
|
||||
role: 'user' | 'model';
|
||||
parts: Array<{ text?: string; functionCall?: any; functionResponse?: any }>;
|
||||
}
|
||||
|
||||
export class GoogleProvider extends BaseProvider {
|
||||
readonly providerId = 'google';
|
||||
private apiKey: string;
|
||||
private baseUrl: string;
|
||||
|
||||
constructor(modelId: string, config: ProviderConfig = {}) {
|
||||
super(modelId, config);
|
||||
this.apiKey = config.apiKey || process.env.GOOGLE_API_KEY || '';
|
||||
this.baseUrl = config.baseUrl || 'https://generativelanguage.googleapis.com/v1beta';
|
||||
|
||||
if (!this.apiKey) {
|
||||
throw new Error('GOOGLE_API_KEY is required');
|
||||
}
|
||||
}
|
||||
|
||||
async generateText(options: GenerateTextOptions): Promise<GenerateTextResult> {
|
||||
return this.withRetry(async () => {
|
||||
const { systemInstruction, contents } = this.formatRequest(options.messages);
|
||||
|
||||
const response = await fetch(
|
||||
`${this.baseUrl}/models/${this.modelId}:generateContent?key=${this.apiKey}`,
|
||||
{
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
systemInstruction: systemInstruction ? { parts: [{ text: systemInstruction }] } : undefined,
|
||||
contents,
|
||||
generationConfig: {
|
||||
temperature: options.temperature ?? 0.7,
|
||||
maxOutputTokens: options.maxTokens,
|
||||
topP: options.topP,
|
||||
stopSequences: options.stop,
|
||||
},
|
||||
tools: options.tools ? this.formatTools(options.tools) : undefined,
|
||||
}),
|
||||
}
|
||||
);
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
throw Object.assign(new Error(error.error?.message || `Google API error: ${response.status}`), {
|
||||
status: response.status,
|
||||
});
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
const candidate = data.candidates?.[0];
|
||||
|
||||
let text = '';
|
||||
const toolCalls: ToolCall[] = [];
|
||||
|
||||
for (const part of candidate?.content?.parts || []) {
|
||||
if (part.text) {
|
||||
text += part.text;
|
||||
} else if (part.functionCall) {
|
||||
toolCalls.push({
|
||||
id: `call_${Math.random().toString(36).substr(2, 9)}`,
|
||||
type: 'function',
|
||||
function: {
|
||||
name: part.functionCall.name,
|
||||
arguments: JSON.stringify(part.functionCall.args || {}),
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
text,
|
||||
toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
|
||||
usage: {
|
||||
promptTokens: data.usageMetadata?.promptTokenCount || 0,
|
||||
completionTokens: data.usageMetadata?.candidatesTokenCount || 0,
|
||||
totalTokens: data.usageMetadata?.totalTokenCount || 0,
|
||||
},
|
||||
finishReason: this.mapFinishReason(candidate?.finishReason),
|
||||
raw: data,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
async streamText(options: StreamTextOptions): Promise<AsyncIterable<StreamChunk>> {
|
||||
const self = this;
|
||||
const { systemInstruction, contents } = this.formatRequest(options.messages);
|
||||
|
||||
return {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
const response = await fetch(
|
||||
`${self.baseUrl}/models/${self.modelId}:streamGenerateContent?key=${self.apiKey}&alt=sse`,
|
||||
{
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
systemInstruction: systemInstruction ? { parts: [{ text: systemInstruction }] } : undefined,
|
||||
contents,
|
||||
generationConfig: {
|
||||
temperature: options.temperature ?? 0.7,
|
||||
maxOutputTokens: options.maxTokens,
|
||||
},
|
||||
tools: options.tools ? self.formatTools(options.tools) : undefined,
|
||||
}),
|
||||
}
|
||||
);
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
throw new Error(error.error?.message || `Google API error: ${response.status}`);
|
||||
}
|
||||
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) throw new Error('No response body');
|
||||
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
const toolCalls: ToolCall[] = [];
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop() || '';
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith('data: ')) continue;
|
||||
const data = line.slice(6);
|
||||
if (data === '[DONE]') continue;
|
||||
|
||||
try {
|
||||
const event = JSON.parse(data);
|
||||
const candidate = event.candidates?.[0];
|
||||
|
||||
for (const part of candidate?.content?.parts || []) {
|
||||
if (part.text) {
|
||||
if (options.onToken) options.onToken(part.text);
|
||||
yield { text: part.text };
|
||||
} else if (part.functionCall) {
|
||||
toolCalls.push({
|
||||
id: `call_${Math.random().toString(36).substr(2, 9)}`,
|
||||
type: 'function',
|
||||
function: {
|
||||
name: part.functionCall.name,
|
||||
arguments: JSON.stringify(part.functionCall.args || {}),
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (candidate?.finishReason) {
|
||||
yield {
|
||||
finishReason: self.mapFinishReason(candidate.finishReason),
|
||||
toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
|
||||
usage: event.usageMetadata ? {
|
||||
promptTokens: event.usageMetadata.promptTokenCount || 0,
|
||||
completionTokens: event.usageMetadata.candidatesTokenCount || 0,
|
||||
totalTokens: event.usageMetadata.totalTokenCount || 0,
|
||||
} : undefined,
|
||||
};
|
||||
}
|
||||
} catch (e) {
|
||||
// Skip malformed JSON
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async generateObject<T = any>(options: GenerateObjectOptions<T>): Promise<GenerateObjectResult<T>> {
|
||||
const { systemInstruction, contents } = this.formatRequest(options.messages);
|
||||
|
||||
return this.withRetry(async () => {
|
||||
const response = await fetch(
|
||||
`${this.baseUrl}/models/${this.modelId}:generateContent?key=${this.apiKey}`,
|
||||
{
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
systemInstruction: systemInstruction
|
||||
? { parts: [{ text: `${systemInstruction}\n\nRespond with valid JSON matching this schema:\n${JSON.stringify(options.schema, null, 2)}` }] }
|
||||
: { parts: [{ text: `Respond with valid JSON matching this schema:\n${JSON.stringify(options.schema, null, 2)}` }] },
|
||||
contents,
|
||||
generationConfig: {
|
||||
temperature: options.temperature ?? 0.7,
|
||||
maxOutputTokens: options.maxTokens,
|
||||
responseMimeType: 'application/json',
|
||||
},
|
||||
}),
|
||||
}
|
||||
);
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
throw new Error(error.error?.message || `Google API error: ${response.status}`);
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
const candidate = data.candidates?.[0];
|
||||
let text = '';
|
||||
|
||||
for (const part of candidate?.content?.parts || []) {
|
||||
if (part.text) text += part.text;
|
||||
}
|
||||
|
||||
let parsed: T;
|
||||
try {
|
||||
parsed = JSON.parse(text);
|
||||
} catch (e) {
|
||||
// Try to extract JSON from text
|
||||
const jsonMatch = text.match(/\{[\s\S]*\}/);
|
||||
if (!jsonMatch) throw new Error(`No JSON found: ${text}`);
|
||||
parsed = JSON.parse(jsonMatch[0]);
|
||||
}
|
||||
|
||||
return {
|
||||
object: parsed,
|
||||
usage: {
|
||||
promptTokens: data.usageMetadata?.promptTokenCount || 0,
|
||||
completionTokens: data.usageMetadata?.candidatesTokenCount || 0,
|
||||
totalTokens: data.usageMetadata?.totalTokenCount || 0,
|
||||
},
|
||||
raw: data,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
private formatRequest(messages: Message[]): { systemInstruction: string | null; contents: GeminiContent[] } {
|
||||
const systemMessages = messages.filter(m => m.role === 'system');
|
||||
const otherMessages = messages.filter(m => m.role !== 'system');
|
||||
const systemInstruction = systemMessages.length > 0
|
||||
? systemMessages.map(m => m.content).join('\n')
|
||||
: null;
|
||||
|
||||
return {
|
||||
systemInstruction,
|
||||
contents: this.formatMessages(otherMessages),
|
||||
};
|
||||
}
|
||||
|
||||
protected formatMessages(messages: Message[]): GeminiContent[] {
|
||||
const result: GeminiContent[] = [];
|
||||
|
||||
for (const msg of messages) {
|
||||
if (msg.role === 'system') continue;
|
||||
|
||||
const role = msg.role === 'assistant' ? 'model' : 'user';
|
||||
|
||||
if (msg.role === 'tool') {
|
||||
result.push({
|
||||
role: 'user',
|
||||
parts: [{
|
||||
functionResponse: {
|
||||
name: msg.name || 'function',
|
||||
response: { result: msg.content },
|
||||
},
|
||||
}],
|
||||
});
|
||||
} else if (msg.role === 'assistant' && msg.tool_calls) {
|
||||
const parts: any[] = [];
|
||||
if (msg.content) parts.push({ text: msg.content });
|
||||
for (const tc of msg.tool_calls) {
|
||||
parts.push({
|
||||
functionCall: {
|
||||
name: tc.function.name,
|
||||
args: JSON.parse(tc.function.arguments),
|
||||
},
|
||||
});
|
||||
}
|
||||
result.push({ role: 'model', parts });
|
||||
} else {
|
||||
result.push({
|
||||
role,
|
||||
parts: [{ text: msg.content || '' }],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
protected formatTools(tools: ToolDefinition[]): any[] {
|
||||
return [{
|
||||
functionDeclarations: tools.map(tool => ({
|
||||
name: tool.name,
|
||||
description: tool.description || `Function ${tool.name}`,
|
||||
parameters: tool.parameters || { type: 'object', properties: {} },
|
||||
})),
|
||||
}];
|
||||
}
|
||||
|
||||
private mapFinishReason(reason: string | null): GenerateTextResult['finishReason'] {
|
||||
switch (reason) {
|
||||
case 'STOP': return 'stop';
|
||||
case 'MAX_TOKENS': return 'length';
|
||||
case 'SAFETY': return 'content_filter';
|
||||
default: return 'stop';
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
/**
|
||||
* LLM Providers - Factory and exports for multi-provider support
|
||||
*/
|
||||
|
||||
export * from './types';
|
||||
export * from './base';
|
||||
export { OpenAIProvider } from './openai';
|
||||
export { AnthropicProvider } from './anthropic';
|
||||
export { GoogleProvider } from './google';
|
||||
|
||||
import type { LLMProvider, ProviderConfig } from './types';
|
||||
import { OpenAIProvider } from './openai';
|
||||
import { AnthropicProvider } from './anthropic';
|
||||
import { GoogleProvider } from './google';
|
||||
|
||||
/**
|
||||
* Provider registry for dynamic provider loading
|
||||
*/
|
||||
const PROVIDER_MAP: Record<string, new (modelId: string, config?: ProviderConfig) => LLMProvider> = {
|
||||
openai: OpenAIProvider,
|
||||
anthropic: AnthropicProvider,
|
||||
google: GoogleProvider,
|
||||
gemini: GoogleProvider, // Alias
|
||||
};
|
||||
|
||||
/**
|
||||
* Parse model string into provider and model ID
|
||||
* Supports formats:
|
||||
* - "provider/model" (e.g., "openai/gpt-4o")
|
||||
* - "model" (defaults to OpenAI, e.g., "gpt-4o-mini")
|
||||
*/
|
||||
export function parseModelString(model: string): { providerId: string; modelId: string } {
|
||||
if (model.includes('/')) {
|
||||
const [providerId, ...rest] = model.split('/');
|
||||
return { providerId: providerId.toLowerCase(), modelId: rest.join('/') };
|
||||
}
|
||||
|
||||
// Default to OpenAI for common model prefixes
|
||||
if (model.startsWith('gpt-') || model.startsWith('o1') || model.startsWith('o3')) {
|
||||
return { providerId: 'openai', modelId: model };
|
||||
}
|
||||
if (model.startsWith('claude-')) {
|
||||
return { providerId: 'anthropic', modelId: model };
|
||||
}
|
||||
if (model.startsWith('gemini-')) {
|
||||
return { providerId: 'google', modelId: model };
|
||||
}
|
||||
|
||||
// Default to OpenAI
|
||||
return { providerId: 'openai', modelId: model };
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a provider instance from a model string
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* const provider = createProvider('openai/gpt-4o');
|
||||
* const provider = createProvider('anthropic/claude-3-5-sonnet-latest');
|
||||
* const provider = createProvider('google/gemini-2.0-flash');
|
||||
* const provider = createProvider('gpt-4o-mini'); // Defaults to OpenAI
|
||||
* ```
|
||||
*/
|
||||
export function createProvider(model: string, config?: ProviderConfig): LLMProvider {
|
||||
const { providerId, modelId } = parseModelString(model);
|
||||
|
||||
const ProviderClass = PROVIDER_MAP[providerId];
|
||||
if (!ProviderClass) {
|
||||
throw new Error(`Unknown provider: ${providerId}. Available providers: ${Object.keys(PROVIDER_MAP).join(', ')}`);
|
||||
}
|
||||
|
||||
return new ProviderClass(modelId, config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the default provider (OpenAI with gpt-4o-mini)
|
||||
*/
|
||||
export function getDefaultProvider(config?: ProviderConfig): LLMProvider {
|
||||
return createProvider('openai/gpt-4o-mini', config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a provider is available (has required API key)
|
||||
*/
|
||||
export function isProviderAvailable(providerId: string): boolean {
|
||||
switch (providerId.toLowerCase()) {
|
||||
case 'openai':
|
||||
return !!process.env.OPENAI_API_KEY;
|
||||
case 'anthropic':
|
||||
return !!process.env.ANTHROPIC_API_KEY;
|
||||
case 'google':
|
||||
case 'gemini':
|
||||
return !!process.env.GOOGLE_API_KEY;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get list of available providers
|
||||
*/
|
||||
export function getAvailableProviders(): string[] {
|
||||
return Object.keys(PROVIDER_MAP).filter(isProviderAvailable);
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
/**
|
||||
* OpenAI Provider - Implementation for OpenAI API
|
||||
*/
|
||||
|
||||
import OpenAI from 'openai';
|
||||
import type { ChatCompletionMessageParam, ChatCompletionTool } from 'openai/resources/chat/completions';
|
||||
import { BaseProvider } from './base';
|
||||
import type {
|
||||
ProviderConfig,
|
||||
GenerateTextOptions,
|
||||
GenerateTextResult,
|
||||
StreamTextOptions,
|
||||
StreamChunk,
|
||||
GenerateObjectOptions,
|
||||
GenerateObjectResult,
|
||||
Message,
|
||||
ToolDefinition,
|
||||
ToolCall,
|
||||
} from './types';
|
||||
|
||||
export class OpenAIProvider extends BaseProvider {
|
||||
readonly providerId = 'openai';
|
||||
private client: OpenAI;
|
||||
|
||||
constructor(modelId: string, config: ProviderConfig = {}) {
|
||||
super(modelId, config);
|
||||
this.client = new OpenAI({
|
||||
apiKey: config.apiKey || process.env.OPENAI_API_KEY,
|
||||
baseURL: config.baseUrl,
|
||||
timeout: config.timeout || 60000,
|
||||
maxRetries: 0, // We handle retries ourselves
|
||||
});
|
||||
}
|
||||
|
||||
async generateText(options: GenerateTextOptions): Promise<GenerateTextResult> {
|
||||
return this.withRetry(async () => {
|
||||
const response = await this.client.chat.completions.create({
|
||||
model: this.modelId,
|
||||
messages: this.formatMessages(options.messages),
|
||||
temperature: options.temperature ?? 0.7,
|
||||
max_tokens: options.maxTokens,
|
||||
tools: options.tools ? this.formatTools(options.tools) : undefined,
|
||||
tool_choice: options.toolChoice as any,
|
||||
stop: options.stop,
|
||||
top_p: options.topP,
|
||||
frequency_penalty: options.frequencyPenalty,
|
||||
presence_penalty: options.presencePenalty,
|
||||
});
|
||||
|
||||
const choice = response.choices[0];
|
||||
const message = choice.message;
|
||||
|
||||
return {
|
||||
text: message.content || '',
|
||||
toolCalls: message.tool_calls?.map(tc => ({
|
||||
id: tc.id,
|
||||
type: 'function' as const,
|
||||
function: {
|
||||
name: tc.function.name,
|
||||
arguments: tc.function.arguments,
|
||||
},
|
||||
})),
|
||||
usage: {
|
||||
promptTokens: response.usage?.prompt_tokens || 0,
|
||||
completionTokens: response.usage?.completion_tokens || 0,
|
||||
totalTokens: response.usage?.total_tokens || 0,
|
||||
},
|
||||
finishReason: this.mapFinishReason(choice.finish_reason),
|
||||
raw: response,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
async streamText(options: StreamTextOptions): Promise<AsyncIterable<StreamChunk>> {
|
||||
const self = this;
|
||||
|
||||
return {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
const stream = await self.client.chat.completions.create({
|
||||
model: self.modelId,
|
||||
messages: self.formatMessages(options.messages),
|
||||
temperature: options.temperature ?? 0.7,
|
||||
max_tokens: options.maxTokens,
|
||||
tools: options.tools ? self.formatTools(options.tools) : undefined,
|
||||
tool_choice: options.toolChoice as any,
|
||||
stop: options.stop,
|
||||
stream: true,
|
||||
});
|
||||
|
||||
let toolCalls: ToolCall[] = [];
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const delta = chunk.choices[0]?.delta;
|
||||
|
||||
if (delta?.content) {
|
||||
if (options.onToken) {
|
||||
options.onToken(delta.content);
|
||||
}
|
||||
yield { text: delta.content };
|
||||
}
|
||||
|
||||
if (delta?.tool_calls) {
|
||||
for (const tc of delta.tool_calls) {
|
||||
if (tc.index !== undefined) {
|
||||
if (!toolCalls[tc.index]) {
|
||||
toolCalls[tc.index] = {
|
||||
id: tc.id || '',
|
||||
type: 'function',
|
||||
function: { name: '', arguments: '' },
|
||||
};
|
||||
}
|
||||
if (tc.id) toolCalls[tc.index].id = tc.id;
|
||||
if (tc.function?.name) toolCalls[tc.index].function.name = tc.function.name;
|
||||
if (tc.function?.arguments) toolCalls[tc.index].function.arguments += tc.function.arguments;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (chunk.choices[0]?.finish_reason) {
|
||||
yield {
|
||||
finishReason: chunk.choices[0].finish_reason,
|
||||
toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
|
||||
usage: chunk.usage ? {
|
||||
promptTokens: chunk.usage.prompt_tokens,
|
||||
completionTokens: chunk.usage.completion_tokens,
|
||||
totalTokens: chunk.usage.total_tokens,
|
||||
} : undefined,
|
||||
};
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async generateObject<T = any>(options: GenerateObjectOptions<T>): Promise<GenerateObjectResult<T>> {
|
||||
return this.withRetry(async () => {
|
||||
const response = await this.client.chat.completions.create({
|
||||
model: this.modelId,
|
||||
messages: this.formatMessages(options.messages),
|
||||
temperature: options.temperature ?? 0.7,
|
||||
max_tokens: options.maxTokens,
|
||||
response_format: {
|
||||
type: 'json_schema',
|
||||
json_schema: {
|
||||
name: 'response',
|
||||
schema: this.normalizeSchema(options.schema),
|
||||
strict: true,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const choice = response.choices[0];
|
||||
const content = choice.message.content || '{}';
|
||||
|
||||
let parsed: T;
|
||||
try {
|
||||
parsed = JSON.parse(content);
|
||||
} catch (e) {
|
||||
throw new Error(`Failed to parse JSON response: ${content}`);
|
||||
}
|
||||
|
||||
return {
|
||||
object: parsed,
|
||||
usage: {
|
||||
promptTokens: response.usage?.prompt_tokens || 0,
|
||||
completionTokens: response.usage?.completion_tokens || 0,
|
||||
totalTokens: response.usage?.total_tokens || 0,
|
||||
},
|
||||
raw: response,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
protected formatMessages(messages: Message[]): ChatCompletionMessageParam[] {
|
||||
return messages.map(msg => {
|
||||
if (msg.role === 'tool') {
|
||||
return {
|
||||
role: 'tool' as const,
|
||||
content: msg.content || '',
|
||||
tool_call_id: msg.tool_call_id || '',
|
||||
};
|
||||
}
|
||||
if (msg.role === 'assistant' && msg.tool_calls) {
|
||||
return {
|
||||
role: 'assistant' as const,
|
||||
content: msg.content,
|
||||
tool_calls: msg.tool_calls,
|
||||
};
|
||||
}
|
||||
return {
|
||||
role: msg.role as 'system' | 'user' | 'assistant',
|
||||
content: msg.content || '',
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
protected formatTools(tools: ToolDefinition[]): ChatCompletionTool[] {
|
||||
return tools.map(tool => ({
|
||||
type: 'function' as const,
|
||||
function: {
|
||||
name: tool.name,
|
||||
description: tool.description || `Function ${tool.name}`,
|
||||
parameters: tool.parameters || { type: 'object', properties: {} },
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
private mapFinishReason(reason: string | null): GenerateTextResult['finishReason'] {
|
||||
switch (reason) {
|
||||
case 'stop': return 'stop';
|
||||
case 'length': return 'length';
|
||||
case 'tool_calls': return 'tool_calls';
|
||||
case 'content_filter': return 'content_filter';
|
||||
default: return 'stop';
|
||||
}
|
||||
}
|
||||
|
||||
private normalizeSchema(schema: any): any {
|
||||
// If it's a Zod schema, convert to JSON schema
|
||||
if (schema && typeof schema.parse === 'function' && typeof schema._def === 'object') {
|
||||
// This is a Zod schema - we need zod-to-json-schema
|
||||
// For now, throw an error suggesting to use JSON schema directly
|
||||
throw new Error('Zod schemas require zod-to-json-schema. Please use JSON schema directly or install zod-to-json-schema.');
|
||||
}
|
||||
return schema;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
/**
|
||||
* LLM Provider Types - Core type definitions for the provider abstraction layer
|
||||
*/
|
||||
|
||||
export interface Message {
|
||||
role: 'system' | 'user' | 'assistant' | 'tool';
|
||||
content: string | null;
|
||||
name?: string;
|
||||
tool_call_id?: string;
|
||||
tool_calls?: ToolCall[];
|
||||
}
|
||||
|
||||
export interface ToolCall {
|
||||
id: string;
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
arguments: string;
|
||||
};
|
||||
}
|
||||
|
||||
export interface ToolDefinition {
|
||||
name: string;
|
||||
description?: string;
|
||||
parameters?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface TokenUsage {
|
||||
promptTokens: number;
|
||||
completionTokens: number;
|
||||
totalTokens: number;
|
||||
}
|
||||
|
||||
export interface GenerateTextOptions {
|
||||
messages: Message[];
|
||||
temperature?: number;
|
||||
maxTokens?: number;
|
||||
tools?: ToolDefinition[];
|
||||
toolChoice?: 'auto' | 'none' | 'required' | { type: 'function'; function: { name: string } };
|
||||
stop?: string[];
|
||||
topP?: number;
|
||||
frequencyPenalty?: number;
|
||||
presencePenalty?: number;
|
||||
}
|
||||
|
||||
export interface GenerateTextResult {
|
||||
text: string;
|
||||
toolCalls?: ToolCall[];
|
||||
usage: TokenUsage;
|
||||
finishReason: 'stop' | 'length' | 'tool_calls' | 'content_filter' | 'error';
|
||||
raw?: any;
|
||||
}
|
||||
|
||||
export interface StreamTextOptions extends GenerateTextOptions {
|
||||
onToken?: (token: string) => void;
|
||||
onToolCall?: (toolCall: ToolCall) => void;
|
||||
}
|
||||
|
||||
export interface StreamChunk {
|
||||
text?: string;
|
||||
toolCalls?: ToolCall[];
|
||||
usage?: Partial<TokenUsage>;
|
||||
finishReason?: string;
|
||||
}
|
||||
|
||||
export interface GenerateObjectOptions<T = any> {
|
||||
messages: Message[];
|
||||
schema: Record<string, any> | any; // JSON Schema or Zod schema
|
||||
temperature?: number;
|
||||
maxTokens?: number;
|
||||
maxRetries?: number;
|
||||
}
|
||||
|
||||
export interface GenerateObjectResult<T = any> {
|
||||
object: T;
|
||||
usage: TokenUsage;
|
||||
raw?: any;
|
||||
}
|
||||
|
||||
export interface ProviderConfig {
|
||||
apiKey?: string;
|
||||
baseUrl?: string;
|
||||
maxRetries?: number;
|
||||
timeout?: number;
|
||||
defaultModel?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Base interface for all LLM providers
|
||||
*/
|
||||
export interface LLMProvider {
|
||||
readonly providerId: string;
|
||||
readonly modelId: string;
|
||||
|
||||
generateText(options: GenerateTextOptions): Promise<GenerateTextResult>;
|
||||
streamText(options: StreamTextOptions): Promise<AsyncIterable<StreamChunk>>;
|
||||
generateObject<T = any>(options: GenerateObjectOptions<T>): Promise<GenerateObjectResult<T>>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Provider factory function type
|
||||
*/
|
||||
export type ProviderFactory = (modelId: string, config?: ProviderConfig) => LLMProvider;
|
||||
@@ -0,0 +1,238 @@
|
||||
/**
|
||||
* Observability - Tracing, logging, and metrics hooks
|
||||
*/
|
||||
|
||||
export type SpanKind = 'llm' | 'tool' | 'agent' | 'workflow' | 'custom';
|
||||
export type SpanStatus = 'pending' | 'running' | 'completed' | 'failed';
|
||||
|
||||
export interface SpanData {
|
||||
id: string;
|
||||
traceId: string;
|
||||
parentId?: string;
|
||||
name: string;
|
||||
kind: SpanKind;
|
||||
status: SpanStatus;
|
||||
startTime: number;
|
||||
endTime?: number;
|
||||
attributes: Record<string, any>;
|
||||
events: SpanEvent[];
|
||||
}
|
||||
|
||||
export interface SpanEvent {
|
||||
name: string;
|
||||
timestamp: number;
|
||||
attributes?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface TraceData {
|
||||
id: string;
|
||||
name: string;
|
||||
startTime: number;
|
||||
endTime?: number;
|
||||
status: SpanStatus;
|
||||
spans: SpanData[];
|
||||
metadata: Record<string, any>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Observability Adapter Protocol
|
||||
*/
|
||||
export interface ObservabilityAdapter {
|
||||
startTrace(name: string, metadata?: Record<string, any>): TraceContext;
|
||||
endTrace(traceId: string, status?: SpanStatus): void;
|
||||
startSpan(traceId: string, name: string, kind: SpanKind, parentId?: string): SpanContext;
|
||||
endSpan(spanId: string, status?: SpanStatus, attributes?: Record<string, any>): void;
|
||||
addEvent(spanId: string, name: string, attributes?: Record<string, any>): void;
|
||||
flush(): Promise<void>;
|
||||
}
|
||||
|
||||
export interface TraceContext {
|
||||
traceId: string;
|
||||
startSpan(name: string, kind: SpanKind): SpanContext;
|
||||
end(status?: SpanStatus): void;
|
||||
}
|
||||
|
||||
export interface SpanContext {
|
||||
spanId: string;
|
||||
traceId: string;
|
||||
addEvent(name: string, attributes?: Record<string, any>): void;
|
||||
setAttributes(attributes: Record<string, any>): void;
|
||||
end(status?: SpanStatus): void;
|
||||
}
|
||||
|
||||
/**
|
||||
* In-memory observability adapter for development
|
||||
*/
|
||||
export class MemoryObservabilityAdapter implements ObservabilityAdapter {
|
||||
private traces: Map<string, TraceData> = new Map();
|
||||
private spans: Map<string, SpanData> = new Map();
|
||||
|
||||
startTrace(name: string, metadata: Record<string, any> = {}): TraceContext {
|
||||
const traceId = this.generateId();
|
||||
const trace: TraceData = {
|
||||
id: traceId,
|
||||
name,
|
||||
startTime: Date.now(),
|
||||
status: 'running',
|
||||
spans: [],
|
||||
metadata
|
||||
};
|
||||
this.traces.set(traceId, trace);
|
||||
|
||||
const self = this;
|
||||
return {
|
||||
traceId,
|
||||
startSpan(spanName: string, kind: SpanKind): SpanContext {
|
||||
return self.startSpan(traceId, spanName, kind);
|
||||
},
|
||||
end(status: SpanStatus = 'completed'): void {
|
||||
self.endTrace(traceId, status);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
endTrace(traceId: string, status: SpanStatus = 'completed'): void {
|
||||
const trace = this.traces.get(traceId);
|
||||
if (trace) {
|
||||
trace.endTime = Date.now();
|
||||
trace.status = status;
|
||||
}
|
||||
}
|
||||
|
||||
startSpan(traceId: string, name: string, kind: SpanKind, parentId?: string): SpanContext {
|
||||
const spanId = this.generateId();
|
||||
const span: SpanData = {
|
||||
id: spanId,
|
||||
traceId,
|
||||
parentId,
|
||||
name,
|
||||
kind,
|
||||
status: 'running',
|
||||
startTime: Date.now(),
|
||||
attributes: {},
|
||||
events: []
|
||||
};
|
||||
this.spans.set(spanId, span);
|
||||
|
||||
const trace = this.traces.get(traceId);
|
||||
if (trace) {
|
||||
trace.spans.push(span);
|
||||
}
|
||||
|
||||
const self = this;
|
||||
return {
|
||||
spanId,
|
||||
traceId,
|
||||
addEvent(eventName: string, attributes?: Record<string, any>): void {
|
||||
self.addEvent(spanId, eventName, attributes);
|
||||
},
|
||||
setAttributes(attributes: Record<string, any>): void {
|
||||
const s = self.spans.get(spanId);
|
||||
if (s) {
|
||||
s.attributes = { ...s.attributes, ...attributes };
|
||||
}
|
||||
},
|
||||
end(status: SpanStatus = 'completed'): void {
|
||||
self.endSpan(spanId, status);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
endSpan(spanId: string, status: SpanStatus = 'completed', attributes?: Record<string, any>): void {
|
||||
const span = this.spans.get(spanId);
|
||||
if (span) {
|
||||
span.endTime = Date.now();
|
||||
span.status = status;
|
||||
if (attributes) {
|
||||
span.attributes = { ...span.attributes, ...attributes };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
addEvent(spanId: string, name: string, attributes?: Record<string, any>): void {
|
||||
const span = this.spans.get(spanId);
|
||||
if (span) {
|
||||
span.events.push({
|
||||
name,
|
||||
timestamp: Date.now(),
|
||||
attributes
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async flush(): Promise<void> {
|
||||
// No-op for memory adapter
|
||||
}
|
||||
|
||||
// Utility methods
|
||||
getTrace(traceId: string): TraceData | undefined {
|
||||
return this.traces.get(traceId);
|
||||
}
|
||||
|
||||
getSpan(spanId: string): SpanData | undefined {
|
||||
return this.spans.get(spanId);
|
||||
}
|
||||
|
||||
getAllTraces(): TraceData[] {
|
||||
return Array.from(this.traces.values());
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.traces.clear();
|
||||
this.spans.clear();
|
||||
}
|
||||
|
||||
private generateId(): string {
|
||||
return Math.random().toString(36).substring(2, 15);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Console observability adapter for debugging
|
||||
*/
|
||||
export class ConsoleObservabilityAdapter implements ObservabilityAdapter {
|
||||
private memory = new MemoryObservabilityAdapter();
|
||||
|
||||
startTrace(name: string, metadata?: Record<string, any>): TraceContext {
|
||||
console.log(`[TRACE START] ${name}`, metadata || '');
|
||||
return this.memory.startTrace(name, metadata);
|
||||
}
|
||||
|
||||
endTrace(traceId: string, status?: SpanStatus): void {
|
||||
console.log(`[TRACE END] ${traceId} - ${status || 'completed'}`);
|
||||
this.memory.endTrace(traceId, status);
|
||||
}
|
||||
|
||||
startSpan(traceId: string, name: string, kind: SpanKind, parentId?: string): SpanContext {
|
||||
console.log(` [SPAN START] ${name} (${kind})`);
|
||||
return this.memory.startSpan(traceId, name, kind, parentId);
|
||||
}
|
||||
|
||||
endSpan(spanId: string, status?: SpanStatus, attributes?: Record<string, any>): void {
|
||||
console.log(` [SPAN END] ${spanId} - ${status || 'completed'}`, attributes || '');
|
||||
this.memory.endSpan(spanId, status, attributes);
|
||||
}
|
||||
|
||||
addEvent(spanId: string, name: string, attributes?: Record<string, any>): void {
|
||||
console.log(` [EVENT] ${name}`, attributes || '');
|
||||
this.memory.addEvent(spanId, name, attributes);
|
||||
}
|
||||
|
||||
async flush(): Promise<void> {
|
||||
await this.memory.flush();
|
||||
}
|
||||
}
|
||||
|
||||
// Global observability instance
|
||||
let globalAdapter: ObservabilityAdapter | null = null;
|
||||
|
||||
export function setObservabilityAdapter(adapter: ObservabilityAdapter): void {
|
||||
globalAdapter = adapter;
|
||||
}
|
||||
|
||||
export function getObservabilityAdapter(): ObservabilityAdapter {
|
||||
if (!globalAdapter) {
|
||||
globalAdapter = new MemoryObservabilityAdapter();
|
||||
}
|
||||
return globalAdapter;
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
/**
|
||||
* Session Management - Session, Run, and Trace tracking
|
||||
*/
|
||||
|
||||
import { randomUUID } from 'crypto';
|
||||
|
||||
export interface SessionConfig {
|
||||
id?: string;
|
||||
metadata?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface RunConfig {
|
||||
id?: string;
|
||||
metadata?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface TraceConfig {
|
||||
id?: string;
|
||||
name: string;
|
||||
attributes?: Record<string, any>;
|
||||
}
|
||||
|
||||
export type RunStatus = 'pending' | 'running' | 'completed' | 'failed';
|
||||
export type TraceStatus = 'pending' | 'running' | 'completed' | 'failed';
|
||||
|
||||
export interface Message {
|
||||
id: string;
|
||||
role: 'system' | 'user' | 'assistant' | 'tool';
|
||||
content: string | null;
|
||||
name?: string;
|
||||
tool_call_id?: string;
|
||||
tool_calls?: any[];
|
||||
timestamp: number;
|
||||
runId?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Trace represents a span within a run (e.g., LLM call, tool call)
|
||||
*/
|
||||
export class Trace {
|
||||
readonly id: string;
|
||||
readonly runId: string;
|
||||
readonly name: string;
|
||||
readonly parentId: string | null;
|
||||
readonly startTime: number;
|
||||
endTime: number | null = null;
|
||||
status: TraceStatus = 'pending';
|
||||
attributes: Record<string, any>;
|
||||
private children: Trace[] = [];
|
||||
|
||||
constructor(runId: string, config: TraceConfig, parentId: string | null = null) {
|
||||
this.id = config.id || randomUUID();
|
||||
this.runId = runId;
|
||||
this.name = config.name;
|
||||
this.parentId = parentId;
|
||||
this.startTime = Date.now();
|
||||
this.attributes = config.attributes || {};
|
||||
}
|
||||
|
||||
start(): this {
|
||||
this.status = 'running';
|
||||
return this;
|
||||
}
|
||||
|
||||
complete(attributes?: Record<string, any>): this {
|
||||
this.status = 'completed';
|
||||
this.endTime = Date.now();
|
||||
if (attributes) {
|
||||
this.attributes = { ...this.attributes, ...attributes };
|
||||
}
|
||||
return this;
|
||||
}
|
||||
|
||||
fail(error?: Error): this {
|
||||
this.status = 'failed';
|
||||
this.endTime = Date.now();
|
||||
if (error) {
|
||||
this.attributes.error = error.message;
|
||||
this.attributes.errorStack = error.stack;
|
||||
}
|
||||
return this;
|
||||
}
|
||||
|
||||
get duration(): number | null {
|
||||
if (this.endTime === null) return null;
|
||||
return this.endTime - this.startTime;
|
||||
}
|
||||
|
||||
createChild(config: TraceConfig): Trace {
|
||||
const child = new Trace(this.runId, config, this.id);
|
||||
this.children.push(child);
|
||||
return child;
|
||||
}
|
||||
|
||||
getChildren(): Trace[] {
|
||||
return [...this.children];
|
||||
}
|
||||
|
||||
toJSON(): Record<string, any> {
|
||||
return {
|
||||
id: this.id,
|
||||
runId: this.runId,
|
||||
name: this.name,
|
||||
parentId: this.parentId,
|
||||
startTime: this.startTime,
|
||||
endTime: this.endTime,
|
||||
duration: this.duration,
|
||||
status: this.status,
|
||||
attributes: this.attributes,
|
||||
children: this.children.map(c => c.toJSON()),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Run represents a single execution within a session (e.g., one chat call)
|
||||
*/
|
||||
export class Run {
|
||||
readonly id: string;
|
||||
readonly sessionId: string;
|
||||
readonly startTime: number;
|
||||
endTime: number | null = null;
|
||||
status: RunStatus = 'pending';
|
||||
error: Error | null = null;
|
||||
metadata: Record<string, any>;
|
||||
private traces: Trace[] = [];
|
||||
private messages: Message[] = [];
|
||||
|
||||
constructor(sessionId: string, config: RunConfig = {}) {
|
||||
this.id = config.id || randomUUID();
|
||||
this.sessionId = sessionId;
|
||||
this.startTime = Date.now();
|
||||
this.metadata = config.metadata || {};
|
||||
}
|
||||
|
||||
start(): this {
|
||||
this.status = 'running';
|
||||
return this;
|
||||
}
|
||||
|
||||
complete(metadata?: Record<string, any>): this {
|
||||
this.status = 'completed';
|
||||
this.endTime = Date.now();
|
||||
if (metadata) {
|
||||
this.metadata = { ...this.metadata, ...metadata };
|
||||
}
|
||||
return this;
|
||||
}
|
||||
|
||||
fail(error: Error): this {
|
||||
this.status = 'failed';
|
||||
this.endTime = Date.now();
|
||||
this.error = error;
|
||||
return this;
|
||||
}
|
||||
|
||||
get duration(): number | null {
|
||||
if (this.endTime === null) return null;
|
||||
return this.endTime - this.startTime;
|
||||
}
|
||||
|
||||
createTrace(config: TraceConfig): Trace {
|
||||
const trace = new Trace(this.id, config);
|
||||
this.traces.push(trace);
|
||||
return trace;
|
||||
}
|
||||
|
||||
addMessage(message: Omit<Message, 'id' | 'timestamp' | 'runId'>): Message {
|
||||
const msg: Message = {
|
||||
...message,
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
runId: this.id,
|
||||
};
|
||||
this.messages.push(msg);
|
||||
return msg;
|
||||
}
|
||||
|
||||
getTraces(): Trace[] {
|
||||
return [...this.traces];
|
||||
}
|
||||
|
||||
getMessages(): Message[] {
|
||||
return [...this.messages];
|
||||
}
|
||||
|
||||
toJSON(): Record<string, any> {
|
||||
return {
|
||||
id: this.id,
|
||||
sessionId: this.sessionId,
|
||||
startTime: this.startTime,
|
||||
endTime: this.endTime,
|
||||
duration: this.duration,
|
||||
status: this.status,
|
||||
error: this.error?.message,
|
||||
metadata: this.metadata,
|
||||
traces: this.traces.map(t => t.toJSON()),
|
||||
messageCount: this.messages.length,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Session represents a conversation session with message history
|
||||
*/
|
||||
export class Session {
|
||||
readonly id: string;
|
||||
readonly createdAt: number;
|
||||
metadata: Record<string, any>;
|
||||
private _runs: Run[] = [];
|
||||
private _messages: Message[] = [];
|
||||
|
||||
constructor(config: SessionConfig = {}) {
|
||||
this.id = config.id || randomUUID();
|
||||
this.createdAt = Date.now();
|
||||
this.metadata = config.metadata || {};
|
||||
}
|
||||
|
||||
createRun(config: RunConfig = {}): Run {
|
||||
const run = new Run(this.id, config);
|
||||
this._runs.push(run);
|
||||
return run;
|
||||
}
|
||||
|
||||
get runs(): Run[] {
|
||||
return [...this._runs];
|
||||
}
|
||||
|
||||
get messages(): Message[] {
|
||||
return [...this._messages];
|
||||
}
|
||||
|
||||
addMessage(message: Omit<Message, 'id' | 'timestamp'>): Message {
|
||||
const msg: Message = {
|
||||
...message,
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
this._messages.push(msg);
|
||||
return msg;
|
||||
}
|
||||
|
||||
getMessagesForLLM(): Array<{ role: string; content: string | null; tool_call_id?: string; tool_calls?: any[] }> {
|
||||
return this._messages.map(m => ({
|
||||
role: m.role,
|
||||
content: m.content,
|
||||
...(m.tool_call_id ? { tool_call_id: m.tool_call_id } : {}),
|
||||
...(m.tool_calls ? { tool_calls: m.tool_calls } : {}),
|
||||
}));
|
||||
}
|
||||
|
||||
clearMessages(): void {
|
||||
this._messages = [];
|
||||
}
|
||||
|
||||
toJSON(): Record<string, any> {
|
||||
return {
|
||||
id: this.id,
|
||||
createdAt: this.createdAt,
|
||||
metadata: this.metadata,
|
||||
runCount: this._runs.length,
|
||||
messageCount: this._messages.length,
|
||||
runs: this._runs.map(r => r.toJSON()),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* SessionManager for managing multiple sessions
|
||||
*/
|
||||
export class SessionManager {
|
||||
private sessions: Map<string, Session> = new Map();
|
||||
|
||||
create(config: SessionConfig = {}): Session {
|
||||
const session = new Session(config);
|
||||
this.sessions.set(session.id, session);
|
||||
return session;
|
||||
}
|
||||
|
||||
get(id: string): Session | undefined {
|
||||
return this.sessions.get(id);
|
||||
}
|
||||
|
||||
getOrCreate(id: string, config: Omit<SessionConfig, 'id'> = {}): Session {
|
||||
let session = this.sessions.get(id);
|
||||
if (!session) {
|
||||
session = new Session({ ...config, id });
|
||||
this.sessions.set(id, session);
|
||||
}
|
||||
return session;
|
||||
}
|
||||
|
||||
list(): Session[] {
|
||||
return Array.from(this.sessions.values());
|
||||
}
|
||||
|
||||
delete(id: string): boolean {
|
||||
return this.sessions.delete(id);
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.sessions.clear();
|
||||
}
|
||||
|
||||
toJSON(): Record<string, any> {
|
||||
return {
|
||||
sessionCount: this.sessions.size,
|
||||
sessions: this.list().map(s => s.toJSON()),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Default session manager instance
|
||||
let defaultManager: SessionManager | null = null;
|
||||
|
||||
export function getSessionManager(): SessionManager {
|
||||
if (!defaultManager) {
|
||||
defaultManager = new SessionManager();
|
||||
}
|
||||
return defaultManager;
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
/**
|
||||
* Skills System - Agent Skills standard implementation
|
||||
*/
|
||||
|
||||
import * as fs from 'fs';
|
||||
import * as path from 'path';
|
||||
|
||||
export interface SkillMetadata {
|
||||
name: string;
|
||||
description: string;
|
||||
license?: string;
|
||||
compatibility?: string;
|
||||
allowedTools?: string[];
|
||||
metadata?: Record<string, string>;
|
||||
}
|
||||
|
||||
export interface Skill {
|
||||
metadata: SkillMetadata;
|
||||
instructions: string;
|
||||
path?: string;
|
||||
}
|
||||
|
||||
export interface SkillDiscoveryOptions {
|
||||
paths?: string[];
|
||||
includeSystem?: boolean;
|
||||
includeUser?: boolean;
|
||||
includeProject?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse SKILL.md file
|
||||
*/
|
||||
export function parseSkillFile(content: string): Skill {
|
||||
const frontmatterMatch = content.match(/^---\n([\s\S]*?)\n---\n([\s\S]*)$/);
|
||||
|
||||
if (!frontmatterMatch) {
|
||||
throw new Error('Invalid SKILL.md format: missing frontmatter');
|
||||
}
|
||||
|
||||
const [, frontmatter, instructions] = frontmatterMatch;
|
||||
const metadata = parseYamlFrontmatter(frontmatter);
|
||||
|
||||
if (!metadata.name || !metadata.description) {
|
||||
throw new Error('Invalid SKILL.md: name and description are required');
|
||||
}
|
||||
|
||||
return {
|
||||
metadata: {
|
||||
name: metadata.name,
|
||||
description: metadata.description,
|
||||
license: metadata.license,
|
||||
compatibility: metadata.compatibility,
|
||||
allowedTools: metadata['allowed-tools']?.split(/\s+/),
|
||||
metadata: metadata.metadata
|
||||
},
|
||||
instructions: instructions.trim()
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Simple YAML frontmatter parser
|
||||
*/
|
||||
function parseYamlFrontmatter(yaml: string): Record<string, any> {
|
||||
const result: Record<string, any> = {};
|
||||
const lines = yaml.split('\n');
|
||||
let currentKey = '';
|
||||
let inMetadata = false;
|
||||
const metadataObj: Record<string, string> = {};
|
||||
|
||||
for (const line of lines) {
|
||||
if (line.startsWith('metadata:')) {
|
||||
inMetadata = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (inMetadata && line.startsWith(' ')) {
|
||||
const match = line.match(/^\s+(\w+):\s*(.*)$/);
|
||||
if (match) {
|
||||
metadataObj[match[1]] = match[2].replace(/^["']|["']$/g, '');
|
||||
}
|
||||
continue;
|
||||
} else if (inMetadata && !line.startsWith(' ')) {
|
||||
inMetadata = false;
|
||||
result.metadata = metadataObj;
|
||||
}
|
||||
|
||||
const match = line.match(/^(\S+):\s*(.*)$/);
|
||||
if (match) {
|
||||
currentKey = match[1];
|
||||
result[currentKey] = match[2].replace(/^["']|["']$/g, '');
|
||||
}
|
||||
}
|
||||
|
||||
if (inMetadata) {
|
||||
result.metadata = metadataObj;
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Skill Manager - Load and manage skills
|
||||
*/
|
||||
export class SkillManager {
|
||||
private skills: Map<string, Skill> = new Map();
|
||||
private discoveryPaths: string[] = [];
|
||||
|
||||
constructor(options: SkillDiscoveryOptions = {}) {
|
||||
this.setupDiscoveryPaths(options);
|
||||
}
|
||||
|
||||
private setupDiscoveryPaths(options: SkillDiscoveryOptions): void {
|
||||
if (options.paths) {
|
||||
this.discoveryPaths.push(...options.paths);
|
||||
}
|
||||
|
||||
if (options.includeProject !== false) {
|
||||
this.discoveryPaths.push('./.praison/skills');
|
||||
this.discoveryPaths.push('./.claude/skills');
|
||||
}
|
||||
|
||||
if (options.includeUser !== false) {
|
||||
const home = process.env.HOME || process.env.USERPROFILE || '';
|
||||
if (home) {
|
||||
this.discoveryPaths.push(path.join(home, '.praison/skills'));
|
||||
}
|
||||
}
|
||||
|
||||
if (options.includeSystem) {
|
||||
this.discoveryPaths.push('/etc/praison/skills');
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Discover and load skills from configured paths
|
||||
*/
|
||||
async discover(): Promise<Skill[]> {
|
||||
const discovered: Skill[] = [];
|
||||
|
||||
for (const basePath of this.discoveryPaths) {
|
||||
if (!fs.existsSync(basePath)) continue;
|
||||
|
||||
const entries = fs.readdirSync(basePath, { withFileTypes: true });
|
||||
for (const entry of entries) {
|
||||
if (entry.isDirectory()) {
|
||||
const skillPath = path.join(basePath, entry.name, 'SKILL.md');
|
||||
if (fs.existsSync(skillPath)) {
|
||||
try {
|
||||
const skill = await this.loadSkill(skillPath);
|
||||
skill.path = path.join(basePath, entry.name);
|
||||
discovered.push(skill);
|
||||
this.skills.set(skill.metadata.name, skill);
|
||||
} catch (e) {
|
||||
// Skip invalid skills
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return discovered;
|
||||
}
|
||||
|
||||
/**
|
||||
* Load a skill from a SKILL.md file
|
||||
*/
|
||||
async loadSkill(skillPath: string): Promise<Skill> {
|
||||
const content = fs.readFileSync(skillPath, 'utf-8');
|
||||
return parseSkillFile(content);
|
||||
}
|
||||
|
||||
/**
|
||||
* Register a skill manually
|
||||
*/
|
||||
register(skill: Skill): void {
|
||||
this.skills.set(skill.metadata.name, skill);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a skill by name
|
||||
*/
|
||||
get(name: string): Skill | undefined {
|
||||
return this.skills.get(name);
|
||||
}
|
||||
|
||||
/**
|
||||
* List all loaded skills
|
||||
*/
|
||||
list(): Skill[] {
|
||||
return Array.from(this.skills.values());
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate XML prompt for skills
|
||||
*/
|
||||
generatePrompt(skillNames?: string[]): string {
|
||||
const skillsToInclude = skillNames
|
||||
? skillNames.map(n => this.skills.get(n)).filter(Boolean) as Skill[]
|
||||
: this.list();
|
||||
|
||||
if (skillsToInclude.length === 0) return '';
|
||||
|
||||
const parts = ['<skills>'];
|
||||
for (const skill of skillsToInclude) {
|
||||
parts.push(` <skill name="${skill.metadata.name}">`);
|
||||
parts.push(` <description>${skill.metadata.description}</description>`);
|
||||
parts.push(` <instructions>${skill.instructions}</instructions>`);
|
||||
parts.push(' </skill>');
|
||||
}
|
||||
parts.push('</skills>');
|
||||
|
||||
return parts.join('\n');
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate a skill
|
||||
*/
|
||||
validate(skill: Skill): { valid: boolean; errors: string[] } {
|
||||
const errors: string[] = [];
|
||||
|
||||
if (!skill.metadata.name) {
|
||||
errors.push('Missing required field: name');
|
||||
} else if (!/^[a-z0-9-]{1,64}$/.test(skill.metadata.name)) {
|
||||
errors.push('Invalid name: must be 1-64 lowercase chars, numbers, or hyphens');
|
||||
}
|
||||
|
||||
if (!skill.metadata.description) {
|
||||
errors.push('Missing required field: description');
|
||||
} else if (skill.metadata.description.length > 1024) {
|
||||
errors.push('Description too long: max 1024 characters');
|
||||
}
|
||||
|
||||
if (skill.instructions.length > 5000) {
|
||||
errors.push('Instructions too long: max 5000 characters');
|
||||
}
|
||||
|
||||
return { valid: errors.length === 0, errors };
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a skill manager
|
||||
*/
|
||||
export function createSkillManager(options?: SkillDiscoveryOptions): SkillManager {
|
||||
return new SkillManager(options);
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
/**
|
||||
* Tool Decorator and Registry - Type-safe tool creation and management
|
||||
*/
|
||||
|
||||
export interface ToolParameters {
|
||||
type: 'object';
|
||||
properties: Record<string, {
|
||||
type: string;
|
||||
description?: string;
|
||||
enum?: string[];
|
||||
default?: any;
|
||||
}>;
|
||||
required?: string[];
|
||||
}
|
||||
|
||||
export interface ToolConfig<TParams = any, TResult = any> {
|
||||
name: string;
|
||||
description?: string;
|
||||
parameters?: ToolParameters | any; // Support Zod schemas
|
||||
category?: string;
|
||||
execute: (params: TParams, context?: ToolContext) => Promise<TResult> | TResult;
|
||||
}
|
||||
|
||||
export interface ToolContext {
|
||||
agentName?: string;
|
||||
sessionId?: string;
|
||||
runId?: string;
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
|
||||
export interface ToolDefinition {
|
||||
name: string;
|
||||
description: string;
|
||||
parameters: ToolParameters;
|
||||
category?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Tool class - Represents an executable tool
|
||||
*/
|
||||
export class FunctionTool<TParams = any, TResult = any> {
|
||||
readonly name: string;
|
||||
readonly description: string;
|
||||
readonly parameters: ToolParameters;
|
||||
readonly category?: string;
|
||||
private readonly executeFn: (params: TParams, context?: ToolContext) => Promise<TResult> | TResult;
|
||||
|
||||
constructor(config: ToolConfig<TParams, TResult>) {
|
||||
this.name = config.name;
|
||||
this.description = config.description || `Function ${config.name}`;
|
||||
this.parameters = this.normalizeParameters(config.parameters);
|
||||
this.category = config.category;
|
||||
this.executeFn = config.execute;
|
||||
}
|
||||
|
||||
private normalizeParameters(params: any): ToolParameters {
|
||||
if (!params) {
|
||||
return { type: 'object', properties: {}, required: [] };
|
||||
}
|
||||
|
||||
// Check if it's a Zod schema
|
||||
if (params && typeof params.parse === 'function' && typeof params._def === 'object') {
|
||||
return this.zodToJsonSchema(params);
|
||||
}
|
||||
|
||||
// Already a JSON schema
|
||||
return params;
|
||||
}
|
||||
|
||||
private zodToJsonSchema(zodSchema: any): ToolParameters {
|
||||
// Basic Zod to JSON Schema conversion
|
||||
// For full support, use zod-to-json-schema package
|
||||
const def = zodSchema._def;
|
||||
|
||||
if (def.typeName === 'ZodObject') {
|
||||
const properties: Record<string, any> = {};
|
||||
const required: string[] = [];
|
||||
|
||||
for (const [key, value] of Object.entries(def.shape())) {
|
||||
const fieldDef = (value as any)._def;
|
||||
properties[key] = this.zodFieldToJsonSchema(fieldDef);
|
||||
|
||||
// Check if field is required (not optional)
|
||||
if (fieldDef.typeName !== 'ZodOptional' && fieldDef.typeName !== 'ZodDefault') {
|
||||
required.push(key);
|
||||
}
|
||||
}
|
||||
|
||||
return { type: 'object', properties, required };
|
||||
}
|
||||
|
||||
return { type: 'object', properties: {}, required: [] };
|
||||
}
|
||||
|
||||
private zodFieldToJsonSchema(def: any): any {
|
||||
const typeName = def.typeName;
|
||||
|
||||
switch (typeName) {
|
||||
case 'ZodString':
|
||||
return { type: 'string', description: def.description };
|
||||
case 'ZodNumber':
|
||||
return { type: 'number', description: def.description };
|
||||
case 'ZodBoolean':
|
||||
return { type: 'boolean', description: def.description };
|
||||
case 'ZodArray':
|
||||
return { type: 'array', items: this.zodFieldToJsonSchema(def.type._def) };
|
||||
case 'ZodEnum':
|
||||
return { type: 'string', enum: def.values };
|
||||
case 'ZodOptional':
|
||||
return this.zodFieldToJsonSchema(def.innerType._def);
|
||||
case 'ZodDefault':
|
||||
const inner = this.zodFieldToJsonSchema(def.innerType._def);
|
||||
inner.default = def.defaultValue();
|
||||
return inner;
|
||||
default:
|
||||
return { type: 'string' };
|
||||
}
|
||||
}
|
||||
|
||||
async execute(params: TParams, context?: ToolContext): Promise<TResult> {
|
||||
// Validate parameters if we have a schema
|
||||
// For now, just execute - validation can be added later
|
||||
return this.executeFn(params, context);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the tool definition for LLM
|
||||
*/
|
||||
getDefinition(): ToolDefinition {
|
||||
return {
|
||||
name: this.name,
|
||||
description: this.description,
|
||||
parameters: this.parameters,
|
||||
category: this.category,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get OpenAI-compatible tool format
|
||||
*/
|
||||
toOpenAITool(): { type: 'function'; function: { name: string; description: string; parameters: any } } {
|
||||
return {
|
||||
type: 'function',
|
||||
function: {
|
||||
name: this.name,
|
||||
description: this.description,
|
||||
parameters: this.parameters,
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a tool from a configuration object
|
||||
*/
|
||||
export function tool<TParams = any, TResult = any>(
|
||||
config: ToolConfig<TParams, TResult>
|
||||
): FunctionTool<TParams, TResult> {
|
||||
return new FunctionTool(config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Tool Registry - Manages tool registration and lookup
|
||||
*/
|
||||
export class ToolRegistry {
|
||||
private tools: Map<string, FunctionTool> = new Map();
|
||||
|
||||
register(tool: FunctionTool, options?: { overwrite?: boolean }): this {
|
||||
if (this.tools.has(tool.name) && !options?.overwrite) {
|
||||
throw new Error(`Tool '${tool.name}' is already registered. Use { overwrite: true } to replace.`);
|
||||
}
|
||||
this.tools.set(tool.name, tool);
|
||||
return this;
|
||||
}
|
||||
|
||||
get(name: string): FunctionTool | undefined {
|
||||
return this.tools.get(name);
|
||||
}
|
||||
|
||||
has(name: string): boolean {
|
||||
return this.tools.has(name);
|
||||
}
|
||||
|
||||
list(): FunctionTool[] {
|
||||
return Array.from(this.tools.values());
|
||||
}
|
||||
|
||||
getByCategory(category: string): FunctionTool[] {
|
||||
return this.list().filter(t => t.category === category);
|
||||
}
|
||||
|
||||
getDefinitions(): ToolDefinition[] {
|
||||
return this.list().map(t => t.getDefinition());
|
||||
}
|
||||
|
||||
toOpenAITools(): Array<{ type: 'function'; function: any }> {
|
||||
return this.list().map(t => t.toOpenAITool());
|
||||
}
|
||||
|
||||
delete(name: string): boolean {
|
||||
return this.tools.delete(name);
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.tools.clear();
|
||||
}
|
||||
}
|
||||
|
||||
// Global registry instance
|
||||
let globalRegistry: ToolRegistry | null = null;
|
||||
|
||||
export function getRegistry(): ToolRegistry {
|
||||
if (!globalRegistry) {
|
||||
globalRegistry = new ToolRegistry();
|
||||
}
|
||||
return globalRegistry;
|
||||
}
|
||||
|
||||
export function registerTool(tool: FunctionTool, options?: { overwrite?: boolean }): void {
|
||||
getRegistry().register(tool, options);
|
||||
}
|
||||
|
||||
export function getTool(name: string): FunctionTool | undefined {
|
||||
return getRegistry().get(name);
|
||||
}
|
||||
@@ -19,6 +19,9 @@ export class BaseTool implements Tool {
|
||||
}
|
||||
}
|
||||
|
||||
// Export decorator and registry
|
||||
export * from './decorator';
|
||||
|
||||
// Export all tool modules
|
||||
export * from './arxivTools';
|
||||
export * from './mcpSse';
|
||||
@@ -0,0 +1,300 @@
|
||||
/**
|
||||
* Workflows - Pipeline and orchestration patterns
|
||||
*/
|
||||
|
||||
import { randomUUID } from 'crypto';
|
||||
|
||||
export type StepStatus = 'pending' | 'running' | 'completed' | 'failed' | 'skipped';
|
||||
|
||||
export interface StepResult<T = any> {
|
||||
stepId: string;
|
||||
stepName: string;
|
||||
status: StepStatus;
|
||||
output?: T;
|
||||
error?: Error;
|
||||
duration: number;
|
||||
startedAt: number;
|
||||
completedAt?: number;
|
||||
}
|
||||
|
||||
export interface WorkflowContext {
|
||||
workflowId: string;
|
||||
stepResults: Map<string, StepResult>;
|
||||
metadata: Record<string, any>;
|
||||
get<T = any>(stepName: string): T | undefined;
|
||||
set(key: string, value: any): void;
|
||||
}
|
||||
|
||||
export type StepFunction<TInput = any, TOutput = any> = (
|
||||
input: TInput,
|
||||
context: WorkflowContext
|
||||
) => Promise<TOutput> | TOutput;
|
||||
|
||||
export interface WorkflowStepConfig<TInput = any, TOutput = any> {
|
||||
name: string;
|
||||
execute: StepFunction<TInput, TOutput>;
|
||||
condition?: (context: WorkflowContext) => boolean;
|
||||
onError?: 'fail' | 'skip' | 'retry';
|
||||
maxRetries?: number;
|
||||
timeout?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Workflow Step
|
||||
*/
|
||||
export class WorkflowStep<TInput = any, TOutput = any> {
|
||||
readonly id: string;
|
||||
readonly name: string;
|
||||
readonly execute: StepFunction<TInput, TOutput>;
|
||||
readonly condition?: (context: WorkflowContext) => boolean;
|
||||
readonly onError: 'fail' | 'skip' | 'retry';
|
||||
readonly maxRetries: number;
|
||||
readonly timeout?: number;
|
||||
|
||||
constructor(config: WorkflowStepConfig<TInput, TOutput>) {
|
||||
this.id = randomUUID();
|
||||
this.name = config.name;
|
||||
this.execute = config.execute;
|
||||
this.condition = config.condition;
|
||||
this.onError = config.onError || 'fail';
|
||||
this.maxRetries = config.maxRetries || 0;
|
||||
this.timeout = config.timeout;
|
||||
}
|
||||
|
||||
async run(input: TInput, context: WorkflowContext): Promise<StepResult<TOutput>> {
|
||||
const startedAt = Date.now();
|
||||
|
||||
// Check condition
|
||||
if (this.condition && !this.condition(context)) {
|
||||
return {
|
||||
stepId: this.id,
|
||||
stepName: this.name,
|
||||
status: 'skipped',
|
||||
duration: 0,
|
||||
startedAt,
|
||||
completedAt: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
let lastError: Error | undefined;
|
||||
let attempts = 0;
|
||||
|
||||
while (attempts <= this.maxRetries) {
|
||||
attempts++;
|
||||
try {
|
||||
const output = await this.executeWithTimeout(input, context);
|
||||
const completedAt = Date.now();
|
||||
|
||||
return {
|
||||
stepId: this.id,
|
||||
stepName: this.name,
|
||||
status: 'completed',
|
||||
output,
|
||||
duration: completedAt - startedAt,
|
||||
startedAt,
|
||||
completedAt,
|
||||
};
|
||||
} catch (error: any) {
|
||||
lastError = error;
|
||||
|
||||
if (this.onError === 'skip') {
|
||||
return {
|
||||
stepId: this.id,
|
||||
stepName: this.name,
|
||||
status: 'skipped',
|
||||
error: lastError,
|
||||
duration: Date.now() - startedAt,
|
||||
startedAt,
|
||||
completedAt: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
if (this.onError !== 'retry' || attempts > this.maxRetries) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
stepId: this.id,
|
||||
stepName: this.name,
|
||||
status: 'failed',
|
||||
error: lastError,
|
||||
duration: Date.now() - startedAt,
|
||||
startedAt,
|
||||
completedAt: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
private async executeWithTimeout(input: TInput, context: WorkflowContext): Promise<TOutput> {
|
||||
if (!this.timeout) {
|
||||
return this.execute(input, context);
|
||||
}
|
||||
|
||||
return Promise.race([
|
||||
this.execute(input, context),
|
||||
new Promise<never>((_, reject) => {
|
||||
setTimeout(() => reject(new Error(`Step ${this.name} timed out after ${this.timeout}ms`)), this.timeout);
|
||||
}),
|
||||
]);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a workflow context
|
||||
*/
|
||||
function createContext(workflowId: string): WorkflowContext {
|
||||
const stepResults = new Map<string, StepResult>();
|
||||
const metadata: Record<string, any> = {};
|
||||
|
||||
return {
|
||||
workflowId,
|
||||
stepResults,
|
||||
metadata,
|
||||
get<T = any>(stepName: string): T | undefined {
|
||||
const result = stepResults.get(stepName);
|
||||
return result?.output as T | undefined;
|
||||
},
|
||||
set(key: string, value: any): void {
|
||||
metadata[key] = value;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Workflow - Sequential pipeline execution
|
||||
*/
|
||||
export class Workflow<TInput = any, TOutput = any> {
|
||||
readonly id: string;
|
||||
readonly name: string;
|
||||
private steps: WorkflowStep[] = [];
|
||||
|
||||
constructor(name: string) {
|
||||
this.id = randomUUID();
|
||||
this.name = name;
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a step to the workflow
|
||||
*/
|
||||
addStep<TStepInput = any, TStepOutput = any>(
|
||||
config: WorkflowStepConfig<TStepInput, TStepOutput>
|
||||
): this {
|
||||
this.steps.push(new WorkflowStep(config));
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a step using a simpler syntax
|
||||
*/
|
||||
step<TStepInput = any, TStepOutput = any>(
|
||||
name: string,
|
||||
execute: StepFunction<TStepInput, TStepOutput>
|
||||
): this {
|
||||
return this.addStep({ name, execute });
|
||||
}
|
||||
|
||||
/**
|
||||
* Run the workflow
|
||||
*/
|
||||
async run(input: TInput): Promise<{ output: TOutput | undefined; results: StepResult[]; context: WorkflowContext }> {
|
||||
const context = createContext(this.id);
|
||||
const results: StepResult[] = [];
|
||||
let currentInput: any = input;
|
||||
|
||||
for (const step of this.steps) {
|
||||
const result = await step.run(currentInput, context);
|
||||
results.push(result);
|
||||
context.stepResults.set(step.name, result);
|
||||
|
||||
if (result.status === 'failed') {
|
||||
return { output: undefined, results, context };
|
||||
}
|
||||
|
||||
if (result.status === 'completed') {
|
||||
currentInput = result.output;
|
||||
}
|
||||
}
|
||||
|
||||
const lastResult = results[results.length - 1];
|
||||
return {
|
||||
output: lastResult?.output as TOutput | undefined,
|
||||
results,
|
||||
context,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get step count
|
||||
*/
|
||||
get stepCount(): number {
|
||||
return this.steps.length;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Pipeline - Alias for Workflow
|
||||
*/
|
||||
export const Pipeline = Workflow;
|
||||
|
||||
/**
|
||||
* Parallel execution helper
|
||||
*/
|
||||
export async function parallel<T>(
|
||||
tasks: Array<() => Promise<T>>
|
||||
): Promise<T[]> {
|
||||
return Promise.all(tasks.map(task => task()));
|
||||
}
|
||||
|
||||
/**
|
||||
* Route helper - Execute based on condition
|
||||
*/
|
||||
export async function route<T>(
|
||||
conditions: Array<{ condition: () => boolean; execute: () => Promise<T> }>,
|
||||
defaultExecute?: () => Promise<T>
|
||||
): Promise<T | undefined> {
|
||||
for (const { condition, execute } of conditions) {
|
||||
if (condition()) {
|
||||
return execute();
|
||||
}
|
||||
}
|
||||
return defaultExecute?.();
|
||||
}
|
||||
|
||||
/**
|
||||
* Loop helper - Repeat until condition
|
||||
*/
|
||||
export async function loop<T>(
|
||||
execute: (iteration: number) => Promise<T>,
|
||||
shouldContinue: (result: T, iteration: number) => boolean,
|
||||
maxIterations: number = 100
|
||||
): Promise<T[]> {
|
||||
const results: T[] = [];
|
||||
let iteration = 0;
|
||||
|
||||
while (iteration < maxIterations) {
|
||||
const result = await execute(iteration);
|
||||
results.push(result);
|
||||
|
||||
if (!shouldContinue(result, iteration)) {
|
||||
break;
|
||||
}
|
||||
iteration++;
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
/**
|
||||
* Repeat helper - Execute N times
|
||||
*/
|
||||
export async function repeat<T>(
|
||||
execute: (iteration: number) => Promise<T>,
|
||||
times: number
|
||||
): Promise<T[]> {
|
||||
const results: T[] = [];
|
||||
for (let i = 0; i < times; i++) {
|
||||
results.push(await execute(i));
|
||||
}
|
||||
return results;
|
||||
}
|
||||
@@ -0,0 +1,398 @@
|
||||
#!/usr/bin/env npx ts-node
|
||||
/**
|
||||
* Comprehensive Feature Verification Script
|
||||
* Tests all implemented features with real API keys
|
||||
*/
|
||||
|
||||
import {
|
||||
// Providers
|
||||
createProvider,
|
||||
isProviderAvailable,
|
||||
getAvailableProviders,
|
||||
|
||||
// Session
|
||||
Session,
|
||||
SessionManager,
|
||||
getSessionManager,
|
||||
|
||||
// Tools
|
||||
tool,
|
||||
ToolRegistry,
|
||||
getRegistry,
|
||||
|
||||
// Database
|
||||
MemoryDbAdapter,
|
||||
db,
|
||||
|
||||
// Workflows
|
||||
Workflow,
|
||||
parallel,
|
||||
route,
|
||||
loop,
|
||||
|
||||
// Guardrails
|
||||
guardrail,
|
||||
GuardrailManager,
|
||||
builtinGuardrails,
|
||||
|
||||
// Knowledge Base
|
||||
KnowledgeBase,
|
||||
createKnowledgeBase,
|
||||
|
||||
// Evaluation
|
||||
accuracyEval,
|
||||
performanceEval,
|
||||
reliabilityEval,
|
||||
EvalSuite,
|
||||
|
||||
// Observability
|
||||
MemoryObservabilityAdapter,
|
||||
ConsoleObservabilityAdapter,
|
||||
|
||||
// Skills
|
||||
SkillManager,
|
||||
parseSkillFile,
|
||||
|
||||
// Router
|
||||
RouterAgent,
|
||||
createRouter,
|
||||
routeConditions,
|
||||
|
||||
// Handoff
|
||||
Handoff,
|
||||
handoff,
|
||||
handoffFilters,
|
||||
} from '../../src';
|
||||
|
||||
console.log('='.repeat(60));
|
||||
console.log('PraisonAI TypeScript - Comprehensive Feature Verification');
|
||||
console.log('='.repeat(60));
|
||||
|
||||
let passed = 0;
|
||||
let failed = 0;
|
||||
|
||||
async function test(name: string, fn: () => Promise<void>) {
|
||||
try {
|
||||
await fn();
|
||||
console.log(`✅ ${name}`);
|
||||
passed++;
|
||||
} catch (error: any) {
|
||||
console.log(`❌ ${name}: ${error.message}`);
|
||||
failed++;
|
||||
}
|
||||
}
|
||||
|
||||
async function main() {
|
||||
// 1. Provider Tests
|
||||
console.log('\n--- 1. LLM Providers ---');
|
||||
|
||||
await test('Provider availability check', async () => {
|
||||
const available = getAvailableProviders();
|
||||
if (!available.includes('openai')) throw new Error('OpenAI not available');
|
||||
});
|
||||
|
||||
await test('OpenAI generateText', async () => {
|
||||
const provider = createProvider('openai/gpt-4o-mini');
|
||||
const result = await provider.generateText({
|
||||
messages: [{ role: 'user', content: 'Say "test" only' }]
|
||||
});
|
||||
if (!result.text) throw new Error('No response');
|
||||
});
|
||||
|
||||
await test('OpenAI streaming', async () => {
|
||||
const provider = createProvider('openai/gpt-4o-mini');
|
||||
let text = '';
|
||||
const stream = await provider.streamText({
|
||||
messages: [{ role: 'user', content: 'Say "ok"' }],
|
||||
onToken: (t) => { text += t; }
|
||||
});
|
||||
// Consume the stream
|
||||
for await (const chunk of stream) {
|
||||
if (chunk.text) text += chunk.text;
|
||||
}
|
||||
if (!text) throw new Error('No streamed text');
|
||||
});
|
||||
|
||||
await test('Anthropic generateText', async () => {
|
||||
const provider = createProvider('anthropic/claude-sonnet-4-20250514');
|
||||
const result = await provider.generateText({
|
||||
messages: [{ role: 'user', content: 'Say "test" only' }]
|
||||
});
|
||||
if (!result.text) throw new Error('No response');
|
||||
});
|
||||
|
||||
// 2. Session Tests
|
||||
console.log('\n--- 2. Session Management ---');
|
||||
|
||||
await test('Session creation', async () => {
|
||||
const session = new Session();
|
||||
if (!session.id) throw new Error('No session ID');
|
||||
});
|
||||
|
||||
await test('Session messages', async () => {
|
||||
const session = new Session();
|
||||
session.addMessage({ role: 'user', content: 'Hello' });
|
||||
session.addMessage({ role: 'assistant', content: 'Hi' });
|
||||
const messages = session.getMessagesForLLM();
|
||||
if (messages.length !== 2) throw new Error('Wrong message count');
|
||||
});
|
||||
|
||||
await test('Session runs', async () => {
|
||||
const session = new Session();
|
||||
const run = session.createRun().start();
|
||||
run.complete();
|
||||
if (run.status !== 'completed') throw new Error('Run not completed');
|
||||
});
|
||||
|
||||
await test('SessionManager', async () => {
|
||||
const manager = getSessionManager();
|
||||
const session = manager.create({ id: 'test-session' });
|
||||
const retrieved = manager.get('test-session');
|
||||
if (!retrieved) throw new Error('Session not found');
|
||||
manager.delete('test-session');
|
||||
});
|
||||
|
||||
// 3. Tool Tests
|
||||
console.log('\n--- 3. Tool System ---');
|
||||
|
||||
await test('Tool creation', async () => {
|
||||
const myTool = tool({
|
||||
name: 'test_tool',
|
||||
description: 'A test tool',
|
||||
execute: async () => 'result'
|
||||
});
|
||||
const result = await myTool.execute({});
|
||||
if (result !== 'result') throw new Error('Wrong result');
|
||||
});
|
||||
|
||||
await test('Tool registry', async () => {
|
||||
const registry = new ToolRegistry();
|
||||
registry.register(tool({ name: 'tool1', execute: async () => '1' }));
|
||||
registry.register(tool({ name: 'tool2', execute: async () => '2' }));
|
||||
if (registry.list().length !== 2) throw new Error('Wrong tool count');
|
||||
});
|
||||
|
||||
await test('OpenAI tool format', async () => {
|
||||
const myTool = tool({
|
||||
name: 'calculator',
|
||||
description: 'Calculate',
|
||||
parameters: { type: 'object', properties: { expr: { type: 'string' } } },
|
||||
execute: async () => '42'
|
||||
});
|
||||
const format = myTool.toOpenAITool();
|
||||
if (format.type !== 'function') throw new Error('Wrong format');
|
||||
});
|
||||
|
||||
// 4. Database Tests
|
||||
console.log('\n--- 4. Database Adapters ---');
|
||||
|
||||
await test('Memory adapter', async () => {
|
||||
const adapter = new MemoryDbAdapter();
|
||||
await adapter.connect();
|
||||
await adapter.createSession({ id: 'test', createdAt: Date.now(), updatedAt: Date.now() });
|
||||
const session = await adapter.getSession('test');
|
||||
if (!session) throw new Error('Session not found');
|
||||
await adapter.disconnect();
|
||||
});
|
||||
|
||||
await test('DB factory', async () => {
|
||||
const adapter = db({ type: 'memory' });
|
||||
await adapter.connect();
|
||||
if (!adapter.isConnected()) throw new Error('Not connected');
|
||||
await adapter.disconnect();
|
||||
});
|
||||
|
||||
// 5. Workflow Tests
|
||||
console.log('\n--- 5. Workflows ---');
|
||||
|
||||
await test('Basic workflow', async () => {
|
||||
const workflow = new Workflow<number, number>('test')
|
||||
.step('add', async (n) => n + 10)
|
||||
.step('multiply', async (n) => n * 2);
|
||||
const { output } = await workflow.run(5);
|
||||
if (output !== 30) throw new Error(`Expected 30, got ${output}`);
|
||||
});
|
||||
|
||||
await test('Parallel execution', async () => {
|
||||
const results = await parallel([
|
||||
async () => 1,
|
||||
async () => 2,
|
||||
async () => 3
|
||||
]);
|
||||
if (results.length !== 3) throw new Error('Wrong result count');
|
||||
});
|
||||
|
||||
await test('Route helper', async () => {
|
||||
const result = await route([
|
||||
{ condition: () => false, execute: async () => 'a' },
|
||||
{ condition: () => true, execute: async () => 'b' }
|
||||
]);
|
||||
if (result !== 'b') throw new Error('Wrong route');
|
||||
});
|
||||
|
||||
// 6. Guardrails Tests
|
||||
console.log('\n--- 6. Guardrails ---');
|
||||
|
||||
await test('Basic guardrail', async () => {
|
||||
const g = guardrail({
|
||||
name: 'test',
|
||||
check: (content) => ({ status: content.length > 0 ? 'passed' : 'failed' })
|
||||
});
|
||||
const result = await g.run('hello');
|
||||
if (result.status !== 'passed') throw new Error('Should pass');
|
||||
});
|
||||
|
||||
await test('Built-in guardrails', async () => {
|
||||
const maxLen = builtinGuardrails.maxLength(10);
|
||||
const result = await maxLen.run('hello');
|
||||
if (result.status !== 'passed') throw new Error('Should pass');
|
||||
});
|
||||
|
||||
await test('Guardrail manager', async () => {
|
||||
const manager = new GuardrailManager();
|
||||
manager.add(builtinGuardrails.minLength(1));
|
||||
manager.add(builtinGuardrails.maxLength(100));
|
||||
const { passed } = await manager.runAll('test');
|
||||
if (!passed) throw new Error('Should pass');
|
||||
});
|
||||
|
||||
// 7. Knowledge Base Tests
|
||||
console.log('\n--- 7. Knowledge Base (RAG) ---');
|
||||
|
||||
await test('Knowledge base add/search', async () => {
|
||||
const kb = createKnowledgeBase();
|
||||
await kb.add({ id: 'doc1', content: 'TypeScript is great' });
|
||||
await kb.add({ id: 'doc2', content: 'Python is also great' });
|
||||
const results = await kb.search('TypeScript');
|
||||
if (results.length === 0) throw new Error('No results');
|
||||
});
|
||||
|
||||
await test('Knowledge base context', async () => {
|
||||
const kb = createKnowledgeBase();
|
||||
await kb.add({ id: 'doc1', content: 'Test document' });
|
||||
const results = await kb.search('test');
|
||||
const context = kb.buildContext(results);
|
||||
if (!context.includes('Test document')) throw new Error('Wrong context');
|
||||
});
|
||||
|
||||
// 8. Evaluation Tests
|
||||
console.log('\n--- 8. Evaluation Framework ---');
|
||||
|
||||
await test('Accuracy eval', async () => {
|
||||
const result = await accuracyEval({
|
||||
input: 'test',
|
||||
expectedOutput: 'hello world',
|
||||
actualOutput: 'hello world'
|
||||
});
|
||||
if (!result.passed) throw new Error('Should pass');
|
||||
});
|
||||
|
||||
await test('Performance eval', async () => {
|
||||
const result = await performanceEval({
|
||||
func: async () => 'done',
|
||||
iterations: 3,
|
||||
warmupRuns: 1
|
||||
});
|
||||
if (result.times.length !== 3) throw new Error('Wrong iteration count');
|
||||
});
|
||||
|
||||
await test('Reliability eval', async () => {
|
||||
const result = await reliabilityEval({
|
||||
expectedToolCalls: ['a', 'b'],
|
||||
actualToolCalls: ['a', 'b']
|
||||
});
|
||||
if (!result.passed) throw new Error('Should pass');
|
||||
});
|
||||
|
||||
await test('Eval suite', async () => {
|
||||
const suite = new EvalSuite();
|
||||
await suite.runAccuracy('test', {
|
||||
input: 'x',
|
||||
expectedOutput: 'same',
|
||||
actualOutput: 'same'
|
||||
});
|
||||
const summary = suite.getSummary();
|
||||
if (summary.total !== 1) throw new Error('Wrong count');
|
||||
});
|
||||
|
||||
// 9. Observability Tests
|
||||
console.log('\n--- 9. Observability ---');
|
||||
|
||||
await test('Memory observability', async () => {
|
||||
const adapter = new MemoryObservabilityAdapter();
|
||||
const trace = adapter.startTrace('test');
|
||||
const span = trace.startSpan('operation', 'custom');
|
||||
span.end();
|
||||
trace.end();
|
||||
const traceData = adapter.getTrace(trace.traceId);
|
||||
if (!traceData) throw new Error('Trace not found');
|
||||
});
|
||||
|
||||
// 10. Skills Tests
|
||||
console.log('\n--- 10. Skills System ---');
|
||||
|
||||
await test('Parse skill file', async () => {
|
||||
const content = `---
|
||||
name: test-skill
|
||||
description: A test skill
|
||||
---
|
||||
# Instructions
|
||||
Do something`;
|
||||
const skill = parseSkillFile(content);
|
||||
if (skill.metadata.name !== 'test-skill') throw new Error('Wrong name');
|
||||
});
|
||||
|
||||
await test('Skill manager', async () => {
|
||||
const manager = new SkillManager({ paths: [] });
|
||||
manager.register({
|
||||
metadata: { name: 'manual-skill', description: 'Manual' },
|
||||
instructions: 'Do this'
|
||||
});
|
||||
const skill = manager.get('manual-skill');
|
||||
if (!skill) throw new Error('Skill not found');
|
||||
});
|
||||
|
||||
// 11. Router Tests
|
||||
console.log('\n--- 11. Router Agent ---');
|
||||
|
||||
await test('Route conditions', async () => {
|
||||
const kw = routeConditions.keywords(['math', 'calc']);
|
||||
if (!kw('I need math help')) throw new Error('Should match');
|
||||
if (kw('I need code help')) throw new Error('Should not match');
|
||||
});
|
||||
|
||||
await test('Router creation', async () => {
|
||||
const router = createRouter({
|
||||
routes: [],
|
||||
name: 'TestRouter'
|
||||
});
|
||||
if (router.name !== 'TestRouter') throw new Error('Wrong name');
|
||||
});
|
||||
|
||||
// 12. Handoff Tests
|
||||
console.log('\n--- 12. Handoff System ---');
|
||||
|
||||
await test('Handoff filters', async () => {
|
||||
const topicFilter = handoffFilters.topic(['billing', 'payment']);
|
||||
if (!topicFilter({ messages: [], lastMessage: 'I have a billing question' })) {
|
||||
throw new Error('Should match');
|
||||
}
|
||||
});
|
||||
|
||||
await test('Handoff creation', async () => {
|
||||
const mockAgent = { name: 'SupportAgent', chat: async () => ({ text: 'ok' }) } as any;
|
||||
const h = handoff({ agent: mockAgent, name: 'support_handoff' });
|
||||
if (h.name !== 'support_handoff') throw new Error('Wrong name');
|
||||
});
|
||||
|
||||
// Summary
|
||||
console.log('\n' + '='.repeat(60));
|
||||
console.log(`RESULTS: ${passed} passed, ${failed} failed`);
|
||||
console.log('='.repeat(60));
|
||||
|
||||
if (failed > 0) {
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,150 @@
|
||||
/**
|
||||
* Direct verification script for provider implementation
|
||||
* Run with: npx ts-node tests/development/verify-providers.ts
|
||||
*/
|
||||
|
||||
import { createProvider } from '../../src/llm/providers';
|
||||
|
||||
async function testOpenAI() {
|
||||
console.log('\n=== Testing OpenAI Provider ===');
|
||||
|
||||
if (!process.env.OPENAI_API_KEY) {
|
||||
console.log('OPENAI_API_KEY not set, skipping');
|
||||
return;
|
||||
}
|
||||
|
||||
const provider = createProvider('openai/gpt-4o-mini');
|
||||
|
||||
// Test generateText
|
||||
console.log('\n1. Testing generateText...');
|
||||
const result = await provider.generateText({
|
||||
messages: [
|
||||
{ role: 'system', content: 'Be concise.' },
|
||||
{ role: 'user', content: 'Say hello in one word.' },
|
||||
],
|
||||
maxTokens: 10,
|
||||
temperature: 0,
|
||||
});
|
||||
console.log('Response:', result.text);
|
||||
console.log('Usage:', result.usage);
|
||||
|
||||
// Test streaming
|
||||
console.log('\n2. Testing streamText...');
|
||||
let streamedText = '';
|
||||
const stream = await provider.streamText({
|
||||
messages: [{ role: 'user', content: 'Count 1 to 3' }],
|
||||
maxTokens: 20,
|
||||
onToken: (token) => process.stdout.write(token),
|
||||
});
|
||||
for await (const chunk of stream) {
|
||||
if (chunk.text) streamedText += chunk.text;
|
||||
}
|
||||
console.log('\nStreamed text:', streamedText);
|
||||
|
||||
// Test tool calling
|
||||
console.log('\n3. Testing tool calling...');
|
||||
const toolResult = await provider.generateText({
|
||||
messages: [{ role: 'user', content: 'What is 5+3? Use calculator.' }],
|
||||
tools: [{
|
||||
name: 'calculator',
|
||||
description: 'Math calculator',
|
||||
parameters: {
|
||||
type: 'object',
|
||||
properties: { expression: { type: 'string' } },
|
||||
required: ['expression'],
|
||||
},
|
||||
}],
|
||||
maxTokens: 100,
|
||||
});
|
||||
console.log('Tool calls:', toolResult.toolCalls);
|
||||
|
||||
// Test structured output
|
||||
console.log('\n4. Testing structured output...');
|
||||
const structuredResult = await provider.generateObject({
|
||||
messages: [{ role: 'user', content: 'Extract: Alice is 25 years old' }],
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
name: { type: 'string' },
|
||||
age: { type: 'number' },
|
||||
},
|
||||
required: ['name', 'age'],
|
||||
additionalProperties: false,
|
||||
},
|
||||
maxTokens: 100,
|
||||
});
|
||||
console.log('Structured output:', structuredResult.object);
|
||||
|
||||
console.log('\n✅ OpenAI Provider tests passed!');
|
||||
}
|
||||
|
||||
async function testAnthropic() {
|
||||
console.log('\n=== Testing Anthropic Provider ===');
|
||||
|
||||
if (!process.env.ANTHROPIC_API_KEY) {
|
||||
console.log('ANTHROPIC_API_KEY not set, skipping');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const provider = createProvider('anthropic/claude-sonnet-4-20250514');
|
||||
|
||||
console.log('\n1. Testing generateText...');
|
||||
const result = await provider.generateText({
|
||||
messages: [
|
||||
{ role: 'user', content: 'Say hello in one word.' },
|
||||
],
|
||||
maxTokens: 10,
|
||||
});
|
||||
console.log('Response:', result.text);
|
||||
console.log('Usage:', result.usage);
|
||||
|
||||
console.log('\n✅ Anthropic Provider tests passed!');
|
||||
} catch (error: any) {
|
||||
console.log('Anthropic test failed (may be API key issue):', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
async function testGoogle() {
|
||||
console.log('\n=== Testing Google Provider ===');
|
||||
|
||||
if (!process.env.GOOGLE_API_KEY) {
|
||||
console.log('GOOGLE_API_KEY not set, skipping');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const provider = createProvider('google/gemini-2.0-flash');
|
||||
|
||||
console.log('\n1. Testing generateText...');
|
||||
const result = await provider.generateText({
|
||||
messages: [
|
||||
{ role: 'user', content: 'Say hello in one word.' },
|
||||
],
|
||||
maxTokens: 10,
|
||||
});
|
||||
console.log('Response:', result.text);
|
||||
console.log('Usage:', result.usage);
|
||||
|
||||
console.log('\n✅ Google Provider tests passed!');
|
||||
} catch (error: any) {
|
||||
console.log('Google test failed:', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
async function main() {
|
||||
console.log('Provider Verification Script');
|
||||
console.log('============================');
|
||||
|
||||
try {
|
||||
await testOpenAI();
|
||||
await testAnthropic();
|
||||
await testGoogle();
|
||||
console.log('\n🎉 All provider tests completed!');
|
||||
} catch (error) {
|
||||
console.error('\n❌ Error:', error);
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
main();
|
||||
@@ -0,0 +1,121 @@
|
||||
/**
|
||||
* OpenAI Provider Integration Tests - Real API calls
|
||||
* Run with: PRAISONAI_TEST_REAL_KEYS=1 npm test -- --testPathPattern=integration/llm
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeAll } from '@jest/globals';
|
||||
import { createProvider, OpenAIProvider } from '../../../src/llm/providers';
|
||||
|
||||
const SKIP_REAL_TESTS = !process.env.PRAISONAI_TEST_REAL_KEYS;
|
||||
const HAS_OPENAI_KEY = !!process.env.OPENAI_API_KEY;
|
||||
|
||||
describe('OpenAI Provider Integration', () => {
|
||||
beforeAll(() => {
|
||||
if (SKIP_REAL_TESTS) {
|
||||
console.log('Skipping real API tests. Set PRAISONAI_TEST_REAL_KEYS=1 to enable.');
|
||||
}
|
||||
if (!HAS_OPENAI_KEY) {
|
||||
console.log('OPENAI_API_KEY not set.');
|
||||
}
|
||||
});
|
||||
|
||||
it('should generate text', async () => {
|
||||
if (SKIP_REAL_TESTS || !HAS_OPENAI_KEY) return;
|
||||
|
||||
const provider = createProvider('openai/gpt-4o-mini');
|
||||
const result = await provider.generateText({
|
||||
messages: [
|
||||
{ role: 'system', content: 'You are a helpful assistant. Be concise.' },
|
||||
{ role: 'user', content: 'Say "Hello" and nothing else.' },
|
||||
],
|
||||
maxTokens: 10,
|
||||
temperature: 0,
|
||||
});
|
||||
|
||||
expect(result.text).toBeDefined();
|
||||
expect(result.text.toLowerCase()).toContain('hello');
|
||||
expect(result.usage.promptTokens).toBeGreaterThan(0);
|
||||
expect(result.usage.completionTokens).toBeGreaterThan(0);
|
||||
console.log('OpenAI generateText result:', result.text.substring(0, 50));
|
||||
});
|
||||
|
||||
it('should stream text', async () => {
|
||||
if (SKIP_REAL_TESTS || !HAS_OPENAI_KEY) return;
|
||||
|
||||
const provider = createProvider('openai/gpt-4o-mini');
|
||||
const tokens: string[] = [];
|
||||
|
||||
const stream = await provider.streamText({
|
||||
messages: [
|
||||
{ role: 'user', content: 'Count from 1 to 3, one number per line.' },
|
||||
],
|
||||
maxTokens: 20,
|
||||
onToken: (token) => tokens.push(token),
|
||||
});
|
||||
|
||||
let fullText = '';
|
||||
for await (const chunk of stream) {
|
||||
if (chunk.text) fullText += chunk.text;
|
||||
}
|
||||
|
||||
expect(tokens.length).toBeGreaterThan(0);
|
||||
expect(fullText).toContain('1');
|
||||
console.log('OpenAI streamText tokens:', tokens.length);
|
||||
});
|
||||
|
||||
it('should call tools', async () => {
|
||||
if (SKIP_REAL_TESTS || !HAS_OPENAI_KEY) return;
|
||||
|
||||
const provider = createProvider('openai/gpt-4o-mini');
|
||||
const result = await provider.generateText({
|
||||
messages: [
|
||||
{ role: 'user', content: 'What is 2 + 2? Use the calculator tool.' },
|
||||
],
|
||||
tools: [{
|
||||
name: 'calculator',
|
||||
description: 'Perform math calculations',
|
||||
parameters: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
expression: { type: 'string', description: 'Math expression to evaluate' },
|
||||
},
|
||||
required: ['expression'],
|
||||
},
|
||||
}],
|
||||
maxTokens: 100,
|
||||
});
|
||||
|
||||
expect(result.toolCalls).toBeDefined();
|
||||
expect(result.toolCalls!.length).toBeGreaterThan(0);
|
||||
expect(result.toolCalls![0].function.name).toBe('calculator');
|
||||
console.log('OpenAI tool call:', result.toolCalls![0].function);
|
||||
});
|
||||
|
||||
it('should generate structured output', async () => {
|
||||
if (SKIP_REAL_TESTS || !HAS_OPENAI_KEY) return;
|
||||
|
||||
const provider = createProvider('openai/gpt-4o-mini');
|
||||
const result = await provider.generateObject({
|
||||
messages: [
|
||||
{ role: 'user', content: 'Extract: John is 30 years old and lives in New York.' },
|
||||
],
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
name: { type: 'string' },
|
||||
age: { type: 'number' },
|
||||
city: { type: 'string' },
|
||||
},
|
||||
required: ['name', 'age', 'city'],
|
||||
additionalProperties: false,
|
||||
},
|
||||
maxTokens: 100,
|
||||
});
|
||||
|
||||
expect(result.object).toBeDefined();
|
||||
expect(result.object.name).toBe('John');
|
||||
expect(result.object.age).toBe(30);
|
||||
expect(result.object.city).toBe('New York');
|
||||
console.log('OpenAI structured output:', result.object);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,135 @@
|
||||
/**
|
||||
* Real Provider Integration Tests - Tests with actual API calls
|
||||
* These tests require API keys and are gated by PRAISONAI_TEST_REAL_KEYS env var
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeAll, afterAll } from '@jest/globals';
|
||||
|
||||
// Skip all tests if real keys not enabled
|
||||
const SKIP_REAL_TESTS = !process.env.PRAISONAI_TEST_REAL_KEYS;
|
||||
|
||||
// These imports will be implemented
|
||||
// import { createProvider } from '../../../src/llm/providers';
|
||||
|
||||
describe('Real Provider Integration Tests', () => {
|
||||
beforeAll(() => {
|
||||
if (SKIP_REAL_TESTS) {
|
||||
console.log('Skipping real API tests. Set PRAISONAI_TEST_REAL_KEYS=1 to enable.');
|
||||
}
|
||||
});
|
||||
|
||||
describe('OpenAI Provider', () => {
|
||||
const skipIfNoKey = !process.env.OPENAI_API_KEY;
|
||||
|
||||
it.skip('should generate text with OpenAI', async () => {
|
||||
if (SKIP_REAL_TESTS || skipIfNoKey) return;
|
||||
// const provider = createProvider('openai/gpt-4o-mini');
|
||||
// const result = await provider.generateText({
|
||||
// messages: [{ role: 'user', content: 'Say hello in one word' }],
|
||||
// maxTokens: 10,
|
||||
// });
|
||||
// expect(result.text).toBeDefined();
|
||||
// expect(result.text.length).toBeGreaterThan(0);
|
||||
// // Mask output for security
|
||||
// console.log('OpenAI response length:', result.text.length);
|
||||
});
|
||||
|
||||
it.skip('should stream text with OpenAI', async () => {
|
||||
if (SKIP_REAL_TESTS || skipIfNoKey) return;
|
||||
// const provider = createProvider('openai/gpt-4o-mini');
|
||||
// let tokenCount = 0;
|
||||
// await provider.streamText({
|
||||
// messages: [{ role: 'user', content: 'Count to 3' }],
|
||||
// maxTokens: 20,
|
||||
// onToken: () => tokenCount++,
|
||||
// });
|
||||
// expect(tokenCount).toBeGreaterThan(0);
|
||||
// console.log('OpenAI streamed tokens:', tokenCount);
|
||||
});
|
||||
|
||||
it.skip('should call tools with OpenAI', async () => {
|
||||
if (SKIP_REAL_TESTS || skipIfNoKey) return;
|
||||
// const provider = createProvider('openai/gpt-4o-mini');
|
||||
// const result = await provider.generateText({
|
||||
// messages: [{ role: 'user', content: 'What is 2+2? Use the calculator.' }],
|
||||
// tools: [{
|
||||
// name: 'calculator',
|
||||
// description: 'Perform math',
|
||||
// parameters: { type: 'object', properties: { expression: { type: 'string' } } },
|
||||
// }],
|
||||
// });
|
||||
// expect(result.toolCalls).toBeDefined();
|
||||
// expect(result.toolCalls.length).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Anthropic Provider', () => {
|
||||
const skipIfNoKey = !process.env.ANTHROPIC_API_KEY;
|
||||
|
||||
it.skip('should generate text with Anthropic', async () => {
|
||||
if (SKIP_REAL_TESTS || skipIfNoKey) return;
|
||||
// const provider = createProvider('anthropic/claude-3-5-sonnet-latest');
|
||||
// const result = await provider.generateText({
|
||||
// messages: [{ role: 'user', content: 'Say hello in one word' }],
|
||||
// maxTokens: 10,
|
||||
// });
|
||||
// expect(result.text).toBeDefined();
|
||||
// console.log('Anthropic response length:', result.text.length);
|
||||
});
|
||||
|
||||
it.skip('should stream text with Anthropic', async () => {
|
||||
if (SKIP_REAL_TESTS || skipIfNoKey) return;
|
||||
// const provider = createProvider('anthropic/claude-3-5-sonnet-latest');
|
||||
// let tokenCount = 0;
|
||||
// await provider.streamText({
|
||||
// messages: [{ role: 'user', content: 'Count to 3' }],
|
||||
// maxTokens: 20,
|
||||
// onToken: () => tokenCount++,
|
||||
// });
|
||||
// expect(tokenCount).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Google Provider', () => {
|
||||
const skipIfNoKey = !process.env.GOOGLE_API_KEY;
|
||||
|
||||
it.skip('should generate text with Google', async () => {
|
||||
if (SKIP_REAL_TESTS || skipIfNoKey) return;
|
||||
// const provider = createProvider('google/gemini-2.0-flash');
|
||||
// const result = await provider.generateText({
|
||||
// messages: [{ role: 'user', content: 'Say hello in one word' }],
|
||||
// maxTokens: 10,
|
||||
// });
|
||||
// expect(result.text).toBeDefined();
|
||||
// console.log('Google response length:', result.text.length);
|
||||
});
|
||||
|
||||
it.skip('should stream text with Google', async () => {
|
||||
if (SKIP_REAL_TESTS || skipIfNoKey) return;
|
||||
// const provider = createProvider('google/gemini-2.0-flash');
|
||||
// let tokenCount = 0;
|
||||
// await provider.streamText({
|
||||
// messages: [{ role: 'user', content: 'Count to 3' }],
|
||||
// maxTokens: 20,
|
||||
// onToken: () => tokenCount++,
|
||||
// });
|
||||
// expect(tokenCount).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Token Usage Tracking', () => {
|
||||
it.skip('should track token usage with real API', async () => {
|
||||
if (SKIP_REAL_TESTS || !process.env.OPENAI_API_KEY) return;
|
||||
// const provider = createProvider('openai/gpt-4o-mini');
|
||||
// const result = await provider.generateText({
|
||||
// messages: [{ role: 'user', content: 'Hello' }],
|
||||
// });
|
||||
// expect(result.usage.promptTokens).toBeGreaterThan(0);
|
||||
// expect(result.usage.completionTokens).toBeGreaterThan(0);
|
||||
// console.log('Token usage:', {
|
||||
// prompt: result.usage.promptTokens,
|
||||
// completion: result.usage.completionTokens,
|
||||
// });
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,116 +0,0 @@
|
||||
import { Agent } from '../../../../src/agent/agent';
|
||||
import { Task, Tool } from '../../../../src/types';
|
||||
|
||||
jest.mock('../../../../src/llm/openai');
|
||||
|
||||
describe('Agent', () => {
|
||||
let agent: Agent;
|
||||
|
||||
beforeEach(() => {
|
||||
agent = new Agent();
|
||||
});
|
||||
|
||||
describe('configuration', () => {
|
||||
it('should configure agent with tools', () => {
|
||||
const tool: Tool = {
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
execute: async () => 'test result'
|
||||
};
|
||||
|
||||
agent.configure({ tools: [tool] });
|
||||
const tools = agent.getTools();
|
||||
expect(tools).toHaveLength(1);
|
||||
expect(tools[0].name).toBe('test-tool');
|
||||
});
|
||||
});
|
||||
|
||||
describe('task execution', () => {
|
||||
it('should execute simple task', async () => {
|
||||
const task: Task = {
|
||||
id: 'test-task',
|
||||
name: 'Test Task',
|
||||
description: 'A test task',
|
||||
config: {
|
||||
priority: 1
|
||||
},
|
||||
expected_output: 'test output',
|
||||
agent: 'test-agent',
|
||||
dependencies: [],
|
||||
result: null
|
||||
};
|
||||
|
||||
const result = await agent.execute(task);
|
||||
expect(result).toBeDefined();
|
||||
expect(result.id).toBe('test-task');
|
||||
});
|
||||
|
||||
it('should execute complex task', async () => {
|
||||
const task: Task = {
|
||||
id: 'complex-task',
|
||||
name: 'Complex Task',
|
||||
description: 'A complex task',
|
||||
config: {
|
||||
priority: 1
|
||||
},
|
||||
expected_output: 'test output',
|
||||
agent: 'test-agent',
|
||||
dependencies: [
|
||||
{
|
||||
id: 'dep-task',
|
||||
name: 'Dependency Task',
|
||||
description: 'A dependency task',
|
||||
config: {
|
||||
priority: 1
|
||||
},
|
||||
expected_output: 'dependency output',
|
||||
agent: 'test-agent',
|
||||
dependencies: [],
|
||||
result: 'dependency result'
|
||||
}
|
||||
],
|
||||
result: null
|
||||
};
|
||||
|
||||
const result = await agent.execute(task);
|
||||
expect(result).toBeDefined();
|
||||
expect(result.id).toBe('complex-task');
|
||||
});
|
||||
|
||||
it('should handle task failure', async () => {
|
||||
const task: Task = {
|
||||
id: 'failing-task',
|
||||
name: 'Failing Task',
|
||||
description: 'A task that fails',
|
||||
config: {
|
||||
priority: 1,
|
||||
shouldFail: true
|
||||
},
|
||||
expected_output: 'test output',
|
||||
agent: 'test-agent',
|
||||
dependencies: [],
|
||||
result: null
|
||||
};
|
||||
|
||||
await expect(agent.execute(task)).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('tool management', () => {
|
||||
it('should execute tool', async () => {
|
||||
const tool: Tool = {
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
execute: async () => 'test result'
|
||||
};
|
||||
|
||||
agent.configure({ tools: [tool] });
|
||||
const result = await agent.executeTool('test-tool', 'test input');
|
||||
expect(result).toBe('test result');
|
||||
});
|
||||
|
||||
it('should handle unknown tool', async () => {
|
||||
await expect(agent.executeTool('unknown-tool', 'test input')).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
});
|
||||
Loaded 100 of 180 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user