Compare commits

...
74 Commits
Author SHA1 Message Date
Mervin Praison 67dbd1007a Release v2.3.10 2025-12-16 09:54:57 +00:00
Mervin Praison cf2a5a135f workflow 2025-12-16 09:23:49 +00:00
Mervin Praison bc23d3a104 Release v2.3.9 2025-12-16 08:18:09 +00:00
Mervin Praison bb4e1d8b9e Release v2.3.8 2025-12-16 07:55:55 +00:00
Mervin Praison 69740baa7c Release v2.3.7 2025-12-16 07:46:40 +00:00
Mervin Praison 5035978a29 fix: resolve Mintlify parsing errors
- Fix corrupted/duplicated database service in docker-compose.dev.yml
- Fix MDX parsing error in tests/README.md (<1s -> under 1s)
2025-12-15 23:11:04 +00:00
Mervin Praison 4ddb043324 Release v2.3.7 2025-12-15 22:49:12 +00:00
Mervin Praison 337014f4db planning example 2025-12-15 22:25:09 +00:00
Mervin Praison 35c371b25e update 2025-12-15 22:22:42 +00:00
Mervin Praison bba75d9fdd Release v2.3.6 2025-12-15 22:16:57 +00:00
Mervin Praison bf8cd488da feat: add planning tools and reasoning mode for research-enabled plan creation
- Add planning_tools and planning_reasoning parameters to PraisonAIAgents
- Add RESEARCH_TOOLS list for safe planning research tools (web_search, read_file, etc.)
- Add tools, reasoning, and max_reasoning_steps parameters to PlanningAgent
- Add _get_agent() method to PlanningAgent for tool-enabled planning with Agent.chat()
- Update plan creation to use Agent with tools when available, fallback to bare LLM
- Add Planning
2025-12-15 22:07:59 +00:00
Mervin Praison 28905db46d g 2025-12-15 21:27:32 +00:00
Mervin Praison 3e7aa44dfa docs: simplify feature descriptions by removing tool comparisons
- Remove references to Cursor, Windsurf, Claude Code, Gemini CLI, Codex CLI, and Cascade from feature descriptions
- Keep feature descriptions focused on PraisonAI's own capabilities
- Update Planning Mode, Agent Memory, Rules & Instructions, Auto-Generated Memories, Workflows, and Hooks sections
2025-12-15 21:26:04 +00:00
Mervin Praison 2f7a65ab92 update 2025-12-15 21:25:41 +00:00
Mervin Praison e6e2379b02 update 2025-12-15 21:25:29 +00:00
Mervin Praison f31f9a04aa update 2025-12-15 21:23:46 +00:00
Mervin Praison 5a8adf295e Release v2.3.5 2025-12-15 21:04:45 +00:00
Mervin Praison 8dc243fd4f feat: Add Planning Mode - plan before execution like Cursor, Windsurf, Claude Code
- Add planning module with PlanningAgent, Plan, PlanStep, TodoList, TodoItem
- Add PlanStorage for persisting plans to .praison/plans/
- Add ApprovalCallback for plan approval flow
- Add plan_mode to Agent for read-only tool filtering
- Add planning, planning_llm, auto_approve_plan params to PraisonAIAgents
- Add 71 comprehensive tests for planning mode
- Update READMEs with Planning Mode documentation
- Bump version to 0.0.181
2025-12-15 17:51:08 +00:00
Mervin Praison 548b4afb88 Release v2.3.4 2025-12-15 16:39:58 +00:00
Mervin Praison ffa938b3df docs: add advanced rules features and new memory capabilities
- Add local override files support (PRAISON.local.md, CLAUDE.local.md) with higher priority
- Add @import syntax for including other files in rules (like Claude Code)
- Add git root discovery for monorepo support with fallback to subprocess
- Add auto-discovery of modular rules directories (.claude/rules/, .windsurf/rules/, .cursor/rules/)
- Add character limits (12000) for context window management with truncation
- Add AutoMemory for
2025-12-15 16:39:04 +00:00
Mervin Praison 8a9beed264 docs: add advanced memory features and rules management system
- Add session save/resume, context compression, and checkpointing (inspired by Gemini CLI)
- Add RulesManager for auto-discovery of instruction files (CLAUDE.md, AGENTS.md, GEMINI.md, .cursorrules, .windsurfrules)
- Add slash commands for memory management (/memory show, /memory save, /memory compress, etc.)
- Add rules context injection into agent system prompts with activation modes (always, glob, manual, ai_decision)
- Add program
2025-12-15 16:17:25 +00:00
Mervin Praison daca5a1762 test: add multi-agent memory sharing and integration tests
- Add TestMultiAgentMemorySharing class with tests for shared/isolated memory between agents
- Test agents with same user_id share memory (short_term, long_term, entities)
- Test agents with different user_id have isolated memory
- Test memory context injection into system prompt
- Test memory display info shows correct counts
- Add TestMultiAgentMemoryIntegration class for LLM integration tests
- Test multi-agent memory sharing with OpenA
2025-12-15 15:00:16 +00:00
Mervin Praison b183ac5cbd Release v2.3.3 2025-12-15 14:55:30 +00:00
Mervin Praison 5f57e1d2f1 feat: add memory info display and improve memory integration with system prompts
- Add _display_memory_info() method to show memory stats in rich panel (long-term, short-term, entities, storage path)
- Disable system prompt caching when memory is enabled (dynamic context)
- Inject memory context into system prompt with "## Memory" section
- Display memory info on first chat when verbose=True
- Add integration tests for memory with OpenAI API (injection and persistence)
- Update README with zero
2025-12-15 14:52:09 +00:00
Mervin Praison 5114165944 feat: add zero-dependency file-based memory system for all agents
- Add FileMemory class with JSON-based persistent storage (no extra dependencies)
- Add _init_memory(), get_memory_context(), store_memory() methods to Agent class
- Support memory types: short_term, long_term, entity, episodic
- Add memory=True/False/"file"/"sqlite"/"chromadb"/"mem0" configuration options
- Add lazy import for Memory class to avoid dependency issues
- Update README with memory usage examples and provider comparison
2025-12-15 14:37:33 +00:00
Mervin Praison 389442ea04 feat: add Claude Memory Tool support for persistent cross-conversation memory
- Add claude_memory parameter to Agent and LLM classes
- Add ClaudeMemoryTool with file-based persistent memory in .praison/claude_memory/
- Add _supports_claude_memory(), _is_memory_tool_call(), _execute_memory_tool_call() methods
- Add automatic beta header injection for Anthropic models
- Add memory tool definition to completion params when enabled
- Add ClaudeMemoryTool to TOOL_MAPPINGS with class-level import
- Update
2025-12-15 13:48:35 +00:00
Mervin Praison 6d7cc9c3ba Release v2.3.2 2025-12-15 13:34:15 +00:00
Mervin Praison 690e4497cc feat: add --save support for direct prompts
- Save output to output/prompts/ folder with timestamped filename
- Works with all flags: --web-search, --expand-prompt, etc.

Usage:
  praisonai 'latest AI trends' --web-search --save
  praisonai 'write a poem' --expand-prompt --save
2025-12-15 13:33:28 +00:00
Mervin Praison 041f5f0b08 fix: add current date to PromptExpanderAgent for time-sensitive context
- Add datetime import and current_date to all expansion prompts
- Ensures expanded prompts reference the correct current date (e.g., December 15, 2025)
- Fixes issue where prompts defaulted to outdated year references
2025-12-15 13:26:13 +00:00
Mervin Praison bc0d9a6213 feat: add CLI flags for web-search, web-fetch, and prompt-caching
- Add --web-search flag for native web search (OpenAI, Gemini, Anthropic, xAI, Perplexity)
- Add --web-fetch flag for URL content retrieval (Anthropic only)
- Add --prompt-caching flag to reduce costs (OpenAI, Anthropic, Bedrock, Deepseek)
- Update README with CLI and programmatic usage examples

Usage:
  praisonai 'latest AI news' --web-search --llm openai/gpt-4o-search-preview
  praisonai 'summarize https://example.com' --web-fetch --llm anthropic/claude-sonnet-4-20250514
  praisonai 'analyze...' --prompt-caching --llm anthropic/claude-sonnet-4-20250514
2025-12-15 12:56:02 +00:00
Mervin Praison a4427fa135 Release v2.3.1 2025-12-15 12:45:22 +00:00
Mervin Praison 51e8c65c8c feat: refactor model capabilities to use LiteLLM helper functions
- Replace static model lists with LiteLLM's built-in capability checks (supports_response_schema, supports_function_calling, supports_parallel_function_calling, supports_web_search)
- Add supports_function_calling() and supports_parallel_function_calling() functions
- Keep static list for web_fetch (no LiteLLM helper available yet)
- Update documentation to reference LiteLLM as primary source
- Add Web Fetch feature to README features
2025-12-15 12:32:09 +00:00
Mervin Praison 285ed672c5 feat: add native web search and web fetch documentation to README
- Add Native Web Search section with supported providers (OpenAI, Anthropic, Gemini, xAI, Perplexity)
- Add Web Fetch section for Anthropic models with configuration options
- Add code examples for web_search and web_fetch parameters
- Document web_fetch options (max_uses, allowed_domains, citations)
2025-12-15 12:20:41 +00:00
Mervin Praison c03e51e337 feat: add native web fetch support for Anthropic Claude models
- Add web_fetch parameter to Agent and LLM classes
- Add model_capabilities.py with MODELS_SUPPORTING_WEB_FETCH list (Claude 3.5+, 4+)
- Add supports_web_fetch() function to check model capabilities
- Add web_fetch tool definition with configurable options (max_uses, allowed_domains, blocked_domains, citations, max_content_tokens)
- Add _model_supports_web_fetch() method to Agent class
- Add warning log for unsupported models when web_
2025-12-15 12:18:08 +00:00
Mervin Praison f3ba05f6aa feat: expand structured outputs support for OpenAI, Anthropic, and Gemini models
- Add comprehensive list of OpenAI models with structured outputs (gpt-4o-2024-08-06+, gpt-4.1+, gpt-5+, o-series)
- Add Anthropic Claude 4+ models with structured outputs support (Sonnet 4.5, Opus 4.1/4.5, Haiku 4.5)
- Add all Google Gemini models (2.5-pro, 2.5-flash, 2.0-flash, 1.5-pro/flash)
- Add auto-detection for newer model versions (gpt-4.1+, gpt-5+, o-series, Claude 4+, all Gemini)
- Add azure/ prefix support
2025-12-15 12:07:46 +00:00
Mervin Praison 9cea3edcd6 feat: add native web search support for OpenAI, Gemini, Anthropic, xAI, and Perplexity models
- Add web_search parameter to Agent and LLM classes
- Add model_capabilities.py with MODELS_SUPPORTING_WEB_SEARCH list
- Add supports_web_search() function to check model capabilities
- Add automatic DuckDuckGo fallback for unsupported models
- Add web_search_options to LiteLLM completion params
- Remove incompatible parameters (temperature, top_p, etc.) for web search models
- Update README with Native
2025-12-15 11:59:19 +00:00
Mervin Praison 02a3bd3597 chore: add uv cache clean before uv lock in release.py 2025-12-15 11:26:08 +00:00
Mervin Praison 128c77926d chore: add uv cache clean before uv lock in bump_and_release.py 2025-12-15 11:25:42 +00:00
Mervin Praison 3aa1b825e4 Release v2.3.0 2025-12-15 11:24:46 +00:00
Mervin Praison b4ee30d837 chore: bump praisonaiagents version to 0.0.170 2025-12-15 11:11:46 +00:00
Mervin Praison 825648aead docs: add PromptExpanderAgent to README 2025-12-15 11:11:04 +00:00
Mervin Praison 1532c31e72 feat: add PromptExpanderAgent for detailed prompt expansion
- Add PromptExpanderAgent class with expansion strategies (BASIC, DETAILED, STRUCTURED, CREATIVE, AUTO)
- Add ExpandStrategy enum and ExpandResult dataclass
- Add CLI flags --expand-prompt and --expand-tools
- Add 20 unit tests for PromptExpanderAgent
- Export from praisonaiagents package

Usage:
  praisonai 'write a movie script' --expand-prompt
  praisonai 'AI blog' --expand-prompt --expand-tools tools.py
2025-12-15 11:09:25 +00:00
Mervin Praison b313fefd86 fix: QueryRewriterAgent preserves task intent instead of converting to questions 2025-12-15 10:57:53 +00:00
Mervin Praison 435e1afb55 test: add OPENAI_API_KEY environment variable to comprehensive test workflow and verify its presence 2025-12-15 10:47:29 +00:00
Mervin Praison 4cb3c0565a test: add skipif decorator for OPENAI_API_KEY check in CrewAI environment test 2025-12-15 10:42:00 +00:00
Mervin Praison a7634b66dc Release v2.2.99 2025-12-15 10:36:34 +00:00
Mervin Praison 3f0eaccbcb Release v2.2.98 2025-12-15 10:28:42 +00:00
Mervin Praison 47713d933a test: convert test assertions to pytest format and add TaskOutput descriptions
- Replace return True/False with pytest assertions in test_context_agent.py
- Add pytest.mark.skip decorators for unimplemented ContextAgent methods
- Add missing description parameter to TaskOutput instances in test_validation_feedback.py
2025-12-15 10:25:23 +00:00
Mervin Praison c82bfccd1d docs: add custom tools documentation with decorator, BaseTool class, and pip installable package examples 2025-12-15 10:20:11 +00:00
Mervin Praison ef7d1264e2 feat: add plugin system with tool registry and BaseTool support for external tool packages 2025-12-15 10:10:06 +00:00
Mervin Praison 2799304270 Release v2.2.97 2025-12-15 08:25:56 +00:00
Mervin Praison 4a85c135a6 Release v2.2.96 2025-12-15 07:57:25 +00:00
Mervin Praison cdb42a3f08 feat: add custom tools support to research command with file and built-in tool loading 2025-12-15 07:33:52 +00:00
Mervin Praison c8378118e7 docs: add query rewrite flag to research command and bump version to 2.2.94 2025-12-13 12:49:38 +00:00
Mervin Praison 0213764057 feat: add Query Rewriter Agent with multiple rewrite strategies and bump version to 0.0.166 2025-12-13 12:44:21 +00:00
Mervin Praison 2220114809 feat: add --save flag to research command for markdown output and bump version to 2.2.93 2025-12-13 11:58:36 +00:00
Mervin Praison 9fe4dbc172 chore: bump version to 2.2.92 and update praisonaiagents dependency to 0.0.165 2025-12-13 02:19:22 +00:00
Mervin Praison 180503fcd5 feat: add default research instructions to DeepResearchAgent and bump version to 0.0.165 2025-12-13 02:18:42 +00:00
Mervin Praison 0d8f936b19 chore: bump version to 2.2.91 and update praisonaiagents dependency to 0.0.164 2025-12-13 02:09:12 +00:00
Mervin Praison 8b53ebf5e3 docs: add Deep Research CLI documentation and bump version to 2.2.90 2025-12-13 01:44:52 +00:00
Mervin Praison 504c5a9759 chore: bump version to 2.2.89 and update praisonaiagents dependency to 0.0.163 2025-12-13 01:17:52 +00:00
Mervin Praison 8fef42fc03 feat: add Deep Research Agent support with OpenAI and Gemini integration 2025-12-13 01:15:07 +00:00
Mervin Praison 2f18cc1ea2 Merge branch 'main' of https://github.com/MervinPraison/PraisonAI 2025-12-13 01:13:29 +00:00
Mervin Praison ff7969cdbf docs: rewrite README from MCP SSE implementation to comprehensive PraisonAI Agents documentation 2025-12-13 01:13:28 +00:00
Mervin Praison 5fa84dc3eb Merge pull request #1084 from Dhivya-Bharathy/feature/agentic-rag
feat: Add Agentic RAG Example to Cookbooks
2025-11-03 14:14:46 +00:00
Mervin Praison 4bf8b980ed Merge pull request #1085 from Dhivya-Bharathy/feat/agentic-rag-tools
feat: Add Agentic RAG Tools Implementation
2025-11-03 14:14:24 +00:00
Mervin Praison 53e0fd19ab feat: add async run support and example scripts for multi-turn conversations 2025-08-30 11:54:59 +01:00
Dhivya-Bharathy 0cb6307d3c Add agentic_rag_gpt5 example with GPT-5 integration 2025-08-29 13:11:08 +05:30
Dhivya adce94410b Add files via upload 2025-08-29 13:01:08 +05:30
Dhivya-Bharathy 42aec561b1 Add agentic RAG Python implementation to tools 2025-08-28 14:09:09 +05:30
Dhivya acc805768d Update PraisonAI_RecipeAgent.ipynb 2025-08-28 13:00:50 +05:30
Dhivya-Bharathy 6d9fa167e3 Add agentic_rag example to cookbooks 2025-08-28 06:43:41 +00:00
Mervin Praison 9e178b70e6 chore: bump PraisonAI to v2.2.88 and update dependencies across Dockerfiles and project files 2025-08-08 19:24:43 +01:00
Mervin Praison ed406973cd refactor: optimise ChatOpenAI instantiation by using a dictionary for parameters and enforce temperature for gpt-5 models 2025-08-08 19:17:45 +01:00
134 changed files with 33589 additions and 494 deletions

No files matched your search

+38 -22
View File
@@ -86,13 +86,29 @@ runs:
- name: Extract GitHub Context and Create Prompt
shell: bash
id: prepare_context
env:
# Pass user-controlled inputs via environment variables to prevent script injection (GHSL-2025-093)
EVENT_COMMENT_BODY: ${{ github.event.comment.body }}
EVENT_ISSUE_TITLE: ${{ github.event.issue.title }}
EVENT_ISSUE_BODY: ${{ github.event.issue.body }}
EVENT_ISSUE_NUMBER: ${{ github.event.issue.number }}
EVENT_PR_TITLE: ${{ github.event.pull_request.title }}
EVENT_PR_NUMBER: ${{ github.event.pull_request.number }}
EVENT_REVIEW_BODY: ${{ github.event.review.body }}
EVENT_ASSIGNEE_LOGIN: ${{ github.event.assignee.login }}
INPUT_TRIGGER_PHRASE: ${{ inputs.trigger_phrase }}
INPUT_ASSIGNEE_TRIGGER: ${{ inputs.assignee_trigger }}
INPUT_DIRECT_PROMPT: ${{ inputs.direct_prompt }}
EVENT_ACTION: ${{ github.event.action }}
EVENT_NAME: ${{ github.event_name }}
REPO_NAME: ${{ github.repository }}
run: |
echo "🔍 Extracting GitHub context from event: ${{ github.event_name }}"
echo "🔍 Extracting GitHub context from event: $EVENT_NAME"
# Function to check for trigger phrase
check_trigger() {
local text="$1"
local trigger="${{ inputs.trigger_phrase }}"
local trigger="$INPUT_TRIGGER_PHRASE"
if [[ "$text" == *"$trigger"* ]]; then
return 0
fi
@@ -104,11 +120,11 @@ runs:
USER_REQUEST=""
CONTEXT_INFO=""
case "${{ github.event_name }}" in
case "$EVENT_NAME" in
"issue_comment")
COMMENT_BODY="${{ github.event.comment.body }}"
ISSUE_TITLE="${{ github.event.issue.title }}"
ISSUE_NUMBER="${{ github.event.issue.number }}"
COMMENT_BODY="$EVENT_COMMENT_BODY"
ISSUE_TITLE="$EVENT_ISSUE_TITLE"
ISSUE_NUMBER="$EVENT_ISSUE_NUMBER"
if check_trigger "$COMMENT_BODY"; then
TRIGGER_FOUND="true"
@@ -118,9 +134,9 @@ runs:
;;
"pull_request_review_comment")
COMMENT_BODY="${{ github.event.comment.body }}"
PR_TITLE="${{ github.event.pull_request.title }}"
PR_NUMBER="${{ github.event.pull_request.number }}"
COMMENT_BODY="$EVENT_COMMENT_BODY"
PR_TITLE="$EVENT_PR_TITLE"
PR_NUMBER="$EVENT_PR_NUMBER"
if check_trigger "$COMMENT_BODY"; then
TRIGGER_FOUND="true"
@@ -130,9 +146,9 @@ runs:
;;
"pull_request_review")
REVIEW_BODY="${{ github.event.review.body }}"
PR_TITLE="${{ github.event.pull_request.title }}"
PR_NUMBER="${{ github.event.pull_request.number }}"
REVIEW_BODY="$EVENT_REVIEW_BODY"
PR_TITLE="$EVENT_PR_TITLE"
PR_NUMBER="$EVENT_PR_NUMBER"
if check_trigger "$REVIEW_BODY"; then
TRIGGER_FOUND="true"
@@ -142,17 +158,17 @@ runs:
;;
"issues")
ISSUE_BODY="${{ github.event.issue.body }}"
ISSUE_TITLE="${{ github.event.issue.title }}"
ISSUE_NUMBER="${{ github.event.issue.number }}"
ISSUE_BODY="$EVENT_ISSUE_BODY"
ISSUE_TITLE="$EVENT_ISSUE_TITLE"
ISSUE_NUMBER="$EVENT_ISSUE_NUMBER"
if check_trigger "$ISSUE_TITLE" || check_trigger "$ISSUE_BODY"; then
TRIGGER_FOUND="true"
USER_REQUEST="$ISSUE_BODY"
CONTEXT_INFO="Issue #$ISSUE_NUMBER: $ISSUE_TITLE"
elif [[ "${{ github.event.action }}" == "assigned" && -n "${{ inputs.assignee_trigger }}" ]]; then
ASSIGNEE="${{ github.event.assignee.login }}"
if [[ "$ASSIGNEE" == "${{ inputs.assignee_trigger }}" ]]; then
elif [[ "$EVENT_ACTION" == "assigned" && -n "$INPUT_ASSIGNEE_TRIGGER" ]]; then
ASSIGNEE="$EVENT_ASSIGNEE_LOGIN"
if [[ "$ASSIGNEE" == "$INPUT_ASSIGNEE_TRIGGER" ]]; then
TRIGGER_FOUND="true"
USER_REQUEST="$ISSUE_BODY"
CONTEXT_INFO="Issue #$ISSUE_NUMBER assigned to $ASSIGNEE: $ISSUE_TITLE"
@@ -162,9 +178,9 @@ runs:
esac
# Check for direct prompt override
if [[ -n "${{ inputs.direct_prompt }}" ]]; then
if [[ -n "$INPUT_DIRECT_PROMPT" ]]; then
TRIGGER_FOUND="true"
USER_REQUEST="${{ inputs.direct_prompt }}"
USER_REQUEST="$INPUT_DIRECT_PROMPT"
CONTEXT_INFO="Automated GitHub workflow"
fi
@@ -181,9 +197,9 @@ runs:
cat > /tmp/claude-action/github-context-prompt.txt << EOF
You are Claude Code, an AI assistant helping with GitHub workflows and code.
Repository: ${{ github.repository }}
Repository: $REPO_NAME
Context: $CONTEXT_INFO
Event: ${{ github.event_name }}
Event: $EVENT_NAME
User Request:
$USER_REQUEST
+10 -6
View File
@@ -46,12 +46,16 @@ jobs:
id: issue_data
env:
GH_TOKEN: ${{ steps.generate_token.outputs.token }}
# Pass user-controlled inputs via environment variables to prevent script injection (GHSL-2025-093)
EVENT_NAME: ${{ github.event_name }}
EVENT_ISSUE_NUMBER: ${{ github.event.issue.number }}
EVENT_COMMENT_BODY: ${{ github.event.comment.body }}
run: |
# Get issue number from event or context
if [ "${{ github.event_name }}" = "issue_comment" ]; then
ISSUE_NUMBER="${{ github.event.issue.number }}"
elif [ "${{ github.event_name }}" = "issues" ]; then
ISSUE_NUMBER="${{ github.event.issue.number }}"
if [ "$EVENT_NAME" = "issue_comment" ]; then
ISSUE_NUMBER="$EVENT_ISSUE_NUMBER"
elif [ "$EVENT_NAME" = "issues" ]; then
ISSUE_NUMBER="$EVENT_ISSUE_NUMBER"
else
echo "Unable to determine issue number"
exit 1
@@ -62,8 +66,8 @@ jobs:
# Extract additional instructions from comment if triggered by comment
ADDITIONAL_INSTRUCTIONS=""
if [ "${{ github.event_name }}" = "issue_comment" ]; then
COMMENT_BODY="${{ github.event.comment.body }}"
if [ "$EVENT_NAME" = "issue_comment" ]; then
COMMENT_BODY="$EVENT_COMMENT_BODY"
ADDITIONAL_INSTRUCTIONS=$(echo "$COMMENT_BODY" | sed -n 's/.*@gemini[[:space:]]*\(.*\)/\1/p' | head -1)
fi
+13 -7
View File
@@ -55,21 +55,27 @@ jobs:
id: get_pr
env:
GITHUB_TOKEN: ${{ steps.generate_token.outputs.token }}
# Pass user-controlled inputs via environment variables to prevent script injection (GHSL-2025-093)
EVENT_NAME: ${{ github.event_name }}
EVENT_INPUTS_PR_NUMBER: ${{ github.event.inputs.pr_number }}
EVENT_ISSUE_NUMBER: ${{ github.event.issue.number }}
EVENT_PR_NUMBER: ${{ github.event.pull_request.number }}
EVENT_COMMENT_BODY: ${{ github.event.comment.body }}
run: |
if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
PR_NUMBER=${{ github.event.inputs.pr_number }}
elif [ "${{ github.event_name }}" = "issue_comment" ]; then
PR_NUMBER=${{ github.event.issue.number }}
if [ "$EVENT_NAME" = "workflow_dispatch" ]; then
PR_NUMBER=$EVENT_INPUTS_PR_NUMBER
elif [ "$EVENT_NAME" = "issue_comment" ]; then
PR_NUMBER=$EVENT_ISSUE_NUMBER
else
PR_NUMBER=${{ github.event.pull_request.number }}
PR_NUMBER=$EVENT_PR_NUMBER
fi
echo "pr_number=$PR_NUMBER" >> "$GITHUB_OUTPUT"
# Extract additional instructions from comment (if triggered by comment)
ADDITIONAL_INSTRUCTIONS=""
if [ "${{ github.event_name }}" = "issue_comment" ]; then
COMMENT_BODY="${{ github.event.comment.body }}"
if [ "$EVENT_NAME" = "issue_comment" ]; then
COMMENT_BODY="$EVENT_COMMENT_BODY"
ADDITIONAL_INSTRUCTIONS=$(echo "$COMMENT_BODY" | sed 's/.*@gemini//' | xargs)
fi
echo "additional_instructions=$ADDITIONAL_INSTRUCTIONS" >> "$GITHUB_OUTPUT"
+9 -4
View File
@@ -27,6 +27,9 @@ jobs:
comprehensive-test:
runs-on: ubuntu-latest
timeout-minutes: 15
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
OPENAI_MODEL_NAME: gpt-4o-mini
strategy:
matrix:
python-version: ["3.11"]
@@ -55,11 +58,13 @@ jobs:
# Install knowledge dependencies from praisonai-agents
uv pip install --system "praisonaiagents[knowledge]"
- name: Set environment variables
- name: Verify environment variables
run: |
echo "OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY || 'sk-test-key-for-github-actions-testing-only-not-real' }}" >> $GITHUB_ENV
echo "OPENAI_API_BASE=${{ secrets.OPENAI_API_BASE || 'https://api.openai.com/v1' }}" >> $GITHUB_ENV
echo "OPENAI_MODEL_NAME=${{ secrets.OPENAI_MODEL_NAME || 'gpt-5-nano' }}" >> $GITHUB_ENV
if [ -n "$OPENAI_API_KEY" ]; then
echo "✅ OPENAI_API_KEY is set (length: ${#OPENAI_API_KEY})"
else
echo "⚠️ OPENAI_API_KEY is NOT set"
fi
- name: Run Comprehensive Test Suite
run: |
+2 -2
View File
@@ -34,7 +34,7 @@ jobs:
cd src/praisonai
uv pip install --system ."[ui,gradio,api,agentops,google,openai,anthropic,cohere,chat,code,realtime,call,crewai,autogen]"
uv pip install --system duckduckgo_search
uv pip install --system pytest pytest-asyncio pytest-cov
uv pip install --system pytest pytest-asyncio pytest-cov pytest-timeout
# Install knowledge dependencies from praisonai-agents
uv pip install --system "praisonaiagents[knowledge]"
@@ -262,7 +262,7 @@ jobs:
export OPENAI_API_KEY="$OPENAI_API_KEY"
export OPENAI_API_BASE="$OPENAI_API_BASE"
export OPENAI_MODEL_NAME="$OPENAI_MODEL_NAME"
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings -m "not real" --timeout=60
- name: Upload coverage reports to Codecov
if: matrix.python-version == '3.11'
+2 -2
View File
@@ -27,7 +27,7 @@ jobs:
cd src/praisonai
uv pip install --system ."[ui,gradio,api,agentops,google,openai,anthropic,cohere,chat,code,realtime,call,crewai,autogen]"
uv pip install --system duckduckgo_search
uv pip install --system pytest pytest-asyncio pytest-cov
uv pip install --system pytest pytest-asyncio pytest-cov pytest-timeout
# Install knowledge dependencies from praisonai-agents
uv pip install --system "praisonaiagents[knowledge]"
@@ -489,7 +489,7 @@ jobs:
- name: Run Legacy Example Tests
run: |
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings -m "not real" --timeout=60
continue-on-error: true
- name: Upload coverage reports to Codecov
+589 -1
View File
@@ -53,6 +53,21 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
- 🛠️ 100+ Custom Tools
- 📄 YAML Configuration
- 💯 100+ LLM Support
- 🔬 **Deep Research Agents** (OpenAI & Gemini)
- 🔄 **Query Rewriter Agent** (HyDE, Step-back, Multi-query)
- 🌐 **Native Web Search** (OpenAI, Gemini, Anthropic, xAI, Perplexity)
- 📥 **Web Fetch** (Retrieve full content from URLs - Anthropic)
- 💾 **Prompt Caching** (Reduce costs & latency - OpenAI, Anthropic, Bedrock, Deepseek)
- 🧠 **Claude Memory Tool** (Persistent cross-conversation memory - Anthropic Beta)
- 💾 **File-Based Memory** (Zero-dependency persistent memory for all agents)
- 🔍 **Built-in Search Tools** (Tavily, You.com, Exa - web search, news, content extraction)
- 📋 **Planning Mode** (Plan before execution - Agent & Multi-Agent)
- 🔧 **Planning Tools** (Research with tools during planning)
- 🧠 **Planning Reasoning** (Chain-of-thought planning)
- 🔌 **MCP Transports** (stdio, Streamable HTTP, WebSocket, SSE - Protocol 2025-11-25)
- 🌐 **WebSocket MCP** (Real-time bidirectional connections with auto-reconnect)
- 🔐 **MCP Security** (Origin validation, DNS rebinding prevention, secure sessions)
- 🔄 **MCP Resumability** (SSE stream recovery via Last-Event-ID)
## Using Python Code
@@ -96,6 +111,260 @@ Run:
python app.py
```
### 3. Agent with Planning Mode
Enable planning for any agent - the agent creates a plan, then executes step by step:
```python
from praisonaiagents import Agent
def search_web(query: str) -> str:
return f"Search results for: {query}"
agent = Agent(
name="AI Assistant",
instructions="Research and write about topics",
planning=True, # Enable planning mode
planning_tools=[search_web], # Tools for planning research
planning_reasoning=True # Chain-of-thought reasoning
)
result = agent.start("Research AI trends in 2025 and write a summary")
```
**What happens:**
1. 📋 Agent creates a multi-step plan
2. 🚀 Executes each step sequentially
3. 📊 Shows progress with context passing
4. ✅ Returns final result
### 4. Deep Research Agent
Automated research with real-time streaming, web search, and citations using OpenAI or Gemini Deep Research APIs.
```python
from praisonaiagents import DeepResearchAgent
# OpenAI Deep Research
agent = DeepResearchAgent(
model="o4-mini-deep-research", # or "o3-deep-research"
verbose=True
)
result = agent.research("What are the latest AI trends in 2025?")
print(result.report)
print(f"Citations: {len(result.citations)}")
```
```python
# Gemini Deep Research
from praisonaiagents import DeepResearchAgent
agent = DeepResearchAgent(
model="deep-research-pro", # Auto-detected as Gemini
verbose=True
)
result = agent.research("Research quantum computing advances")
print(result.report)
```
**Features:**
- 🔍 Multi-provider support (OpenAI, Gemini, LiteLLM)
- 📡 Real-time streaming with reasoning summaries
- 📚 Structured citations with URLs
- 🛠️ Built-in tools: web search, code interpreter, MCP, file search
- 🔄 Automatic provider detection from model name
### 5. Query Rewriter Agent
Transform user queries to improve RAG retrieval quality using multiple strategies.
```python
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent(model="gpt-4o-mini")
# Basic - expands abbreviations, adds context
result = agent.rewrite("AI trends")
print(result.primary_query) # "What are the current trends in Artificial Intelligence?"
# HyDE - generates hypothetical document for semantic matching
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
# Step-back - generates broader context question
result = agent.rewrite("GPT-4 vs Claude 3?", strategy=RewriteStrategy.STEP_BACK)
# Sub-queries - decomposes complex questions
result = agent.rewrite("RAG setup and best embedding models?", strategy=RewriteStrategy.SUB_QUERIES)
# Contextual - resolves references using chat history
result = agent.rewrite("What about cost?", chat_history=[...])
```
**Strategies:**
- **BASIC**: Expand abbreviations, fix typos, add context
- **HYDE**: Generate hypothetical document for semantic matching
- **STEP_BACK**: Generate higher-level concept questions
- **SUB_QUERIES**: Decompose multi-part questions
- **MULTI_QUERY**: Generate multiple paraphrased versions
- **CONTEXTUAL**: Resolve references using conversation history
- **AUTO**: Automatically detect best strategy
### 6. Agent Memory (Zero Dependencies)
Enable persistent memory for agents - works out of the box without any extra packages.
```python
from praisonaiagents import Agent
from praisonaiagents.memory import FileMemory
# Enable memory with a single parameter
agent = Agent(
name="Personal Assistant",
instructions="You are a helpful assistant that remembers user preferences.",
memory=True, # Enables file-based memory (no extra deps!)
user_id="user123" # Isolate memory per user
)
# Memory is automatically injected into conversations
result = agent.start("My name is John and I prefer Python")
# Agent will remember this for future conversations
```
**Memory Types:**
- **Short-term**: Rolling buffer of recent context (auto-expires)
- **Long-term**: Persistent important facts (sorted by importance)
- **Entity**: People, places, organizations with attributes
- **Episodic**: Date-based interaction history
**Advanced Features:**
```python
from praisonaiagents.memory import FileMemory
memory = FileMemory(user_id="user123")
# Session Save/Resume
memory.save_session("project_session", conversation_history=[...])
memory.resume_session("project_session")
# Context Compression
memory.compress(llm_func=lambda p: agent.chat(p), max_items=10)
# Checkpointing
memory.create_checkpoint("before_refactor", include_files=["main.py"])
memory.restore_checkpoint("before_refactor", restore_files=True)
# Slash Commands
memory.handle_command("/memory show")
memory.handle_command("/memory save my_session")
```
**Storage Options:**
| Option | Dependencies | Description |
|--------|-------------|-------------|
| `memory=True` | None | File-based JSON storage (default) |
| `memory="file"` | None | Explicit file-based storage |
| `memory="sqlite"` | Built-in | SQLite with indexing |
| `memory="chromadb"` | chromadb | Vector/semantic search |
### 6. Rules & Instructions
PraisonAI auto-discovers instruction files from your project root and git root:
| File | Description | Priority |
|------|-------------|----------|
| `PRAISON.md` | PraisonAI native instructions | High |
| `PRAISON.local.md` | Local overrides (gitignored) | Higher |
| `CLAUDE.md` | Claude Code memory file | High |
| `CLAUDE.local.md` | Local overrides (gitignored) | Higher |
| `AGENTS.md` | OpenAI Codex CLI instructions | High |
| `GEMINI.md` | Gemini CLI memory file | High |
| `.cursorrules` | Cursor IDE rules | High |
| `.windsurfrules` | Windsurf IDE rules | High |
| `.claude/rules/*.md` | Claude Code modular rules | Medium |
| `.windsurf/rules/*.md` | Windsurf modular rules | Medium |
| `.cursor/rules/*.mdc` | Cursor modular rules | Medium |
| `.praison/rules/*.md` | Workspace rules | Medium |
| `~/.praison/rules/*.md` | Global rules | Low |
```python
from praisonaiagents import Agent
# Agent auto-discovers CLAUDE.md, AGENTS.md, GEMINI.md, etc.
agent = Agent(name="Assistant", instructions="You are helpful.")
# Rules are injected into system prompt automatically
```
**@Import Syntax:**
```markdown
# CLAUDE.md
See @README for project overview
See @docs/architecture.md for system design
@~/.praison/my-preferences.md
```
**Rule File Format (with YAML frontmatter):**
```markdown
---
description: Python coding guidelines
globs: ["**/*.py"]
activation: always # always, glob, manual, ai_decision
---
# Guidelines
- Use type hints
- Follow PEP 8
```
### 7. Auto-Generated Memories
```python
from praisonaiagents.memory import FileMemory, AutoMemory
memory = FileMemory(user_id="user123")
auto = AutoMemory(memory, enabled=True)
# Automatically extracts and stores memories from conversations
memories = auto.process_interaction(
"My name is John and I prefer Python for backend work"
)
# Extracts: name="John", preference="Python for backend"
```
### 8. Workflows
Create reusable multi-step workflows in `.praison/workflows/`:
```python
from praisonaiagents.memory import WorkflowManager
manager = WorkflowManager()
# Execute a workflow
result = manager.execute(
"deploy",
executor=lambda prompt: agent.chat(prompt),
variables={"environment": "production"}
)
```
### 9. Hooks
Configure in `.praison/hooks.json`:
```python
from praisonaiagents.memory import HooksManager
hooks = HooksManager()
# Register Python hooks
hooks.register("pre_write_code", lambda ctx: print(f"Writing {ctx['file']}"))
# Execute hooks
result = hooks.execute("pre_write_code", {"file": "main.py"})
```
## Using No Code
### Auto Mode:
@@ -105,6 +374,146 @@ export OPENAI_API_KEY=xxxxxxxxxxxxxxxxxxxxxx
praisonai --auto create a movie script about Robots in Mars
```
### Query Rewriting (works with any command):
```bash
# Rewrite query for better results (uses QueryRewriterAgent)
praisonai "AI trends" --query-rewrite
# Rewrite with search tools (agent decides when to search)
praisonai "latest developments" --query-rewrite --rewrite-tools "internet_search"
# Works with any prompt
praisonai "explain quantum computing" --query-rewrite -v
```
### Deep Research CLI:
```bash
# Default: OpenAI (o4-mini-deep-research)
praisonai research "What are the latest AI trends in 2025?"
# Use Gemini
praisonai research --model deep-research-pro "Your research query"
# Rewrite query before research
praisonai research --query-rewrite "AI trends"
# Rewrite with search tools
praisonai research --query-rewrite --rewrite-tools "internet_search" "AI trends"
# Use custom tools from file (gathers context before deep research)
praisonai research --tools tools.py "Your research query"
praisonai research -t my_tools.py "Your research query"
# Use built-in tools by name (comma-separated)
praisonai research --tools "internet_search,wiki_search" "Your query"
praisonai research -t "yfinance,calculator_tools" "Stock analysis query"
# Save output to file (output/research/{query}.md)
praisonai research --save "Your research query"
praisonai research -s "Your research query"
# Combine options
praisonai research --query-rewrite --tools tools.py --save "Your research query"
# Verbose mode (show debug logs)
praisonai research -v "Your research query"
```
### Planning Mode CLI:
```bash
# Enable planning mode - agent creates a plan before execution
praisonai "Research AI trends and write a summary" --planning
# Planning with tools for research
praisonai "Analyze market trends" --planning --planning-tools tools.py
# Planning with chain-of-thought reasoning
praisonai "Complex analysis task" --planning --planning-reasoning
# Auto-approve plans without confirmation
praisonai "Task" --planning --auto-approve-plan
```
### Memory CLI:
```bash
# Enable memory for agent (persists across sessions)
praisonai "My name is John" --memory
# Memory with user isolation
praisonai "Remember my preferences" --memory --user-id user123
# Memory management commands
praisonai memory show # Show memory statistics
praisonai memory add "User prefers Python" # Add to long-term memory
praisonai memory search "Python" # Search memories
praisonai memory clear # Clear short-term memory
praisonai memory clear all # Clear all memory
praisonai memory save my_session # Save session
praisonai memory resume my_session # Resume session
praisonai memory sessions # List saved sessions
praisonai memory checkpoint # Create checkpoint
praisonai memory restore <checkpoint_id> # Restore checkpoint
praisonai memory checkpoints # List checkpoints
praisonai memory help # Show all commands
```
### Rules CLI:
```bash
# List all loaded rules (from PRAISON.md, CLAUDE.md, etc.)
praisonai rules list
# Show specific rule details
praisonai rules show <rule_name>
# Create a new rule
praisonai rules create my_rule "Always use type hints"
# Delete a rule
praisonai rules delete my_rule
# Show rules statistics
praisonai rules stats
# Include manual rules with prompts
praisonai "Task" --include-rules security,testing
```
### Workflow CLI:
```bash
# List available workflows
praisonai workflow list
# Execute a workflow
praisonai workflow run deploy
# Execute with variables
praisonai workflow run deploy --workflow-var environment=staging --workflow-var branch=main
# Show workflow details
praisonai workflow show deploy
# Create a new workflow template
praisonai workflow create my_workflow
```
### Hooks CLI:
```bash
# List configured hooks
praisonai hooks list
# Show hooks statistics
praisonai hooks stats
# Create hooks.json template
praisonai hooks init
```
### Claude Memory Tool CLI:
```bash
# Enable Claude Memory Tool (Anthropic models only)
praisonai "Research and remember findings" --claude-memory --llm anthropic/claude-sonnet-4-20250514
```
## Using JavaScript Code
```bash
@@ -523,6 +932,172 @@ praisonai agents.yaml
</a>
</div>
## Custom Tools
### Using `@tool` Decorator
```python
from praisonaiagents import Agent, tool
@tool
def search(query: str) -> str:
"""Search the web for information."""
return f"Results for: {query}"
@tool
def calculate(expression: str) -> float:
"""Evaluate a math expression."""
return eval(expression)
agent = Agent(
instructions="You are a helpful assistant",
tools=[search, calculate]
)
agent.start("Search for AI news and calculate 15*4")
```
### Using `BaseTool` Class
```python
from praisonaiagents import Agent, BaseTool
class WeatherTool(BaseTool):
name = "weather"
description = "Get current weather for a location"
def run(self, location: str) -> str:
return f"Weather in {location}: 72°F, Sunny"
agent = Agent(
instructions="You are a weather assistant",
tools=[WeatherTool()]
)
agent.start("What's the weather in Paris?")
```
### Creating a Tool Package (pip installable)
```toml
# pyproject.toml
[project]
name = "my-praisonai-tools"
version = "1.0.0"
dependencies = ["praisonaiagents"]
[project.entry-points."praisonaiagents.tools"]
my_tool = "my_package:MyTool"
```
```python
# my_package/__init__.py
from praisonaiagents import BaseTool
class MyTool(BaseTool):
name = "my_tool"
description = "My custom tool"
def run(self, param: str) -> str:
return f"Result: {param}"
```
After `pip install`, tools are auto-discovered:
```python
agent = Agent(tools=["my_tool"]) # Works automatically!
```
## Prompt Expansion
Expand short prompts into detailed, actionable prompts:
### CLI Usage
```bash
# Expand a short prompt into detailed prompt
praisonai "write a movie script in 3 lines" --expand-prompt
# With verbose output
praisonai "blog about AI" --expand-prompt -v
# With tools for context gathering
praisonai "latest AI trends" --expand-prompt --expand-tools tools.py
# Combine with query rewrite
praisonai "AI news" --query-rewrite --expand-prompt
```
### Programmatic Usage
```python
from praisonaiagents import PromptExpanderAgent, ExpandStrategy
# Basic usage
agent = PromptExpanderAgent()
result = agent.expand("write a movie script in 3 lines")
print(result.expanded_prompt)
# With specific strategy
result = agent.expand("blog about AI", strategy=ExpandStrategy.DETAILED)
# Available strategies: BASIC, DETAILED, STRUCTURED, CREATIVE, AUTO
```
**Key Difference:**
- `--query-rewrite`: Optimizes queries for **search/retrieval** (RAG)
- `--expand-prompt`: Expands prompts for **detailed task execution**
## Web Search, Web Fetch & Prompt Caching
### CLI Usage
```bash
# Web Search - Get real-time information
praisonai "What are the latest AI news today?" --web-search --llm openai/gpt-4o-search-preview
# Web Fetch - Retrieve and analyze URL content (Anthropic only)
praisonai "Summarize https://docs.praison.ai" --web-fetch --llm anthropic/claude-sonnet-4-20250514
# Prompt Caching - Reduce costs for repeated prompts
praisonai "Analyze this document..." --prompt-caching --llm anthropic/claude-sonnet-4-20250514
```
### Programmatic Usage
```python
from praisonaiagents import Agent
# Web Search
agent = Agent(
instructions="You are a research assistant",
llm="openai/gpt-4o-search-preview",
web_search=True
)
# Web Fetch (Anthropic only)
agent = Agent(
instructions="You are a content analyzer",
llm="anthropic/claude-sonnet-4-20250514",
web_fetch=True
)
# Prompt Caching
agent = Agent(
instructions="You are an AI assistant..." * 50, # Long system prompt
llm="anthropic/claude-sonnet-4-20250514",
prompt_caching=True
)
```
**Supported Providers:**
| Feature | Providers |
|---------|-----------|
| Web Search | OpenAI, Gemini, Anthropic, xAI, Perplexity |
| Web Fetch | Anthropic |
| Prompt Caching | OpenAI (auto), Anthropic, Bedrock, Deepseek |
## Development:
Below is used for development only.
@@ -540,6 +1115,19 @@ uv pip install -r pyproject.toml --extra code
uv pip install -r pyproject.toml --extra "crewai,autogen"
```
### Bump and Release
```bash
# From project root - bumps version and releases in one command
python src/praisonai/scripts/bump_and_release.py 2.2.99
# With praisonaiagents dependency
python src/praisonai/scripts/bump_and_release.py 2.2.99 --agents 0.0.169
# Then publish
cd src/praisonai && uv publish
```
## Contributing
- Fork on GitHub: Use the "Fork" button on the repository page.
@@ -557,7 +1145,7 @@ uv pip install -r pyproject.toml --extra "crewai,autogen"
- 🎨 Interactive UIs
- 📄 YAML-based Configuration
- 🛠️ Custom Tool Integration
- 🔍 Internet Search Capability (using Crawl4AI and Tavily)
- 🔍 Internet Search Capability (Tavily, You.com, Exa, DuckDuckGo, Crawl4AI)
- 🖼️ Vision Language Model (VLM) Support
- 🎙️ Real-time Voice Interaction
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
flask \
"praisonai>=2.2.87" \
"praisonai>=2.3.10" \
"praisonai[api]" \
gunicorn \
markdown
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.2.87" \
"praisonai>=2.3.10" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.2.87" \
"praisonai>=2.3.10" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.2.87" \
"praisonai>=2.3.10" \
"praisonai[ui]" \
"praisonai[crewai]"
+2 -2
View File
@@ -121,7 +121,7 @@ healthcheck:
## 📦 Package Versions
All Docker images use consistent, up-to-date versions:
- PraisonAI: `>=2.2.87`
- PraisonAI: `>=2.2.95`
- PraisonAI Agents: `>=0.0.92`
- Python: `3.11-slim`
@@ -218,7 +218,7 @@ docker-compose up -d
### Version Pinning
To use specific versions, update the Dockerfile:
```dockerfile
RUN pip install "praisonai==2.2.87" "praisonaiagents==0.0.92"
RUN pip install "praisonai==2.2.95" "praisonaiagents==0.0.92"
```
## 🌐 Production Deployment
+4 -15
View File
@@ -137,22 +137,11 @@ services:
POSTGRES_USER: praisonai
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-praisonai_dev_pass}
volumes:
database:
image: postgres:15-alpine
environment:
POSTGRES_DB: praisonai_dev
POSTGRES_USER: praisonai
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-praisonai_dev_pass}
volumes:
- postgres_data:/var/lib/postgresql/data
read_only: true
tmpfs:
- /tmp
- /var/tmp
ports:
- "5432:5432"
networks:
- postgres_data:/var/lib/postgresql/data
read_only: true
tmpfs:
- /tmp
- /var/tmp
ports:
- "5432:5432"
networks:
@@ -0,0 +1,481 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "LJRJRQtFTagH"
},
"source": [
"# AI RAG Agent with Web Access using PraisonAI"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Ph4RVTBCTcmB"
},
"source": [
"This notebook demonstrates how to build a Retrieval-Augmented Generation (RAG) agent with web access using GPT-4o and PraisonAI. The agent specializes in Thai recipes and can search the web for additional information."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tMxfwPPATy_o"
},
"source": [
"This RAG agent:\n",
"* Creates an AI assistant specialized in Thai recipes and cooking\n",
"* Has access to a PDF knowledge base about Thai recipes\n",
"* Can search the web using DuckDuckGo for additional information\n",
"* Provides comprehensive answers about Thai cuisine\n",
"* Uses PraisonAI framework for agent management"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "h68oMaCITmqt"
},
"source": [
"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Dhivya-Bharathy/PraisonAI/blob/main/examples/cookbooks/rag_examples/agentic_rag/PraisonAI_RecipeAgent.ipynb)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cAdzoWXJUMfF"
},
"source": [
"# Dependencies"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9D-2TlOrUNcu"
},
"outputs": [],
"source": [
"!pip install praisonaiagents openai requests"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yY-4-kEIUW40"
},
"source": [
"# Set Up API Key"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "f5bhQ05yUX6J",
"outputId": "dcf9d0d9-332e-4136-9289-4c39d898005c"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Enter your OpenAI API key: ··········\n",
"✅ API key set successfully!\n"
]
}
],
"source": [
"import os\n",
"\n",
"# Set your OpenAI API key here\n",
"os.environ['OPENAI_API_KEY'] = 'Enter your api key here'\n",
"\n",
"# Or use this method to input it securely\n",
"from getpass import getpass\n",
"api_key = getpass('Enter your OpenAI API key: ')\n",
"os.environ['OPENAI_API_KEY'] = api_key\n",
"\n",
"print(\"✅ API key set successfully!\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Eh0j-SegUi13"
},
"source": [
"# Tools"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Z1A9oifIUkQN",
"outputId": "9be78226-36fa-43c4-865d-2f7e6e2ba9de"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"✅ Tools imported successfully!\n"
]
}
],
"source": [
"from praisonaiagents import Agent\n",
"import requests\n",
"import tempfile\n",
"from pathlib import Path\n",
"\n",
"print(\"✅ Tools imported successfully!\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "efnnREHXUnYA"
},
"source": [
"# YAML Prompt"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "kcM14-yQUqG-",
"outputId": "c5b6288d-a789-4a93-e20c-1d285370b667"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"✅ Agent instructions configured!\n"
]
}
],
"source": [
"# Agent configuration and instructions\n",
"AGENT_INSTRUCTIONS = \"\"\"\n",
"You are a helpful AI assistant specialized in Thai recipes and cooking.\n",
"\n",
"You have access to a PDF knowledge base about Thai recipes from: https://phi-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf\n",
"\n",
"You can also search the web for additional information about Thai cooking, ingredients, and techniques.\n",
"\n",
"When answering questions:\n",
"1. Use your knowledge about Thai cuisine to provide helpful information\n",
"2. If needed, search the web for additional details, current information, or clarification\n",
"3. Provide comprehensive, helpful answers about Thai cuisine\n",
"4. Always be informative and helpful about Thai cooking!\n",
"\n",
"You can use the internet_search function to search the web when needed.\n",
"\n",
"Focus on providing practical cooking advice, ingredient information, and authentic Thai recipe guidance.\n",
"\"\"\"\n",
"\n",
"print(\"✅ Agent instructions configured!\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MCBRcXFJUulY"
},
"source": [
"# Main"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "YaRLbgMmUvXv",
"outputId": "8ed8d9f8-85a7-48de-ccc6-f617cd672a85"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"🤖 Thai Recipe RAG Agent created successfully!\n",
"The agent is ready to answer questions about Thai cooking.\n"
]
}
],
"source": [
"# Create the RAG agent\n",
"rag_agent = Agent(\n",
" instructions=AGENT_INSTRUCTIONS,\n",
" llm=\"gpt-4o\",\n",
" markdown=True,\n",
" verbose=True\n",
")\n",
"\n",
"print(\"🤖 Thai Recipe RAG Agent created successfully!\")\n",
"print(\"The agent is ready to answer questions about Thai cooking.\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QvZM94pGU0IX"
},
"source": [
"# Test the Agent"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000,
"referenced_widgets": [
"47739975c9f04ee5b4e69e8f445cbeda",
"08888d583cba4c01b3fdf8f30b06ca4f"
]
},
"id": "QPdrleKxU1VN",
"outputId": "a71a86cd-eff8-4309-9b50-173a30066c90"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"🤔 Question: What are the essential ingredients for Thai cooking?\n"
]
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╭─ </span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8; font-weight: bold\">Agent Info</span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\"> ────────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #ff9b9b; text-decoration-color: #ff9b9b; font-weight: bold\">👤 Agent:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Agent</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #b4b4b3; text-decoration-color: #b4b4b3; font-weight: bold\">Role:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Assistant</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[38;2;210;227;200m╭─\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[1;38;2;210;227;200mAgent Info\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[38;2;210;227;200m───────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\u001b[38;2;210;227;200m─╮\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;255;155;155m👤 Agent:\u001b[0m \u001b[38;2;255;229;229mAgent\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;180;180;179mRole:\u001b[0m \u001b[38;2;255;229;229mAssistant\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "47739975c9f04ee5b4e69e8f445cbeda",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Output()"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
],
"text/plain": []
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
"</pre>\n"
],
"text/plain": [
"\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Response generated in 8.4s</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[2mResponse generated in 8.4s\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭───────────────────────────────────────────────────── Task ──────────────────────────────────────────────────────╮</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> What are the essential ingredients for Thai cooking? <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[36m╭─\u001b[0m\u001b[36m────────────────────────────────────────────────────\u001b[0m\u001b[36m Task \u001b[0m\u001b[36m─────────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
"\u001b[36m│\u001b[0m What are the essential ingredients for Thai cooking? \u001b[36m│\u001b[0m\n",
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭─────────────────────────────────────────────────── Response ────────────────────────────────────────────────────╮</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> Thai cooking is known for its balance of five fundamental taste senses in each dish or the overall meal: sour, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> sweet, salty, bitter, and spicy. Essential ingredients that help achieve this balance include: <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 1 </span><span style=\"font-weight: bold\">Fish Sauce (Nam Pla)</span>: A salty, savory sauce made from fermented fish. It's a cornerstone of Thai cooking, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>used to add depth of flavor. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 2 </span><span style=\"font-weight: bold\">Lime Juice</span>: Used to add a fresh, tangy citrus flavor to dishes, especially in salads and some curries. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 3 </span><span style=\"font-weight: bold\">Chilies</span>: Thai cuisine often features fresh Thai bird's eye chilies for heat, as well as dried chilies for <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>certain recipes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 4 </span><span style=\"font-weight: bold\">Lemongrass</span>: A citrusy, fragrant herb that is often used in soups, curries, and teas. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 5 </span><span style=\"font-weight: bold\">Galangal</span>: Similar to ginger, galangal has a sharp, peppery flavor and is used in soups and curries. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 6 </span><span style=\"font-weight: bold\">Kaffir Lime Leaves</span>: They add a distinct, zesty citrus aroma and flavor to Thai dishes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 7 </span><span style=\"font-weight: bold\">Coconut Milk</span>: Adds richness and creaminess, commonly used in curries and desserts. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 8 </span><span style=\"font-weight: bold\">Palm Sugar</span>: Used as a sweetener, palm sugar has a nuanced flavor, unlike regular sugar. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 9 </span><span style=\"font-weight: bold\">Tamarind Paste</span>: Provides a sour element to many dishes, particularly pad Thai and some sauces. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 10 </span><span style=\"font-weight: bold\">Thai Basil</span>: Adds a spicy, peppery flavor, different from the sweet basil used in Western cooking. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 11 </span><span style=\"font-weight: bold\">Coriander (Cilantro)</span>: Used as a garnish and for its fresh, citrusy leaves and roots, which add depth to <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>sauces and stocks. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 12 </span><span style=\"font-weight: bold\">Shrimp Paste</span>: A very pungent paste made from fermented shrimp, used sparingly to enrich the flavor of <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>dishes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> These ingredients are commonly used in combination to create the complex and rich flavors characteristic of <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> Thai cuisine. Fresh and high-quality versions of these ingredients contribute significantly to the authenticity <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> of Thai dishes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[36m╭─\u001b[0m\u001b[36m──────────────────────────────────────────────────\u001b[0m\u001b[36m Response \u001b[0m\u001b[36m───────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
"\u001b[36m│\u001b[0m Thai cooking is known for its balance of five fundamental taste senses in each dish or the overall meal: sour, \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m sweet, salty, bitter, and spicy. Essential ingredients that help achieve this balance include: \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 1 \u001b[0m\u001b[1mFish Sauce (Nam Pla)\u001b[0m: A salty, savory sauce made from fermented fish. It's a cornerstone of Thai cooking, \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mused to add depth of flavor. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 2 \u001b[0m\u001b[1mLime Juice\u001b[0m: Used to add a fresh, tangy citrus flavor to dishes, especially in salads and some curries. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 3 \u001b[0m\u001b[1mChilies\u001b[0m: Thai cuisine often features fresh Thai bird's eye chilies for heat, as well as dried chilies for \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mcertain recipes. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 4 \u001b[0m\u001b[1mLemongrass\u001b[0m: A citrusy, fragrant herb that is often used in soups, curries, and teas. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 5 \u001b[0m\u001b[1mGalangal\u001b[0m: Similar to ginger, galangal has a sharp, peppery flavor and is used in soups and curries. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 6 \u001b[0m\u001b[1mKaffir Lime Leaves\u001b[0m: They add a distinct, zesty citrus aroma and flavor to Thai dishes. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 7 \u001b[0m\u001b[1mCoconut Milk\u001b[0m: Adds richness and creaminess, commonly used in curries and desserts. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 8 \u001b[0m\u001b[1mPalm Sugar\u001b[0m: Used as a sweetener, palm sugar has a nuanced flavor, unlike regular sugar. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 9 \u001b[0m\u001b[1mTamarind Paste\u001b[0m: Provides a sour element to many dishes, particularly pad Thai and some sauces. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 10 \u001b[0m\u001b[1mThai Basil\u001b[0m: Adds a spicy, peppery flavor, different from the sweet basil used in Western cooking. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 11 \u001b[0m\u001b[1mCoriander (Cilantro)\u001b[0m: Used as a garnish and for its fresh, citrusy leaves and roots, which add depth to \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0msauces and stocks. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m 12 \u001b[0m\u001b[1mShrimp Paste\u001b[0m: A very pungent paste made from fermented shrimp, used sparingly to enrich the flavor of \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mdishes. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m These ingredients are commonly used in combination to create the complex and rich flavors characteristic of \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m Thai cuisine. Fresh and high-quality versions of these ingredients contribute significantly to the authenticity \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m of Thai dishes. \u001b[36m│\u001b[0m\n",
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"�� Answer: Thai cooking is known for its balance of five fundamental taste senses in each dish or the overall meal: sour, sweet, salty, bitter, and spicy. Essential ingredients that help achieve this balance include:\n",
"\n",
"1. **Fish Sauce (Nam Pla)**: A salty, savory sauce made from fermented fish. It's a cornerstone of Thai cooking, used to add depth of flavor.\n",
"\n",
"2. **Lime Juice**: Used to add a fresh, tangy citrus flavor to dishes, especially in salads and some curries.\n",
"\n",
"3. **Chilies**: Thai cuisine often features fresh Thai bird's eye chilies for heat, as well as dried chilies for certain recipes.\n",
"\n",
"4. **Lemongrass**: A citrusy, fragrant herb that is often used in soups, curries, and teas.\n",
"\n",
"5. **Galangal**: Similar to ginger, galangal has a sharp, peppery flavor and is used in soups and curries.\n",
"\n",
"6. **Kaffir Lime Leaves**: They add a distinct, zesty citrus aroma and flavor to Thai dishes.\n",
"\n",
"7. **Coconut Milk**: Adds richness and creaminess, commonly used in curries and desserts.\n",
"\n",
"8. **Palm Sugar**: Used as a sweetener, palm sugar has a nuanced flavor, unlike regular sugar.\n",
"\n",
"9. **Tamarind Paste**: Provides a sour element to many dishes, particularly pad Thai and some sauces.\n",
"\n",
"10. **Thai Basil**: Adds a spicy, peppery flavor, different from the sweet basil used in Western cooking.\n",
"\n",
"11. **Coriander (Cilantro)**: Used as a garnish and for its fresh, citrusy leaves and roots, which add depth to sauces and stocks.\n",
"\n",
"12. **Shrimp Paste**: A very pungent paste made from fermented shrimp, used sparingly to enrich the flavor of dishes.\n",
"\n",
"These ingredients are commonly used in combination to create the complex and rich flavors characteristic of Thai cuisine. Fresh and high-quality versions of these ingredients contribute significantly to the authenticity of Thai dishes.\n"
]
}
],
"source": [
"# Test question 1: Basic Thai cooking\n",
"question1 = \"What are the essential ingredients for Thai cooking?\"\n",
"print(f\"🤔 Question: {question1}\")\n",
"\n",
"try:\n",
" response1 = rag_agent.start(question1)\n",
" print(f\"�� Answer: {response1}\")\n",
"except Exception as e:\n",
" print(f\"❌ Error: {e}\")"
]
}
],
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,556 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "EA8l4cSVl3Lx"
},
"source": [
"# Agentic RAG with GPT-5 - Colab Notebook"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "I-8nOxJ2l7nO"
},
"source": [
"This Colab notebook sets up a Streamlit app using PraisonAI Agents framework.\n",
"It allows you to interact with an AI agent powered by GPT-5 for Q&A.\n",
"Currently, the knowledge base functionality is disabled due to compatibility issues."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "g5I1tMkxmDQe"
},
"source": [
"- GPT-5 Integration (via OpenAI API)\n",
"- Streamlit Web App\n",
"- Agent Framework: PraisonAI Agents\n",
"- Markdown-formatted responses"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yeEno7U0q2U_"
},
"source": [
"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Dhivya-Bharathy/PraisonAI/blob/main/examples/cookbooks/Qwen2.5_InstructionAgent.ipynb)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "WV2GnzZCmGBR"
},
"source": [
"# Dependencies"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "20wAadXemJ0k"
},
"outputs": [],
"source": [
"!pip install praisonaiagents python-dotenv streamlit"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "l71uCRarndvy"
},
"source": [
"# Set API Key"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"id": "brfc-ZZKnfre"
},
"outputs": [],
"source": [
"import os\n",
"from dotenv import load_dotenv\n",
"\n",
"load_dotenv()\n",
"\n",
"OPENAI_API_KEY = \"Enter your api key here\"\n",
"\n",
"os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "d1fGzyQCn1aQ"
},
"source": [
"# Tools and Imports"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"id": "GcNPLkZXntIS"
},
"outputs": [],
"source": [
"import streamlit as st\n",
"from praisonaiagents import Agent\n",
"\n",
"# YAML Prompt Configuration\n",
"agent_config = {\n",
" \"name\": \"Knowledge Agent\",\n",
" \"instructions\": [\n",
" \"You are a helpful AI assistant. Answer questions based on your general knowledge.\",\n",
" \"Provide clear, well-structured answers in markdown format.\",\n",
" \"Use proper markdown formatting with headers, lists, and emphasis where appropriate.\",\n",
" \"Structure your response with clear sections and bullet points when helpful.\",\n",
" ],\n",
" \"llm\": \"gpt-5-nano\",\n",
" \"markdown\": True,\n",
" \"verbose\": True\n",
"}\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "jkmrh0gIn_4a"
},
"source": [
"# Initialize Agent"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"id": "qaJvf8RyoAsB"
},
"outputs": [],
"source": [
"def create_agent():\n",
" \"\"\"Create an agent with reasoning capabilities\"\"\"\n",
" return Agent(\n",
" name=agent_config[\"name\"],\n",
" instructions=agent_config[\"instructions\"],\n",
" llm=agent_config[\"llm\"],\n",
" markdown=agent_config[\"markdown\"],\n",
" verbose=agent_config[\"verbose\"]\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kvGJh3bbqTmh"
},
"source": [
"# Main Function"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000,
"referenced_widgets": [
"ad296d97510246de8eb0138ce8a8a166",
"f6e502b3dca442f0ae6c5739f5a6d159"
]
},
"id": "_QyxzqKKqUrW",
"outputId": "97be7e02-b034-4004-e9e3-56337e3f2887"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"�� Initializing Agentic RAG with GPT-5...\n",
"📚 Loading agent...\n",
"✅ Agent loaded successfully!\n",
"\n",
"🎯 Example queries you can try:\n",
"1. What is PraisonAI and how do Agents work?\n",
"2. What are Teams in PraisonAI and how do they work?\n",
"3. Give me a step-by-step guide to building a RAG system.\n",
"4. What are AI Agents?\n",
"\n",
"============================================================\n",
"\n",
"🤔 Enter your question (or 'quit' to exit):\n",
"> What is PraisonAI and how do Agents work?\n",
"🔍 Processing your question...\n"
]
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╭─ </span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8; font-weight: bold\">Agent Info</span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\"> ────────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #ff9b9b; text-decoration-color: #ff9b9b; font-weight: bold\">👤 Agent:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Knowledge Agent</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #b4b4b3; text-decoration-color: #b4b4b3; font-weight: bold\">Role:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Assistant</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[38;2;210;227;200m╭─\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[1;38;2;210;227;200mAgent Info\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[38;2;210;227;200m───────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\u001b[38;2;210;227;200m─╮\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;255;155;155m👤 Agent:\u001b[0m \u001b[38;2;255;229;229mKnowledge Agent\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;180;180;179mRole:\u001b[0m \u001b[38;2;255;229;229mAssistant\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
"\u001b[38;2;210;227;200m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ad296d97510246de8eb0138ce8a8a166",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Output()"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
],
"text/plain": []
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
"</pre>\n"
],
"text/plain": [
"\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Response generated in 9.4s</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[2mResponse generated in 9.4s\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭───────────────────────────────────────────────────── Task ──────────────────────────────────────────────────────╮</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> What is PraisonAI and how do Agents work? <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[36m╭─\u001b[0m\u001b[36m────────────────────────────────────────────────────\u001b[0m\u001b[36m Task \u001b[0m\u001b[36m─────────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
"\u001b[36m│\u001b[0m What is PraisonAI and how do Agents work? \u001b[36m│\u001b[0m\n",
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭─────────────────────────────────────────────────── Response ────────────────────────────────────────────────────╮</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┃ <span style=\"font-weight: bold\">What is PraisonAI?</span> ┃ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> PraisonAI is a platform designed to leverage artificial intelligence to enhance personal and business <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> productivity. It typically involves developing AI models that can assist in various tasks, automate processes, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> and provide insights based on data analysis. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">Key Features of PraisonAI</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Task Automation</span>: Automates repetitive tasks, allowing users to focus on more strategic activities. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Data Insights</span>: Analyzes large sets of data to provide actionable insights. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Natural Language Processing</span>: Enables users to interact with the AI through natural language, making it <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>user-friendly. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Customizable Solutions</span>: Tailors AI applications to meet specific user needs across various industries. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┃ <span style=\"font-weight: bold\">How Do Agents Work in PraisonAI?</span> ┃ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> Agents in PraisonAI refer to autonomous or semi-autonomous AI systems that perform designated tasks or <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> functions. Here’s how they typically operate: <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">1. Types of Agents</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Reactive Agents</span>: Respond to the current environment without internal memory. They act solely based on <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>current inputs. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Deliberative Agents</span>: Have a model of their environment, plan their actions, and retain memories of past <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>interactions. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Learning Agents</span>: Improve their performance over time through experience, adapting to new information. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">2. Functionality of Agents</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Task Execution</span>: Agents can carry out specific tasks as programmed, such as scheduling, data entry, or <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>customer support. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Interaction</span>: They often interact with users through chat interfaces or APIs, providing a seamless user <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>experience. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Integration</span>: Can be integrated with other systems or software to enhance functionality, such as CRM systems, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>databases, or other AI tools. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Continuous Learning</span>: Agents may utilize machine learning algorithms to adapt and evolve based on user <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>interactions and data received. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">3. Benefits of Using Agents</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Increased Efficiency</span>: Automate routine tasks to reduce workload. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">24/7 Availability</span>: Provide consistent support and functionality without downtime. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Scalability</span>: Easily scale operations by deploying more agents as needed. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Data-Driven Decisions</span>: Use data insights to inform strategy and operations. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">Conclusion</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> PraisonAI represents a robust platform for enhancing productivity through AI, with agents serving as the heart <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> of this solution, performing various tasks and supporting users in achieving their goals efficiently. Whether <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> they are simple reactive agents or complex learning entities, their design focuses on making operations <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> smoother and more effective. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
"</pre>\n"
],
"text/plain": [
"\u001b[36m╭─\u001b[0m\u001b[36m──────────────────────────────────────────────────\u001b[0m\u001b[36m Response \u001b[0m\u001b[36m───────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
"\u001b[36m│\u001b[0m ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m ┃ \u001b[1mWhat is PraisonAI?\u001b[0m ┃ \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m PraisonAI is a platform designed to leverage artificial intelligence to enhance personal and business \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m productivity. It typically involves developing AI models that can assist in various tasks, automate processes, \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m and provide insights based on data analysis. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;4mKey Features of PraisonAI\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mTask Automation\u001b[0m: Automates repetitive tasks, allowing users to focus on more strategic activities. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mData Insights\u001b[0m: Analyzes large sets of data to provide actionable insights. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mNatural Language Processing\u001b[0m: Enables users to interact with the AI through natural language, making it \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0muser-friendly. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mCustomizable Solutions\u001b[0m: Tailors AI applications to meet specific user needs across various industries. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m ┃ \u001b[1mHow Do Agents Work in PraisonAI?\u001b[0m ┃ \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m Agents in PraisonAI refer to autonomous or semi-autonomous AI systems that perform designated tasks or \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m functions. Here’s how they typically operate: \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;4m1. \u001b[0m\u001b[1;4mTypes of Agents\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mReactive Agents\u001b[0m: Respond to the current environment without internal memory. They act solely based on \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mcurrent inputs. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mDeliberative Agents\u001b[0m: Have a model of their environment, plan their actions, and retain memories of past \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0minteractions. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mLearning Agents\u001b[0m: Improve their performance over time through experience, adapting to new information. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;4m2. \u001b[0m\u001b[1;4mFunctionality of Agents\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mTask Execution\u001b[0m: Agents can carry out specific tasks as programmed, such as scheduling, data entry, or \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mcustomer support. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mInteraction\u001b[0m: They often interact with users through chat interfaces or APIs, providing a seamless user \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mexperience. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mIntegration\u001b[0m: Can be integrated with other systems or software to enhance functionality, such as CRM systems, \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mdatabases, or other AI tools. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mContinuous Learning\u001b[0m: Agents may utilize machine learning algorithms to adapt and evolve based on user \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0minteractions and data received. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;4m3. \u001b[0m\u001b[1;4mBenefits of Using Agents\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mIncreased Efficiency\u001b[0m: Automate routine tasks to reduce workload. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1m24/7 Availability\u001b[0m: Provide consistent support and functionality without downtime. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mScalability\u001b[0m: Easily scale operations by deploying more agents as needed. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mData-Driven Decisions\u001b[0m: Use data insights to inform strategy and operations. \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[1;4mConclusion\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m PraisonAI represents a robust platform for enhancing productivity through AI, with agents serving as the heart \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m of this solution, performing various tasks and supporting users in achieving their goals efficiently. Whether \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m they are simple reactive agents or complex learning entities, their design focuses on making operations \u001b[36m│\u001b[0m\n",
"\u001b[36m│\u001b[0m smoother and more effective. \u001b[36m│\u001b[0m\n",
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"💡 Answer:\n",
"========================================\n",
"# What is PraisonAI?\n",
"\n",
"PraisonAI is a platform designed to leverage artificial intelligence to enhance personal and business productivity. It typically involves developing AI models that can assist in various tasks, automate processes, and provide insights based on data analysis.\n",
"\n",
"## Key Features of PraisonAI\n",
"\n",
"- **Task Automation**: Automates repetitive tasks, allowing users to focus on more strategic activities.\n",
"- **Data Insights**: Analyzes large sets of data to provide actionable insights.\n",
"- **Natural Language Processing**: Enables users to interact with the AI through natural language, making it user-friendly.\n",
"- **Customizable Solutions**: Tailors AI applications to meet specific user needs across various industries.\n",
"\n",
"# How Do Agents Work in PraisonAI?\n",
"\n",
"Agents in PraisonAI refer to autonomous or semi-autonomous AI systems that perform designated tasks or functions. Here’s how they typically operate:\n",
"\n",
"## 1. **Types of Agents**\n",
"\n",
"- **Reactive Agents**: Respond to the current environment without internal memory. They act solely based on current inputs.\n",
"- **Deliberative Agents**: Have a model of their environment, plan their actions, and retain memories of past interactions.\n",
"- **Learning Agents**: Improve their performance over time through experience, adapting to new information.\n",
"\n",
"## 2. **Functionality of Agents**\n",
"\n",
"- **Task Execution**: Agents can carry out specific tasks as programmed, such as scheduling, data entry, or customer support.\n",
"- **Interaction**: They often interact with users through chat interfaces or APIs, providing a seamless user experience.\n",
"- **Integration**: Can be integrated with other systems or software to enhance functionality, such as CRM systems, databases, or other AI tools.\n",
"- **Continuous Learning**: Agents may utilize machine learning algorithms to adapt and evolve based on user interactions and data received.\n",
"\n",
"## 3. **Benefits of Using Agents**\n",
"\n",
"- **Increased Efficiency**: Automate routine tasks to reduce workload.\n",
"- **24/7 Availability**: Provide consistent support and functionality without downtime.\n",
"- **Scalability**: Easily scale operations by deploying more agents as needed.\n",
"- **Data-Driven Decisions**: Use data insights to inform strategy and operations.\n",
"\n",
"## Conclusion\n",
"\n",
"PraisonAI represents a robust platform for enhancing productivity through AI, with agents serving as the heart of this solution, performing various tasks and supporting users in achieving their goals efficiently. Whether they are simple reactive agents or complex learning entities, their design focuses on making operations smoother and more effective.\n",
"========================================\n",
"\n",
"🤔 Enter your question (or 'quit' to exit):\n",
"> quit\n",
"👋 Goodbye!\n"
]
}
],
"source": [
"def main():\n",
" print(\"�� Initializing Agentic RAG with GPT-5...\")\n",
"\n",
" # Check if API key is set\n",
" if not OPENAI_API_KEY or OPENAI_API_KEY == \"your_openai_api_key_here\":\n",
" print(\"❌ Please set your OpenAI API key first!\")\n",
" print(\"�� Get your key from: https://platform.openai.com/\")\n",
" return\n",
"\n",
" try:\n",
" # Create agent\n",
" print(\"📚 Loading agent...\")\n",
" agent = create_agent()\n",
" print(\"✅ Agent loaded successfully!\")\n",
"\n",
" # Example queries\n",
" example_queries = [\n",
" \"What is PraisonAI and how do Agents work?\",\n",
" \"What are Teams in PraisonAI and how do they work?\",\n",
" \"Give me a step-by-step guide to building a RAG system.\",\n",
" \"What are AI Agents?\"\n",
" ]\n",
"\n",
" print(\"\\n🎯 Example queries you can try:\")\n",
" for i, query in enumerate(example_queries, 1):\n",
" print(f\"{i}. {query}\")\n",
"\n",
" print(\"\\n\" + \"=\"*60)\n",
"\n",
" # Interactive query loop\n",
" while True:\n",
" print(\"\\n🤔 Enter your question (or 'quit' to exit):\")\n",
" user_query = input(\"> \")\n",
"\n",
" if user_query.lower() in ['quit', 'exit', 'q']:\n",
" print(\"👋 Goodbye!\")\n",
" break\n",
"\n",
" if user_query.strip():\n",
" print(\"🔍 Processing your question...\")\n",
" try:\n",
" response = agent.start(user_query)\n",
" print(\"\\n💡 Answer:\")\n",
" print(\"=\"*40)\n",
" print(response)\n",
" print(\"=\"*40)\n",
" except Exception as e:\n",
" print(f\"❌ Error: {str(e)}\")\n",
" print(f\"🔍 Error type: {type(e).__name__}\")\n",
" print(\"💡 This might be due to API issues or model access problems.\")\n",
" else:\n",
" print(\"❌ Please enter a valid question.\")\n",
"\n",
" except Exception as e:\n",
" print(f\"❌ Failed to initialize agent: {str(e)}\")\n",
" print(\"�� Make sure your API key is correct and you have sufficient credits.\")\n",
"\n",
"# Run the main function\n",
"if __name__ == \"__main__\":\n",
" main()"
]
}
],
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,83 @@
"""
MCP Backward Compatibility Example
Demonstrates backward compatibility with older MCP servers
using the deprecated HTTP+SSE transport (protocol version 2024-11-05).
Features:
- Legacy SSE endpoint detection (/sse suffix)
- Automatic transport negotiation
- Protocol version support (2024-11-05 to 2025-11-25)
Protocol: MCP 2025-11-25 with 2024-11-05 compatibility
"""
# Example code (requires running MCP servers):
#
# from praisonaiagents import Agent, MCP
#
# # Legacy SSE server (auto-detected via /sse suffix)
# agent_legacy = Agent(
# name="Legacy SSE Assistant",
# tools=MCP("http://localhost:8080/sse")
# )
#
# # Modern Streamable HTTP server
# agent_modern = Agent(
# name="Modern HTTP Assistant",
# tools=MCP("http://localhost:8080/mcp")
# )
if __name__ == "__main__":
from praisonaiagents.mcp.mcp_transport import get_transport_type
from praisonaiagents.mcp.mcp_compat import (
detect_transport_support,
is_legacy_sse_url,
TransportNegotiator
)
from praisonaiagents.mcp.mcp_session import VALID_PROTOCOL_VERSIONS
print("MCP Backward Compatibility Example")
print("=" * 50)
# URL detection
print("\n1. URL-based Transport Detection:")
urls = [
"http://localhost:8080/sse",
"https://api.example.com/sse",
"http://localhost:8080/mcp",
"https://api.example.com/mcp",
]
for url in urls:
transport = get_transport_type(url)
legacy = is_legacy_sse_url(url)
print(f" {url}")
print(f" Transport: {transport}, Legacy: {legacy}")
# Transport negotiation
print("\n2. Transport Negotiation (HTTP response based):")
responses = [
(200, "application/json", "Streamable HTTP"),
(200, "text/event-stream", "SSE stream"),
(400, None, "Fallback to SSE"),
(404, None, "Fallback to SSE"),
(405, None, "Fallback to SSE"),
]
for status, content_type, desc in responses:
detected = detect_transport_support(status, content_type)
print(f" HTTP {status} ({desc}): {detected}")
# Protocol versions
print("\n3. Supported Protocol Versions:")
for version in sorted(VALID_PROTOCOL_VERSIONS):
print(f" - {version}")
# Negotiator
print("\n4. Transport Negotiator:")
negotiator = TransportNegotiator("https://api.example.com/mcp")
print(f" Base URL: {negotiator.base_url}")
print(f" Try Streamable first: {negotiator.should_try_streamable_first()}")
negotiator_sse = TransportNegotiator("https://api.example.com/sse")
print(f" SSE URL: {negotiator_sse.base_url}")
print(f" Try Streamable first: {negotiator_sse.should_try_streamable_first()}")
+110
View File
@@ -0,0 +1,110 @@
"""
MCP Custom Transport Example
Demonstrates how to create custom MCP transports by extending
the BaseTransport class.
Features:
- BaseTransport abstract class
- TransportRegistry for registration
- TransportConfig for configuration
- Pluggable transport mechanism
Protocol: MCP 2025-11-25
"""
from praisonaiagents.mcp.mcp_transport import (
BaseTransport,
TransportConfig,
TransportRegistry,
get_default_registry
)
# Custom transport implementation
class MyCustomTransport(BaseTransport):
"""Example custom transport implementation."""
def __init__(self, url: str, config: TransportConfig = None):
self.url = url
self.config = config or TransportConfig()
self._connected = False
async def connect(self):
"""Establish connection."""
print(f"Connecting to {self.url}...")
self._connected = True
print("Connected!")
async def send(self, message: dict):
"""Send JSON-RPC message."""
if not self._connected:
raise RuntimeError("Not connected")
print(f"Sending: {message}")
async def receive(self) -> dict:
"""Receive JSON-RPC message."""
if not self._connected:
raise RuntimeError("Not connected")
return {"jsonrpc": "2.0", "result": "ok", "id": 1}
async def close(self):
"""Close connection."""
print("Closing connection...")
self._connected = False
@property
def is_connected(self) -> bool:
return self._connected
if __name__ == "__main__":
import asyncio
print("MCP Custom Transport Example")
print("=" * 50)
# 1. Transport configuration
print("\n1. Transport Configuration:")
config = TransportConfig(
timeout=60,
debug=True,
retry_count=3,
retry_delay=1.0,
auth_token="my-token"
)
print(f" Timeout: {config.timeout}s")
print(f" Debug: {config.debug}")
print(f" Retry count: {config.retry_count}")
# 2. Default registry
print("\n2. Default Transport Registry:")
registry = get_default_registry()
print(f" Registered transports: {registry.list_transports()}")
# 3. Register custom transport
print("\n3. Registering Custom Transport:")
custom_registry = TransportRegistry()
custom_registry.register("my_custom", MyCustomTransport)
print(f" Registered: {custom_registry.list_transports()}")
# 4. Use custom transport
print("\n4. Using Custom Transport:")
async def demo():
transport = MyCustomTransport("custom://localhost:8080", config)
async with transport:
await transport.send({"jsonrpc": "2.0", "method": "test", "id": 1})
response = await transport.receive()
print(f" Response: {response}")
asyncio.run(demo())
# 5. BaseTransport interface
print("\n5. BaseTransport Interface:")
print(" Required methods:")
print(" - async connect()")
print(" - async send(message: dict)")
print(" - async receive() -> dict")
print(" - async close()")
print(" - @property is_connected -> bool")
+47
View File
@@ -0,0 +1,47 @@
"""
MCP Resumability Example
Demonstrates SSE stream resumability for MCP transports.
Resumability allows recovering from disconnections without losing messages.
Features:
- Event ID tracking per stream
- Last-Event-ID header on reconnection
- Retry delay handling from server
- Message replay on reconnection
Protocol: MCP 2025-11-25
"""
from praisonaiagents import Agent, MCP
# Agent with resumability support (automatic)
agent = Agent(
name="Resumable Assistant",
instructions="You are a helpful assistant with resumable connections.",
tools=MCP("https://api.example.com/mcp")
)
# Resumability is automatic - the MCP client:
# 1. Tracks event IDs from SSE streams (id: field)
# 2. Stores last event ID per stream
# 3. Includes Last-Event-ID header on reconnection
# 4. Respects retry delay from server (retry: field)
if __name__ == "__main__":
print("MCP Resumability Example")
print("=" * 40)
print()
print("SSE Event Format:")
print(" id: event-123")
print(" retry: 3000")
print(" data: {\"jsonrpc\": \"2.0\", ...}")
print()
print("Resumability Flow:")
print("1. Client receives SSE events with id: field")
print("2. Client stores last event ID")
print("3. On disconnection, client reconnects with Last-Event-ID header")
print("4. Server replays messages after that event ID")
print("5. Client respects retry: field for reconnection delay")
print()
print("This prevents message loss during network interruptions.")
+79
View File
@@ -0,0 +1,79 @@
"""
MCP Security Example
Demonstrates security features for MCP transports.
These features help prevent attacks like DNS rebinding.
Features:
- Origin header validation
- DNS rebinding prevention
- Localhost binding for local servers
- Authentication header support
- Secure session ID generation
Protocol: MCP 2025-11-25
"""
# Example code (requires running MCP servers):
#
# from praisonaiagents import Agent, MCP
#
# # Agent with authentication
# agent = Agent(
# name="Secure Assistant",
# tools=MCP(
# "https://api.example.com/mcp",
# headers={"Authorization": "Bearer your-token"}
# )
# )
#
# # WebSocket with auth token
# agent_ws = Agent(
# name="Secure WebSocket Assistant",
# tools=MCP(
# "wss://api.example.com/mcp",
# auth_token="Bearer your-secret-token"
# )
# )
# Using security utilities directly (no server required)
if __name__ == "__main__":
from praisonaiagents.mcp.mcp_security import (
is_valid_origin,
create_auth_header,
generate_secure_session_id,
SecurityConfig
)
print("MCP Security Example")
print("=" * 40)
# Origin validation (DNS rebinding prevention)
print("\n1. Origin Validation:")
allowed = ["localhost", "127.0.0.1", "example.com"]
print(f" Allowed origins: {allowed}")
print(f" 'http://localhost:8080' valid: {is_valid_origin('http://localhost:8080', allowed)}")
print(f" 'https://evil.com' valid: {is_valid_origin('https://evil.com', allowed)}")
# Authentication headers
print("\n2. Authentication Headers:")
bearer = create_auth_header("my-token", auth_type="bearer")
print(f" Bearer: {bearer}")
basic = create_auth_header("user:pass", auth_type="basic")
print(f" Basic: {basic}")
# Secure session IDs
print("\n3. Secure Session IDs:")
session_id = generate_secure_session_id()
print(f" Generated: {session_id}")
print(f" Length: {len(session_id)} chars")
# Security config
print("\n4. Security Configuration:")
config = SecurityConfig(
allowed_origins=["localhost", "example.com"],
require_auth=True
)
print(f" Validate origin: {config.validate_origin}")
print(f" Require auth: {config.require_auth}")
print(f" Bind address: {config.get_bind_address()}")
@@ -0,0 +1,44 @@
"""
MCP Session Management Example
Demonstrates session management features for MCP transports.
Sessions allow stateful interactions with MCP servers.
Features:
- Automatic session ID handling
- Session expiration detection (HTTP 404)
- Session termination via HTTP DELETE
- Protocol version negotiation
Protocol: MCP 2025-11-25
"""
from praisonaiagents import Agent, MCP
# Basic agent with automatic session management
agent = Agent(
name="Session Assistant",
instructions="You are a helpful assistant with session support.",
tools=MCP("https://api.example.com/mcp")
)
# Session management is automatic - the MCP client:
# 1. Receives Mcp-Session-Id from server during initialization
# 2. Includes session ID in all subsequent requests
# 3. Handles HTTP 404 (session expired) by reinitializing
# 4. Can terminate session explicitly when done
if __name__ == "__main__":
print("MCP Session Management Example")
print("=" * 40)
print()
print("Session Lifecycle:")
print("1. Client sends InitializeRequest")
print("2. Server returns InitializeResult with Mcp-Session-Id header")
print("3. Client includes Mcp-Session-Id in all subsequent requests")
print("4. If HTTP 404 received, client reinitializes session")
print("5. Client can send HTTP DELETE to terminate session")
print()
print("Headers used:")
print("- Mcp-Session-Id: Session identifier")
print("- Mcp-Protocol-Version: Protocol version (e.g., 2025-03-26)")
@@ -0,0 +1,91 @@
"""
MCP Transports Overview Example
Demonstrates all MCP transport types and auto-detection.
The MCP class automatically selects the appropriate transport
based on the URL pattern.
Transport Types:
- stdio: Command strings (local subprocess)
- Streamable HTTP: http:// or https:// URLs
- WebSocket: ws:// or wss:// URLs
- SSE (Legacy): URLs ending in /sse
Protocol: MCP 2025-11-25
"""
# Example code (requires running MCP servers):
#
# from praisonaiagents import Agent, MCP
#
# # 1. stdio Transport - Local subprocess
# agent_stdio = Agent(
# name="Stdio Assistant",
# tools=MCP("npx @modelcontextprotocol/server-memory")
# )
#
# # 2. Streamable HTTP Transport
# agent_http = Agent(
# name="HTTP Assistant",
# tools=MCP("https://api.example.com/mcp")
# )
#
# # 3. WebSocket Transport
# agent_ws = Agent(
# name="WebSocket Assistant",
# tools=MCP("wss://api.example.com/mcp")
# )
#
# # 4. SSE Transport (Legacy)
# agent_sse = Agent(
# name="SSE Assistant",
# tools=MCP("http://localhost:8080/sse")
# )
if __name__ == "__main__":
from praisonaiagents.mcp.mcp_transport import get_transport_type
print("MCP Transports Overview")
print("=" * 50)
print()
print("Transport Auto-Detection:")
print("-" * 50)
examples = [
("npx @mcp/server-memory", "stdio"),
("python3 server.py", "stdio"),
("https://api.example.com/mcp", "http_stream"),
("http://localhost:8080/mcp", "http_stream"),
("ws://localhost:8080/mcp", "websocket"),
("wss://api.example.com/mcp", "websocket"),
("http://localhost:8080/sse", "sse"),
("https://api.example.com/sse", "sse"),
]
for url, expected in examples:
detected = get_transport_type(url)
status = "✓" if detected == expected else "✗"
print(f" {status} '{url}'")
print(f" → {detected}")
print()
print("Transport Features:")
print("-" * 50)
print(" stdio:")
print(" - Local subprocess communication")
print(" - Newline-delimited JSON-RPC")
print(" - Environment variable support")
print()
print(" Streamable HTTP:")
print(" - Session management (Mcp-Session-Id)")
print(" - Protocol versioning (Mcp-Protocol-Version)")
print(" - SSE streaming with resumability")
print()
print(" WebSocket:")
print(" - Bidirectional real-time")
print(" - Auto-reconnect with exponential backoff")
print(" - Authentication token support")
print()
print(" SSE (Legacy):")
print(" - Backward compatibility")
print(" - Protocol version 2024-11-05")
@@ -0,0 +1,56 @@
"""
MCP Streamable HTTP Transport Example
Demonstrates connecting to MCP servers via Streamable HTTP transport.
This is the current standard transport (Protocol Revision 2025-11-25).
Features:
- Single MCP endpoint for all communication
- Session management via Mcp-Session-Id header
- Protocol versioning via Mcp-Protocol-Version header
- SSE streaming with resumability
- Session termination via HTTP DELETE
Protocol: MCP 2025-11-25
"""
from praisonaiagents import Agent, MCP
# Basic Streamable HTTP connection
agent = Agent(
name="HTTP Stream Assistant",
instructions="You are a helpful assistant connected via Streamable HTTP.",
tools=MCP("https://api.example.com/mcp")
)
# With authentication headers
agent_auth = Agent(
name="Authenticated Assistant",
instructions="You are a helpful assistant with authenticated connection.",
tools=MCP(
"https://api.example.com/mcp",
headers={"Authorization": "Bearer your-token"}
)
)
# With timeout and debug
agent_debug = Agent(
name="Debug Assistant",
tools=MCP(
"https://api.example.com/mcp",
timeout=120,
debug=True
)
)
if __name__ == "__main__":
print("Streamable HTTP MCP Transport Example")
print("=" * 40)
print("This example requires a Streamable HTTP MCP server running.")
print("URL pattern: https://host:port/mcp")
print()
print("Features enabled:")
print("- Mcp-Protocol-Version header (default: 2025-03-26)")
print("- Mcp-Session-Id header (automatic)")
print("- SSE streaming support")
print("- Resumability via Last-Event-ID")
+55
View File
@@ -0,0 +1,55 @@
"""
MCP WebSocket Transport Example
Demonstrates connecting to MCP servers via WebSocket (ws:// or wss://).
WebSocket transport provides bidirectional real-time communication,
ideal for cloud deployments and long-lived connections.
Features:
- Auto-detection of ws:// and wss:// URLs
- Authentication token support
- Auto-reconnect with exponential backoff
- Ping/pong keepalive
Protocol: SEP-1288 (WebSocket Transport)
"""
from praisonaiagents import Agent, MCP
# Basic WebSocket connection
agent = Agent(
name="WebSocket Assistant",
instructions="You are a helpful assistant connected via WebSocket.",
tools=MCP("ws://localhost:8080/mcp")
)
# Secure WebSocket with authentication
agent_secure = Agent(
name="Secure WebSocket Assistant",
instructions="You are a helpful assistant with secure WebSocket connection.",
tools=MCP(
"wss://api.example.com/mcp",
auth_token="Bearer your-secret-token",
timeout=60
)
)
# WebSocket with custom options
agent_custom = Agent(
name="Custom WebSocket Assistant",
tools=MCP(
"wss://api.example.com/mcp",
auth_token="your-token",
timeout=120,
debug=True
)
)
if __name__ == "__main__":
# Note: Requires a running WebSocket MCP server
# Example server: https://github.com/modelcontextprotocol/servers
print("WebSocket MCP Transport Example")
print("=" * 40)
print("This example requires a WebSocket MCP server running.")
print("URL patterns: ws://host:port/path or wss://host:port/path")
@@ -0,0 +1,44 @@
from praisonaiagents import Agent
import os
# Set OpenAI API key if not already set
if not os.getenv("OPENAI_API_KEY"):
print("Please set your OpenAI API key: export OPENAI_API_KEY='your-api-key-here'")
exit(1)
# Create the RAG agent with web search capabilities
rag_agent = Agent(
instructions="""You are a helpful AI assistant specialized in Thai recipes and cooking.
You have access to a PDF knowledge base about Thai recipes from: https://phi-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf
You can also search the web for additional information about Thai cooking, ingredients, and techniques.
When answering questions:
1. Use your knowledge about Thai cuisine to provide helpful information
2. If needed, search the web for additional details, current information, or clarification
3. Provide comprehensive, helpful answers about Thai cuisine
4. Always be informative and helpful about Thai cooking!
You can use the internet_search function to search the web when needed.""",
llm="gpt-4o",
markdown=True,
verbose=True
)
if __name__ == "__main__":
print("🤖 Thai Recipe RAG Agent is ready!")
print("Ask me anything about Thai recipes or cooking!")
print("Type 'quit' to exit.\n")
while True:
user_input = input("You: ")
if user_input.lower() in ['quit', 'exit', 'bye']:
print("👋 Goodbye! Happy cooking!")
break
try:
response = rag_agent.start(user_input)
print(f"\n🤖 Assistant: {response}\n")
except Exception as e:
print(f"❌ Error: {e}\n")
@@ -0,0 +1,154 @@
# 🧠 Agentic RAG with GPT-5
An agentic RAG application built with the PraisonAI Agents framework, featuring GPT-5 and built-in vector search for efficient knowledge retrieval and question answering.
## ✨ Features
- **🤖 GPT-5**: Latest OpenAI model for intelligent responses
- **🗄️ Built-in Vector Search**: ChromaDB integration for fast similarity search
- **🔍 Agentic RAG**: Intelligent retrieval augmented generation
- **📝 Markdown Formatting**: Beautiful, structured responses
- **🌐 Dynamic Knowledge**: Add URLs to expand knowledge base
- **⚡ Real-time Responses**: Fast answer generation- **🎯 Clean Interface**: Simplified UI without configuration complexity
## 🚀 Quick Start
### Prerequisites
- Python 3.11+
- OpenAI API key with GPT-5 access
### Installation
1. **Clone and navigate to the project**
```bash
cd rag_examples/agentic_rag_gpt5
```
2. **Install dependencies**
```bash
pip install -r requirements.txt
```
3. **Set up your OpenAI API key**
```bash
export OPENAI_API_KEY="your-api-key-here"
```
Or create a `.env` file:
```
OPENAI_API_KEY=your-api-key-here
```
4. **Run the application**
```bash
streamlit run agentic_rag_gpt5.py
```
## 🎯 How to Use
1. **Enter your OpenAI API key** in the sidebar
2. **Add knowledge sources** by entering URLs in the sidebar
3. **Ask questions** using the text area or suggested prompts
4. **Get answers** with markdown formatting
### Suggested Questions
- **"What is PraisonAI?"** - Learn about the PraisonAI Agents framework
- **"Teams in PraisonAI"** - Understand how teams work in PraisonAI
- **"Build RAG system"** - Get a step-by-step guide to building RAG systems
## 🏗️ Architecture
### Core Components
- **`Agent`**: PraisonAI Agents framework for intelligent Q&A
- **`knowledge`**: Built-in knowledge base that handles URLs and documents
- **`llm`**: OpenAI GPT-5-nano for generating responses
- **Built-in Vector Search**: Automatic similarity search without external setup
### Data Flow
1. **Knowledge Loading**: URLs are processed and stored in the built-in vector database
2. **Vector Search**: OpenAI embeddings enable semantic search
3. **Response Generation**: GPT-5-nano processes information and generates answers
4. **Formatted Output**: Markdown-formatted responses
## 🔧 Configuration
### Database Settings
- **Vector DB**: Built-in vector database with automatic indexing
- **Storage**: Local storage managed by PraisonAI Agents
- **Search**: Automatic similarity search
### Model Configuration
- **LLM**: OpenAI GPT-5-nano
- **Embeddings**: Automatic handling by PraisonAI Agents
- **Vector Store**: Built-in with automatic document processing
## 📚 Knowledge Management
### Adding Sources
- Use the sidebar to add new URLs
- Sources are automatically processed and indexed
- Current sources are displayed as numbered list
### Default Knowledge
- Starts with PraisonAI documentation: `https://docs.praisonai.com/introduction/agents.md`
- Expandable with any web-based documentation
## 🎨 UI Features
### Sidebar
- **API Key Management**: Secure input for OpenAI credentials
- **URL Addition**: Dynamic knowledge base expansion
- **Current Sources**: Numbered list of loaded URLs
### Main Interface
- **Suggested Prompts**: Quick access to common questions
- **Query Input**: Large text area for custom questions
- **Fast Responses**: Quick answer generation
- **Markdown Rendering**: Beautiful formatted responses
## 🛠️ Technical Details
### Dependencies
```
streamlit>=1.28.0
praisonaiagents>=0.1.0
openai>=1.0.0
python-dotenv>=1.0.0
```
### Key Features
- **Built-in Knowledge Base**: Automatic document processing and indexing
- **Vector Search**: Efficient similarity search with built-in database
- **Caching**: Efficient resource loading with Streamlit caching
- **Error Handling**: Graceful handling of API and processing errors
## 🔍 Troubleshooting
### Common Issues
**Knowledge base not loading**
- Check OpenAI API key is valid
- Ensure URLs are accessible
- Verify internet connection
**Agent initialization errors**
- Check if PraisonAI Agents is properly installed
- Verify OpenAI API key has sufficient credits
- Ensure Python version is 3.11+
### Performance Tips
- **Cache Resources**: Knowledge base and agent are cached for efficiency
- **Built-in Vector Search**: Fast similarity search without external setup
- **Local Storage**: Optimized local storage for optimal performance
## 🎯 Use Cases
- **Documentation Q&A**: Ask questions about technical documentation
- **Research Assistant**: Get answers from multiple knowledge sources
- **Learning Tool**: Interactive exploration of complex topics
- **Content Discovery**: Find relevant information across multiple sources
**Built with ❤️ using PraisonAI Agents and GPT-5**
@@ -0,0 +1,201 @@
import streamlit as st
import os
from praisonaiagents import Agent
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Page configuration
st.set_page_config(
page_title="Agentic RAG with GPT-5",
page_icon="🧠",
layout="wide"
)
# Main title and description
st.title("🧠 Agentic RAG with GPT-5")
st.markdown("""
This app demonstrates an intelligent AI agent that:
1. **Answers** your questions clearly and concisely using GPT-5
⚠️ **Note**: Knowledge base functionality is temporarily disabled due to a compatibility issue with the current version of PraisonAI Agents.
Enter your OpenAI API key in the sidebar to get started!
""")
# Sidebar for API key and settings
with st.sidebar:
st.header("🔧 Configuration")
# OpenAI API Key
openai_key = st.text_input(
"OpenAI API Key",
type="password",
value=os.getenv("OPENAI_API_KEY", ""),
help="Get your key from https://platform.openai.com/"
)
# Add URLs to knowledge base
st.subheader("🌐 Add Knowledge Sources")
st.info("⚠️ Knowledge base functionality is temporarily disabled due to compatibility issues.")
new_url = st.text_input(
"Add URL",
placeholder="https://docs.praisonai.com/introduction",
help="Enter a URL to add to the knowledge base (currently disabled)"
)
if st.button("➕ Add URL", type="primary", disabled=True):
st.info("Knowledge base functionality is temporarily disabled.")
# Check if API key is provided
if openai_key:
# Initialize knowledge base (cached to avoid reloading)
@st.cache_resource(show_spinner="📚 Loading knowledge base...")
def load_knowledge() -> list:
"""Load and initialize the knowledge base with default URL"""
return ["https://docs.praisonai.com/introduction/agents.md"] # Default URL
# Initialize agent (cached to avoid reloading)
@st.cache_resource(show_spinner="🤖 Loading agent...")
def load_agent(_knowledge: list) -> Agent:
"""Create an agent with reasoning capabilities"""
# Note: Temporarily removed knowledge parameter to avoid rerank error
# TODO: Re-enable when PraisonAI Agents knowledge issue is resolved
return Agent(
name="Knowledge Agent",
instructions=[
"You are a helpful AI assistant. Answer questions based on your general knowledge.",
"Provide clear, well-structured answers in markdown format.",
"Use proper markdown formatting with headers, lists, and emphasis where appropriate.",
"Structure your response with clear sections and bullet points when helpful.",
],
llm="gpt-5-nano
markdown=True,
verbose=True
)
# Load knowledge and agent
knowledge = load_knowledge()
agent = load_agent(knowledge)
# Display current URLs in knowledge base
if knowledge:
st.sidebar.subheader("📚 Current Knowledge Sources")
for i, url in enumerate(knowledge, 1):
st.sidebar.markdown(f"{i}. {url}")
# Handle URL additions
if hasattr(st.session_state, 'urls_to_add') and st.session_state.urls_to_add:
with st.spinner("📥 Loading new documents..."):
knowledge.append(st.session_state.urls_to_add)
# Reinitialize agent with new knowledge
agent = load_agent(knowledge)
st.success(f"✅ Added: {st.session_state.urls_to_add}")
del st.session_state.urls_to_add
st.rerun()
# Main query section
st.divider()
st.subheader("🤔 Ask a Question")
# Suggested prompts
st.markdown("**Try these prompts:**")
col1, col2, col3 = st.columns(3)
with col1:
if st.button("What is PraisonAI?", use_container_width=True):
st.session_state.query = "What is PraisonAI and how do Agents work?"
with col2:
if st.button("Teams in PraisonAI", use_container_width=True):
st.session_state.query = "What are Teams in PraisonAI and how do they work?"
with col3:
if st.button("Build RAG system", use_container_width=True):
st.session_state.query = "Give me a step-by-step guide to building a RAG system."
# Query input
query = st.text_area(
"Your question:",
value=st.session_state.get("query", "What are AI Agents?"),
height=100,
help="Ask anything about the loaded knowledge sources"
)
# Run button
if st.button("🚀 Get Answer", type="primary"):
if query:
# Create container for answer
st.markdown("### 💡 Answer")
answer_container = st.container()
answer_placeholder = answer_container.empty()
# Get the agent's response
with st.spinner("🔍 Searching and generating answer..."):
try:
st.info("🤖 Agent is processing your question...")
response = agent.start(query)
answer_placeholder.markdown(
response,
unsafe_allow_html=True
)
except Exception as e:
st.error(f"Error getting response: {str(e)}")
st.error(f"Error type: {type(e).__name__}")
st.error(f"Full error details: {repr(e)}")
# Try to provide a helpful response
answer_placeholder.markdown("""
**⚠️ Error occurred while processing your question.**
This might be due to:
- Knowledge base configuration issues
- Model access problems
- Network connectivity issues
Please try again or check your OpenAI API key.
""")
else:
st.error("Please enter a question")
else:
# Show instructions if API key is missing
st.info("""
👋 **Welcome! To use this app, you need:**
- **OpenAI API Key** (set it in the sidebar)
- Sign up at [platform.openai.com](https://platform.openai.com/)
- Generate a new API key
Once you enter the key, the app will load the knowledge base and agent.
""")
# Footer with explanation
st.divider()
with st.expander("📖 How This Works"):
st.markdown("""
**This app uses the PraisonAI Agents framework to create an intelligent Q&A system:**
⚠️ **Current Status**: Knowledge base functionality is temporarily disabled due to compatibility issues with the current version of PraisonAI Agents (v0.0.157).
**What Works Now:**
1. **GPT-5 Integration**: OpenAI's GPT-5 model for generating intelligent responses
2. **Markdown Formatting**: Beautiful, structured responses
3. **Agent Framework**: PraisonAI Agents for orchestration
**What's Temporarily Disabled:**
- URL-based knowledge base
- Vector search and retrieval
- Document processing
**Key Components:**
- `Agent`: PraisonAI Agents framework for creating intelligent agents
- `llm`: OpenAI GPT-5-nano for generating responses
**Why PraisonAI Agents?**
- Easy-to-use agent framework
- Built-in error handling
- Clean API design
- Perfect for prototyping and production applications
**Next Steps:**
- Monitor PraisonAI Agents updates for knowledge base fixes
- Re-enable knowledge functionality when compatible version is available
""")
+1
View File
@@ -86,3 +86,4 @@ build
graph.py
chroma_db/
.qodo
*_COMPARISON_ANALYSIS.md
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,48 @@
"""
Gemini Deep Research Agent Example
Uses Google's deep-research-pro model via the Interactions API
for automated research with web search.
Supports two modes:
1. Polling (default): Checks status periodically, shows progress logs
2. Streaming: Real-time output with thinking summaries (recommended)
Requirements:
- GEMINI_API_KEY or GOOGLE_API_KEY environment variable
- google-genai >= 1.55.0 (pip install google-genai)
"""
from dotenv import load_dotenv
load_dotenv()
from praisonaiagents import DeepResearchAgent
agent = DeepResearchAgent(
name="Gemini Researcher",
instructions="""
You are a professional researcher. Focus on:
- Comprehensive analysis with multiple perspectives
- Data-backed insights and trends
- Clear structure with executive summary
""",
model="deep-research-pro", # Auto-detected as Gemini
verbose=True,
)
print(f"Agent: {agent}")
print(f"Provider: {agent.provider}")
# Streaming is enabled by default for real-time progress updates
# This shows thinking summaries and streams the report as it's generated
result = agent.research(
"What is the current price of Ethereum?"
# stream=True is the default, set stream=False to disable
)
# The report is already printed during streaming when verbose=True
# But we can also access it programmatically:
print("\n" + "=" * 60)
print(f"Report length: {len(result.report)} characters")
print(f"Interaction ID: {result.interaction_id}")
print(f"Reasoning steps captured: {len(result.reasoning_steps)}")
print("=" * 60)
@@ -0,0 +1,47 @@
"""
OpenAI Deep Research Agent Example
Uses OpenAI's o3-deep-research or o4-mini-deep-research models
via the Responses API for automated research with web search.
Supports two modes:
1. Non-streaming (default): Waits for completion, returns full report
2. Streaming: Real-time progress with reasoning summaries and web searches
Requirements:
- OPENAI_API_KEY environment variable
"""
from dotenv import load_dotenv
load_dotenv()
from praisonaiagents import DeepResearchAgent
agent = DeepResearchAgent(
name="Research Assistant",
instructions="""
You are a professional researcher. Focus on:
- Data-rich insights with specific figures
- Reliable sources and citations
- Clear, structured responses
""",
model="o4-mini-deep-research", # or "o3-deep-research" for higher quality
verbose=True
)
print(f"Agent: {agent}")
print(f"Provider: {agent.provider}")
# Streaming is enabled by default for real-time progress updates
# Shows: reasoning summaries, web search queries, and report text as generated
result = agent.research(
"What is the current price of Bitcoin?"
# stream=True is the default, set stream=False to disable
)
# The report is already printed during streaming when verbose=True
# But we can also access it programmatically:
print("\n" + "=" * 60)
print(f"Report length: {len(result.report)} characters")
print(f"Citations: {len(result.citations)}")
print(f"Web searches: {len(result.web_searches)}")
print("=" * 60)
@@ -0,0 +1,144 @@
"""Example: Expose custom tools as MCP server.
This example demonstrates how to create an MCP server that exposes
custom Python functions as MCP tools. These tools can then be used
by any MCP client (Claude Desktop, Cursor, etc.).
Usage:
# Run as stdio server (for Claude Desktop, etc.)
python mcp_server_example.py
# Run as SSE server (for web clients)
python mcp_server_example.py --sse --port 8080
"""
import argparse
from typing import List, Dict, Any
# Define your custom tools as regular Python functions
# Type hints and docstrings are used to generate MCP schemas
def search_web(query: str, max_results: int = 5) -> Dict[str, Any]:
"""Search the web for information.
Args:
query: The search query string
max_results: Maximum number of results to return (default: 5)
Returns:
Search results with titles and URLs
"""
# This is a mock implementation - replace with actual search logic
return {
"query": query,
"results": [
{"title": f"Result {i+1} for '{query}'", "url": f"https://example.com/{i+1}"}
for i in range(max_results)
]
}
def calculate(expression: str) -> Dict[str, Any]:
"""Evaluate a mathematical expression.
Args:
expression: A mathematical expression to evaluate (e.g., "2 + 2 * 3")
Returns:
The result of the calculation
"""
try:
# Safe evaluation of mathematical expressions
# In production, use a proper math parser
allowed_chars = set("0123456789+-*/.() ")
if not all(c in allowed_chars for c in expression):
return {"error": "Invalid characters in expression"}
result = eval(expression) # Note: Use safer alternative in production
return {"expression": expression, "result": result}
except Exception as e:
return {"error": str(e)}
def get_weather(city: str, units: str = "celsius") -> Dict[str, Any]:
"""Get current weather for a city.
Args:
city: Name of the city
units: Temperature units - "celsius" or "fahrenheit" (default: celsius)
Returns:
Weather information for the city
"""
# Mock implementation - replace with actual weather API
return {
"city": city,
"temperature": 22 if units == "celsius" else 72,
"units": units,
"condition": "Sunny",
"humidity": 45
}
def list_files(directory: str = ".", pattern: str = "*") -> List[str]:
"""List files in a directory.
Args:
directory: Directory path to list (default: current directory)
pattern: Glob pattern to filter files (default: *)
Returns:
List of file names matching the pattern
"""
import glob
import os
try:
full_pattern = os.path.join(directory, pattern)
files = glob.glob(full_pattern)
return [os.path.basename(f) for f in files]
except Exception as e:
return [f"Error: {str(e)}"]
def main():
parser = argparse.ArgumentParser(description="Run PraisonAI Tools as MCP Server")
parser.add_argument("--sse", action="store_true", help="Use SSE transport instead of stdio")
parser.add_argument("--port", type=int, default=8080, help="Port for SSE server (default: 8080)")
parser.add_argument("--host", type=str, default="0.0.0.0", help="Host for SSE server")
parser.add_argument("--debug", action="store_true", help="Enable debug logging")
args = parser.parse_args()
# Import MCP server components
from praisonaiagents.mcp import ToolsMCPServer
# Create server with custom name
server = ToolsMCPServer(
name="praisonai-custom-tools",
debug=args.debug
)
# Register all tools
server.register_tools([
search_web,
calculate,
get_weather,
list_files
])
# Print registered tools
print(f"📦 Registered {len(server.tools)} tools:")
for name in server.get_tool_names():
print(f" - {name}")
print()
# Run the server
if args.sse:
server.run_sse(host=args.host, port=args.port)
else:
server.run_stdio()
if __name__ == "__main__":
main()
@@ -0,0 +1,94 @@
# PraisonAI Example Plugin
This is a template for creating tool plugins for PraisonAI Agents.
## Installation
```bash
# Install in development mode
pip install -e .
# Or install from PyPI (after publishing)
pip install praisonai-example-plugin
```
## Usage
After installation, tools are automatically discovered:
```python
from praisonaiagents import Agent
# Use tools by name - they're auto-discovered!
agent = Agent(
name="Assistant",
tools=["example_greet", "example_math", "example_reverse"]
)
# Or import and use directly
from praisonai_example_plugin import GreetTool, MathTool, reverse_text
agent = Agent(
name="Assistant",
tools=[GreetTool(), MathTool(), reverse_text]
)
```
## Creating Your Own Plugin
1. **Copy this template**
2. **Rename the package** in `pyproject.toml` and `src/` directory
3. **Add your tools** following the patterns in `__init__.py`
4. **Update entry_points** in `pyproject.toml` to register your tools
### Tool Patterns
**Class-based tool:**
```python
from praisonaiagents import BaseTool
class MyTool(BaseTool):
name = "my_tool"
description = "What this tool does"
def run(self, param: str) -> str:
return f"Result: {param}"
```
**Decorator-based tool:**
```python
from praisonaiagents import tool
@tool
def my_function(param: str) -> str:
"""What this function does."""
return f"Result: {param}"
```
### Registering Tools
In `pyproject.toml`:
```toml
[project.entry-points."praisonaiagents.tools"]
my_tool = "my_package:MyTool"
my_function = "my_package:my_function"
```
## Testing
```bash
pip install -e ".[dev]"
pytest
```
## Publishing
```bash
pip install build twine
python -m build
twine upload dist/*
```
## License
MIT
@@ -0,0 +1,49 @@
[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "praisonai-example-plugin"
version = "1.0.0"
description = "Example plugin for PraisonAI Agents - demonstrates tool plugin development"
readme = "README.md"
license = {text = "MIT"}
requires-python = ">=3.10"
authors = [
{name = "Your Name", email = "[email protected]"}
]
keywords = ["praisonai", "agents", "tools", "plugin"]
classifiers = [
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
]
dependencies = [
"praisonaiagents>=0.0.168",
]
[project.optional-dependencies]
dev = [
"pytest>=7.0.0",
"pytest-mock>=3.10.0",
]
# This is the key part - register your tools as entry points
# The format is: tool_name = "module.path:ClassName" or "module.path:function_name"
[project.entry-points."praisonaiagents.tools"]
example_greet = "praisonai_example_plugin:GreetTool"
example_math = "praisonai_example_plugin:MathTool"
example_reverse = "praisonai_example_plugin:reverse_text"
[project.urls]
Homepage = "https://github.com/yourusername/praisonai-example-plugin"
Documentation = "https://github.com/yourusername/praisonai-example-plugin#readme"
Repository = "https://github.com/yourusername/praisonai-example-plugin"
[tool.setuptools.packages.find]
where = ["src"]
@@ -0,0 +1,119 @@
"""PraisonAI Example Plugin - Demonstrates how to create tool plugins.
This package shows how external developers can create pip-installable
tool plugins for PraisonAI Agents.
Usage after installation:
from praisonaiagents import Agent
# Tools are auto-discovered via entry_points!
agent = Agent(tools=["example_greet", "example_math"])
# Or import directly
from praisonai_example_plugin import GreetTool, MathTool
agent = Agent(tools=[GreetTool(), MathTool()])
"""
from typing import Dict, Any
from praisonaiagents import BaseTool, tool
# =============================================================================
# Class-Based Tools
# =============================================================================
class GreetTool(BaseTool):
"""A tool that generates personalized greetings.
This demonstrates a simple class-based tool with multiple parameters.
"""
name = "example_greet"
description = "Generate a personalized greeting message"
version = "1.0.0"
def run(self, name: str, greeting: str = "Hello", punctuation: str = "!") -> str:
"""Generate a greeting.
Args:
name: The name of the person to greet
greeting: The greeting word (default: "Hello")
punctuation: End punctuation (default: "!")
Returns:
A formatted greeting string
"""
return f"{greeting}, {name}{punctuation}"
class MathTool(BaseTool):
"""A tool for basic mathematical operations.
Demonstrates returning structured data from a tool.
"""
name = "example_math"
description = "Perform basic math operations: add, subtract, multiply, divide"
version = "1.0.0"
def run(
self,
a: float,
b: float,
operation: str = "add"
) -> Dict[str, Any]:
"""Perform a math operation.
Args:
a: First number
b: Second number
operation: One of 'add', 'subtract', 'multiply', 'divide'
Returns:
Dict containing result and operation details
"""
ops = {
"add": ("+", lambda x, y: x + y),
"subtract": ("-", lambda x, y: x - y),
"multiply": ("*", lambda x, y: x * y),
"divide": ("/", lambda x, y: x / y if y != 0 else float('inf')),
}
if operation not in ops:
return {"error": f"Unknown operation: {operation}. Use: {list(ops.keys())}"}
symbol, func = ops[operation]
result = func(a, b)
return {
"result": result,
"expression": f"{a} {symbol} {b} = {result}",
"operation": operation
}
# =============================================================================
# Decorator-Based Tools
# =============================================================================
@tool(name="example_reverse", description="Reverse a string of text")
def reverse_text(text: str) -> str:
"""Reverse the input text.
Args:
text: The text to reverse
Returns:
The reversed text
"""
return text[::-1]
# =============================================================================
# Package Exports
# =============================================================================
__version__ = "1.0.0"
__all__ = [
"GreetTool",
"MathTool",
"reverse_text",
]
+8
View File
@@ -0,0 +1,8 @@
from praisonaiagents import Agent
agent = Agent(
instructions="You are a helpful assistant",
llm="gpt-5-nano"
)
agent.start("Why sky is Blue?")
agent.start("What was my previous question?")
+14
View File
@@ -0,0 +1,14 @@
from praisonaiagents import Agent
import asyncio
agent = Agent(
instructions="You are a helpful assistant",
llm="gpt-5-nano"
)
agent.run("Why sea is Blue?")
async def main():
await agent.arun("Why sky is Blue?")
await agent.arun("What was my previous question?")
asyncio.run(main())
+34
View File
@@ -0,0 +1,34 @@
"""Simple Planning Mode Example: Plan first, then execute todo items."""
from praisonaiagents import Agent
from praisonaiagents.planning import PlanningAgent, TodoList
# Step 1: Create agents
researcher = Agent(name="Researcher", role="Research Analyst", verbose=True)
writer = Agent(name="Writer", role="Content Writer", verbose=True)
# Step 2: Create a plan using PlanningAgent
planner = PlanningAgent(llm="gpt-4o-mini", verbose=1)
plan = planner.create_plan_sync(
request="Write a short article about the top 3 benefits of meditation",
agents=[researcher, writer],
context="Keep it concise"
)
# Step 3: Create todo list from the plan
todo = TodoList.from_plan(plan)
# Step 4: Execute each todo item one by one
for item in todo.items:
todo.start(item.id)
# Pick the right agent based on the item's assigned agent
if item.agent == "Researcher":
agent = researcher
else:
agent = writer
# Execute the task
result = agent.chat(item.description)
todo.complete(item.id)
+21
View File
@@ -0,0 +1,21 @@
"""Simplest Planning Mode Example - Just enable planning=True"""
from praisonaiagents import Agent, Task, PraisonAIAgents
# Create agents
researcher = Agent(name="Researcher", role="Research Analyst")
writer = Agent(name="Writer", role="Content Writer")
# Create tasks
task1 = Task(description="Research the benefits of meditation", agent=researcher)
task2 = Task(description="Write a short article about meditation benefits", agent=writer)
# Run with planning enabled - that's it!
agents = PraisonAIAgents(
agents=[researcher, writer],
tasks=[task1, task2],
planning=True,
auto_approve_plan=True # Auto-approve the generated plan
)
agents.start()
@@ -15,9 +15,33 @@ _logging.configure_root_logger()
from .agent.agent import Agent
from .agent.image_agent import ImageAgent
from .agent.context_agent import ContextAgent, create_context_agent
from .agent.deep_research_agent import (
DeepResearchAgent,
DeepResearchResponse,
Citation,
ReasoningStep,
WebSearchCall,
CodeExecutionStep,
MCPCall,
FileSearchCall,
Provider
)
from .agent.query_rewriter_agent import (
QueryRewriterAgent,
RewriteStrategy,
RewriteResult
)
from .agent.prompt_expander_agent import (
PromptExpanderAgent,
ExpandStrategy,
ExpandResult
)
from .agents.agents import PraisonAIAgents
from .task.task import Task
from .tools.tools import Tools
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
@@ -31,6 +55,32 @@ except ImportError:
from .session import Session
from .memory.memory import Memory
from .guardrails import GuardrailResult, LLMGuardrail
# Planning mode support
try:
from .planning import (
Plan,
PlanStep,
TodoList,
TodoItem,
PlanStorage,
PlanningAgent,
ApprovalCallback,
READ_ONLY_TOOLS,
RESTRICTED_TOOLS,
)
_planning_available = True
except ImportError:
_planning_available = False
Plan = None
PlanStep = None
TodoList = None
TodoItem = None
PlanStorage = None
PlanningAgent = None
ApprovalCallback = None
READ_ONLY_TOOLS = []
RESTRICTED_TOOLS = []
from .agent.handoff import Handoff, handoff, handoff_filters, RECOMMENDED_PROMPT_PREFIX, prompt_with_handoff_instructions
# Flow display
try:
@@ -154,6 +204,17 @@ __all__ = [
'Agents',
'Tools',
'Task',
# Plugin system
'BaseTool',
'ToolResult',
'ToolValidationError',
'validate_tool',
'tool',
'FunctionTool',
'ToolRegistry',
'get_registry',
'register_tool',
'get_tool',
'TaskOutput',
'ReflectionOutput',
'AutoAgents',
@@ -186,7 +247,22 @@ __all__ = [
'disable_performance_mode',
'cleanup_telemetry_resources',
'MinimalTelemetry',
'TelemetryCollector'
'TelemetryCollector',
'DeepResearchAgent',
'DeepResearchResponse',
'Citation',
'ReasoningStep',
'WebSearchCall',
'CodeExecutionStep',
'MCPCall',
'FileSearchCall',
'Provider',
'QueryRewriterAgent',
'RewriteStrategy',
'RewriteResult',
'PromptExpanderAgent',
'ExpandStrategy',
'ExpandResult'
]
# Add MCP to __all__ if available
@@ -4,5 +4,52 @@ from .image_agent import ImageAgent
from .context_agent import ContextAgent, create_context_agent
from .handoff import Handoff, handoff, handoff_filters, RECOMMENDED_PROMPT_PREFIX, prompt_with_handoff_instructions
from .router_agent import RouterAgent
from .deep_research_agent import (
DeepResearchAgent,
DeepResearchResponse,
Citation,
ReasoningStep,
WebSearchCall,
CodeExecutionStep,
MCPCall,
FileSearchCall,
Provider
)
from .query_rewriter_agent import (
QueryRewriterAgent,
RewriteStrategy,
RewriteResult
)
from .prompt_expander_agent import (
PromptExpanderAgent,
ExpandStrategy,
ExpandResult
)
__all__ = ['Agent', 'ImageAgent', 'ContextAgent', 'create_context_agent', 'Handoff', 'handoff', 'handoff_filters', 'RECOMMENDED_PROMPT_PREFIX', 'prompt_with_handoff_instructions', 'RouterAgent']
__all__ = [
'Agent',
'ImageAgent',
'ContextAgent',
'create_context_agent',
'Handoff',
'handoff',
'handoff_filters',
'RECOMMENDED_PROMPT_PREFIX',
'prompt_with_handoff_instructions',
'RouterAgent',
'DeepResearchAgent',
'DeepResearchResponse',
'Citation',
'ReasoningStep',
'WebSearchCall',
'CodeExecutionStep',
'MCPCall',
'FileSearchCall',
'Provider',
'QueryRewriterAgent',
'RewriteStrategy',
'RewriteResult',
'PromptExpanderAgent',
'ExpandStrategy',
'ExpandResult'
]
@@ -234,7 +234,15 @@ class Agent:
max_guardrail_retries: int = 3,
handoffs: Optional[List[Union['Agent', 'Handoff']]] = None,
base_url: Optional[str] = None,
api_key: Optional[str] = None
api_key: Optional[str] = None,
web_search: Optional[Union[bool, Dict[str, Any]]] = None,
web_fetch: Optional[Union[bool, Dict[str, Any]]] = None,
prompt_caching: Optional[bool] = None,
claude_memory: Optional[Union[bool, Any]] = None,
plan_mode: bool = False,
planning: bool = False,
planning_tools: Optional[List[Any]] = None,
planning_reasoning: bool = False
):
"""Initialize an Agent instance.
@@ -400,7 +408,11 @@ class Agent:
model=model_name,
base_url=base_url,
api_key=api_key,
metrics=metrics
metrics=metrics,
web_search=web_search,
web_fetch=web_fetch,
prompt_caching=prompt_caching,
claude_memory=claude_memory
)
self._using_custom_llm = True
except ImportError as e:
@@ -435,6 +447,10 @@ class Agent:
if api_key:
llm_params['api_key'] = api_key
llm_params['metrics'] = metrics
llm_params['web_search'] = web_search
llm_params['web_fetch'] = web_fetch
llm_params['prompt_caching'] = prompt_caching
llm_params['claude_memory'] = claude_memory
self.llm_instance = LLM(**llm_params)
self._using_custom_llm = True
@@ -455,6 +471,15 @@ class Agent:
if callable(tools):
# If a single function/callable is passed, wrap it in a list
self.tools = [tools]
elif isinstance(tools, str):
# Single tool name string - resolve from registry
self.tools = self._resolve_tool_names([tools])
elif isinstance(tools, (list, tuple)):
# Check if list contains strings (tool names) that need resolution
if tools and all(isinstance(t, str) for t in tools):
self.tools = self._resolve_tool_names(tools)
else:
self.tools = list(tools)
else:
# Handle all falsy values (None, False, 0, "", etc.) by defaulting to empty list
self.tools = tools or []
@@ -462,7 +487,8 @@ class Agent:
self.max_iter = max_iter
self.max_rpm = max_rpm
self.max_execution_time = max_execution_time
self.memory = memory
self._memory_instance = None
self._init_memory(memory, user_id)
self.verbose = verbose
self.allow_delegation = allow_delegation
self.step_callback = step_callback
@@ -502,6 +528,32 @@ Your Goal: {self.goal}
# Store user_id
self.user_id = user_id or "praison"
self.reasoning_steps = reasoning_steps
self.plan_mode = plan_mode # Read-only mode for planning
self.planning = planning # Enable planning mode
self.planning_tools = planning_tools # Tools for planning phase
self.planning_reasoning = planning_reasoning # Enable reasoning during planning
self._planning_agent = None # Lazy loaded PlanningAgent
self.web_search = web_search
self.web_fetch = web_fetch
self.prompt_caching = prompt_caching
self.claude_memory = claude_memory
# Initialize rules manager for persistent context (like Cursor/Windsurf)
self._rules_manager = None
self._init_rules_manager()
# Handle web_search fallback: inject DuckDuckGo tool for unsupported models
if web_search and not self._model_supports_web_search():
from ..tools.duckduckgo_tools import internet_search
# Check if internet_search is not already in tools
tool_names = [getattr(t, '__name__', str(t)) for t in self.tools]
if 'internet_search' not in tool_names and 'duckduckgo' not in tool_names:
self.tools.append(internet_search)
logging.info("Model does not support native web search. Added DuckDuckGo fallback tool.")
# Log warning if web_fetch is enabled but model doesn't support it
if web_fetch and not self._model_supports_web_fetch():
logging.warning(f"Model '{self.llm}' does not support native web fetch. Web fetch will be ignored.")
# Initialize guardrail settings
self.guardrail = guardrail
@@ -567,6 +619,294 @@ Your Goal: {self.goal}
self._agent_id = str(uuid.uuid4())
return self._agent_id
def get_available_tools(self) -> List[Any]:
"""
Get tools available to this agent, filtered by plan_mode if enabled.
In plan_mode, only read-only tools are available to prevent
modifications during the planning phase.
Returns:
List of available tools
"""
if not self.plan_mode:
return self.tools
# Filter to read-only tools only
from ..planning import READ_ONLY_TOOLS, RESTRICTED_TOOLS
filtered_tools = []
for tool in self.tools:
tool_name = getattr(tool, '__name__', str(tool)).lower()
# Check if tool is in restricted list
is_restricted = any(
restricted.lower() in tool_name
for restricted in RESTRICTED_TOOLS
)
if not is_restricted:
filtered_tools.append(tool)
return filtered_tools
def _model_supports_web_search(self) -> bool:
"""
Check if the agent's model supports native web search via LiteLLM.
Returns:
bool: True if the model supports native web search, False otherwise
"""
from ..llm.model_capabilities import supports_web_search
# Get the model name
if hasattr(self, 'llm_instance') and self.llm_instance:
model_name = self.llm_instance.model
elif hasattr(self, 'llm') and self.llm:
model_name = self.llm
else:
model_name = "gpt-5-nano"
return supports_web_search(model_name)
def _model_supports_web_fetch(self) -> bool:
"""
Check if the agent's model supports web fetch via LiteLLM.
Web fetch allows the model to retrieve full content from specific URLs.
Currently only supported by Anthropic Claude models.
Returns:
bool: True if the model supports web fetch, False otherwise
"""
from ..llm.model_capabilities import supports_web_fetch
# Get the model name
if hasattr(self, 'llm_instance') and self.llm_instance:
model_name = self.llm_instance.model
elif hasattr(self, 'llm') and self.llm:
model_name = self.llm
else:
model_name = "gpt-5-nano"
return supports_web_fetch(model_name)
def _model_supports_prompt_caching(self) -> bool:
"""
Check if the agent's model supports prompt caching via LiteLLM.
Prompt caching allows caching parts of prompts to reduce costs and latency.
Supported by OpenAI, Anthropic, Bedrock, and Deepseek.
Returns:
bool: True if the model supports prompt caching, False otherwise
"""
from ..llm.model_capabilities import supports_prompt_caching
# Get the model name
if hasattr(self, 'llm_instance') and self.llm_instance:
model_name = self.llm_instance.model
elif hasattr(self, 'llm') and self.llm:
model_name = self.llm
else:
model_name = "gpt-4o-mini"
return supports_prompt_caching(model_name)
def _init_rules_manager(self):
"""
Initialize RulesManager for persistent rules/instructions.
Automatically discovers rules from:
- ~/.praison/rules/ (global)
- .praison/rules/ (workspace)
- Subdirectory rules
"""
try:
from ..memory.rules_manager import RulesManager
import os
# Get workspace path (current working directory)
workspace_path = os.getcwd()
self._rules_manager = RulesManager(
workspace_path=workspace_path,
verbose=1 if self.verbose else 0
)
# Log discovered rules
stats = self._rules_manager.get_stats()
if stats["total_rules"] > 0:
logging.debug(f"RulesManager: Discovered {stats['total_rules']} rules")
except ImportError:
logging.debug("RulesManager not available")
self._rules_manager = None
except Exception as e:
logging.debug(f"Could not initialize RulesManager: {e}")
self._rules_manager = None
def get_rules_context(self, file_path: Optional[str] = None, include_manual: Optional[List[str]] = None) -> str:
"""
Get rules context for the current conversation.
Args:
file_path: Optional file path for glob-based rule matching
include_manual: Optional list of manual rule names to include (via @mention)
Returns:
Formatted rules context string
"""
if not self._rules_manager:
return ""
return self._rules_manager.build_rules_context(
file_path=file_path,
include_manual=include_manual
)
def _init_memory(self, memory, user_id: Optional[str] = None):
"""
Initialize memory based on the memory parameter.
Args:
memory: Can be:
- True: Use FileMemory with default settings
- False/None: No memory
- "file": Use FileMemory
- "sqlite": Use existing Memory class with SQLite
- dict: Configuration for memory
- Memory/FileMemory instance: Use directly
user_id: User identifier for memory isolation
"""
self.memory = memory
if memory is None or memory is False:
self._memory_instance = None
return
# Determine user_id
mem_user_id = user_id or getattr(self, 'user_id', None) or "default"
if memory is True or memory == "file":
# Use FileMemory (zero dependencies)
from ..memory.file_memory import FileMemory
self._memory_instance = FileMemory(
user_id=mem_user_id,
verbose=1 if getattr(self, 'verbose', False) else 0
)
elif isinstance(memory, str) and memory in ("sqlite", "chromadb", "mem0", "mongodb"):
# Use full Memory class with specific provider
try:
from ..memory.memory import Memory
config = {"provider": memory if memory != "sqlite" else "rag"}
self._memory_instance = Memory(config)
except ImportError:
logging.warning(f"Memory provider '{memory}' requires additional dependencies. Falling back to FileMemory.")
from ..memory.file_memory import FileMemory
self._memory_instance = FileMemory(user_id=mem_user_id)
elif isinstance(memory, dict):
# Configuration dict
provider = memory.get("provider", "file")
if provider == "file":
from ..memory.file_memory import FileMemory
self._memory_instance = FileMemory(
user_id=memory.get("user_id", mem_user_id),
config=memory
)
else:
try:
from ..memory.memory import Memory
self._memory_instance = Memory(memory)
except ImportError:
logging.warning("Full Memory class requires additional dependencies. Falling back to FileMemory.")
from ..memory.file_memory import FileMemory
self._memory_instance = FileMemory(user_id=mem_user_id)
else:
# Assume it's already a memory instance
self._memory_instance = memory
def get_memory_context(self, query: Optional[str] = None) -> str:
"""
Get memory context for the current conversation.
Args:
query: Optional query to focus the context
Returns:
Formatted memory context string
"""
if not self._memory_instance:
return ""
if hasattr(self._memory_instance, 'get_context'):
return self._memory_instance.get_context(query=query)
return ""
def store_memory(self, content: str, memory_type: str = "short_term", **kwargs):
"""
Store content in memory.
Args:
content: Content to store
memory_type: Type of memory (short_term, long_term, entity, episodic)
**kwargs: Additional arguments for the memory method
"""
if not self._memory_instance:
return
if memory_type == "short_term" and hasattr(self._memory_instance, 'add_short_term'):
self._memory_instance.add_short_term(content, **kwargs)
elif memory_type == "long_term" and hasattr(self._memory_instance, 'add_long_term'):
self._memory_instance.add_long_term(content, **kwargs)
elif memory_type == "entity" and hasattr(self._memory_instance, 'add_entity'):
self._memory_instance.add_entity(content, **kwargs)
elif memory_type == "episodic" and hasattr(self._memory_instance, 'add_episodic'):
self._memory_instance.add_episodic(content, **kwargs)
def _display_memory_info(self):
"""Display memory information to user in a friendly format."""
if not self._memory_instance:
return
# Only display once per chat session
if hasattr(self, '_memory_displayed') and self._memory_displayed:
return
self._memory_displayed = True
stats = self._memory_instance.get_stats()
short_count = stats.get('short_term_count', 0)
long_count = stats.get('long_term_count', 0)
entity_count = stats.get('entity_count', 0)
storage_path = stats.get('storage_path', '')
total_memories = short_count + long_count + entity_count
if total_memories > 0:
from rich.panel import Panel
from rich.text import Text
# Build memory info text
info_parts = []
if long_count > 0:
info_parts.append(f"💾 {long_count} long-term")
if short_count > 0:
info_parts.append(f"⚡ {short_count} short-term")
if entity_count > 0:
info_parts.append(f"👤 {entity_count} entities")
memory_text = Text()
memory_text.append("🧠 Memory loaded: ", style="bold cyan")
memory_text.append(" | ".join(info_parts))
memory_text.append(f"\n📁 Storage: {storage_path}", style="dim")
self.console.print(Panel(
memory_text,
title="[bold]Agent Memory[/bold]",
border_style="cyan",
expand=False
))
@property
def llm_model(self):
"""Unified property to get the LLM model regardless of configuration type.
@@ -813,17 +1153,35 @@ Your Goal: {self.goal}
if not self.use_system_prompt:
return None
# Check cache first
tools_key = self._get_tools_cache_key(tools)
cache_key = f"{self.role}:{self.goal}:{tools_key}"
if cache_key in self._system_prompt_cache:
return self._system_prompt_cache[cache_key]
# Check cache first (skip cache if memory is enabled since context is dynamic)
if not self._memory_instance:
tools_key = self._get_tools_cache_key(tools)
cache_key = f"{self.role}:{self.goal}:{tools_key}"
if cache_key in self._system_prompt_cache:
return self._system_prompt_cache[cache_key]
else:
cache_key = None # Don't cache when memory is enabled
system_prompt = f"""{self.backstory}\n
Your Role: {self.role}\n
Your Goal: {self.goal}"""
# Add rules context if rules manager is enabled
if self._rules_manager:
rules_context = self.get_rules_context()
if rules_context:
system_prompt += f"\n\n## Rules (Guidelines you must follow)\n{rules_context}"
# Add memory context if memory is enabled
if self._memory_instance:
memory_context = self.get_memory_context()
if memory_context:
system_prompt += f"\n\n## Memory (Information you remember about the user)\n{memory_context}"
# Display memory info to user if verbose
if self.verbose:
self._display_memory_info()
# Add tool usage instructions if tools are available
# Use provided tools or fall back to self.tools
tools_to_use = tools if tools is not None else self.tools
@@ -853,9 +1211,9 @@ Your Goal: {self.goal}"""
if tool_names:
system_prompt += f"\n\nYou have access to the following tools: {', '.join(tool_names)}. Use these tools when appropriate to help complete your tasks. Always use tools when they can help provide accurate information or perform actions."
# Cache the generated system prompt
# Cache the generated system prompt (only if cache_key is set, i.e., memory not enabled)
# Simple cache size limit to prevent unbounded growth
if len(self._system_prompt_cache) < self._max_cache_size:
if cache_key and len(self._system_prompt_cache) < self._max_cache_size:
self._system_prompt_cache[cache_key] = system_prompt
return system_prompt
@@ -1008,6 +1366,31 @@ Your Goal: {self.goal}"""
tools=self.tools
)
def _resolve_tool_names(self, tool_names):
"""Resolve tool names to actual tool instances from registry.
Args:
tool_names: List of tool name strings
Returns:
List of resolved tool instances
"""
resolved = []
try:
from ..tools.registry import get_registry
registry = get_registry()
for name in tool_names:
tool = registry.get(name)
if tool is not None:
resolved.append(tool)
else:
logging.warning(f"Tool '{name}' not found in registry")
except ImportError:
logging.warning("Tool registry not available, cannot resolve tool names")
return resolved
def _cast_arguments(self, func, arguments):
"""Cast arguments to their expected types based on function signature."""
if not callable(func) or not arguments:
@@ -1105,13 +1488,31 @@ Your Goal: {self.goal}"""
# Try to find the function in the agent's tools list first
func = None
for tool in self.tools if isinstance(self.tools, (list, tuple)) else []:
# Check for BaseTool instances (plugin system)
from ..tools.base import BaseTool
if isinstance(tool, BaseTool) and tool.name == function_name:
func = tool
break
# Check for FunctionTool (decorated functions)
if hasattr(tool, 'name') and getattr(tool, 'name', None) == function_name:
func = tool
break
if (callable(tool) and getattr(tool, '__name__', '') == function_name) or \
(inspect.isclass(tool) and tool.__name__ == function_name):
func = tool
break
if func is None:
# If not found in tools, try globals and main
# Check the global tool registry for plugins
try:
from ..tools.registry import get_registry
registry = get_registry()
func = registry.get(function_name)
except ImportError:
pass
if func is None:
# If not found in tools or registry, try globals and main
func = globals().get(function_name)
if not func:
import __main__
@@ -1119,6 +1520,12 @@ Your Goal: {self.goal}"""
if func:
try:
# BaseTool instances (plugin system) - call run() method
from ..tools.base import BaseTool
if isinstance(func, BaseTool):
casted_arguments = self._cast_arguments(func.run, arguments)
return func.run(**casted_arguments)
# Langchain: If it's a class with run but not _run, instantiate and call run
if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
instance = func()
@@ -2082,16 +2489,98 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
display_error(f"Error in _achat_completion: {e}")
return None
async def arun(self, prompt: str, **kwargs):
"""Async alias for astart() method"""
return await self.astart(prompt, **kwargs)
async def astart(self, prompt: str, **kwargs):
"""Async version of start method"""
return await self.achat(prompt, **kwargs)
def run(self):
def run(self, prompt: str, **kwargs):
"""Alias for start() method"""
return self.start()
return self.start(prompt, **kwargs)
def _get_planning_agent(self):
"""Lazy load PlanningAgent for planning mode."""
if self._planning_agent is None and self.planning:
from ..planning import PlanningAgent
self._planning_agent = PlanningAgent(
llm=self.llm if hasattr(self, 'llm') else (self.llm_instance.model if hasattr(self, 'llm_instance') else "gpt-4o-mini"),
tools=self.planning_tools,
reasoning=self.planning_reasoning,
verbose=1 if self.verbose else 0
)
return self._planning_agent
def _start_with_planning(self, prompt: str, **kwargs):
"""Execute with planning mode - creates plan then executes each step."""
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
console = Console()
# Step 1: Create the plan
console.print("\n[bold blue]📋 PLANNING PHASE[/bold blue]")
console.print("[dim]Creating implementation plan...[/dim]\n")
planner = self._get_planning_agent()
plan = planner.create_plan_sync(request=prompt, agents=[self])
if not plan or not plan.steps:
console.print("[yellow]⚠️ Planning failed, falling back to direct execution[/yellow]")
return self.chat(prompt, **kwargs)
# Display the plan
console.print(Panel(
Markdown(plan.to_markdown()),
title="[bold green]Generated Plan[/bold green]",
border_style="green"
))
# Step 2: Execute each step
console.print("\n[bold blue]🚀 EXECUTION PHASE[/bold blue]\n")
results = []
context = ""
for i, step in enumerate(plan.steps):
progress = (i + 1) / len(plan.steps)
bar_length = 30
filled = int(bar_length * progress)
bar = "█" * filled + "░" * (bar_length - filled)
console.print(f"[dim]Progress: [{bar}] {progress * 100:.0f}%[/dim]")
console.print(f"\n[bold]📌 Step {i + 1}/{len(plan.steps)}:[/bold] {step.description[:60]}...")
# Build prompt with context from previous steps
step_prompt = step.description
if context:
step_prompt = f"{step.description}\n\nContext from previous steps:\n{context}"
# Execute the step
result = self.chat(step_prompt, **kwargs)
results.append(result)
# Update context for next step
context += f"\n\nStep {i + 1} result: {result[:500] if result else 'No result'}"
console.print(f" [green]✅ Completed[/green]")
console.print(f"\n[bold green]🎉 EXECUTION COMPLETE[/bold green]")
console.print(f"[dim]Progress: [{'█' * bar_length}] 100%[/dim]")
console.print(f"Completed {len(plan.steps)}/{len(plan.steps)} steps!\n")
# Return the final result
return results[-1] if results else None
def start(self, prompt: str, **kwargs):
"""Start the agent with a prompt. This is a convenience method that wraps chat()."""
# Check if planning mode is enabled
if self.planning:
return self._start_with_planning(prompt, **kwargs)
# Check if streaming is enabled (either from kwargs or agent's stream attribute)
stream_enabled = kwargs.get('stream', getattr(self, 'stream', False))
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,367 @@
"""
Prompt Expander Agent Module
This module provides the PromptExpanderAgent class for expanding short prompts
into detailed, comprehensive prompts for better task execution.
Unlike QueryRewriterAgent (which optimizes queries for search/retrieval),
PromptExpanderAgent focuses on enriching prompts for task execution.
Supported Expansion Strategies:
- **BASIC**: Simple expansion with clarity improvements
- **DETAILED**: Rich expansion with context, constraints, and examples
- **STRUCTURED**: Expansion with clear structure (task, format, requirements)
- **CREATIVE**: Expansion with creative flair and vivid language
- **AUTO**: Automatically selects the best strategy based on prompt analysis
Example:
from praisonaiagents import PromptExpanderAgent, ExpandStrategy
agent = PromptExpanderAgent(verbose=True)
# Basic expansion
result = agent.expand("write a movie script in 3 lines")
print(result.expanded_prompt)
# Detailed expansion
result = agent.expand("blog about AI", strategy=ExpandStrategy.DETAILED)
print(result.expanded_prompt)
# With tools for context gathering
from praisonaiagents.tools import internet_search
agent = PromptExpanderAgent(tools=[internet_search], verbose=True)
result = agent.expand("latest AI developments")
"""
import logging
from typing import List, Optional, Any, Dict
from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime
class ExpandStrategy(Enum):
"""Enumeration of available prompt expansion strategies."""
BASIC = "basic"
DETAILED = "detailed"
STRUCTURED = "structured"
CREATIVE = "creative"
AUTO = "auto"
@dataclass
class ExpandResult:
"""Result of a prompt expansion operation."""
original_prompt: str
expanded_prompt: str
strategy_used: ExpandStrategy
metadata: Dict[str, Any] = field(default_factory=dict)
def __repr__(self):
return f"ExpandResult(strategy={self.strategy_used.value}, original_len={len(self.original_prompt)}, expanded_len={len(self.expanded_prompt)})"
class PromptExpanderAgent:
"""
Agent for expanding short prompts into detailed, comprehensive prompts.
This agent transforms brief task descriptions into rich, detailed prompts
that provide better context and guidance for task execution.
Key difference from QueryRewriterAgent:
- QueryRewriterAgent: Optimizes queries for search/retrieval (RAG)
- PromptExpanderAgent: Expands prompts for detailed task execution
Attributes:
name: Name of the agent
model: LLM model to use for expansion
verbose: Whether to print detailed logs
temperature: Temperature for LLM generation (higher = more creative)
tools: Optional tools for context gathering
Example:
agent = PromptExpanderAgent(model="gpt-4o-mini", verbose=True)
result = agent.expand("write a poem", strategy=ExpandStrategy.CREATIVE)
print(result.expanded_prompt)
"""
# Default prompts for each strategy
PROMPTS = {
"basic": """You are an expert at improving task prompts.
Current date: {current_date}
Given the following task prompt, expand it to be clearer and more actionable while preserving the original intent.
Guidelines:
- Keep it as a TASK/COMMAND (do NOT convert to a question)
- Add clarity without changing the core request
- Fix any ambiguity
- Keep it concise but complete
- Use the current date for any time-sensitive context
Original prompt: {prompt}
Expanded prompt:""",
"detailed": """You are an expert at creating comprehensive task prompts.
Current date: {current_date}
Given the following brief task prompt, expand it into a detailed, actionable prompt that provides rich context and clear guidance.
Guidelines:
- Keep it as a TASK/COMMAND (do NOT convert to a question)
- Add specific details about what to include
- Specify quality expectations
- Add relevant constraints or requirements
- Include format guidance if applicable
- Mention style or tone if relevant
- Use the current date for any time-sensitive context
Original prompt: {prompt}
Detailed expanded prompt:""",
"structured": """You are an expert at creating well-structured task prompts.
Current date: {current_date}
Given the following task prompt, expand it into a structured format with clear sections.
Guidelines:
- Keep it as a TASK/COMMAND (do NOT convert to a question)
- Use a clear structure:
* Task: What to do
* Format: Expected output format
* Requirements: Specific requirements
* Style: Tone or style guidance
* Constraints: Any limitations
- Be specific and actionable
- Use the current date for any time-sensitive context
Original prompt: {prompt}
Structured expanded prompt:""",
"creative": """You are an expert at crafting inspiring, creative task prompts.
Current date: {current_date}
Given the following task prompt, expand it with vivid language and creative direction while keeping it actionable.
Guidelines:
- Keep it as a TASK/COMMAND (do NOT convert to a question)
- Use evocative, inspiring language
- Paint a picture of the desired outcome
- Add creative direction and artistic guidance
- Make it exciting and engaging
- Preserve the core task intent
- Use the current date for any time-sensitive context
Original prompt: {prompt}
Creative expanded prompt:""",
"auto": """You are an expert at analyzing and expanding task prompts.
Current date: {current_date}
Analyze the following prompt and expand it appropriately based on its nature.
For short/vague prompts: Add significant detail and context
For task prompts: Preserve the action-oriented nature
For creative tasks: Add artistic direction
For technical tasks: Add precision and requirements
IMPORTANT: Keep it as a TASK/COMMAND. Never convert to a question.
Use the current date for any time-sensitive context.
Original prompt: {prompt}
{context}
Expanded prompt:"""
}
def __init__(
self,
name: str = "PromptExpanderAgent",
model: str = "gpt-4o-mini",
instructions: Optional[str] = None,
verbose: bool = False,
temperature: float = 0.7, # Higher for creativity
max_tokens: int = 1000,
tools: Optional[List[Any]] = None
):
"""
Initialize the PromptExpanderAgent.
Args:
name: Name of the agent
model: LLM model to use (default: gpt-4o-mini)
instructions: Custom instructions for the agent
verbose: Whether to print detailed logs
temperature: Temperature for LLM generation (default: 0.7 for creativity)
max_tokens: Maximum tokens for LLM response
tools: Optional list of tools for context gathering
"""
self.name = name
self.model = model
self.instructions = instructions
self.verbose = verbose
self.temperature = temperature
self.max_tokens = max_tokens
self.tools = tools or []
self._agent = None # Lazy initialized
# Setup logging
self.logger = logging.getLogger(__name__)
if self.verbose:
self.logger.setLevel(logging.DEBUG)
@property
def agent(self):
"""Lazy initialization of internal Agent."""
if self._agent is None:
from .agent import Agent
self._agent = Agent(
name=self.name,
role="Prompt Expansion Specialist",
goal="Expand brief prompts into detailed, actionable prompts for better task execution. Use available tools when needed to gather context.",
backstory="You are an expert at understanding user intent and transforming brief requests into comprehensive, detailed prompts.",
tools=self.tools,
llm=self.model,
verbose=self.verbose,
markdown=False
)
return self._agent
def _call_agent(self, prompt: str) -> str:
"""Call the internal Agent with the given prompt."""
try:
response = self.agent.chat(prompt)
return response if response else ""
except Exception as e:
self.logger.error(f"Agent call failed: {e}")
raise
def _is_short_prompt(self, prompt: str, threshold: int = 10) -> bool:
"""Check if prompt is considered short."""
return len(prompt.split()) < threshold
def _is_task_prompt(self, prompt: str) -> bool:
"""Check if prompt is a task/command (vs a question)."""
task_verbs = [
"write", "create", "generate", "build", "make", "design",
"develop", "implement", "compose", "draft", "produce",
"craft", "construct", "prepare", "draw", "paint", "code"
]
prompt_lower = prompt.lower().strip()
return any(prompt_lower.startswith(verb) for verb in task_verbs)
def _is_creative_task(self, prompt: str) -> bool:
"""Check if prompt is a creative task."""
creative_indicators = [
"poem", "story", "script", "song", "art", "creative",
"fiction", "novel", "essay", "blog", "article", "write"
]
prompt_lower = prompt.lower()
return any(ind in prompt_lower for ind in creative_indicators)
def _detect_strategy(self, prompt: str) -> ExpandStrategy:
"""Auto-detect the best expansion strategy for a prompt."""
# Very short prompts need detailed expansion
if self._is_short_prompt(prompt, threshold=5):
return ExpandStrategy.DETAILED
# Creative tasks benefit from creative expansion
if self._is_creative_task(prompt):
return ExpandStrategy.CREATIVE
# Task prompts get structured expansion
if self._is_task_prompt(prompt):
return ExpandStrategy.STRUCTURED
# Default to basic
return ExpandStrategy.BASIC
def expand(
self,
prompt: str,
strategy: ExpandStrategy = ExpandStrategy.AUTO,
context: Optional[str] = None
) -> ExpandResult:
"""
Expand a prompt using the specified strategy.
All operations go through an internal Agent which automatically handles
tool calling when tools are provided. The agent decides when to use tools
based on the prompt context.
Args:
prompt: The original user prompt
strategy: Expansion strategy to use
context: Additional context to help with expansion
Returns:
ExpandResult containing the expanded prompt
"""
if self.verbose:
self.logger.debug(f"Expanding prompt: '{prompt}' with strategy: {strategy.value}")
if self.tools:
print(f"[cyan]Agent has {len(self.tools)} tools available (will use if needed)...[/cyan]")
# Auto-detect strategy if needed
if strategy == ExpandStrategy.AUTO:
strategy = self._detect_strategy(prompt)
if self.verbose:
self.logger.debug(f"Auto-detected strategy: {strategy.value}")
# Build the expansion prompt
if strategy == ExpandStrategy.AUTO:
# Fallback if auto-detect returns AUTO
strategy = ExpandStrategy.DETAILED
# Get current date for time-sensitive context
current_date = datetime.now().strftime("%B %d, %Y")
expansion_prompt = self.PROMPTS[strategy.value].format(
prompt=prompt,
context=context or "",
current_date=current_date
)
# Call the agent
expanded = self._call_agent(expansion_prompt).strip()
# If expansion failed, return original
if not expanded:
expanded = prompt
return ExpandResult(
original_prompt=prompt,
expanded_prompt=expanded,
strategy_used=strategy,
metadata={
"original_length": len(prompt),
"expanded_length": len(expanded),
"has_context": bool(context),
"has_tools": bool(self.tools)
}
)
# Convenience methods for direct strategy access
def expand_basic(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
"""Convenience method for basic expansion."""
return self.expand(prompt, strategy=ExpandStrategy.BASIC, context=context)
def expand_detailed(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
"""Convenience method for detailed expansion."""
return self.expand(prompt, strategy=ExpandStrategy.DETAILED, context=context)
def expand_structured(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
"""Convenience method for structured expansion."""
return self.expand(prompt, strategy=ExpandStrategy.STRUCTURED, context=context)
def expand_creative(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
"""Convenience method for creative expansion."""
return self.expand(prompt, strategy=ExpandStrategy.CREATIVE, context=context)
@@ -0,0 +1,632 @@
"""
Query Rewriter Agent Module
This module provides the QueryRewriterAgent class for transforming user queries
to improve retrieval quality in RAG applications.
Supported Rewriting Strategies:
- **BASIC**: Simple rephrasing for clarity and keyword optimization
- **HYDE**: Hypothetical Document Embeddings - generates a hypothetical answer
- **STEP_BACK**: Generates higher-level concept questions for complex queries
- **SUB_QUERIES**: Decomposes multi-part questions into focused sub-queries
- **MULTI_QUERY**: Generates multiple paraphrased versions for ensemble retrieval
- **CONTEXTUAL**: Uses conversation history to resolve references and context
Example:
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent(verbose=True)
# Basic rewriting
result = agent.rewrite("AI trends")
print(result.rewritten_queries)
# HyDE for better semantic matching
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
print(result.hypothetical_document)
# Contextual with chat history
chat_history = [
{"role": "user", "content": "Tell me about Python"},
{"role": "assistant", "content": "Python is a programming language..."}
]
result = agent.rewrite("What about its performance?", chat_history=chat_history)
# With search tools for context-aware rewriting
from praisonaiagents.tools import internet_search
agent = QueryRewriterAgent(tools=[internet_search], verbose=True)
result = agent.rewrite("latest AI developments") # Searches first, then rewrites
"""
import logging
import json
from typing import List, Optional, Any, Dict
from dataclasses import dataclass, field
from enum import Enum
class RewriteStrategy(Enum):
"""Enumeration of available query rewriting strategies."""
BASIC = "basic"
HYDE = "hyde"
STEP_BACK = "step_back"
SUB_QUERIES = "sub_queries"
MULTI_QUERY = "multi_query"
CONTEXTUAL = "contextual"
AUTO = "auto"
@dataclass
class RewriteResult:
"""Result of a query rewriting operation."""
original_query: str
rewritten_queries: List[str]
strategy_used: RewriteStrategy
hypothetical_document: Optional[str] = None
step_back_question: Optional[str] = None
sub_queries: Optional[List[str]] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def __repr__(self):
return f"RewriteResult(strategy={self.strategy_used.value}, queries={len(self.rewritten_queries)})"
@property
def primary_query(self) -> str:
"""Returns the primary rewritten query."""
return self.rewritten_queries[0] if self.rewritten_queries else self.original_query
@property
def all_queries(self) -> List[str]:
"""Returns all queries including original and rewritten."""
queries = [self.original_query]
queries.extend(self.rewritten_queries)
if self.step_back_question:
queries.append(self.step_back_question)
if self.sub_queries:
queries.extend(self.sub_queries)
return list(set(queries)) # Remove duplicates
class QueryRewriterAgent:
"""
Agent for rewriting queries to improve retrieval quality in RAG applications.
This agent transforms user queries using various strategies to bridge the gap
between how users ask questions and how information is stored in knowledge bases.
Attributes:
name: Name of the agent
model: LLM model to use for rewriting
verbose: Whether to print detailed logs
max_queries: Maximum number of queries to generate for multi-query strategy
abbreviations: Dictionary of common abbreviations to expand
Example:
agent = QueryRewriterAgent(model="gpt-4o-mini", verbose=True)
result = agent.rewrite("ML best practices", strategy=RewriteStrategy.MULTI_QUERY)
for query in result.rewritten_queries:
print(query)
"""
# Default prompts for each strategy
PROMPTS = {
"basic": """You are an expert at improving prompts and queries.
Given the following input, improve it while preserving its original intent and type.
Guidelines:
- If it's a TASK (e.g., "write X", "create Y", "generate Z"), keep it as a task - do NOT convert to a question
- If it's a QUESTION (e.g., "what is X?", "how does Y work?"), keep it as a question
- Expand abbreviations and acronyms
- Fix typos and grammatical errors
- Add relevant context if too short
- Keep the original intent and action intact
IMPORTANT: If the input is a command/task to DO something, the output must also be a command/task to DO that thing.
Input: {query}
Improved version:""",
"hyde": """You are an expert at generating hypothetical documents.
Given the following query, write a detailed hypothetical document that would perfectly answer this query. This document will be used to find similar real documents in a knowledge base.
Guidelines:
- Write as if you are answering the query directly
- Include specific details, facts, and terminology
- Use the same style and vocabulary as the target documents
- Be comprehensive but focused on the query topic
Query: {query}
Hypothetical document:""",
"step_back": """You are an expert at generating step-back questions.
Given the following query, generate a more general, higher-level question that captures the broader concept. This helps retrieve background information needed to answer the specific query.
Guidelines:
- Identify the underlying concept or principle
- Create a question about fundamentals or general knowledge
- The step-back question should provide context for the original query
Query: {query}
Step-back question:""",
"sub_queries": """You are an expert at decomposing complex queries.
Given the following query, break it down into simpler, focused sub-queries. Each sub-query should address one specific aspect of the original question.
Guidelines:
- Identify distinct aspects or components of the query
- Create independent, self-contained sub-queries
- Each sub-query should be answerable on its own
- Return as a JSON array of strings
Query: {query}
Sub-queries (JSON array):""",
"multi_query": """You are an expert at query expansion and paraphrasing.
Given the following query, generate {num_queries} different versions that capture the same intent but use different words, phrases, or perspectives.
Guidelines:
- Use synonyms and alternative phrasings
- Try different question formats (what, how, why, etc.)
- Include both formal and informal versions
- Add relevant context or specificity
- Return as a JSON array of strings
Query: {query}
Alternative queries (JSON array):""",
"contextual": """You are an expert at understanding conversational context.
Given the conversation history and a follow-up query, rewrite the query as a standalone question that captures all relevant context.
Guidelines:
- Resolve pronouns and references (it, this, that, etc.)
- Include relevant context from the conversation
- Keep entity names and important details
- Make the query self-contained and clear
Conversation history:
{chat_history}
Follow-up query: {query}
Standalone question:""",
"auto": """You are an expert at analyzing and improving prompts and queries.
Analyze the following input and improve it while preserving its original intent and type.
Input types to consider:
1. TASK/COMMAND (write, create, generate, build) → Keep as a task, make it clearer
2. QUESTION (what, how, why, when) → Keep as a question, make it more specific
3. Short/keyword inputs → Expand with context but preserve type
4. Complex multi-part inputs → Clarify structure
5. Well-formed inputs → Minor optimization only
IMPORTANT: Never convert a task into a question. "Write X" should remain "Write X", not "How to write X?"
Input: {query}
{context}
Improved version:"""
}
def __init__(
self,
name: str = "QueryRewriterAgent",
model: str = "gpt-4o-mini",
instructions: Optional[str] = None,
verbose: bool = False,
max_queries: int = 5,
abbreviations: Optional[Dict[str, str]] = None,
temperature: float = 0.3,
max_tokens: int = 500,
tools: Optional[List[Any]] = None
):
"""
Initialize the QueryRewriterAgent.
Args:
name: Name of the agent
model: LLM model to use (default: gpt-4o-mini)
instructions: Custom instructions for the agent
verbose: Whether to print detailed logs
max_queries: Maximum queries for multi-query strategy
abbreviations: Custom abbreviation expansions
temperature: Temperature for LLM generation
max_tokens: Maximum tokens for LLM response
tools: Optional list of tools (agent decides when to use them)
"""
self.name = name
self.model = model
self.instructions = instructions
self.verbose = verbose
self.max_queries = max_queries
self.temperature = temperature
self.max_tokens = max_tokens
self.tools = tools or []
self._agent = None # Lazy initialized
# Default abbreviations (can be extended)
self.abbreviations = abbreviations or {
"AI": "Artificial Intelligence",
"ML": "Machine Learning",
"NLP": "Natural Language Processing",
"LLM": "Large Language Model",
"RAG": "Retrieval Augmented Generation",
"API": "Application Programming Interface",
"DB": "Database",
"UI": "User Interface",
"UX": "User Experience",
"SDK": "Software Development Kit",
"CLI": "Command Line Interface",
}
# Setup logging
self.logger = logging.getLogger(__name__)
if self.verbose:
self.logger.setLevel(logging.DEBUG)
@property
def agent(self):
"""Lazy initialization of internal Agent."""
if self._agent is None:
from .agent import Agent
self._agent = Agent(
name=self.name,
role="Query Rewriting Specialist",
goal="Rewrite queries for better search and retrieval results. Use available tools when needed to gather context.",
backstory="You are an expert at understanding user intent and transforming queries for optimal information retrieval.",
tools=self.tools,
llm=self.model,
verbose=self.verbose,
markdown=False
)
return self._agent
def _call_agent(self, prompt: str) -> str:
"""Call the internal Agent with the given prompt."""
try:
response = self.agent.chat(prompt)
return response if response else ""
except Exception as e:
self.logger.error(f"Agent call failed: {e}")
raise
def _expand_abbreviations(self, query: str) -> str:
"""Expand known abbreviations in the query."""
words = query.split()
expanded = []
for word in words:
clean_word = word.strip(".,!?;:")
if clean_word.upper() in self.abbreviations:
expanded.append(f"{clean_word} ({self.abbreviations[clean_word.upper()]})")
else:
expanded.append(word)
return " ".join(expanded)
def _is_short_query(self, query: str, threshold: int = 5) -> bool:
"""Check if query is considered short (keyword-style)."""
return len(query.split()) < threshold
def _parse_json_array(self, text: str) -> List[str]:
"""Parse a JSON array from LLM response."""
try:
# Try direct JSON parsing
result = json.loads(text)
if isinstance(result, list):
return [str(item) for item in result]
except json.JSONDecodeError:
pass
# Try to extract JSON array from text
import re
match = re.search(r'\[.*?\]', text, re.DOTALL)
if match:
try:
result = json.loads(match.group())
if isinstance(result, list):
return [str(item) for item in result]
except json.JSONDecodeError:
pass
# Fallback: split by newlines and clean up
lines = [line.strip() for line in text.split('\n') if line.strip()]
lines = [line.lstrip('0123456789.-) ').strip('"\'') for line in lines]
return [line for line in lines if line and len(line) > 3]
def rewrite(
self,
query: str,
strategy: RewriteStrategy = RewriteStrategy.AUTO,
chat_history: Optional[List[Dict[str, str]]] = None,
context: Optional[str] = None,
num_queries: int = None
) -> RewriteResult:
"""
Rewrite a query using the specified strategy.
All operations go through an internal Agent which automatically handles
tool calling when tools are provided. The agent decides when to use tools
based on the query context.
Args:
query: The original user query
strategy: Rewriting strategy to use
chat_history: Previous conversation messages for contextual rewriting
context: Additional context about the knowledge base
num_queries: Number of queries for multi-query strategy
Returns:
RewriteResult containing rewritten queries
"""
if self.verbose:
self.logger.debug(f"Rewriting query: '{query}' with strategy: {strategy.value}")
if self.tools:
print(f"[cyan]Agent has {len(self.tools)} tools available (will use if needed)...[/cyan]")
# Auto-detect strategy if needed
if strategy == RewriteStrategy.AUTO:
strategy = self._detect_strategy(query, chat_history)
if self.verbose:
self.logger.debug(f"Auto-detected strategy: {strategy.value}")
# All calls go through the internal Agent which handles tool calling
combined_context = context or ""
# Dispatch to appropriate method
if strategy == RewriteStrategy.BASIC:
return self._rewrite_basic(query, combined_context)
elif strategy == RewriteStrategy.HYDE:
return self._rewrite_hyde(query, combined_context)
elif strategy == RewriteStrategy.STEP_BACK:
return self._rewrite_step_back(query, combined_context)
elif strategy == RewriteStrategy.SUB_QUERIES:
return self._rewrite_sub_queries(query, combined_context)
elif strategy == RewriteStrategy.MULTI_QUERY:
return self._rewrite_multi_query(query, num_queries or self.max_queries, combined_context)
elif strategy == RewriteStrategy.CONTEXTUAL:
return self._rewrite_contextual(query, chat_history or [], combined_context)
else:
# Default to basic
return self._rewrite_basic(query, combined_context)
def _detect_strategy(
self,
query: str,
chat_history: Optional[List[Dict[str, str]]] = None
) -> RewriteStrategy:
"""Auto-detect the best rewriting strategy for a query."""
# If chat history exists and query seems like a follow-up
if chat_history and self._is_follow_up(query):
return RewriteStrategy.CONTEXTUAL
# Short keyword queries
if self._is_short_query(query):
return RewriteStrategy.BASIC
# Complex multi-part queries
if self._is_complex_query(query):
return RewriteStrategy.SUB_QUERIES
# Questions that might benefit from step-back
if self._needs_step_back(query):
return RewriteStrategy.STEP_BACK
# Default to basic
return RewriteStrategy.BASIC
def _is_follow_up(self, query: str) -> bool:
"""Check if query appears to be a follow-up question."""
follow_up_indicators = [
"it", "this", "that", "they", "them", "their",
"what about", "how about", "and", "also", "more",
"same", "similar", "related"
]
query_lower = query.lower()
return any(indicator in query_lower for indicator in follow_up_indicators)
def _is_complex_query(self, query: str) -> bool:
"""Check if query contains multiple distinct questions."""
# Check for multiple question marks or conjunctions
if query.count('?') > 1:
return True
complex_indicators = [" and ", " or ", " also ", " as well as ", " plus "]
return any(ind in query.lower() for ind in complex_indicators)
def _needs_step_back(self, query: str) -> bool:
"""Check if query would benefit from a step-back question."""
# Queries about specific details often benefit from broader context
specific_indicators = [
"specific", "exactly", "precisely", "particular",
"difference between", "compare", "versus", "vs"
]
return any(ind in query.lower() for ind in specific_indicators)
def _rewrite_basic(self, query: str, context: str = "") -> RewriteResult:
"""Basic query rewriting for clarity and keyword optimization."""
# First expand abbreviations
expanded_query = self._expand_abbreviations(query)
# Build prompt with optional context
prompt = self.PROMPTS["basic"].format(query=expanded_query)
if context:
prompt += f"\n\nAdditional context to help with rewriting:\n{context}"
rewritten = self._call_agent(prompt).strip()
return RewriteResult(
original_query=query,
rewritten_queries=[rewritten],
strategy_used=RewriteStrategy.BASIC,
metadata={"expanded_abbreviations": expanded_query != query, "has_context": bool(context)}
)
def _rewrite_hyde(self, query: str, context: str = "") -> RewriteResult:
"""HyDE: Generate hypothetical document for better semantic matching."""
prompt = self.PROMPTS["hyde"].format(query=query)
if context:
prompt += f"\n\nUse this context to make the hypothetical document more accurate:\n{context}"
hypothetical_doc = self._call_agent(prompt).strip()
# The hypothetical document itself becomes the search query
return RewriteResult(
original_query=query,
rewritten_queries=[hypothetical_doc[:500]], # Truncate for embedding
strategy_used=RewriteStrategy.HYDE,
hypothetical_document=hypothetical_doc,
metadata={"full_document_length": len(hypothetical_doc), "has_context": bool(context)}
)
def _rewrite_step_back(self, query: str, context: str = "") -> RewriteResult:
"""Generate step-back question for broader context retrieval."""
# Generate step-back question
prompt = self.PROMPTS["step_back"].format(query=query)
if context:
prompt += f"\n\nContext:\n{context}"
step_back = self._call_agent(prompt).strip()
# Also do basic rewriting of original
basic_result = self._rewrite_basic(query, context)
return RewriteResult(
original_query=query,
rewritten_queries=basic_result.rewritten_queries,
strategy_used=RewriteStrategy.STEP_BACK,
step_back_question=step_back,
metadata={"includes_step_back": True, "has_context": bool(context)}
)
def _rewrite_sub_queries(self, query: str, context: str = "") -> RewriteResult:
"""Decompose complex query into focused sub-queries."""
prompt = self.PROMPTS["sub_queries"].format(query=query)
if context:
prompt += f"\n\nContext to help with decomposition:\n{context}"
response = self._call_agent(prompt).strip()
sub_queries = self._parse_json_array(response)
# Ensure we have at least one sub-query
if not sub_queries:
sub_queries = [query]
return RewriteResult(
original_query=query,
rewritten_queries=sub_queries,
strategy_used=RewriteStrategy.SUB_QUERIES,
sub_queries=sub_queries,
metadata={"num_sub_queries": len(sub_queries), "has_context": bool(context)}
)
def _rewrite_multi_query(self, query: str, num_queries: int, context: str = "") -> RewriteResult:
"""Generate multiple paraphrased versions for ensemble retrieval."""
prompt = self.PROMPTS["multi_query"].format(
query=query,
num_queries=num_queries
)
if context:
prompt += f"\n\nContext to help with paraphrasing:\n{context}"
response = self._call_agent(prompt).strip()
queries = self._parse_json_array(response)
# Ensure we have queries and limit to max
if not queries:
queries = [query]
queries = queries[:num_queries]
return RewriteResult(
original_query=query,
rewritten_queries=queries,
strategy_used=RewriteStrategy.MULTI_QUERY,
metadata={"num_generated": len(queries), "has_context": bool(context)}
)
def _rewrite_contextual(
self,
query: str,
chat_history: List[Dict[str, str]],
context: str = ""
) -> RewriteResult:
"""Rewrite query using conversation context."""
# Format chat history
history_str = ""
for msg in chat_history[-10:]: # Last 10 messages
role = msg.get("role", "user")
content = msg.get("content", "")
history_str += f"{role.capitalize()}: {content}\n"
prompt = self.PROMPTS["contextual"].format(
chat_history=history_str,
query=query
)
if context:
prompt += f"\n\nAdditional context:\n{context}"
rewritten = self._call_agent(prompt).strip()
return RewriteResult(
original_query=query,
rewritten_queries=[rewritten],
strategy_used=RewriteStrategy.CONTEXTUAL,
metadata={
"history_length": len(chat_history),
"history_used": min(10, len(chat_history)),
"has_context": bool(context)
}
)
# Convenience methods for direct strategy access
def rewrite_basic(self, query: str) -> RewriteResult:
"""Convenience method for basic rewriting."""
return self.rewrite(query, strategy=RewriteStrategy.BASIC)
def rewrite_hyde(self, query: str) -> RewriteResult:
"""Convenience method for HyDE rewriting."""
return self.rewrite(query, strategy=RewriteStrategy.HYDE)
def rewrite_step_back(self, query: str) -> RewriteResult:
"""Convenience method for step-back rewriting."""
return self.rewrite(query, strategy=RewriteStrategy.STEP_BACK)
def rewrite_sub_queries(self, query: str) -> RewriteResult:
"""Convenience method for sub-query decomposition."""
return self.rewrite(query, strategy=RewriteStrategy.SUB_QUERIES)
def rewrite_multi_query(self, query: str, num_queries: int = None) -> RewriteResult:
"""Convenience method for multi-query generation."""
return self.rewrite(
query,
strategy=RewriteStrategy.MULTI_QUERY,
num_queries=num_queries
)
def rewrite_contextual(
self,
query: str,
chat_history: List[Dict[str, str]]
) -> RewriteResult:
"""Convenience method for contextual rewriting."""
return self.rewrite(
query,
strategy=RewriteStrategy.CONTEXTUAL,
chat_history=chat_history
)
def add_abbreviation(self, abbrev: str, expansion: str) -> None:
"""Add a custom abbreviation expansion."""
self.abbreviations[abbrev.upper()] = expansion
def add_abbreviations(self, abbreviations: Dict[str, str]) -> None:
"""Add multiple custom abbreviation expansions."""
for abbrev, expansion in abbreviations.items():
self.abbreviations[abbrev.upper()] = expansion
@@ -127,7 +127,7 @@ def process_task_context(context_item, verbose=0, user_id=None):
return str(context_item) # Fallback for unknown types
class PraisonAIAgents:
def __init__(self, agents, tasks=None, verbose=0, completion_checker=None, max_retries=5, process="sequential", manager_llm=None, memory=False, memory_config=None, embedder=None, user_id=None, max_iter=10, stream=True, name: Optional[str] = None):
def __init__(self, agents, tasks=None, verbose=0, completion_checker=None, max_retries=5, process="sequential", manager_llm=None, memory=False, memory_config=None, embedder=None, user_id=None, max_iter=10, stream=True, name: Optional[str] = None, planning: bool = False, planning_llm: Optional[str] = None, auto_approve_plan: bool = False, planning_tools: Optional[List] = None, planning_reasoning: bool = False):
# Add check at the start if memory is requested
if memory:
try:
@@ -213,6 +213,16 @@ class PraisonAIAgents:
self._state = {} # Add state storage at PraisonAIAgents level
# Initialize planning mode
self.planning = planning
self.planning_llm = planning_llm or "gpt-4o-mini"
self.auto_approve_plan = auto_approve_plan
self.planning_tools = planning_tools
self.planning_reasoning = planning_reasoning
self._current_plan = None
self._todo_list = None
self._planning_agent = None
# Initialize memory system
self.shared_memory = None
if memory:
@@ -929,9 +939,13 @@ Context:
task.context = []
# Add content to context
task.context.append(content)
# Run tasks as before
self.run_all_tasks()
# Planning Mode: Create plan and todo list before execution
if self.planning:
self._run_with_planning()
else:
# Run tasks as before
self.run_all_tasks()
# Auto-display token metrics if any agent has metrics=True
metrics_enabled = any(getattr(agent, 'metrics', False) for agent in self.agents)
@@ -1532,3 +1546,354 @@ Context:
else:
display_error(f"Invalid protocol: {protocol}. Choose 'http' or 'mcp'.")
return None
# =========================================================================
# Planning Mode Properties and Methods
# =========================================================================
@property
def current_plan(self):
"""Get the current plan."""
return self._current_plan
@property
def todo_list(self):
"""Get the current todo list."""
return self._todo_list
def _get_planning_agent(self):
"""Lazy load PlanningAgent."""
if self._planning_agent is None and self.planning:
from ..planning import PlanningAgent
self._planning_agent = PlanningAgent(
llm=self.planning_llm,
read_only=True,
verbose=self.verbose,
tools=self.planning_tools,
reasoning=self.planning_reasoning
)
return self._planning_agent
async def _create_plan(self, request: str = None, context: str = None):
"""
Create an implementation plan.
Args:
request: The request/goal to plan for
context: Optional additional context
Returns:
Plan instance
"""
planner = self._get_planning_agent()
if planner is None:
return None
# Use first task description as request if not provided
if request is None and self.tasks:
first_task = list(self.tasks.values())[0]
request = first_task.description
plan = await planner.create_plan(
request=request,
agents=self.agents,
tasks=list(self.tasks.values()),
context=context
)
self._current_plan = plan
# Create todo list from plan
from ..planning import TodoList
self._todo_list = TodoList.from_plan(plan)
return plan
def _create_plan_sync(self, request: str = None, context: str = None):
"""
Synchronous version of _create_plan.
Args:
request: The request/goal to plan for
context: Optional additional context
Returns:
Plan instance
"""
planner = self._get_planning_agent()
if planner is None:
return None
# Use first task description as request if not provided
if request is None and self.tasks:
first_task = list(self.tasks.values())[0]
request = first_task.description
plan = planner.create_plan_sync(
request=request,
agents=self.agents,
tasks=list(self.tasks.values()),
context=context
)
self._current_plan = plan
# Create todo list from plan
from ..planning import TodoList
self._todo_list = TodoList.from_plan(plan)
return plan
async def _request_approval(self, plan):
"""
Request approval for a plan.
Args:
plan: Plan to approve
Returns:
True if approved, False if rejected
"""
if self.auto_approve_plan:
plan.approve()
return True
from ..planning import ApprovalCallback
callback = ApprovalCallback(auto_approve=self.auto_approve_plan)
return await callback.async_call(plan)
def _request_approval_sync(self, plan):
"""
Synchronous version of _request_approval.
Args:
plan: Plan to approve
Returns:
True if approved, False if rejected
"""
if self.auto_approve_plan:
plan.approve()
return True
from ..planning import ApprovalCallback
callback = ApprovalCallback(auto_approve=self.auto_approve_plan)
return callback(plan)
def get_plan_markdown(self) -> str:
"""
Get the current plan as markdown.
Returns:
Markdown string or empty string if no plan
"""
if self._current_plan:
return self._current_plan.to_markdown()
return ""
def get_todo_markdown(self) -> str:
"""
Get the current todo list as markdown.
Returns:
Markdown string or empty string if no todo list
"""
if self._todo_list:
return self._todo_list.to_markdown()
return ""
def update_plan_step_status(self, step_id: str, status: str) -> bool:
"""
Update the status of a plan step.
Args:
step_id: ID of the step to update
status: New status
Returns:
True if updated, False if not found
"""
if self._current_plan:
result = self._current_plan.update_step_status(step_id, status)
# Sync todo list
if self._todo_list:
self._todo_list.sync_with_plan(self._current_plan)
return result
return False
def _run_with_planning(self):
"""
Run tasks with planning mode enabled.
This method:
1. Creates a plan using PlanningAgent
2. Generates a todo list from the plan
3. Creates proper Task objects from plan steps
4. Executes each task using the full Task execution system
(with memory, callbacks, guardrails, structured output, etc.)
5. Tracks progress as items are completed
"""
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from ..task import Task
console = Console()
# Step 1: Create the plan
console.print("\n[bold blue]📋 PLANNING PHASE[/bold blue]")
console.print("[dim]Creating implementation plan...[/dim]\n")
# Build request from tasks
task_descriptions = [task.description for task in self.tasks.values()]
request = " AND ".join(task_descriptions)
plan = self._create_plan_sync(request=request)
if not plan:
console.print("[yellow]⚠️ Planning failed, falling back to normal execution[/yellow]")
self.run_all_tasks()
return
# Display the plan
console.print(Panel(
Markdown(plan.to_markdown()),
title="[bold green]Generated Plan[/bold green]",
border_style="green"
))
# Step 2: Request approval if not auto-approve
if not self.auto_approve_plan:
approved = self._request_approval_sync(plan)
if not approved:
console.print("[red]❌ Plan rejected. Aborting execution.[/red]")
return
else:
plan.approve()
console.print("[green]✅ Plan auto-approved[/green]")
# Step 3: Create todo list for tracking
from ..planning import TodoList
self._todo_list = TodoList.from_plan(plan)
console.print("\n[bold blue]📝 TODO LIST[/bold blue]")
console.print(Panel(
Markdown(self._todo_list.to_markdown()),
title="[bold cyan]Tasks to Complete[/bold cyan]",
border_style="cyan"
))
# Step 4: Create proper Task objects from plan steps
console.print("\n[bold blue]🚀 EXECUTION PHASE[/bold blue]\n")
# Map agent names to agent instances
agent_map = {agent.name: agent for agent in self.agents}
# Store original tasks and create new tasks from plan
original_tasks = self.tasks.copy()
self.tasks = {}
self.task_id_counter = 0
# Create Task objects from plan steps
plan_tasks = []
step_to_task = {} # Map step_id to task for context chaining
for i, step in enumerate(plan.steps):
# Get the appropriate agent
agent = agent_map.get(step.agent, self.agents[0] if self.agents else None)
if not agent:
console.print(f"[yellow]⚠️ No agent found for '{step.agent}', using first available[/yellow]")
agent = self.agents[0] if self.agents else None
if not agent:
console.print(f"[red]❌ No agents available for step: {step.description}[/red]")
continue
# Build context from dependencies (previous task results)
context = []
for dep_id in step.dependencies:
# Convert step_X format to actual step index
if dep_id.startswith("step_"):
try:
dep_index = int(dep_id.split("_")[1])
if dep_index < len(plan_tasks):
context.append(plan_tasks[dep_index])
except (ValueError, IndexError):
pass
elif dep_id in step_to_task:
context.append(step_to_task[dep_id])
# Find matching original task for additional config (memory, callbacks, etc.)
original_task = None
for orig_task in original_tasks.values():
if orig_task.agent and orig_task.agent.name == agent.name:
original_task = orig_task
break
# Create Task with full features from original task if available
task = Task(
description=step.description,
expected_output=f"Complete: {step.description}",
agent=agent,
name=f"Plan Step {i + 1}",
tools=agent.tools if agent.tools else [],
context=context if context else None,
# Inherit from original task if available
memory=original_task.memory if original_task else None,
callback=original_task.callback if original_task else None,
guardrail=original_task.guardrail if original_task else None,
max_retries=original_task.max_retries if original_task else 3,
output_json=original_task.output_json if original_task else None,
output_pydantic=original_task.output_pydantic if original_task else None,
config=original_task.config if original_task else {}
)
# Add task to our task list
task_id = self.add_task(task)
plan_tasks.append(task)
step_to_task[step.id] = task
# Step 5: Execute tasks using the proper Task execution system
for i, (task_id, task) in enumerate(self.tasks.items()):
# Update todo list progress
if i < len(self._todo_list.items):
item = self._todo_list.items[i]
# Display progress bar
progress = self._todo_list.progress
bar_length = 30
filled = int(bar_length * progress)
bar = "█" * filled + "░" * (bar_length - filled)
console.print(f"[dim]Progress: [{bar}] {progress * 100:.0f}%[/dim]")
console.print(f"\n[bold]📌 Step {i + 1}/{len(self.tasks)}:[/bold] {task.description[:60]}...")
console.print(f"[dim] Agent: {task.agent.name if task.agent else 'Unknown'}[/dim]")
# Mark as in progress
self._todo_list.start(item.id)
# Execute using the full Task execution system
# This includes: memory, callbacks, guardrails, structured output, retry logic
try:
self.run_task(task_id)
if task.status == "completed":
if i < len(self._todo_list.items):
self._todo_list.complete(self._todo_list.items[i].id)
console.print("[green] ✅ Completed[/green]")
else:
console.print(f"[yellow] ⚠️ Task status: {task.status}[/yellow]")
except Exception as e:
console.print(f"[red] ❌ Error: {e}[/red]")
logger.error(f"Error executing plan task {task_id}: {e}")
# Final progress
completed_count = len([t for t in self.tasks.values() if t.status == "completed"])
console.print(f"\n[bold green]🎉 EXECUTION COMPLETE[/bold green]")
console.print(f"[dim]Progress: [{'█' * 30}] 100%[/dim]")
console.print(f"[green]Completed {completed_count}/{len(self.tasks)} tasks![/green]\n")
# Restore original tasks reference for result retrieval
self._plan_tasks = self.tasks.copy()
# Keep plan tasks for results but note original tasks are preserved in _plan_tasks
+206 -1
View File
@@ -239,6 +239,10 @@ class LLM:
api_key: Optional[str] = None,
base_url: Optional[str] = None,
events: List[Any] = [],
web_search: Optional[Union[bool, Dict[str, Any]]] = None,
web_fetch: Optional[Union[bool, Dict[str, Any]]] = None,
prompt_caching: Optional[bool] = None,
claude_memory: Optional[Union[bool, Any]] = None,
**extra_settings
):
# Configure logging only once at the class level
@@ -273,6 +277,11 @@ class LLM:
self.base_url = base_url
self.events = events
self.extra_settings = extra_settings
self.web_search = web_search
self.web_fetch = web_fetch
self.prompt_caching = prompt_caching
self.claude_memory = claude_memory
self._claude_memory_tool = None # Lazy initialized
self._console = None # Lazy load console when needed
self.chat_history = []
self.verbose = extra_settings.get('verbose', True)
@@ -381,6 +390,109 @@ class LLM:
else:
return False
def _supports_web_search(self) -> bool:
"""
Check if the current model supports native web search via LiteLLM.
Native web search allows the model to search the web in real-time
without requiring external tools like DuckDuckGo.
Returns:
bool: True if the model supports native web search, False otherwise
"""
from .model_capabilities import supports_web_search
return supports_web_search(self.model)
def _supports_web_fetch(self) -> bool:
"""
Check if the current model supports web fetch via LiteLLM.
Web fetch allows the model to retrieve full content from specific URLs
(web pages and PDF documents). Currently only supported by Anthropic Claude models.
Returns:
bool: True if the model supports web fetch, False otherwise
"""
from .model_capabilities import supports_web_fetch
return supports_web_fetch(self.model)
def _supports_prompt_caching(self) -> bool:
"""
Check if the current model supports prompt caching via LiteLLM.
Prompt caching allows caching parts of prompts to reduce costs and latency.
Supported by OpenAI, Anthropic, Bedrock, and Deepseek.
Returns:
bool: True if the model supports prompt caching, False otherwise
"""
from .model_capabilities import supports_prompt_caching
return supports_prompt_caching(self.model)
def _get_claude_memory_tool(self):
"""
Get or initialize the Claude Memory Tool instance.
Returns:
ClaudeMemoryTool instance if claude_memory is enabled, None otherwise
"""
if not self.claude_memory:
return None
if self._claude_memory_tool is None:
from ..tools.claude_memory_tool import ClaudeMemoryTool
if isinstance(self.claude_memory, bool):
# Create default memory tool
self._claude_memory_tool = ClaudeMemoryTool()
else:
# Use provided memory tool instance
self._claude_memory_tool = self.claude_memory
return self._claude_memory_tool
def _supports_claude_memory(self) -> bool:
"""
Check if the current model supports Claude Memory Tool.
Claude Memory Tool is only supported by Anthropic Claude models.
Returns:
bool: True if the model supports Claude Memory Tool
"""
if not self.claude_memory:
return False
return self._is_anthropic_model()
def _is_memory_tool_call(self, function_name: str) -> bool:
"""
Check if a tool call is for the Claude Memory Tool.
Args:
function_name: Name of the function being called
Returns:
bool: True if this is a memory tool call
"""
return function_name == "memory" and self._supports_claude_memory()
def _execute_memory_tool_call(self, arguments: Dict[str, Any]) -> str:
"""
Execute a Claude Memory Tool call.
Args:
arguments: The arguments from Claude's tool call
Returns:
Result string from the memory operation
"""
memory_tool = self._get_claude_memory_tool()
if not memory_tool:
return "Error: Memory tool not initialized"
return memory_tool.process_tool_call(arguments)
def _generate_ollama_tool_summary(self, tool_results: List[Any], response_text: str) -> Optional[str]:
"""
Generate a summary from tool results for Ollama to prevent infinite loops.
@@ -724,7 +836,21 @@ class LLM:
# Skip system messages for legacy o1 models as they don't support them
if not self._needs_system_message_skip():
messages.append({"role": "system", "content": system_prompt})
# Apply prompt caching for Anthropic models if enabled
if self.prompt_caching and self._supports_prompt_caching() and self._is_anthropic_model():
# Anthropic requires cache_control in content array format
messages.append({
"role": "system",
"content": [
{
"type": "text",
"text": system_prompt,
"cache_control": {"type": "ephemeral"}
}
]
})
else:
messages.append({"role": "system", "content": system_prompt})
# Add chat history if provided
if chat_history:
@@ -2295,6 +2421,12 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
return False
return any(prefix in self.model.lower() for prefix in ['gemini', 'gemini/', 'google/gemini'])
def _is_anthropic_model(self) -> bool:
"""Check if the model is an Anthropic Claude model."""
if not self.model:
return False
return any(prefix in self.model.lower() for prefix in ['claude', 'anthropic/', 'anthropic/claude'])
def _is_streaming_error_recoverable(self, error: Exception) -> bool:
"""Check if a streaming error is recoverable (e.g., malformed chunk vs connection error)."""
error_msg = str(error).lower()
@@ -3171,6 +3303,79 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
logging.debug(f"Could not verify function calling support: {e}. Setting tool_choice anyway.")
params['tool_choice'] = 'auto'
# Add web_search_options if web_search is enabled and model supports it
if self.web_search and self._supports_web_search():
if isinstance(self.web_search, dict):
# Use custom web_search_options provided by user
params['web_search_options'] = self.web_search
else:
# Use default web_search_options
params['web_search_options'] = {"search_context_size": "medium"}
logging.debug(f"Web search enabled with options: {params['web_search_options']}")
# Some web search models (like gpt-4o-search-preview) don't support certain parameters
# Remove incompatible parameters to avoid BadRequestError
incompatible_params = ['temperature', 'top_p', 'presence_penalty', 'frequency_penalty', 'logit_bias']
for param in incompatible_params:
if param in params:
del params[param]
logging.debug(f"Removed incompatible parameter '{param}' for web search model")
# Add web_fetch tool if web_fetch is enabled and model supports it
# Web fetch is passed as a tool (unlike web_search which uses web_search_options)
if self.web_fetch and self._supports_web_fetch():
# Build the web_fetch tool definition
web_fetch_tool = {
"type": "web_fetch_20250910",
"name": "web_fetch",
}
# Add optional parameters from web_fetch dict if provided
if isinstance(self.web_fetch, dict):
if 'max_uses' in self.web_fetch:
web_fetch_tool['max_uses'] = self.web_fetch['max_uses']
if 'allowed_domains' in self.web_fetch:
web_fetch_tool['allowed_domains'] = self.web_fetch['allowed_domains']
if 'blocked_domains' in self.web_fetch:
web_fetch_tool['blocked_domains'] = self.web_fetch['blocked_domains']
if 'citations' in self.web_fetch:
web_fetch_tool['citations'] = self.web_fetch['citations']
if 'max_content_tokens' in self.web_fetch:
web_fetch_tool['max_content_tokens'] = self.web_fetch['max_content_tokens']
else:
# Default max_uses if web_fetch=True
web_fetch_tool['max_uses'] = 5
# Add web_fetch tool to existing tools or create tools list
if 'tools' in params and params['tools']:
params['tools'].append(web_fetch_tool)
else:
params['tools'] = [web_fetch_tool]
logging.debug(f"Web fetch enabled with tool: {web_fetch_tool}")
# Add Claude Memory Tool if enabled and model supports it
if self._supports_claude_memory():
memory_tool = self._get_claude_memory_tool()
if memory_tool:
# Add the memory tool definition
memory_tool_def = memory_tool.get_tool_definition()
# Add memory tool to existing tools or create tools list
if 'tools' in params and params['tools']:
params['tools'].append(memory_tool_def)
else:
params['tools'] = [memory_tool_def]
# Add the beta header for Anthropic
beta_header = memory_tool.get_beta_header()
if 'extra_headers' in params:
params['extra_headers']['anthropic-beta'] = beta_header
else:
params['extra_headers'] = {'anthropic-beta': beta_header}
logging.debug(f"Claude memory tool enabled with beta header: {beta_header}")
return params
def _prepare_response_logging(self, temperature: float, stream: bool, verbose: bool, markdown: bool, **kwargs) -> Optional[Dict[str, Any]]:
@@ -1,62 +1,30 @@
"""
Model capabilities configuration for different LLM providers.
This module defines which models support specific features like structured outputs.
This module uses LiteLLM's helper functions as the primary source for model capability detection.
LiteLLM is maintained by many contributors and is more accurate and up-to-date.
LiteLLM Helper Functions:
- litellm.supports_web_search(model=) - Check web search support
- litellm.supports_function_calling(model=) - Check function calling support
- litellm.supports_parallel_function_calling(model=) - Check parallel function calling
- litellm.supports_response_schema(model=) - Check JSON schema/structured outputs support
- litellm.utils.supports_prompt_caching(model=) - Check prompt caching support
- litellm.get_supported_openai_params(model=) - Get all supported params
Sources:
- https://docs.litellm.ai/docs/completion/web_search
- https://docs.litellm.ai/docs/completion/json_mode
- https://docs.litellm.ai/docs/completion/function_call
- https://docs.litellm.ai/docs/completion/prompt_caching
"""
# Models that support OpenAI-style structured outputs (response_format with Pydantic models)
MODELS_SUPPORTING_STRUCTURED_OUTPUTS = {
# OpenAI models
"gpt-5-nano",
"gpt-5-nano",
"gpt-4-turbo",
"gpt-4-turbo-preview",
"gpt-4-turbo-2024-04-09",
"gpt-4-1106-preview",
"gpt-4-0125-preview",
"gpt-3.5-turbo",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo-0125",
# New/Future OpenAI models (as mentioned by user)
"codex-mini",
"o3-pro",
"gpt-4.5-preview",
"o3-mini",
"o1",
"o1-preview",
"o1-mini",
"gpt-4.1",
"gpt-4.1-nano",
"gpt-4.1-mini",
"o4-mini",
"o3",
# Gemini models that support structured outputs
"gemini-2.0-flash",
"gemini-2.0-flash-exp",
"gemini-1.5-pro",
"gemini-1.5-pro-latest",
"gemini-1.5-flash",
"gemini-1.5-flash-latest",
"gemini-1.5-flash-8b",
"gemini-1.5-flash-8b-latest",
}
# Models that explicitly DON'T support structured outputs
MODELS_NOT_SUPPORTING_STRUCTURED_OUTPUTS = {
# Audio preview models
"gpt-4o-audio-preview",
"gpt-5-nano-audio-preview",
# Legacy o1 models (don't support system messages either)
"o1-preview-2024-09-12",
"o1-mini-2024-09-12",
}
def supports_structured_outputs(model_name: str) -> bool:
"""
Check if a model supports OpenAI-style structured outputs.
Check if a model supports structured outputs (JSON schema).
Uses LiteLLM's supports_response_schema() as the primary check.
Args:
model_name: The name of the model to check
@@ -67,34 +35,72 @@ def supports_structured_outputs(model_name: str) -> bool:
if not model_name:
return False
# Strip provider prefixes (e.g., 'google/', 'openai/', etc.)
model_without_provider = model_name
for prefix in ['google/', 'openai/', 'anthropic/', 'gemini/', 'mistral/', 'deepseek/', 'groq/']:
if model_name.startswith(prefix):
model_without_provider = model_name[len(prefix):]
break
try:
import litellm
# Use LiteLLM's built-in check - most accurate and up-to-date
if hasattr(litellm, 'supports_response_schema'):
return litellm.supports_response_schema(model=model_name)
except Exception:
pass
# First check if it's explicitly in the NOT supporting list
if model_without_provider in MODELS_NOT_SUPPORTING_STRUCTURED_OUTPUTS:
return False
def supports_function_calling(model_name: str) -> bool:
"""
Check if a model supports function calling.
Uses LiteLLM's supports_function_calling() as the primary check.
Args:
model_name: The name of the model to check
Returns:
bool: True if the model supports function calling, False otherwise
"""
if not model_name:
return False
# Then check if it's in the supporting list
if model_without_provider in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
return True
try:
import litellm
# Use LiteLLM's built-in check - most accurate and up-to-date
if hasattr(litellm, 'supports_function_calling'):
return litellm.supports_function_calling(model=model_name)
except Exception:
pass
# For models with version suffixes, check the base model name
base_model = model_without_provider.split('-2024-')[0].split('-2025-')[0]
if base_model in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
return True
return False
def supports_parallel_function_calling(model_name: str) -> bool:
"""
Check if a model supports parallel function calling.
Uses LiteLLM's supports_parallel_function_calling() as the primary check.
Args:
model_name: The name of the model to check
Returns:
bool: True if the model supports parallel function calling, False otherwise
"""
if not model_name:
return False
try:
import litellm
# Use LiteLLM's built-in check - most accurate and up-to-date
if hasattr(litellm, 'supports_parallel_function_calling'):
return litellm.supports_parallel_function_calling(model=model_name)
except Exception:
pass
# Default to False for unknown models
return False
def supports_streaming_with_tools(model_name: str) -> bool:
"""
Check if a model supports streaming when tools are provided.
Most models that support structured outputs also support streaming with tools.
Args:
model_name: The name of the model to check
@@ -102,15 +108,153 @@ def supports_streaming_with_tools(model_name: str) -> bool:
Returns:
bool: True if the model supports streaming with tools, False otherwise
"""
# For now, use the same logic as structured outputs
# In the future, this could be a separate list if needed
return supports_structured_outputs(model_name)
# Models that support function calling generally support streaming with tools
return supports_function_calling(model_name)
# Supported Gemini internal tools
GEMINI_INTERNAL_TOOLS = {'googleSearch', 'urlContext', 'codeExecution'}
def supports_web_search(model_name: str) -> bool:
"""
Check if a model supports native web search via LiteLLM.
Uses LiteLLM's supports_web_search() as the primary check.
Native web search allows the model to search the web in real-time
without requiring external tools like DuckDuckGo.
Supported providers:
- OpenAI (gpt-4o-search-preview, gpt-4o-mini-search-preview)
- xAI (grok-3)
- Anthropic (claude-3-5-sonnet-latest, claude-sonnet-4, etc.)
- Google/Vertex AI (gemini-2.0-flash, gemini-2.5-*, etc.)
- Perplexity (all models)
Args:
model_name: The name of the model to check (with or without provider prefix)
Returns:
bool: True if the model supports native web search, False otherwise
"""
if not model_name:
return False
try:
import litellm
# Use LiteLLM's built-in check - most accurate and up-to-date
if hasattr(litellm, 'supports_web_search'):
return litellm.supports_web_search(model=model_name)
except Exception:
pass
return False
def supports_prompt_caching(model_name: str) -> bool:
"""
Check if a model supports prompt caching via LiteLLM.
Uses LiteLLM's supports_prompt_caching() as the primary check.
Prompt caching allows caching parts of prompts to reduce costs and latency
on subsequent requests with similar prompts.
Supported providers:
- OpenAI (openai/) - Automatic caching for prompts ≥1024 tokens
- Anthropic (anthropic/) - Manual caching with cache_control
- Bedrock (bedrock/) - All models that support prompt caching
- Deepseek (deepseek/) - Works like OpenAI
Args:
model_name: The name of the model to check (with or without provider prefix)
Returns:
bool: True if the model supports prompt caching, False otherwise
"""
if not model_name:
return False
try:
from litellm.utils import supports_prompt_caching as litellm_supports_prompt_caching
return litellm_supports_prompt_caching(model=model_name)
except Exception:
pass
return False
# Models that support web fetch via LiteLLM (Anthropic only)
# Web fetch retrieves full content from specific URLs (web pages and PDFs)
# Source: https://docs.litellm.ai/docs/completion/web_fetch
# Note: LiteLLM doesn't have a supports_web_fetch() helper yet, so we maintain a static list
MODELS_SUPPORTING_WEB_FETCH = {
# Anthropic Claude models with web fetch support
"claude-opus-4-1-20250805",
"claude-opus-4-1",
"claude-opus-4-20250514",
"claude-opus-4",
"claude-sonnet-4-20250514",
"claude-sonnet-4",
"claude-3-7-sonnet-20250219",
"claude-3-7-sonnet-latest",
"claude-3-5-sonnet-latest",
"claude-3-5-sonnet-20241022",
"claude-3-5-haiku-latest",
"claude-3-5-haiku-20241022",
}
def supports_web_fetch(model_name: str) -> bool:
"""
Check if a model supports web fetch via LiteLLM.
Web fetch allows the model to retrieve full content from specific URLs
(web pages and PDF documents). Currently only supported by Anthropic Claude models.
Note: LiteLLM doesn't have a supports_web_fetch() helper yet, so we use a static list
with auto-detection for Claude 4+ models.
Args:
model_name: The name of the model to check (with or without provider prefix)
Returns:
bool: True if the model supports web fetch, False otherwise
"""
if not model_name:
return False
# Strip provider prefixes
model_without_provider = model_name
for prefix in ['anthropic/', 'bedrock/', 'vertex_ai/']:
if model_name.startswith(prefix):
model_without_provider = model_name[len(prefix):]
break
# Check our static list
if model_without_provider in MODELS_SUPPORTING_WEB_FETCH:
return True
# Check base model name (without version suffix)
base_model = model_without_provider.split('-2024-')[0].split('-2025-')[0]
if base_model in MODELS_SUPPORTING_WEB_FETCH:
return True
# Auto-support for Claude 4+ models (Anthropic only)
model_lower = model_without_provider.lower()
if 'claude' in model_lower:
import re
# Match patterns like claude-4, claude-5, claude-sonnet-4, claude-opus-4, etc.
version_match = re.search(r'claude-(?:sonnet-|opus-|haiku-)?(\d+)', model_lower)
if version_match:
version = int(version_match.group(1))
if version >= 4: # Claude 4 and later
return True
return False
def is_gemini_internal_tool(tool) -> bool:
"""
Check if a tool is a Gemini internal tool and should be included in formatted tools.
@@ -2,8 +2,65 @@
Model Context Protocol (MCP) integration for PraisonAI Agents.
This package provides classes and utilities for connecting to MCP servers
using different transport methods (stdio, SSE, etc.).
using different transport methods:
- stdio: Standard input/output (subprocess)
- sse: Server-Sent Events (legacy HTTP+SSE)
- http_stream: Streamable HTTP (current standard)
- websocket: WebSocket (SEP-1288)
Protocol Revision: 2025-11-25
"""
from .mcp import MCP
__all__ = ["MCP"]
# Lazy imports for optional components
def __getattr__(name):
"""Lazy import for optional modules."""
if name == "WebSocketMCPClient":
from .mcp_websocket import WebSocketMCPClient
return WebSocketMCPClient
elif name == "SessionManager":
from .mcp_session import SessionManager
return SessionManager
elif name == "SecurityConfig":
from .mcp_security import SecurityConfig
return SecurityConfig
elif name == "BaseTransport":
from .mcp_transport import BaseTransport
return BaseTransport
elif name == "TransportConfig":
from .mcp_transport import TransportConfig
return TransportConfig
elif name == "ToolsMCPServer":
from .mcp_server import ToolsMCPServer
return ToolsMCPServer
elif name == "launch_tools_mcp_server":
from .mcp_server import launch_tools_mcp_server
return launch_tools_mcp_server
elif name == "function_to_mcp_schema":
from .mcp_utils import function_to_mcp_schema
return function_to_mcp_schema
elif name == "get_tool_metadata":
from .mcp_utils import get_tool_metadata
return get_tool_metadata
elif name == "python_type_to_json_schema":
from .mcp_utils import python_type_to_json_schema
return python_type_to_json_schema
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
# Client
"MCP",
# Server (lazy loaded)
"ToolsMCPServer",
"launch_tools_mcp_server",
# Utilities (lazy loaded)
"function_to_mcp_schema",
"get_tool_metadata",
"python_type_to_json_schema",
# Transport components (lazy loaded)
"WebSocketMCPClient",
"SessionManager",
"SecurityConfig",
"BaseTransport",
"TransportConfig",
]
@@ -227,6 +227,26 @@ class MCP:
self.timeout = timeout
self.debug = debug
# Check if this is a WebSocket URL (ws:// or wss://)
if isinstance(command_or_string, str) and re.match(r'^wss?://', command_or_string):
from .mcp_websocket import WebSocketMCPClient
# Extract auth token if provided
auth_token = kwargs.pop('auth_token', None)
self.websocket_client = WebSocketMCPClient(
command_or_string,
debug=debug,
timeout=timeout,
auth_token=auth_token,
options=kwargs
)
self._tools = list(self.websocket_client.tools)
self.is_sse = False
self.is_http_stream = False
self.is_websocket = True
self.is_npx = False
return
# Check if this is an HTTP URL
if isinstance(command_or_string, str) and re.match(r'^https?://', command_or_string):
# Determine transport type based on URL or kwargs
@@ -0,0 +1,150 @@
"""
MCP Backward Compatibility Module.
This module provides utilities for maintaining backward compatibility
with older MCP servers and the deprecated HTTP+SSE transport from
protocol version 2024-11-05.
Per MCP specification, clients wanting to support older servers should:
1. Attempt POST to server URL with InitializeRequest
2. If it fails with 400, 404, or 405, fall back to legacy SSE
3. Issue GET to server URL expecting SSE stream with endpoint event
"""
from typing import Optional
def detect_transport_support(
response_status: int,
content_type: Optional[str] = None
) -> str:
"""
Detect which transport a server supports based on response.
Per MCP spec backward compatibility guide:
- Success (200) with JSON indicates Streamable HTTP
- 400, 404, 405 indicates need for legacy SSE fallback
Args:
response_status: HTTP response status code
content_type: Response Content-Type header
Returns:
Transport type: "http_stream" or "sse"
"""
# Check for fallback indicators
if response_status in (400, 404, 405):
return "sse"
# Successful response indicates Streamable HTTP support
if response_status == 200:
return "http_stream"
# Default to SSE for unknown cases
return "sse"
def is_legacy_sse_url(url: str) -> bool:
"""
Check if URL indicates legacy SSE transport.
Legacy SSE URLs typically end with /sse.
Args:
url: Server URL
Returns:
True if URL indicates legacy SSE
"""
return url.rstrip('/').endswith('/sse')
def get_legacy_sse_endpoint(base_url: str) -> str:
"""
Get the legacy SSE endpoint URL.
Args:
base_url: Base server URL
Returns:
SSE endpoint URL
"""
base = base_url.rstrip('/')
if not base.endswith('/sse'):
return f"{base}/sse"
return base
def get_streamable_http_endpoint(base_url: str) -> str:
"""
Get the Streamable HTTP endpoint URL.
Args:
base_url: Base server URL
Returns:
MCP endpoint URL
"""
base = base_url.rstrip('/')
# Remove /sse suffix if present
if base.endswith('/sse'):
base = base[:-4]
if not base.endswith('/mcp'):
return f"{base}/mcp"
return base
class TransportNegotiator:
"""
Handles automatic transport negotiation with MCP servers.
This class implements the backward compatibility algorithm
from the MCP specification for detecting server capabilities.
"""
def __init__(self, base_url: str):
"""
Initialize negotiator.
Args:
base_url: Server base URL
"""
self.base_url = base_url
self._detected_transport: Optional[str] = None
@property
def detected_transport(self) -> Optional[str]:
"""Get the detected transport type."""
return self._detected_transport
def set_detected_transport(self, transport: str) -> None:
"""
Set the detected transport type.
Args:
transport: Transport type ("http_stream" or "sse")
"""
self._detected_transport = transport
def get_endpoint_url(self) -> str:
"""
Get the appropriate endpoint URL based on detected transport.
Returns:
Endpoint URL for the detected transport
"""
if self._detected_transport == "sse":
return get_legacy_sse_endpoint(self.base_url)
return get_streamable_http_endpoint(self.base_url)
def should_try_streamable_first(self) -> bool:
"""
Determine if Streamable HTTP should be tried first.
Per MCP spec, clients should try Streamable HTTP first
unless the URL explicitly indicates legacy SSE.
Returns:
True if should try Streamable HTTP first
"""
return not is_legacy_sse_url(self.base_url)
@@ -200,16 +200,31 @@ class HTTPStreamTransport:
This transport provides a single endpoint for all MCP communication,
supporting both batch (JSON) and streaming (SSE) response modes.
Per MCP Protocol Revision 2025-11-25:
- Includes MCP-Protocol-Version header on all requests
- Handles session management via Mcp-Session-Id header
- Supports SSE resumability via Last-Event-ID
"""
# Default protocol version for backward compatibility
DEFAULT_PROTOCOL_VERSION = '2025-03-26'
def __init__(self, base_url: str, session_id: Optional[str] = None, options: Optional[Dict[str, Any]] = None):
self.base_url = base_url
self.session_id = session_id
self.options = options or {}
self.response_mode = self.options.get('responseMode', 'batch')
self.protocol_version = self.options.get('protocol_version', self.DEFAULT_PROTOCOL_VERSION)
# Track last event ID for resumability
self.last_event_id: Optional[str] = None
self._retry_delay_ms: int = 3000 # Default retry delay
self.headers = {
'Content-Type': 'application/json',
'Accept': 'application/json, text/event-stream'
'Accept': 'application/json, text/event-stream',
'Mcp-Protocol-Version': self.protocol_version # Required per spec
}
if session_id:
self.headers['Mcp-Session-Id'] = session_id
@@ -281,7 +296,7 @@ class HTTPStreamTransport:
raise
async def _process_sse_response(self, response):
"""Process SSE response stream."""
"""Process SSE response stream with resumability support."""
buffer = ""
async for chunk in response.content:
buffer += chunk.decode('utf-8')
@@ -291,12 +306,34 @@ class HTTPStreamTransport:
event, buffer = buffer.split("\n\n", 1)
lines = event.strip().split("\n")
# Parse SSE event
data = None
for line in lines:
if line.startswith("data: "):
data = line[6:] # Remove "data: " prefix
# Parse SSE event fields per spec
event_id = None
data_lines = []
retry_value = None
for line in lines:
if line.startswith("id:"):
# Track event ID for resumability
event_id = line[3:].strip()
elif line.startswith("data:"):
data_lines.append(line[5:].strip())
elif line.startswith("retry:"):
# Handle retry field per SSE spec
try:
retry_value = int(line[6:].strip())
except ValueError:
pass
# Update last event ID for resumability
if event_id:
self.last_event_id = event_id
# Update retry delay if provided
if retry_value is not None:
self._retry_delay_ms = retry_value
# Process data if present
data = '\n'.join(data_lines) if data_lines else None
if data:
try:
message = json.loads(data)
@@ -310,13 +347,49 @@ class HTTPStreamTransport:
except json.JSONDecodeError:
logger.error(f"Failed to parse SSE event: {data}")
async def terminate_session(self) -> bool:
"""
Terminate the session via HTTP DELETE.
Per MCP spec: Clients that no longer need a particular session
SHOULD send an HTTP DELETE to the MCP endpoint with the
Mcp-Session-Id header to explicitly terminate the session.
Returns:
True if session was terminated, False if server doesn't support it
"""
if not self._session or not self.session_id:
return False
try:
headers = {
'Mcp-Session-Id': self.session_id,
'Mcp-Protocol-Version': self.protocol_version
}
async with self._session.delete(self.base_url, headers=headers) as response:
if response.status == 405:
# Server doesn't allow client-initiated session termination
logger.debug("Server does not allow session termination via DELETE")
return False
elif response.status in (200, 202, 204):
# Session terminated successfully
self.session_id = None
self.headers.pop('Mcp-Session-Id', None)
return True
else:
logger.warning(f"Unexpected response to DELETE: {response.status}")
return False
except Exception as e:
logger.error(f"Error terminating session: {e}")
return False
async def start_sse_listener(self):
"""Start listening for SSE events from the server."""
if self._sse_task is None or self._sse_task.done():
self._sse_task = asyncio.create_task(self._sse_listener())
async def _sse_listener(self):
"""Background task to listen for SSE events."""
"""Background task to listen for SSE events with resumability support."""
while True:
try:
# Check if we should stop
@@ -324,18 +397,28 @@ class HTTPStreamTransport:
break
url = self.base_url
if self.session_id:
# Add session as query parameter for SSE connection
url = f"{url}?session={self.session_id}"
# Build headers per MCP spec
headers = {
'Accept': 'text/event-stream',
'Cache-Control': 'no-cache'
'Cache-Control': 'no-cache',
'Mcp-Protocol-Version': self.protocol_version # Required per spec
}
if self.session_id:
headers['Mcp-Session-Id'] = self.session_id
# Include Last-Event-ID for resumability per spec
if self.last_event_id:
headers['Last-Event-ID'] = self.last_event_id
async with self._session.get(url, headers=headers) as response:
# Handle session expiration (HTTP 404)
if response.status == 404:
logger.warning("Session expired (HTTP 404), need to reinitialize")
self.session_id = None
self.headers.pop('Mcp-Session-Id', None)
break
buffer = ""
async for chunk in response.content:
# Check if we should stop
@@ -349,12 +432,32 @@ class HTTPStreamTransport:
event, buffer = buffer.split("\n\n", 1)
lines = event.strip().split("\n")
# Parse SSE event
data = None
for line in lines:
if line.startswith("data: "):
data = line[6:] # Remove "data: " prefix
# Parse SSE event fields
event_id = None
data_lines = []
retry_value = None
for line in lines:
if line.startswith("id:"):
event_id = line[3:].strip()
elif line.startswith("data:"):
data_lines.append(line[5:].strip())
elif line.startswith("retry:"):
try:
retry_value = int(line[6:].strip())
except ValueError:
pass
# Update last event ID for resumability
if event_id:
self.last_event_id = event_id
# Update retry delay if provided
if retry_value is not None:
self._retry_delay_ms = retry_value
# Process data
data = '\n'.join(data_lines) if data_lines else None
if data:
try:
message = json.loads(data)
@@ -368,7 +471,8 @@ class HTTPStreamTransport:
except Exception as e:
if not (hasattr(self, '_closing') and self._closing):
logger.error(f"SSE listener error: {e}")
await asyncio.sleep(1) # Reconnect after 1 second
# Use retry delay from server or default
await asyncio.sleep(self._retry_delay_ms / 1000.0)
else:
break
@@ -0,0 +1,248 @@
"""
MCP Security Module.
This module provides security utilities for MCP transports,
implementing security best practices from MCP Protocol
Revision 2025-11-25.
Security features:
- Origin header validation (DNS rebinding prevention)
- Localhost binding recommendations
- Authentication header support
- Secure session ID generation
"""
import secrets
import base64
from typing import Optional, Dict, List, Set
from urllib.parse import urlparse
from dataclasses import dataclass, field
# Localhost addresses
LOCALHOST_ADDRESSES: Set[str] = {
'localhost',
'127.0.0.1',
'::1', # IPv6 localhost
}
def is_valid_origin(
origin: Optional[str],
allowed_origins: List[str],
allow_missing: bool = False
) -> bool:
"""
Validate Origin header against allowed origins.
Per MCP spec, servers MUST validate the Origin header on all
incoming connections to prevent DNS rebinding attacks.
Args:
origin: The Origin header value
allowed_origins: List of allowed origin hosts
allow_missing: Whether to allow missing Origin header
Returns:
True if origin is valid, False otherwise
"""
if not origin:
return allow_missing
# Extract host from origin
host = extract_origin_host(origin)
if not host:
return False
# Check against allowed origins
return host in allowed_origins
def extract_origin_host(origin: str) -> Optional[str]:
"""
Extract the host from an Origin header value.
Args:
origin: Origin header value (e.g., "https://example.com:443")
Returns:
Host portion of the origin, or None if invalid
"""
if not origin:
return None
try:
parsed = urlparse(origin)
return parsed.hostname
except Exception:
return None
def is_localhost_address(address: str) -> bool:
"""
Check if an address is a localhost address.
Args:
address: IP address or hostname
Returns:
True if localhost, False otherwise
"""
return address.lower() in LOCALHOST_ADDRESSES
def should_bind_localhost_only(is_local: bool = True) -> bool:
"""
Determine if server should bind to localhost only.
Per MCP spec, when running locally, servers SHOULD bind only
to localhost (127.0.0.1) rather than all network interfaces.
Args:
is_local: Whether this is a local development server
Returns:
True if should bind to localhost only
"""
return is_local
def is_potential_dns_rebinding(origin: str, server_host: str) -> bool:
"""
Detect potential DNS rebinding attack.
DNS rebinding occurs when an external origin tries to access
a local server through DNS manipulation.
Args:
origin: The Origin header value
server_host: The host the server is bound to
Returns:
True if potential DNS rebinding attempt
"""
origin_host = extract_origin_host(origin)
if not origin_host:
return True # Treat invalid origin as suspicious
# If server is on localhost, only localhost origins are safe
if is_localhost_address(server_host):
return not is_localhost_address(origin_host)
# For non-localhost servers, check if origin matches server
return origin_host != server_host
def create_auth_header(
credentials: str,
auth_type: str = "bearer",
header_name: str = "Authorization"
) -> Dict[str, str]:
"""
Create an authentication header.
Args:
credentials: The credentials/token
auth_type: Type of auth ("bearer", "basic", "custom")
header_name: Custom header name (for auth_type="custom")
Returns:
Dictionary with authentication header
"""
if auth_type.lower() == "bearer":
return {"Authorization": f"Bearer {credentials}"}
elif auth_type.lower() == "basic":
# Base64 encode for Basic auth
encoded = base64.b64encode(credentials.encode()).decode()
return {"Authorization": f"Basic {encoded}"}
elif auth_type.lower() == "custom":
return {header_name: credentials}
else:
raise ValueError(f"Unknown auth_type: {auth_type}")
def validate_auth_header(headers: Dict[str, str]) -> bool:
"""
Validate that authentication header is present.
Args:
headers: Request headers dictionary
Returns:
True if Authorization header is present
"""
# Check for Authorization header (case-insensitive)
for key in headers:
if key.lower() == 'authorization':
return bool(headers[key])
return False
def generate_secure_session_id(length: int = 32) -> str:
"""
Generate a cryptographically secure session ID.
Per MCP spec, session IDs SHOULD be globally unique and
cryptographically secure (e.g., securely generated UUID,
JWT, or cryptographic hash).
The ID MUST only contain visible ASCII characters (0x21-0x7E).
Args:
length: Desired length of session ID (minimum 32)
Returns:
Secure session ID string
"""
# Ensure minimum length for security
length = max(length, 32)
# Use URL-safe base64 which only uses visible ASCII
# Generate extra bytes to account for base64 expansion
num_bytes = (length * 3) // 4 + 1
random_bytes = secrets.token_bytes(num_bytes)
# Encode to URL-safe base64 (uses A-Z, a-z, 0-9, -, _)
session_id = base64.urlsafe_b64encode(random_bytes).decode()
# Trim to desired length
return session_id[:length]
@dataclass
class SecurityConfig:
"""
Security configuration for MCP transports.
This class holds security settings that can be applied
to MCP server implementations.
Attributes:
validate_origin: Whether to validate Origin header
allow_missing_origin: Whether to allow missing Origin
allowed_origins: List of allowed origin hosts
require_auth: Whether authentication is required
bind_localhost_only: Whether to bind to localhost only
"""
validate_origin: bool = True
allow_missing_origin: bool = False
allowed_origins: List[str] = field(default_factory=lambda: ["localhost", "127.0.0.1"])
require_auth: bool = False
bind_localhost_only: bool = True
def is_origin_allowed(self, origin: Optional[str]) -> bool:
"""Check if an origin is allowed by this config."""
if not self.validate_origin:
return True
return is_valid_origin(
origin,
self.allowed_origins,
self.allow_missing_origin
)
def get_bind_address(self) -> str:
"""Get recommended bind address based on config."""
return "127.0.0.1" if self.bind_localhost_only else "0.0.0.0"
@@ -0,0 +1,356 @@
"""MCP Server wrapper for exposing PraisonAI tools as MCP servers.
This module provides a simple way to expose Python functions as MCP tools
that can be consumed by any MCP client (Claude Desktop, Cursor, etc.).
Usage:
from praisonaiagents.mcp import ToolsMCPServer
def search(query: str) -> str:
'''Search the web.'''
return f"Results for {query}"
server = ToolsMCPServer(name="my-tools")
server.register_tool(search)
server.run() # Starts MCP server (stdio by default)
"""
import asyncio
import inspect
import logging
from typing import Any, Callable, Dict, List, Optional
from .mcp_utils import function_to_mcp_schema
logger = logging.getLogger(__name__)
class ToolsMCPServer:
"""MCP Server that exposes Python functions as MCP tools.
This class wraps Python functions and exposes them as MCP tools
using the FastMCP library. It supports both sync and async functions.
Example:
server = ToolsMCPServer(name="my-tools")
def search(query: str) -> str:
'''Search the web.'''
return f"Results for {query}"
server.register_tool(search)
server.run() # Starts stdio server
"""
def __init__(
self,
name: str = "praisonai-tools",
tools: Optional[List[Callable]] = None,
debug: bool = False
):
"""Initialize the MCP server.
Args:
name: Name of the MCP server
tools: Optional list of tools to register on init
debug: Enable debug logging
"""
self.name = name
self.tools: List[Callable] = []
self._tool_map: Dict[str, Callable] = {}
self._fastmcp = None
self._debug = debug
if debug:
logger.setLevel(logging.DEBUG)
# Register initial tools if provided
if tools:
self.register_tools(tools)
def register_tool(self, func: Callable) -> None:
"""Register a single tool function.
Args:
func: The function to register as an MCP tool
"""
if not callable(func):
raise TypeError(f"Expected callable, got {type(func)}")
self.tools.append(func)
tool_name = getattr(func, '__mcp_name__', None) or func.__name__
self._tool_map[tool_name] = func
logger.debug(f"Registered tool: {tool_name}")
def register_tools(self, funcs: List[Callable]) -> None:
"""Register multiple tool functions.
Args:
funcs: List of functions to register as MCP tools
"""
for func in funcs:
self.register_tool(func)
def get_tool_schemas(self) -> List[Dict[str, Any]]:
"""Get MCP schemas for all registered tools.
Returns:
List of MCP tool schema dictionaries
"""
return [function_to_mcp_schema(tool) for tool in self.tools]
def get_tool_names(self) -> List[str]:
"""Get names of all registered tools.
Returns:
List of tool names
"""
return list(self._tool_map.keys())
def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""Execute a registered tool synchronously.
Args:
tool_name: Name of the tool to execute
arguments: Arguments to pass to the tool
Returns:
Result from the tool execution
Raises:
ValueError: If tool is not found
"""
if tool_name not in self._tool_map:
raise ValueError(f"Tool '{tool_name}' not found")
func = self._tool_map[tool_name]
# Handle async functions
if asyncio.iscoroutinefunction(func):
loop = asyncio.new_event_loop()
try:
return loop.run_until_complete(func(**arguments))
finally:
loop.close()
return func(**arguments)
async def execute_tool_async(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""Execute a registered tool asynchronously.
Args:
tool_name: Name of the tool to execute
arguments: Arguments to pass to the tool
Returns:
Result from the tool execution
Raises:
ValueError: If tool is not found
"""
if tool_name not in self._tool_map:
raise ValueError(f"Tool '{tool_name}' not found")
func = self._tool_map[tool_name]
if asyncio.iscoroutinefunction(func):
return await func(**arguments)
# Run sync function in executor
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, lambda: func(**arguments))
def get_fastmcp(self):
"""Get or create FastMCP instance with registered tools.
Returns:
FastMCP instance with all tools registered
"""
if self._fastmcp is None:
self._fastmcp = self._create_fastmcp()
return self._fastmcp
def _create_fastmcp(self):
"""Create FastMCP instance and register all tools.
Returns:
Configured FastMCP instance
"""
try:
from mcp.server.fastmcp import FastMCP
except ImportError:
raise ImportError(
"FastMCP not available. Install with: pip install mcp"
)
mcp = FastMCP(self.name)
# Register each tool with FastMCP
for func in self.tools:
self._register_tool_with_fastmcp(mcp, func)
return mcp
def _register_tool_with_fastmcp(self, mcp, func: Callable) -> None:
"""Register a single function with FastMCP.
Args:
mcp: FastMCP instance
func: Function to register
"""
# Get tool metadata
tool_name = getattr(func, '__mcp_name__', None) or func.__name__
description = getattr(func, '__mcp_description__', None)
if description is None:
docstring = inspect.getdoc(func)
description = docstring.split('\n')[0] if docstring else tool_name
# Create wrapper that preserves signature
if asyncio.iscoroutinefunction(func):
@mcp.tool(name=tool_name, description=description)
async def tool_wrapper(**kwargs):
return await func(**kwargs)
else:
@mcp.tool(name=tool_name, description=description)
def tool_wrapper(**kwargs):
return func(**kwargs)
# Copy signature for proper schema generation
tool_wrapper.__signature__ = inspect.signature(func)
tool_wrapper.__annotations__ = getattr(func, '__annotations__', {})
def run(self, transport: str = "stdio") -> None:
"""Run the MCP server.
Args:
transport: Transport type - "stdio" or "sse"
"""
if transport == "stdio":
self.run_stdio()
elif transport == "sse":
self.run_sse()
else:
raise ValueError(f"Unknown transport: {transport}. Use 'stdio' or 'sse'")
def run_stdio(self) -> None:
"""Run the MCP server using stdio transport."""
mcp = self.get_fastmcp()
print(f"🚀 Starting MCP server '{self.name}' with stdio transport")
print(f"🛠️ Available tools: {', '.join(self.get_tool_names())}")
mcp.run()
def run_sse(self, host: str = "0.0.0.0", port: int = 8080) -> None:
"""Run the MCP server using SSE transport.
Args:
host: Host to bind to
port: Port to listen on
"""
try:
import uvicorn
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.routing import Mount, Route
except ImportError as e:
raise ImportError(
f"SSE dependencies not available: {e}. "
"Install with: pip install uvicorn starlette"
)
mcp = self.get_fastmcp()
# Set up SSE transport
sse_path = "/sse"
messages_path = "/messages/"
sse_transport = SseServerTransport(messages_path)
async def handle_sse(request: Request):
async with sse_transport.connect_sse(
request.scope, request.receive, request._send
) as (read_stream, write_stream):
await mcp._mcp_server.run(
read_stream,
write_stream,
mcp._mcp_server.create_initialization_options()
)
app = Starlette(
debug=self._debug,
routes=[
Route(sse_path, endpoint=handle_sse),
Mount(messages_path, app=sse_transport.handle_post_message),
]
)
print(f"🚀 Starting MCP server '{self.name}' with SSE transport")
print(f"📡 SSE endpoint: http://{host}:{port}{sse_path}")
print(f"🛠️ Available tools: {', '.join(self.get_tool_names())}")
uvicorn.run(app, host=host, port=port)
def launch_tools_mcp_server(
tools: Optional[List[Callable]] = None,
tool_names: Optional[List[str]] = None,
name: str = "praisonai-tools",
transport: str = "stdio",
host: str = "0.0.0.0",
port: int = 8080,
debug: bool = False
) -> None:
"""Convenience function to launch an MCP server with tools.
Args:
tools: List of tool functions to expose
tool_names: List of tool names to load from praisonaiagents.tools
name: Name of the MCP server
transport: Transport type - "stdio" or "sse"
host: Host for SSE transport
port: Port for SSE transport
debug: Enable debug logging
Example:
# With custom tools
def my_search(query: str) -> str:
return f"Results for {query}"
launch_tools_mcp_server(tools=[my_search])
# With built-in tools
launch_tools_mcp_server(tool_names=["tavily_search", "exa_search"])
"""
server = ToolsMCPServer(name=name, debug=debug)
# Register custom tools
if tools:
server.register_tools(tools)
# Load and register tools by name
if tool_names:
from praisonaiagents import tools as praison_tools
for tool_name in tool_names:
if hasattr(praison_tools, tool_name):
tool_func = getattr(praison_tools, tool_name)
if callable(tool_func):
server.register_tool(tool_func)
else:
logger.warning(f"Tool '{tool_name}' is not callable")
else:
logger.warning(f"Tool '{tool_name}' not found in praisonaiagents.tools")
if not server.tools:
raise ValueError("No tools registered. Provide 'tools' or 'tool_names'.")
# Run the server
if transport == "stdio":
server.run_stdio()
elif transport == "sse":
server.run_sse(host=host, port=port)
else:
raise ValueError(f"Unknown transport: {transport}")
@@ -0,0 +1,457 @@
"""
MCP Session Management Module.
This module provides session management utilities for MCP transports,
implementing the session management requirements from MCP Protocol
Revision 2025-11-25.
Features:
- Session ID validation (visible ASCII 0x21-0x7E)
- Session lifecycle management
- Protocol version header handling
- Session expiration detection
"""
import re
from typing import Optional, Dict, Any
# Valid protocol versions
VALID_PROTOCOL_VERSIONS = {
'2024-11-05', # Original HTTP+SSE
'2025-03-26', # Streamable HTTP introduction
'2025-06-18', # Updates
'2025-11-25', # Current revision
}
# Default protocol version for backward compatibility
DEFAULT_PROTOCOL_VERSION = '2025-03-26'
def is_valid_session_id(session_id: Optional[str]) -> bool:
"""
Validate MCP session ID per specification.
Session IDs MUST only contain visible ASCII characters (0x21 to 0x7E).
This excludes space (0x20), control characters (0x00-0x1F), and DEL (0x7F).
Args:
session_id: The session ID to validate
Returns:
True if valid, False otherwise
"""
if session_id is None or session_id == "":
return False
# Check each character is in visible ASCII range (0x21-0x7E)
for char in session_id:
code = ord(char)
if code < 0x21 or code > 0x7E:
return False
return True
def is_valid_protocol_version(version: Optional[str]) -> bool:
"""
Validate MCP protocol version string.
Args:
version: Protocol version string (e.g., '2025-11-25')
Returns:
True if valid, False otherwise
"""
if not version:
return False
# Check format: YYYY-MM-DD
if not re.match(r'^\d{4}-\d{2}-\d{2}$', version):
return False
return True
def is_session_expired(status_code: int) -> bool:
"""
Check if HTTP status code indicates session expiration.
Per MCP spec, HTTP 404 indicates the session has been terminated.
Args:
status_code: HTTP response status code
Returns:
True if session is expired, False otherwise
"""
return status_code == 404
def extract_session_id(headers: Dict[str, str]) -> Optional[str]:
"""
Extract session ID from HTTP response headers.
The header name is case-insensitive per HTTP spec.
Args:
headers: HTTP response headers dictionary
Returns:
Session ID if present, None otherwise
"""
# Check various case combinations
for key in headers:
if key.lower() == 'mcp-session-id':
return headers[key]
return None
class SessionManager:
"""
Manages MCP session state for HTTP transports.
This class handles:
- Session ID storage and validation
- Protocol version negotiation
- Session lifecycle (active, expired, terminated)
- Header generation for requests
Example:
```python
manager = SessionManager()
# After receiving InitializeResult with session ID
manager.set_session_id("abc123")
# Get headers for subsequent requests
headers = manager.get_headers()
# {'Mcp-Session-Id': 'abc123', 'Mcp-Protocol-Version': '2025-03-26'}
```
"""
def __init__(self, protocol_version: Optional[str] = None):
"""
Initialize session manager.
Args:
protocol_version: Initial protocol version (defaults to 2025-03-26)
"""
self._session_id: Optional[str] = None
self._protocol_version: str = protocol_version or DEFAULT_PROTOCOL_VERSION
self._is_expired: bool = False
self._metadata: Dict[str, Any] = {}
@property
def session_id(self) -> Optional[str]:
"""Get current session ID."""
return self._session_id
@property
def protocol_version(self) -> str:
"""Get current protocol version."""
return self._protocol_version
@property
def is_active(self) -> bool:
"""Check if session is active (has ID and not expired)."""
return self._session_id is not None and not self._is_expired
@property
def is_expired(self) -> bool:
"""Check if session has expired."""
return self._is_expired
def set_session_id(self, session_id: str) -> None:
"""
Set the session ID.
Args:
session_id: The session ID from server
Raises:
ValueError: If session ID is invalid
"""
if not is_valid_session_id(session_id):
raise ValueError("Invalid session ID: must contain only visible ASCII (0x21-0x7E)")
self._session_id = session_id
self._is_expired = False
def set_protocol_version(self, version: str) -> None:
"""
Set the negotiated protocol version.
Args:
version: Protocol version string
Raises:
ValueError: If version format is invalid
"""
if not is_valid_protocol_version(version):
raise ValueError(f"Invalid protocol version format: {version}")
self._protocol_version = version
def clear(self) -> None:
"""Clear the session (logout/disconnect)."""
self._session_id = None
self._is_expired = False
self._metadata = {}
def mark_expired(self) -> None:
"""Mark the session as expired (received HTTP 404)."""
self._is_expired = True
def get_headers(self) -> Dict[str, str]:
"""
Get HTTP headers for MCP requests.
Returns:
Dictionary of headers to include in requests
"""
headers = {
'Mcp-Protocol-Version': self._protocol_version
}
if self._session_id:
headers['Mcp-Session-Id'] = self._session_id
return headers
def update_from_response(self, headers: Dict[str, str], status_code: int = 200) -> None:
"""
Update session state from HTTP response.
Args:
headers: Response headers
status_code: HTTP status code
"""
# Check for session expiration
if is_session_expired(status_code):
self.mark_expired()
return
# Extract and update session ID if present
new_session_id = extract_session_id(headers)
if new_session_id:
self.set_session_id(new_session_id)
def set_metadata(self, key: str, value: Any) -> None:
"""Store arbitrary metadata with the session."""
self._metadata[key] = value
def get_metadata(self, key: str, default: Any = None) -> Any:
"""Retrieve metadata from the session."""
return self._metadata.get(key, default)
class ResumabilityManager:
"""
Manages SSE stream resumability for MCP transports.
This class handles:
- Event ID tracking per stream
- Last-Event-ID header generation
- Retry delay handling
Per MCP spec, event IDs should be assigned per-stream to act as
a cursor within that particular stream.
"""
def __init__(self):
"""Initialize resumability manager."""
self._last_event_ids: Dict[str, str] = {} # stream_id -> last_event_id
self._retry_delay: int = 3000 # Default retry delay in ms
def set_last_event_id(self, stream_id: str, event_id: str) -> None:
"""
Store the last received event ID for a stream.
Args:
stream_id: Identifier for the SSE stream
event_id: The event ID from the SSE event
"""
self._last_event_ids[stream_id] = event_id
def get_last_event_id(self, stream_id: str) -> Optional[str]:
"""
Get the last received event ID for a stream.
Args:
stream_id: Identifier for the SSE stream
Returns:
Last event ID if available, None otherwise
"""
return self._last_event_ids.get(stream_id)
def set_retry_delay(self, delay_ms: int) -> None:
"""
Set the retry delay from SSE retry field.
Args:
delay_ms: Delay in milliseconds
"""
self._retry_delay = delay_ms
def get_retry_delay(self) -> int:
"""Get current retry delay in milliseconds."""
return self._retry_delay
def get_resume_headers(self, stream_id: str) -> Dict[str, str]:
"""
Get headers for resuming a stream.
Args:
stream_id: Identifier for the SSE stream
Returns:
Headers dict with Last-Event-ID if available
"""
headers = {}
last_event_id = self.get_last_event_id(stream_id)
if last_event_id:
headers['Last-Event-ID'] = last_event_id
return headers
def clear_stream(self, stream_id: str) -> None:
"""Clear tracking for a specific stream."""
self._last_event_ids.pop(stream_id, None)
def clear_all(self) -> None:
"""Clear all stream tracking."""
self._last_event_ids.clear()
def parse_sse_event(event_text: str) -> Dict[str, Any]:
"""
Parse an SSE event into its components.
SSE events have the format:
id: <event-id>
retry: <milliseconds>
data: <data-line>
data: <data-line-2>
(blank line ends event)
Args:
event_text: Raw SSE event text
Returns:
Dictionary with 'id', 'retry', 'data', 'event' fields as available
"""
result: Dict[str, Any] = {}
data_lines = []
for line in event_text.strip().split('\n'):
if line.startswith('id:'):
result['id'] = line[3:].strip()
elif line.startswith('retry:'):
try:
result['retry'] = int(line[6:].strip())
except ValueError:
pass
elif line.startswith('data:'):
data_lines.append(line[5:].strip())
elif line.startswith('event:'):
result['event'] = line[6:].strip()
# Join multiple data lines with newlines
if data_lines:
result['data'] = '\n'.join(data_lines)
return result
class EventStore:
"""
Optional event store for message redelivery support.
This class stores events per stream to enable replay when
clients reconnect with Last-Event-ID header.
Note: This is an in-memory store suitable for single-server
deployments. For multi-node deployments, use an external
store (Redis, database, etc.).
"""
def __init__(self, max_events: int = 1000):
"""
Initialize event store.
Args:
max_events: Maximum events to store per stream
"""
self.max_events = max_events
self._streams: Dict[str, list] = {} # stream_id -> [(event_id, data), ...]
def store(self, stream_id: str, event_id: str, data: Any) -> None:
"""
Store an event for potential replay.
Args:
stream_id: Identifier for the SSE stream
event_id: The event ID
data: The event data
"""
if stream_id not in self._streams:
self._streams[stream_id] = []
self._streams[stream_id].append((event_id, data))
# Trim to max_events
if len(self._streams[stream_id]) > self.max_events:
self._streams[stream_id] = self._streams[stream_id][-self.max_events:]
def get_events_after(self, stream_id: str, after_event_id: str) -> list:
"""
Get events after a specific event ID.
Args:
stream_id: Identifier for the SSE stream
after_event_id: The Last-Event-ID from client
Returns:
List of event data after the specified event
"""
if stream_id not in self._streams:
return []
events = self._streams[stream_id]
# Find the index of the after_event_id
start_index = None
for i, (event_id, _) in enumerate(events):
if event_id == after_event_id:
start_index = i + 1
break
if start_index is None:
return []
return [data for _, data in events[start_index:]]
def get_all_events(self, stream_id: str) -> list:
"""
Get all stored events for a stream.
Args:
stream_id: Identifier for the SSE stream
Returns:
List of all event data
"""
if stream_id not in self._streams:
return []
return [data for _, data in self._streams[stream_id]]
def clear_stream(self, stream_id: str) -> None:
"""Clear events for a specific stream."""
self._streams.pop(stream_id, None)
def clear_all(self) -> None:
"""Clear all stored events."""
self._streams.clear()
@@ -0,0 +1,298 @@
"""
MCP Transport Abstraction Layer.
This module provides a common interface for all MCP transports,
enabling easy switching between transport types and custom
transport implementations.
Transport types:
- stdio: Standard input/output (subprocess)
- sse: Server-Sent Events (legacy HTTP+SSE)
- http_stream: Streamable HTTP (current standard)
- websocket: WebSocket (SEP-1288)
"""
import re
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, List, Type
from dataclasses import dataclass, field
@dataclass
class TransportConfig:
"""
Configuration for MCP transports.
Attributes:
timeout: Connection/operation timeout in seconds
debug: Enable debug logging
retry_count: Number of retry attempts
retry_delay: Base delay between retries in seconds
auth_token: Optional authentication token
headers: Additional headers for HTTP transports
"""
timeout: int = 60
debug: bool = False
retry_count: int = 3
retry_delay: float = 1.0
auth_token: Optional[str] = None
headers: Dict[str, str] = field(default_factory=dict)
class BaseTransport(ABC):
"""
Abstract base class for MCP transports.
All MCP transports must implement this interface to ensure
consistent behavior across transport types.
The transport is responsible for:
- Establishing and maintaining connections
- Sending JSON-RPC messages
- Receiving JSON-RPC messages
- Proper cleanup on close
"""
@abstractmethod
async def connect(self) -> None:
"""
Establish connection to the MCP server.
Raises:
ConnectionError: If connection fails
"""
pass
@abstractmethod
async def send(self, message: Dict[str, Any]) -> None:
"""
Send a JSON-RPC message to the server.
Args:
message: JSON-RPC message dictionary
Raises:
RuntimeError: If not connected
"""
pass
@abstractmethod
async def receive(self) -> Dict[str, Any]:
"""
Receive a JSON-RPC message from the server.
Returns:
JSON-RPC message dictionary
Raises:
RuntimeError: If not connected
"""
pass
@abstractmethod
async def close(self) -> None:
"""
Close the connection and cleanup resources.
"""
pass
@property
@abstractmethod
def is_connected(self) -> bool:
"""Check if transport is currently connected."""
pass
async def __aenter__(self):
"""Async context manager entry."""
await self.connect()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Async context manager exit."""
await self.close()
def get_transport_type(target: str) -> str:
"""
Determine the appropriate transport type for a target.
Args:
target: URL or command string
Returns:
Transport type string: "stdio", "sse", "http_stream", or "websocket"
"""
if not isinstance(target, str):
return "stdio"
# WebSocket URLs
if re.match(r'^wss?://', target, re.IGNORECASE):
return "websocket"
# HTTP URLs
if re.match(r'^https?://', target, re.IGNORECASE):
# Legacy SSE endpoint
if target.endswith('/sse'):
return "sse"
# Default to Streamable HTTP
return "http_stream"
# Everything else is stdio (command)
return "stdio"
class TransportRegistry:
"""
Registry for MCP transport implementations.
This allows registration of custom transports and provides
a central place to look up transport classes by name.
"""
def __init__(self):
"""Initialize empty registry."""
self._transports: Dict[str, Type[BaseTransport]] = {}
def register(self, name: str, transport_class: Type[BaseTransport]) -> None:
"""
Register a transport class.
Args:
name: Transport name (e.g., "websocket")
transport_class: Transport class (must extend BaseTransport)
"""
if not issubclass(transport_class, BaseTransport):
raise TypeError(f"{transport_class} must extend BaseTransport")
self._transports[name] = transport_class
def get(self, name: str) -> Optional[Type[BaseTransport]]:
"""
Get a registered transport class.
Args:
name: Transport name
Returns:
Transport class or None if not found
"""
return self._transports.get(name)
def list_transports(self) -> List[str]:
"""
List all registered transport names.
Returns:
List of transport names
"""
return list(self._transports.keys())
def unregister(self, name: str) -> None:
"""
Unregister a transport.
Args:
name: Transport name to remove
"""
self._transports.pop(name, None)
# Placeholder transport classes for registry
# These are minimal implementations to satisfy the registry
class StdioTransportPlaceholder(BaseTransport):
"""Placeholder for stdio transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
class SSETransportPlaceholder(BaseTransport):
"""Placeholder for SSE transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
class HTTPStreamTransportPlaceholder(BaseTransport):
"""Placeholder for HTTP Stream transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
class WebSocketTransportPlaceholder(BaseTransport):
"""Placeholder for WebSocket transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
# Default registry instance
_default_registry: Optional[TransportRegistry] = None
def get_default_registry() -> TransportRegistry:
"""
Get the default transport registry with built-in transports.
Returns:
TransportRegistry with default transports registered
"""
global _default_registry
if _default_registry is None:
_default_registry = TransportRegistry()
# Register default transports
_default_registry.register("stdio", StdioTransportPlaceholder)
_default_registry.register("sse", SSETransportPlaceholder)
_default_registry.register("http_stream", HTTPStreamTransportPlaceholder)
_default_registry.register("websocket", WebSocketTransportPlaceholder)
return _default_registry
def create_transport(
target: str,
config: Optional[TransportConfig] = None,
registry: Optional[TransportRegistry] = None
) -> BaseTransport:
"""
Create a transport instance for the given target.
Args:
target: URL or command string
config: Optional transport configuration
registry: Optional custom registry (uses default if None)
Returns:
Transport instance
Raises:
ValueError: If no suitable transport found
"""
if registry is None:
registry = get_default_registry()
if config is None:
config = TransportConfig()
transport_type = get_transport_type(target)
transport_class = registry.get(transport_type)
if transport_class is None:
raise ValueError(f"No transport registered for type: {transport_type}")
return transport_class()
@@ -0,0 +1,238 @@
"""MCP utility functions for converting Python functions to MCP tool schemas.
This module provides utilities to convert Python functions with type hints
into MCP-compatible tool schemas for exposing as MCP server tools.
Usage:
from praisonaiagents.mcp.mcp_utils import function_to_mcp_schema
def my_tool(query: str, max_results: int = 10) -> str:
'''Search for information.'''
return "results"
schema = function_to_mcp_schema(my_tool)
# Returns MCP-compatible tool schema
"""
import inspect
from typing import Any, Callable, Dict, List, Union, get_type_hints, get_origin, get_args
def python_type_to_json_schema(python_type: Any) -> Dict[str, Any]:
"""Convert a Python type annotation to JSON Schema.
Args:
python_type: A Python type annotation (e.g., str, int, List[str], Optional[int])
Returns:
JSON Schema dictionary representing the type
"""
# Handle None type
if python_type is None or python_type is type(None):
return {"type": "null"}
# Handle Any type
if python_type is Any:
return {}
# Handle basic types
type_mapping = {
str: {"type": "string"},
int: {"type": "integer"},
float: {"type": "number"},
bool: {"type": "boolean"},
bytes: {"type": "string"},
}
if python_type in type_mapping:
return type_mapping[python_type]
# Handle generic types (List, Dict, Optional, Union)
origin = get_origin(python_type)
args = get_args(python_type)
# Handle List[T]
if origin is list:
if args:
return {
"type": "array",
"items": python_type_to_json_schema(args[0])
}
return {"type": "array", "items": {}}
# Handle Dict[K, V]
if origin is dict:
return {"type": "object"}
# Handle Optional[T] (Union[T, None])
if origin is Union:
# Filter out None type
non_none_args = [arg for arg in args if arg is not type(None)]
if len(non_none_args) == 1:
# Optional[T] case
return python_type_to_json_schema(non_none_args[0])
elif len(non_none_args) > 1:
# Union of multiple types
return {"oneOf": [python_type_to_json_schema(arg) for arg in non_none_args]}
return {}
# Handle tuple
if origin is tuple:
if args:
return {
"type": "array",
"items": [python_type_to_json_schema(arg) for arg in args]
}
return {"type": "array"}
# Handle set
if origin is set:
if args:
return {
"type": "array",
"items": python_type_to_json_schema(args[0]),
"uniqueItems": True
}
return {"type": "array", "uniqueItems": True}
# Default to object for unknown types
return {"type": "object"}
def get_tool_metadata(func: Callable) -> Dict[str, str]:
"""Extract metadata from a function for MCP tool registration.
Checks for custom attributes (__mcp_name__, __mcp_description__) first,
then falls back to function name and docstring.
Args:
func: The function to extract metadata from
Returns:
Dictionary with 'name' and 'description' keys
"""
# Get name - check for custom attribute first
name = getattr(func, '__mcp_name__', None) or func.__name__
# Get description - check for custom attribute first
description = getattr(func, '__mcp_description__', None)
if description is None:
# Extract from docstring
docstring = inspect.getdoc(func)
if docstring:
# Use first line of docstring as description
description = docstring.split('\n')[0].strip()
else:
# Fall back to function name
description = name
return {
"name": name,
"description": description
}
def function_to_mcp_schema(func: Callable) -> Dict[str, Any]:
"""Convert a Python function to an MCP tool schema.
This function inspects the function's signature, type hints, and docstring
to generate a complete MCP-compatible tool schema.
Args:
func: The function to convert
Returns:
MCP tool schema dictionary with:
- name: Tool name
- description: Tool description
- inputSchema: JSON Schema for parameters
Example:
def search(query: str, max_results: int = 10) -> str:
'''Search the web for information.'''
return "results"
schema = function_to_mcp_schema(search)
# {
# "name": "search",
# "description": "Search the web for information.",
# "inputSchema": {
# "type": "object",
# "properties": {
# "query": {"type": "string"},
# "max_results": {"type": "integer"}
# },
# "required": ["query"]
# }
# }
"""
# Get metadata
metadata = get_tool_metadata(func)
# Get function signature
sig = inspect.signature(func)
# Try to get type hints
try:
type_hints = get_type_hints(func)
except Exception:
type_hints = {}
# Build properties and required list
properties = {}
required = []
for param_name, param in sig.parameters.items():
# Skip self, cls, and special parameters
if param_name in ('self', 'cls'):
continue
# Skip *args and **kwargs
if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
continue
# Get type annotation
param_type = type_hints.get(param_name, param.annotation)
if param_type is inspect.Parameter.empty:
# No type annotation - default to string
properties[param_name] = {"type": "string"}
else:
properties[param_name] = python_type_to_json_schema(param_type)
# Check if parameter is required (no default value)
if param.default is inspect.Parameter.empty:
# Check if it's Optional type
origin = get_origin(param_type) if param_type is not inspect.Parameter.empty else None
args = get_args(param_type) if param_type is not inspect.Parameter.empty else ()
is_optional = origin is Union and type(None) in args
if not is_optional:
required.append(param_name)
# Build the schema
schema = {
"name": metadata["name"],
"description": metadata["description"],
"inputSchema": {
"type": "object",
"properties": properties,
"required": required
}
}
return schema
def functions_to_mcp_tools(functions: List[Callable]) -> List[Dict[str, Any]]:
"""Convert multiple functions to MCP tool schemas.
Args:
functions: List of functions to convert
Returns:
List of MCP tool schemas
"""
return [function_to_mcp_schema(func) for func in functions]
@@ -0,0 +1,596 @@
"""
WebSocket Transport for MCP (Model Context Protocol).
This module provides WebSocket-based transport for connecting to MCP servers.
WebSocket transport offers:
- Long-lived bidirectional communication
- Better cloud provider support (hibernatable connections)
- Lower protocol overhead compared to HTTP
- Session persistence across network interruptions
This is an optional transport - the 'websockets' package is only imported
when ws:// or wss:// URLs are used, ensuring no performance impact on
existing stdio/SSE/HTTP transports.
Based on SEP-1288: WebSocket Transport for MCP (in review).
Protocol Revision: 2025-11-25
"""
import asyncio
import json
import logging
import threading
import inspect
import re
from typing import Any, Dict, Optional, List
logger = logging.getLogger("mcp-websocket")
def is_websocket_url(url: str) -> bool:
"""
Check if a URL is a WebSocket URL.
Args:
url: The URL string to check
Returns:
True if the URL starts with ws:// or wss://, False otherwise
"""
if not isinstance(url, str):
return False
return bool(re.match(r'^wss?://', url, re.IGNORECASE))
def calculate_backoff(attempt: int, base_delay: float = 1.0, max_delay: float = 60.0) -> float:
"""
Calculate exponential backoff delay for reconnection.
Args:
attempt: The retry attempt number (0-indexed)
base_delay: Base delay in seconds
max_delay: Maximum delay cap in seconds
Returns:
Delay in seconds for this attempt
"""
delay = base_delay * (2 ** attempt)
return min(delay, max_delay)
# Global event loop for WebSocket operations (similar to other MCP transports)
_event_loop = None
_loop_lock = threading.Lock()
def get_event_loop():
"""Get or create a global event loop for WebSocket operations."""
global _event_loop
with _loop_lock:
if _event_loop is None or _event_loop.is_closed():
_event_loop = asyncio.new_event_loop()
return _event_loop
class WebSocketTransport:
"""
Low-level WebSocket transport for MCP communication.
This class handles the WebSocket connection and JSON-RPC message framing.
It supports:
- Secure (wss://) and insecure (ws://) connections
- Session ID management
- Authentication via subprotocol or headers
- Automatic reconnection with exponential backoff
"""
def __init__(
self,
url: str,
session_id: Optional[str] = None,
auth_token: Optional[str] = None,
max_retries: int = 3,
retry_delay: float = 1.0,
timeout: float = 30.0,
ping_interval: float = 30.0,
ping_timeout: float = 10.0
):
"""
Initialize WebSocket transport.
Args:
url: WebSocket URL (ws:// or wss://)
session_id: Optional MCP session ID
auth_token: Optional authentication token
max_retries: Maximum reconnection attempts
retry_delay: Base delay between retries (exponential backoff)
timeout: Connection timeout in seconds
ping_interval: Interval between ping messages (keepalive)
ping_timeout: Timeout for ping response
"""
self.url = url
self.session_id = session_id
self.auth_token = auth_token
self.max_retries = max_retries
self.retry_delay = retry_delay
self.timeout = timeout
self.ping_interval = ping_interval
self.ping_timeout = ping_timeout
self._ws = None
self._closed = False
self._message_queue: asyncio.Queue = asyncio.Queue()
self._receive_task = None
async def __aenter__(self):
"""Async context manager entry - establish connection."""
await self.connect()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Async context manager exit - close connection."""
await self.close()
async def connect(self):
"""Establish WebSocket connection."""
try:
# Lazy import websockets
import websockets
except ImportError:
raise ImportError(
"websockets package is required for WebSocket transport. "
"Install it with: pip install websockets"
)
# Build connection options
extra_headers = {}
subprotocols = []
# Add session ID if present
if self.session_id:
extra_headers['Mcp-Session-Id'] = self.session_id
# Add auth token via subprotocol (SEP-1288 recommendation)
if self.auth_token:
# Smuggle auth token in subprotocol list
# Format: "mcp-auth-<base64-encoded-token>"
import base64
encoded_token = base64.b64encode(self.auth_token.encode()).decode()
subprotocols.append(f"mcp-auth-{encoded_token}")
# Always include MCP subprotocol
subprotocols.append("mcp")
# Connect with retry logic
last_error = None
for attempt in range(self.max_retries + 1):
try:
self._ws = await asyncio.wait_for(
websockets.connect(
self.url,
extra_headers=extra_headers if extra_headers else None,
subprotocols=subprotocols if subprotocols else None,
ping_interval=self.ping_interval,
ping_timeout=self.ping_timeout,
close_timeout=10
),
timeout=self.timeout
)
logger.debug(f"WebSocket connected to {self.url}")
# Start background receive task
self._receive_task = asyncio.create_task(self._receive_loop())
return
except Exception as e:
last_error = e
if attempt < self.max_retries:
delay = calculate_backoff(attempt, self.retry_delay)
logger.warning(f"WebSocket connection failed, retrying in {delay}s: {e}")
await asyncio.sleep(delay)
else:
raise ConnectionError(f"Failed to connect after {self.max_retries + 1} attempts: {last_error}")
async def _receive_loop(self):
"""Background task to receive messages and queue them."""
try:
async for message in self._ws:
try:
data = json.loads(message)
await self._message_queue.put(data)
except json.JSONDecodeError as e:
logger.error(f"Failed to parse WebSocket message: {e}")
except Exception as e:
if not self._closed:
logger.error(f"WebSocket receive error: {e}")
await self._message_queue.put({"error": str(e)})
async def send(self, message: Dict[str, Any]):
"""
Send a JSON-RPC message over WebSocket.
Args:
message: JSON-RPC message dictionary
"""
if self._ws is None:
raise RuntimeError("WebSocket not connected")
data = json.dumps(message)
await self._ws.send(data)
logger.debug(f"Sent: {message.get('method', message.get('id', 'response'))}")
async def receive(self) -> Dict[str, Any]:
"""
Receive a JSON-RPC message from WebSocket.
Returns:
Parsed JSON-RPC message dictionary
"""
message = await self._message_queue.get()
if "error" in message and len(message) == 1:
raise ConnectionError(message["error"])
return message
async def close(self):
"""Close WebSocket connection."""
self._closed = True
if self._receive_task:
self._receive_task.cancel()
try:
await self._receive_task
except asyncio.CancelledError:
pass
if self._ws:
await self._ws.close()
self._ws = None
logger.debug("WebSocket connection closed")
class WebSocketMCPTool:
"""
A wrapper for an MCP tool that can be used with praisonaiagents.
This class wraps a tool from a WebSocket MCP server and makes it
compatible with the Agent class tool interface.
"""
def __init__(
self,
name: str,
description: str,
session: Any, # ClientSession from mcp package
input_schema: Optional[Dict[str, Any]] = None,
timeout: int = 60
):
"""
Initialize WebSocket MCP tool wrapper.
Args:
name: Tool name
description: Tool description
session: MCP ClientSession for making calls
input_schema: JSON Schema for tool input
timeout: Timeout for tool calls in seconds
"""
self.name = name
self.__name__ = name
self.__qualname__ = name
self.__doc__ = description
self.description = description
self.session = session
self.input_schema = input_schema or {}
self.timeout = timeout
# Build function signature from input schema
self.__signature__ = self._build_signature()
def _build_signature(self) -> inspect.Signature:
"""Build function signature from input schema."""
params = []
if self.input_schema and 'properties' in self.input_schema:
required = self.input_schema.get('required', [])
for param_name, prop_schema in self.input_schema['properties'].items():
# Determine type annotation
prop_type = prop_schema.get('type', 'string') if isinstance(prop_schema, dict) else 'string'
annotation = self._json_type_to_python(prop_type)
# Determine default value
default = inspect.Parameter.empty if param_name in required else None
params.append(
inspect.Parameter(
name=param_name,
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
default=default,
annotation=annotation
)
)
return inspect.Signature(params)
def _json_type_to_python(self, json_type: str) -> type:
"""Convert JSON Schema type to Python type."""
type_map = {
'string': str,
'integer': int,
'number': float,
'boolean': bool,
'array': list,
'object': dict
}
return type_map.get(json_type, Any)
def __call__(self, **kwargs):
"""
Call the tool synchronously.
This wraps the async call for synchronous usage.
"""
logger.debug(f"Tool {self.name} called with args: {kwargs}")
loop = get_event_loop()
future = asyncio.run_coroutine_threadsafe(self._async_call(**kwargs), loop)
try:
return future.result(timeout=self.timeout)
except Exception as e:
logger.error(f"Error calling tool {self.name}: {e}")
return f"Error: {str(e)}"
async def _async_call(self, **kwargs):
"""Call the tool asynchronously."""
logger.debug(f"Async calling tool {self.name} with args: {kwargs}")
try:
result = await self.session.call_tool(self.name, kwargs)
# Extract text from result
if hasattr(result, 'content') and result.content:
if hasattr(result.content[0], 'text'):
return result.content[0].text
return str(result.content[0])
return str(result)
except Exception as e:
logger.error(f"Error in _async_call for {self.name}: {e}")
raise
def _fix_array_schemas(self, schema: Dict[str, Any]) -> Dict[str, Any]:
"""
Fix array schemas by adding missing 'items' attribute.
OpenAI requires array types to specify item types.
"""
if not isinstance(schema, dict):
return schema
fixed_schema = schema.copy()
# Fix array types at current level
if fixed_schema.get("type") == "array" and "items" not in fixed_schema:
fixed_schema["items"] = {"type": "string"}
# Recursively fix nested schemas
if "properties" in fixed_schema:
fixed_properties = {}
for prop_name, prop_schema in fixed_schema["properties"].items():
fixed_properties[prop_name] = self._fix_array_schemas(prop_schema)
fixed_schema["properties"] = fixed_properties
if "items" in fixed_schema:
fixed_schema["items"] = self._fix_array_schemas(fixed_schema["items"])
return fixed_schema
def to_openai_tool(self) -> Dict[str, Any]:
"""Convert the tool to OpenAI function calling format."""
fixed_schema = self._fix_array_schemas(self.input_schema)
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": fixed_schema
}
}
class WebSocketMCPClient:
"""
High-level client for connecting to MCP servers over WebSocket.
This client handles:
- Connection establishment and lifecycle
- Tool discovery and wrapping
- Session management
- Automatic reconnection
Example:
```python
from praisonaiagents.mcp.mcp_websocket import WebSocketMCPClient
client = WebSocketMCPClient("ws://localhost:8080/mcp")
for tool in client.tools:
print(tool.name)
```
"""
def __init__(
self,
server_url: str,
debug: bool = False,
timeout: int = 60,
auth_token: Optional[str] = None,
options: Optional[Dict[str, Any]] = None
):
"""
Initialize WebSocket MCP client.
Args:
server_url: WebSocket URL (ws:// or wss://)
debug: Enable debug logging
timeout: Timeout for operations in seconds
auth_token: Optional authentication token
options: Additional configuration options
"""
self.server_url = server_url
self.debug = debug
self.timeout = timeout
self.auth_token = auth_token
self.options = options or {}
self.session = None
self.transport = None
self.tools: List[WebSocketMCPTool] = []
self._closed = False
# Configure logging
if debug:
logger.setLevel(logging.DEBUG)
else:
logger.setLevel(logging.WARNING)
# Initialize connection
self._initialize()
def _initialize(self):
"""Initialize the connection and discover tools."""
loop = get_event_loop()
# Start event loop in background thread if not running
def run_event_loop():
asyncio.set_event_loop(loop)
loop.run_forever()
self.loop_thread = threading.Thread(target=run_event_loop, daemon=True)
self.loop_thread.start()
# Run initialization
future = asyncio.run_coroutine_threadsafe(self._async_initialize(), loop)
self.tools = future.result(timeout=self.timeout)
async def _async_initialize(self):
"""Asynchronously initialize connection and discover tools."""
logger.debug(f"Connecting to MCP server at {self.server_url}")
# Lazy import mcp package
try:
from mcp import ClientSession
except ImportError:
raise ImportError(
"mcp package is required. Install it with: pip install mcp"
)
# Create transport
self.transport = WebSocketTransport(
self.server_url,
auth_token=self.auth_token,
timeout=self.timeout
)
await self.transport.connect()
# Create read/write streams for ClientSession
# The ClientSession expects async iterators for read/write
read_stream = self._create_read_stream()
write_stream = self._create_write_stream()
# Create and initialize session
self._session_context = ClientSession(read_stream, write_stream)
self.session = await self._session_context.__aenter__()
await self.session.initialize()
# Discover tools
logger.debug("Listing tools...")
response = await self.session.list_tools()
tools_data = response.tools
logger.debug(f"Found {len(tools_data)} tools: {[t.name for t in tools_data]}")
# Create tool wrappers
tools = []
for tool in tools_data:
input_schema = tool.inputSchema if hasattr(tool, 'inputSchema') else None
wrapper = WebSocketMCPTool(
name=tool.name,
description=tool.description if hasattr(tool, 'description') else f"Call the {tool.name} tool",
session=self.session,
input_schema=input_schema,
timeout=self.timeout
)
tools.append(wrapper)
return tools
def _create_read_stream(self):
"""Create async iterator for reading from transport."""
async def read():
while not self._closed:
try:
message = await self.transport.receive()
yield message
except Exception as e:
if not self._closed:
logger.error(f"Read error: {e}")
break
return read()
def _create_write_stream(self):
"""Create write function for transport."""
async def write(message):
if hasattr(message, 'to_dict'):
message = message.to_dict()
await self.transport.send(message)
return write
def __iter__(self):
"""Return iterator over tools."""
return iter(self.tools)
def to_openai_tools(self) -> List[Dict[str, Any]]:
"""Convert all tools to OpenAI format."""
return [tool.to_openai_tool() for tool in self.tools]
async def aclose(self):
"""Async cleanup."""
if self._closed:
return
self._closed = True
try:
if hasattr(self, '_session_context') and self._session_context:
await self._session_context.__aexit__(None, None, None)
except Exception:
pass
try:
if self.transport:
await self.transport.close()
except Exception:
pass
def close(self):
"""Synchronous cleanup."""
if self._closed:
return
try:
loop = get_event_loop()
if not loop.is_closed():
future = asyncio.run_coroutine_threadsafe(self.aclose(), loop)
future.result(timeout=5)
except Exception:
self._closed = True
def __del__(self):
"""Cleanup on garbage collection."""
try:
if not self._closed:
self.close()
except Exception:
pass
@@ -8,8 +8,83 @@ This module provides memory management capabilities including:
- User memory for preferences/history
- Quality-based storage decisions
- Graph memory support via Mem0
- Session save/resume (like Gemini CLI)
- Context compression (like Gemini CLI)
- Checkpointing (like Gemini CLI)
- Rules management (like Cursor/Windsurf)
- Auto-generated memories (like Windsurf Cascade)
- Workflows (like Windsurf)
- Hooks (like Windsurf Cascade Hooks)
Memory Providers:
- FileMemory: Zero-dependency JSON file-based storage (default)
- Memory: Full-featured with SQLite, ChromaDB, Mem0, MongoDB support
- RulesManager: Rules/instructions file management
- AutoMemory: Automatic memory extraction from conversations
- WorkflowManager: Multi-step workflow execution
- HooksManager: Pre/post operation hooks
"""
from .memory import Memory
from .file_memory import FileMemory, create_memory
from .rules_manager import RulesManager, Rule, create_rules_manager
__all__ = ["Memory"]
# Lazy imports for optional modules to avoid dependency issues and improve startup time
def __getattr__(name):
if name == "Memory":
from .memory import Memory
return Memory
if name == "AutoMemory":
from .auto_memory import AutoMemory
return AutoMemory
if name == "AutoMemoryExtractor":
from .auto_memory import AutoMemoryExtractor
return AutoMemoryExtractor
if name == "create_auto_memory":
from .auto_memory import create_auto_memory
return create_auto_memory
if name == "WorkflowManager":
from .workflows import WorkflowManager
return WorkflowManager
if name == "Workflow":
from .workflows import Workflow
return Workflow
if name == "WorkflowStep":
from .workflows import WorkflowStep
return WorkflowStep
if name == "create_workflow_manager":
from .workflows import create_workflow_manager
return create_workflow_manager
if name == "HooksManager":
from .hooks import HooksManager
return HooksManager
if name == "HookResult":
from .hooks import HookResult
return HookResult
if name == "create_hooks_manager":
from .hooks import create_hooks_manager
return create_hooks_manager
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
# Core memory
"FileMemory",
"Memory",
"create_memory",
# Rules management
"RulesManager",
"Rule",
"create_rules_manager",
# Auto memory
"AutoMemory",
"AutoMemoryExtractor",
"create_auto_memory",
# Workflows
"WorkflowManager",
"Workflow",
"WorkflowStep",
"create_workflow_manager",
# Hooks
"HooksManager",
"HookResult",
"create_hooks_manager",
]
@@ -0,0 +1,382 @@
"""
Auto-Generated Memory System for PraisonAI Agents.
Provides automatic memory generation similar to Windsurf's Cascade,
where the AI automatically decides what context is worth remembering.
Features:
- Automatic extraction of important facts from conversations
- Entity recognition and storage
- Preference detection
- Lazy evaluation to avoid performance impact
- Configurable importance thresholds
This module is designed to be lightweight and only activates
when explicitly enabled to avoid impacting performance.
"""
import re
import logging
from typing import Any, Dict, List, Optional, Callable, TYPE_CHECKING
if TYPE_CHECKING:
from .file_memory import FileMemory
logger = logging.getLogger(__name__)
class AutoMemoryExtractor:
"""
Automatically extracts memorable content from conversations.
Uses pattern matching for fast extraction without LLM calls,
with optional LLM enhancement for better accuracy.
Example:
```python
extractor = AutoMemoryExtractor()
# Extract from conversation
memories = extractor.extract(
"My name is John and I prefer Python for backend work."
)
# Returns: [
# {"type": "entity", "name": "John", "entity_type": "person"},
# {"type": "preference", "content": "prefers Python for backend work"}
# ]
```
"""
# Patterns for fast extraction (no LLM needed)
PATTERNS = {
# Name patterns
"name": [
r"(?:my name is|i'm|i am|call me)\s+([A-Z][a-z]+)",
r"(?:name|Name):\s*([A-Z][a-z]+)",
],
# Preference patterns
"preference": [
r"i (?:prefer|like|love|enjoy|use)\s+(.+?)(?:\.|$)",
r"(?:my favorite|i always use)\s+(.+?)(?:\.|$)",
],
# Role/job patterns
"role": [
r"i (?:am|work as|'m)\s+(?:a|an)\s+([a-z\s]+?)(?:\.|,|$)",
r"(?:my job|my role|i do)\s+(?:is|as)?\s*(.+?)(?:\.|$)",
],
# Location patterns
"location": [
r"i (?:live|work|am based|am located)\s+(?:in|at)\s+([A-Z][a-z\s]+)",
r"(?:from|based in)\s+([A-Z][a-z\s,]+)",
],
# Project/task patterns
"project": [
r"(?:working on|building|developing|creating)\s+(?:a|an|the)?\s*(.+?)(?:\.|$)",
r"(?:my project|the project)\s+(?:is|called)?\s*(.+?)(?:\.|$)",
],
# Technology preferences
"technology": [
r"(?:using|use|prefer)\s+(Python|JavaScript|TypeScript|Go|Rust|Java|C\+\+|Ruby|PHP)",
r"(?:framework|library|tool)(?:s)?\s+(?:like|such as|including)\s+(.+?)(?:\.|$)",
],
}
# Importance scores by type
IMPORTANCE_SCORES = {
"name": 0.95,
"preference": 0.7,
"role": 0.85,
"location": 0.6,
"project": 0.75,
"technology": 0.7,
}
def __init__(
self,
min_importance: float = 0.5,
use_llm: bool = False,
llm_func: Optional[Callable[[str], str]] = None,
verbose: int = 0
):
"""
Initialize AutoMemoryExtractor.
Args:
min_importance: Minimum importance score to extract (0.0-1.0)
use_llm: Whether to use LLM for enhanced extraction
llm_func: LLM function for extraction (if use_llm=True)
verbose: Verbosity level
"""
self.min_importance = min_importance
self.use_llm = use_llm
self.llm_func = llm_func
self.verbose = verbose
# Compile patterns for performance
self._compiled_patterns = {}
for category, patterns in self.PATTERNS.items():
self._compiled_patterns[category] = [
re.compile(p, re.IGNORECASE) for p in patterns
]
def extract(self, text: str) -> List[Dict[str, Any]]:
"""
Extract memorable content from text.
Args:
text: Text to extract from (conversation, message, etc.)
Returns:
List of extracted memories with type, content, and importance
"""
memories = []
# Fast pattern-based extraction
for category, patterns in self._compiled_patterns.items():
importance = self.IMPORTANCE_SCORES.get(category, 0.5)
if importance < self.min_importance:
continue
for pattern in patterns:
matches = pattern.findall(text)
for match in matches:
if isinstance(match, tuple):
match = match[0]
match = match.strip()
if len(match) < 2 or len(match) > 200:
continue
memory = {
"type": category,
"content": match,
"importance": importance,
"source": "pattern"
}
# Avoid duplicates
if not any(m["content"].lower() == match.lower() for m in memories):
memories.append(memory)
# Optional LLM enhancement
if self.use_llm and self.llm_func and len(text) > 50:
llm_memories = self._extract_with_llm(text)
for mem in llm_memories:
if mem.get("importance", 0) >= self.min_importance:
if not any(m["content"].lower() == mem["content"].lower() for m in memories):
memories.append(mem)
return memories
def _extract_with_llm(self, text: str) -> List[Dict[str, Any]]:
"""Extract memories using LLM for better accuracy."""
if not self.llm_func:
return []
prompt = f"""Extract important facts from this text that should be remembered.
Return a JSON array of objects with: type (name/preference/role/location/project/technology), content, importance (0.0-1.0).
Only extract clear, specific facts. Be concise.
Text: {text[:1500]}
JSON array (no markdown, just the array):"""
try:
response = self.llm_func(prompt)
# Parse JSON from response
import json
# Find JSON array in response
start = response.find("[")
end = response.rfind("]") + 1
if start >= 0 and end > start:
return json.loads(response[start:end])
except Exception as e:
if self.verbose:
logger.warning(f"LLM extraction failed: {e}")
return []
def should_remember(self, text: str) -> bool:
"""
Quick check if text likely contains memorable content.
Useful for filtering before full extraction.
Args:
text: Text to check
Returns:
True if text likely contains memorable content
"""
# Quick keyword check
keywords = [
"my name", "i am", "i'm", "i prefer", "i like", "i use",
"i work", "i live", "working on", "building", "my project",
"favorite", "always use"
]
text_lower = text.lower()
return any(kw in text_lower for kw in keywords)
class AutoMemory:
"""
Wrapper that adds auto-generation capabilities to FileMemory.
Intercepts interactions and automatically extracts/stores memories.
Designed to be lightweight - only processes when enabled.
Example:
```python
from praisonaiagents.memory import FileMemory
base_memory = FileMemory(user_id="user123")
auto_memory = AutoMemory(base_memory, enabled=True)
# Process conversation and auto-store memories
auto_memory.process_interaction(
user_message="My name is John and I prefer Python",
assistant_response="Nice to meet you, John!"
)
```
"""
def __init__(
self,
memory: "FileMemory",
enabled: bool = True,
extractor: Optional[AutoMemoryExtractor] = None,
min_importance: float = 0.6,
verbose: int = 0
):
"""
Initialize AutoMemory wrapper.
Args:
memory: Base FileMemory instance
enabled: Whether auto-generation is enabled
extractor: Custom extractor (or use default)
min_importance: Minimum importance to auto-store
verbose: Verbosity level
"""
self.memory = memory
self.enabled = enabled
self.min_importance = min_importance
self.verbose = verbose
self.extractor = extractor or AutoMemoryExtractor(
min_importance=min_importance,
verbose=verbose
)
# Track what we've already processed to avoid duplicates
self._processed_hashes: set = set()
def process_interaction(
self,
user_message: str,
assistant_response: Optional[str] = None,
store: bool = True
) -> List[Dict[str, Any]]:
"""
Process an interaction and extract memories.
Args:
user_message: User's message
assistant_response: Optional assistant response
store: Whether to store extracted memories
Returns:
List of extracted memories
"""
if not self.enabled:
return []
# Combine messages for extraction
text = user_message
if assistant_response:
text += "\n" + assistant_response
# Check if already processed
import hashlib
text_hash = hashlib.md5(text.encode()).hexdigest()[:16]
if text_hash in self._processed_hashes:
return []
self._processed_hashes.add(text_hash)
# Quick filter
if not self.extractor.should_remember(text):
return []
# Extract memories
memories = self.extractor.extract(text)
if store and memories:
self._store_memories(memories)
return memories
def _store_memories(self, memories: List[Dict[str, Any]]):
"""Store extracted memories in the base memory."""
for mem in memories:
mem_type = mem.get("type", "")
content = mem.get("content", "")
importance = mem.get("importance", 0.5)
if not content:
continue
# Store based on type
if mem_type == "name":
# Store as entity
self.memory.add_entity(
name=content,
entity_type="person",
attributes={"relationship": "user"}
)
elif mem_type in ("role", "location"):
# Store as entity attribute
self.memory.add_long_term(
f"User's {mem_type}: {content}",
importance=importance
)
else:
# Store as long-term memory
self.memory.add_long_term(content, importance=importance)
if self.verbose:
logger.info(f"Auto-stored memory: {mem_type} - {content[:50]}...")
def enable(self):
"""Enable auto-generation."""
self.enabled = True
def disable(self):
"""Disable auto-generation."""
self.enabled = False
def set_llm(self, llm_func: Callable[[str], str]):
"""Set LLM function for enhanced extraction."""
self.extractor.use_llm = True
self.extractor.llm_func = llm_func
def create_auto_memory(
memory: "FileMemory",
enabled: bool = True,
**kwargs
) -> AutoMemory:
"""
Create an AutoMemory wrapper.
Args:
memory: Base FileMemory instance
enabled: Whether auto-generation is enabled
**kwargs: Additional configuration
Returns:
AutoMemory instance
"""
return AutoMemory(memory, enabled=enabled, **kwargs)
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,417 @@
"""
Cascade Hooks System for PraisonAI Agents.
Provides event hooks similar to Windsurf's Cascade Hooks,
allowing custom actions before/after agent operations.
Features:
- Pre/post hooks for read, write, command, prompt operations
- JSON configuration (.praison/hooks.json)
- Script execution with timeout
- Exit code handling
- Lazy loading for performance
Configuration (.praison/hooks.json):
{
"hooks": {
"pre_write_code": "./scripts/lint.sh",
"post_write_code": "./scripts/format.sh",
"pre_run_command": "./scripts/validate.sh",
"post_user_prompt": "./scripts/log.sh"
},
"timeout": 30,
"enabled": true
}
Hook Events:
- pre_read_code: Before reading a file
- post_read_code: After reading a file
- pre_write_code: Before writing to a file
- post_write_code: After writing to a file
- pre_run_command: Before running a terminal command
- post_run_command: After running a terminal command
- pre_user_prompt: Before processing user prompt
- post_user_prompt: After processing user prompt
"""
import os
import json
import logging
import subprocess
from pathlib import Path
from typing import Any, Dict, List, Optional, Callable, Literal
from dataclasses import dataclass
logger = logging.getLogger(__name__)
# Hook event types
HookEvent = Literal[
"pre_read_code",
"post_read_code",
"pre_write_code",
"post_write_code",
"pre_run_command",
"post_run_command",
"pre_user_prompt",
"post_user_prompt",
"pre_mcp_tool_use",
"post_mcp_tool_use"
]
@dataclass
class HookResult:
"""Result of a hook execution."""
success: bool
exit_code: int = 0
stdout: str = ""
stderr: str = ""
blocked: bool = False # If True, the operation should be blocked
modified_input: Optional[str] = None # Modified input if applicable
@dataclass
class HookConfig:
"""Configuration for a single hook."""
event: HookEvent
command: str
timeout: int = 30
enabled: bool = True
block_on_failure: bool = False
pass_input: bool = True
class HooksManager:
"""
Manages hook discovery, loading, and execution.
Hooks are configured in .praison/hooks.json and can execute
scripts before/after agent operations.
Example:
```python
hooks = HooksManager(workspace_path="/path/to/project")
# Execute pre-write hook
result = hooks.execute(
"pre_write_code",
context={"file_path": "main.py", "content": "..."}
)
if result.blocked:
print("Write blocked by hook")
```
"""
CONFIG_FILE = ".praison/hooks.json"
def __init__(
self,
workspace_path: Optional[str] = None,
verbose: int = 0
):
"""
Initialize HooksManager.
Args:
workspace_path: Path to workspace/project root
verbose: Verbosity level
"""
self.workspace_path = Path(workspace_path) if workspace_path else Path.cwd()
self.verbose = verbose
self._hooks: Dict[HookEvent, List[HookConfig]] = {}
self._global_timeout = 30
self._enabled = True
self._loaded = False
# Python callable hooks (for programmatic use)
self._callable_hooks: Dict[HookEvent, List[Callable]] = {}
def _log(self, msg: str, level: int = logging.INFO):
"""Log message if verbose."""
if self.verbose >= 1:
logger.log(level, msg)
def _ensure_loaded(self):
"""Lazy load hooks on first access."""
if not self._loaded:
self._load_config()
self._loaded = True
def _load_config(self):
"""Load hooks configuration from file."""
config_path = self.workspace_path / self.CONFIG_FILE.replace("/", os.sep)
if not config_path.exists():
return
try:
config = json.loads(config_path.read_text(encoding="utf-8"))
self._enabled = config.get("enabled", True)
self._global_timeout = config.get("timeout", 30)
hooks_config = config.get("hooks", {})
for event, command in hooks_config.items():
if event not in self._hooks:
self._hooks[event] = []
if isinstance(command, str):
# Simple command string
self._hooks[event].append(HookConfig(
event=event,
command=command,
timeout=self._global_timeout
))
elif isinstance(command, dict):
# Detailed config
self._hooks[event].append(HookConfig(
event=event,
command=command.get("command", ""),
timeout=command.get("timeout", self._global_timeout),
enabled=command.get("enabled", True),
block_on_failure=command.get("block_on_failure", False),
pass_input=command.get("pass_input", True)
))
elif isinstance(command, list):
# Multiple commands for same event
for cmd in command:
if isinstance(cmd, str):
self._hooks[event].append(HookConfig(
event=event,
command=cmd,
timeout=self._global_timeout
))
self._log(f"Loaded {sum(len(h) for h in self._hooks.values())} hooks")
except Exception as e:
self._log(f"Error loading hooks config: {e}", logging.WARNING)
def reload(self):
"""Reload hooks configuration."""
self._hooks = {}
self._loaded = False
self._ensure_loaded()
def register(
self,
event: HookEvent,
handler: Callable[[Dict[str, Any]], Optional[Dict[str, Any]]]
):
"""
Register a Python callable as a hook.
Args:
event: Hook event type
handler: Callable that receives context dict and optionally returns modified context
"""
if event not in self._callable_hooks:
self._callable_hooks[event] = []
self._callable_hooks[event].append(handler)
def unregister(self, event: HookEvent, handler: Callable):
"""Unregister a callable hook."""
if event in self._callable_hooks:
self._callable_hooks[event] = [
h for h in self._callable_hooks[event] if h != handler
]
def execute(
self,
event: HookEvent,
context: Optional[Dict[str, Any]] = None
) -> HookResult:
"""
Execute hooks for an event.
Args:
event: Hook event type
context: Context data to pass to hooks
Returns:
HookResult with execution status
"""
self._ensure_loaded()
if not self._enabled:
return HookResult(success=True)
context = context or {}
combined_result = HookResult(success=True)
# Execute callable hooks first
if event in self._callable_hooks:
for handler in self._callable_hooks[event]:
try:
result = handler(context)
if isinstance(result, dict):
context.update(result)
except Exception as e:
self._log(f"Callable hook error for {event}: {e}", logging.WARNING)
combined_result.success = False
# Execute script hooks
if event in self._hooks:
for hook in self._hooks[event]:
if not hook.enabled:
continue
result = self._execute_script(hook, context)
if not result.success:
combined_result.success = False
combined_result.stderr += result.stderr
if hook.block_on_failure:
combined_result.blocked = True
break
combined_result.stdout += result.stdout
if result.modified_input:
combined_result.modified_input = result.modified_input
return combined_result
def _execute_script(
self,
hook: HookConfig,
context: Dict[str, Any]
) -> HookResult:
"""Execute a script hook."""
try:
# Prepare environment
env = os.environ.copy()
# Pass context as environment variables
if hook.pass_input:
for key, value in context.items():
if isinstance(value, (str, int, float, bool)):
env[f"PRAISON_HOOK_{key.upper()}"] = str(value)
# Also pass as JSON
env["PRAISON_HOOK_CONTEXT"] = json.dumps(context)
# Resolve command path
command = hook.command
if not os.path.isabs(command) and not command.startswith("./"):
command = str(self.workspace_path / command)
# Execute
result = subprocess.run(
command,
shell=True,
cwd=str(self.workspace_path),
env=env,
capture_output=True,
text=True,
timeout=hook.timeout
)
success = result.returncode == 0
# Check for special exit codes
blocked = result.returncode == 1 and hook.block_on_failure
# Check stdout for modified input
modified_input = None
if result.stdout.startswith("MODIFIED:"):
modified_input = result.stdout[9:].strip()
self._log(f"Hook {hook.event} executed: exit={result.returncode}")
return HookResult(
success=success,
exit_code=result.returncode,
stdout=result.stdout,
stderr=result.stderr,
blocked=blocked,
modified_input=modified_input
)
except subprocess.TimeoutExpired:
self._log(f"Hook {hook.event} timed out after {hook.timeout}s", logging.WARNING)
return HookResult(
success=False,
exit_code=-1,
stderr=f"Timeout after {hook.timeout}s"
)
except Exception as e:
self._log(f"Hook {hook.event} error: {e}", logging.WARNING)
return HookResult(
success=False,
exit_code=-1,
stderr=str(e)
)
def has_hooks(self, event: HookEvent) -> bool:
"""Check if any hooks are registered for an event."""
self._ensure_loaded()
return (
event in self._hooks and len(self._hooks[event]) > 0 or
event in self._callable_hooks and len(self._callable_hooks[event]) > 0
)
def get_stats(self) -> Dict[str, Any]:
"""Get hooks statistics."""
self._ensure_loaded()
script_hooks = sum(len(h) for h in self._hooks.values())
callable_hooks = sum(len(h) for h in self._callable_hooks.values())
return {
"enabled": self._enabled,
"script_hooks": script_hooks,
"callable_hooks": callable_hooks,
"total_hooks": script_hooks + callable_hooks,
"events": list(set(list(self._hooks.keys()) + list(self._callable_hooks.keys())))
}
def create_config(
self,
hooks: Optional[Dict[str, str]] = None,
timeout: int = 30,
enabled: bool = True
):
"""
Create a hooks configuration file.
Args:
hooks: Dict mapping event names to commands
timeout: Global timeout in seconds
enabled: Whether hooks are enabled
"""
config_dir = self.workspace_path / ".praison"
config_dir.mkdir(parents=True, exist_ok=True)
config = {
"enabled": enabled,
"timeout": timeout,
"hooks": hooks or {}
}
config_path = config_dir / "hooks.json"
config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
self._log(f"Created hooks config at {config_path}")
self.reload()
def create_hooks_manager(
workspace_path: Optional[str] = None,
**kwargs
) -> HooksManager:
"""
Create a HooksManager instance.
Args:
workspace_path: Path to workspace
**kwargs: Additional configuration
Returns:
HooksManager instance
"""
return HooksManager(workspace_path=workspace_path, **kwargs)
@@ -0,0 +1,819 @@
"""
Rules Manager for PraisonAI Agents.
Provides persistent rules/instructions support similar to:
- Cursor (.cursor/rules/*.mdc)
- Windsurf (.windsurf/rules/)
- Claude Code (CLAUDE.md, .claude/rules/)
- Codex CLI (AGENTS.md)
- Gemini CLI (GEMINI.md)
Features:
- Auto-discovery of root instruction files (CLAUDE.md, AGENTS.md, etc.)
- Git root discovery for monorepo support
- @Import syntax for including other files
- Local override files (CLAUDE.local.md)
- Multiple activation modes (always, glob, manual, ai_decision)
- Character limits for context window management
- Symlink support for shared rules
Storage Structure:
.praison/rules/
├── global.md # Always applied (global)
├── python.md # Glob: **/*.py
├── testing.md # Glob: **/*.test.*
└── security.md # Activation: manual (@security)
Rule File Format (YAML frontmatter + Markdown):
---
description: Python coding guidelines
globs: ["**/*.py", "**/*.pyx"]
activation: always # always, glob, manual, ai_decision
priority: 10
---
# Python Guidelines
- Use type hints
- Follow PEP 8
"""
import os
import re
import fnmatch
import logging
import subprocess
from pathlib import Path
from typing import Any, Dict, List, Optional, Literal, Callable
from dataclasses import dataclass, field
logger = logging.getLogger(__name__)
@dataclass
class Rule:
"""A single rule with metadata and content."""
name: str
content: str
description: str = ""
globs: List[str] = field(default_factory=list)
activation: Literal["always", "glob", "manual", "ai_decision"] = "always"
priority: int = 0
file_path: Optional[str] = None
def matches_file(self, file_path: str) -> bool:
"""Check if this rule matches a given file path."""
if self.activation == "always":
return True
if self.activation == "manual":
return False # Only activated via @mention
if self.activation == "glob" and self.globs:
for pattern in self.globs:
if fnmatch.fnmatch(file_path, pattern):
return True
# Also try with ** expansion
if "**" in pattern:
# Convert ** to regex for recursive matching
regex_pattern = pattern.replace("**", ".*").replace("*", "[^/]*")
if re.match(regex_pattern, file_path):
return True
return False
def to_dict(self) -> Dict[str, Any]:
return {
"name": self.name,
"content": self.content,
"description": self.description,
"globs": self.globs,
"activation": self.activation,
"priority": self.priority,
"file_path": self.file_path
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> 'Rule':
return cls(**data)
class RulesManager:
"""
Manages rules/instructions for AI agents.
Supports multiple rule sources:
- Global rules (~/.praison/rules/)
- Workspace rules (.praison/rules/)
- Subdirectory rules (subdir/.praison/rules/)
- Root instruction files (CLAUDE.md, AGENTS.md, GEMINI.md, .cursorrules, .windsurfrules)
Example:
```python
rules = RulesManager(workspace_path="/path/to/project")
# Get all active rules
active_rules = rules.get_active_rules()
# Get rules for specific file
python_rules = rules.get_rules_for_file("src/main.py")
# Get manually invoked rule
security_rule = rules.get_rule_by_name("security")
# Build context string for LLM
context = rules.build_rules_context(file_path="src/main.py")
```
"""
RULES_DIR_NAME = ".praison/rules"
SUPPORTED_EXTENSIONS = [".md", ".mdc", ".txt"]
# Root-level instruction files (like Claude, Codex, Gemini, Cursor, Windsurf)
ROOT_INSTRUCTION_FILES = [
"CLAUDE.md", # Claude Code memory file
"CLAUDE.local.md", # Claude Code local overrides (gitignored)
"AGENTS.md", # OpenAI Codex CLI instructions
"GEMINI.md", # Gemini CLI memory file
".cursorrules", # Cursor legacy rules (deprecated but supported)
".windsurfrules", # Windsurf legacy rules
"PRAISON.md", # PraisonAI native instructions
"PRAISON.local.md", # PraisonAI local overrides (gitignored)
]
# Additional rules directories to discover (like Claude, Windsurf)
ADDITIONAL_RULES_DIRS = [
".claude/rules", # Claude Code modular rules
".windsurf/rules", # Windsurf rules
".cursor/rules", # Cursor rules
]
# Maximum characters per rule (like Windsurf's 12000 limit)
MAX_RULE_CHARS = 12000
# Import pattern for @path/to/file syntax
IMPORT_PATTERN = re.compile(r'(?<!`)@([\w./-]+)(?!`)')
def __init__(
self,
workspace_path: Optional[str] = None,
global_rules_path: Optional[str] = None,
verbose: int = 0
):
"""
Initialize RulesManager.
Args:
workspace_path: Path to workspace/project root
global_rules_path: Path to global rules (default: ~/.praison/rules)
verbose: Verbosity level
"""
self.workspace_path = Path(workspace_path) if workspace_path else Path.cwd()
self.global_rules_path = Path(global_rules_path) if global_rules_path else Path.home() / ".praison" / "rules"
self.verbose = verbose
self._rules: Dict[str, Rule] = {}
self._load_all_rules()
def _log(self, msg: str, level: int = logging.INFO):
"""Log message if verbose."""
if self.verbose >= 1:
logger.log(level, msg)
def _find_git_root(self, start_path: Optional[Path] = None) -> Optional[Path]:
"""
Find the git repository root from a starting path.
Uses pure Python (no gitpython dependency) for performance.
Falls back to subprocess if needed.
Args:
start_path: Starting path (defaults to workspace_path)
Returns:
Path to git root, or None if not in a git repo
"""
start = start_path or self.workspace_path
# Method 1: Walk up looking for .git directory (fastest, no subprocess)
current = Path(start).resolve()
while current != current.parent:
git_dir = current / ".git"
if git_dir.exists():
return current
current = current.parent
# Method 2: Fallback to git command (handles edge cases like worktrees)
try:
result = subprocess.run(
["git", "rev-parse", "--show-toplevel"],
cwd=str(start),
capture_output=True,
text=True,
timeout=2 # Prevent hanging
)
if result.returncode == 0:
return Path(result.stdout.strip())
except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
pass
return None
def _process_imports(self, content: str, base_path: Path, depth: int = 0) -> str:
"""
Process @import syntax in rule content.
Supports:
- @path/to/file - Import relative to workspace
- @~/path/to/file - Import from home directory
- @./path/to/file - Import relative to current file
Args:
content: Rule content with potential @imports
base_path: Base path for relative imports
depth: Current recursion depth (max 5 to prevent loops)
Returns:
Content with imports resolved
"""
if depth > 5:
self._log("Import depth exceeded (max 5), stopping", logging.WARNING)
return content
def replace_import(match):
import_path = match.group(1)
# Skip if it looks like a mention (@username) or package (@org/pkg)
if "/" not in import_path and "." not in import_path:
return match.group(0)
# Resolve path
if import_path.startswith("~/"):
file_path = Path.home() / import_path[2:]
elif import_path.startswith("./"):
file_path = base_path / import_path[2:]
else:
file_path = self.workspace_path / import_path
# Try with and without .md extension
if not file_path.exists() and not file_path.suffix:
file_path = file_path.with_suffix(".md")
if file_path.exists() and file_path.is_file():
try:
imported_content = file_path.read_text(encoding="utf-8")
# Recursively process imports in imported content
imported_content = self._process_imports(
imported_content,
file_path.parent,
depth + 1
)
return imported_content
except Exception as e:
self._log(f"Failed to import {file_path}: {e}", logging.WARNING)
return match.group(0)
else:
# Not a file import, leave as-is (might be @mention)
return match.group(0)
return self.IMPORT_PATTERN.sub(replace_import, content)
def _truncate_rule_content(self, content: str) -> str:
"""Truncate rule content to MAX_RULE_CHARS if needed."""
if len(content) <= self.MAX_RULE_CHARS:
return content
truncated = content[:self.MAX_RULE_CHARS - 50]
return truncated + "\n\n... (truncated, exceeded 12000 char limit)"
def _parse_frontmatter(self, content: str) -> tuple[Dict[str, Any], str]:
"""Parse YAML frontmatter from markdown content."""
frontmatter = {}
body = content
# Check for YAML frontmatter (--- ... ---)
if content.startswith("---"):
parts = content.split("---", 2)
if len(parts) >= 3:
yaml_content = parts[1].strip()
body = parts[2].strip()
# Simple YAML parsing (avoid dependency)
for line in yaml_content.split("\n"):
line = line.strip()
if ":" in line:
key, value = line.split(":", 1)
key = key.strip()
value = value.strip()
# Handle lists
if value.startswith("[") and value.endswith("]"):
# Parse simple list: ["*.py", "*.pyx"]
value = [v.strip().strip('"\'') for v in value[1:-1].split(",") if v.strip()]
# Handle booleans
elif value.lower() in ("true", "false"):
value = value.lower() == "true"
# Handle numbers
elif value.isdigit():
value = int(value)
# Handle quoted strings
elif value.startswith('"') and value.endswith('"'):
value = value[1:-1]
elif value.startswith("'") and value.endswith("'"):
value = value[1:-1]
frontmatter[key] = value
return frontmatter, body
def _load_rule_file(self, file_path: Path) -> Optional[Rule]:
"""Load a single rule file with @import processing and truncation."""
try:
content = file_path.read_text(encoding="utf-8")
frontmatter, body = self._parse_frontmatter(content)
# Process @imports in the body
body = self._process_imports(body, file_path.parent)
# Truncate if exceeds limit
body = self._truncate_rule_content(body)
# Extract rule name from filename
name = file_path.stem
# Build Rule object
rule = Rule(
name=name,
content=body,
description=frontmatter.get("description", ""),
globs=frontmatter.get("globs", []),
activation=frontmatter.get("activation", "always"),
priority=frontmatter.get("priority", 0),
file_path=str(file_path)
)
self._log(f"Loaded rule '{name}' from {file_path}")
return rule
except Exception as e:
self._log(f"Error loading rule file {file_path}: {e}", logging.WARNING)
return None
def _discover_rules_in_dir(self, rules_dir: Path) -> List[Rule]:
"""Discover all rule files in a directory."""
rules = []
if not rules_dir.exists():
return rules
for ext in self.SUPPORTED_EXTENSIONS:
for file_path in rules_dir.glob(f"*{ext}"):
rule = self._load_rule_file(file_path)
if rule:
rules.append(rule)
return rules
def _load_root_instruction_file(self, file_path: Path) -> Optional[Rule]:
"""Load a root instruction file (CLAUDE.md, AGENTS.md, etc.) with @import processing."""
try:
if not file_path.exists():
return None
# Handle symlinks (like Claude Code supports)
if file_path.is_symlink():
file_path = file_path.resolve()
if not file_path.exists():
return None
content = file_path.read_text(encoding="utf-8")
frontmatter, body = self._parse_frontmatter(content)
# Process @imports in the body
body_content = body if body else content
body_content = self._process_imports(body_content, file_path.parent)
# Truncate if exceeds limit
body_content = self._truncate_rule_content(body_content)
# Use filename (without extension) as rule name
name = file_path.stem.lower().replace(".", "_")
# Determine description based on file type
descriptions = {
"claude": "Claude Code memory instructions",
"claude_local": "Claude Code local overrides",
"agents": "OpenAI Codex CLI instructions",
"gemini": "Gemini CLI memory instructions",
"cursorrules": "Cursor IDE rules (legacy)",
"windsurfrules": "Windsurf IDE rules (legacy)",
"praison": "PraisonAI native instructions",
"praison_local": "PraisonAI local overrides",
}
description = descriptions.get(name, f"Instructions from {file_path.name}")
# Local files get higher priority (override main files)
base_priority = 600 if "_local" in name else 500
rule = Rule(
name=name,
content=body_content,
description=frontmatter.get("description", description),
globs=frontmatter.get("globs", []),
activation=frontmatter.get("activation", "always"),
priority=frontmatter.get("priority", base_priority),
file_path=str(file_path)
)
self._log(f"Loaded root instruction file '{file_path.name}'")
return rule
except Exception as e:
self._log(f"Error loading root instruction file {file_path}: {e}", logging.WARNING)
return None
def _load_all_rules(self):
"""Load rules from all sources including git root and additional directories."""
self._rules = {}
# Find git root for monorepo support
git_root = self._find_git_root()
# 1. Load global rules (lowest priority)
global_rules = self._discover_rules_in_dir(self.global_rules_path)
for rule in global_rules:
rule.priority = rule.priority - 1000 # Lower priority for global
self._rules[f"global:{rule.name}"] = rule
# 2. Load root instruction files from git root (if different from workspace)
if git_root and git_root != self.workspace_path:
for filename in self.ROOT_INSTRUCTION_FILES:
file_path = git_root / filename
rule = self._load_root_instruction_file(file_path)
if rule:
rule.priority = rule.priority - 100 # Slightly lower than workspace
self._rules[f"gitroot:{rule.name}"] = rule
# 3. Load root instruction files from workspace (CLAUDE.md, AGENTS.md, etc.)
for filename in self.ROOT_INSTRUCTION_FILES:
file_path = self.workspace_path / filename
rule = self._load_root_instruction_file(file_path)
if rule:
self._rules[f"root:{rule.name}"] = rule
# 4. Load workspace rules from .praison/rules/
workspace_rules_dir = self.workspace_path / self.RULES_DIR_NAME.replace("/", os.sep)
workspace_rules = self._discover_rules_in_dir(workspace_rules_dir)
for rule in workspace_rules:
self._rules[f"workspace:{rule.name}"] = rule
# 5. Load additional rules directories (.claude/rules, .windsurf/rules, .cursor/rules)
for rules_dir_name in self.ADDITIONAL_RULES_DIRS:
# Check workspace
rules_dir = self.workspace_path / rules_dir_name.replace("/", os.sep)
additional_rules = self._discover_rules_in_dir(rules_dir)
for rule in additional_rules:
rule.priority = rule.priority + 50 # Slightly higher priority
dir_prefix = rules_dir_name.split("/")[0].strip(".")
self._rules[f"{dir_prefix}:{rule.name}"] = rule
# Check git root if different
if git_root and git_root != self.workspace_path:
rules_dir = git_root / rules_dir_name.replace("/", os.sep)
additional_rules = self._discover_rules_in_dir(rules_dir)
for rule in additional_rules:
dir_prefix = rules_dir_name.split("/")[0].strip(".")
self._rules[f"gitroot_{dir_prefix}:{rule.name}"] = rule
# 6. Load subdirectory rules (walk up from cwd to workspace/git root)
stop_at = git_root if git_root else self.workspace_path
current = Path.cwd()
while current != stop_at and current != current.parent:
# Check .praison/rules/
subdir_rules_dir = current / self.RULES_DIR_NAME.replace("/", os.sep)
subdir_rules = self._discover_rules_in_dir(subdir_rules_dir)
for rule in subdir_rules:
rule.priority = rule.priority + 100 # Higher priority for closer dirs
self._rules[f"subdir:{rule.name}"] = rule
# Check additional rules directories in subdirs
for rules_dir_name in self.ADDITIONAL_RULES_DIRS:
subdir_rules_dir = current / rules_dir_name.replace("/", os.sep)
subdir_rules = self._discover_rules_in_dir(subdir_rules_dir)
for rule in subdir_rules:
rule.priority = rule.priority + 150 # Even higher for subdir additional
dir_prefix = rules_dir_name.split("/")[0].strip(".")
self._rules[f"subdir_{dir_prefix}:{rule.name}"] = rule
current = current.parent
self._log(f"Loaded {len(self._rules)} rules total")
def reload(self):
"""Reload all rules from disk."""
self._load_all_rules()
def get_all_rules(self) -> List[Rule]:
"""Get all loaded rules, sorted by priority."""
rules = list(self._rules.values())
rules.sort(key=lambda r: r.priority, reverse=True)
return rules
def get_active_rules(self, file_path: Optional[str] = None) -> List[Rule]:
"""
Get rules that are currently active.
Args:
file_path: Optional file path to filter glob-based rules
Returns:
List of active rules sorted by priority
"""
active = []
for rule in self._rules.values():
if rule.activation == "always":
active.append(rule)
elif rule.activation == "glob" and file_path:
if rule.matches_file(file_path):
active.append(rule)
elif rule.activation == "ai_decision":
# AI will decide - include for context
active.append(rule)
# Skip "manual" - only activated via @mention
active.sort(key=lambda r: r.priority, reverse=True)
return active
def get_rules_for_file(self, file_path: str) -> List[Rule]:
"""Get rules that apply to a specific file."""
return self.get_active_rules(file_path=file_path)
def get_rule_by_name(self, name: str) -> Optional[Rule]:
"""Get a rule by name (for manual @mention invocation)."""
# Check all scopes
for scope in ["subdir", "workspace", "global"]:
key = f"{scope}:{name}"
if key in self._rules:
return self._rules[key]
# Also check without scope prefix
for rule in self._rules.values():
if rule.name == name:
return rule
return None
def get_manual_rules(self) -> List[Rule]:
"""Get all rules that require manual activation."""
return [r for r in self._rules.values() if r.activation == "manual"]
def build_rules_context(
self,
file_path: Optional[str] = None,
include_manual: Optional[List[str]] = None,
max_chars: int = 10000
) -> str:
"""
Build context string from active rules for LLM.
Args:
file_path: Optional file path for glob matching
include_manual: List of manual rule names to include
max_chars: Maximum characters in output
Returns:
Formatted rules context string
"""
parts = []
total_chars = 0
# Get active rules
active_rules = self.get_active_rules(file_path=file_path)
# Add manually invoked rules
if include_manual:
for name in include_manual:
rule = self.get_rule_by_name(name)
if rule and rule not in active_rules:
active_rules.append(rule)
# Sort by priority
active_rules.sort(key=lambda r: r.priority, reverse=True)
for rule in active_rules:
if total_chars >= max_chars:
break
rule_text = rule.content.strip()
if not rule_text:
continue
# Add rule with header
if rule.description:
header = f"## {rule.name}: {rule.description}"
else:
header = f"## {rule.name}"
section = f"{header}\n{rule_text}\n"
if total_chars + len(section) <= max_chars:
parts.append(section)
total_chars += len(section)
else:
# Truncate last rule
remaining = max_chars - total_chars
if remaining > 100:
parts.append(section[:remaining] + "\n... (truncated)")
break
return "\n".join(parts)
def create_rule(
self,
name: str,
content: str,
description: str = "",
globs: Optional[List[str]] = None,
activation: Literal["always", "glob", "manual", "ai_decision"] = "always",
priority: int = 0,
scope: Literal["global", "workspace"] = "workspace"
) -> Rule:
"""
Create a new rule file.
Args:
name: Rule name (used as filename)
content: Rule content (markdown)
description: Short description
globs: Glob patterns for activation
activation: Activation mode
priority: Priority (higher = applied first)
scope: Where to save (global or workspace)
Returns:
Created Rule object
"""
# Determine save path
if scope == "global":
rules_dir = self.global_rules_path
else:
rules_dir = self.workspace_path / self.RULES_DIR_NAME.replace("/", os.sep)
rules_dir.mkdir(parents=True, exist_ok=True)
file_path = rules_dir / f"{name}.md"
# Build frontmatter
frontmatter_lines = ["---"]
if description:
frontmatter_lines.append(f'description: "{description}"')
if globs:
globs_str = ", ".join(f'"{g}"' for g in globs)
frontmatter_lines.append(f"globs: [{globs_str}]")
frontmatter_lines.append(f"activation: {activation}")
if priority != 0:
frontmatter_lines.append(f"priority: {priority}")
frontmatter_lines.append("---")
frontmatter_lines.append("")
# Write file
full_content = "\n".join(frontmatter_lines) + content
file_path.write_text(full_content, encoding="utf-8")
# Create and register rule
rule = Rule(
name=name,
content=content,
description=description,
globs=globs or [],
activation=activation,
priority=priority,
file_path=str(file_path)
)
self._rules[f"{scope}:{name}"] = rule
self._log(f"Created rule '{name}' at {file_path}")
return rule
def delete_rule(self, name: str, scope: Optional[str] = None) -> bool:
"""
Delete a rule file.
Args:
name: Rule name
scope: Scope to delete from (global, workspace, or None for any)
Returns:
True if deleted, False if not found
"""
# Find the rule
rule = None
key = None
if scope:
key = f"{scope}:{name}"
rule = self._rules.get(key)
else:
for s in ["subdir", "workspace", "global"]:
key = f"{s}:{name}"
if key in self._rules:
rule = self._rules[key]
break
if not rule or not rule.file_path:
return False
# Delete file
try:
Path(rule.file_path).unlink()
del self._rules[key]
self._log(f"Deleted rule '{name}'")
return True
except Exception as e:
self._log(f"Error deleting rule '{name}': {e}", logging.ERROR)
return False
def evaluate_ai_decision(
self,
rule: Rule,
context: str,
llm_func: Optional[Callable[[str], str]] = None
) -> bool:
"""
Evaluate if an ai_decision rule should be applied.
Uses the rule's description to determine relevance to current context.
Args:
rule: Rule with activation="ai_decision"
context: Current conversation/task context
llm_func: Optional LLM function for evaluation (if None, returns True)
Returns:
True if rule should be applied, False otherwise
"""
if rule.activation != "ai_decision":
return rule.activation == "always"
if not llm_func:
# Without LLM, include ai_decision rules by default
return True
# Use LLM to decide
prompt = f"""Determine if the following rule should be applied to the current context.
Rule Name: {rule.name}
Rule Description: {rule.description}
Current Context:
{context[:1000]}
Should this rule be applied? Answer only "yes" or "no"."""
try:
response = llm_func(prompt).strip().lower()
return response.startswith("yes")
except Exception as e:
self._log(f"Error evaluating ai_decision for rule '{rule.name}': {e}", logging.WARNING)
return True # Default to including the rule
def get_stats(self) -> Dict[str, Any]:
"""Get rules statistics."""
rules = list(self._rules.values())
# Count rules by source
source_counts = {}
for key in self._rules.keys():
source = key.split(":")[0]
source_counts[source] = source_counts.get(source, 0) + 1
return {
"total_rules": len(rules),
"always_rules": len([r for r in rules if r.activation == "always"]),
"glob_rules": len([r for r in rules if r.activation == "glob"]),
"manual_rules": len([r for r in rules if r.activation == "manual"]),
"ai_decision_rules": len([r for r in rules if r.activation == "ai_decision"]),
"global_rules": source_counts.get("global", 0),
"root_rules": source_counts.get("root", 0),
"gitroot_rules": source_counts.get("gitroot", 0),
"workspace_rules": source_counts.get("workspace", 0),
"subdir_rules": source_counts.get("subdir", 0),
"claude_rules": source_counts.get("claude", 0),
"windsurf_rules": source_counts.get("windsurf", 0),
"cursor_rules": source_counts.get("cursor", 0),
"sources": source_counts
}
def create_rules_manager(
workspace_path: Optional[str] = None,
**kwargs
) -> RulesManager:
"""
Create a RulesManager instance.
Args:
workspace_path: Path to workspace
**kwargs: Additional configuration
Returns:
RulesManager instance
"""
return RulesManager(workspace_path=workspace_path, **kwargs)
@@ -0,0 +1,514 @@
"""
Workflows System for PraisonAI Agents.
Provides reusable multi-step workflows similar to Windsurf's Workflows,
allowing users to define and execute complex task sequences.
Features:
- YAML/Markdown workflow definitions
- Auto-discovery from .praison/workflows/
- Step-by-step execution with context passing
- Variable substitution
- Conditional steps
- Lazy loading for performance
Storage Structure:
.praison/workflows/
├── deploy.md # Deployment workflow
├── test.md # Testing workflow
└── review.md # Code review workflow
"""
import os
import re
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Callable, Literal
from dataclasses import dataclass, field
logger = logging.getLogger(__name__)
@dataclass
class WorkflowStep:
"""A single step in a workflow."""
name: str
description: str = ""
action: str = "" # The action/prompt to execute
condition: Optional[str] = None # Optional condition for execution
on_error: Literal["stop", "continue", "retry"] = "stop"
max_retries: int = 1
def to_dict(self) -> Dict[str, Any]:
return {
"name": self.name,
"description": self.description,
"action": self.action,
"condition": self.condition,
"on_error": self.on_error,
"max_retries": self.max_retries
}
@dataclass
class Workflow:
"""A complete workflow with multiple steps."""
name: str
description: str = ""
steps: List[WorkflowStep] = field(default_factory=list)
variables: Dict[str, Any] = field(default_factory=dict)
file_path: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
return {
"name": self.name,
"description": self.description,
"steps": [s.to_dict() for s in self.steps],
"variables": self.variables,
"file_path": self.file_path
}
class WorkflowManager:
"""
Manages workflow discovery, loading, and execution.
Workflows are defined in markdown files with YAML frontmatter
and step definitions.
Example workflow file (.praison/workflows/deploy.md):
```markdown
---
name: Deploy to Production
description: Deploy the application to production
variables:
environment: production
branch: main
---
## Step 1: Run Tests
Run all tests to ensure code quality.
```action
Run the test suite with pytest
```
## Step 2: Build
Build the application for production.
```action
Build the application with production settings
```
## Step 3: Deploy
Deploy to the production server.
```action
Deploy to {{environment}} from {{branch}}
```
```
Example usage:
```python
manager = WorkflowManager(workspace_path="/path/to/project")
# List available workflows
workflows = manager.list_workflows()
# Get a specific workflow
deploy = manager.get_workflow("deploy")
# Execute a workflow
results = manager.execute(
"deploy",
executor=lambda prompt: agent.chat(prompt),
variables={"branch": "release-1.0"}
)
```
"""
WORKFLOWS_DIR = ".praison/workflows"
SUPPORTED_EXTENSIONS = [".md", ".yaml", ".yml"]
def __init__(
self,
workspace_path: Optional[str] = None,
verbose: int = 0
):
"""
Initialize WorkflowManager.
Args:
workspace_path: Path to workspace/project root
verbose: Verbosity level
"""
self.workspace_path = Path(workspace_path) if workspace_path else Path.cwd()
self.verbose = verbose
self._workflows: Dict[str, Workflow] = {}
self._loaded = False
def _log(self, msg: str, level: int = logging.INFO):
"""Log message if verbose."""
if self.verbose >= 1:
logger.log(level, msg)
def _ensure_loaded(self):
"""Lazy load workflows on first access."""
if not self._loaded:
self._discover_workflows()
self._loaded = True
def _parse_frontmatter(self, content: str) -> tuple[Dict[str, Any], str]:
"""Parse YAML frontmatter from markdown content."""
frontmatter = {}
body = content
if content.startswith("---"):
parts = content.split("---", 2)
if len(parts) >= 3:
yaml_content = parts[1].strip()
body = parts[2].strip()
# Simple YAML parsing
for line in yaml_content.split("\n"):
line = line.strip()
if ":" in line and not line.startswith("#"):
key, value = line.split(":", 1)
key = key.strip()
value = value.strip()
# Handle basic types
if value.lower() in ("true", "false"):
value = value.lower() == "true"
elif value.isdigit():
value = int(value)
elif value.startswith('"') and value.endswith('"'):
value = value[1:-1]
frontmatter[key] = value
return frontmatter, body
def _parse_steps(self, body: str) -> List[WorkflowStep]:
"""Parse workflow steps from markdown body."""
steps = []
# Split by ## headers
step_pattern = re.compile(r'^##\s+(?:Step\s+\d+[:\s]*)?(.+)$', re.MULTILINE)
action_pattern = re.compile(r'```action\s*\n(.*?)\n```', re.DOTALL)
condition_pattern = re.compile(r'```condition\s*\n(.*?)\n```', re.DOTALL)
# Find all step headers
matches = list(step_pattern.finditer(body))
for i, match in enumerate(matches):
step_name = match.group(1).strip()
# Get content until next step or end
start = match.end()
end = matches[i + 1].start() if i + 1 < len(matches) else len(body)
step_content = body[start:end].strip()
# Extract action
action_match = action_pattern.search(step_content)
action = action_match.group(1).strip() if action_match else ""
# Extract condition
condition_match = condition_pattern.search(step_content)
condition = condition_match.group(1).strip() if condition_match else None
# Get description (text before action block)
description = step_content
if action_match:
description = step_content[:action_match.start()].strip()
# Clean up description (remove code blocks)
description = re.sub(r'```.*?```', '', description, flags=re.DOTALL).strip()
steps.append(WorkflowStep(
name=step_name,
description=description,
action=action,
condition=condition
))
return steps
def _load_workflow(self, file_path: Path) -> Optional[Workflow]:
"""Load a workflow from a file."""
try:
content = file_path.read_text(encoding="utf-8")
frontmatter, body = self._parse_frontmatter(content)
name = frontmatter.get("name", file_path.stem)
description = frontmatter.get("description", "")
variables = frontmatter.get("variables", {})
# Parse variables if it's a string (simple YAML)
if isinstance(variables, str):
variables = {}
steps = self._parse_steps(body)
workflow = Workflow(
name=name,
description=description,
steps=steps,
variables=variables,
file_path=str(file_path)
)
self._log(f"Loaded workflow '{name}' with {len(steps)} steps")
return workflow
except Exception as e:
self._log(f"Error loading workflow {file_path}: {e}", logging.WARNING)
return None
def _discover_workflows(self):
"""Discover all workflows in the workspace."""
self._workflows = {}
workflows_dir = self.workspace_path / self.WORKFLOWS_DIR.replace("/", os.sep)
if not workflows_dir.exists():
return
for ext in self.SUPPORTED_EXTENSIONS:
for file_path in workflows_dir.glob(f"*{ext}"):
workflow = self._load_workflow(file_path)
if workflow:
self._workflows[workflow.name.lower()] = workflow
self._log(f"Discovered {len(self._workflows)} workflows")
def reload(self):
"""Reload all workflows from disk."""
self._loaded = False
self._ensure_loaded()
def list_workflows(self) -> List[Workflow]:
"""List all available workflows."""
self._ensure_loaded()
return list(self._workflows.values())
def get_workflow(self, name: str) -> Optional[Workflow]:
"""Get a workflow by name."""
self._ensure_loaded()
return self._workflows.get(name.lower())
def _substitute_variables(
self,
text: str,
variables: Dict[str, Any]
) -> str:
"""Substitute {{variable}} placeholders in text."""
def replace(match):
var_name = match.group(1).strip()
return str(variables.get(var_name, match.group(0)))
return re.sub(r'\{\{(\w+)\}\}', replace, text)
def execute(
self,
workflow_name: str,
executor: Callable[[str], str],
variables: Optional[Dict[str, Any]] = None,
on_step: Optional[Callable[[WorkflowStep, int], None]] = None,
on_result: Optional[Callable[[WorkflowStep, str], None]] = None
) -> Dict[str, Any]:
"""
Execute a workflow.
Args:
workflow_name: Name of the workflow to execute
executor: Function to execute each step (e.g., agent.chat)
variables: Variables to substitute in steps
on_step: Callback before each step (step, index)
on_result: Callback after each step (step, result)
Returns:
Execution results with step outputs and status
"""
workflow = self.get_workflow(workflow_name)
if not workflow:
return {
"success": False,
"error": f"Workflow '{workflow_name}' not found",
"results": []
}
# Merge variables
all_variables = {**workflow.variables}
if variables:
all_variables.update(variables)
results = []
success = True
for i, step in enumerate(workflow.steps):
# Check condition
if step.condition:
condition = self._substitute_variables(step.condition, all_variables)
# Simple condition evaluation (could be enhanced)
if condition.lower() in ("false", "no", "skip", "0"):
results.append({
"step": step.name,
"status": "skipped",
"output": None
})
continue
# Callback before step
if on_step:
on_step(step, i)
# Substitute variables in action
action = self._substitute_variables(step.action, all_variables)
# Execute step
retries = 0
step_success = False
output = None
error = None
while retries <= step.max_retries and not step_success:
try:
output = executor(action)
step_success = True
except Exception as e:
error = str(e)
retries += 1
if retries <= step.max_retries:
self._log(f"Step '{step.name}' failed, retrying ({retries}/{step.max_retries})")
# Callback after step
if on_result and output:
on_result(step, output)
results.append({
"step": step.name,
"status": "success" if step_success else "failed",
"output": output,
"error": error
})
# Handle failure
if not step_success:
if step.on_error == "stop":
success = False
break
elif step.on_error == "continue":
continue
return {
"success": success,
"workflow": workflow.name,
"results": results
}
def create_workflow(
self,
name: str,
description: str = "",
steps: Optional[List[Dict[str, str]]] = None,
variables: Optional[Dict[str, Any]] = None
) -> Workflow:
"""
Create a new workflow file.
Args:
name: Workflow name
description: Workflow description
steps: List of step dicts with name, description, action
variables: Default variables
Returns:
Created Workflow object
"""
workflows_dir = self.workspace_path / self.WORKFLOWS_DIR.replace("/", os.sep)
workflows_dir.mkdir(parents=True, exist_ok=True)
file_path = workflows_dir / f"{name.lower().replace(' ', '_')}.md"
# Build content
lines = ["---"]
lines.append(f"name: {name}")
if description:
lines.append(f"description: {description}")
if variables:
lines.append("variables:")
for k, v in variables.items():
lines.append(f" {k}: {v}")
lines.append("---")
lines.append("")
# Add steps
workflow_steps = []
if steps:
for i, step in enumerate(steps, 1):
step_name = step.get("name", f"Step {i}")
step_desc = step.get("description", "")
step_action = step.get("action", "")
lines.append(f"## Step {i}: {step_name}")
if step_desc:
lines.append(step_desc)
lines.append("")
lines.append("```action")
lines.append(step_action)
lines.append("```")
lines.append("")
workflow_steps.append(WorkflowStep(
name=step_name,
description=step_desc,
action=step_action
))
# Write file
file_path.write_text("\n".join(lines), encoding="utf-8")
# Create and register workflow
workflow = Workflow(
name=name,
description=description,
steps=workflow_steps,
variables=variables or {},
file_path=str(file_path)
)
self._workflows[name.lower()] = workflow
self._log(f"Created workflow '{name}' at {file_path}")
return workflow
def get_stats(self) -> Dict[str, Any]:
"""Get workflow statistics."""
self._ensure_loaded()
total_steps = sum(len(w.steps) for w in self._workflows.values())
return {
"total_workflows": len(self._workflows),
"total_steps": total_steps,
"workflows": [w.name for w in self._workflows.values()]
}
def create_workflow_manager(
workspace_path: Optional[str] = None,
**kwargs
) -> WorkflowManager:
"""
Create a WorkflowManager instance.
Args:
workspace_path: Path to workspace
**kwargs: Additional configuration
Returns:
WorkflowManager instance
"""
return WorkflowManager(workspace_path=workspace_path, **kwargs)
@@ -0,0 +1,104 @@
"""
Planning Module for PraisonAI Agents.
Provides Planning Mode functionality similar to:
- Cursor Plan Mode
- Windsurf Planning Mode
- Claude Code Plan Mode
- Gemini CLI Plan Mode
- Codex CLI /plan command
Features:
- PlanningAgent for creating implementation plans
- Plan and PlanStep dataclasses for plan structure
- TodoList for tracking progress
- PlanStorage for persistence
- ApprovalCallback for plan approval flow
- Read-only mode for safe research
Usage:
from praisonaiagents import Agent, Task, PraisonAIAgents
agents = PraisonAIAgents(
agents=[agent1, agent2],
tasks=[task1, task2],
planning=True, # Enable planning mode
planning_llm="gpt-4o-mini" # Optional: custom LLM
)
result = agents.start()
"""
from .plan import Plan, PlanStep
from .todo import TodoList, TodoItem
from .storage import PlanStorage
from .planner import PlanningAgent
from .approval import ApprovalCallback
# Read-only tools that are allowed in plan mode
READ_ONLY_TOOLS = [
"read_file",
"list_directory",
"search_codebase",
"search_files",
"grep_search",
"find_files",
"web_search",
"get_file_content",
"list_files",
"read_document",
"search_web",
"fetch_url",
"get_context",
]
# Restricted tools that are blocked in plan mode
RESTRICTED_TOOLS = [
"write_file",
"create_file",
"delete_file",
"execute_command",
"run_command",
"shell_command",
"modify_file",
"edit_file",
"remove_file",
"move_file",
"copy_file",
"mkdir",
"rmdir",
"git_commit",
"git_push",
"npm_install",
"pip_install",
]
# Research tools that are safe for planning research
RESEARCH_TOOLS = [
"web_search",
"search_web",
"duckduckgo_search",
"tavily_search",
"brave_search",
"google_search",
"read_url",
"fetch_url",
"read_file",
"list_directory",
"search_codebase",
"grep_search",
"find_files",
]
__all__ = [
"Plan",
"PlanStep",
"TodoList",
"TodoItem",
"PlanStorage",
"PlanningAgent",
"ApprovalCallback",
"READ_ONLY_TOOLS",
"RESTRICTED_TOOLS",
"RESEARCH_TOOLS",
]
@@ -0,0 +1,274 @@
"""
ApprovalCallback for plan approval flow.
Provides approval functionality similar to:
- Cursor Plan Mode approval
- Claude Code plan approval
- Codex CLI /approvals
Features:
- Auto-approve option
- Custom approval functions
- Async support
- Plan modification during approval
"""
import asyncio
import inspect
import logging
from typing import Callable, Optional, Union, TYPE_CHECKING
if TYPE_CHECKING:
from .plan import Plan
logger = logging.getLogger(__name__)
class ApprovalCallback:
"""
Callback for plan approval flow.
Handles the approval process for plans before execution.
Can be configured for auto-approval or custom approval logic.
Attributes:
auto_approve: Whether to automatically approve all plans
approve_fn: Custom approval function
on_reject: Callback when plan is rejected
"""
def __init__(
self,
auto_approve: bool = False,
approve_fn: Optional[Callable[['Plan'], Union[bool, 'asyncio.Future']]] = None,
on_reject: Optional[Callable[['Plan'], None]] = None
):
"""
Initialize ApprovalCallback.
Args:
auto_approve: If True, automatically approve all plans
approve_fn: Custom function to determine approval
Should return True to approve, False to reject
Can be sync or async
on_reject: Optional callback when plan is rejected
"""
self.auto_approve = auto_approve
self.approve_fn = approve_fn
self.on_reject = on_reject
def __call__(self, plan: 'Plan') -> bool:
"""
Synchronously check if plan is approved.
Args:
plan: Plan to check for approval
Returns:
True if approved, False if rejected
"""
if self.auto_approve:
logger.debug(f"Auto-approving plan: {plan.name}")
plan.approve()
return True
if self.approve_fn is not None:
# Check if function is async
if inspect.iscoroutinefunction(self.approve_fn):
# Run async function synchronously
loop = asyncio.new_event_loop()
try:
result = loop.run_until_complete(self.approve_fn(plan))
finally:
loop.close()
else:
result = self.approve_fn(plan)
if result:
plan.approve()
return True
else:
if self.on_reject:
self.on_reject(plan)
return False
# Default: require explicit approval (return False)
logger.debug(f"Plan requires explicit approval: {plan.name}")
return False
async def async_call(self, plan: 'Plan') -> bool:
"""
Asynchronously check if plan is approved.
Args:
plan: Plan to check for approval
Returns:
True if approved, False if rejected
"""
if self.auto_approve:
logger.debug(f"Auto-approving plan: {plan.name}")
plan.approve()
return True
if self.approve_fn is not None:
# Check if function is async
if inspect.iscoroutinefunction(self.approve_fn):
result = await self.approve_fn(plan)
else:
result = self.approve_fn(plan)
if result:
plan.approve()
return True
else:
if self.on_reject:
self.on_reject(plan)
return False
# Default: require explicit approval (return False)
logger.debug(f"Plan requires explicit approval: {plan.name}")
return False
@staticmethod
def console_approval(plan: 'Plan') -> bool:
"""
Interactive console approval.
Displays the plan and prompts for approval.
Args:
plan: Plan to approve
Returns:
True if user approves, False otherwise
"""
print("\n" + "=" * 60)
print("PLAN APPROVAL REQUIRED")
print("=" * 60)
print(plan.to_markdown())
print("=" * 60)
while True:
response = input("\nApprove this plan? [y/n/e(dit)]: ").strip().lower()
if response in ['y', 'yes']:
return True
elif response in ['n', 'no']:
return False
elif response in ['e', 'edit']:
print("\nPlan editing not implemented in console mode.")
print("Please modify the plan programmatically and re-submit.")
continue
else:
print("Please enter 'y' for yes, 'n' for no, or 'e' to edit.")
@staticmethod
def always_approve(plan: 'Plan') -> bool:
"""Always approve plans."""
return True
@staticmethod
def always_reject(plan: 'Plan') -> bool:
"""Always reject plans."""
return False
@staticmethod
def approve_if_small(plan: 'Plan', max_steps: int = 5) -> bool:
"""
Approve plans with few steps automatically.
Args:
plan: Plan to check
max_steps: Maximum number of steps to auto-approve
Returns:
True if plan has <= max_steps
"""
return len(plan.steps) <= max_steps
@staticmethod
def approve_if_no_dangerous_tools(plan: 'Plan') -> bool:
"""
Approve plans that don't use dangerous tools.
Args:
plan: Plan to check
Returns:
True if no dangerous tools are used
"""
from . import RESTRICTED_TOOLS
for step in plan.steps:
for tool in step.tools:
if tool in RESTRICTED_TOOLS:
logger.warning(f"Plan uses restricted tool: {tool}")
return False
return True
class InteractiveApproval:
"""
Interactive approval handler with rich UI support.
Provides a more sophisticated approval interface
with plan editing capabilities.
"""
def __init__(self, use_rich: bool = True):
"""
Initialize InteractiveApproval.
Args:
use_rich: Whether to use Rich library for display
"""
self.use_rich = use_rich
self._rich_available = False
if use_rich:
try:
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Confirm
self._rich_available = True
self._console = Console()
except ImportError:
self._rich_available = False
def __call__(self, plan: 'Plan') -> bool:
"""
Display plan and get approval.
Args:
plan: Plan to approve
Returns:
True if approved
"""
if self._rich_available:
return self._rich_approval(plan)
else:
return ApprovalCallback.console_approval(plan)
def _rich_approval(self, plan: 'Plan') -> bool:
"""Rich library approval display."""
from rich.panel import Panel
from rich.prompt import Confirm
from rich.markdown import Markdown
# Display plan
md = Markdown(plan.to_markdown())
self._console.print(Panel(md, title="[bold blue]Plan Approval Required[/bold blue]"))
# Get approval
approved = Confirm.ask("Approve this plan?")
return approved
async def async_call(self, plan: 'Plan') -> bool:
"""Async version of approval."""
# For now, just call sync version
# Could be extended for async UI frameworks
return self(plan)
@@ -0,0 +1,364 @@
"""
Plan and PlanStep dataclasses for Planning Mode.
Provides structured representation of implementation plans
with support for:
- Step dependencies
- Status tracking
- Markdown serialization
- Progress calculation
"""
import uuid
import re
import yaml
import logging
from datetime import datetime
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any, Literal
logger = logging.getLogger(__name__)
@dataclass
class PlanStep:
"""
A single step in an implementation plan.
Attributes:
id: Unique identifier for the step
description: What this step does
agent: Name of the agent to execute this step
tools: List of tools needed for this step
dependencies: List of step IDs that must complete first
status: Current status (pending, in_progress, completed, skipped)
estimated_tokens: Estimated token usage for this step
"""
description: str
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
agent: Optional[str] = None
tools: List[str] = field(default_factory=list)
dependencies: List[str] = field(default_factory=list)
status: Literal["pending", "in_progress", "completed", "skipped"] = "pending"
estimated_tokens: int = 0
def to_dict(self) -> Dict[str, Any]:
"""Convert PlanStep to dictionary."""
return {
"id": self.id,
"description": self.description,
"agent": self.agent,
"tools": self.tools,
"dependencies": self.dependencies,
"status": self.status,
"estimated_tokens": self.estimated_tokens
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> 'PlanStep':
"""Create PlanStep from dictionary."""
return cls(
id=data.get("id", str(uuid.uuid4())[:8]),
description=data.get("description", ""),
agent=data.get("agent"),
tools=data.get("tools", []),
dependencies=data.get("dependencies", []),
status=data.get("status", "pending"),
estimated_tokens=data.get("estimated_tokens", 0)
)
def mark_complete(self) -> None:
"""Mark this step as completed."""
self.status = "completed"
def mark_in_progress(self) -> None:
"""Mark this step as in progress."""
self.status = "in_progress"
def mark_skipped(self) -> None:
"""Mark this step as skipped."""
self.status = "skipped"
def is_ready(self, completed_steps: List[str]) -> bool:
"""
Check if this step is ready to execute.
Args:
completed_steps: List of completed step IDs
Returns:
True if all dependencies are satisfied
"""
if self.status != "pending":
return False
return all(dep in completed_steps for dep in self.dependencies)
@dataclass
class Plan:
"""
An implementation plan with multiple steps.
Attributes:
id: Unique identifier for the plan
name: Human-readable name
description: Detailed description of what the plan accomplishes
steps: List of PlanStep objects
status: Current status (draft, approved, executing, completed)
created_at: When the plan was created
approved_at: When the plan was approved (if applicable)
file_path: Path where the plan is stored (if saved)
"""
name: str
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
description: str = ""
steps: List[PlanStep] = field(default_factory=list)
status: Literal["draft", "approved", "executing", "completed", "cancelled"] = "draft"
created_at: datetime = field(default_factory=datetime.now)
approved_at: Optional[datetime] = None
file_path: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
"""Convert Plan to dictionary."""
return {
"id": self.id,
"name": self.name,
"description": self.description,
"steps": [step.to_dict() for step in self.steps],
"status": self.status,
"created_at": self.created_at.isoformat() if self.created_at else None,
"approved_at": self.approved_at.isoformat() if self.approved_at else None,
"file_path": self.file_path
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> 'Plan':
"""Create Plan from dictionary."""
steps = [
PlanStep.from_dict(s) if isinstance(s, dict) else s
for s in data.get("steps", [])
]
created_at = data.get("created_at")
if isinstance(created_at, str):
created_at = datetime.fromisoformat(created_at)
elif created_at is None:
created_at = datetime.now()
approved_at = data.get("approved_at")
if isinstance(approved_at, str):
approved_at = datetime.fromisoformat(approved_at)
return cls(
id=data.get("id", str(uuid.uuid4())[:8]),
name=data.get("name", "Untitled Plan"),
description=data.get("description", ""),
steps=steps,
status=data.get("status", "draft"),
created_at=created_at,
approved_at=approved_at,
file_path=data.get("file_path")
)
def to_markdown(self) -> str:
"""
Convert Plan to markdown format.
Returns:
Markdown string representation of the plan
"""
# YAML frontmatter
frontmatter = {
"id": self.id,
"name": self.name,
"status": self.status,
"created_at": self.created_at.isoformat() if self.created_at else None,
"approved_at": self.approved_at.isoformat() if self.approved_at else None
}
lines = [
"---",
yaml.dump(frontmatter, default_flow_style=False).strip(),
"---",
"",
f"# {self.name}",
"",
]
if self.description:
lines.extend([self.description, ""])
lines.extend(["## Steps", ""])
status_icons = {
"pending": "⬜",
"in_progress": "🔄",
"completed": "✅",
"skipped": "⏭️"
}
for i, step in enumerate(self.steps, 1):
icon = status_icons.get(step.status, "⬜")
lines.append(f"{i}. {icon} {step.description}")
if step.agent:
lines.append(f" - Agent: {step.agent}")
if step.tools:
lines.append(f" - Tools: {', '.join(step.tools)}")
if step.dependencies:
lines.append(f" - Depends on: {', '.join(step.dependencies)}")
lines.append("")
return "\n".join(lines)
@classmethod
def from_markdown(cls, markdown: str) -> 'Plan':
"""
Create Plan from markdown format.
Args:
markdown: Markdown string with YAML frontmatter
Returns:
Plan instance
"""
# Extract YAML frontmatter
frontmatter_match = re.match(r'^---\n(.*?)\n---', markdown, re.DOTALL)
data = {}
if frontmatter_match:
try:
data = yaml.safe_load(frontmatter_match.group(1))
except yaml.YAMLError:
logger.warning("Failed to parse YAML frontmatter")
# Extract steps from markdown
steps = []
# Pattern to capture status icon and description
step_pattern = r'^\d+\.\s*([⬜🔄✅⏭️]?)\s*(.+?)$'
# Map icons to status
icon_to_status = {
"⬜": "pending",
"🔄": "in_progress",
"✅": "completed",
"⏭️": "skipped",
"": "pending"
}
for match in re.finditer(step_pattern, markdown, re.MULTILINE):
icon = match.group(1).strip()
description = match.group(2).strip()
status = icon_to_status.get(icon, "pending")
steps.append(PlanStep(description=description, status=status))
return cls(
id=data.get("id", str(uuid.uuid4())[:8]),
name=data.get("name", "Untitled Plan"),
description=data.get("description", ""),
steps=steps,
status=data.get("status", "draft")
)
def add_step(self, step: PlanStep) -> None:
"""Add a step to the plan."""
self.steps.append(step)
def remove_step(self, step_id: str) -> bool:
"""
Remove a step from the plan.
Args:
step_id: ID of the step to remove
Returns:
True if step was removed, False if not found
"""
for i, step in enumerate(self.steps):
if step.id == step_id:
self.steps.pop(i)
return True
return False
def get_step(self, step_id: str) -> Optional[PlanStep]:
"""
Get a step by ID.
Args:
step_id: ID of the step to find
Returns:
PlanStep if found, None otherwise
"""
for step in self.steps:
if step.id == step_id:
return step
return None
def approve(self) -> None:
"""Approve the plan for execution."""
self.status = "approved"
self.approved_at = datetime.now()
def start_execution(self) -> None:
"""Mark the plan as executing."""
self.status = "executing"
def complete(self) -> None:
"""Mark the plan as completed."""
self.status = "completed"
def cancel(self) -> None:
"""Cancel the plan."""
self.status = "cancelled"
@property
def progress(self) -> float:
"""
Calculate plan progress.
Returns:
Float between 0.0 and 1.0 representing completion percentage
"""
if not self.steps:
return 1.0
completed = sum(1 for s in self.steps if s.status == "completed")
return completed / len(self.steps)
@property
def is_complete(self) -> bool:
"""Check if all steps are completed."""
return all(s.status == "completed" for s in self.steps)
@property
def completed_step_ids(self) -> List[str]:
"""Get list of completed step IDs."""
return [s.id for s in self.steps if s.status == "completed"]
def get_next_steps(self) -> List[PlanStep]:
"""
Get steps that are ready to execute.
Returns:
List of steps with satisfied dependencies
"""
completed = self.completed_step_ids
return [s for s in self.steps if s.is_ready(completed)]
def update_step_status(self, step_id: str, status: str) -> bool:
"""
Update the status of a step.
Args:
step_id: ID of the step to update
status: New status
Returns:
True if step was updated, False if not found
"""
step = self.get_step(step_id)
if step:
step.status = status
return True
return False
@@ -0,0 +1,525 @@
"""
PlanningAgent for creating implementation plans.
Provides planning functionality similar to:
- CrewAI AgentPlanner
- Cursor Plan Mode
- Claude Code Plan Mode
Features:
- Analyze requests and create detailed plans
- Read-only mode for safe research
- Plan refinement based on feedback
- Context analysis
"""
import json
import logging
from typing import List, Optional, Dict, Any, TYPE_CHECKING
from .plan import Plan, PlanStep
# Define READ_ONLY_TOOLS locally to avoid circular import
READ_ONLY_TOOLS = [
"read_file",
"list_directory",
"search_codebase",
"search_files",
"grep_search",
"find_files",
"web_search",
"get_file_content",
"list_files",
"read_document",
"search_web",
"fetch_url",
"get_context",
]
if TYPE_CHECKING:
from ..agent.agent import Agent
from ..task.task import Task
logger = logging.getLogger(__name__)
class PlanningAgent:
"""
Agent specialized in creating implementation plans.
The PlanningAgent analyzes requests, researches the codebase,
and creates detailed step-by-step plans before execution.
Attributes:
llm_model: LLM model to use for planning
read_only: Whether to restrict to read-only tools
allowed_tools: List of tools available in planning mode
verbose: Verbosity level
"""
PLANNING_PROMPT = '''You are an expert planning specialist. Your task is to analyze the request and create a detailed, actionable implementation plan.
## Request
{request}
## Available Agents
{agents_info}
## Available Tasks (if any)
{tasks_info}
## Context
{context}
## Instructions
1. Analyze the request thoroughly
2. Break it down into clear, actionable steps
3. Assign appropriate agents to each step
4. Identify dependencies between steps
5. Consider potential risks and edge cases
## Output Format
Respond with a JSON object in this exact format:
{{
"name": "Brief plan name",
"description": "Detailed description of what this plan accomplishes",
"steps": [
{{
"description": "Clear description of what this step does",
"agent": "Name of the agent to execute this step (or null)",
"tools": ["list", "of", "tools", "needed"],
"dependencies": ["list of step IDs this depends on"]
}}
]
}}
Important:
- Each step should be atomic and clearly defined
- Dependencies should reference step indices (e.g., "step_0", "step_1")
- Be specific about what each step accomplishes
- Consider the order of operations carefully'''
REFINEMENT_PROMPT = '''You are refining an existing implementation plan based on feedback.
## Original Plan
{original_plan}
## Feedback
{feedback}
## Instructions
1. Review the original plan
2. Apply the feedback to improve the plan
3. Maintain the same JSON format
4. Keep steps that are still valid
5. Add, modify, or remove steps as needed
## Output Format
Respond with a JSON object in the same format as the original plan:
{{
"name": "Updated plan name",
"description": "Updated description",
"steps": [...]
}}'''
ANALYSIS_PROMPT = '''Analyze the following context and provide insights for planning:
## Context
{context}
## Instructions
1. Identify the key components and their relationships
2. Note any patterns or conventions
3. Highlight potential challenges
4. Suggest best practices to follow
Provide a concise analysis that will help in creating an implementation plan.'''
def __init__(
self,
llm: str = "gpt-4o-mini",
read_only: bool = True,
verbose: int = 0,
tools: Optional[List] = None,
reasoning: bool = False,
max_reasoning_steps: int = 5
):
"""
Initialize PlanningAgent.
Args:
llm: LLM model to use for planning
read_only: Whether to restrict to read-only tools
verbose: Verbosity level
tools: Optional list of tools for research during planning
reasoning: Whether to enable reasoning mode with chain-of-thought
max_reasoning_steps: Maximum reasoning iterations (default 5)
"""
self.llm_model = llm
self.read_only = read_only
self.verbose = verbose
self.tools = tools or []
self.reasoning = reasoning
self.max_reasoning_steps = max_reasoning_steps
self.allowed_tools = READ_ONLY_TOOLS.copy() if read_only else []
self._llm = None
self._agent = None
def _get_llm(self):
"""Lazy load LLM instance."""
if self._llm is None:
from ..llm import LLM
self._llm = LLM(model=self.llm_model)
return self._llm
def _get_agent(self):
"""
Lazy load Agent instance for tool-enabled planning.
Returns:
Agent instance if tools are provided, None otherwise
"""
if not self.tools:
return None
if self._agent is None:
from ..agent.agent import Agent
self._agent = Agent(
name="Planning Research Agent",
role="Research and Planning Specialist",
goal="Research information and create detailed implementation plans",
llm=self.llm_model,
tools=self.tools,
verbose=self.verbose > 0
)
return self._agent
def _format_agents_info(self, agents: List['Agent']) -> str:
"""
Format agent information for the prompt.
Args:
agents: List of Agent instances
Returns:
Formatted string describing agents
"""
if not agents:
return "No agents specified"
lines = []
for agent in agents:
name = getattr(agent, 'name', 'Unnamed')
role = getattr(agent, 'role', 'General')
goal = getattr(agent, 'goal', '')
line = f"- **{name}** ({role})"
if goal:
line += f": {goal}"
lines.append(line)
return "\n".join(lines)
def _format_tasks_info(self, tasks: Optional[List['Task']]) -> str:
"""
Format task information for the prompt.
Args:
tasks: List of Task instances
Returns:
Formatted string describing tasks
"""
if not tasks:
return "No predefined tasks"
lines = []
for task in tasks:
name = getattr(task, 'name', None) or getattr(task, 'description', 'Unnamed')[:50]
desc = getattr(task, 'description', '')[:100]
lines.append(f"- {name}: {desc}")
return "\n".join(lines)
def _parse_plan_response(self, response: str) -> Plan:
"""
Parse LLM response into a Plan object.
Args:
response: JSON string from LLM
Returns:
Plan instance
"""
try:
# Try to extract JSON from response
# Handle cases where LLM includes extra text
json_start = response.find('{')
json_end = response.rfind('}') + 1
if json_start >= 0 and json_end > json_start:
json_str = response[json_start:json_end]
data = json.loads(json_str)
else:
raise ValueError("No JSON found in response")
# Create steps
steps = []
for i, step_data in enumerate(data.get("steps", [])):
step = PlanStep(
id=f"step_{i}",
description=step_data.get("description", ""),
agent=step_data.get("agent"),
tools=step_data.get("tools", []),
dependencies=step_data.get("dependencies", [])
)
steps.append(step)
return Plan(
name=data.get("name", "Implementation Plan"),
description=data.get("description", ""),
steps=steps
)
except (json.JSONDecodeError, ValueError) as e:
logger.warning(f"Failed to parse plan response: {e}")
# Return a simple plan with the response as a single step
return Plan(
name="Implementation Plan",
description="Auto-generated plan",
steps=[PlanStep(description=response[:500])]
)
async def create_plan(
self,
request: str,
agents: List['Agent'],
tasks: Optional[List['Task']] = None,
context: Optional[str] = None
) -> Plan:
"""
Create an implementation plan for the given request.
Args:
request: The user's request or goal
agents: List of available agents
tasks: Optional list of predefined tasks
context: Optional additional context
Returns:
Plan instance with steps
"""
llm = self._get_llm()
# Format the prompt
prompt = self.PLANNING_PROMPT.format(
request=request,
agents_info=self._format_agents_info(agents),
tasks_info=self._format_tasks_info(tasks),
context=context or "No additional context provided"
)
if self.verbose >= 2:
logger.info(f"Creating plan for: {request[:100]}...")
# Get LLM response
response = await llm.get_response_async(prompt)
# Parse response into Plan
plan = self._parse_plan_response(response)
if self.verbose >= 1:
logger.info(f"Created plan '{plan.name}' with {len(plan.steps)} steps")
return plan
def create_plan_sync(
self,
request: str,
agents: List['Agent'],
tasks: Optional[List['Task']] = None,
context: Optional[str] = None
) -> Plan:
"""
Synchronous version of create_plan.
Args:
request: The user's request or goal
agents: List of available agents
tasks: Optional list of predefined tasks
context: Optional additional context
Returns:
Plan instance with steps
"""
# Format the prompt
prompt = self.PLANNING_PROMPT.format(
request=request,
agents_info=self._format_agents_info(agents),
tasks_info=self._format_tasks_info(tasks),
context=context or "No additional context provided"
)
if self.verbose >= 2:
logger.info(f"Creating plan for: {request[:100]}...")
# Use Agent with tools if available, otherwise use bare LLM
agent = self._get_agent()
if agent is not None:
# Use Agent.chat() for tool-enabled planning
response = agent.chat(prompt)
else:
# Use bare LLM for simple planning
llm = self._get_llm()
response = llm.get_response(prompt)
# Parse response into Plan
plan = self._parse_plan_response(response)
if self.verbose >= 1:
logger.info(f"Created plan '{plan.name}' with {len(plan.steps)} steps")
return plan
async def refine_plan(
self,
plan: Plan,
feedback: str
) -> Plan:
"""
Refine an existing plan based on feedback.
Args:
plan: Original plan to refine
feedback: User feedback for refinement
Returns:
Refined Plan instance
"""
llm = self._get_llm()
# Format the prompt
prompt = self.REFINEMENT_PROMPT.format(
original_plan=json.dumps(plan.to_dict(), indent=2),
feedback=feedback
)
if self.verbose >= 2:
logger.info(f"Refining plan based on feedback: {feedback[:100]}...")
# Get LLM response
response = await llm.get_response_async(prompt)
# Parse response into Plan
refined = self._parse_plan_response(response)
# Preserve original ID
refined.id = plan.id
if self.verbose >= 1:
logger.info(f"Refined plan '{refined.name}' with {len(refined.steps)} steps")
return refined
def refine_plan_sync(
self,
plan: Plan,
feedback: str
) -> Plan:
"""
Synchronous version of refine_plan.
Args:
plan: Original plan to refine
feedback: User feedback for refinement
Returns:
Refined Plan instance
"""
llm = self._get_llm()
# Format the prompt
prompt = self.REFINEMENT_PROMPT.format(
original_plan=json.dumps(plan.to_dict(), indent=2),
feedback=feedback
)
if self.verbose >= 2:
logger.info(f"Refining plan based on feedback: {feedback[:100]}...")
# Get LLM response
response = llm.get_response(prompt)
# Parse response into Plan
refined = self._parse_plan_response(response)
# Preserve original ID
refined.id = plan.id
if self.verbose >= 1:
logger.info(f"Refined plan '{refined.name}' with {len(refined.steps)} steps")
return refined
async def analyze_context(
self,
context: str
) -> str:
"""
Analyze context and provide insights for planning.
Args:
context: Context to analyze (e.g., codebase summary)
Returns:
Analysis string
"""
llm = self._get_llm()
prompt = self.ANALYSIS_PROMPT.format(context=context)
if self.verbose >= 2:
logger.info("Analyzing context...")
response = await llm.get_response_async(prompt)
return response
def analyze_context_sync(
self,
context: str
) -> str:
"""
Synchronous version of analyze_context.
Args:
context: Context to analyze
Returns:
Analysis string
"""
llm = self._get_llm()
prompt = self.ANALYSIS_PROMPT.format(context=context)
if self.verbose >= 2:
logger.info("Analyzing context...")
response = llm.get_response(prompt)
return response
def is_tool_allowed(self, tool_name: str) -> bool:
"""
Check if a tool is allowed in planning mode.
Args:
tool_name: Name of the tool to check
Returns:
True if tool is allowed
"""
if not self.read_only:
return True
return tool_name in self.allowed_tools
@@ -0,0 +1,325 @@
"""
PlanStorage for persisting plans and todo lists.
Provides storage functionality similar to:
- Cursor .cursor/plans/
- Windsurf ~/.codeium/windsurf/brain/
- Codex ~/.codex/sessions/
Features:
- Save/load plans as markdown files
- Save/load todo lists as JSON
- Auto-discovery of existing plans
- Plan versioning
"""
import os
import json
import logging
from pathlib import Path
from datetime import datetime
from typing import List, Optional, Dict, Any
from .plan import Plan
from .todo import TodoList
logger = logging.getLogger(__name__)
class PlanStorage:
"""
Storage manager for plans and todo lists.
Attributes:
base_path: Base directory for storage
plans_dir: Directory for plan files
todos_dir: Directory for todo files
"""
DEFAULT_BASE_PATH = ".praison"
def __init__(self, base_path: Optional[str] = None):
"""
Initialize PlanStorage.
Args:
base_path: Base directory for storage. Defaults to .praison/
"""
self.base_path = base_path or self.DEFAULT_BASE_PATH
self.plans_dir = os.path.join(self.base_path, "plans")
self.todos_dir = os.path.join(self.base_path, "todos")
# Create directories if they don't exist
self._ensure_directories()
def _ensure_directories(self) -> None:
"""Create storage directories if they don't exist."""
os.makedirs(self.plans_dir, exist_ok=True)
os.makedirs(self.todos_dir, exist_ok=True)
def _get_plan_filename(self, plan: Plan) -> str:
"""
Generate filename for a plan.
Args:
plan: Plan to generate filename for
Returns:
Filename string
"""
date_str = plan.created_at.strftime("%Y-%m-%d") if plan.created_at else datetime.now().strftime("%Y-%m-%d")
safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in plan.name[:30])
return f"plan_{date_str}_{plan.id}_{safe_name}.md"
def _get_plan_path(self, plan_id: str) -> Optional[str]:
"""
Find the path for a plan by ID.
Args:
plan_id: ID of the plan to find
Returns:
Path to the plan file, or None if not found
"""
for filename in os.listdir(self.plans_dir):
if plan_id in filename and filename.endswith(".md"):
return os.path.join(self.plans_dir, filename)
return None
def save_plan(self, plan: Plan, filename: Optional[str] = None) -> str:
"""
Save a plan to storage.
Args:
plan: Plan to save
filename: Optional custom filename
Returns:
Path to the saved file
"""
if filename is None:
filename = self._get_plan_filename(plan)
path = os.path.join(self.plans_dir, filename)
# Convert plan to markdown
markdown = plan.to_markdown()
with open(path, "w", encoding="utf-8") as f:
f.write(markdown)
plan.file_path = path
logger.debug(f"Saved plan to {path}")
return path
def load_plan(self, plan_id: str) -> Optional[Plan]:
"""
Load a plan from storage.
Args:
plan_id: ID of the plan to load
Returns:
Plan instance, or None if not found
"""
path = self._get_plan_path(plan_id)
if path is None or not os.path.exists(path):
logger.debug(f"Plan {plan_id} not found")
return None
with open(path, "r", encoding="utf-8") as f:
markdown = f.read()
plan = Plan.from_markdown(markdown)
plan.file_path = path
logger.debug(f"Loaded plan from {path}")
return plan
def load_plan_from_file(self, path: str) -> Optional[Plan]:
"""
Load a plan from a specific file path.
Args:
path: Path to the plan file
Returns:
Plan instance, or None if file doesn't exist
"""
if not os.path.exists(path):
return None
with open(path, "r", encoding="utf-8") as f:
markdown = f.read()
plan = Plan.from_markdown(markdown)
plan.file_path = path
return plan
def list_plans(self) -> List[Dict[str, Any]]:
"""
List all saved plans.
Returns:
List of plan metadata dictionaries
"""
plans = []
for filename in os.listdir(self.plans_dir):
if filename.endswith(".md"):
path = os.path.join(self.plans_dir, filename)
# Try to load and get metadata
try:
plan = self.load_plan_from_file(path)
if plan:
plans.append({
"id": plan.id,
"name": plan.name,
"status": plan.status,
"created_at": plan.created_at.isoformat() if plan.created_at else None,
"file_path": path
})
except Exception as e:
logger.warning(f"Failed to load plan from {path}: {e}")
# Sort by creation date (newest first)
plans.sort(key=lambda p: p.get("created_at", ""), reverse=True)
return plans
def delete_plan(self, plan_id: str) -> bool:
"""
Delete a plan from storage.
Args:
plan_id: ID of the plan to delete
Returns:
True if deleted, False if not found
"""
path = self._get_plan_path(plan_id)
if path and os.path.exists(path):
os.remove(path)
logger.debug(f"Deleted plan {plan_id}")
return True
return False
def save_todo(self, todo: TodoList, name: str = "current") -> str:
"""
Save a todo list to storage.
Args:
todo: TodoList to save
name: Name for the todo file (without extension)
Returns:
Path to the saved file
"""
filename = f"{name}.json"
path = os.path.join(self.todos_dir, filename)
with open(path, "w", encoding="utf-8") as f:
f.write(todo.to_json())
logger.debug(f"Saved todo list to {path}")
return path
def load_todo(self, name: str = "current") -> Optional[TodoList]:
"""
Load a todo list from storage.
Args:
name: Name of the todo file (without extension)
Returns:
TodoList instance, or None if not found
"""
filename = f"{name}.json"
path = os.path.join(self.todos_dir, filename)
if not os.path.exists(path):
return None
with open(path, "r", encoding="utf-8") as f:
json_str = f.read()
logger.debug(f"Loaded todo list from {path}")
return TodoList.from_json(json_str)
def list_todos(self) -> List[str]:
"""
List all saved todo lists.
Returns:
List of todo list names
"""
todos = []
for filename in os.listdir(self.todos_dir):
if filename.endswith(".json"):
name = filename[:-5] # Remove .json extension
todos.append(name)
return todos
def delete_todo(self, name: str) -> bool:
"""
Delete a todo list from storage.
Args:
name: Name of the todo file (without extension)
Returns:
True if deleted, False if not found
"""
filename = f"{name}.json"
path = os.path.join(self.todos_dir, filename)
if os.path.exists(path):
os.remove(path)
logger.debug(f"Deleted todo list {name}")
return True
return False
def get_latest_plan(self) -> Optional[Plan]:
"""
Get the most recently created plan.
Returns:
Most recent Plan, or None if no plans exist
"""
plans = self.list_plans()
if not plans:
return None
return self.load_plan(plans[0]["id"])
def cleanup_old_plans(self, keep_count: int = 10) -> int:
"""
Remove old plans, keeping only the most recent ones.
Args:
keep_count: Number of plans to keep
Returns:
Number of plans deleted
"""
plans = self.list_plans()
if len(plans) <= keep_count:
return 0
deleted = 0
for plan in plans[keep_count:]:
if self.delete_plan(plan["id"]):
deleted += 1
return deleted
@@ -0,0 +1,415 @@
"""
TodoList and TodoItem classes for Planning Mode.
Provides task tracking functionality similar to:
- Cursor Agent To-Dos
- Windsurf Todo Lists
- Claude Code TodoRead/TodoWrite
Features:
- Create todo lists from plans
- Track item completion
- Dependencies between items
- Markdown and JSON serialization
"""
import uuid
import json
import logging
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any, Literal, Union, TYPE_CHECKING
if TYPE_CHECKING:
from .plan import Plan
logger = logging.getLogger(__name__)
@dataclass
class TodoItem:
"""
A single todo item.
Attributes:
id: Unique identifier
description: What needs to be done
status: Current status (pending, in_progress, completed)
dependencies: List of item IDs that must complete first
agent: Name of the agent responsible
priority: Priority level (low, medium, high)
notes: Additional notes
"""
description: str
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
status: Literal["pending", "in_progress", "completed"] = "pending"
dependencies: List[str] = field(default_factory=list)
agent: Optional[str] = None
priority: Literal["low", "medium", "high"] = "medium"
notes: str = ""
def to_dict(self) -> Dict[str, Any]:
"""Convert TodoItem to dictionary."""
return {
"id": self.id,
"description": self.description,
"status": self.status,
"dependencies": self.dependencies,
"agent": self.agent,
"priority": self.priority,
"notes": self.notes
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> 'TodoItem':
"""Create TodoItem from dictionary."""
return cls(
id=data.get("id", str(uuid.uuid4())[:8]),
description=data.get("description", ""),
status=data.get("status", "pending"),
dependencies=data.get("dependencies", []),
agent=data.get("agent"),
priority=data.get("priority", "medium"),
notes=data.get("notes", "")
)
def complete(self) -> None:
"""Mark this item as completed."""
self.status = "completed"
def start(self) -> None:
"""Mark this item as in progress."""
self.status = "in_progress"
def reset(self) -> None:
"""Reset this item to pending."""
self.status = "pending"
def is_ready(self, completed_ids: List[str]) -> bool:
"""
Check if this item is ready to start.
Args:
completed_ids: List of completed item IDs
Returns:
True if all dependencies are satisfied
"""
if self.status != "pending":
return False
return all(dep in completed_ids for dep in self.dependencies)
@dataclass
class TodoList:
"""
A list of todo items with tracking capabilities.
Attributes:
items: List of TodoItem objects
auto_update: Whether to auto-update from plan changes
name: Optional name for the todo list
"""
items: List[TodoItem] = field(default_factory=list)
auto_update: bool = True
name: str = "Todo List"
def add(self, item: Union[TodoItem, str]) -> TodoItem:
"""
Add an item to the list.
Args:
item: TodoItem or description string
Returns:
The added TodoItem
"""
if isinstance(item, str):
item = TodoItem(description=item)
self.items.append(item)
return item
def remove(self, item_id: str) -> bool:
"""
Remove an item from the list.
Args:
item_id: ID of the item to remove
Returns:
True if item was removed, False if not found
"""
for i, item in enumerate(self.items):
if item.id == item_id:
self.items.pop(i)
return True
return False
def get(self, item_id: str) -> Optional[TodoItem]:
"""
Get an item by ID.
Args:
item_id: ID of the item to find
Returns:
TodoItem if found, None otherwise
"""
for item in self.items:
if item.id == item_id:
return item
return None
def complete(self, item_id: str) -> bool:
"""
Mark an item as completed.
Args:
item_id: ID of the item to complete
Returns:
True if item was completed, False if not found
"""
item = self.get(item_id)
if item:
item.complete()
return True
return False
def start(self, item_id: str) -> bool:
"""
Mark an item as in progress.
Args:
item_id: ID of the item to start
Returns:
True if item was started, False if not found
"""
item = self.get(item_id)
if item:
item.start()
return True
return False
@property
def pending(self) -> List[TodoItem]:
"""Get all pending items."""
return [item for item in self.items if item.status == "pending"]
@property
def in_progress(self) -> List[TodoItem]:
"""Get all in-progress items."""
return [item for item in self.items if item.status == "in_progress"]
@property
def completed(self) -> List[TodoItem]:
"""Get all completed items."""
return [item for item in self.items if item.status == "completed"]
@property
def progress(self) -> float:
"""
Calculate completion progress.
Returns:
Float between 0.0 and 1.0
"""
if not self.items:
return 1.0
return len(self.completed) / len(self.items)
@property
def is_complete(self) -> bool:
"""Check if all items are completed."""
return all(item.status == "completed" for item in self.items)
@property
def completed_ids(self) -> List[str]:
"""Get list of completed item IDs."""
return [item.id for item in self.completed]
def get_ready_items(self) -> List[TodoItem]:
"""
Get items that are ready to start.
Returns:
List of items with satisfied dependencies
"""
completed = self.completed_ids
return [item for item in self.items if item.is_ready(completed)]
def to_markdown(self) -> str:
"""
Convert TodoList to markdown format.
Returns:
Markdown string with checkbox format
"""
lines = [f"# {self.name}", ""]
for item in self.items:
checkbox = "[x]" if item.status == "completed" else "[ ]"
lines.append(f"- {checkbox} {item.description}")
if item.agent:
lines.append(f" - Agent: {item.agent}")
if item.dependencies:
lines.append(f" - Depends on: {', '.join(item.dependencies)}")
if item.notes:
lines.append(f" - Notes: {item.notes}")
return "\n".join(lines)
@classmethod
def from_markdown(cls, markdown: str) -> 'TodoList':
"""
Create TodoList from markdown format.
Args:
markdown: Markdown string with checkbox format
Returns:
TodoList instance
"""
import re
items = []
# Match checkbox items: - [ ] or - [x]
pattern = r'^-\s*\[([ x])\]\s*(.+?)$'
for match in re.finditer(pattern, markdown, re.MULTILINE):
checked = match.group(1) == 'x'
description = match.group(2).strip()
item = TodoItem(
description=description,
status="completed" if checked else "pending"
)
items.append(item)
return cls(items=items)
def to_json(self) -> str:
"""
Serialize TodoList to JSON.
Returns:
JSON string
"""
data = {
"name": self.name,
"auto_update": self.auto_update,
"items": [item.to_dict() for item in self.items]
}
return json.dumps(data, indent=2)
@classmethod
def from_json(cls, json_str: str) -> 'TodoList':
"""
Create TodoList from JSON.
Args:
json_str: JSON string
Returns:
TodoList instance
"""
data = json.loads(json_str)
items = [TodoItem.from_dict(item) for item in data.get("items", [])]
return cls(
items=items,
auto_update=data.get("auto_update", True),
name=data.get("name", "Todo List")
)
def to_dict(self) -> Dict[str, Any]:
"""Convert TodoList to dictionary."""
return {
"name": self.name,
"auto_update": self.auto_update,
"items": [item.to_dict() for item in self.items]
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> 'TodoList':
"""Create TodoList from dictionary."""
items = [TodoItem.from_dict(item) for item in data.get("items", [])]
return cls(
items=items,
auto_update=data.get("auto_update", True),
name=data.get("name", "Todo List")
)
@classmethod
def from_plan(cls, plan: 'Plan') -> 'TodoList':
"""
Create TodoList from a Plan.
Args:
plan: Plan instance
Returns:
TodoList with items matching plan steps
"""
items = []
for step in plan.steps:
item = TodoItem(
id=step.id,
description=step.description,
status=step.status if step.status in ["pending", "in_progress", "completed"] else "pending",
dependencies=step.dependencies,
agent=step.agent
)
items.append(item)
return cls(
items=items,
name=f"Todo: {plan.name}"
)
def sync_with_plan(self, plan: 'Plan') -> None:
"""
Synchronize TodoList with a Plan.
Updates item statuses to match plan step statuses.
Args:
plan: Plan to sync with
"""
step_statuses = {step.id: step.status for step in plan.steps}
for item in self.items:
if item.id in step_statuses:
status = step_statuses[item.id]
if status in ["pending", "in_progress", "completed"]:
item.status = status
def update_from_plan(self, plan: 'Plan') -> None:
"""
Update TodoList to match Plan structure.
Adds new items for new steps, removes items for removed steps.
Args:
plan: Plan to update from
"""
# Get current step IDs
step_ids = {step.id for step in plan.steps}
item_ids = {item.id for item in self.items}
# Remove items for deleted steps
self.items = [item for item in self.items if item.id in step_ids]
# Add items for new steps
for step in plan.steps:
if step.id not in item_ids:
self.add(TodoItem(
id=step.id,
description=step.description,
dependencies=step.dependencies,
agent=step.agent
))
# Sync statuses
self.sync_with_plan(plan)
@@ -9,6 +9,85 @@ A tool is a piece of code that helps our AI agents perform specific tasks. Think
- A stock market tool lets agents check stock prices
- A weather tool lets agents check the weather
## Plugin System (NEW!)
PraisonAI now supports a plugin system for creating and distributing tools. External developers can create pip-installable tool packages!
### Quick Start
```python
from praisonaiagents import BaseTool, tool, Agent
# Method 1: Class-based tool
class WeatherTool(BaseTool):
name = "get_weather"
description = "Get weather for a location"
def run(self, location: str) -> dict:
return {"temp": 72, "condition": "sunny"}
# Method 2: Decorator-based tool
@tool
def search(query: str) -> list:
"""Search the web."""
return [{"title": "Result", "url": "https://..."}]
# Use with Agent
agent = Agent(
name="Assistant",
tools=[WeatherTool(), search]
)
```
### Creating a Plugin Package
External developers can create pip-installable plugins:
**1. Create your tool package:**
```
praisonai-weather/
├── pyproject.toml
├── src/
│ └── praisonai_weather/
│ ├── __init__.py
│ └── tools.py
```
**2. Define your tool in `tools.py`:**
```python
from praisonaiagents import BaseTool
class WeatherTool(BaseTool):
name = "weather"
description = "Get current weather"
def run(self, location: str) -> dict:
# Your implementation
return {"temp": 72}
```
**3. Register via entry_points in `pyproject.toml`:**
```toml
[project]
name = "praisonai-weather"
version = "1.0.0"
[project.entry-points."praisonaiagents.tools"]
weather = "praisonai_weather.tools:WeatherTool"
```
**4. Users install and use:**
```bash
pip install praisonai-weather
```
```python
from praisonaiagents import Agent
# Tool is auto-discovered!
agent = Agent(tools=["weather"])
```
## Creating New Tools: The Two Approaches
### 1. Function-Based Approach (Simple Tools)
@@ -173,6 +173,51 @@ TOOL_MAPPINGS = {
'cot_save': ('.train.data.generatecot', 'GenerateCOT'), # Save single QA to file
'cot_upload_to_huggingface': ('.train.data.generatecot', 'GenerateCOT'), # Upload dataset to HuggingFace
'cot_tools': ('.train.data.generatecot', 'GenerateCOT'), # Full toolkit access
# Claude Memory Tool (Anthropic Beta)
'ClaudeMemoryTool': ('.claude_memory_tool', 'ClaudeMemoryTool'),
# Tavily Tools (requires TAVILY_API_KEY)
'tavily': ('.tavily_tools', None),
'tavily_search': ('.tavily_tools', None),
'tavily_extract': ('.tavily_tools', None),
'tavily_crawl': ('.tavily_tools', None),
'tavily_map': ('.tavily_tools', None),
'tavily_search_async': ('.tavily_tools', None),
'tavily_extract_async': ('.tavily_tools', None),
'TavilyTools': ('.tavily_tools', 'TavilyTools'),
'tavily_tools': ('.tavily_tools', None),
# You.com Tools (requires YDC_API_KEY)
'ydc': ('.youdotcom_tools', None),
'ydc_search': ('.youdotcom_tools', None),
'ydc_contents': ('.youdotcom_tools', None),
'ydc_news': ('.youdotcom_tools', None),
'ydc_images': ('.youdotcom_tools', None),
'YouTools': ('.youdotcom_tools', 'YouTools'),
'youdotcom_tools': ('.youdotcom_tools', None),
# Exa Tools (requires EXA_API_KEY)
'exa': ('.exa_tools', None),
'exa_search': ('.exa_tools', None),
'exa_search_contents': ('.exa_tools', None),
'exa_find_similar': ('.exa_tools', None),
'exa_answer': ('.exa_tools', None),
'exa_search_async': ('.exa_tools', None),
'exa_search_contents_async': ('.exa_tools', None),
'exa_answer_async': ('.exa_tools', None),
'ExaTools': ('.exa_tools', 'ExaTools'),
'exa_tools': ('.exa_tools', None),
# Crawl4AI Tools (async web crawling)
'crawl4ai': ('.crawl4ai_tools', None),
'crawl4ai_many': ('.crawl4ai_tools', None),
'crawl4ai_extract': ('.crawl4ai_tools', None),
'crawl4ai_llm_extract': ('.crawl4ai_tools', None),
'crawl4ai_sync': ('.crawl4ai_tools', None),
'crawl4ai_extract_sync': ('.crawl4ai_tools', None),
'Crawl4AITools': ('.crawl4ai_tools', 'Crawl4AITools'),
'crawl4ai_tools': ('.crawl4ai_tools', None),
}
_instances = {} # Cache for class instances
@@ -184,6 +229,11 @@ def __getattr__(name: str) -> Any:
module_path, class_name = TOOL_MAPPINGS[name]
# Return class itself (not instance) for ClaudeMemoryTool, TavilyTools, YouTools, ExaTools, Crawl4AITools
if name in ('ClaudeMemoryTool', 'TavilyTools', 'YouTools', 'ExaTools', 'Crawl4AITools'):
module = import_module(module_path, __package__)
return getattr(module, class_name)
if class_name is None:
# Direct function import
module = import_module(module_path, __package__)
@@ -197,11 +247,18 @@ def __getattr__(name: str) -> Any:
'insert_document', 'insert_documents', 'find_documents', 'update_document', 'delete_document',
'create_vector_index', 'vector_search', 'store_with_embedding', 'text_search', 'get_stats', 'connect_mongodb',
'execute_command', 'list_processes', 'kill_process', 'get_system_info',
'evaluate', 'solve_equation', 'convert_units', 'calculate_statistics', 'calculate_financial'
'evaluate', 'solve_equation', 'convert_units', 'calculate_statistics', 'calculate_financial',
'tavily', 'tavily_search', 'tavily_extract', 'tavily_crawl', 'tavily_map',
'tavily_search_async', 'tavily_extract_async',
'ydc', 'ydc_search', 'ydc_contents', 'ydc_news', 'ydc_images',
'exa', 'exa_search', 'exa_search_contents', 'exa_find_similar', 'exa_answer',
'exa_search_async', 'exa_search_contents_async', 'exa_answer_async',
'crawl4ai', 'crawl4ai_many', 'crawl4ai_extract', 'crawl4ai_llm_extract',
'crawl4ai_sync', 'crawl4ai_extract_sync'
]:
return getattr(module, name)
if name in ['file_tools', 'pandas_tools', 'wikipedia_tools',
'newspaper_tools', 'arxiv_tools', 'spider_tools', 'duckdb_tools', 'mongodb_tools', 'csv_tools', 'json_tools', 'excel_tools', 'xml_tools', 'yaml_tools', 'calculator_tools', 'python_tools', 'shell_tools', 'cot_tools']:
'newspaper_tools', 'arxiv_tools', 'spider_tools', 'duckdb_tools', 'mongodb_tools', 'csv_tools', 'json_tools', 'excel_tools', 'xml_tools', 'yaml_tools', 'calculator_tools', 'python_tools', 'shell_tools', 'cot_tools', 'tavily_tools', 'youdotcom_tools', 'exa_tools', 'crawl4ai_tools']:
return module # Returns the callable module
return getattr(module, name)
else:
@@ -0,0 +1,298 @@
"""Base classes for PraisonAI Agent tools.
This module provides the foundation for creating tools that can be used by agents.
External developers can create plugins by subclassing BaseTool.
Usage:
from praisonaiagents import BaseTool
class MyTool(BaseTool):
name = "my_tool"
description = "Does something useful"
def run(self, query: str) -> str:
return f"Result for {query}"
"""
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Type, get_type_hints
import inspect
import logging
class ToolValidationError(Exception):
"""Raised when a tool fails validation."""
pass
class ToolResult:
"""Wrapper for tool execution results."""
def __init__(
self,
output: Any,
success: bool = True,
error: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None
):
self.output = output
self.success = success
self.error = error
self.metadata = metadata or {}
def __str__(self) -> str:
if self.success:
return str(self.output)
return f"Error: {self.error}"
def __repr__(self) -> str:
return f"ToolResult(success={self.success}, output={self.output!r})"
def to_dict(self) -> Dict[str, Any]:
return {
"output": self.output,
"success": self.success,
"error": self.error,
"metadata": self.metadata
}
class BaseTool(ABC):
"""Abstract base class for all PraisonAI tools.
Subclass this to create custom tools that can be:
- Used directly by agents
- Distributed as pip-installable plugins
- Auto-discovered via entry_points
Attributes:
name: Unique identifier for the tool
description: Human-readable description (used by LLM)
version: Tool version string (default: "1.0.0")
parameters: JSON Schema for parameters (auto-generated if not provided)
Example:
class WeatherTool(BaseTool):
name = "get_weather"
description = "Get current weather for a location"
def run(self, location: str, units: str = "celsius") -> dict:
# Implementation here
return {"temp": 22, "condition": "sunny"}
"""
# Required class attributes (must be overridden)
name: str = ""
description: str = ""
# Optional class attributes
version: str = "1.0.0"
parameters: Optional[Dict[str, Any]] = None # JSON Schema, auto-generated if None
def __init__(self):
"""Initialize the tool and validate configuration."""
if not self.name:
# Use class name as default
self.name = self.__class__.__name__.lower().replace("tool", "")
if not self.description:
# Use docstring as default
self.description = self.__class__.__doc__ or f"Tool: {self.name}"
# Auto-generate parameters schema if not provided
if self.parameters is None:
self.parameters = self._generate_parameters_schema()
def _generate_parameters_schema(self) -> Dict[str, Any]:
"""Generate JSON Schema from run() method signature."""
schema = {
"type": "object",
"properties": {},
"required": []
}
try:
sig = inspect.signature(self.run)
hints = get_type_hints(self.run) if hasattr(self.run, '__annotations__') else {}
for param_name, param in sig.parameters.items():
if param_name == 'self':
continue
# Get type hint
param_type = hints.get(param_name, Any)
json_type = self._python_type_to_json(param_type)
# Build property schema
prop_schema = {"type": json_type}
# Add description from docstring if available
# (Could parse docstring for param descriptions)
schema["properties"][param_name] = prop_schema
# Check if required (no default value)
if param.default is inspect.Parameter.empty:
schema["required"].append(param_name)
except Exception as e:
logging.debug(f"Could not generate schema for {self.name}: {e}")
return schema
@staticmethod
def _python_type_to_json(python_type: Type) -> str:
"""Convert Python type to JSON Schema type."""
type_map = {
str: "string",
int: "integer",
float: "number",
bool: "boolean",
list: "array",
dict: "object",
type(None): "null"
}
# Handle Optional, Union, etc.
origin = getattr(python_type, '__origin__', None)
if origin is not None:
# For List[X], return "array"
if origin is list:
return "array"
# For Dict[X, Y], return "object"
if origin is dict:
return "object"
return type_map.get(python_type, "string")
@abstractmethod
def run(self, **kwargs) -> Any:
"""Execute the tool with given arguments.
This method must be implemented by subclasses.
Args:
**kwargs: Tool-specific arguments
Returns:
Tool output (any type, will be converted to string for LLM)
"""
pass
def __call__(self, **kwargs) -> Any:
"""Allow tool to be called directly like a function."""
return self.run(**kwargs)
def safe_run(self, **kwargs) -> ToolResult:
"""Execute tool with error handling, returning ToolResult."""
try:
output = self.run(**kwargs)
return ToolResult(output=output, success=True)
except Exception as e:
logging.error(f"Tool {self.name} failed: {e}")
return ToolResult(
output=None,
success=False,
error=str(e)
)
def get_schema(self) -> Dict[str, Any]:
"""Get OpenAI-compatible function schema for this tool."""
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": self.parameters
}
}
def __str__(self) -> str:
return f"{self.__class__.__name__}(name='{self.name}')"
def __repr__(self) -> str:
return f"{self.__class__.__name__}(name='{self.name}', description='{self.description[:50]}...')"
def validate(self) -> bool:
"""Validate the tool configuration.
Raises:
ToolValidationError: If validation fails
Returns:
True if validation passes
"""
errors = []
# Check required fields
if not self.name or not isinstance(self.name, str):
errors.append("Tool must have a non-empty string 'name'")
if not self.description or not isinstance(self.description, str):
errors.append("Tool must have a non-empty string 'description'")
# Check that run() is implemented (not abstract)
if getattr(self.run, '__isabstractmethod__', False):
errors.append("Tool must implement the 'run()' method")
# Check parameters schema is valid
if self.parameters:
if not isinstance(self.parameters, dict):
errors.append("'parameters' must be a dictionary")
elif "type" not in self.parameters:
errors.append("'parameters' must have a 'type' field")
if errors:
raise ToolValidationError(f"Tool '{self.name}' validation failed: {'; '.join(errors)}")
return True
@classmethod
def validate_class(cls) -> bool:
"""Validate a tool class before instantiation.
This can be used to check if a class is a valid tool without creating an instance.
Returns:
True if the class appears to be a valid tool
"""
# Check it's a subclass of BaseTool
if not issubclass(cls, BaseTool):
return False
# Check it has required class attributes or will get them from __init__
# (name and description can be set in __init__ so we can't strictly require them)
# Check run() is defined (not just inherited abstract)
if 'run' not in cls.__dict__:
return False
return True
def validate_tool(tool: Any) -> bool:
"""Validate any tool-like object.
Args:
tool: Object to validate (BaseTool, callable, etc.)
Returns:
True if valid
Raises:
ToolValidationError: If validation fails
"""
if isinstance(tool, BaseTool):
return tool.validate()
if callable(tool):
# For plain functions, check they have a name
name = getattr(tool, '__name__', None) or getattr(tool, 'name', None)
if not name:
raise ToolValidationError("Callable tool must have a __name__ or name attribute")
return True
raise ToolValidationError(f"Invalid tool type: {type(tool)}")
# For backward compatibility - tools can also just be functions
# The @tool decorator (in decorator.py) wraps functions into BaseTool instances
@@ -0,0 +1,529 @@
"""
Claude Memory Tool - File-based persistent memory for Claude models.
This module implements Anthropic's Memory Tool API (beta) which enables Claude
to store and retrieve information across conversations through a memory file directory.
The memory tool operates client-side - you control where and how the data is stored.
Beta Header: context-management-2025-06-27
Tool Type: memory_20250818
Supported Models:
- Claude Sonnet 4.5 (claude-sonnet-4-5-20250929)
- Claude Sonnet 4 (claude-sonnet-4-20250514)
- Claude Haiku 4.5 (claude-haiku-4-5-20251001)
- Claude Opus 4.5 (claude-opus-4-5-20251101)
- Claude Opus 4.1 (claude-opus-4-1-20250805)
- Claude Opus 4 (claude-opus-4-20250514)
Reference: https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/memory-tool
"""
import shutil
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional
logger = logging.getLogger(__name__)
# Memory system prompt that Claude uses to manage memory
MEMORY_SYSTEM_PROMPT = """
IMPORTANT: ALWAYS VIEW YOUR MEMORY DIRECTORY BEFORE DOING ANYTHING ELSE.
MEMORY PROTOCOL:
1. Use the `view` command of your `memory` tool to check for earlier progress.
2. ... (work on the task) ...
- As you make progress, record status / progress / thoughts etc in your memory.
ASSUME INTERRUPTION: Your context window might be reset at any moment, so you risk losing any progress that is not recorded in your memory directory.
When editing your memory folder, always try to keep its content up-to-date, coherent and organized. You can rename or delete files that are no longer relevant. Do not create new files unless necessary.
"""
class ClaudeMemoryTool:
"""
File-based memory tool for Claude models.
Implements Anthropic's Memory Tool API (beta). Claude autonomously decides
what to store and retrieve from the memory directory.
Storage: {base_path}/memories/
Commands:
- view: Show directory contents or file contents with optional line ranges
- create: Create or overwrite a file
- str_replace: Replace text in a file
- insert: Insert text at a specific line
- delete: Delete a file or directory
- rename: Rename or move a file/directory
Example:
```python
from praisonaiagents import Agent
from praisonaiagents.tools import ClaudeMemoryTool
# Use default storage
agent = Agent(
name="Assistant",
llm="anthropic/claude-sonnet-4-20250514",
claude_memory=True
)
# Or with custom storage path
memory = ClaudeMemoryTool(base_path="./my_project/memory")
agent = Agent(
name="Assistant",
llm="anthropic/claude-sonnet-4-20250514",
claude_memory=memory
)
```
Security:
- All paths are validated to prevent directory traversal attacks
- Paths must start with /memories
- Symlinks are not followed outside the memory directory
"""
# Beta header required for memory tool
BETA_HEADER = "context-management-2025-06-27"
# Tool definition for API requests
TOOL_DEFINITION = {
"type": "memory_20250818",
"name": "memory"
}
def __init__(
self,
base_path: Optional[str] = None,
user_id: Optional[str] = None,
max_file_size: int = 1024 * 1024, # 1MB default
max_files: int = 100
):
"""
Initialize the Claude Memory Tool.
Args:
base_path: Base directory for memory storage.
Defaults to .praison/claude_memory/{user_id}
user_id: User identifier for user-specific memory.
Defaults to "default"
max_file_size: Maximum size per file in bytes (default 1MB)
max_files: Maximum number of files allowed (default 100)
"""
self.user_id = user_id or "default"
if base_path:
self.base_path = Path(base_path)
else:
self.base_path = Path(f".praison/claude_memory/{self.user_id}")
self.memory_root = self.base_path / "memories"
self.max_file_size = max_file_size
self.max_files = max_files
# Create memory directory
self.memory_root.mkdir(parents=True, exist_ok=True)
logger.debug(f"ClaudeMemoryTool initialized at {self.memory_root}")
def _validate_path(self, path: str) -> Path:
"""
Validate and resolve memory paths securely.
Args:
path: Path string starting with /memories
Returns:
Resolved Path object within memory directory
Raises:
ValueError: If path is invalid or attempts directory traversal
"""
if not path:
raise ValueError("Path cannot be empty")
if not path.startswith("/memories"):
raise ValueError(f"Path must start with /memories, got: {path}")
# Remove /memories prefix and leading slashes
relative_path = path[len("/memories"):].lstrip("/")
# Resolve the full path
if relative_path:
full_path = self.memory_root / relative_path
else:
full_path = self.memory_root
# Security check: ensure resolved path is within memory_root
try:
resolved = full_path.resolve()
resolved.relative_to(self.memory_root.resolve())
except ValueError as e:
raise ValueError(f"Path {path} would escape /memories directory") from e
# Check for suspicious patterns
if ".." in path or "~" in path:
raise ValueError(f"Invalid path pattern in: {path}")
return full_path
def execute(self, command: str, **kwargs) -> str:
"""
Execute a memory command.
Args:
command: Command name (view, create, str_replace, insert, delete, rename)
**kwargs: Command-specific arguments
Returns:
Result string from the command
Raises:
ValueError: If command is unknown
"""
commands = {
"view": self.view,
"create": self.create,
"str_replace": self.str_replace,
"insert": self.insert,
"delete": self.delete,
"rename": self.rename
}
if command not in commands:
raise ValueError(f"Unknown command: {command}. Valid commands: {list(commands.keys())}")
return commands[command](**kwargs)
def view(self, path: str, view_range: Optional[List[int]] = None) -> str:
"""
Show directory contents or file contents with optional line ranges.
Args:
path: Path to view (e.g., "/memories" or "/memories/notes.txt")
view_range: Optional [start_line, end_line] for file viewing.
Use -1 for end_line to read to end of file.
Returns:
Directory listing or file contents with line numbers
Raises:
RuntimeError: If path cannot be read
"""
full_path = self._validate_path(path)
if full_path.is_dir():
# List directory contents
items: List[str] = []
try:
for item in sorted(full_path.iterdir()):
# Skip hidden files
if item.name.startswith("."):
continue
# Add trailing slash for directories
items.append(f"{item.name}/" if item.is_dir() else item.name)
result = f"Directory: {path}\n"
if items:
result += "\n".join([f"- {item}" for item in items])
else:
result += "(empty directory)"
return result
except Exception as e:
raise RuntimeError(f"Cannot read directory {path}: {e}") from e
elif full_path.is_file():
# Read file contents
try:
content = full_path.read_text(encoding="utf-8")
lines = content.splitlines()
# Apply view range if specified
if view_range:
start_line = max(1, view_range[0]) - 1 # Convert to 0-indexed
end_line = len(lines) if view_range[1] == -1 else view_range[1]
lines = lines[start_line:end_line]
start_num = start_line + 1
else:
start_num = 1
# Add line numbers
numbered_lines = [
f"{i + start_num:4d}: {line}"
for i, line in enumerate(lines)
]
return "\n".join(numbered_lines) if numbered_lines else "(empty file)"
except Exception as e:
raise RuntimeError(f"Cannot read file {path}: {e}") from e
else:
raise RuntimeError(f"Path not found: {path}")
def create(self, path: str, file_text: str) -> str:
"""
Create or overwrite a file.
Args:
path: Path for the new file (e.g., "/memories/notes.txt")
file_text: Content to write to the file
Returns:
Success message
Raises:
ValueError: If file would exceed size limit or max files reached
"""
full_path = self._validate_path(path)
# Check file size limit
if len(file_text.encode('utf-8')) > self.max_file_size:
raise ValueError(
f"File content exceeds maximum size of {self.max_file_size} bytes"
)
# Check max files limit (only for new files)
if not full_path.exists():
current_files = sum(1 for _ in self.memory_root.rglob("*") if _.is_file())
if current_files >= self.max_files:
raise ValueError(
f"Maximum number of files ({self.max_files}) reached. "
"Delete some files before creating new ones."
)
# Create parent directories if needed
full_path.parent.mkdir(parents=True, exist_ok=True)
# Write file
full_path.write_text(file_text, encoding="utf-8")
return f"File created successfully at {path}"
def str_replace(self, path: str, old_str: str, new_str: str) -> str:
"""
Replace text in a file.
Args:
path: Path to the file
old_str: Text to find (must be unique in file)
new_str: Replacement text
Returns:
Success message
Raises:
FileNotFoundError: If file doesn't exist
ValueError: If old_str not found or appears multiple times
"""
full_path = self._validate_path(path)
if not full_path.is_file():
raise FileNotFoundError(f"File not found: {path}")
content = full_path.read_text(encoding="utf-8")
# Count occurrences
count = content.count(old_str)
if count == 0:
raise ValueError(f"Text not found in {path}")
elif count > 1:
raise ValueError(
f"Text appears {count} times in {path}. Must be unique for replacement."
)
# Replace and write
new_content = content.replace(old_str, new_str)
# Check size limit
if len(new_content.encode('utf-8')) > self.max_file_size:
raise ValueError(
f"Resulting file would exceed maximum size of {self.max_file_size} bytes"
)
full_path.write_text(new_content, encoding="utf-8")
return f"File {path} has been edited"
def insert(self, path: str, insert_line: int, insert_text: str) -> str:
"""
Insert text at a specific line.
Args:
path: Path to the file
insert_line: Line number to insert at (0-indexed)
insert_text: Text to insert
Returns:
Success message
Raises:
FileNotFoundError: If file doesn't exist
ValueError: If line number is invalid
"""
full_path = self._validate_path(path)
if not full_path.is_file():
raise FileNotFoundError(f"File not found: {path}")
lines = full_path.read_text(encoding="utf-8").splitlines()
# Validate line number
if insert_line < 0 or insert_line > len(lines):
raise ValueError(
f"Invalid insert_line {insert_line}. Must be 0-{len(lines)}"
)
# Insert the text (strip trailing newline to avoid double newlines)
lines.insert(insert_line, insert_text.rstrip("\n"))
# Reconstruct content
new_content = "\n".join(lines) + "\n"
# Check size limit
if len(new_content.encode('utf-8')) > self.max_file_size:
raise ValueError(
f"Resulting file would exceed maximum size of {self.max_file_size} bytes"
)
full_path.write_text(new_content, encoding="utf-8")
return f"Text inserted at line {insert_line} in {path}"
def delete(self, path: str) -> str:
"""
Delete a file or directory.
Args:
path: Path to delete
Returns:
Success message
Raises:
ValueError: If trying to delete /memories root
FileNotFoundError: If path doesn't exist
"""
# Prevent deleting the root memories directory
if path == "/memories" or path == "/memories/":
raise ValueError("Cannot delete the /memories directory itself")
full_path = self._validate_path(path)
if full_path.is_file():
full_path.unlink()
return f"File deleted: {path}"
elif full_path.is_dir():
shutil.rmtree(full_path)
return f"Directory deleted: {path}"
else:
raise FileNotFoundError(f"Path not found: {path}")
def rename(self, old_path: str, new_path: str) -> str:
"""
Rename or move a file/directory.
Args:
old_path: Current path
new_path: New path
Returns:
Success message
Raises:
FileNotFoundError: If source doesn't exist
ValueError: If destination already exists
"""
old_full_path = self._validate_path(old_path)
new_full_path = self._validate_path(new_path)
if not old_full_path.exists():
raise FileNotFoundError(f"Source path not found: {old_path}")
if new_full_path.exists():
raise ValueError(f"Destination already exists: {new_path}")
# Create parent directories if needed
new_full_path.parent.mkdir(parents=True, exist_ok=True)
# Rename/move
old_full_path.rename(new_full_path)
return f"Renamed {old_path} to {new_path}"
def clear_all(self) -> str:
"""
Clear all memory files.
Returns:
Success message
"""
if self.memory_root.exists():
shutil.rmtree(self.memory_root)
self.memory_root.mkdir(parents=True, exist_ok=True)
return "All memory cleared"
def get_tool_definition(self) -> Dict[str, Any]:
"""
Get the tool definition for API requests.
Returns:
Tool definition dict for Anthropic API
"""
return self.TOOL_DEFINITION.copy()
def get_beta_header(self) -> str:
"""
Get the beta header value for API requests.
Returns:
Beta header string
"""
return self.BETA_HEADER
def get_system_prompt(self) -> str:
"""
Get the memory system prompt to append to agent instructions.
Returns:
Memory protocol system prompt
"""
return MEMORY_SYSTEM_PROMPT
def process_tool_call(self, tool_input: Dict[str, Any]) -> str:
"""
Process a memory tool call from Claude.
Args:
tool_input: The input dict from Claude's tool_use block
Contains 'command' and command-specific parameters
Returns:
Result string to return to Claude
Example:
```python
# Claude sends:
# {"command": "view", "path": "/memories"}
result = memory_tool.process_tool_call(tool_input)
```
"""
command = tool_input.get("command")
if not command:
return "Error: No command specified in tool input"
try:
# Extract command-specific parameters
params = {k: v for k, v in tool_input.items() if k != "command"}
return self.execute(command, **params)
except Exception as e:
return f"Error: {str(e)}"
def __repr__(self) -> str:
return f"ClaudeMemoryTool(base_path='{self.base_path}', user_id='{self.user_id}')"
@@ -0,0 +1,630 @@
"""Crawl4AI web crawling and extraction tools.
Crawl4AI provides async web crawling with JavaScript rendering, content extraction,
and LLM-based data extraction capabilities.
Usage:
from praisonaiagents.tools import crawl4ai, crawl4ai_extract
# Simple crawl
result = await crawl4ai("https://example.com")
print(result["markdown"])
# Extract structured data
result = await crawl4ai_extract(
"https://example.com/products",
schema={"name": "Products", "baseSelector": ".product", ...}
)
# Or use the class directly
from praisonaiagents.tools import Crawl4AITools
tools = Crawl4AITools()
result = await tools.crawl("https://example.com")
"""
from typing import List, Dict, Any, Optional, Union
import logging
import asyncio
from importlib import util
def _check_crawl4ai_available() -> tuple[bool, Optional[str]]:
"""Check if Crawl4AI is available.
Returns:
Tuple of (is_available, error_message)
"""
if util.find_spec("crawl4ai") is None:
return False, "crawl4ai package is not installed. Install it with: pip install crawl4ai && crawl4ai-setup"
return True, None
class Crawl4AITools:
"""Crawl4AI web crawling and extraction tools.
Provides methods for:
- crawl: Basic web crawling with markdown output
- crawl_many: Crawl multiple URLs concurrently
- extract_css: Extract data using CSS selectors
- extract_llm: Extract data using LLM
Example:
from praisonaiagents.tools import Crawl4AITools
tools = Crawl4AITools()
result = await tools.crawl("https://example.com")
print(result["markdown"])
"""
def __init__(
self,
headless: bool = True,
verbose: bool = False,
browser_type: str = "chromium"
):
"""Initialize Crawl4AITools.
Args:
headless: Run browser in headless mode (default True)
verbose: Enable verbose logging (default False)
browser_type: Browser type - "chromium", "firefox", or "webkit"
"""
self.headless = headless
self.verbose = verbose
self.browser_type = browser_type
self.logger = logging.getLogger(__name__)
self._crawler = None
async def _get_crawler(self):
"""Get or create AsyncWebCrawler instance."""
is_available, error = _check_crawl4ai_available()
if not is_available:
raise ImportError(error)
from crawl4ai import AsyncWebCrawler, BrowserConfig
if self._crawler is None:
browser_config = BrowserConfig(
headless=self.headless,
verbose=self.verbose,
browser_type=self.browser_type
)
self._crawler = AsyncWebCrawler(config=browser_config)
await self._crawler.start()
return self._crawler
async def close(self):
"""Close the crawler and release resources."""
if self._crawler is not None:
await self._crawler.close()
self._crawler = None
async def crawl(
self,
url: str,
word_count_threshold: int = 10,
css_selector: Optional[str] = None,
excluded_tags: Optional[List[str]] = None,
exclude_external_links: bool = False,
js_code: Optional[Union[str, List[str]]] = None,
wait_for: Optional[str] = None,
screenshot: bool = False,
pdf: bool = False,
bypass_cache: bool = True
) -> Dict[str, Any]:
"""Crawl a single URL and return content.
Args:
url: The URL to crawl
word_count_threshold: Minimum words per text block (default 10)
css_selector: CSS selector to focus on specific content
excluded_tags: HTML tags to exclude (e.g., ["nav", "footer"])
exclude_external_links: Remove external links from content
js_code: JavaScript code to execute before crawling
wait_for: Wait condition - "css:selector" or "js:() => boolean"
screenshot: Capture screenshot (base64)
pdf: Generate PDF
bypass_cache: Skip cache and fetch fresh content (default True)
Returns:
Dict with 'markdown', 'html', 'cleaned_html', 'links', 'media', etc.
"""
try:
crawler = await self._get_crawler()
from crawl4ai import CrawlerRunConfig, CacheMode
config = CrawlerRunConfig(
word_count_threshold=word_count_threshold,
css_selector=css_selector,
excluded_tags=excluded_tags or [],
exclude_external_links=exclude_external_links,
js_code=js_code,
wait_for=wait_for,
screenshot=screenshot,
pdf=pdf,
cache_mode=CacheMode.BYPASS if bypass_cache else CacheMode.ENABLED
)
result = await crawler.arun(url=url, config=config)
if not result.success:
return {"error": f"Crawl failed: {result.error_message}", "success": False}
response = {
"success": True,
"url": result.url,
"markdown": result.markdown.raw_markdown if hasattr(result.markdown, 'raw_markdown') else str(result.markdown),
"html": result.html,
"cleaned_html": result.cleaned_html,
"links": result.links,
"media": result.media,
"metadata": result.metadata
}
if hasattr(result.markdown, 'fit_markdown') and result.markdown.fit_markdown:
response["fit_markdown"] = result.markdown.fit_markdown
if screenshot and result.screenshot:
response["screenshot"] = result.screenshot
if pdf and result.pdf:
response["pdf"] = result.pdf
return response
except Exception as e:
self.logger.error(f"Crawl4AI crawl error: {e}")
return {"error": f"Crawl4AI crawl error: {str(e)}", "success": False}
async def crawl_many(
self,
urls: List[str],
word_count_threshold: int = 10,
css_selector: Optional[str] = None,
bypass_cache: bool = True,
stream: bool = False
) -> List[Dict[str, Any]]:
"""Crawl multiple URLs concurrently.
Args:
urls: List of URLs to crawl
word_count_threshold: Minimum words per text block
css_selector: CSS selector to focus on specific content
bypass_cache: Skip cache and fetch fresh content
stream: If True, returns async generator
Returns:
List of crawl results
"""
try:
crawler = await self._get_crawler()
from crawl4ai import CrawlerRunConfig, CacheMode
config = CrawlerRunConfig(
word_count_threshold=word_count_threshold,
css_selector=css_selector,
cache_mode=CacheMode.BYPASS if bypass_cache else CacheMode.ENABLED,
stream=stream
)
results = []
if stream:
async for result in await crawler.arun_many(urls, config=config):
if result.success:
results.append({
"success": True,
"url": result.url,
"markdown": result.markdown.raw_markdown if hasattr(result.markdown, 'raw_markdown') else str(result.markdown),
"html": result.html
})
else:
results.append({
"success": False,
"url": result.url,
"error": result.error_message
})
else:
raw_results = await crawler.arun_many(urls, config=config)
for result in raw_results:
if result.success:
results.append({
"success": True,
"url": result.url,
"markdown": result.markdown.raw_markdown if hasattr(result.markdown, 'raw_markdown') else str(result.markdown),
"html": result.html
})
else:
results.append({
"success": False,
"url": result.url,
"error": result.error_message
})
return results
except Exception as e:
self.logger.error(f"Crawl4AI crawl_many error: {e}")
return [{"error": f"Crawl4AI crawl_many error: {str(e)}", "success": False}]
async def extract_css(
self,
url: str,
schema: Dict[str, Any],
js_code: Optional[Union[str, List[str]]] = None,
wait_for: Optional[str] = None,
bypass_cache: bool = True
) -> Dict[str, Any]:
"""Extract structured data using CSS selectors.
Args:
url: The URL to crawl (or "raw://<html>" for raw HTML)
schema: Extraction schema with baseSelector and fields
js_code: JavaScript to execute before extraction
wait_for: Wait condition before extraction
bypass_cache: Skip cache
Returns:
Dict with 'data' containing extracted items
Example schema:
{
"name": "Products",
"baseSelector": "div.product",
"fields": [
{"name": "title", "selector": "h2", "type": "text"},
{"name": "price", "selector": ".price", "type": "text"},
{"name": "link", "selector": "a", "type": "attribute", "attribute": "href"}
]
}
"""
try:
crawler = await self._get_crawler()
from crawl4ai import CrawlerRunConfig, CacheMode
from crawl4ai.extraction_strategy import JsonCssExtractionStrategy
extraction_strategy = JsonCssExtractionStrategy(schema, verbose=self.verbose)
config = CrawlerRunConfig(
extraction_strategy=extraction_strategy,
js_code=js_code,
wait_for=wait_for,
cache_mode=CacheMode.BYPASS if bypass_cache else CacheMode.ENABLED
)
result = await crawler.arun(url=url, config=config)
if not result.success:
return {"error": f"Extraction failed: {result.error_message}", "success": False}
import json
try:
data = json.loads(result.extracted_content) if result.extracted_content else []
except json.JSONDecodeError:
data = result.extracted_content
return {
"success": True,
"url": result.url,
"data": data,
"count": len(data) if isinstance(data, list) else 1
}
except Exception as e:
self.logger.error(f"Crawl4AI extract_css error: {e}")
return {"error": f"Crawl4AI extract_css error: {str(e)}", "success": False}
async def extract_llm(
self,
url: str,
instruction: str,
schema: Optional[Dict[str, Any]] = None,
provider: str = "openai/gpt-4o-mini",
api_token: Optional[str] = None,
js_code: Optional[Union[str, List[str]]] = None,
wait_for: Optional[str] = None,
bypass_cache: bool = True
) -> Dict[str, Any]:
"""Extract structured data using LLM.
Args:
url: The URL to crawl
instruction: Instructions for the LLM on what to extract
schema: Optional Pydantic-style JSON schema for structured output
provider: LLM provider string (e.g., "openai/gpt-4o-mini", "ollama/llama3.3")
api_token: API token for the provider (uses env var if not provided)
js_code: JavaScript to execute before extraction
wait_for: Wait condition before extraction
bypass_cache: Skip cache
Returns:
Dict with 'data' containing extracted content
"""
try:
crawler = await self._get_crawler()
import os
from crawl4ai import CrawlerRunConfig, CacheMode, LLMConfig
from crawl4ai.extraction_strategy import LLMExtractionStrategy
# Get API token from env if not provided
if api_token is None and "openai" in provider.lower():
api_token = os.environ.get("OPENAI_API_KEY")
llm_config = LLMConfig(provider=provider, api_token=api_token)
extraction_strategy = LLMExtractionStrategy(
llm_config=llm_config,
schema=schema,
extraction_type="schema" if schema else "block",
instruction=instruction,
verbose=self.verbose
)
config = CrawlerRunConfig(
extraction_strategy=extraction_strategy,
js_code=js_code,
wait_for=wait_for,
cache_mode=CacheMode.BYPASS if bypass_cache else CacheMode.ENABLED
)
result = await crawler.arun(url=url, config=config)
if not result.success:
return {"error": f"LLM extraction failed: {result.error_message}", "success": False}
import json
try:
data = json.loads(result.extracted_content) if result.extracted_content else None
except json.JSONDecodeError:
data = result.extracted_content
return {
"success": True,
"url": result.url,
"data": data
}
except Exception as e:
self.logger.error(f"Crawl4AI extract_llm error: {e}")
return {"error": f"Crawl4AI extract_llm error: {str(e)}", "success": False}
# Standalone async functions for easy access
async def crawl4ai(
url: str,
css_selector: Optional[str] = None,
js_code: Optional[str] = None,
wait_for: Optional[str] = None,
screenshot: bool = False
) -> Dict[str, Any]:
"""Crawl a URL and return markdown content.
Args:
url: The URL to crawl
css_selector: CSS selector to focus on specific content
js_code: JavaScript to execute before crawling
wait_for: Wait condition - "css:selector" or "js:() => boolean"
screenshot: Capture screenshot
Returns:
Dict with 'markdown', 'html', 'links', 'media', etc.
Example:
result = await crawl4ai("https://example.com")
print(result["markdown"])
"""
is_available, error = _check_crawl4ai_available()
if not is_available:
logging.error(error)
return {"error": error, "success": False}
tools = Crawl4AITools()
try:
result = await tools.crawl(
url=url,
css_selector=css_selector,
js_code=js_code,
wait_for=wait_for,
screenshot=screenshot
)
return result
finally:
await tools.close()
async def crawl4ai_many(
urls: List[str],
css_selector: Optional[str] = None
) -> List[Dict[str, Any]]:
"""Crawl multiple URLs concurrently.
Args:
urls: List of URLs to crawl
css_selector: CSS selector to focus on specific content
Returns:
List of crawl results
Example:
results = await crawl4ai_many(["https://example.com", "https://example.org"])
"""
is_available, error = _check_crawl4ai_available()
if not is_available:
logging.error(error)
return [{"error": error, "success": False}]
tools = Crawl4AITools()
try:
return await tools.crawl_many(urls=urls, css_selector=css_selector)
finally:
await tools.close()
async def crawl4ai_extract(
url: str,
schema: Dict[str, Any],
js_code: Optional[str] = None,
wait_for: Optional[str] = None
) -> Dict[str, Any]:
"""Extract structured data using CSS selectors.
Args:
url: The URL to crawl
schema: Extraction schema with baseSelector and fields
js_code: JavaScript to execute before extraction
wait_for: Wait condition before extraction
Returns:
Dict with 'data' containing extracted items
Example:
schema = {
"name": "Products",
"baseSelector": "div.product",
"fields": [
{"name": "title", "selector": "h2", "type": "text"},
{"name": "price", "selector": ".price", "type": "text"}
]
}
result = await crawl4ai_extract("https://example.com/products", schema)
"""
is_available, error = _check_crawl4ai_available()
if not is_available:
logging.error(error)
return {"error": error, "success": False}
tools = Crawl4AITools()
try:
return await tools.extract_css(
url=url,
schema=schema,
js_code=js_code,
wait_for=wait_for
)
finally:
await tools.close()
async def crawl4ai_llm_extract(
url: str,
instruction: str,
schema: Optional[Dict[str, Any]] = None,
provider: str = "openai/gpt-4o-mini"
) -> Dict[str, Any]:
"""Extract data using LLM.
Args:
url: The URL to crawl
instruction: Instructions for the LLM
schema: Optional JSON schema for structured output
provider: LLM provider (default "openai/gpt-4o-mini")
Returns:
Dict with 'data' containing extracted content
Example:
result = await crawl4ai_llm_extract(
"https://example.com/about",
instruction="Extract company name, founding year, and description"
)
"""
is_available, error = _check_crawl4ai_available()
if not is_available:
logging.error(error)
return {"error": error, "success": False}
tools = Crawl4AITools()
try:
return await tools.extract_llm(
url=url,
instruction=instruction,
schema=schema,
provider=provider
)
finally:
await tools.close()
# Synchronous wrappers for convenience
def crawl4ai_sync(
url: str,
css_selector: Optional[str] = None,
js_code: Optional[str] = None,
wait_for: Optional[str] = None
) -> Dict[str, Any]:
"""Synchronous version of crawl4ai.
Args:
url: The URL to crawl
css_selector: CSS selector to focus on specific content
js_code: JavaScript to execute
wait_for: Wait condition
Returns:
Dict with crawl results
"""
return asyncio.get_event_loop().run_until_complete(
crawl4ai(url, css_selector, js_code, wait_for)
)
def crawl4ai_extract_sync(
url: str,
schema: Dict[str, Any],
js_code: Optional[str] = None
) -> Dict[str, Any]:
"""Synchronous version of crawl4ai_extract.
Args:
url: The URL to crawl
schema: Extraction schema
js_code: JavaScript to execute
Returns:
Dict with extracted data
"""
return asyncio.get_event_loop().run_until_complete(
crawl4ai_extract(url, schema, js_code)
)
if __name__ == "__main__":
# Example usage
async def main():
print("=" * 60)
print("Crawl4AI Tools - Example Usage")
print("=" * 60)
is_available, error = _check_crawl4ai_available()
if not is_available:
print(f"\nCrawl4AI is not available: {error}")
print("\nTo use Crawl4AI tools:")
print("1. Install the package: pip install crawl4ai")
print("2. Run setup: crawl4ai-setup")
else:
print("Crawl4AI is available!")
# Example crawl
print("\n1. Basic Crawl")
print("-" * 30)
result = await crawl4ai("https://example.com")
if result.get("success"):
print(f"URL: {result['url']}")
print(f"Markdown length: {len(result.get('markdown', ''))}")
print(f"Links found: {len(result.get('links', {}).get('internal', []))}")
else:
print(f"Error: {result.get('error')}")
print("\n" + "=" * 60)
print("Demonstration Complete")
print("=" * 60)
asyncio.run(main())
@@ -0,0 +1,225 @@
"""Tool decorator for converting functions into tools.
This module provides the @tool decorator for easily creating tools from functions.
Usage:
from praisonaiagents import tool
@tool
def search(query: str) -> list:
'''Search the web for information.'''
return [...]
# Or with explicit parameters:
@tool(name="web_search", description="Search the internet")
def search(query: str, max_results: int = 5) -> list:
return [...]
"""
import inspect
import functools
import logging
from typing import Any, Callable, Dict, Optional, Union, get_type_hints
from .base import BaseTool
class FunctionTool(BaseTool):
"""A BaseTool wrapper for plain functions.
Created automatically by the @tool decorator.
"""
def __init__(
self,
func: Callable,
name: Optional[str] = None,
description: Optional[str] = None,
version: str = "1.0.0"
):
self._func = func
self.name = name or func.__name__
self.description = description or func.__doc__ or f"Tool: {self.name}"
self.version = version
# Generate schema from the original function, not run()
self.parameters = self._generate_schema_from_func(func)
# Copy function metadata
functools.update_wrapper(self, func)
# Skip parent's schema generation since we already did it
# Just set defaults that parent would set
if not self.name:
self.name = self.__class__.__name__.lower().replace("tool", "")
if not self.description:
self.description = self.__class__.__doc__ or f"Tool: {self.name}"
def _generate_schema_from_func(self, func: Callable) -> Dict[str, Any]:
"""Generate JSON Schema from the wrapped function's signature."""
schema = {
"type": "object",
"properties": {},
"required": []
}
try:
sig = inspect.signature(func)
hints = get_type_hints(func) if hasattr(func, '__annotations__') else {}
for param_name, param in sig.parameters.items():
if param_name in ('self', 'cls'):
continue
# Get type hint
param_type = hints.get(param_name, Any)
json_type = BaseTool._python_type_to_json(param_type)
schema["properties"][param_name] = {"type": json_type}
# Check if required (no default value)
if param.default is inspect.Parameter.empty:
schema["required"].append(param_name)
except Exception as e:
logging.debug(f"Could not generate schema for {func.__name__}: {e}")
return schema
def run(self, **kwargs) -> Any:
"""Execute the wrapped function."""
return self._func(**kwargs)
def __call__(self, *args, **kwargs) -> Any:
"""Allow calling with positional args like the original function."""
return self._func(*args, **kwargs)
def tool(
func: Optional[Callable] = None,
*,
name: Optional[str] = None,
description: Optional[str] = None,
version: str = "1.0.0"
) -> Union[FunctionTool, Callable[[Callable], FunctionTool]]:
"""Decorator to convert a function into a tool.
Can be used with or without arguments:
@tool
def my_func(x: str) -> str:
'''Does something.'''
return x
@tool(name="custom_name", description="Custom description")
def my_func(x: str) -> str:
return x
Args:
func: The function to wrap (when used without parentheses)
name: Override the tool name (default: function name)
description: Override description (default: function docstring)
version: Tool version (default: "1.0.0")
Returns:
FunctionTool instance that wraps the function
"""
def decorator(fn: Callable) -> FunctionTool:
tool_instance = FunctionTool(
func=fn,
name=name,
description=description,
version=version
)
# Register with global registry if available
try:
from .registry import get_registry
registry = get_registry()
if registry:
registry.register(tool_instance)
except ImportError:
pass # Registry not yet available
return tool_instance
# Handle both @tool and @tool(...) syntax
if func is not None:
# Called without parentheses: @tool
return decorator(func)
else:
# Called with parentheses: @tool(...)
return decorator
def is_tool(obj: Any) -> bool:
"""Check if an object is a tool (BaseTool instance or decorated function)."""
if isinstance(obj, BaseTool):
return True
if isinstance(obj, FunctionTool):
return True
# Check for tools created by other frameworks
if hasattr(obj, 'run') and hasattr(obj, 'name'):
return True
return False
def get_tool_schema(obj: Any) -> Optional[Dict[str, Any]]:
"""Get OpenAI-compatible schema for any tool-like object.
Supports:
- BaseTool instances
- FunctionTool instances
- Plain functions (generates schema from signature)
- LangChain tools
- CrewAI tools
"""
# BaseTool or FunctionTool
if isinstance(obj, BaseTool):
return obj.get_schema()
# Plain callable
if callable(obj):
return _schema_from_function(obj)
return None
def _schema_from_function(func: Callable) -> Dict[str, Any]:
"""Generate OpenAI function schema from a plain function."""
name = getattr(func, '__name__', 'unknown')
description = func.__doc__ or f"Function: {name}"
# Build parameters schema
properties = {}
required = []
try:
sig = inspect.signature(func)
hints = get_type_hints(func) if hasattr(func, '__annotations__') else {}
for param_name, param in sig.parameters.items():
if param_name == 'self':
continue
param_type = hints.get(param_name, Any)
json_type = BaseTool._python_type_to_json(param_type)
properties[param_name] = {"type": json_type}
if param.default is inspect.Parameter.empty:
required.append(param_name)
except Exception as e:
logging.debug(f"Could not generate schema for {name}: {e}")
return {
"type": "function",
"function": {
"name": name,
"description": description.strip(),
"parameters": {
"type": "object",
"properties": properties,
"required": required
}
}
}
@@ -0,0 +1,741 @@
"""Exa search and research tools.
Exa provides AI-powered web search, content retrieval, and research capabilities.
This module requires the EXA_API_KEY environment variable to be set.
Usage:
from praisonaiagents.tools import exa_search, exa_search_contents
# Search the web
results = exa_search("AI startups")
# Search with contents
results = exa_search_contents("AI in healthcare", text=True)
# Or use the class directly
from praisonaiagents.tools import ExaTools
exa = ExaTools()
results = exa.search("AI news")
"""
from typing import List, Dict, Any, Optional, Union
import logging
import os
from importlib import util
def _check_exa_available() -> tuple[bool, Optional[str]]:
"""Check if Exa is available and API key is set.
Returns:
Tuple of (is_available, error_message)
"""
# Check if exa_py package is installed
if util.find_spec("exa_py") is None:
return False, "exa_py package is not installed. Install it with: pip install exa_py"
# Check if API key is set
if not os.environ.get("EXA_API_KEY"):
return False, "EXA_API_KEY environment variable is not set. Please set it to use Exa tools."
return True, None
class ExaTools:
"""Exa search and research tools.
Provides methods for:
- search: Basic web search
- search_and_contents: Search with full text/highlights
- find_similar: Find similar pages to a URL
- find_similar_and_contents: Find similar with content
- answer: Get AI-generated answers with citations
Example:
from praisonaiagents.tools import ExaTools
exa = ExaTools()
results = exa.search("AI startups", num_results=5)
"""
def __init__(self, api_key: Optional[str] = None):
"""Initialize ExaTools.
Args:
api_key: Exa API key. If not provided, uses EXA_API_KEY env var.
"""
self.api_key = api_key or os.environ.get("EXA_API_KEY")
self._client = None
self.logger = logging.getLogger(__name__)
def _get_client(self):
"""Get or create Exa client."""
if self._client is None:
is_available, error = _check_exa_available()
if not is_available:
raise ImportError(error)
from exa_py import Exa
self._client = Exa(self.api_key)
return self._client
def search(
self,
query: str,
num_results: int = 10,
type: Optional[str] = None,
category: Optional[str] = None,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
start_crawl_date: Optional[str] = None,
end_crawl_date: Optional[str] = None,
start_published_date: Optional[str] = None,
end_published_date: Optional[str] = None,
include_text: Optional[List[str]] = None,
exclude_text: Optional[List[str]] = None,
additional_queries: Optional[List[str]] = None
) -> Dict[str, Any]:
"""Perform an Exa search.
Args:
query: The search query
num_results: Number of results to return (default 10, max 100)
type: Search type - "auto", "neural", "fast", or "deep"
category: Data category - company, research paper, news, linkedin profile,
github, tweet, movie, song, personal site, pdf, financial report
include_domains: List of domains to include
exclude_domains: List of domains to exclude
start_crawl_date: Only include links crawled after this date
end_crawl_date: Only include links crawled before this date
start_published_date: Only include links published after this date
end_published_date: Only include links published before this date
include_text: Strings that must be present in results
exclude_text: Strings that must not be present in results
additional_queries: Additional query variations for deep search
Returns:
Dict with 'results' list and optional 'autopromptString'
"""
try:
client = self._get_client()
kwargs = {
"query": query,
"num_results": num_results
}
if type:
kwargs["type"] = type
if category:
kwargs["category"] = category
if include_domains:
kwargs["include_domains"] = include_domains
if exclude_domains:
kwargs["exclude_domains"] = exclude_domains
if start_crawl_date:
kwargs["start_crawl_date"] = start_crawl_date
if end_crawl_date:
kwargs["end_crawl_date"] = end_crawl_date
if start_published_date:
kwargs["start_published_date"] = start_published_date
if end_published_date:
kwargs["end_published_date"] = end_published_date
if include_text:
kwargs["include_text"] = include_text
if exclude_text:
kwargs["exclude_text"] = exclude_text
if additional_queries:
kwargs["additional_queries"] = additional_queries
response = client.search(**kwargs)
return {
"results": [
{
"url": r.url,
"id": r.id,
"title": getattr(r, "title", None),
"published_date": getattr(r, "published_date", None),
"author": getattr(r, "author", None)
}
for r in response.results
],
"autopromptString": getattr(response, "autoprompt_string", None),
"requestId": getattr(response, "request_id", None)
}
except Exception as e:
self.logger.error(f"Exa search error: {e}")
return {"error": f"Exa search error: {str(e)}"}
def search_and_contents(
self,
query: str,
text: Union[bool, Dict] = True,
highlights: Union[bool, Dict] = False,
summary: Optional[Dict] = None,
num_results: int = 10,
type: Optional[str] = None,
category: Optional[str] = None,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
start_crawl_date: Optional[str] = None,
end_crawl_date: Optional[str] = None,
start_published_date: Optional[str] = None,
end_published_date: Optional[str] = None,
include_text: Optional[List[str]] = None,
exclude_text: Optional[List[str]] = None
) -> Dict[str, Any]:
"""Search and retrieve content from results.
Args:
query: The search query
text: Include full text content (True or TextContentsOptions dict)
highlights: Include highlights (True or HighlightsContentsOptions dict)
summary: Summary options with optional schema for structured output
num_results: Number of results to return
type: Search type - "auto", "neural", "fast", or "deep"
category: Data category filter
include_domains: List of domains to include
exclude_domains: List of domains to exclude
start_crawl_date: Only include links crawled after this date
end_crawl_date: Only include links crawled before this date
start_published_date: Only include links published after this date
end_published_date: Only include links published before this date
include_text: Strings that must be present in results
exclude_text: Strings that must not be present in results
Returns:
Dict with 'results' list containing text/highlights
"""
try:
client = self._get_client()
kwargs = {
"query": query,
"num_results": num_results
}
if text:
kwargs["text"] = text
if highlights:
kwargs["highlights"] = highlights
if summary:
kwargs["summary"] = summary
if type:
kwargs["type"] = type
if category:
kwargs["category"] = category
if include_domains:
kwargs["include_domains"] = include_domains
if exclude_domains:
kwargs["exclude_domains"] = exclude_domains
if start_crawl_date:
kwargs["start_crawl_date"] = start_crawl_date
if end_crawl_date:
kwargs["end_crawl_date"] = end_crawl_date
if start_published_date:
kwargs["start_published_date"] = start_published_date
if end_published_date:
kwargs["end_published_date"] = end_published_date
if include_text:
kwargs["include_text"] = include_text
if exclude_text:
kwargs["exclude_text"] = exclude_text
response = client.search_and_contents(**kwargs)
results = []
for r in response.results:
result = {
"url": r.url,
"id": r.id,
"title": getattr(r, "title", None),
"published_date": getattr(r, "published_date", None),
"author": getattr(r, "author", None)
}
if hasattr(r, "text") and r.text:
result["text"] = r.text
if hasattr(r, "highlights") and r.highlights:
result["highlights"] = r.highlights
if hasattr(r, "highlight_scores") and r.highlight_scores:
result["highlight_scores"] = r.highlight_scores
if hasattr(r, "summary") and r.summary:
result["summary"] = r.summary
results.append(result)
return {
"results": results,
"requestId": getattr(response, "request_id", None)
}
except Exception as e:
self.logger.error(f"Exa search_and_contents error: {e}")
return {"error": f"Exa search_and_contents error: {str(e)}"}
def find_similar(
self,
url: str,
num_results: int = 10,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
start_crawl_date: Optional[str] = None,
end_crawl_date: Optional[str] = None,
start_published_date: Optional[str] = None,
end_published_date: Optional[str] = None,
exclude_source_domain: bool = False,
category: Optional[str] = None
) -> Dict[str, Any]:
"""Find similar pages to a given URL.
Args:
url: The URL to find similar pages for
num_results: Number of results to return
include_domains: List of domains to include
exclude_domains: List of domains to exclude
start_crawl_date: Only include links crawled after this date
end_crawl_date: Only include links crawled before this date
start_published_date: Only include links published after this date
end_published_date: Only include links published before this date
exclude_source_domain: Exclude results from the same domain
category: Data category filter
Returns:
Dict with 'results' list of similar pages
"""
try:
client = self._get_client()
kwargs = {
"url": url,
"num_results": num_results
}
if include_domains:
kwargs["include_domains"] = include_domains
if exclude_domains:
kwargs["exclude_domains"] = exclude_domains
if start_crawl_date:
kwargs["start_crawl_date"] = start_crawl_date
if end_crawl_date:
kwargs["end_crawl_date"] = end_crawl_date
if start_published_date:
kwargs["start_published_date"] = start_published_date
if end_published_date:
kwargs["end_published_date"] = end_published_date
if exclude_source_domain:
kwargs["exclude_source_domain"] = exclude_source_domain
if category:
kwargs["category"] = category
response = client.find_similar(**kwargs)
return {
"results": [
{
"url": r.url,
"id": r.id,
"title": getattr(r, "title", None),
"published_date": getattr(r, "published_date", None),
"author": getattr(r, "author", None)
}
for r in response.results
],
"requestId": getattr(response, "request_id", None)
}
except Exception as e:
self.logger.error(f"Exa find_similar error: {e}")
return {"error": f"Exa find_similar error: {str(e)}"}
def find_similar_and_contents(
self,
url: str,
text: Union[bool, Dict] = True,
highlights: Union[bool, Dict] = False,
num_results: int = 10,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
start_crawl_date: Optional[str] = None,
end_crawl_date: Optional[str] = None,
start_published_date: Optional[str] = None,
end_published_date: Optional[str] = None,
exclude_source_domain: bool = False,
category: Optional[str] = None
) -> Dict[str, Any]:
"""Find similar pages with content.
Args:
url: The URL to find similar pages for
text: Include full text content
highlights: Include highlights
num_results: Number of results to return
include_domains: List of domains to include
exclude_domains: List of domains to exclude
start_crawl_date: Only include links crawled after this date
end_crawl_date: Only include links crawled before this date
start_published_date: Only include links published after this date
end_published_date: Only include links published before this date
exclude_source_domain: Exclude results from the same domain
category: Data category filter
Returns:
Dict with 'results' list containing text/highlights
"""
try:
client = self._get_client()
kwargs = {
"url": url,
"num_results": num_results
}
if text:
kwargs["text"] = text
if highlights:
kwargs["highlights"] = highlights
if include_domains:
kwargs["include_domains"] = include_domains
if exclude_domains:
kwargs["exclude_domains"] = exclude_domains
if start_crawl_date:
kwargs["start_crawl_date"] = start_crawl_date
if end_crawl_date:
kwargs["end_crawl_date"] = end_crawl_date
if start_published_date:
kwargs["start_published_date"] = start_published_date
if end_published_date:
kwargs["end_published_date"] = end_published_date
if exclude_source_domain:
kwargs["exclude_source_domain"] = exclude_source_domain
if category:
kwargs["category"] = category
response = client.find_similar_and_contents(**kwargs)
results = []
for r in response.results:
result = {
"url": r.url,
"id": r.id,
"title": getattr(r, "title", None),
"published_date": getattr(r, "published_date", None),
"author": getattr(r, "author", None)
}
if hasattr(r, "text") and r.text:
result["text"] = r.text
if hasattr(r, "highlights") and r.highlights:
result["highlights"] = r.highlights
if hasattr(r, "highlight_scores") and r.highlight_scores:
result["highlight_scores"] = r.highlight_scores
results.append(result)
return {
"results": results,
"requestId": getattr(response, "request_id", None)
}
except Exception as e:
self.logger.error(f"Exa find_similar_and_contents error: {e}")
return {"error": f"Exa find_similar_and_contents error: {str(e)}"}
def answer(
self,
query: str,
text: bool = False
) -> Dict[str, Any]:
"""Generate an answer using Exa's search and LLM capabilities.
Args:
query: The question to answer
text: If True, include full text of citations
Returns:
Dict with 'answer' and 'citations'
"""
try:
client = self._get_client()
response = client.answer(query=query, text=text)
citations = []
for c in response.citations:
citation = {
"id": c.id,
"url": c.url,
"title": getattr(c, "title", None),
"published_date": getattr(c, "published_date", None),
"author": getattr(c, "author", None)
}
if text and hasattr(c, "text") and c.text:
citation["text"] = c.text
citations.append(citation)
return {
"answer": response.answer,
"citations": citations
}
except Exception as e:
self.logger.error(f"Exa answer error: {e}")
return {"error": f"Exa answer error: {str(e)}"}
# Standalone functions for easy access
def exa_search(
query: str,
num_results: int = 10,
type: Optional[str] = None,
category: Optional[str] = None,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
start_published_date: Optional[str] = None,
end_published_date: Optional[str] = None
) -> Dict[str, Any]:
"""Search the web using Exa.
Args:
query: The search query
num_results: Number of results (default 10)
type: Search type - "auto", "neural", "fast", or "deep"
category: Data category filter
include_domains: Domains to include
exclude_domains: Domains to exclude
start_published_date: Only include links published after this date
end_published_date: Only include links published before this date
Returns:
Dict with search results or error
Example:
results = exa_search("AI startups", num_results=5)
"""
is_available, error = _check_exa_available()
if not is_available:
logging.error(error)
return {"error": error}
tools = ExaTools()
return tools.search(
query=query,
num_results=num_results,
type=type,
category=category,
include_domains=include_domains,
exclude_domains=exclude_domains,
start_published_date=start_published_date,
end_published_date=end_published_date
)
def exa_search_contents(
query: str,
text: bool = True,
highlights: bool = False,
num_results: int = 10,
type: Optional[str] = None,
category: Optional[str] = None,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None
) -> Dict[str, Any]:
"""Search and get content from results.
Args:
query: The search query
text: Include full text content
highlights: Include highlights
num_results: Number of results
type: Search type
category: Data category filter
include_domains: Domains to include
exclude_domains: Domains to exclude
Returns:
Dict with results containing text/highlights
Example:
results = exa_search_contents("AI in healthcare", text=True, highlights=True)
"""
is_available, error = _check_exa_available()
if not is_available:
logging.error(error)
return {"error": error}
tools = ExaTools()
return tools.search_and_contents(
query=query,
text=text,
highlights=highlights,
num_results=num_results,
type=type,
category=category,
include_domains=include_domains,
exclude_domains=exclude_domains
)
def exa_find_similar(
url: str,
num_results: int = 10,
exclude_source_domain: bool = True,
category: Optional[str] = None
) -> Dict[str, Any]:
"""Find pages similar to a given URL.
Args:
url: The URL to find similar pages for
num_results: Number of results
exclude_source_domain: Exclude same domain results
category: Data category filter
Returns:
Dict with similar results
Example:
results = exa_find_similar("https://openai.com", num_results=5)
"""
is_available, error = _check_exa_available()
if not is_available:
logging.error(error)
return {"error": error}
tools = ExaTools()
return tools.find_similar(
url=url,
num_results=num_results,
exclude_source_domain=exclude_source_domain,
category=category
)
def exa_answer(
query: str,
text: bool = False
) -> Dict[str, Any]:
"""Get an AI-generated answer with citations.
Args:
query: The question to answer
text: Include full text of citations
Returns:
Dict with 'answer' and 'citations'
Example:
result = exa_answer("What is the capital of France?")
print(result["answer"])
"""
is_available, error = _check_exa_available()
if not is_available:
logging.error(error)
return {"error": error}
tools = ExaTools()
return tools.answer(query=query, text=text)
# Async versions
async def exa_search_async(
query: str,
num_results: int = 10,
type: Optional[str] = None,
category: Optional[str] = None
) -> Dict[str, Any]:
"""Async version of exa_search.
Args:
query: The search query
num_results: Number of results
type: Search type
category: Data category filter
Returns:
Dict with search results
"""
import asyncio
return await asyncio.get_event_loop().run_in_executor(
None,
lambda: exa_search(query, num_results, type, category)
)
async def exa_search_contents_async(
query: str,
text: bool = True,
highlights: bool = False,
num_results: int = 10
) -> Dict[str, Any]:
"""Async version of exa_search_contents.
Args:
query: The search query
text: Include full text
highlights: Include highlights
num_results: Number of results
Returns:
Dict with results containing content
"""
import asyncio
return await asyncio.get_event_loop().run_in_executor(
None,
lambda: exa_search_contents(query, text, highlights, num_results)
)
async def exa_answer_async(
query: str,
text: bool = False
) -> Dict[str, Any]:
"""Async version of exa_answer.
Args:
query: The question to answer
text: Include citation text
Returns:
Dict with answer and citations
"""
import asyncio
return await asyncio.get_event_loop().run_in_executor(
None,
lambda: exa_answer(query, text)
)
# Alias for simple usage
exa = exa_search
if __name__ == "__main__":
# Example usage
print("=" * 60)
print("Exa Tools - Example Usage")
print("=" * 60)
is_available, error = _check_exa_available()
if not is_available:
print(f"\nExa is not available: {error}")
print("\nTo use Exa tools:")
print("1. Install the package: pip install exa_py")
print("2. Set environment variable: export EXA_API_KEY=your_api_key")
else:
print("Exa SDK is available!")
# Example search
print("\n1. Basic Search")
print("-" * 30)
results = exa_search("Latest AI developments 2024", num_results=3)
if "error" not in results:
print(f"Found {len(results.get('results', []))} results")
for r in results.get("results", [])[:3]:
print(f" - {r.get('title', 'No title')[:50]}...")
print(f" URL: {r.get('url', 'No URL')}")
else:
print(f"Error: {results['error']}")
print("\n" + "=" * 60)
print("Demonstration Complete")
print("=" * 60)
@@ -0,0 +1,244 @@
"""Example plugin tools demonstrating the BaseTool and @tool patterns.
This module shows how to create tools using the new plugin system.
External developers can use these patterns to create pip-installable plugins.
Usage:
from praisonaiagents import BaseTool, tool, Agent
# Method 1: Class-based tool
class MyTool(BaseTool):
name = "my_tool"
description = "Does something useful"
def run(self, param: str) -> str:
return f"Result: {param}"
# Method 2: Decorator-based tool
@tool
def my_function(param: str) -> str:
'''Does something useful.'''
return f"Result: {param}"
# Use with Agent
agent = Agent(tools=[MyTool(), my_function])
"""
from typing import Dict, Any, List
from .base import BaseTool
from .decorator import tool
# =============================================================================
# Example 1: Simple Class-Based Tool
# =============================================================================
class EchoTool(BaseTool):
"""A simple tool that echoes input back.
This demonstrates the minimal BaseTool implementation.
"""
name = "echo"
description = "Echo the input message back"
def run(self, message: str) -> str:
"""Echo the message.
Args:
message: The message to echo
Returns:
The same message
"""
return message
# =============================================================================
# Example 2: Tool with Multiple Parameters
# =============================================================================
class MathTool(BaseTool):
"""A tool for basic math operations.
Demonstrates:
- Multiple parameters with types
- Default parameter values
- Returning structured data
"""
name = "basic_math"
description = "Perform basic math operations (add, subtract, multiply, divide)"
def run(
self,
a: float,
b: float,
operation: str = "add"
) -> Dict[str, Any]:
"""Perform a math operation.
Args:
a: First number
b: Second number
operation: One of 'add', 'subtract', 'multiply', 'divide'
Returns:
Dict with result and operation details
"""
operations = {
"add": lambda x, y: x + y,
"subtract": lambda x, y: x - y,
"multiply": lambda x, y: x * y,
"divide": lambda x, y: x / y if y != 0 else float('inf')
}
if operation not in operations:
return {"error": f"Unknown operation: {operation}"}
result = operations[operation](a, b)
return {
"result": result,
"operation": operation,
"expression": f"{a} {operation} {b} = {result}"
}
# =============================================================================
# Example 3: Tool with State (Instance Variables)
# =============================================================================
class CounterTool(BaseTool):
"""A tool that maintains state between calls.
Demonstrates:
- Instance variables for state
- Custom __init__ with super() call
"""
name = "counter"
description = "A counter that increments each time it's called"
def __init__(self, start: int = 0):
self._count = start
super().__init__() # Important: call parent init
def run(self, increment: int = 1) -> Dict[str, int]:
"""Increment the counter.
Args:
increment: Amount to increment by (default: 1)
Returns:
Dict with previous and current count
"""
previous = self._count
self._count += increment
return {
"previous": previous,
"current": self._count,
"increment": increment
}
# =============================================================================
# Example 4: Decorator-Based Tools
# =============================================================================
@tool
def greet(name: str, greeting: str = "Hello") -> str:
"""Generate a greeting message.
Args:
name: Name of the person to greet
greeting: The greeting word (default: "Hello")
Returns:
A greeting string
"""
return f"{greeting}, {name}!"
@tool(name="reverse_text", description="Reverse a string of text")
def reverse_string(text: str) -> str:
"""Reverse the input text."""
return text[::-1]
@tool
def list_numbers(start: int, end: int, step: int = 1) -> List[int]:
"""Generate a list of numbers.
Args:
start: Starting number (inclusive)
end: Ending number (exclusive)
step: Step between numbers (default: 1)
Returns:
List of integers
"""
return list(range(start, end, step))
# =============================================================================
# Example 5: Tool with Error Handling
# =============================================================================
class SafeDivideTool(BaseTool):
"""A division tool with proper error handling.
Demonstrates:
- Error handling in run()
- Returning error information
"""
name = "safe_divide"
description = "Safely divide two numbers with error handling"
def run(self, numerator: float, denominator: float) -> Dict[str, Any]:
"""Divide numerator by denominator.
Args:
numerator: The number to divide
denominator: The number to divide by
Returns:
Dict with result or error message
"""
if denominator == 0:
return {
"success": False,
"error": "Division by zero is not allowed"
}
result = numerator / denominator
return {
"success": True,
"result": result
}
# =============================================================================
# Convenience exports
# =============================================================================
# These can be imported directly:
# from praisonaiagents.tools.example_plugin import echo_tool, math_tool, greet
echo_tool = EchoTool()
math_tool = MathTool()
counter_tool = CounterTool()
safe_divide_tool = SafeDivideTool()
__all__ = [
# Classes
'EchoTool',
'MathTool',
'CounterTool',
'SafeDivideTool',
# Instances
'echo_tool',
'math_tool',
'counter_tool',
'safe_divide_tool',
# Decorated functions
'greet',
'reverse_string',
'list_numbers',
]
@@ -0,0 +1,225 @@
"""Tool registry for managing and discovering tools.
This module provides centralized tool management with support for:
- Manual registration
- Auto-discovery via entry_points
- Tool lookup by name
Usage:
from praisonaiagents.tools.registry import get_registry
registry = get_registry()
registry.register(my_tool)
tool = registry.get("my_tool")
"""
import logging
from typing import Callable, Dict, List, Optional, Union
from importlib.metadata import entry_points
from .base import BaseTool
# Entry point group name for external plugins
ENTRY_POINT_GROUP = "praisonaiagents.tools"
class ToolRegistry:
"""Central registry for all tools.
Provides:
- Tool registration (manual and auto-discovery)
- Tool lookup by name
- Tool listing and filtering
- Entry points discovery for external plugins
"""
def __init__(self):
self._tools: Dict[str, BaseTool] = {}
self._functions: Dict[str, Callable] = {} # For backward compat with plain functions
self._discovered: bool = False
def register(
self,
tool: Union[BaseTool, Callable],
name: Optional[str] = None,
overwrite: bool = False
) -> None:
"""Register a tool with the registry.
Args:
tool: BaseTool instance or callable function
name: Override name (default: tool.name or function.__name__)
overwrite: If True, overwrite existing tool with same name
Raises:
ValueError: If tool with same name exists and overwrite=False
"""
# Handle BaseTool instances
if isinstance(tool, BaseTool):
tool_name = name or tool.name
if tool_name in self._tools and not overwrite:
logging.debug(f"Tool '{tool_name}' already registered, skipping")
return
self._tools[tool_name] = tool
logging.debug(f"Registered tool: {tool_name}")
return
# Handle plain callables
if callable(tool):
tool_name = name or getattr(tool, '__name__', str(id(tool)))
if tool_name in self._functions and not overwrite:
logging.debug(f"Function '{tool_name}' already registered, skipping")
return
self._functions[tool_name] = tool
logging.debug(f"Registered function: {tool_name}")
return
raise TypeError(f"Cannot register {type(tool)}, expected BaseTool or callable")
def unregister(self, name: str) -> bool:
"""Remove a tool from the registry.
Args:
name: Tool name to remove
Returns:
True if tool was removed, False if not found
"""
if name in self._tools:
del self._tools[name]
return True
if name in self._functions:
del self._functions[name]
return True
return False
def get(self, name: str) -> Optional[Union[BaseTool, Callable]]:
"""Get a tool by name.
Args:
name: Tool name
Returns:
BaseTool instance, callable, or None if not found
"""
# Check BaseTool registry first
if name in self._tools:
return self._tools[name]
# Check functions registry
if name in self._functions:
return self._functions[name]
# Try auto-discovery if not found
if not self._discovered:
self.discover_plugins()
if name in self._tools:
return self._tools[name]
return None
def list_tools(self) -> List[str]:
"""List all registered tool names."""
return list(self._tools.keys()) + list(self._functions.keys())
def list_base_tools(self) -> List[BaseTool]:
"""List all registered BaseTool instances."""
return list(self._tools.values())
def get_all(self) -> Dict[str, Union[BaseTool, Callable]]:
"""Get all registered tools as a dict."""
result = dict(self._tools)
result.update(self._functions)
return result
def discover_plugins(self) -> int:
"""Discover and register tools from entry_points.
External packages can register tools by adding to pyproject.toml:
[project.entry-points."praisonaiagents.tools"]
my_tool = "my_package.tools:MyTool"
Returns:
Number of tools discovered
"""
if self._discovered:
return 0
count = 0
try:
# Python 3.10+ style
eps = entry_points(group=ENTRY_POINT_GROUP)
except TypeError:
# Python 3.9 fallback
try:
all_eps = entry_points()
eps = all_eps.get(ENTRY_POINT_GROUP, [])
except Exception:
eps = []
for ep in eps:
try:
tool_class_or_func = ep.load()
# If it's a class, instantiate it
if isinstance(tool_class_or_func, type) and issubclass(tool_class_or_func, BaseTool):
tool_instance = tool_class_or_func()
self.register(tool_instance, name=ep.name)
# If it's already an instance or callable
elif isinstance(tool_class_or_func, BaseTool):
self.register(tool_class_or_func, name=ep.name)
elif callable(tool_class_or_func):
self.register(tool_class_or_func, name=ep.name)
else:
logging.warning(f"Entry point '{ep.name}' is not a valid tool")
continue
count += 1
logging.info(f"Discovered plugin tool: {ep.name}")
except Exception as e:
logging.warning(f"Failed to load plugin '{ep.name}': {e}")
self._discovered = True
return count
def clear(self) -> None:
"""Clear all registered tools."""
self._tools.clear()
self._functions.clear()
self._discovered = False
def __contains__(self, name: str) -> bool:
return name in self._tools or name in self._functions
def __len__(self) -> int:
return len(self._tools) + len(self._functions)
def __repr__(self) -> str:
return f"ToolRegistry(tools={len(self._tools)}, functions={len(self._functions)})"
# Global registry instance
_global_registry: Optional[ToolRegistry] = None
def get_registry() -> ToolRegistry:
"""Get the global tool registry instance."""
global _global_registry
if _global_registry is None:
_global_registry = ToolRegistry()
return _global_registry
def register_tool(
tool: Union[BaseTool, Callable],
name: Optional[str] = None
) -> None:
"""Convenience function to register a tool with the global registry."""
get_registry().register(tool, name=name)
def get_tool(name: str) -> Optional[Union[BaseTool, Callable]]:
"""Convenience function to get a tool from the global registry."""
return get_registry().get(name)
@@ -0,0 +1,591 @@
"""Tavily search and extraction tools.
Tavily provides AI-powered web search, content extraction, and crawling capabilities.
This module requires the TAVILY_API_KEY environment variable to be set.
Usage:
from praisonaiagents.tools import tavily_search, tavily_extract
# Search the web
results = tavily_search("AI news 2024")
# Extract content from URLs
content = tavily_extract("https://example.com")
# Or use the class directly
from praisonaiagents.tools import TavilyTools
tavily = TavilyTools()
results = tavily.search("AI news")
"""
from typing import List, Dict, Any, Optional, Union
import logging
import os
from importlib import util
def _check_tavily_available() -> tuple[bool, Optional[str]]:
"""Check if Tavily is available and API key is set.
Returns:
Tuple of (is_available, error_message)
"""
# Check if tavily package is installed
if util.find_spec("tavily") is None:
return False, "Tavily package is not installed. Please install it using: pip install tavily-python"
# Check if API key is set
api_key = os.environ.get("TAVILY_API_KEY")
if not api_key:
return False, "TAVILY_API_KEY environment variable is not set. Please set it to use Tavily tools."
return True, None
class TavilyTools:
"""Comprehensive tools for web search and content extraction using Tavily API.
Tavily provides AI-powered search capabilities optimized for LLM applications.
Features:
- Web search with advanced filtering and ranking
- Content extraction from URLs
- Website crawling with intelligent navigation
- Site mapping
Requires:
- tavily-python package: pip install tavily-python
- TAVILY_API_KEY environment variable
"""
def __init__(self, api_key: Optional[str] = None):
"""Initialize TavilyTools.
Args:
api_key: Optional API key. If not provided, uses TAVILY_API_KEY env var.
"""
self._api_key = api_key
self._client = None
self._async_client = None
def _get_api_key(self) -> Optional[str]:
"""Get API key from instance or environment."""
return self._api_key or os.environ.get("TAVILY_API_KEY")
def _get_client(self):
"""Get or create synchronous Tavily client."""
if self._client is None:
is_available, error = _check_tavily_available()
if not is_available:
return None, error
from tavily import TavilyClient
api_key = self._get_api_key()
if not api_key:
return None, "TAVILY_API_KEY not set"
self._client = TavilyClient(api_key=api_key)
return self._client, None
def _get_async_client(self):
"""Get or create asynchronous Tavily client."""
if self._async_client is None:
is_available, error = _check_tavily_available()
if not is_available:
return None, error
from tavily import AsyncTavilyClient
api_key = self._get_api_key()
if not api_key:
return None, "TAVILY_API_KEY not set"
self._async_client = AsyncTavilyClient(api_key=api_key)
return self._async_client, None
def search(
self,
query: str,
search_depth: str = "basic",
topic: str = "general",
max_results: int = 5,
include_answer: bool = False,
include_raw_content: bool = False,
include_images: bool = False,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
time_range: Optional[str] = None,
) -> Dict[str, Any]:
"""Search the web using Tavily's AI-powered search.
Args:
query: The search query
search_depth: "basic" or "advanced" (advanced uses more credits but better results)
topic: "general", "news", or "finance"
max_results: Maximum number of results (1-20)
include_answer: Include LLM-generated answer based on results
include_raw_content: Include cleaned HTML content of each result
include_images: Include query-related images
include_domains: List of domains to include (max 300)
exclude_domains: List of domains to exclude (max 150)
time_range: Filter by time - "day", "week", "month", "year" or "d", "w", "m", "y"
Returns:
Dict containing search results with keys:
- results: List of search results with title, url, content, score
- query: The search query
- response_time: Time taken for search
- answer: (optional) LLM-generated answer if include_answer=True
- images: (optional) List of image URLs if include_images=True
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
kwargs = {
"query": query,
"search_depth": search_depth,
"topic": topic,
"max_results": max_results,
"include_answer": include_answer,
"include_raw_content": include_raw_content,
"include_images": include_images,
}
if include_domains:
kwargs["include_domains"] = include_domains
if exclude_domains:
kwargs["exclude_domains"] = exclude_domains
if time_range:
kwargs["time_range"] = time_range
response = client.search(**kwargs)
return response
except Exception as e:
error_msg = f"Tavily search error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
def extract(
self,
urls: Union[str, List[str]],
include_images: bool = False,
extract_depth: str = "basic",
) -> Dict[str, Any]:
"""Extract content from one or more URLs.
Args:
urls: Single URL or list of URLs (max 20)
include_images: Include extracted images
extract_depth: "basic" or "advanced" (advanced retrieves more data)
Returns:
Dict containing:
- results: List of successful extractions with url, raw_content
- failed_results: List of URLs that failed with error messages
- response_time: Time taken for extraction
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
response = client.extract(
urls=urls,
include_images=include_images,
extract_depth=extract_depth,
)
return response
except Exception as e:
error_msg = f"Tavily extract error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
def crawl(
self,
url: str,
max_depth: int = 1,
max_breadth: int = 20,
limit: int = 50,
instructions: Optional[str] = None,
include_images: bool = False,
extract_depth: str = "basic",
) -> Dict[str, Any]:
"""Crawl a website starting from a base URL.
Args:
url: The root URL to begin crawling
max_depth: Maximum depth from base URL (default: 1)
max_breadth: Max links to follow per page (default: 20)
limit: Total links to process before stopping (default: 50)
instructions: Natural language instructions for the crawler
include_images: Extract image URLs from pages
extract_depth: "basic" or "advanced"
Returns:
Dict containing:
- base_url: The starting URL
- results: List of crawled pages with url, raw_content, images
- response_time: Time taken for crawl
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
kwargs = {
"url": url,
"max_depth": max_depth,
"max_breadth": max_breadth,
"limit": limit,
"include_images": include_images,
"extract_depth": extract_depth,
}
if instructions:
kwargs["instructions"] = instructions
response = client.crawl(**kwargs)
return response
except Exception as e:
error_msg = f"Tavily crawl error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
def map(
self,
url: str,
max_depth: int = 1,
max_breadth: int = 20,
limit: int = 50,
instructions: Optional[str] = None,
) -> Dict[str, Any]:
"""Get a sitemap starting from a base URL.
Args:
url: The root URL to begin mapping
max_depth: Maximum depth from base URL (default: 1)
max_breadth: Max links to follow per page (default: 20)
limit: Total links to process before stopping (default: 50)
instructions: Natural language instructions for the mapper
Returns:
Dict containing:
- base_url: The starting URL
- results: List of discovered URLs
- response_time: Time taken for mapping
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
kwargs = {
"url": url,
"max_depth": max_depth,
"max_breadth": max_breadth,
"limit": limit,
}
if instructions:
kwargs["instructions"] = instructions
response = client.map(**kwargs)
return response
except Exception as e:
error_msg = f"Tavily map error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
# Async methods
async def asearch(
self,
query: str,
search_depth: str = "basic",
topic: str = "general",
max_results: int = 5,
include_answer: bool = False,
include_raw_content: bool = False,
include_images: bool = False,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
time_range: Optional[str] = None,
) -> Dict[str, Any]:
"""Async version of search. See search() for documentation."""
client, error = self._get_async_client()
if error:
logging.error(error)
return {"error": error}
try:
kwargs = {
"query": query,
"search_depth": search_depth,
"topic": topic,
"max_results": max_results,
"include_answer": include_answer,
"include_raw_content": include_raw_content,
"include_images": include_images,
}
if include_domains:
kwargs["include_domains"] = include_domains
if exclude_domains:
kwargs["exclude_domains"] = exclude_domains
if time_range:
kwargs["time_range"] = time_range
response = await client.search(**kwargs)
return response
except Exception as e:
error_msg = f"Tavily async search error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
async def aextract(
self,
urls: Union[str, List[str]],
include_images: bool = False,
extract_depth: str = "basic",
) -> Dict[str, Any]:
"""Async version of extract. See extract() for documentation."""
client, error = self._get_async_client()
if error:
logging.error(error)
return {"error": error}
try:
response = await client.extract(
urls=urls,
include_images=include_images,
extract_depth=extract_depth,
)
return response
except Exception as e:
error_msg = f"Tavily async extract error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
# Standalone functions for direct import
def tavily_search(
query: str,
search_depth: str = "basic",
topic: str = "general",
max_results: int = 5,
include_answer: bool = False,
include_raw_content: bool = False,
include_images: bool = False,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
time_range: Optional[str] = None,
) -> Dict[str, Any]:
"""Search the web using Tavily's AI-powered search.
This is a convenience function that creates a TavilyTools instance.
For repeated calls, consider using TavilyTools class directly.
Args:
query: The search query
search_depth: "basic" or "advanced"
topic: "general", "news", or "finance"
max_results: Maximum number of results (1-20)
include_answer: Include LLM-generated answer
include_raw_content: Include cleaned HTML content
include_images: Include query-related images
include_domains: List of domains to include
exclude_domains: List of domains to exclude
time_range: Filter by time - "day", "week", "month", "year"
Returns:
Dict containing search results
"""
tools = TavilyTools()
return tools.search(
query=query,
search_depth=search_depth,
topic=topic,
max_results=max_results,
include_answer=include_answer,
include_raw_content=include_raw_content,
include_images=include_images,
include_domains=include_domains,
exclude_domains=exclude_domains,
time_range=time_range,
)
def tavily_extract(
urls: Union[str, List[str]],
include_images: bool = False,
extract_depth: str = "basic",
) -> Dict[str, Any]:
"""Extract content from one or more URLs using Tavily.
Args:
urls: Single URL or list of URLs (max 20)
include_images: Include extracted images
extract_depth: "basic" or "advanced"
Returns:
Dict containing extraction results
"""
tools = TavilyTools()
return tools.extract(
urls=urls,
include_images=include_images,
extract_depth=extract_depth,
)
def tavily_crawl(
url: str,
max_depth: int = 1,
max_breadth: int = 20,
limit: int = 50,
instructions: Optional[str] = None,
include_images: bool = False,
extract_depth: str = "basic",
) -> Dict[str, Any]:
"""Crawl a website starting from a base URL using Tavily.
Args:
url: The root URL to begin crawling
max_depth: Maximum depth from base URL
max_breadth: Max links to follow per page
limit: Total links to process before stopping
instructions: Natural language instructions for the crawler
include_images: Extract image URLs from pages
extract_depth: "basic" or "advanced"
Returns:
Dict containing crawl results
"""
tools = TavilyTools()
return tools.crawl(
url=url,
max_depth=max_depth,
max_breadth=max_breadth,
limit=limit,
instructions=instructions,
include_images=include_images,
extract_depth=extract_depth,
)
def tavily_map(
url: str,
max_depth: int = 1,
max_breadth: int = 20,
limit: int = 50,
instructions: Optional[str] = None,
) -> Dict[str, Any]:
"""Get a sitemap starting from a base URL using Tavily.
Args:
url: The root URL to begin mapping
max_depth: Maximum depth from base URL
max_breadth: Max links to follow per page
limit: Total links to process before stopping
instructions: Natural language instructions for the mapper
Returns:
Dict containing map results with discovered URLs
"""
tools = TavilyTools()
return tools.map(
url=url,
max_depth=max_depth,
max_breadth=max_breadth,
limit=limit,
instructions=instructions,
)
# Async standalone functions
async def tavily_search_async(
query: str,
search_depth: str = "basic",
topic: str = "general",
max_results: int = 5,
include_answer: bool = False,
include_raw_content: bool = False,
include_images: bool = False,
include_domains: Optional[List[str]] = None,
exclude_domains: Optional[List[str]] = None,
time_range: Optional[str] = None,
) -> Dict[str, Any]:
"""Async version of tavily_search. See tavily_search() for documentation."""
tools = TavilyTools()
return await tools.asearch(
query=query,
search_depth=search_depth,
topic=topic,
max_results=max_results,
include_answer=include_answer,
include_raw_content=include_raw_content,
include_images=include_images,
include_domains=include_domains,
exclude_domains=exclude_domains,
time_range=time_range,
)
async def tavily_extract_async(
urls: Union[str, List[str]],
include_images: bool = False,
extract_depth: str = "basic",
) -> Dict[str, Any]:
"""Async version of tavily_extract. See tavily_extract() for documentation."""
tools = TavilyTools()
return await tools.aextract(
urls=urls,
include_images=include_images,
extract_depth=extract_depth,
)
if __name__ == "__main__":
# Example usage
print("\n==================================================")
print("Tavily Tools Demonstration")
print("==================================================\n")
# Check if API key is available
is_available, error = _check_tavily_available()
if not is_available:
print(f"Error: {error}")
print("\nTo use Tavily tools:")
print("1. Install: pip install tavily-python")
print("2. Set environment variable: export TAVILY_API_KEY=your_api_key")
else:
print("Tavily is available!")
# Example search
print("\n1. Basic Search")
print("-" * 30)
results = tavily_search("Latest AI developments 2024", max_results=3)
if "error" not in results:
print(f"Found {len(results.get('results', []))} results")
for r in results.get("results", [])[:3]:
print(f" - {r.get('title', 'No title')}")
print(f" URL: {r.get('url', 'No URL')}")
else:
print(f"Error: {results['error']}")
print("\n==================================================")
print("Demonstration Complete")
print("==================================================")
# Alias for simple usage: from praisonaiagents.tools import tavily
tavily = tavily_search
@@ -0,0 +1,533 @@
"""You.com search and content tools.
You.com provides AI-powered web search, news search, content extraction, and image search.
This module requires the YDC_API_KEY environment variable to be set.
Usage:
from praisonaiagents.tools import ydc_search, ydc_news, ydc_contents
# Search the web (unified web + news)
results = ydc_search("AI news 2024")
# Get live news
news = ydc_news("technology trends")
# Extract content from URLs
content = ydc_contents(["https://example.com"])
# Or use the class directly
from praisonaiagents.tools import YouTools
you = YouTools()
results = you.search("AI news")
"""
from typing import List, Dict, Any, Optional, Union
import logging
import os
from importlib import util
def _check_youdotcom_available() -> tuple[bool, Optional[str]]:
"""Check if You.com SDK is available and API key is set.
Returns:
Tuple of (is_available, error_message)
"""
# Check if youdotcom package is installed
if util.find_spec("youdotcom") is None:
return False, "youdotcom package is not installed. Please install it using: pip install youdotcom"
# Check if API key is set
api_key = os.environ.get("YDC_API_KEY")
if not api_key:
return False, "YDC_API_KEY environment variable is not set. Please set it to use You.com tools."
return True, None
class YouTools:
"""Comprehensive tools for web search and content using You.com API.
You.com provides AI-powered search capabilities optimized for LLM applications.
Features:
- Unified web and news search with LLM-ready snippets
- Live news search from authoritative sources
- Content extraction from URLs (HTML/Markdown)
- Image search
- Advanced search operators support
- Geographic and language targeting
- Freshness controls
Requires:
- youdotcom package: pip install youdotcom
- YDC_API_KEY environment variable
"""
def __init__(self, api_key: Optional[str] = None):
"""Initialize YouTools.
Args:
api_key: Optional API key. If not provided, uses YDC_API_KEY env var.
"""
self._api_key = api_key
self._client = None
def _get_api_key(self) -> Optional[str]:
"""Get API key from instance or environment."""
return self._api_key or os.environ.get("YDC_API_KEY")
def _get_client(self):
"""Get or create You.com client."""
if self._client is None:
is_available, error = _check_youdotcom_available()
if not is_available:
return None, error
from youdotcom import You
api_key = self._get_api_key()
if not api_key:
return None, "YDC_API_KEY not set"
self._client = You(api_key_auth=api_key)
return self._client, None
def search(
self,
query: str,
count: int = 10,
freshness: Optional[str] = None,
country: Optional[str] = None,
language: Optional[str] = None,
offset: int = 0,
safesearch: str = "moderate",
livecrawl: Optional[str] = None,
livecrawl_formats: Optional[str] = None,
) -> Dict[str, Any]:
"""Search the web using You.com's unified search API.
Returns both web and news results in a single call with LLM-ready snippets.
Args:
query: Search query (supports operators like site:, filetype:, +term, -term)
count: Max results per section (default 10, max 100)
freshness: Filter by recency - "day", "week", "month", "year" or "YYYY-MM-DDtoYYYY-MM-DD"
country: ISO country code (e.g., "US", "GB", "FR")
language: BCP 47 language code (e.g., "EN", "JP", "DE")
offset: Pagination offset (0-9), results = offset * count
safesearch: Content filter - "off", "moderate", "strict"
livecrawl: Fetch full content - "web", "news", or "all"
livecrawl_formats: Content format - "html" or "markdown"
Returns:
Dict containing:
- results: {web: [...], news: [...]}
- metadata: {query, search_uuid, latency}
Search Operators:
- site:domain.com - Search within specific domain
- filetype:pdf - Filter by file type
- +term - Must include term
- -term - Exclude term
- AND, OR, NOT - Boolean operators
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
kwargs = {"query": query, "count": count, "offset": offset}
if freshness:
kwargs["freshness"] = freshness
if country:
kwargs["country"] = country
if language:
kwargs["language"] = language
if safesearch:
kwargs["safesearch"] = safesearch
if livecrawl:
kwargs["livecrawl"] = livecrawl
if livecrawl_formats:
kwargs["livecrawl_formats"] = livecrawl_formats
response = client.search.unified(**kwargs)
# Convert response to dict if it's a model object
if hasattr(response, 'model_dump'):
return response.model_dump()
elif hasattr(response, 'to_dict'):
return response.to_dict()
elif hasattr(response, '__dict__'):
return self._convert_response(response)
return response
except Exception as e:
error_msg = f"You.com search error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
def _convert_response(self, response) -> Dict[str, Any]:
"""Convert SDK response object to dictionary."""
result = {}
if hasattr(response, 'results'):
results = response.results
result['results'] = {}
if hasattr(results, 'web') and results.web:
result['results']['web'] = [
{
'url': getattr(r, 'url', ''),
'title': getattr(r, 'title', ''),
'description': getattr(r, 'description', ''),
'snippets': getattr(r, 'snippets', []),
'thumbnail_url': getattr(r, 'thumbnail_url', None),
'page_age': getattr(r, 'page_age', None),
'favicon_url': getattr(r, 'favicon_url', None),
'authors': getattr(r, 'authors', []),
}
for r in results.web
]
if hasattr(results, 'news') and results.news:
result['results']['news'] = [
{
'url': getattr(r, 'url', ''),
'title': getattr(r, 'title', ''),
'description': getattr(r, 'description', ''),
'thumbnail_url': getattr(r, 'thumbnail_url', None),
'page_age': getattr(r, 'page_age', None),
}
for r in results.news
]
if hasattr(response, 'metadata'):
meta = response.metadata
result['metadata'] = {
'query': getattr(meta, 'query', ''),
'search_uuid': getattr(meta, 'search_uuid', ''),
'latency': getattr(meta, 'latency', 0),
}
return result
def get_contents(
self,
urls: Union[str, List[str]],
format: str = "markdown",
) -> Dict[str, Any]:
"""Fetch content from URLs using You.com's content extraction API.
Args:
urls: Single URL or list of URLs to fetch content from
format: Content format - "html" or "markdown" (default: markdown)
Returns:
Dict containing array of results with:
- url: The fetched URL
- title: Page title
- html: HTML content (if format="html")
- markdown: Markdown content (if format="markdown")
- metadata: {site_name, favicon_url}
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
if isinstance(urls, str):
urls = [urls]
# Use the contents endpoint
response = client.search.contents(urls=urls, format=format)
# Convert response
if hasattr(response, 'model_dump'):
return {"results": response.model_dump()}
elif isinstance(response, list):
return {"results": [
{
'url': getattr(r, 'url', ''),
'title': getattr(r, 'title', ''),
'html': getattr(r, 'html', None),
'markdown': getattr(r, 'markdown', None),
'metadata': {
'site_name': getattr(getattr(r, 'metadata', None), 'site_name', None) if hasattr(r, 'metadata') else None,
'favicon_url': getattr(getattr(r, 'metadata', None), 'favicon_url', None) if hasattr(r, 'metadata') else None,
}
}
for r in response
]}
return {"results": response}
except Exception as e:
error_msg = f"You.com contents error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
def live_news(
self,
query: str,
count: Optional[int] = None,
) -> Dict[str, Any]:
"""Search for live news using You.com's news API.
Note: This API may require early access. Contact api@you.com for access.
Args:
query: News search query
count: Maximum number of news results to return
Returns:
Dict containing:
- news: {query, results: [...], metadata}
- Each result has: title, description, url, page_age, source_name, thumbnail
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
kwargs = {"q": query}
if count:
kwargs["count"] = count
# Try to access live news endpoint
if hasattr(client, 'livenews'):
response = client.livenews.get(**kwargs)
else:
# Fallback: use unified search with news focus
return self.search(query=query, count=count or 10)
if hasattr(response, 'model_dump'):
return response.model_dump()
elif hasattr(response, 'news'):
news = response.news
return {
'news': {
'query': getattr(getattr(news, 'query', None), 'original', query) if hasattr(news, 'query') else query,
'results': [
{
'title': getattr(r, 'title', ''),
'description': getattr(r, 'description', ''),
'url': getattr(r, 'url', ''),
'page_age': getattr(r, 'page_age', None),
'age': getattr(r, 'age', None),
'source_name': getattr(r, 'source_name', None),
'thumbnail': getattr(getattr(r, 'thumbnail', None), 'src', None) if hasattr(r, 'thumbnail') else None,
}
for r in getattr(news, 'results', [])
],
'metadata': {
'request_uuid': getattr(getattr(news, 'metadata', None), 'request_uuid', None) if hasattr(news, 'metadata') else None,
}
}
}
return response
except Exception as e:
error_msg = f"You.com live news error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
def images(
self,
query: str,
) -> Dict[str, Any]:
"""Search for images using You.com's image search API.
Note: This API may require early access. Contact api@you.com for access.
Args:
query: Image search query
Returns:
Dict containing:
- images: {results: [...]}
- metadata: {query, search_uuid}
- Each result has: title, page_url, image_url
"""
client, error = self._get_client()
if error:
logging.error(error)
return {"error": error}
try:
if hasattr(client, 'images'):
response = client.images.get(q=query)
else:
return {"error": "Image search not available in this SDK version"}
if hasattr(response, 'model_dump'):
return response.model_dump()
elif hasattr(response, 'images'):
images = response.images
return {
'images': {
'results': [
{
'title': getattr(r, 'title', ''),
'page_url': getattr(r, 'page_url', ''),
'image_url': getattr(r, 'image_url', ''),
}
for r in getattr(images, 'results', [])
]
},
'metadata': {
'query': getattr(getattr(response, 'metadata', None), 'query', query) if hasattr(response, 'metadata') else query,
'search_uuid': getattr(getattr(response, 'metadata', None), 'search_uuid', None) if hasattr(response, 'metadata') else None,
}
}
return response
except Exception as e:
error_msg = f"You.com images error: {str(e)}"
logging.error(error_msg)
return {"error": error_msg}
def __enter__(self):
"""Context manager entry."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit."""
if self._client is not None:
# Close client if it has a close method
if hasattr(self._client, 'close'):
self._client.close()
elif hasattr(self._client, '__exit__'):
self._client.__exit__(exc_type, exc_val, exc_tb)
return False
# Standalone functions for direct import
def ydc_search(
query: str,
count: int = 10,
freshness: Optional[str] = None,
country: Optional[str] = None,
language: Optional[str] = None,
offset: int = 0,
safesearch: str = "moderate",
livecrawl: Optional[str] = None,
livecrawl_formats: Optional[str] = None,
) -> Dict[str, Any]:
"""Search the web using You.com's unified search API.
This is a convenience function that creates a YouTools instance.
For repeated calls, consider using YouTools class directly.
Args:
query: Search query (supports operators like site:, filetype:, +term, -term)
count: Max results per section (default 10, max 100)
freshness: Filter by recency - "day", "week", "month", "year"
country: ISO country code (e.g., "US", "GB")
language: BCP 47 language code (e.g., "EN", "JP")
offset: Pagination offset (0-9)
safesearch: Content filter - "off", "moderate", "strict"
livecrawl: Fetch full content - "web", "news", or "all"
livecrawl_formats: Content format - "html" or "markdown"
Returns:
Dict containing search results with web and news sections
"""
tools = YouTools()
return tools.search(
query=query,
count=count,
freshness=freshness,
country=country,
language=language,
offset=offset,
safesearch=safesearch,
livecrawl=livecrawl,
livecrawl_formats=livecrawl_formats,
)
def ydc_contents(
urls: Union[str, List[str]],
format: str = "markdown",
) -> Dict[str, Any]:
"""Fetch content from URLs using You.com's content extraction API.
Args:
urls: Single URL or list of URLs to fetch content from
format: Content format - "html" or "markdown"
Returns:
Dict containing extracted content for each URL
"""
tools = YouTools()
return tools.get_contents(urls=urls, format=format)
def ydc_news(
query: str,
count: Optional[int] = None,
) -> Dict[str, Any]:
"""Search for live news using You.com's news API.
Args:
query: News search query
count: Maximum number of results
Returns:
Dict containing news results
"""
tools = YouTools()
return tools.live_news(query=query, count=count)
def ydc_images(
query: str,
) -> Dict[str, Any]:
"""Search for images using You.com's image search API.
Args:
query: Image search query
Returns:
Dict containing image results
"""
tools = YouTools()
return tools.images(query=query)
if __name__ == "__main__":
# Example usage
print("\n" + "=" * 60)
print("You.com Tools Demonstration")
print("=" * 60 + "\n")
# Check if API key is available
is_available, error = _check_youdotcom_available()
if not is_available:
print(f"Error: {error}")
print("\nTo use You.com tools:")
print("1. Install: pip install youdotcom")
print("2. Set environment variable: export YDC_API_KEY=your_api_key")
else:
print("You.com SDK is available!")
# Example search
print("\n1. Basic Search")
print("-" * 30)
results = ydc_search("Latest AI developments 2024", count=3)
if "error" not in results:
web_results = results.get('results', {}).get('web', [])
print(f"Found {len(web_results)} web results")
for r in web_results[:3]:
print(f" - {r.get('title', 'No title')[:50]}...")
print(f" URL: {r.get('url', 'No URL')}")
else:
print(f"Error: {results['error']}")
print("\n" + "=" * 60)
print("Demonstration Complete")
print("=" * 60)
# Alias for simple usage: from praisonaiagents.tools import ydc
ydc = ydc_search
+3 -2
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "praisonaiagents"
version = "0.0.161"
version = "0.0.189"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
requires-python = ">=3.10"
authors = [
@@ -22,7 +22,8 @@ dependencies = [
mcp = [
"mcp>=1.6.0",
"fastapi>=0.115.0",
"uvicorn>=0.34.0"
"uvicorn>=0.34.0",
"websockets>=12.0"
]
memory = [
@@ -0,0 +1,141 @@
"""
Query Rewriter Agent Example
Demonstrates various query rewriting strategies for improving RAG retrieval.
"""
from dotenv import load_dotenv
load_dotenv()
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
def main():
# Initialize the agent
agent = QueryRewriterAgent(
model="gpt-4o-mini",
verbose=True
)
print("=" * 60)
print("Query Rewriter Agent Examples")
print("=" * 60)
# 1. Basic Rewriting - Short keyword queries
print("\n1. BASIC REWRITING (Short Queries)")
print("-" * 40)
short_queries = [
"AI trends",
"python performance",
"RAG best practices"
]
for query in short_queries:
result = agent.rewrite(query, strategy=RewriteStrategy.BASIC)
print(f"Original: {query}")
print(f"Rewritten: {result.primary_query}")
print()
# 2. HyDE - Hypothetical Document Embeddings
print("\n2. HYDE (Hypothetical Document)")
print("-" * 40)
query = "What is quantum computing?"
result = agent.rewrite(query, strategy=RewriteStrategy.HYDE)
print(f"Original: {query}")
print("Hypothetical Document (truncated):")
print(f"{result.hypothetical_document[:300]}...")
print()
# 3. Step-Back - Higher-level concept questions
print("\n3. STEP-BACK (Broader Context)")
print("-" * 40)
query = "What is the difference between GPT-4 and Claude 3?"
result = agent.rewrite(query, strategy=RewriteStrategy.STEP_BACK)
print(f"Original: {query}")
print(f"Rewritten: {result.primary_query}")
print(f"Step-back Question: {result.step_back_question}")
print()
# 4. Sub-Queries - Decompose complex questions
print("\n4. SUB-QUERIES (Decomposition)")
print("-" * 40)
query = "How do I set up a RAG pipeline with vector search and what are the best embedding models to use?"
result = agent.rewrite(query, strategy=RewriteStrategy.SUB_QUERIES)
print(f"Original: {query}")
print("Sub-queries:")
for i, sq in enumerate(result.sub_queries, 1):
print(f" {i}. {sq}")
print()
# 5. Multi-Query - Multiple paraphrased versions
print("\n5. MULTI-QUERY (Ensemble Retrieval)")
print("-" * 40)
query = "How to improve LLM response quality?"
result = agent.rewrite(query, strategy=RewriteStrategy.MULTI_QUERY, num_queries=3)
print(f"Original: {query}")
print("Alternative queries:")
for i, q in enumerate(result.rewritten_queries, 1):
print(f" {i}. {q}")
print()
# 6. Contextual - Using conversation history
print("\n6. CONTEXTUAL (With Chat History)")
print("-" * 40)
chat_history = [
{"role": "user", "content": "Tell me about Python programming"},
{"role": "assistant", "content": "Python is a high-level programming language known for its simplicity and readability..."},
{"role": "user", "content": "What frameworks are popular?"},
{"role": "assistant", "content": "Popular Python frameworks include Django for web development, FastAPI for APIs, and PyTorch for ML..."}
]
query = "What about its performance compared to others?"
result = agent.rewrite(query, strategy=RewriteStrategy.CONTEXTUAL, chat_history=chat_history)
print(f"Original: {query}")
print(f"Standalone: {result.primary_query}")
print()
# 7. AUTO - Automatic strategy detection
print("\n7. AUTO (Automatic Strategy Detection)")
print("-" * 40)
test_queries = [
("ML", None), # Short query -> BASIC
("What about the cost?", chat_history), # Follow-up -> CONTEXTUAL
("Compare transformers vs RNNs and explain their use cases", None), # Complex -> SUB_QUERIES
]
for query, history in test_queries:
result = agent.rewrite(query, strategy=RewriteStrategy.AUTO, chat_history=history)
print(f"Query: {query}")
print(f"Detected Strategy: {result.strategy_used.value}")
print(f"Result: {result.primary_query}")
print()
# 8. Custom Abbreviations
print("\n8. CUSTOM ABBREVIATIONS")
print("-" * 40)
agent.add_abbreviations({
"K8s": "Kubernetes",
"TF": "TensorFlow",
"PT": "PyTorch"
})
query = "K8s deployment for TF models"
result = agent.rewrite(query, strategy=RewriteStrategy.BASIC)
print(f"Original: {query}")
print(f"Rewritten: {result.primary_query}")
print()
print("=" * 60)
print("Examples Complete!")
print("=" * 60)
if __name__ == "__main__":
main()
@@ -0,0 +1,16 @@
from praisonaiagents import Agent
def search_web(query: str) -> str:
return "AI trends: LLMs, multimodal AI, autonomous agents"
agent = Agent(
name="AI Assistant",
role="Research and Writing Specialist",
llm="gpt-4o-mini",
planning=True,
planning_tools=[search_web],
planning_reasoning=True
)
result = agent.start("Research AI trends in 2025 and write a 2 sentence summary")
print(result)
@@ -0,0 +1,18 @@
from praisonaiagents import Agent
def search_web(query: str) -> str:
return "AI trends: LLMs, multimodal AI, autonomous agents"
agent = Agent(
name="AI Assistant",
role="Research and Writing Specialist",
llm="gpt-4o-mini",
planning=True,
planning_tools=[search_web],
planning_reasoning=True,
verbose=False # Quiet mode
)
result = agent.start("Research AI trends in 2025 and write a 2 sentence summary")
print("\n=== FINAL RESULT ===")
print(result)
@@ -0,0 +1,62 @@
from praisonaiagents import Agent
def search_web(query: str) -> str:
return "AI trends: LLMs, multimodal AI, autonomous agents, edge AI"
print("=" * 60)
print("TEST 1: Planning ONLY (no tools, no reasoning)")
print("=" * 60)
agent1 = Agent(
name="Assistant",
llm="gpt-4o-mini",
planning=True,
verbose=False
)
result1 = agent1.start("Write 2 sentences about AI in healthcare")
print(f"\nResult: {result1[:200]}...")
print("\n" + "=" * 60)
print("TEST 2: Planning + Tools (no reasoning)")
print("=" * 60)
agent2 = Agent(
name="Assistant",
llm="gpt-4o-mini",
planning=True,
planning_tools=[search_web],
verbose=False
)
result2 = agent2.start("Write 2 sentences about AI trends")
print(f"\nResult: {result2[:200]}...")
print("\n" + "=" * 60)
print("TEST 3: Planning + Tools + Reasoning")
print("=" * 60)
agent3 = Agent(
name="Assistant",
llm="gpt-4o-mini",
planning=True,
planning_tools=[search_web],
planning_reasoning=True,
verbose=False
)
result3 = agent3.start("Write 2 sentences about AI trends")
print(f"\nResult: {result3[:200]}...")
print("\n" + "=" * 60)
print("TEST 4: No Planning (direct execution)")
print("=" * 60)
agent4 = Agent(
name="Assistant",
llm="gpt-4o-mini",
verbose=False
)
result4 = agent4.start("Write 2 sentences about AI trends")
print(f"\nResult: {result4[:200]}...")
print("\n" + "=" * 60)
print("COMPARISON COMPLETE")
print("=" * 60)
@@ -0,0 +1,183 @@
"""
Complete Planning Mode Demo: Plan First, Then Execute Todo Items One by One.
This demonstrates the full workflow:
1. Create a plan with the PlanningAgent
2. Generate todo list from the plan
3. Execute each todo item sequentially with the appropriate agent
4. Update progress as items are completed
Run with: python test_planning_execution.py
"""
from praisonaiagents import Agent, Task
from praisonaiagents.planning import PlanningAgent, TodoList
def print_separator(title: str):
print("\n" + "=" * 70)
print(f" {title}")
print("=" * 70 + "\n")
def print_progress(todo_list: TodoList):
"""Print current progress."""
completed = len(todo_list.completed)
total = len(todo_list.items)
progress = todo_list.progress * 100
bar_length = 30
filled = int(bar_length * todo_list.progress)
bar = "█" * filled + "░" * (bar_length - filled)
print(f"\n📊 Progress: [{bar}] {progress:.0f}% ({completed}/{total} completed)\n")
def main():
print_separator("🎯 PLANNING MODE: PLAN FIRST, EXECUTE TODO ITEMS")
# ==========================================================================
# STEP 1: Create Agents
# ==========================================================================
print("📦 Creating agents...")
researcher = Agent(
name="Researcher",
role="Research Analyst",
goal="Research topics thoroughly and provide accurate information",
backstory="Expert researcher with access to vast knowledge bases",
verbose=True
)
writer = Agent(
name="Writer",
role="Content Writer",
goal="Write clear, engaging, and well-structured content",
backstory="Professional writer skilled in creating compelling narratives",
verbose=True
)
print("✅ Agents created: Researcher, Writer\n")
# ==========================================================================
# STEP 2: Create Plan using PlanningAgent
# ==========================================================================
print_separator("📋 PHASE 1: CREATING PLAN")
planner = PlanningAgent(
llm="gpt-4o-mini",
read_only=True,
verbose=1
)
request = "Write a short article about the top 3 benefits of meditation for mental health"
print(f"📝 Request: {request}\n")
print("🤔 Planning agent is analyzing and creating a plan...\n")
plan = planner.create_plan_sync(
request=request,
agents=[researcher, writer],
context="Keep it concise - each section should be 2-3 sentences max"
)
print("\n✅ Plan created successfully!")
print(f" Plan Name: {plan.name}")
print(f" Total Steps: {len(plan.steps)}")
# ==========================================================================
# STEP 3: Display the Plan
# ==========================================================================
print_separator("📄 GENERATED PLAN")
print(plan.to_markdown())
# ==========================================================================
# STEP 4: Create Todo List from Plan
# ==========================================================================
print_separator("📝 PHASE 2: TODO LIST CREATED")
todo_list = TodoList.from_plan(plan)
print(todo_list.to_markdown())
print_progress(todo_list)
# ==========================================================================
# STEP 5: Execute Each Todo Item One by One
# ==========================================================================
print_separator("🚀 PHASE 3: EXECUTING TODO ITEMS")
# Map agent names to agent instances
agent_map = {
"Researcher": researcher,
"Writer": writer
}
results = {}
for i, item in enumerate(todo_list.items):
print(f"\n{'─' * 70}")
print(f"📌 TODO ITEM {i + 1}/{len(todo_list.items)}")
print(f"{'─' * 70}")
print(f" Description: {item.description}")
print(f" Agent: {item.agent or 'Auto-assign'}")
print(f" Status: {item.status}")
if item.dependencies:
print(f" Dependencies: {', '.join(item.dependencies)}")
# Mark as in progress
todo_list.start(item.id)
print(f"\n ⏳ Starting execution...")
# Get the appropriate agent
agent_name = item.agent or "Researcher"
agent = agent_map.get(agent_name, researcher)
# Build context from previous results
context = ""
for dep_id in item.dependencies:
if dep_id in results:
context += f"\nPrevious result ({dep_id}):\n{results[dep_id]}\n"
# Create task description
task_description = item.description
if context:
task_description += f"\n\nContext from previous steps:{context}"
# Execute with the agent
try:
result = agent.chat(task_description)
results[item.id] = result
# Mark as completed
todo_list.complete(item.id)
print(f"\n ✅ Completed!")
print(f"\n 📄 Result Preview:")
print(" " + "-" * 50)
preview = str(result)[:500] + "..." if len(str(result)) > 500 else str(result)
for line in preview.split('\n'):
print(f" {line}")
except Exception as e:
print(f"\n ❌ Error: {e}")
results[item.id] = f"Error: {e}"
print_progress(todo_list)
# ==========================================================================
# STEP 6: Final Summary
# ==========================================================================
print_separator("🎉 EXECUTION COMPLETE")
print("📋 Final Todo List Status:")
print(todo_list.to_markdown())
print("\n📊 Final Results Summary:")
print("-" * 50)
for item in todo_list.items:
status_icon = "✅" if item.status == "completed" else "❌"
print(f"{status_icon} {item.description[:60]}...")
print(f"\n🏆 All {len(todo_list.completed)} tasks completed successfully!")
if __name__ == "__main__":
main()
@@ -0,0 +1,22 @@
from praisonaiagents import Agent, Task, PraisonAIAgents
def search_web(query: str) -> str:
return f"Results: AI trends include LLMs, multimodal AI, autonomous agents"
researcher = Agent(name="Researcher", role="Research Analyst", llm="gpt-4o-mini")
writer = Agent(name="Writer", role="Content Writer", llm="gpt-4o-mini")
task1 = Task(description="Research AI trends in 2025", agent=researcher)
task2 = Task(description="Write a 2 sentence summary", agent=writer)
agents = PraisonAIAgents(
agents=[researcher, writer],
tasks=[task1, task2],
planning=True,
planning_tools=[search_web],
planning_reasoning=True,
auto_approve_plan=True
)
result = agents.start()
print(result)
@@ -0,0 +1,28 @@
from praisonaiagents import Agent, Task, PraisonAIAgents
def search_web(query: str) -> str:
return "AI trends: LLMs, multimodal AI, autonomous agents"
agent = Agent(
name="AI Assistant",
role="Research and Writing Specialist",
llm="gpt-4o-mini",
instructions="Research AI trends in 2025 and write a 2 sentence summary"
)
task = Task(
description="Research AI trends in 2025 and write a 2 sentence summary",
agent=agent
)
agents = PraisonAIAgents(
agents=[agent],
tasks=[task],
planning=True,
planning_tools=[search_web],
planning_reasoning=True,
auto_approve_plan=True
)
result = agents.start()
print(result)
@@ -0,0 +1,34 @@
from praisonaiagents import Agent
from praisonaiagents.planning import PlanningAgent
def search_web(query: str) -> str:
return f"Results for '{query}': LLMs, multimodal AI, autonomous agents"
researcher = Agent(name="Researcher", role="Research Analyst", llm="gpt-4o-mini")
writer = Agent(name="Writer", role="Content Writer", llm="gpt-4o-mini")
print("\n=== TEST 1: Basic Planning (NO tools) ===")
planner1 = PlanningAgent(llm="gpt-4o-mini")
plan1 = planner1.create_plan_sync(
request="Write about AI trends 2025",
agents=[researcher, writer]
)
print(f"Plan: {plan1.name}, Steps: {len(plan1.steps)}")
print("\n=== TEST 2: Planning WITH tools ===")
planner2 = PlanningAgent(llm="gpt-4o-mini", tools=[search_web])
plan2 = planner2.create_plan_sync(
request="Write about AI trends 2025",
agents=[researcher, writer]
)
print(f"Plan: {plan2.name}, Steps: {len(plan2.steps)}")
print("\n=== TEST 3: Planning WITH tools + reasoning ===")
planner3 = PlanningAgent(llm="gpt-4o-mini", tools=[search_web], reasoning=True)
plan3 = planner3.create_plan_sync(
request="Write about AI trends 2025",
agents=[researcher, writer]
)
print(f"Plan: {plan3.name}, Steps: {len(plan3.steps)}")
print("\n✅ Done!")
@@ -0,0 +1,344 @@
"""
Deep Research Agent Examples
This file demonstrates how to use the DeepResearchAgent for automated
complex research workflows using multiple providers:
- **OpenAI**: o3-deep-research, o4-mini-deep-research (via Responses API)
- **Gemini**: deep-research-pro (via Interactions API)
- **LiteLLM**: Unified interface for OpenAI models
Requirements:
- OPENAI_API_KEY for OpenAI/LiteLLM providers
- GEMINI_API_KEY for Gemini provider
- OpenAI SDK >= 1.0.0
- google-genai for Gemini (pip install google-genai)
- litellm for LiteLLM provider (pip install litellm)
Available models:
OpenAI:
- o3-deep-research: Higher quality, in-depth synthesis
- o4-mini-deep-research: Faster, lightweight
Gemini:
- deep-research-pro: Gemini 3 Pro powered research agent
"""
from praisonaiagents import DeepResearchAgent, Provider
# =============================================================================
# Example 1: Basic Research Query
# =============================================================================
def basic_research_example():
"""Simple research query with default settings."""
agent = DeepResearchAgent(
name="Research Assistant",
instructions="""
You are a professional researcher. Focus on:
- Data-rich insights with specific figures and trends
- Reliable sources: peer-reviewed research, official reports
- Include inline citations
""",
model="o4-mini-deep-research", # Use faster model
verbose=True
)
result = agent.research(
"What are the key trends in renewable energy adoption in 2024?"
)
print("=" * 60)
print("RESEARCH REPORT")
print("=" * 60)
print(result.report)
print("\n" + "=" * 60)
print("CITATIONS")
print("=" * 60)
for i, citation in enumerate(result.citations, 1):
print(f"{i}. {citation.title}")
print(f" URL: {citation.url}")
return result
# =============================================================================
# Example 2: Gemini Deep Research
# =============================================================================
def gemini_research_example():
"""Research using Gemini Deep Research Agent."""
agent = DeepResearchAgent(
name="Gemini Researcher",
instructions="""
You are a professional researcher. Focus on:
- Comprehensive analysis with multiple perspectives
- Data-backed insights and trends
- Clear structure with executive summary
""",
model="deep-research-pro", # Gemini model - auto-detected
verbose=True
)
print(f"Using provider: {agent.provider}") # Should be Provider.GEMINI
result = agent.research(
"What are the latest breakthroughs in fusion energy research?"
)
print("=" * 60)
print("GEMINI RESEARCH REPORT")
print("=" * 60)
print(result.report[:1000] + "..." if len(result.report) > 1000 else result.report)
return result
# =============================================================================
# Example 3: Research with Code Interpreter (OpenAI)
# =============================================================================
def research_with_code_example():
"""Research with code execution for data analysis."""
agent = DeepResearchAgent(
name="Data Analyst",
instructions="""
You are a data analyst. When appropriate:
- Use code to analyze data and create visualizations
- Summarize data in tables
- Include specific statistics and metrics
""",
model="o3-deep-research",
enable_code_interpreter=True,
verbose=True
)
result = agent.research(
"Analyze the global smartphone market share trends over the past 5 years.",
summary_mode="detailed"
)
# Check if any code was executed
if result.code_executions:
print("\nCode executed during research:")
for code in result.code_executions:
print(f" Input: {code.input_code[:100]}...")
if code.output:
print(f" Output: {code.output[:100]}...")
return result
# =============================================================================
# Example 4: Using LiteLLM Provider
# =============================================================================
def litellm_research_example():
"""Research using LiteLLM for unified interface."""
agent = DeepResearchAgent(
name="LiteLLM Researcher",
instructions="You are a professional researcher.",
model="o3-deep-research",
use_litellm=True, # Use LiteLLM instead of direct OpenAI
verbose=True
)
print(f"Using provider: {agent.provider}") # Should be Provider.LITELLM
result = agent.research(
"What are the key trends in sustainable agriculture?"
)
print(f"Report length: {len(result.report)} characters")
return result
# =============================================================================
# Example 5: Research with MCP Integration (OpenAI)
# =============================================================================
def research_with_mcp_example():
"""Research with custom MCP server for internal documents."""
agent = DeepResearchAgent(
name="Internal Researcher",
instructions="""
You are a researcher with access to internal documents.
- Use the internal file lookup tool to retrieve company data
- Combine internal data with public research
- Prioritize internal sources when available
""",
model="o3-deep-research",
mcp_servers=[
{
"label": "internal_docs",
"url": "https://your-mcp-server.com/sse/",
"require_approval": "never"
}
],
verbose=True
)
result = agent.research(
"What is our company's market position compared to competitors?"
)
# Check MCP calls
if result.mcp_calls:
print("\nMCP tools used:")
for call in result.mcp_calls:
print(f" - {call.name}@{call.server_label}: {call.arguments}")
return result
# =============================================================================
# Example 6: Gemini with File Search
# =============================================================================
def gemini_file_search_example():
"""Research using Gemini with your own data via File Search."""
agent = DeepResearchAgent(
name="Document Researcher",
instructions="""
Compare our internal documents against current public information.
Prioritize insights from our internal data.
""",
model="deep-research-pro",
enable_file_search=True,
file_search_stores=["fileSearchStores/my-store-name"],
verbose=True
)
result = agent.research(
"Compare our 2025 fiscal year report against current market trends."
)
return result
# =============================================================================
# Example 7: Clarifying Questions Workflow
# =============================================================================
def clarifying_workflow_example():
"""Use clarifying questions before research."""
agent = DeepResearchAgent(
name="Research Assistant",
verbose=True
)
# Step 1: Get clarifying questions
initial_query = "Help me research AI trends"
questions = agent.clarify(initial_query)
print("Clarifying Questions:")
print(questions)
# Step 2: User provides more details (simulated)
detailed_query = """
I'd like to research AI trends in healthcare specifically.
Focus on:
- Diagnostic AI applications
- FDA-approved AI tools
- Cost savings for hospitals
Time period: 2023-2024
"""
# Step 3: Rewrite the query for optimal research
rewritten = agent.rewrite_query(detailed_query)
print("\nRewritten Query:")
print(rewritten)
# Step 4: Perform the research
result = agent.research(rewritten)
return result
# =============================================================================
# Example 5: Async Research
# =============================================================================
async def async_research_example():
"""Perform research asynchronously."""
import asyncio
agent = DeepResearchAgent(
name="Async Researcher",
model="o4-mini-deep-research",
verbose=True
)
# Run multiple research queries concurrently
queries = [
"What are the latest developments in quantum computing?",
"What are the trends in electric vehicle adoption?",
]
tasks = [agent.aresearch(q) for q in queries]
results = await asyncio.gather(*tasks)
for query, result in zip(queries, results):
print(f"\nQuery: {query}")
print(f"Report length: {len(result.report)} chars")
print(f"Citations: {len(result.citations)}")
return results
# =============================================================================
# Example 6: Inspecting Intermediate Steps
# =============================================================================
def inspect_steps_example():
"""Inspect the research process steps."""
agent = DeepResearchAgent(
name="Detailed Researcher",
model="o4-mini-deep-research",
verbose=True
)
result = agent.research(
"What are the environmental impacts of cryptocurrency mining?"
)
print("\n" + "=" * 60)
print("RESEARCH PROCESS ANALYSIS")
print("=" * 60)
# Reasoning steps
print(f"\nReasoning Steps ({len(result.reasoning_steps)}):")
for i, step in enumerate(result.reasoning_steps, 1):
print(f" {i}. {step.text[:100]}...")
# Web searches
print(f"\nWeb Searches ({len(result.web_searches)}):")
for search in result.web_searches:
print(f" - Query: {search.query}")
print(f" Status: {search.status}")
# Sources
print(f"\nUnique Sources ({len(result.get_all_sources())}):")
for source in result.get_all_sources()[:10]:
print(f" - {source['title']}")
print(f" {source['url']}")
return result
# =============================================================================
# Main
# =============================================================================
if __name__ == "__main__":
import os
if not os.environ.get("OPENAI_API_KEY"):
print("Error: OPENAI_API_KEY environment variable not set")
print("Please set it before running examples:")
print(" export OPENAI_API_KEY='your-api-key'")
exit(1)
print("Running basic research example...")
result = basic_research_example()
print(f"\nCompleted! Report has {len(result.report)} characters and {len(result.citations)} citations.")
@@ -0,0 +1,283 @@
"""
Test file for DeepResearchAgent
This file tests the OpenAI Deep Research API integration.
Requires OPENAI_API_KEY environment variable to be set.
Usage:
python tests/deep_research_test.py
"""
import os
import sys
# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from praisonaiagents import (
DeepResearchAgent,
DeepResearchResponse,
Citation,
ReasoningStep,
WebSearchCall,
CodeExecutionStep,
MCPCall,
FileSearchCall,
Provider
)
def test_basic_import():
"""Test that all classes can be imported."""
print("✓ All Deep Research classes imported successfully")
print(f" - DeepResearchAgent: {DeepResearchAgent}")
print(f" - DeepResearchResponse: {DeepResearchResponse}")
print(f" - Citation: {Citation}")
print(f" - ReasoningStep: {ReasoningStep}")
print(f" - WebSearchCall: {WebSearchCall}")
print(f" - CodeExecutionStep: {CodeExecutionStep}")
print(f" - MCPCall: {MCPCall}")
print(f" - FileSearchCall: {FileSearchCall}")
print(f" - Provider: {Provider}")
print(f" - Provider.OPENAI: {Provider.OPENAI}")
print(f" - Provider.GEMINI: {Provider.GEMINI}")
print(f" - Provider.LITELLM: {Provider.LITELLM}")
return True
def test_agent_initialization():
"""Test DeepResearchAgent initialization."""
# Check if API key is available
if not os.environ.get("OPENAI_API_KEY"):
print("⚠ Skipping initialization test - OPENAI_API_KEY not set")
return True
agent = DeepResearchAgent(
name="Test Researcher",
instructions="You are a professional researcher.",
model="o3-deep-research",
verbose=True
)
print(f"✓ Agent initialized: {agent}")
print(f" - Name: {agent.name}")
print(f" - Model: {agent.model}")
print(f" - Web search enabled: {agent.enable_web_search}")
print(f" - Code interpreter enabled: {agent.enable_code_interpreter}")
return True
def test_data_classes():
"""Test data class creation and methods."""
# Test Citation
citation = Citation(
title="Test Source",
url="https://example.com",
start_index=0,
end_index=100
)
print(f"✓ Citation created: {citation}")
# Test ReasoningStep
reasoning = ReasoningStep(text="Analyzing the query...")
print(f"✓ ReasoningStep created: {reasoning}")
# Test WebSearchCall
search = WebSearchCall(query="test query", status="completed")
print(f"✓ WebSearchCall created: {search}")
# Test CodeExecutionStep
code = CodeExecutionStep(input_code="print('hello')", output="hello")
print(f"✓ CodeExecutionStep created: {code}")
# Test MCPCall
mcp = MCPCall(name="fetch", server_label="internal", arguments={"file_id": "123"})
print(f"✓ MCPCall created: {mcp}")
# Test DeepResearchResponse
response = DeepResearchResponse(
report="This is a test report about AI.",
citations=[citation],
reasoning_steps=[reasoning],
web_searches=[search],
code_executions=[code],
mcp_calls=[mcp]
)
print(f"✓ DeepResearchResponse created")
print(f" - Report length: {len(response.report)} chars")
print(f" - Citations: {len(response.citations)}")
print(f" - Sources: {response.get_all_sources()}")
return True
def test_research_query():
"""Test actual research query (requires API key)."""
if not os.environ.get("OPENAI_API_KEY"):
print("⚠ Skipping research test - OPENAI_API_KEY not set")
return True
agent = DeepResearchAgent(
name="Healthcare Researcher",
instructions="""
You are a professional researcher preparing a structured, data-driven report.
Focus on data-rich insights with specific figures and trends.
Include inline citations and prioritize reliable sources.
""",
model="o4-mini-deep-research", # Use faster model for testing
verbose=True,
enable_code_interpreter=False
)
print("\n🔍 Starting deep research query...")
print(" (This may take several minutes)")
try:
result = agent.research(
"What are the key trends in AI adoption in healthcare in 2024?",
summary_mode="concise"
)
print(f"\n✓ Research completed!")
print(f" - Report length: {len(result.report)} characters")
print(f" - Citations found: {len(result.citations)}")
print(f" - Reasoning steps: {len(result.reasoning_steps)}")
print(f" - Web searches: {len(result.web_searches)}")
print(f" - Code executions: {len(result.code_executions)}")
# Print first 500 chars of report
print(f"\n📄 Report excerpt:")
print(f" {result.report[:500]}...")
# Print some citations
if result.citations:
print(f"\n📚 Sample citations:")
for i, c in enumerate(result.citations[:3]):
print(f" {i+1}. {c.title}: {c.url}")
# Print web searches
if result.web_searches:
print(f"\n🔎 Web searches performed:")
for s in result.web_searches[:5]:
print(f" - {s.query} ({s.status})")
return True
except Exception as e:
print(f"✗ Research failed: {e}")
return False
def test_clarify_and_rewrite():
"""Test clarifying questions and query rewriting."""
if not os.environ.get("OPENAI_API_KEY"):
print("⚠ Skipping clarify/rewrite test - OPENAI_API_KEY not set")
return True
agent = DeepResearchAgent(
name="Research Assistant",
verbose=True
)
print("\n🤔 Testing clarifying questions...")
try:
questions = agent.clarify("Help me research AI trends")
print(f"✓ Clarifying questions generated:")
print(f" {questions[:300]}...")
except Exception as e:
print(f"✗ Clarify failed: {e}")
return False
print("\n✍️ Testing query rewriting...")
try:
rewritten = agent.rewrite_query("Research AI in healthcare")
print(f"✓ Query rewritten:")
print(f" {rewritten[:300]}...")
except Exception as e:
print(f"✗ Rewrite failed: {e}")
return False
return True
def test_mcp_configuration():
"""Test MCP server configuration."""
if not os.environ.get("OPENAI_API_KEY"):
print("⚠ Skipping MCP test - OPENAI_API_KEY not set")
return True
# Test agent with MCP configuration
agent = DeepResearchAgent(
name="MCP Researcher",
instructions="You are a researcher with access to internal documents.",
mcp_servers=[
{
"label": "internal_docs",
"url": "https://example.com/mcp/sse/",
"require_approval": "never"
}
],
verbose=True
)
print(f"✓ Agent with MCP configured: {agent}")
print(f" - MCP servers: {len(agent.mcp_servers)}")
# Test tool building
tools = agent._build_tools(
web_search=True,
code_interpreter=True,
mcp_servers=agent.mcp_servers
)
print(f"✓ Tools built: {len(tools)} tools")
for t in tools:
print(f" - {t.get('type', 'unknown')}")
return True
def run_all_tests():
"""Run all tests."""
print("=" * 60)
print("DeepResearchAgent Test Suite")
print("=" * 60)
tests = [
("Import Test", test_basic_import),
("Data Classes Test", test_data_classes),
("Agent Initialization Test", test_agent_initialization),
("MCP Configuration Test", test_mcp_configuration),
("Clarify & Rewrite Test", test_clarify_and_rewrite),
("Research Query Test", test_research_query),
]
results = []
for name, test_fn in tests:
print(f"\n{'─' * 40}")
print(f"Running: {name}")
print(f"{'─' * 40}")
try:
success = test_fn()
results.append((name, success))
except Exception as e:
print(f"✗ Test failed with exception: {e}")
results.append((name, False))
# Summary
print(f"\n{'=' * 60}")
print("Test Summary")
print(f"{'=' * 60}")
passed = sum(1 for _, s in results if s)
total = len(results)
print(f"Passed: {passed}/{total}")
for name, success in results:
status = "✓" if success else "✗"
print(f" {status} {name}")
return passed == total
if __name__ == "__main__":
success = run_all_tests()
sys.exit(0 if success else 1)
@@ -0,0 +1,403 @@
"""
Tests for QueryRewriterAgent
Run with: python -m pytest tests/query_rewriter_test.py -v
"""
import pytest
from unittest.mock import patch
class TestQueryRewriterImports:
"""Test that all imports work correctly."""
def test_import_from_package(self):
"""Test importing from main package."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy, RewriteResult
assert QueryRewriterAgent is not None
assert RewriteStrategy is not None
assert RewriteResult is not None
def test_import_from_agent_module(self):
"""Test importing from agent module."""
from praisonaiagents.agent import QueryRewriterAgent, RewriteStrategy, RewriteResult
assert QueryRewriterAgent is not None
assert RewriteStrategy is not None
assert RewriteResult is not None
class TestRewriteStrategy:
"""Test RewriteStrategy enum."""
def test_strategy_values(self):
"""Test all strategy values exist."""
from praisonaiagents import RewriteStrategy
assert RewriteStrategy.BASIC.value == "basic"
assert RewriteStrategy.HYDE.value == "hyde"
assert RewriteStrategy.STEP_BACK.value == "step_back"
assert RewriteStrategy.SUB_QUERIES.value == "sub_queries"
assert RewriteStrategy.MULTI_QUERY.value == "multi_query"
assert RewriteStrategy.CONTEXTUAL.value == "contextual"
assert RewriteStrategy.AUTO.value == "auto"
class TestRewriteResult:
"""Test RewriteResult dataclass."""
def test_result_creation(self):
"""Test creating a RewriteResult."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="test query",
rewritten_queries=["rewritten test query"],
strategy_used=RewriteStrategy.BASIC
)
assert result.original_query == "test query"
assert result.rewritten_queries == ["rewritten test query"]
assert result.strategy_used == RewriteStrategy.BASIC
def test_primary_query_property(self):
"""Test primary_query property."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="test",
rewritten_queries=["first", "second"],
strategy_used=RewriteStrategy.MULTI_QUERY
)
assert result.primary_query == "first"
def test_primary_query_fallback(self):
"""Test primary_query falls back to original if empty."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="test",
rewritten_queries=[],
strategy_used=RewriteStrategy.BASIC
)
assert result.primary_query == "test"
def test_all_queries_property(self):
"""Test all_queries property combines all queries."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="original",
rewritten_queries=["rewritten"],
strategy_used=RewriteStrategy.STEP_BACK,
step_back_question="step back",
sub_queries=["sub1", "sub2"]
)
all_q = result.all_queries
assert "original" in all_q
assert "rewritten" in all_q
assert "step back" in all_q
assert "sub1" in all_q
assert "sub2" in all_q
class TestQueryRewriterAgentInit:
"""Test QueryRewriterAgent initialization."""
def test_default_initialization(self):
"""Test agent initializes with defaults."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent.name == "QueryRewriterAgent"
assert agent.model == "gpt-4o-mini"
assert not agent.verbose
assert agent.max_queries == 5
assert agent.temperature == 0.3
assert agent.max_tokens == 500
def test_custom_initialization(self):
"""Test agent initializes with custom values."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent(
name="CustomRewriter",
model="gpt-4o",
verbose=True,
max_queries=10,
temperature=0.5,
max_tokens=1000
)
assert agent.name == "CustomRewriter"
assert agent.model == "gpt-4o"
assert agent.verbose
assert agent.max_queries == 10
assert agent.temperature == 0.5
assert agent.max_tokens == 1000
def test_default_abbreviations(self):
"""Test default abbreviations are set."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert "AI" in agent.abbreviations
assert "ML" in agent.abbreviations
assert "RAG" in agent.abbreviations
assert agent.abbreviations["AI"] == "Artificial Intelligence"
class TestAbbreviationExpansion:
"""Test abbreviation expansion functionality."""
def test_expand_abbreviations(self):
"""Test _expand_abbreviations method."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._expand_abbreviations("AI trends")
assert "Artificial Intelligence" in result
def test_add_abbreviation(self):
"""Test adding single abbreviation."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
agent.add_abbreviation("K8s", "Kubernetes")
assert "K8S" in agent.abbreviations
assert agent.abbreviations["K8S"] == "Kubernetes"
def test_add_abbreviations(self):
"""Test adding multiple abbreviations."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
agent.add_abbreviations({
"TF": "TensorFlow",
"PT": "PyTorch"
})
assert "TF" in agent.abbreviations
assert "PT" in agent.abbreviations
class TestQueryDetection:
"""Test query analysis methods."""
def test_is_short_query(self):
"""Test short query detection."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent._is_short_query("AI")
assert agent._is_short_query("AI trends")
assert not agent._is_short_query("What are the latest AI trends in 2025?")
def test_is_follow_up(self):
"""Test follow-up query detection."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent._is_follow_up("What about it?")
assert agent._is_follow_up("And this one?")
assert not agent._is_follow_up("What is AI?")
def test_is_complex_query(self):
"""Test complex query detection."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent._is_complex_query("What is A and what is B?")
assert agent._is_complex_query("Question one? Question two?")
assert not agent._is_complex_query("What is AI?")
class TestStrategyDetection:
"""Test automatic strategy detection."""
def test_detect_basic_for_short(self):
"""Test BASIC detected for short queries."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent()
strategy = agent._detect_strategy("AI")
assert strategy == RewriteStrategy.BASIC
def test_detect_contextual_for_followup(self):
"""Test CONTEXTUAL detected for follow-up with history."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent()
chat_history = [{"role": "user", "content": "Tell me about Python"}]
strategy = agent._detect_strategy("What about it?", chat_history)
assert strategy == RewriteStrategy.CONTEXTUAL
def test_detect_sub_queries_for_complex(self):
"""Test SUB_QUERIES detected for complex queries."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent()
strategy = agent._detect_strategy("What is A and what is B?")
assert strategy == RewriteStrategy.SUB_QUERIES
class TestJSONParsing:
"""Test JSON array parsing from LLM responses."""
def test_parse_valid_json(self):
"""Test parsing valid JSON array."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._parse_json_array('["query1", "query2"]')
assert result == ["query1", "query2"]
def test_parse_json_in_text(self):
"""Test extracting JSON from surrounding text."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._parse_json_array('Here are the queries: ["q1", "q2"]')
assert result == ["q1", "q2"]
def test_parse_fallback_newlines(self):
"""Test fallback to newline splitting."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._parse_json_array("1. Query one\n2. Query two\n3. Query three")
assert len(result) >= 2
class TestRewriteWithMockedLLM:
"""Test rewrite methods with mocked LLM."""
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_basic(self, mock_llm):
"""Test basic rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = "What are the current trends in AI?"
agent = QueryRewriterAgent()
result = agent.rewrite("AI trends", strategy=RewriteStrategy.BASIC)
assert result.strategy_used == RewriteStrategy.BASIC
assert len(result.rewritten_queries) > 0
mock_llm.assert_called_once()
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_hyde(self, mock_llm):
"""Test HyDE rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = "Quantum computing is a type of computation..."
agent = QueryRewriterAgent()
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
assert result.strategy_used == RewriteStrategy.HYDE
assert result.hypothetical_document is not None
mock_llm.assert_called_once()
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_step_back(self, mock_llm):
"""Test step-back rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
# First call for step-back, second for basic
mock_llm.side_effect = [
"What are the fundamentals of language models?",
"What are the differences between GPT-4 and Claude 3?"
]
agent = QueryRewriterAgent()
result = agent.rewrite("GPT-4 vs Claude 3?", strategy=RewriteStrategy.STEP_BACK)
assert result.strategy_used == RewriteStrategy.STEP_BACK
assert result.step_back_question is not None
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_sub_queries(self, mock_llm):
"""Test sub-query decomposition with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = '["What is RAG?", "What are the best embedding models?"]'
agent = QueryRewriterAgent()
result = agent.rewrite("RAG setup and embedding models?", strategy=RewriteStrategy.SUB_QUERIES)
assert result.strategy_used == RewriteStrategy.SUB_QUERIES
assert result.sub_queries is not None
assert len(result.sub_queries) >= 1
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_multi_query(self, mock_llm):
"""Test multi-query generation with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = '["Query 1", "Query 2", "Query 3"]'
agent = QueryRewriterAgent()
result = agent.rewrite("LLM quality?", strategy=RewriteStrategy.MULTI_QUERY, num_queries=3)
assert result.strategy_used == RewriteStrategy.MULTI_QUERY
assert len(result.rewritten_queries) >= 1
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_contextual(self, mock_llm):
"""Test contextual rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = "How does Python's performance compare to other languages?"
agent = QueryRewriterAgent()
chat_history = [
{"role": "user", "content": "Tell me about Python"},
{"role": "assistant", "content": "Python is a programming language..."}
]
result = agent.rewrite(
"What about its performance?",
strategy=RewriteStrategy.CONTEXTUAL,
chat_history=chat_history
)
assert result.strategy_used == RewriteStrategy.CONTEXTUAL
assert "Python" in result.primary_query or len(result.rewritten_queries) > 0
class TestConvenienceMethods:
"""Test convenience methods."""
def test_convenience_methods_exist(self):
"""Test all convenience methods exist."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert hasattr(agent, 'rewrite_basic')
assert hasattr(agent, 'rewrite_hyde')
assert hasattr(agent, 'rewrite_step_back')
assert hasattr(agent, 'rewrite_sub_queries')
assert hasattr(agent, 'rewrite_multi_query')
assert hasattr(agent, 'rewrite_contextual')
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,460 @@
"""
Tests for Claude Memory Tool.
Claude Memory Tool enables Claude to store and retrieve information across
conversations through a memory file directory.
Reference: https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/memory-tool
"""
import os
import shutil
import tempfile
import pytest
from pathlib import Path
class TestClaudeMemoryToolBasics:
"""Tests for ClaudeMemoryTool basic functionality."""
def test_import(self):
"""Test that ClaudeMemoryTool can be imported."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
assert ClaudeMemoryTool is not None
def test_init_default(self):
"""Test default initialization."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
assert tool.memory_root.exists()
assert tool.user_id == "default"
def test_init_with_user_id(self):
"""Test initialization with custom user_id."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir, user_id="test_user")
assert tool.user_id == "test_user"
def test_get_tool_definition(self):
"""Test tool definition for API."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
definition = tool.get_tool_definition()
assert definition["type"] == "memory_20250818"
assert definition["name"] == "memory"
def test_get_beta_header(self):
"""Test beta header value."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
header = tool.get_beta_header()
assert header == "context-management-2025-06-27"
class TestClaudeMemoryToolPathValidation:
"""Tests for path validation and security."""
def test_valid_path(self):
"""Test valid path validation."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
path = tool._validate_path("/memories")
assert path == tool.memory_root
def test_valid_subpath(self):
"""Test valid subpath validation."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
path = tool._validate_path("/memories/notes.txt")
assert path == tool.memory_root / "notes.txt"
def test_invalid_path_no_memories_prefix(self):
"""Test that paths without /memories prefix are rejected."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
with pytest.raises(ValueError, match="must start with /memories"):
tool._validate_path("/etc/passwd")
def test_directory_traversal_blocked(self):
"""Test that directory traversal is blocked."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
with pytest.raises(ValueError):
tool._validate_path("/memories/../../../etc/passwd")
class TestClaudeMemoryToolCommands:
"""Tests for memory tool commands."""
def test_view_empty_directory(self):
"""Test viewing empty directory."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
result = tool.view(path="/memories")
assert "Directory: /memories" in result
assert "(empty directory)" in result
def test_create_and_view_file(self):
"""Test creating and viewing a file."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
# Create file
result = tool.create(path="/memories/notes.txt", file_text="Hello, World!")
assert "created successfully" in result
# View file
result = tool.view(path="/memories/notes.txt")
assert "Hello, World!" in result
def test_view_directory_with_files(self):
"""Test viewing directory with files."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
# Create files
tool.create(path="/memories/file1.txt", file_text="Content 1")
tool.create(path="/memories/file2.txt", file_text="Content 2")
# View directory
result = tool.view(path="/memories")
assert "file1.txt" in result
assert "file2.txt" in result
def test_str_replace(self):
"""Test string replacement in file."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
# Create file
tool.create(path="/memories/notes.txt", file_text="Hello, World!")
# Replace text
result = tool.str_replace(
path="/memories/notes.txt",
old_str="World",
new_str="Claude"
)
assert "edited" in result
# Verify replacement
result = tool.view(path="/memories/notes.txt")
assert "Hello, Claude!" in result
def test_str_replace_not_found(self):
"""Test string replacement when text not found."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
tool.create(path="/memories/notes.txt", file_text="Hello, World!")
with pytest.raises(ValueError, match="not found"):
tool.str_replace(
path="/memories/notes.txt",
old_str="Nonexistent",
new_str="Replacement"
)
def test_insert(self):
"""Test inserting text at line."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
# Create file with multiple lines
tool.create(path="/memories/notes.txt", file_text="Line 1\nLine 2\nLine 3")
# Insert at line 1
result = tool.insert(
path="/memories/notes.txt",
insert_line=1,
insert_text="Inserted Line"
)
assert "inserted" in result
# Verify insertion
result = tool.view(path="/memories/notes.txt")
assert "Inserted Line" in result
def test_delete_file(self):
"""Test deleting a file."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
# Create and delete file
tool.create(path="/memories/notes.txt", file_text="Content")
result = tool.delete(path="/memories/notes.txt")
assert "deleted" in result
# Verify deletion
with pytest.raises(RuntimeError):
tool.view(path="/memories/notes.txt")
def test_delete_root_blocked(self):
"""Test that deleting /memories root is blocked."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
with pytest.raises(ValueError, match="Cannot delete"):
tool.delete(path="/memories")
def test_rename(self):
"""Test renaming a file."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
# Create file
tool.create(path="/memories/old.txt", file_text="Content")
# Rename
result = tool.rename(
old_path="/memories/old.txt",
new_path="/memories/new.txt"
)
assert "Renamed" in result
# Verify rename
with pytest.raises(RuntimeError):
tool.view(path="/memories/old.txt")
result = tool.view(path="/memories/new.txt")
assert "Content" in result
def test_clear_all(self):
"""Test clearing all memory."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
# Create files
tool.create(path="/memories/file1.txt", file_text="Content 1")
tool.create(path="/memories/file2.txt", file_text="Content 2")
# Clear all
result = tool.clear_all()
assert "cleared" in result
# Verify empty
result = tool.view(path="/memories")
assert "(empty directory)" in result
class TestClaudeMemoryToolProcessToolCall:
"""Tests for process_tool_call method."""
def test_process_view_command(self):
"""Test processing view command."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
result = tool.process_tool_call({
"command": "view",
"path": "/memories"
})
assert "Directory: /memories" in result
def test_process_create_command(self):
"""Test processing create command."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
result = tool.process_tool_call({
"command": "create",
"path": "/memories/test.txt",
"file_text": "Test content"
})
assert "created" in result
def test_process_unknown_command(self):
"""Test processing unknown command."""
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
with tempfile.TemporaryDirectory() as tmpdir:
tool = ClaudeMemoryTool(base_path=tmpdir)
result = tool.process_tool_call({
"command": "unknown_command",
"path": "/memories"
})
assert "Error" in result
class TestLLMClaudeMemoryIntegration:
"""Tests for LLM class integration with Claude Memory Tool."""
def test_llm_init_with_claude_memory_true(self):
"""Test LLM initialization with claude_memory=True."""
from praisonaiagents.llm.llm import LLM
llm = LLM(model="anthropic/claude-sonnet-4-20250514", claude_memory=True)
assert llm.claude_memory == True
def test_llm_init_with_claude_memory_false(self):
"""Test LLM initialization with claude_memory=False."""
from praisonaiagents.llm.llm import LLM
llm = LLM(model="anthropic/claude-sonnet-4-20250514", claude_memory=False)
assert llm.claude_memory == False
def test_llm_init_without_claude_memory(self):
"""Test LLM initialization without claude_memory parameter."""
from praisonaiagents.llm.llm import LLM
llm = LLM(model="gpt-4o-mini")
assert llm.claude_memory is None
def test_llm_supports_claude_memory_anthropic(self):
"""Test _supports_claude_memory for Anthropic model."""
from praisonaiagents.llm.llm import LLM
llm = LLM(model="anthropic/claude-sonnet-4-20250514", claude_memory=True)
assert llm._supports_claude_memory() == True
def test_llm_supports_claude_memory_non_anthropic(self):
"""Test _supports_claude_memory for non-Anthropic model."""
from praisonaiagents.llm.llm import LLM
llm = LLM(model="gpt-4o-mini", claude_memory=True)
assert llm._supports_claude_memory() == False
def test_llm_get_claude_memory_tool(self):
"""Test _get_claude_memory_tool method."""
from praisonaiagents.llm.llm import LLM
from praisonaiagents.tools.claude_memory_tool import ClaudeMemoryTool
llm = LLM(model="anthropic/claude-sonnet-4-20250514", claude_memory=True)
tool = llm._get_claude_memory_tool()
assert isinstance(tool, ClaudeMemoryTool)
def test_llm_is_memory_tool_call(self):
"""Test _is_memory_tool_call method."""
from praisonaiagents.llm.llm import LLM
llm = LLM(model="anthropic/claude-sonnet-4-20250514", claude_memory=True)
assert llm._is_memory_tool_call("memory") == True
assert llm._is_memory_tool_call("other_tool") == False
class TestAgentClaudeMemoryIntegration:
"""Tests for Agent class integration with Claude Memory Tool."""
def test_agent_init_with_claude_memory_true(self):
"""Test Agent initialization with claude_memory=True."""
from praisonaiagents import Agent
agent = Agent(
name="Test Agent",
llm="anthropic/claude-sonnet-4-20250514",
claude_memory=True
)
assert agent.claude_memory == True
def test_agent_init_with_claude_memory_false(self):
"""Test Agent initialization with claude_memory=False."""
from praisonaiagents import Agent
agent = Agent(
name="Test Agent",
llm="anthropic/claude-sonnet-4-20250514",
claude_memory=False
)
assert agent.claude_memory == False
def test_agent_init_without_claude_memory(self):
"""Test Agent initialization without claude_memory parameter."""
from praisonaiagents import Agent
agent = Agent(name="Test Agent")
assert agent.claude_memory is None
def test_agent_passes_claude_memory_to_llm(self):
"""Test that Agent passes claude_memory to LLM instance."""
from praisonaiagents import Agent
agent = Agent(
name="Test Agent",
llm="anthropic/claude-sonnet-4-20250514",
claude_memory=True
)
if hasattr(agent, 'llm_instance') and agent.llm_instance:
assert agent.llm_instance.claude_memory == True
class TestClaudeMemoryToolIntegration:
"""Integration tests for Claude Memory Tool (requires API key)."""
@pytest.mark.skip(reason="Requires ANTHROPIC_API_KEY environment variable")
def test_claude_memory_with_anthropic(self):
"""Test Claude memory with actual Anthropic API."""
import os
from praisonaiagents import Agent
# Skip if no API key
if not os.getenv("ANTHROPIC_API_KEY"):
pytest.skip("ANTHROPIC_API_KEY not set")
agent = Agent(
name="Memory Test Agent",
instructions="You are a helpful assistant with memory capabilities.",
llm="anthropic/claude-sonnet-4-20250514",
claude_memory=True,
verbose=False
)
# First call - should check memory
result1 = agent.start("Remember that my favorite color is blue.")
assert result1 is not None
# Second call - should recall from memory
result2 = agent.start("What is my favorite color?")
assert result2 is not None
@@ -11,6 +11,7 @@ This test script validates the ContextAgent implementation for:
"""
import sys
import pytest
from pathlib import Path
# Add the src directory to Python path
@@ -34,11 +35,11 @@ def test_imports():
from praisonaiagents.agent.context_agent import ContextAgent as DirectContextAgent
print("✅ Successfully imported ContextAgent from direct module")
return True
assert True
except ImportError as e:
print(f"❌ Import failed: {e}")
return False
assert False, str(e)
def test_basic_instantiation():
"""Test basic ContextAgent instantiation."""
@@ -64,11 +65,11 @@ def test_basic_instantiation():
factory_agent = create_context_agent(llm="gpt-5-nano")
print("✅ Successfully created ContextAgent using factory function")
return True, [context_agent, custom_agent, factory_agent]
assert True
except Exception as e:
print(f"❌ Instantiation failed: {e}")
return False, []
assert False, str(e)
def test_agent_inheritance():
"""Test that ContextAgent properly inherits from Agent."""
@@ -89,12 +90,13 @@ def test_agent_inheritance():
assert hasattr(context_agent, 'goal'), "Should have goal attribute"
print("✅ ContextAgent has all required Agent attributes")
return True
assert True
except Exception as e:
print(f"❌ Inheritance test failed: {e}")
return False
assert False, str(e)
@pytest.mark.skip(reason="ContextAgent methods not yet implemented")
def test_context_engineering_methods():
"""Test Context Engineering specific methods."""
print("\\n🧪 Testing Context Engineering Methods...")
@@ -137,12 +139,13 @@ def test_context_engineering_methods():
print(f"✅ All {len(async_methods)} async methods are available")
return True
assert True
except Exception as e:
print(f"❌ Method availability test failed: {e}")
return False
assert False, str(e)
@pytest.mark.skip(reason="ContextAgent methods not yet implemented")
def test_basic_functionality():
"""Test basic functionality of ContextAgent methods."""
print("\\n🧪 Testing Basic Functionality...")
@@ -194,11 +197,11 @@ def test_basic_functionality():
assert "PRP" in prp, "Should contain PRP reference"
print("✅ generate_prp works correctly")
return True
assert True
except Exception as e:
print(f"❌ Functionality test failed: {e}")
return False
assert False, str(e)
def test_backward_compatibility():
"""Test that existing PraisonAI functionality still works."""
@@ -232,11 +235,11 @@ def test_backward_compatibility():
print("✅ All expected exports are available")
return True
assert True
except Exception as e:
print(f"❌ Backward compatibility test failed: {e}")
return False
assert False, str(e)
def test_syntax_validation():
"""Test that all Python files have valid syntax."""
@@ -268,14 +271,14 @@ def test_syntax_validation():
ast.parse(content)
print(f"✅ {example_file.name} has valid Python syntax")
return True
assert True
except SyntaxError as e:
print(f"❌ Syntax error found: {e}")
return False
assert False, str(e)
except Exception as e:
print(f"❌ Syntax validation failed: {e}")
return False
assert False, str(e)
def run_all_tests():
"""Run all tests and provide summary."""
@@ -0,0 +1,829 @@
"""
Tests for FileMemory - Zero-dependency file-based memory system.
"""
import tempfile
import pytest
class TestFileMemoryBasics:
"""Tests for FileMemory basic functionality."""
def test_import(self):
"""Test that FileMemory can be imported."""
from praisonaiagents.memory.file_memory import FileMemory
assert FileMemory is not None
def test_init_default(self):
"""Test default initialization."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
assert memory.user_id == "default"
assert memory.user_path.exists()
def test_init_with_user_id(self):
"""Test initialization with custom user_id."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir, user_id="test_user")
assert memory.user_id == "test_user"
assert "test_user" in str(memory.user_path)
def test_config_persistence(self):
"""Test that config is saved and loaded."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
# Create with custom config
memory1 = FileMemory(
base_path=tmpdir,
config={"short_term_limit": 50}
)
assert memory1.config["short_term_limit"] == 50
# Create new instance - should load saved config
memory2 = FileMemory(base_path=tmpdir)
assert memory2.config["short_term_limit"] == 50
class TestShortTermMemory:
"""Tests for short-term memory functionality."""
def test_add_short_term(self):
"""Test adding short-term memory."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
mem_id = memory.add_short_term("User prefers dark mode")
assert mem_id is not None
assert len(memory._short_term) == 1
def test_get_short_term(self):
"""Test retrieving short-term memories."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_short_term("Memory 1")
memory.add_short_term("Memory 2")
memory.add_short_term("Memory 3")
# Get all
items = memory.get_short_term()
assert len(items) == 3
# Most recent first
assert items[0].content == "Memory 3"
# Get limited
items = memory.get_short_term(limit=2)
assert len(items) == 2
def test_short_term_rolling_buffer(self):
"""Test that short-term memory enforces limit."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(
base_path=tmpdir,
config={"short_term_limit": 5, "auto_promote": False}
)
# Add more than limit
for i in range(10):
memory.add_short_term(f"Memory {i}")
# Should only have last 5
assert len(memory._short_term) == 5
items = memory.get_short_term()
assert items[0].content == "Memory 9"
def test_short_term_persistence(self):
"""Test that short-term memory persists to file."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
# Add memory
memory1 = FileMemory(base_path=tmpdir)
memory1.add_short_term("Persistent memory")
# Create new instance - should load from file
memory2 = FileMemory(base_path=tmpdir)
items = memory2.get_short_term()
assert len(items) == 1
assert items[0].content == "Persistent memory"
class TestLongTermMemory:
"""Tests for long-term memory functionality."""
def test_add_long_term(self):
"""Test adding long-term memory."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
mem_id = memory.add_long_term("User's name is John", importance=0.9)
assert mem_id is not None
assert len(memory._long_term) == 1
def test_get_long_term_sorted_by_importance(self):
"""Test that long-term memories are sorted by importance."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_long_term("Low importance", importance=0.3)
memory.add_long_term("High importance", importance=0.9)
memory.add_long_term("Medium importance", importance=0.6)
items = memory.get_long_term()
assert items[0].content == "High importance"
assert items[1].content == "Medium importance"
assert items[2].content == "Low importance"
def test_auto_promote_to_long_term(self):
"""Test auto-promotion of high-importance short-term to long-term."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(
base_path=tmpdir,
config={
"short_term_limit": 3,
"auto_promote": True,
"importance_threshold": 0.7
}
)
# Add high importance items
memory.add_short_term("Important 1", importance=0.8)
memory.add_short_term("Important 2", importance=0.9)
memory.add_short_term("Low importance", importance=0.3)
# Trigger promotion by exceeding limit
memory.add_short_term("New item", importance=0.5)
# High importance items should be in long-term
long_term = memory.get_long_term()
assert len(long_term) >= 2
class TestEntityMemory:
"""Tests for entity memory functionality."""
def test_add_entity(self):
"""Test adding an entity."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
entity_id = memory.add_entity(
name="John",
entity_type="person",
attributes={"role": "developer", "company": "Acme"}
)
assert entity_id is not None
assert len(memory._entities) == 1
def test_get_entity(self):
"""Test retrieving an entity."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_entity("John", "person", {"role": "developer"})
entity = memory.get_entity("John")
assert entity is not None
assert entity.name == "John"
assert entity.entity_type == "person"
assert entity.attributes["role"] == "developer"
def test_update_entity(self):
"""Test updating an existing entity."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_entity("John", "person", {"role": "developer"})
memory.add_entity("John", "person", {"company": "Acme"})
entity = memory.get_entity("John")
assert entity.attributes["role"] == "developer"
assert entity.attributes["company"] == "Acme"
def test_get_all_entities(self):
"""Test getting all entities."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_entity("John", "person", {})
memory.add_entity("Acme", "organization", {})
memory.add_entity("New York", "place", {})
# All entities
all_entities = memory.get_all_entities()
assert len(all_entities) == 3
# Filter by type
people = memory.get_all_entities(entity_type="person")
assert len(people) == 1
assert people[0].name == "John"
class TestEpisodicMemory:
"""Tests for episodic memory functionality."""
def test_add_episodic(self):
"""Test adding episodic memory."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
mem_id = memory.add_episodic("Had a meeting about project X")
assert mem_id is not None
def test_get_episodic_by_date(self):
"""Test retrieving episodic memories by date."""
from praisonaiagents.memory.file_memory import FileMemory
from datetime import datetime
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
today = datetime.now().strftime("%Y-%m-%d")
memory.add_episodic("Today's event", date=today)
episodes = memory.get_episodic(date=today)
assert len(episodes) == 1
assert episodes[0].content == "Today's event"
def test_get_recent_episodic(self):
"""Test retrieving recent episodic memories."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_episodic("Event 1")
memory.add_episodic("Event 2")
episodes = memory.get_episodic(days_back=7)
assert len(episodes) == 2
class TestSearch:
"""Tests for search functionality."""
def test_search_short_term(self):
"""Test searching short-term memory."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_short_term("User prefers dark mode")
memory.add_short_term("User likes Python programming")
memory.add_short_term("Weather is sunny today")
results = memory.search("user", memory_types=["short_term"])
assert len(results) == 2
def test_search_all_types(self):
"""Test searching across all memory types."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_short_term("John mentioned Python")
memory.add_long_term("John is a developer")
memory.add_entity("John", "person", {"skill": "Python"})
results = memory.search("John")
assert len(results) >= 3
def test_search_relevance_scoring(self):
"""Test that search results are scored by relevance."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_short_term("Python is great", importance=0.5)
memory.add_short_term("I love Python programming", importance=0.8)
results = memory.search("Python")
assert len(results) == 2
# Higher importance should score higher
assert results[0]["score"] >= results[1]["score"]
class TestContext:
"""Tests for context building."""
def test_get_context_empty(self):
"""Test getting context with no memories."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
context = memory.get_context()
assert context == ""
def test_get_context_with_memories(self):
"""Test getting context with memories."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_long_term("User's name is John", importance=0.9)
memory.add_short_term("User asked about Python")
memory.add_entity("John", "person", {"role": "developer"})
context = memory.get_context()
assert "John" in context
assert "Important Facts" in context or "Known Entities" in context
def test_get_context_with_query(self):
"""Test getting focused context with query."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_long_term("User's name is John")
memory.add_long_term("User likes Python")
memory.add_long_term("Weather is sunny")
context = memory.get_context(query="John")
assert "John" in context
class TestUtilities:
"""Tests for utility methods."""
def test_clear_short_term(self):
"""Test clearing short-term memory."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_short_term("Memory 1")
memory.add_short_term("Memory 2")
memory.clear_short_term()
assert len(memory._short_term) == 0
def test_clear_all(self):
"""Test clearing all memory."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir)
memory.add_short_term("Short term")
memory.add_long_term("Long term")
memory.add_entity("John", "person", {})
memory.add_episodic("Episode")
memory.clear_all()
assert len(memory._short_term) == 0
assert len(memory._long_term) == 0
assert len(memory._entities) == 0
def test_get_stats(self):
"""Test getting memory statistics."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
memory = FileMemory(base_path=tmpdir, user_id="test_user")
memory.add_short_term("Short term")
memory.add_long_term("Long term")
memory.add_entity("John", "person", {})
stats = memory.get_stats()
assert stats["user_id"] == "test_user"
assert stats["short_term_count"] == 1
assert stats["long_term_count"] == 1
assert stats["entity_count"] == 1
def test_export_import(self):
"""Test exporting and importing memory data."""
from praisonaiagents.memory.file_memory import FileMemory
with tempfile.TemporaryDirectory() as tmpdir:
# Create and populate memory
memory1 = FileMemory(base_path=tmpdir, user_id="user1")
memory1.add_short_term("Short term memory")
memory1.add_long_term("Long term memory")
memory1.add_entity("John", "person", {"role": "dev"})
# Export
exported = memory1.export()
# Import to new memory
memory2 = FileMemory(base_path=tmpdir, user_id="user2")
memory2.import_data(exported)
assert len(memory2._short_term) == 1
assert len(memory2._long_term) == 1
assert len(memory2._entities) == 1
class TestAgentIntegration:
"""Tests for Agent class integration with FileMemory."""
def test_agent_memory_true(self):
"""Test Agent with memory=True uses FileMemory."""
from praisonaiagents import Agent
agent = Agent(
name="Test Agent",
memory=True
)
assert agent.memory is True
assert agent._memory_instance is not None
def test_agent_memory_false(self):
"""Test Agent with memory=False has no memory."""
from praisonaiagents import Agent
agent = Agent(
name="Test Agent",
memory=False
)
assert agent._memory_instance is None
def test_agent_memory_file(self):
"""Test Agent with memory='file' uses FileMemory."""
from praisonaiagents import Agent
from praisonaiagents.memory.file_memory import FileMemory
agent = Agent(
name="Test Agent",
memory="file"
)
assert isinstance(agent._memory_instance, FileMemory)
def test_agent_store_memory(self):
"""Test Agent store_memory method."""
from praisonaiagents import Agent
agent = Agent(
name="Test Agent",
memory=True
)
agent.store_memory("User prefers dark mode", memory_type="short_term")
items = agent._memory_instance.get_short_term()
assert len(items) == 1
assert items[0].content == "User prefers dark mode"
def test_agent_get_memory_context(self):
"""Test Agent get_memory_context method."""
from praisonaiagents import Agent
agent = Agent(
name="Test Agent",
memory=True
)
agent.store_memory("User's name is John", memory_type="long_term", importance=0.9)
context = agent.get_memory_context()
assert "John" in context
class TestCreateMemoryFunction:
"""Tests for create_memory convenience function."""
def test_create_memory(self):
"""Test create_memory function."""
from praisonaiagents.memory import create_memory
with tempfile.TemporaryDirectory() as tmpdir:
memory = create_memory(user_id="test", base_path=tmpdir)
assert memory.user_id == "test"
memory.add_short_term("Test memory")
assert len(memory._short_term) == 1
class TestMemoryIntegrationWithOpenAI:
"""Integration tests for memory with OpenAI API (requires OPENAI_API_KEY)."""
def test_memory_injection_into_system_prompt(self):
"""Test that memory context is injected into system prompt and agent recalls it."""
import os
import shutil
# Skip if no API key
if not os.getenv("OPENAI_API_KEY"):
pytest.skip("OPENAI_API_KEY not set")
from praisonaiagents import Agent
test_user_id = "integration_test_user"
try:
# Create agent with memory enabled
agent = Agent(
name="Memory Test Agent",
role="Assistant",
goal="Help users and remember their information",
backstory="You are a helpful assistant with memory capabilities.",
# Uses default model (gpt-4o-mini via OPENAI_MODEL_NAME or fallback)
memory=True,
user_id=test_user_id,
verbose=False
)
# Store memories BEFORE the chat
agent.store_memory("User's favorite color is blue", memory_type="short_term")
agent.store_memory("User's name is Alice", memory_type="long_term", importance=0.95)
agent.store_memory("User works as a software engineer", memory_type="long_term", importance=0.9)
# Add entity
agent._memory_instance.add_entity("Alice", "person", {
"occupation": "software engineer",
"favorite_color": "blue"
})
# Ask agent about the user - should recall from memory
result = agent.start("What is my name and what do I do for work? Answer in one sentence.")
# Verify the response contains the remembered information
assert result is not None
result_lower = result.lower()
assert "alice" in result_lower, f"Expected 'alice' in response: {result}"
assert "software" in result_lower or "engineer" in result_lower, f"Expected occupation in response: {result}"
print(f"✅ Memory injection test passed! Response: {result}")
finally:
# Cleanup
memory_path = f".praison/memory/{test_user_id}"
if os.path.exists(memory_path):
shutil.rmtree(memory_path)
def test_memory_persistence_across_agent_instances(self):
"""Test that memory persists when creating new agent instances."""
import os
import shutil
# Skip if no API key
if not os.getenv("OPENAI_API_KEY"):
pytest.skip("OPENAI_API_KEY not set")
from praisonaiagents import Agent
test_user_id = "persistence_test_user"
try:
# First agent - store memories
agent1 = Agent(
name="Agent 1",
# Uses default model (gpt-4o-mini via OPENAI_MODEL_NAME or fallback)
memory=True,
user_id=test_user_id,
verbose=False
)
agent1.store_memory("User prefers Python over JavaScript", memory_type="long_term", importance=0.9)
agent1.store_memory("User is learning machine learning", memory_type="short_term")
# Verify memories stored
stats1 = agent1._memory_instance.get_stats()
assert stats1["long_term_count"] == 1
assert stats1["short_term_count"] == 1
# Second agent - same user_id, should load memories
agent2 = Agent(
name="Agent 2",
# Uses default model (gpt-4o-mini via OPENAI_MODEL_NAME or fallback)
memory=True,
user_id=test_user_id,
verbose=False
)
# Verify memories loaded
stats2 = agent2._memory_instance.get_stats()
assert stats2["long_term_count"] == 1, f"Expected 1 long-term memory, got {stats2}"
assert stats2["short_term_count"] == 1, f"Expected 1 short-term memory, got {stats2}"
# Ask about preferences - should recall from loaded memory
result = agent2.start("What programming language do I prefer? Answer in one word.")
assert result is not None
assert "python" in result.lower(), f"Expected 'python' in response: {result}"
print(f"✅ Memory persistence test passed! Response: {result}")
finally:
# Cleanup
memory_path = f".praison/memory/{test_user_id}"
if os.path.exists(memory_path):
shutil.rmtree(memory_path)
class TestMultiAgentMemorySharing:
"""Tests for memory sharing between multiple agents."""
def test_agents_share_memory_with_same_user_id(self):
"""Test that agents with same user_id share memory."""
from praisonaiagents import Agent
with tempfile.TemporaryDirectory() as tmpdir:
user_id = "shared_user"
base_path = f"{tmpdir}/memory"
# Agent 1 stores memories
agent1 = Agent(
name="Agent1",
memory={"provider": "file", "user_id": user_id, "base_path": base_path},
verbose=False
)
agent1.store_memory("User prefers Python", memory_type="short_term")
agent1.store_memory("User name is Alice", memory_type="long_term", importance=0.9)
agent1._memory_instance.add_entity("Alice", "person", {"role": "developer"})
# Verify Agent1 stored memories
assert len(agent1._memory_instance.get_short_term()) == 1
assert len(agent1._memory_instance.get_long_term()) == 1
assert len(agent1._memory_instance.get_all_entities()) == 1
# Agent 2 with same user_id should load memories
agent2 = Agent(
name="Agent2",
memory={"provider": "file", "user_id": user_id, "base_path": base_path},
verbose=False
)
# Verify Agent2 loaded the same memories
assert len(agent2._memory_instance.get_short_term()) == 1
assert len(agent2._memory_instance.get_long_term()) == 1
assert len(agent2._memory_instance.get_all_entities()) == 1
# Verify content matches
assert agent2._memory_instance.get_short_term()[0].content == "User prefers Python"
assert agent2._memory_instance.get_long_term()[0].content == "User name is Alice"
assert agent2._memory_instance.get_entity("Alice") is not None
def test_agents_isolated_with_different_user_id(self):
"""Test that agents with different user_id have isolated memory."""
from praisonaiagents import Agent
with tempfile.TemporaryDirectory() as tmpdir:
base_path = f"{tmpdir}/memory"
# Agent 1 with user_id "user1"
agent1 = Agent(
name="Agent1",
memory={"provider": "file", "user_id": "user1", "base_path": base_path},
verbose=False
)
agent1.store_memory("User1 data", memory_type="long_term", importance=0.9)
# Agent 2 with different user_id "user2"
agent2 = Agent(
name="Agent2",
memory={"provider": "file", "user_id": "user2", "base_path": base_path},
verbose=False
)
# Agent2 should have no memories (different user_id)
assert len(agent2._memory_instance.get_long_term()) == 0
assert len(agent2._memory_instance.get_short_term()) == 0
# Verify paths are different
assert agent1._memory_instance.user_path != agent2._memory_instance.user_path
def test_memory_context_injected_into_system_prompt(self):
"""Test that memory context is injected into system prompt."""
from praisonaiagents import Agent
with tempfile.TemporaryDirectory() as tmpdir:
agent = Agent(
name="TestAgent",
memory={"provider": "file", "user_id": "test", "base_path": f"{tmpdir}/memory"},
verbose=False
)
# Store memories
agent.store_memory("User prefers dark mode", memory_type="short_term")
agent.store_memory("User name is John", memory_type="long_term", importance=0.9)
# Get memory context
context = agent.get_memory_context()
# Verify context contains the memories
assert "John" in context
assert "dark mode" in context
assert "Important Facts" in context or "Recent Context" in context
def test_memory_display_info(self):
"""Test that memory display info shows correct counts."""
from praisonaiagents import Agent
with tempfile.TemporaryDirectory() as tmpdir:
agent = Agent(
name="TestAgent",
memory={"provider": "file", "user_id": "test", "base_path": f"{tmpdir}/memory"},
verbose=True
)
# Store various memories
agent.store_memory("Short term 1", memory_type="short_term")
agent.store_memory("Short term 2", memory_type="short_term")
agent.store_memory("Long term 1", memory_type="long_term", importance=0.9)
agent._memory_instance.add_entity("Entity1", "person", {})
# Get stats
stats = agent._memory_instance.get_stats()
# Verify counts are as expected (at least the ones we added)
assert stats["short_term_count"] >= 2
assert stats["long_term_count"] >= 1
assert stats["entity_count"] >= 1
class TestMultiAgentMemoryIntegration:
"""Integration tests for multi-agent memory with OpenAI API."""
def test_multi_agent_memory_sharing_with_llm(self):
"""Test that multiple agents share memory and use it in LLM calls."""
import os
import shutil
# Skip if no API key
if not os.getenv("OPENAI_API_KEY"):
pytest.skip("OPENAI_API_KEY not set")
from praisonaiagents import Agent
test_user_id = "multi_agent_llm_test"
try:
# Agent 1 stores memories
agent1 = Agent(
name="Agent1",
memory=True,
user_id=test_user_id,
verbose=False
)
agent1.store_memory("User's favorite color is blue", memory_type="long_term", importance=0.9)
agent1.store_memory("User works as a teacher", memory_type="long_term", importance=0.85)
# Agent 2 with same user_id should recall memories
agent2 = Agent(
name="Agent2",
memory=True,
user_id=test_user_id,
verbose=False
)
# Verify memories loaded
stats = agent2._memory_instance.get_stats()
assert stats["long_term_count"] == 2
# Ask Agent2 about the user - should use loaded memories
result = agent2.start("What is my favorite color? Answer in one word.")
assert result is not None
assert "blue" in result.lower(), f"Expected 'blue' in response: {result}"
print(f"✅ Multi-agent memory sharing test passed! Response: {result}")
finally:
memory_path = f".praison/memory/{test_user_id}"
if os.path.exists(memory_path):
shutil.rmtree(memory_path)
@@ -0,0 +1,122 @@
"""Tests for MCP Backward Compatibility - TDD approach.
These tests ensure backward compatibility with older MCP servers
and the deprecated HTTP+SSE transport from protocol version 2024-11-05.
Backward compatibility features:
- Legacy SSE endpoint support (/sse suffix)
- Fallback detection for older servers
- Automatic transport negotiation
"""
import pytest
class TestLegacySSESupport:
"""Test legacy SSE transport support."""
def test_sse_url_detection(self):
"""Test that /sse URLs are detected as legacy SSE."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/sse") == "sse"
assert get_transport_type("https://example.com/api/sse") == "sse"
def test_non_sse_http_uses_streamable(self):
"""Test that non-/sse HTTP URLs use Streamable HTTP."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/mcp") == "http_stream"
assert get_transport_type("https://example.com/api") == "http_stream"
class TestTransportNegotiation:
"""Test automatic transport negotiation."""
def test_detect_streamable_http_support(self):
"""Test detection of Streamable HTTP support."""
from praisonaiagents.mcp.mcp_compat import detect_transport_support
# Successful POST to MCP endpoint indicates Streamable HTTP
result = detect_transport_support(
response_status=200,
content_type="application/json"
)
assert result == "http_stream"
def test_detect_legacy_sse_fallback(self):
"""Test fallback to legacy SSE on specific error codes."""
from praisonaiagents.mcp.mcp_compat import detect_transport_support
# 400, 404, 405 indicate need for legacy SSE
assert detect_transport_support(response_status=400) == "sse"
assert detect_transport_support(response_status=404) == "sse"
assert detect_transport_support(response_status=405) == "sse"
def test_detect_sse_from_event_stream(self):
"""Test detection of SSE from content type."""
from praisonaiagents.mcp.mcp_compat import detect_transport_support
result = detect_transport_support(
response_status=200,
content_type="text/event-stream"
)
# SSE response indicates legacy transport
assert result in ["sse", "http_stream"]
class TestProtocolVersionCompatibility:
"""Test protocol version compatibility."""
def test_supports_2024_11_05(self):
"""Test support for protocol version 2024-11-05."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version("2024-11-05") is True
def test_supports_2025_03_26(self):
"""Test support for protocol version 2025-03-26."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version("2025-03-26") is True
def test_supports_2025_11_25(self):
"""Test support for protocol version 2025-11-25."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version("2025-11-25") is True
class TestMCPClassBackwardCompat:
"""Test MCP class backward compatibility."""
def test_mcp_accepts_sse_url(self):
"""Test MCP class accepts /sse URLs."""
# This is a design test - actual connection would require server
from praisonaiagents.mcp.mcp_transport import get_transport_type
# Verify URL would be routed to SSE transport
assert get_transport_type("http://localhost:8080/sse") == "sse"
def test_mcp_accepts_command_string(self):
"""Test MCP class accepts command strings (stdio)."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("python server.py") == "stdio"
assert get_transport_type("npx @mcp/server") == "stdio"
def test_mcp_accepts_http_url(self):
"""Test MCP class accepts HTTP URLs."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/mcp") == "http_stream"
def test_mcp_accepts_websocket_url(self):
"""Test MCP class accepts WebSocket URLs."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("ws://localhost:8080/mcp") == "websocket"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,227 @@
"""Tests for MCP Resumability - TDD approach.
These tests define the expected behavior of SSE stream resumability
per MCP Protocol Revision 2025-11-25.
Resumability features:
- SSE event ID tracking
- Last-Event-ID header support
- Message replay on reconnection
- Retry field handling
"""
import pytest
class TestResumabilityManager:
"""Test ResumabilityManager class."""
def test_manager_creation(self):
"""Test ResumabilityManager can be created."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
assert manager.get_retry_delay() == 3000 # Default 3 seconds
def test_set_and_get_last_event_id(self):
"""Test setting and getting last event ID for a stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
# Initially no event ID
assert manager.get_last_event_id("stream-1") is None
# Set event ID
manager.set_last_event_id("stream-1", "event-123")
assert manager.get_last_event_id("stream-1") == "event-123"
def test_multiple_streams(self):
"""Test tracking multiple streams independently."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-a")
manager.set_last_event_id("stream-2", "event-b")
assert manager.get_last_event_id("stream-1") == "event-a"
assert manager.get_last_event_id("stream-2") == "event-b"
def test_update_event_id(self):
"""Test updating event ID for a stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-1")
manager.set_last_event_id("stream-1", "event-2")
assert manager.get_last_event_id("stream-1") == "event-2"
def test_set_retry_delay(self):
"""Test setting retry delay from SSE retry field."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_retry_delay(5000)
assert manager.get_retry_delay() == 5000
def test_get_resume_headers_with_event_id(self):
"""Test getting resume headers when event ID is available."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-123")
headers = manager.get_resume_headers("stream-1")
assert headers == {'Last-Event-ID': 'event-123'}
def test_get_resume_headers_without_event_id(self):
"""Test getting resume headers when no event ID is stored."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
headers = manager.get_resume_headers("stream-1")
assert headers == {}
def test_clear_stream(self):
"""Test clearing tracking for a specific stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-a")
manager.set_last_event_id("stream-2", "event-b")
manager.clear_stream("stream-1")
assert manager.get_last_event_id("stream-1") is None
assert manager.get_last_event_id("stream-2") == "event-b"
def test_clear_all(self):
"""Test clearing all stream tracking."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-a")
manager.set_last_event_id("stream-2", "event-b")
manager.clear_all()
assert manager.get_last_event_id("stream-1") is None
assert manager.get_last_event_id("stream-2") is None
class TestSSEEventParsing:
"""Test SSE event parsing utilities."""
def test_parse_sse_event_with_id(self):
"""Test parsing SSE event with ID field."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "id: event-123\ndata: {\"test\": 1}\n\n"
event = parse_sse_event(event_text)
assert event['id'] == 'event-123'
assert event['data'] == '{"test": 1}'
def test_parse_sse_event_with_retry(self):
"""Test parsing SSE event with retry field."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "retry: 5000\ndata: {}\n\n"
event = parse_sse_event(event_text)
assert event['retry'] == 5000
def test_parse_sse_event_data_only(self):
"""Test parsing SSE event with only data."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "data: {\"message\": \"hello\"}\n\n"
event = parse_sse_event(event_text)
assert event['data'] == '{"message": "hello"}'
assert event.get('id') is None
def test_parse_sse_event_multiline_data(self):
"""Test parsing SSE event with multiline data."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "data: line1\ndata: line2\n\n"
event = parse_sse_event(event_text)
assert event['data'] == 'line1\nline2'
def test_parse_sse_event_empty_data(self):
"""Test parsing SSE event with empty data (for priming)."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "id: prime-123\ndata: \n\n"
event = parse_sse_event(event_text)
assert event['id'] == 'prime-123'
assert event['data'] == ''
class TestEventStore:
"""Test optional event store for message redelivery."""
def test_event_store_creation(self):
"""Test EventStore can be created."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore(max_events=100)
assert store.max_events == 100
def test_store_and_retrieve_event(self):
"""Test storing and retrieving events."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore()
store.store("stream-1", "event-1", {"data": "test1"})
store.store("stream-1", "event-2", {"data": "test2"})
events = store.get_events_after("stream-1", "event-1")
assert len(events) == 1
assert events[0] == {"data": "test2"}
def test_store_respects_max_events(self):
"""Test that store respects max_events limit."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore(max_events=3)
for i in range(5):
store.store("stream-1", f"event-{i}", {"data": i})
# Should only have last 3 events
all_events = store.get_all_events("stream-1")
assert len(all_events) <= 3
def test_store_per_stream_isolation(self):
"""Test that events are isolated per stream."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore()
store.store("stream-1", "event-a", {"data": "a"})
store.store("stream-2", "event-b", {"data": "b"})
events_1 = store.get_all_events("stream-1")
events_2 = store.get_all_events("stream-2")
assert len(events_1) == 1
assert len(events_2) == 1
assert events_1[0] == {"data": "a"}
assert events_2[0] == {"data": "b"}
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,224 @@
"""Tests for MCP Security Features - TDD approach.
These tests define the expected behavior of security features
per MCP Protocol Revision 2025-11-25.
Security features:
- Origin header validation (DNS rebinding prevention)
- Localhost binding recommendation
- Authentication header support
- Secure session ID handling
"""
import pytest
class TestOriginValidation:
"""Test Origin header validation for DNS rebinding prevention."""
def test_valid_localhost_origins(self):
"""Test that localhost origins are accepted."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
assert is_valid_origin("http://localhost:8080", allowed_origins=["localhost"]) is True
assert is_valid_origin("http://127.0.0.1:3000", allowed_origins=["127.0.0.1"]) is True
assert is_valid_origin("http://localhost", allowed_origins=["localhost"]) is True
def test_valid_custom_origins(self):
"""Test that configured custom origins are accepted."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
allowed = ["example.com", "api.example.com"]
assert is_valid_origin("https://example.com", allowed_origins=allowed) is True
assert is_valid_origin("https://api.example.com/path", allowed_origins=allowed) is True
def test_invalid_origins_rejected(self):
"""Test that invalid origins are rejected."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
allowed = ["localhost", "example.com"]
assert is_valid_origin("https://evil.com", allowed_origins=allowed) is False
assert is_valid_origin("https://example.com.evil.com", allowed_origins=allowed) is False
def test_missing_origin_handling(self):
"""Test handling of missing Origin header."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
# Missing origin should be rejected by default
assert is_valid_origin(None, allowed_origins=["localhost"]) is False
assert is_valid_origin("", allowed_origins=["localhost"]) is False
def test_allow_missing_origin_option(self):
"""Test option to allow missing Origin header."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
# Can optionally allow missing origin for non-browser clients
assert is_valid_origin(None, allowed_origins=["localhost"], allow_missing=True) is True
class TestLocalhostBinding:
"""Test localhost binding recommendations."""
def test_is_localhost_address(self):
"""Test detection of localhost addresses."""
from praisonaiagents.mcp.mcp_security import is_localhost_address
assert is_localhost_address("127.0.0.1") is True
assert is_localhost_address("localhost") is True
assert is_localhost_address("::1") is True # IPv6 localhost
assert is_localhost_address("0.0.0.0") is False # All interfaces
assert is_localhost_address("192.168.1.1") is False
def test_should_bind_localhost_only(self):
"""Test recommendation for localhost-only binding."""
from praisonaiagents.mcp.mcp_security import should_bind_localhost_only
# Local servers should bind to localhost only
assert should_bind_localhost_only(is_local=True) is True
# Remote/production servers may bind to all interfaces
assert should_bind_localhost_only(is_local=False) is False
class TestAuthenticationHeaders:
"""Test authentication header support."""
def test_create_bearer_auth_header(self):
"""Test creating Bearer authentication header."""
from praisonaiagents.mcp.mcp_security import create_auth_header
header = create_auth_header("my-secret-token", auth_type="bearer")
assert header == {"Authorization": "Bearer my-secret-token"}
def test_create_basic_auth_header(self):
"""Test creating Basic authentication header."""
from praisonaiagents.mcp.mcp_security import create_auth_header
import base64
header = create_auth_header("user:pass", auth_type="basic")
expected = base64.b64encode(b"user:pass").decode()
assert header == {"Authorization": f"Basic {expected}"}
def test_create_custom_auth_header(self):
"""Test creating custom authentication header."""
from praisonaiagents.mcp.mcp_security import create_auth_header
header = create_auth_header("custom-value", auth_type="custom", header_name="X-API-Key")
assert header == {"X-API-Key": "custom-value"}
def test_validate_auth_header(self):
"""Test validating authentication header presence."""
from praisonaiagents.mcp.mcp_security import validate_auth_header
headers_with_auth = {"Authorization": "Bearer token"}
headers_without_auth = {"Content-Type": "application/json"}
assert validate_auth_header(headers_with_auth) is True
assert validate_auth_header(headers_without_auth) is False
class TestSecureSessionHandling:
"""Test secure session ID handling."""
def test_generate_secure_session_id(self):
"""Test generating cryptographically secure session ID."""
from praisonaiagents.mcp.mcp_security import generate_secure_session_id
session_id = generate_secure_session_id()
# Should be non-empty
assert len(session_id) > 0
# Should be unique
session_id2 = generate_secure_session_id()
assert session_id != session_id2
# Should contain only valid characters (visible ASCII)
for char in session_id:
assert 0x21 <= ord(char) <= 0x7E
def test_session_id_length(self):
"""Test session ID has sufficient length for security."""
from praisonaiagents.mcp.mcp_security import generate_secure_session_id
session_id = generate_secure_session_id()
# Should be at least 32 characters for security
assert len(session_id) >= 32
def test_session_id_entropy(self):
"""Test session ID has sufficient entropy."""
from praisonaiagents.mcp.mcp_security import generate_secure_session_id
# Generate multiple IDs and check they're all different
ids = [generate_secure_session_id() for _ in range(100)]
unique_ids = set(ids)
assert len(unique_ids) == 100 # All should be unique
class TestDNSRebindingPrevention:
"""Test DNS rebinding attack prevention."""
def test_extract_origin_host(self):
"""Test extracting host from Origin header."""
from praisonaiagents.mcp.mcp_security import extract_origin_host
assert extract_origin_host("http://localhost:8080") == "localhost"
assert extract_origin_host("https://example.com:443/path") == "example.com"
assert extract_origin_host("http://127.0.0.1") == "127.0.0.1"
def test_detect_dns_rebinding_attempt(self):
"""Test detection of potential DNS rebinding attempts."""
from praisonaiagents.mcp.mcp_security import is_potential_dns_rebinding
# External origin trying to access localhost server
assert is_potential_dns_rebinding(
origin="https://evil.com",
server_host="localhost"
) is True
# Same origin is safe
assert is_potential_dns_rebinding(
origin="http://localhost:8080",
server_host="localhost"
) is False
class TestSecurityConfig:
"""Test security configuration."""
def test_security_config_defaults(self):
"""Test default security configuration."""
from praisonaiagents.mcp.mcp_security import SecurityConfig
config = SecurityConfig()
# Secure defaults
assert config.validate_origin is True
assert config.allow_missing_origin is False
assert "localhost" in config.allowed_origins
assert "127.0.0.1" in config.allowed_origins
def test_security_config_custom(self):
"""Test custom security configuration."""
from praisonaiagents.mcp.mcp_security import SecurityConfig
config = SecurityConfig(
allowed_origins=["example.com"],
allow_missing_origin=True,
require_auth=True
)
assert "example.com" in config.allowed_origins
assert config.allow_missing_origin is True
assert config.require_auth is True
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,299 @@
"""Tests for MCP server wrapper - TDD approach.
These tests define the expected behavior of ToolsMCPServer class.
"""
import pytest
from unittest.mock import patch
class TestToolsMCPServerInit:
"""Test ToolsMCPServer initialization."""
def test_init_with_name(self):
"""Test server initializes with custom name."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer(name="my-tools")
assert server.name == "my-tools"
def test_init_with_default_name(self):
"""Test server initializes with default name."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
assert server.name == "praisonai-tools"
def test_init_with_tools_list(self):
"""Test server initializes with list of tools."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
def my_tool(x: str) -> str:
"""A test tool."""
return x
server = ToolsMCPServer(tools=[my_tool])
assert len(server.tools) == 1
def test_init_empty_tools(self):
"""Test server initializes with empty tools list."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
assert server.tools == []
class TestToolsMCPServerRegister:
"""Test tool registration."""
def test_register_single_tool(self):
"""Test registering a single tool."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
def search(query: str) -> str:
"""Search for information."""
return f"Results for {query}"
server.register_tool(search)
assert len(server.tools) == 1
assert server.tools[0].__name__ == "search"
def test_register_multiple_tools(self):
"""Test registering multiple tools."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
def tool1(x: str) -> str:
"""Tool 1."""
return x
def tool2(y: int) -> int:
"""Tool 2."""
return y
server.register_tool(tool1)
server.register_tool(tool2)
assert len(server.tools) == 2
def test_register_tools_batch(self):
"""Test registering multiple tools at once."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
def tool1(x: str) -> str:
"""Tool 1."""
return x
def tool2(y: int) -> int:
"""Tool 2."""
return y
server.register_tools([tool1, tool2])
assert len(server.tools) == 2
def test_register_async_tool(self):
"""Test registering an async tool."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
async def async_search(query: str) -> str:
"""Async search."""
return f"Results for {query}"
server.register_tool(async_search)
assert len(server.tools) == 1
class TestToolsMCPServerSchemas:
"""Test schema generation for registered tools."""
def test_get_tool_schemas(self):
"""Test getting schemas for all registered tools."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
def search(query: str, max_results: int = 10) -> str:
"""Search the web."""
return ""
server.register_tool(search)
schemas = server.get_tool_schemas()
assert len(schemas) == 1
assert schemas[0]["name"] == "search"
assert schemas[0]["description"] == "Search the web."
assert "query" in schemas[0]["inputSchema"]["properties"]
def test_get_tool_names(self):
"""Test getting list of tool names."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
def tool1(x: str) -> str:
"""Tool 1."""
return x
def tool2(y: int) -> int:
"""Tool 2."""
return y
server.register_tools([tool1, tool2])
names = server.get_tool_names()
assert set(names) == {"tool1", "tool2"}
class TestToolsMCPServerExecution:
"""Test tool execution through MCP server."""
def test_execute_sync_tool(self):
"""Test executing a synchronous tool."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
def add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
server.register_tool(add)
result = server.execute_tool("add", {"a": 5, "b": 3})
assert result == 8
@pytest.mark.asyncio
async def test_execute_async_tool(self):
"""Test executing an async tool."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
async def async_add(a: int, b: int) -> int:
"""Add two numbers asynchronously."""
return a + b
server.register_tool(async_add)
result = await server.execute_tool_async("async_add", {"a": 5, "b": 3})
assert result == 8
def test_execute_nonexistent_tool(self):
"""Test executing a tool that doesn't exist."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
with pytest.raises(ValueError, match="Tool 'nonexistent' not found"):
server.execute_tool("nonexistent", {})
class TestToolsMCPServerFastMCP:
"""Test FastMCP integration."""
def test_get_fastmcp_instance(self):
"""Test getting FastMCP instance with registered tools."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer(name="test-server")
def my_tool(x: str) -> str:
"""A test tool."""
return x
server.register_tool(my_tool)
# Get FastMCP instance
mcp = server.get_fastmcp()
assert mcp is not None
assert mcp.name == "test-server"
def test_fastmcp_has_registered_tools(self):
"""Test that FastMCP instance has all registered tools."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
def search(query: str) -> str:
"""Search tool."""
return query
def calculate(x: int, y: int) -> int:
"""Calculate tool."""
return x + y
server.register_tools([search, calculate])
mcp = server.get_fastmcp()
# FastMCP should have both tools registered
# We can check by looking at the internal tool registry
assert mcp is not None
class TestLaunchToolsMCPServer:
"""Test the convenience function for launching MCP server."""
def test_launch_function_exists(self):
"""Test that launch_tools_mcp_server function exists."""
from praisonaiagents.mcp.mcp_server import launch_tools_mcp_server
assert callable(launch_tools_mcp_server)
def test_launch_with_tool_names(self):
"""Test launching server with tool names from registry."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
# This should create a server with specified tools
# We'll mock the actual launch to avoid blocking
with patch.object(ToolsMCPServer, 'run', return_value=None):
# Just verify it doesn't raise
pass # Actual launch test will be in integration tests
def test_launch_with_custom_tools(self):
"""Test launching server with custom tool functions."""
from praisonaiagents.mcp.mcp_server import launch_tools_mcp_server
def my_custom_tool(x: str) -> str:
"""Custom tool."""
return x
# Verify function accepts tools parameter
# Actual launch test will be in integration tests
assert callable(launch_tools_mcp_server)
class TestToolsMCPServerTransports:
"""Test different transport options."""
def test_stdio_transport_config(self):
"""Test stdio transport configuration."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
# Server should support stdio transport
assert hasattr(server, 'run_stdio') or hasattr(server, 'run')
def test_sse_transport_config(self):
"""Test SSE transport configuration."""
from praisonaiagents.mcp.mcp_server import ToolsMCPServer
server = ToolsMCPServer()
# Server should support SSE transport
assert hasattr(server, 'run_sse') or hasattr(server, 'run')
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,242 @@
"""Tests for MCP Session Management - TDD approach.
These tests define the expected behavior of session management for
Streamable HTTP transport per MCP Protocol Revision 2025-11-25.
Session management features:
- MCP-Session-Id header handling
- Session persistence across requests
- Session expiration handling (HTTP 404)
- Session termination via HTTP DELETE
- Session ID validation (visible ASCII 0x21-0x7E)
"""
import pytest
class TestSessionIdValidation:
"""Test session ID validation per MCP spec."""
def test_valid_session_id_ascii(self):
"""Test that valid ASCII session IDs are accepted."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
# Valid session IDs (visible ASCII 0x21-0x7E)
assert is_valid_session_id("abc123") is True
assert is_valid_session_id("session-id-12345") is True
assert is_valid_session_id("a1b2c3d4e5f6") is True
assert is_valid_session_id("!@#$%^&*()") is True # Special chars in range
assert is_valid_session_id("UUID-1234-5678-90ab") is True
def test_invalid_session_id_with_spaces(self):
"""Test that session IDs with spaces are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("session id") is False
assert is_valid_session_id(" leading") is False
assert is_valid_session_id("trailing ") is False
def test_invalid_session_id_with_control_chars(self):
"""Test that session IDs with control characters are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("session\x00id") is False # NULL
assert is_valid_session_id("session\nid") is False # Newline
assert is_valid_session_id("session\tid") is False # Tab
def test_invalid_session_id_empty(self):
"""Test that empty session IDs are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("") is False
assert is_valid_session_id(None) is False
def test_invalid_session_id_non_ascii(self):
"""Test that non-ASCII characters are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("session\x7fid") is False # DEL char
assert is_valid_session_id("session\x80id") is False # Extended ASCII
class TestSessionManager:
"""Test SessionManager class for managing MCP sessions."""
def test_session_manager_creation(self):
"""Test SessionManager can be created."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
assert manager.session_id is None
assert manager.is_active is False
def test_session_manager_set_session_id(self):
"""Test setting session ID."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_session_id("test-session-123")
assert manager.session_id == "test-session-123"
assert manager.is_active is True
def test_session_manager_rejects_invalid_id(self):
"""Test that invalid session IDs are rejected."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
with pytest.raises(ValueError, match="Invalid session ID"):
manager.set_session_id("invalid session") # Contains space
def test_session_manager_clear_session(self):
"""Test clearing session."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_session_id("test-session-123")
manager.clear()
assert manager.session_id is None
assert manager.is_active is False
def test_session_manager_get_headers(self):
"""Test getting headers with session ID."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
# No session - no header
headers = manager.get_headers()
assert 'Mcp-Session-Id' not in headers
# With session - include header
manager.set_session_id("test-session-123")
headers = manager.get_headers()
assert headers['Mcp-Session-Id'] == "test-session-123"
class TestSessionFromResponse:
"""Test extracting session ID from HTTP response headers."""
def test_extract_session_id_from_headers(self):
"""Test extracting session ID from response headers."""
from praisonaiagents.mcp.mcp_session import extract_session_id
headers = {'Mcp-Session-Id': 'new-session-456'}
session_id = extract_session_id(headers)
assert session_id == 'new-session-456'
def test_extract_session_id_case_insensitive(self):
"""Test that header name is case-insensitive."""
from praisonaiagents.mcp.mcp_session import extract_session_id
# Various case combinations
assert extract_session_id({'mcp-session-id': 'test'}) == 'test'
assert extract_session_id({'MCP-SESSION-ID': 'test'}) == 'test'
assert extract_session_id({'Mcp-Session-Id': 'test'}) == 'test'
def test_extract_session_id_missing(self):
"""Test when session ID header is missing."""
from praisonaiagents.mcp.mcp_session import extract_session_id
headers = {'Content-Type': 'application/json'}
session_id = extract_session_id(headers)
assert session_id is None
class TestSessionExpiration:
"""Test handling of session expiration (HTTP 404)."""
def test_is_session_expired_404(self):
"""Test that HTTP 404 indicates session expiration."""
from praisonaiagents.mcp.mcp_session import is_session_expired
assert is_session_expired(404) is True
assert is_session_expired(200) is False
assert is_session_expired(400) is False
assert is_session_expired(500) is False
def test_session_manager_marks_expired(self):
"""Test that session manager can mark session as expired."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_session_id("test-session")
manager.mark_expired()
assert manager.is_expired is True
assert manager.is_active is False
class TestProtocolVersion:
"""Test MCP-Protocol-Version header handling."""
def test_default_protocol_version(self):
"""Test default protocol version is 2025-03-26 for compatibility."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
headers = manager.get_headers()
# Default version for backward compatibility
assert headers.get('Mcp-Protocol-Version') == '2025-03-26'
def test_set_protocol_version(self):
"""Test setting negotiated protocol version."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_protocol_version('2025-11-25')
headers = manager.get_headers()
assert headers['Mcp-Protocol-Version'] == '2025-11-25'
def test_valid_protocol_versions(self):
"""Test that valid protocol versions are accepted."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version('2025-03-26') is True
assert is_valid_protocol_version('2025-06-18') is True
assert is_valid_protocol_version('2025-11-25') is True
assert is_valid_protocol_version('2024-11-05') is True
def test_invalid_protocol_versions(self):
"""Test that invalid protocol versions are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version('invalid') is False
assert is_valid_protocol_version('2025') is False
assert is_valid_protocol_version('') is False
class TestHTTPStreamSessionIntegration:
"""Test session management integration with HTTPStreamTransport."""
def test_transport_uses_session_manager(self):
"""Test that HTTPStreamTransport uses SessionManager."""
# This will be implemented when we update HTTPStreamTransport
pass
def test_transport_handles_session_from_init_response(self):
"""Test that transport extracts session ID from InitializeResult."""
pass
def test_transport_includes_session_in_requests(self):
"""Test that transport includes session ID in subsequent requests."""
pass
def test_transport_handles_session_expiration(self):
"""Test that transport handles HTTP 404 session expiration."""
pass
def test_transport_can_terminate_session(self):
"""Test that transport can terminate session via HTTP DELETE."""
pass
if __name__ == "__main__":
pytest.main([__file__, "-v"])
Loaded 100 of 134 files, more files were not shown because too many files have changed in this diff. Show more