Compare commits

...
32 Commits
Author SHA1 Message Date
Mervin Praison 214add68d9 Release v2.3.12 2025-12-16 15:03:23 +00:00
Mervin Praison 6d2cb9676d Update README.md to consolidate and enhance the Supported Providers section, adding new AI providers and their corresponding example links for improved documentation and user guidance. 2025-12-16 11:54:00 +00:00
Mervin Praison c519cb1379 Add supported providers section to README.md, listing various AI providers with example links for each, enhancing documentation and usability. 2025-12-16 11:52:19 +00:00
Mervin Praison 356dc4909a Update README.md to refine feature table, enhancing clarity by renaming 'Automated AI Agents Creation' to 'Single Agent' and reorganizing agent types for better categorization. 2025-12-16 11:25:42 +00:00
Mervin Praison dc278ad07b Enhance README.md with new agent types and advanced features, including Image Generation, Data Analyst, and Finance Agents, along with updates to existing feature descriptions and examples. 2025-12-16 11:22:23 +00:00
Mervin Praison 5cc53eb4bb Add advanced features section to README.md, detailing new capabilities including Deep Research Agents, Query Rewriter Agent, and various MCP enhancements. 2025-12-16 11:13:38 +00:00
Mervin Praison b72576d086 Refactor feature table in README.md for improved clarity and consistency, updating formatting and links for AI agent capabilities. 2025-12-16 11:11:04 +00:00
Mervin Praison 30314b1573 Update README.md to include detailed feature table with example code and documentation links for each AI agent capability. 2025-12-16 11:05:51 +00:00
Mervin Praison f721bd5e9a Add Fast Context feature to README highlights 2025-12-16 10:52:37 +00:00
Mervin Praison d00ea3e3c3 Release v2.3.11 2025-12-16 10:45:36 +00:00
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
146 changed files with 31252 additions and 212 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"
+468 -43
View File
@@ -32,34 +32,100 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
## Key Features
- 🤖 Automated AI Agents Creation
- 🔄 Self Reflection AI Agents
- 🧠 Reasoning AI Agents
- 👁️ Multi Modal AI Agents
- 🤝 Multi Agent Collaboration
- 🎭 AI Agent Workflow
- 📚 Add Custom Knowledge
- 🧠 Agents with Short and Long Term Memory
- 📄 Chat with PDF Agents
- 💻 Code Interpreter Agents
- 📚 RAG Agents
- 🤔 Async & Parallel Processing
- 🔄 Auto Agents
- 🔢 Math Agents
- 🎯 Structured Output Agents
- 🔗 LangChain Integrated Agents
- 📞 Callback Agents
- 🤏 Mini AI Agents
- 🛠️ 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)
| Feature | Code | Docs |
|---------|:----:|:----:|
| 🚀 Single Agent | [Example](examples/python/agents/single-agent.py) | [📖](https://docs.praison.ai/agents/single) |
| 🤝 Multi Agents | [Example](examples/python/general/mini_agents_example.py) | [📖](https://docs.praison.ai/concepts/agents) |
| 🤖 Auto Agents | [Example](examples/python/general/auto_agents_example.py) | [📖](https://docs.praison.ai/features/autoagents) |
| 🔄 Self Reflection AI Agents | [Example](examples/python/concepts/self-reflection-details.py) | [📖](https://docs.praison.ai/features/selfreflection) |
| 🧠 Reasoning AI Agents | [Example](examples/python/concepts/reasoning-extraction.py) | [📖](https://docs.praison.ai/features/reasoning) |
| 👁️ Multi Modal AI Agents | [Example](examples/python/general/multimodal.py) | [📖](https://docs.praison.ai/features/multimodal) |
| 🎭 AI Agent Workflow | [Example](examples/python/stateful/workflow-state-example.py) | [📖](https://docs.praison.ai/features/workflows) |
| 📚 Add Custom Knowledge | [Example](examples/python/concepts/knowledge-agents.py) | [📖](https://docs.praison.ai/features/knowledge) |
| 🧠 Memory (Short & Long Term) | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
| 📄 Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [📖](https://docs.praison.ai/features/chat-with-pdf) |
| 💻 Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [📖](https://docs.praison.ai/features/codeagent) |
| 📚 RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [📖](https://docs.praison.ai/features/rag) |
| 🤔 Async & Parallel Processing | [Example](examples/python/general/async_example.py) | [📖](https://docs.praison.ai/features/async) |
| 🔢 Math Agents | [Example](examples/python/agents/math-agent.py) | [📖](https://docs.praison.ai/features/mathagent) |
| 🎯 Structured Output Agents | [Example](examples/python/general/structured_agents_example.py) | [📖](https://docs.praison.ai/features/structured) |
| 🔗 LangChain Integrated Agents | [Example](examples/python/general/langchain_example.py) | [📖](https://docs.praison.ai/features/langchain) |
| 📞 Callback Agents | [Example](examples/python/general/advanced-callback-systems.py) | [📖](https://docs.praison.ai/features/callbacks) |
| 🛠️ 100+ Custom Tools | [Example](examples/python/general/tools_example.py) | [📖](https://docs.praison.ai/tools/tools) |
| 📄 YAML Configuration | [Example](examples/cookbooks/yaml/secondary_market_research_agents.yaml) | [📖](https://docs.praison.ai/developers/agents-playbook) |
| 💯 100+ LLM Support | [Example](examples/python/providers/openai/openai_gpt4_example.py) | [📖](https://docs.praison.ai/models) |
| 🔬 Deep Research Agents | [Example](examples/python/agents/research-agent.py) | [📖](https://docs.praison.ai/agents/deep-research) |
| 🔄 Query Rewriter Agent | [Example](#5-query-rewriter-agent) | [📖](https://docs.praison.ai/agents/query-rewriter) |
| 🌐 Native Web Search | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/agents/websearch) |
| 📥 Web Fetch (Anthropic) | [Example](#web-search-web-fetch--prompt-caching) | [📖](https://docs.praison.ai/features/model-capabilities) |
| 💾 Prompt Caching | [Example](#web-search-web-fetch--prompt-caching) | [📖](https://docs.praison.ai/features/model-capabilities) |
| 🧠 Claude Memory Tool | [Example](#claude-memory-tool-cli) | [📖](https://docs.praison.ai/features/claude-memory-tool) |
| 💾 File-Based Memory | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
| 🔍 Built-in Search Tools | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/tools/tavily) |
| 📋 Planning Mode | [Example](examples/python/agents/planning-agent.py) | [📖](https://docs.praison.ai/features/planning-mode) |
| 🔧 Planning Tools | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
| 🧠 Planning Reasoning | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
| 🔌 MCP Transports | [Example](examples/python/mcp/mcp-transports-overview.py) | [📖](https://docs.praison.ai/mcp/transports) |
| 🌐 WebSocket MCP | [Example](examples/python/mcp/websocket-mcp.py) | [📖](https://docs.praison.ai/mcp/sse-transport) |
| 🔐 MCP Security | [Example](examples/python/mcp/mcp-security.py) | [📖](https://docs.praison.ai/mcp/transports) |
| 🔄 MCP Resumability | [Example](examples/python/mcp/mcp-resumability.py) | [📖](https://docs.praison.ai/mcp/sse-transport) |
| ⚡ Fast Context | [Example](examples/context/00_agent_fast_context_basic.py) | [📖](https://docs.praison.ai/features/fast-context) |
| 🖼️ Image Generation Agent | [Example](examples/python/image/image-agent.py) | [📖](https://docs.praison.ai/features/image-generation) |
| 📷 Image to Text Agent | [Example](examples/python/agents/image-to-text-agent.py) | [📖](https://docs.praison.ai/agents/image-to-text) |
| 🎬 Video Agent | [Example](examples/python/agents/video-agent.py) | [📖](https://docs.praison.ai/agents/video) |
| 📊 Data Analyst Agent | [Example](examples/python/agents/data-analyst-agent.py) | [📖](https://docs.praison.ai/agents/data-analyst) |
| 💰 Finance Agent | [Example](examples/python/agents/finance-agent.py) | [📖](https://docs.praison.ai/agents/finance) |
| 🛒 Shopping Agent | [Example](examples/python/agents/shopping-agent.py) | [📖](https://docs.praison.ai/agents/shopping) |
| ⭐ Recommendation Agent | [Example](examples/python/agents/recommendation-agent.py) | [📖](https://docs.praison.ai/agents/recommendation) |
| 📖 Wikipedia Agent | [Example](examples/python/agents/wikipedia-agent.py) | [📖](https://docs.praison.ai/agents/wikipedia) |
| 💻 Programming Agent | [Example](examples/python/agents/programming-agent.py) | [📖](https://docs.praison.ai/agents/programming) |
| 📝 Markdown Agent | [Example](examples/python/agents/markdown-agent.py) | [📖](https://docs.praison.ai/agents/markdown) |
| 📝 Prompt Expander Agent | [Example](#prompt-expansion) | [📖](https://docs.praison.ai/agents/prompt-expander) |
| 🔀 Router Agent | [Example](examples/python/agents/router-agent-cost-optimization.py) | [📖](https://docs.praison.ai/features/routing) |
| ⛓️ Prompt Chaining | [Example](examples/python/general/prompt_chaining.py) | [📖](https://docs.praison.ai/features/promptchaining) |
| 🔍 Evaluator Optimiser | [Example](examples/python/general/evaluator-optimiser.py) | [📖](https://docs.praison.ai/features/evaluator-optimiser) |
| 👷 Orchestrator Workers | [Example](examples/python/general/orchestrator-workers.py) | [📖](https://docs.praison.ai/features/orchestrator-worker) |
| ⚡ Parallelisation | [Example](examples/python/general/parallelisation.py) | [📖](https://docs.praison.ai/features/parallelisation) |
| 🔁 Repetitive Agents | [Example](examples/python/concepts/repetitive-agents.py) | [📖](https://docs.praison.ai/features/repetitive) |
| 🤝 Agent Handoffs | [Example](examples/python/handoff/handoff_basic.py) | [📖](https://docs.praison.ai/features/handoffs) |
| 🛡️ Guardrails | [Example](examples/python/guardrails/comprehensive-guardrails-example.py) | [📖](https://docs.praison.ai/features/guardrails) |
| 💬 Sessions Management | [Example](examples/python/sessions/comprehensive-session-management.py) | [📖](https://docs.praison.ai/features/sessions) |
| ✅ Human Approval | [Example](examples/python/general/human_approval_example.py) | [📖](https://docs.praison.ai/features/approval) |
| 🔄 Stateful Agents | [Example](examples/python/stateful/workflow-state-example.py) | [📖](https://docs.praison.ai/features/stateful-agents) |
| 🤖 Autonomous Workflow | [Example](examples/python/general/autonomous-agent.py) | [📖](https://docs.praison.ai/features/autonomous-workflow) |
| 📜 Rules & Instructions | [Example](#6-rules--instructions) | [📖](https://docs.praison.ai/features/rules) |
| 🪝 Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
| 📈 Telemetry | [Example](examples/python/telemetry/production-telemetry-example.py) | [📖](https://docs.praison.ai/features/telemetry) |
| 📹 Camera Integration | [Example](examples/python/camera/) | [📖](https://docs.praison.ai/features/camera-integration) |
## Supported Providers
| Provider | Example |
|----------|:-------:|
| OpenAI | [Example](examples/python/providers/openai/openai_gpt4_example.py) |
| Anthropic | [Example](examples/python/providers/anthropic/anthropic_claude_example.py) |
| Google Gemini | [Example](examples/python/providers/google/google_gemini_example.py) |
| Ollama | [Example](examples/python/providers/ollama/ollama-agents.py) |
| Groq | [Example](examples/python/providers/groq/kimi_with_groq_example.py) |
| DeepSeek | [Example](examples/python/providers/deepseek/deepseek_example.py) |
| xAI Grok | [Example](examples/python/providers/xai/xai_grok_example.py) |
| Mistral | [Example](examples/python/providers/mistral/mistral_example.py) |
| Cohere | [Example](examples/python/providers/cohere/cohere_example.py) |
| Perplexity | [Example](examples/python/providers/perplexity/perplexity_example.py) |
| Fireworks | [Example](examples/python/providers/fireworks/fireworks_example.py) |
| Together AI | [Example](examples/python/providers/together/together_ai_example.py) |
| OpenRouter | [Example](examples/python/providers/openrouter/openrouter_example.py) |
| HuggingFace | [Example](examples/python/providers/huggingface/huggingface_example.py) |
| Azure OpenAI | [Example](examples/python/providers/azure/azure_openai_example.py) |
| AWS Bedrock | [Example](examples/python/providers/aws/aws_bedrock_example.py) |
| Google Vertex | [Example](examples/python/providers/vertex/vertex_example.py) |
| Databricks | [Example](examples/python/providers/databricks/databricks_example.py) |
| Cloudflare | [Example](examples/python/providers/cloudflare/cloudflare_example.py) |
| AI21 | [Example](examples/python/providers/ai21/ai21_example.py) |
| Replicate | [Example](examples/python/providers/replicate/replicate_example.py) |
| SageMaker | [Example](examples/python/providers/sagemaker/sagemaker_example.py) |
| Moonshot | [Example](examples/python/providers/moonshot/moonshot_example.py) |
| vLLM | [Example](examples/python/providers/vllm/vllm_example.py) |
## Using Python Code
@@ -103,7 +169,34 @@ Run:
python app.py
```
### 3. Deep Research Agent
### 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.
@@ -141,7 +234,7 @@ print(result.report)
- 🛠️ Built-in tools: web search, code interpreter, MCP, file search
- 🔄 Automatic provider detection from model name
### 4. Query Rewriter Agent
### 5. Query Rewriter Agent
Transform user queries to improve RAG retrieval quality using multiple strategies.
@@ -176,12 +269,13 @@ result = agent.rewrite("What about cost?", chat_history=[...])
- **CONTEXTUAL**: Resolve references using conversation history
- **AUTO**: Automatically detect best strategy
### 5. Agent Memory (Zero Dependencies)
### 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(
@@ -191,17 +285,9 @@ agent = Agent(
user_id="user123" # Isolate memory per user
)
# Store memories programmatically
agent.store_memory("User prefers dark mode", memory_type="short_term")
agent.store_memory("User's name is John", memory_type="long_term", importance=0.9)
# Memory is automatically injected into conversations
result = agent.start("What do you know about me?")
# Agent will recall: "Your name is John and you prefer dark mode"
# Memory persists across sessions - create new agent with same user_id
agent2 = Agent(name="Assistant", memory=True, user_id="user123")
# agent2 automatically loads all previous memories!
result = agent.start("My name is John and I prefer Python")
# Agent will remember this for future conversations
```
**Memory Types:**
@@ -210,6 +296,28 @@ agent2 = Agent(name="Assistant", memory=True, user_id="user123")
- **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 |
|--------|-------------|-------------|
@@ -218,6 +326,103 @@ agent2 = Agent(name="Assistant", memory=True, user_id="user123")
| `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:
@@ -272,6 +477,101 @@ praisonai research --query-rewrite --tools tools.py --save "Your research query"
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
@@ -690,7 +990,6 @@ praisonai agents.yaml
</a>
</div>
## Custom Tools
### Using `@tool` Decorator
@@ -856,6 +1155,73 @@ agent = Agent(
| Prompt Caching | OpenAI (auto), Anthropic, Bedrock, Deepseek |
## MCP (Model Context Protocol)
PraisonAI supports MCP Protocol Revision 2025-11-25 with multiple transports.
### MCP Client (Consume MCP Servers)
```python
from praisonaiagents import Agent, MCP
# stdio - Local NPX/Python servers
agent = Agent(tools=MCP("npx @modelcontextprotocol/server-memory"))
# Streamable HTTP - Production servers
agent = Agent(tools=MCP("https://api.example.com/mcp"))
# WebSocket - Real-time bidirectional
agent = Agent(tools=MCP("wss://api.example.com/mcp", auth_token="token"))
# SSE (Legacy) - Backward compatibility
agent = Agent(tools=MCP("http://localhost:8080/sse"))
# With environment variables
agent = Agent(
tools=MCP(
command="npx",
args=["-y", "@modelcontextprotocol/server-brave-search"],
env={"BRAVE_API_KEY": "your-key"}
)
)
```
### MCP Server (Expose Tools as MCP Server)
Expose your Python functions as MCP tools for Claude Desktop, Cursor, and other MCP clients:
```python
from praisonaiagents.mcp import ToolsMCPServer
def search_web(query: str, max_results: int = 5) -> dict:
"""Search the web for information."""
return {"results": [f"Result for {query}"]}
def calculate(expression: str) -> dict:
"""Evaluate a mathematical expression."""
return {"result": eval(expression)}
# Create and run MCP server
server = ToolsMCPServer(name="my-tools")
server.register_tools([search_web, calculate])
server.run() # stdio for Claude Desktop
# server.run_sse(host="0.0.0.0", port=8080) # SSE for web clients
```
### MCP Features
| Feature | Description |
|---------|-------------|
| **Session Management** | Automatic `Mcp-Session-Id` handling |
| **Protocol Versioning** | `Mcp-Protocol-Version` header |
| **Resumability** | SSE stream recovery via `Last-Event-ID` |
| **Security** | Origin validation, DNS rebinding prevention |
| **WebSocket** | Auto-reconnect with exponential backoff |
## Development:
Below is used for development only.
@@ -896,6 +1262,65 @@ cd src/praisonai && uv publish
- Submit a pull request via GitHub's web interface.
- Await feedback from project maintainers.
## Advanced Features
**Research & Intelligence:**
- 🔬 **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 Expander Agent** (Expand short prompts into detailed instructions)
**Memory & Caching:**
- 💾 **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 & Workflows:**
- 📋 **Planning Mode** (Plan before execution - Agent & Multi-Agent)
- 🔧 **Planning Tools** (Research with tools during planning)
- 🧠 **Planning Reasoning** (Chain-of-thought planning)
- ⛓️ **Prompt Chaining** (Sequential prompt workflows with gates)
- 🔍 **Evaluator Optimiser** (Generate and optimize through iterative feedback)
- 👷 **Orchestrator Workers** (Distribute tasks among specialized workers)
- ⚡ **Parallelisation** (Execute tasks in parallel for improved performance)
- 🔁 **Repetitive Agents** (Handle repetitive tasks through automated loops)
- 🤖 **Autonomous Workflow** (Monitor, act, adapt based on environment feedback)
**Agent Types:**
- 🖼️ **Image Generation Agent** (Create images from text descriptions)
- 📷 **Image to Text Agent** (Extract text and descriptions from images)
- 🎬 **Video Agent** (Analyze and process video content)
- 📊 **Data Analyst Agent** (Analyze data and generate insights)
- 💰 **Finance Agent** (Financial analysis and recommendations)
- 🛒 **Shopping Agent** (Price comparison and shopping assistance)
- ⭐ **Recommendation Agent** (Personalized recommendations)
- 📖 **Wikipedia Agent** (Search and extract Wikipedia information)
- 💻 **Programming Agent** (Code development and analysis)
- 📝 **Markdown Agent** (Generate and format Markdown content)
- 🔀 **Router Agent** (Dynamic task routing with cost optimization)
**MCP Protocol:**
- 🔌 **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)
**Safety & Control:**
- 🤝 **Agent Handoffs** (Transfer context between specialized agents)
- 🛡️ **Guardrails** (Input/output validation and safety checks)
- ✅ **Human Approval** (Require human confirmation for critical actions)
- 💬 **Sessions Management** (Isolated conversation contexts)
- 🔄 **Stateful Agents** (Maintain state across interactions)
**Developer Tools:**
- ⚡ **Fast Context** (Rapid parallel code search - 10-20x faster than traditional methods)
- 📜 **Rules & Instructions** (Auto-discover CLAUDE.md, AGENTS.md, GEMINI.md)
- 🪝 **Hooks** (Pre/post operation hooks for custom logic)
- 📈 **Telemetry** (Track agent performance and usage)
- 📹 **Camera Integration** (Capture and analyze camera input)
## Other Features
- 🔄 Use CrewAI or AG2 (Formerly AutoGen) Framework
@@ -903,7 +1328,7 @@ cd src/praisonai && uv publish
- 🎨 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.3.3" \
"praisonai>=2.3.12" \
"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.3.3" \
"praisonai>=2.3.12" \
"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.3.3" \
"praisonai>=2.3.12" \
"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.3.3" \
"praisonai>=2.3.12" \
"praisonai[ui]" \
"praisonai[crewai]"
+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,89 @@
"""
Fast Context with Agent - Basic Example
This is the simplest way to use Fast Context with an Agent.
Just set fast_context=True and the agent will automatically
have access to rapid code search capabilities.
Fast Context provides:
- 10-20x faster code search than traditional methods
- Parallel execution (8 concurrent searches)
- Automatic caching for repeated queries
- Seamless integration with Agent workflows
"""
from praisonaiagents import Agent
def main():
print("=" * 60)
print("Agent with Fast Context - Basic Example")
print("=" * 60)
# Create an agent with Fast Context enabled
# This is the simplest configuration
agent = Agent(
name="CodeAssistant",
instructions="You are a helpful code assistant.",
fast_context=True, # Enable Fast Context
fast_context_path=".", # Search in current directory
verbose=False
)
print("\n✓ Agent created with Fast Context enabled")
print(f" - fast_context_enabled: {agent.fast_context_enabled}")
print(f" - fast_context_path: {agent.fast_context_path}")
print(f" - fast_context_model: {agent.fast_context_model}")
# Use delegate_to_fast_context to search code
print("\n" + "-" * 60)
print("Searching for 'class Agent' in the codebase...")
print("-" * 60)
context = agent.delegate_to_fast_context("class Agent")
if context:
print(f"\n✓ Found relevant code!")
print(f" Context length: {len(context)} characters")
print("\nPreview of found context:")
print("-" * 40)
# Show first 500 characters
preview = context[:500]
print(preview)
if len(context) > 500:
print("...")
else:
print("\n✗ No relevant code found")
# Another search example
print("\n" + "-" * 60)
print("Searching for 'def chat' in the codebase...")
print("-" * 60)
context = agent.delegate_to_fast_context("def chat")
if context:
print(f"\n✓ Found relevant code!")
print(f" Context length: {len(context)} characters")
else:
print("\n✗ No relevant code found")
# Access FastContext directly for more control
print("\n" + "-" * 60)
print("Direct FastContext access for advanced usage...")
print("-" * 60)
fc = agent.fast_context
if fc:
result = fc.search("import")
print(f"\n✓ Direct search for 'import':")
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms")
print(f" From cache: {result.from_cache}")
print("\n" + "=" * 60)
print("Fast Context makes code search instant for your agents!")
print("=" * 60)
if __name__ == "__main__":
main()
+84
View File
@@ -0,0 +1,84 @@
"""
Example 1: Basic Fast Context Usage
This example demonstrates the basic usage of Fast Context for rapid code search.
Fast Context is 10-20x faster than traditional sequential search because it:
- Executes up to 8 tool calls in parallel
- Limits search to 4 turns maximum
- Uses a lightweight model optimized for search
"""
from praisonaiagents.context.fast import FastContext
import time
def main():
print("=" * 70)
print("Fast Context - Basic Usage Example")
print("=" * 70)
# Create FastContext instance pointing to the praisonai-agents codebase
fc = FastContext(
workspace_path="/Users/praison/praisonai-package/src/praisonai-agents"
)
# Example 1: Simple pattern search
print("\n1. Simple Pattern Search")
print("-" * 40)
start = time.perf_counter()
result = fc.search("def chat")
elapsed = (time.perf_counter() - start) * 1000
print(f" Query: 'def chat'")
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms (total: {elapsed:.0f}ms)")
print(f" Turns used: {result.turns_used}")
print(f" Tool calls: {result.total_tool_calls}")
if result.files:
print("\n Top matches:")
for f in result.files[:3]:
print(f" - {f.path}")
for lr in f.line_ranges[:2]:
print(f" Lines {lr.start}-{lr.end}")
# Example 2: Search for class definitions
print("\n2. Search for Class Definitions")
print("-" * 40)
result = fc.search("class Agent")
print(f" Query: 'class Agent'")
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms")
if result.files:
print("\n Files containing 'class Agent':")
for f in result.files[:5]:
print(f" - {f.path} (relevance: {f.relevance_score:.2f})")
# Example 3: Search for imports
print("\n3. Search for Specific Imports")
print("-" * 40)
result = fc.search("from typing import")
print(f" Query: 'from typing import'")
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms")
# Example 4: Get formatted context string
print("\n4. Get Formatted Context for Agent")
print("-" * 40)
context = fc.get_context_for_agent("authentication", max_files=3)
print(f" Context length: {len(context)} characters")
print(f" Preview:")
preview = context[:300].replace('\n', '\n ')
print(f" {preview}...")
print("\n" + "=" * 70)
print("Fast Context provides rapid, parallel code search!")
print("=" * 70)
if __name__ == "__main__":
main()
@@ -0,0 +1,113 @@
"""
Example 2: Parallel vs Sequential Search Comparison
This example demonstrates the speed advantage of Fast Context's parallel
execution compared to traditional sequential search.
Key insight: Fast Context executes up to 8 searches simultaneously,
making it 5-10x faster for complex queries.
"""
import time
import os
from praisonaiagents.context.fast import FastContext
from praisonaiagents.context.fast.search_tools import grep_search, glob_search
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def sequential_search(queries: list) -> dict:
"""Traditional sequential search - one query at a time."""
results = {}
for query in queries:
results[query] = grep_search(
search_path=WORKSPACE,
pattern=query,
max_results=20
)
return results
def parallel_search_with_fast_context(queries: list) -> dict:
"""Fast Context parallel search - multiple queries at once."""
fc = FastContext(workspace_path=WORKSPACE)
results = {}
for query in queries:
result = fc.search(query)
results[query] = result
return results
def main():
print("=" * 70)
print("Parallel vs Sequential Search Comparison")
print("=" * 70)
# Define multiple search queries
queries = [
"def __init__",
"class Agent",
"async def",
"import logging",
"from typing",
]
print(f"\nSearching for {len(queries)} patterns in the codebase...")
print(f"Workspace: {WORKSPACE}")
# Sequential search
print("\n1. Sequential Search (Traditional)")
print("-" * 40)
start = time.perf_counter()
seq_results = sequential_search(queries)
seq_time = (time.perf_counter() - start) * 1000
total_seq_matches = sum(len(r) for r in seq_results.values())
print(f" Total matches: {total_seq_matches}")
print(f" Time: {seq_time:.0f}ms")
# Fast Context search (with caching disabled for fair comparison)
print("\n2. Fast Context Search (Parallel)")
print("-" * 40)
fc = FastContext(workspace_path=WORKSPACE, cache_enabled=False)
start = time.perf_counter()
fc_results = {}
for query in queries:
fc_results[query] = fc.search(query)
fc_time = (time.perf_counter() - start) * 1000
total_fc_matches = sum(r.total_files for r in fc_results.values())
print(f" Total files found: {total_fc_matches}")
print(f" Time: {fc_time:.0f}ms")
# Comparison
print("\n3. Comparison")
print("-" * 40)
speedup = seq_time / fc_time if fc_time > 0 else float('inf')
print(f" Sequential: {seq_time:.0f}ms")
print(f" Fast Context: {fc_time:.0f}ms")
print(f" Speedup: {speedup:.1f}x faster")
# Detailed results
print("\n4. Results by Query")
print("-" * 40)
for query in queries:
seq_count = len(seq_results[query])
fc_count = fc_results[query].total_files
fc_ms = fc_results[query].search_time_ms
print(f" '{query}':")
print(f" Sequential: {seq_count} matches")
print(f" Fast Context: {fc_count} files in {fc_ms}ms")
print("\n" + "=" * 70)
print(f"Fast Context is {speedup:.1f}x faster due to parallel execution!")
print("=" * 70)
if __name__ == "__main__":
main()
+132
View File
@@ -0,0 +1,132 @@
"""
Example 3: Caching Performance Demonstration
This example shows how Fast Context's caching dramatically improves
performance for repeated queries.
Key benefits:
- First query: Full search execution
- Subsequent queries: Instant cache hit (0ms search time)
- Configurable TTL (time-to-live)
"""
import time
from praisonaiagents.context.fast import FastContext
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def main():
print("=" * 70)
print("Fast Context Caching Performance")
print("=" * 70)
# Create FastContext with caching enabled
fc = FastContext(
workspace_path=WORKSPACE,
cache_enabled=True,
cache_ttl=300 # 5 minutes
)
query = "def execute"
print(f"\nQuery: '{query}'")
print(f"Cache TTL: 300 seconds")
# First search - cache miss
print("\n1. First Search (Cache Miss)")
print("-" * 40)
start = time.perf_counter()
result1 = fc.search(query)
elapsed1 = (time.perf_counter() - start) * 1000
print(f" Files found: {result1.total_files}")
print(f" From cache: {result1.from_cache}")
print(f" Search time: {result1.search_time_ms}ms")
print(f" Total time: {elapsed1:.1f}ms")
# Second search - cache hit
print("\n2. Second Search (Cache Hit)")
print("-" * 40)
start = time.perf_counter()
result2 = fc.search(query)
elapsed2 = (time.perf_counter() - start) * 1000
print(f" Files found: {result2.total_files}")
print(f" From cache: {result2.from_cache}")
print(f" Search time: {result2.search_time_ms}ms")
print(f" Total time: {elapsed2:.1f}ms")
# Third search - still cached
print("\n3. Third Search (Still Cached)")
print("-" * 40)
start = time.perf_counter()
result3 = fc.search(query)
elapsed3 = (time.perf_counter() - start) * 1000
print(f" Files found: {result3.total_files}")
print(f" From cache: {result3.from_cache}")
print(f" Total time: {elapsed3:.1f}ms")
# Clear cache and search again
print("\n4. After Cache Clear")
print("-" * 40)
fc.clear_cache()
print(" Cache cleared!")
start = time.perf_counter()
result4 = fc.search(query)
elapsed4 = (time.perf_counter() - start) * 1000
print(f" Files found: {result4.total_files}")
print(f" From cache: {result4.from_cache}")
print(f" Search time: {result4.search_time_ms}ms")
print(f" Total time: {elapsed4:.1f}ms")
# Performance comparison
print("\n5. Performance Summary")
print("-" * 40)
cache_speedup = elapsed1 / elapsed2 if elapsed2 > 0 else float('inf')
print(f" First search (no cache): {elapsed1:.1f}ms")
print(f" Cached searches: ~{elapsed2:.1f}ms")
print(f" Cache speedup: {cache_speedup:.0f}x faster")
# Multiple different queries
print("\n6. Multiple Query Caching")
print("-" * 40)
queries = ["class Agent", "async def", "import os", "def __init__"]
# First pass - populate cache
print(" First pass (populating cache):")
start = time.perf_counter()
for q in queries:
fc.search(q)
first_pass = (time.perf_counter() - start) * 1000
print(f" Time: {first_pass:.0f}ms")
# Second pass - all cached
print(" Second pass (all cached):")
start = time.perf_counter()
for q in queries:
r = fc.search(q)
assert r.from_cache, f"Expected cache hit for '{q}'"
second_pass = (time.perf_counter() - start) * 1000
print(f" Time: {second_pass:.1f}ms")
multi_speedup = first_pass / second_pass if second_pass > 0 else float('inf')
print(f" Speedup: {multi_speedup:.0f}x faster")
print("\n" + "=" * 70)
print("Caching provides instant results for repeated queries!")
print("=" * 70)
if __name__ == "__main__":
main()
+112
View File
@@ -0,0 +1,112 @@
"""
Example 4: Agent Integration with Fast Context
This example shows how to use Fast Context integrated with the Agent class
for automatic code search delegation.
Benefits:
- Seamless integration with existing Agent workflows
- Automatic context injection for code-related queries
- Configurable via Agent parameters or environment variables
"""
from praisonaiagents import Agent
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def main():
print("=" * 70)
print("Agent Integration with Fast Context")
print("=" * 70)
# Create an agent with Fast Context enabled
print("\n1. Creating Agent with Fast Context")
print("-" * 40)
agent = Agent(
name="CodeAssistant",
instructions="You are a helpful code assistant that can search and analyze code.",
fast_context=True,
fast_context_path=WORKSPACE,
fast_context_model="gpt-4o-mini",
fast_context_max_turns=4,
fast_context_parallelism=8,
fast_context_timeout=30.0,
verbose=False
)
print(f" Agent name: {agent.name}")
print(f" Fast Context enabled: {agent.fast_context_enabled}")
print(f" Fast Context path: {agent.fast_context_path}")
print(f" Fast Context model: {agent.fast_context_model}")
print(f" Max turns: {agent.fast_context_max_turns}")
print(f" Parallelism: {agent.fast_context_parallelism}")
# Use delegate_to_fast_context for code search
print("\n2. Delegating Code Search to Fast Context")
print("-" * 40)
queries = [
"find the Agent class definition",
"search for authentication handlers",
"locate memory management code",
]
for query in queries:
print(f"\n Query: '{query}'")
context = agent.delegate_to_fast_context(query)
if context:
# Show preview of context
lines = context.split('\n')
preview_lines = lines[:5]
print(f" Found context ({len(context)} chars):")
for line in preview_lines:
print(f" {line[:60]}{'...' if len(line) > 60 else ''}")
if len(lines) > 5:
print(f" ... and {len(lines) - 5} more lines")
else:
print(" No relevant code found")
# Access the FastContext instance directly
print("\n3. Direct FastContext Access")
print("-" * 40)
fc = agent.fast_context
if fc:
result = fc.search("def chat")
print(f" Direct search for 'def chat':")
print(f" Files found: {result.total_files}")
print(f" Search time: {result.search_time_ms}ms")
for f in result.files[:3]:
print(f" - {f.path}")
# Show configuration options
print("\n4. Configuration Options")
print("-" * 40)
print(" Agent parameters:")
print(" - fast_context=True # Enable Fast Context")
print(" - fast_context_path='.' # Workspace path")
print(" - fast_context_model='gpt-4o-mini' # Search model")
print(" - fast_context_max_turns=4 # Max search turns")
print(" - fast_context_parallelism=8 # Parallel calls")
print(" - fast_context_timeout=30.0 # Timeout (seconds)")
print("\n Environment variables:")
print(" - FAST_CONTEXT_MODEL")
print(" - FAST_CONTEXT_MAX_TURNS")
print(" - FAST_CONTEXT_PARALLELISM")
print(" - FAST_CONTEXT_TIMEOUT")
print(" - FAST_CONTEXT_CACHE")
print(" - FAST_CONTEXT_CACHE_TTL")
print("\n" + "=" * 70)
print("Fast Context integrates seamlessly with Agent for code search!")
print("=" * 70)
if __name__ == "__main__":
main()
@@ -0,0 +1,121 @@
"""
Example 5: File and Symbol Indexing
This example demonstrates the file and symbol indexing capabilities
that power Fast Context's rapid search.
Features:
- FileIndexer: Index file metadata for fast lookups
- SymbolIndexer: Extract code symbols (functions, classes, imports)
- Support for Python, JavaScript, TypeScript, Go, Rust, Java
"""
import time
from praisonaiagents.context.fast.indexer import (
FileIndexer,
SymbolIndexer,
SymbolType
)
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def main():
print("=" * 70)
print("File and Symbol Indexing")
print("=" * 70)
# File Indexer
print("\n1. File Indexer")
print("-" * 40)
start = time.perf_counter()
file_indexer = FileIndexer(workspace_path=WORKSPACE)
file_count = file_indexer.index()
file_time = (time.perf_counter() - start) * 1000
print(f" Indexed {file_count} files in {file_time:.0f}ms")
stats = file_indexer.get_stats()
print(f" Total size: {stats['total_size_mb']} MB")
print(f" Extensions: {list(stats['extensions'].keys())[:5]}")
# Find files by pattern
print("\n Finding files by pattern:")
py_files = file_indexer.find_by_pattern("**/*.py")
print(f" - Python files (*.py): {len(py_files)}")
agent_files = file_indexer.find_by_pattern("**/agent*.py")
print(f" - Agent files (agent*.py): {len(agent_files)}")
# Find files by name
print("\n Finding files by name:")
init_files = file_indexer.find_by_name("__init__.py", exact=True)
print(f" - __init__.py files: {len(init_files)}")
agent_matches = file_indexer.find_by_name("agent")
print(f" - Files containing 'agent': {len(agent_matches)}")
# Symbol Indexer
print("\n2. Symbol Indexer")
print("-" * 40)
start = time.perf_counter()
symbol_indexer = SymbolIndexer(workspace_path=WORKSPACE)
symbol_count = symbol_indexer.index()
symbol_time = (time.perf_counter() - start) * 1000
print(f" Indexed {symbol_count} symbols in {symbol_time:.0f}ms")
stats = symbol_indexer.get_stats()
print(f" Files indexed: {stats['total_files']}")
print(f" Symbols by type:")
for sym_type, count in stats['by_type'].items():
print(f" - {sym_type}: {count}")
# Find symbols by name
print("\n Finding symbols by name:")
agent_symbols = symbol_indexer.find_by_name("Agent", exact=True)
print(f" - 'Agent' (exact): {len(agent_symbols)} symbols")
for sym in agent_symbols[:3]:
print(f" {sym.symbol_type.value}: {sym.name} in {sym.file_path}:{sym.line_number}")
chat_symbols = symbol_indexer.find_by_name("chat")
print(f" - 'chat' (partial): {len(chat_symbols)} symbols")
for sym in chat_symbols[:3]:
print(f" {sym.symbol_type.value}: {sym.name} in {sym.file_path}:{sym.line_number}")
# Find symbols by type
print("\n Finding symbols by type:")
classes = symbol_indexer.find_by_type(SymbolType.CLASS)
print(f" - Classes: {len(classes)}")
for cls in classes[:5]:
print(f" class {cls.name} in {cls.file_path}:{cls.line_number}")
functions = symbol_indexer.find_by_type(SymbolType.FUNCTION)
print(f" - Functions: {len(functions)}")
methods = symbol_indexer.find_by_type(SymbolType.METHOD)
print(f" - Methods: {len(methods)}")
imports = symbol_indexer.find_by_type(SymbolType.IMPORT)
print(f" - Imports: {len(imports)}")
# Performance comparison
print("\n3. Performance Summary")
print("-" * 40)
print(f" File indexing: {file_time:.0f}ms for {file_count} files")
print(f" Symbol indexing: {symbol_time:.0f}ms for {symbol_count} symbols")
print(f" Files per second: {file_count / (file_time / 1000):.0f}")
print(f" Symbols per second: {symbol_count / (symbol_time / 1000):.0f}")
print("\n" + "=" * 70)
print("Indexing enables instant lookups for files and symbols!")
print("=" * 70)
if __name__ == "__main__":
main()
+163
View File
@@ -0,0 +1,163 @@
"""
Example 6: Context Injection for Agents
This example demonstrates how Fast Context results can be formatted
and injected into agent prompts for better code understanding.
Features:
- Multiple format styles (markdown, xml, plain)
- Token budget management
- Precision over recall (avoid context pollution)
"""
from praisonaiagents.context.fast import FastContext
from praisonaiagents.context.fast.context_injector import (
ContextInjector,
InjectionConfig,
inject_fast_context
)
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def main():
print("=" * 70)
print("Context Injection for Agents")
print("=" * 70)
# Search for code
fc = FastContext(workspace_path=WORKSPACE)
result = fc.search("class Agent")
print(f"\nSearch results: {result.total_files} files found")
# Markdown format (default)
print("\n1. Markdown Format")
print("-" * 40)
config_md = InjectionConfig(
max_tokens=2000,
max_files=3,
max_lines_per_file=20,
format_style="markdown",
include_line_numbers=True,
include_file_content=True,
prioritize_precision=True
)
injector_md = ContextInjector(config_md)
md_context = injector_md.format_context(result)
print(f" Format: Markdown")
print(f" Length: {len(md_context)} chars (~{len(md_context)//4} tokens)")
print(f" Preview:")
for line in md_context.split('\n')[:10]:
print(f" {line[:70]}")
print(" ...")
# XML format
print("\n2. XML Format")
print("-" * 40)
config_xml = InjectionConfig(
max_tokens=2000,
max_files=3,
format_style="xml",
prioritize_precision=True
)
injector_xml = ContextInjector(config_xml)
xml_context = injector_xml.format_context(result)
print(f" Format: XML")
print(f" Length: {len(xml_context)} chars")
print(f" Preview:")
for line in xml_context.split('\n')[:10]:
print(f" {line[:70]}")
print(" ...")
# Plain text format
print("\n3. Plain Text Format")
print("-" * 40)
config_plain = InjectionConfig(
max_tokens=2000,
max_files=3,
format_style="plain"
)
injector_plain = ContextInjector(config_plain)
plain_context = injector_plain.format_context(result)
print(f" Format: Plain")
print(f" Length: {len(plain_context)} chars")
print(f" Preview:")
for line in plain_context.split('\n')[:8]:
print(f" {line[:70]}")
print(" ...")
# Inject into prompts
print("\n4. Injecting into Prompts")
print("-" * 40)
system_prompt = "You are a helpful code assistant."
user_message = "Explain how the Agent class works."
injected = inject_fast_context(
result=result,
system_prompt=system_prompt,
user_message=user_message,
max_tokens=1500,
format_style="markdown"
)
print(f" Original system prompt: {len(system_prompt)} chars")
print(f" Injected system prompt: {len(injected['system_prompt'])} chars")
print(f" User message unchanged: {injected['user_message'] == user_message}")
# Token budget management
print("\n5. Token Budget Management")
print("-" * 40)
budgets = [500, 1000, 2000, 4000]
for budget in budgets:
config = InjectionConfig(max_tokens=budget, max_files=10)
injector = ContextInjector(config)
context = injector.format_context(result)
estimated_tokens = len(context) // 4
print(f" Budget: {budget} tokens -> {len(context)} chars (~{estimated_tokens} tokens)")
# Precision mode
print("\n6. Precision vs Recall")
print("-" * 40)
# High precision (fewer, more relevant results)
config_precision = InjectionConfig(
max_tokens=2000,
max_files=5,
prioritize_precision=True
)
injector_precision = ContextInjector(config_precision)
precision_context = injector_precision.format_context(result)
# Lower precision (more results)
config_recall = InjectionConfig(
max_tokens=2000,
max_files=10,
prioritize_precision=False
)
injector_recall = ContextInjector(config_recall)
recall_context = injector_recall.format_context(result)
print(f" High precision mode: {len(precision_context)} chars")
print(f" High recall mode: {len(recall_context)} chars")
print(f" Precision mode is more focused for better agent understanding")
print("\n" + "=" * 70)
print("Context injection formats search results for optimal agent use!")
print("=" * 70)
if __name__ == "__main__":
main()
+213
View File
@@ -0,0 +1,213 @@
"""
Example 7: Search Tools Deep Dive
This example demonstrates the individual search tools that power Fast Context:
- grep_search: Pattern-based search with regex support
- glob_search: File pattern matching
- read_file: Read file contents with line ranges
- list_directory: Directory listing with filtering
"""
import time
from praisonaiagents.context.fast.search_tools import (
grep_search,
glob_search,
read_file,
list_directory
)
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def main():
print("=" * 70)
print("Search Tools Deep Dive")
print("=" * 70)
# grep_search
print("\n1. grep_search - Pattern Search")
print("-" * 40)
# Simple pattern
start = time.perf_counter()
results = grep_search(
search_path=WORKSPACE,
pattern="def __init__",
max_results=10
)
elapsed = (time.perf_counter() - start) * 1000
print(f" Pattern: 'def __init__'")
print(f" Matches: {len(results)}")
print(f" Time: {elapsed:.0f}ms")
if results:
print(f" First match: {results[0]['path']}:{results[0]['line_number']}")
# Regex pattern
results = grep_search(
search_path=WORKSPACE,
pattern=r"class \w+Agent",
is_regex=True,
max_results=10
)
print(f"\n Regex: 'class \\w+Agent'")
print(f" Matches: {len(results)}")
for r in results[:3]:
print(f" - {r['path']}:{r['line_number']}: {r['content'][:50]}...")
# Case sensitive
results_sensitive = grep_search(
search_path=WORKSPACE,
pattern="Agent",
case_sensitive=True,
max_results=50
)
results_insensitive = grep_search(
search_path=WORKSPACE,
pattern="agent",
case_sensitive=False,
max_results=50
)
print(f"\n Case sensitive 'Agent': {len(results_sensitive)} matches")
print(f" Case insensitive 'agent': {len(results_insensitive)} matches")
# With context lines
results = grep_search(
search_path=WORKSPACE,
pattern="class Agent:",
context_lines=2,
max_results=3
)
print(f"\n With 2 context lines:")
if results:
print(f" {results[0]['path']}:")
if results[0].get('context'):
for line in results[0]['context'].split('\n')[:5]:
print(f" {line[:60]}")
# glob_search
print("\n2. glob_search - File Pattern Matching")
print("-" * 40)
# Find Python files
start = time.perf_counter()
results = glob_search(
search_path=WORKSPACE,
pattern="**/*.py",
max_results=100
)
elapsed = (time.perf_counter() - start) * 1000
print(f" Pattern: '**/*.py'")
print(f" Files found: {len(results)}")
print(f" Time: {elapsed:.0f}ms")
# Find specific files
results = glob_search(
search_path=WORKSPACE,
pattern="**/agent*.py",
max_results=20
)
print(f"\n Pattern: '**/agent*.py'")
print(f" Files found: {len(results)}")
for r in results[:5]:
print(f" - {r['path']} ({r['size']} bytes)")
# Find test files
results = glob_search(
search_path=WORKSPACE,
pattern="**/test_*.py",
max_results=50
)
print(f"\n Pattern: '**/test_*.py'")
print(f" Test files: {len(results)}")
# read_file
print("\n3. read_file - Read File Contents")
print("-" * 40)
# Read entire file
result = read_file(
filepath=f"{WORKSPACE}/praisonaiagents/__init__.py",
max_lines=20
)
print(f" File: praisonaiagents/__init__.py")
print(f" Success: {result['success']}")
print(f" Lines: {result.get('total_lines', 'N/A')}")
if result['success']:
lines = result['content'].split('\n')[:5]
print(f" Preview:")
for line in lines:
print(f" {line[:60]}")
# Read specific line range
result = read_file(
filepath=f"{WORKSPACE}/praisonaiagents/agent/agent.py",
start_line=1,
end_line=20
)
print(f"\n File: agent/agent.py (lines 1-20)")
print(f" Success: {result['success']}")
if result['success']:
lines = result['content'].split('\n')[:5]
print(f" Preview:")
for line in lines:
print(f" {line[:60]}")
# Read with context
result = read_file(
filepath=f"{WORKSPACE}/praisonaiagents/agent/agent.py",
start_line=100,
end_line=105,
context_lines=3
)
print(f"\n File: agent/agent.py (lines 100-105 with 3 context lines)")
print(f" Success: {result['success']}")
# list_directory
print("\n4. list_directory - Directory Listing")
print("-" * 40)
# List root
result = list_directory(
dir_path=WORKSPACE,
recursive=False,
max_entries=20
)
print(f" Directory: {WORKSPACE}")
print(f" Success: {result['success']}")
print(f" Entries: {len(result.get('entries', []))}")
dirs = [e for e in result.get('entries', []) if e.get('is_dir')]
files = [e for e in result.get('entries', []) if not e.get('is_dir')]
print(f" Directories: {len(dirs)}")
print(f" Files: {len(files)}")
# List recursive with depth limit
result = list_directory(
dir_path=f"{WORKSPACE}/praisonaiagents",
recursive=True,
max_depth=2,
max_entries=50
)
print(f"\n Directory: praisonaiagents (recursive, depth=2)")
print(f" Entries: {len(result.get('entries', []))}")
# Show directory structure
print(f" Structure:")
for entry in result.get('entries', [])[:10]:
prefix = "📁" if entry.get('is_dir') else "📄"
print(f" {prefix} {entry['name']}")
print("\n" + "=" * 70)
print("These search tools power Fast Context's rapid code search!")
print("=" * 70)
if __name__ == "__main__":
main()
@@ -0,0 +1,227 @@
"""
Example 8: Comprehensive Benchmark
This example provides a comprehensive benchmark comparing Fast Context
performance against traditional search approaches.
Metrics measured:
- Search latency
- Throughput (queries per second)
- Cache effectiveness
- Parallel execution speedup
"""
import time
import statistics
from praisonaiagents.context.fast import FastContext
from praisonaiagents.context.fast.search_tools import grep_search
from praisonaiagents.context.fast.indexer import FileIndexer, SymbolIndexer
WORKSPACE = "/Users/praison/praisonai-package/src/praisonai-agents"
def benchmark_search_latency():
"""Benchmark search latency."""
print("\n1. Search Latency Benchmark")
print("-" * 50)
fc = FastContext(workspace_path=WORKSPACE, cache_enabled=False)
queries = [
"def __init__",
"class Agent",
"import logging",
"async def",
"from typing import",
]
# Warm up
fc.search("warmup")
# Benchmark each query
latencies = []
for query in queries:
times = []
for _ in range(3): # 3 runs per query
start = time.perf_counter()
fc.search(query)
elapsed = (time.perf_counter() - start) * 1000
times.append(elapsed)
avg = statistics.mean(times)
latencies.append(avg)
print(f" '{query}': {avg:.1f}ms (min: {min(times):.1f}, max: {max(times):.1f})")
print(f"\n Average latency: {statistics.mean(latencies):.1f}ms")
print(f" Median latency: {statistics.median(latencies):.1f}ms")
print(f" P95 latency: {sorted(latencies)[int(len(latencies)*0.95)]:.1f}ms")
def benchmark_throughput():
"""Benchmark queries per second."""
print("\n2. Throughput Benchmark")
print("-" * 50)
fc = FastContext(workspace_path=WORKSPACE, cache_enabled=True)
queries = ["def", "class", "import", "return", "if", "for", "while", "try"]
# Cold run (no cache)
fc.clear_cache()
start = time.perf_counter()
for query in queries:
fc.search(query)
cold_time = time.perf_counter() - start
cold_qps = len(queries) / cold_time
print(f" Cold (no cache):")
print(f" Queries: {len(queries)}")
print(f" Time: {cold_time*1000:.0f}ms")
print(f" Throughput: {cold_qps:.1f} queries/second")
# Warm run (cached)
start = time.perf_counter()
for _ in range(10): # 10 iterations
for query in queries:
fc.search(query)
warm_time = time.perf_counter() - start
warm_qps = (len(queries) * 10) / warm_time
print(f"\n Warm (cached):")
print(f" Queries: {len(queries) * 10}")
print(f" Time: {warm_time*1000:.0f}ms")
print(f" Throughput: {warm_qps:.0f} queries/second")
print(f"\n Cache speedup: {warm_qps/cold_qps:.0f}x")
def benchmark_cache_effectiveness():
"""Benchmark cache hit rate and speedup."""
print("\n3. Cache Effectiveness Benchmark")
print("-" * 50)
fc = FastContext(workspace_path=WORKSPACE, cache_enabled=True)
fc.clear_cache()
queries = ["Agent", "chat", "memory", "tool", "execute"]
# First pass - cache misses
miss_times = []
for query in queries:
start = time.perf_counter()
result = fc.search(query)
elapsed = (time.perf_counter() - start) * 1000
miss_times.append(elapsed)
assert not result.from_cache
# Second pass - cache hits
hit_times = []
for query in queries:
start = time.perf_counter()
result = fc.search(query)
elapsed = (time.perf_counter() - start) * 1000
hit_times.append(elapsed)
assert result.from_cache
avg_miss = statistics.mean(miss_times)
avg_hit = statistics.mean(hit_times)
print(f" Cache miss (first query): {avg_miss:.1f}ms avg")
print(f" Cache hit (repeat query): {avg_hit:.2f}ms avg")
print(f" Cache speedup: {avg_miss/avg_hit:.0f}x faster")
print(f" Time saved per cached query: {avg_miss - avg_hit:.1f}ms")
def benchmark_parallel_execution():
"""Benchmark parallel vs sequential execution."""
print("\n4. Parallel Execution Benchmark")
print("-" * 50)
queries = ["def", "class", "import", "return", "async"]
# Sequential (one at a time)
start = time.perf_counter()
for query in queries:
grep_search(WORKSPACE, query, max_results=20)
seq_time = (time.perf_counter() - start) * 1000
# Parallel (Fast Context)
fc = FastContext(workspace_path=WORKSPACE, cache_enabled=False)
start = time.perf_counter()
for query in queries:
fc.search(query)
par_time = (time.perf_counter() - start) * 1000
print(f" Sequential: {seq_time:.0f}ms")
print(f" Parallel (Fast Context): {par_time:.0f}ms")
print(f" Speedup: {seq_time/par_time:.1f}x")
def benchmark_indexing():
"""Benchmark indexing performance."""
print("\n5. Indexing Performance Benchmark")
print("-" * 50)
# File indexing
start = time.perf_counter()
file_indexer = FileIndexer(workspace_path=WORKSPACE)
file_count = file_indexer.index()
file_time = (time.perf_counter() - start) * 1000
print(f" File Indexer:")
print(f" Files: {file_count}")
print(f" Time: {file_time:.0f}ms")
print(f" Rate: {file_count/(file_time/1000):.0f} files/second")
# Symbol indexing
start = time.perf_counter()
symbol_indexer = SymbolIndexer(workspace_path=WORKSPACE)
symbol_count = symbol_indexer.index()
symbol_time = (time.perf_counter() - start) * 1000
print(f"\n Symbol Indexer:")
print(f" Symbols: {symbol_count}")
print(f" Time: {symbol_time:.0f}ms")
print(f" Rate: {symbol_count/(symbol_time/1000):.0f} symbols/second")
# Lookup performance
start = time.perf_counter()
for _ in range(100):
file_indexer.find_by_pattern("**/*.py")
lookup_time = (time.perf_counter() - start) * 1000
print(f"\n Lookup Performance:")
print(f" 100 pattern lookups: {lookup_time:.1f}ms")
print(f" Per lookup: {lookup_time/100:.2f}ms")
def main():
print("=" * 70)
print("Fast Context Comprehensive Benchmark")
print("=" * 70)
print(f"\nWorkspace: {WORKSPACE}")
benchmark_search_latency()
benchmark_throughput()
benchmark_cache_effectiveness()
benchmark_parallel_execution()
benchmark_indexing()
print("\n" + "=" * 70)
print("Summary")
print("=" * 70)
print("""
Fast Context provides:
✓ Sub-100ms search latency for most queries
✓ 100+ queries/second with caching
✓ 100x+ speedup for cached queries
✓ 2-5x speedup from parallel execution
✓ 1000+ files/second indexing rate
These benchmarks demonstrate why Fast Context is ideal for
rapid code search in AI agent workflows.
""")
if __name__ == "__main__":
main()
@@ -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.")
@@ -0,0 +1,82 @@
"""MCP Schema Converter Example
Demonstrates how to convert Python functions to MCP-compatible schemas.
Usage:
python mcp-schema-converter.py
"""
from typing import List, Dict, Any, Optional
from praisonaiagents.mcp import function_to_mcp_schema, get_tool_metadata
# Example functions with various type annotations
def simple_search(query: str) -> str:
"""Search for information."""
return f"Results for {query}"
def advanced_search(
query: str,
max_results: int = 10,
include_images: bool = False,
filters: Optional[Dict[str, str]] = None
) -> Dict[str, Any]:
"""Advanced search with multiple options.
Args:
query: The search query
max_results: Maximum results to return
include_images: Whether to include images
filters: Optional filters to apply
"""
return {}
def process_items(items: List[str], batch_size: int = 5) -> List[Dict[str, Any]]:
"""Process a list of items.
Args:
items: List of items to process
batch_size: Size of each batch
"""
return []
async def async_fetch(url: str, timeout: float = 30.0) -> Dict[str, Any]:
"""Fetch data from URL asynchronously.
Args:
url: The URL to fetch
timeout: Request timeout in seconds
"""
return {}
def main():
functions = [simple_search, advanced_search, process_items, async_fetch]
print("=" * 60)
print("MCP Schema Converter Demo")
print("=" * 60)
for func in functions:
print(f"\n### Function: {func.__name__}")
# Get metadata
metadata = get_tool_metadata(func)
print(f"Name: {metadata['name']}")
print(f"Description: {metadata['description']}")
# Get full schema
schema = function_to_mcp_schema(func)
print("\nMCP Schema:")
import json
print(json.dumps(schema, indent=2))
print("-" * 40)
if __name__ == "__main__":
main()
+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,79 @@
"""MCP Tools Server with Async Functions
Example showing how to expose async Python functions as MCP tools.
Usage:
python mcp-tools-server-async.py
"""
import asyncio
from praisonaiagents.mcp import ToolsMCPServer
async def async_search(query: str, delay: float = 0.1) -> dict:
"""Async search function.
Args:
query: The search query
delay: Simulated delay in seconds
"""
await asyncio.sleep(delay) # Simulate async operation
return {
"query": query,
"results": [f"Async result for {query}"]
}
async def async_fetch_url(url: str) -> dict:
"""Fetch content from a URL asynchronously.
Args:
url: The URL to fetch
"""
# Mock implementation - use aiohttp in production
await asyncio.sleep(0.1)
return {
"url": url,
"status": 200,
"content": f"Content from {url}"
}
async def async_process_data(data: str, iterations: int = 3) -> dict:
"""Process data asynchronously with multiple iterations.
Args:
data: The data to process
iterations: Number of processing iterations
"""
results = []
for i in range(iterations):
await asyncio.sleep(0.05)
results.append(f"Iteration {i+1}: processed {data}")
return {
"input": data,
"iterations": iterations,
"results": results
}
# Sync function can also be mixed with async
def sync_calculate(a: int, b: int) -> dict:
"""Synchronous calculation (can be mixed with async tools)."""
return {"sum": a + b, "product": a * b}
if __name__ == "__main__":
server = ToolsMCPServer(name="async-tools-server")
# Register both async and sync tools
server.register_tools([
async_search,
async_fetch_url,
async_process_data,
sync_calculate
])
print(f"Registered tools: {server.get_tool_names()}")
server.run()
@@ -0,0 +1,27 @@
"""MCP Server with Built-in PraisonAI Tools
Expose PraisonAI's built-in tools as an MCP server.
Usage:
# Set required API keys first
export TAVILY_API_KEY=your_key
python mcp-tools-server-builtin.py
"""
from praisonaiagents.mcp import launch_tools_mcp_server
if __name__ == "__main__":
# Expose built-in PraisonAI tools as MCP server
# This allows external MCP clients to use these tools
launch_tools_mcp_server(
tool_names=[
"tavily_search", # Web search via Tavily
"duckduckgo_search", # Web search via DuckDuckGo
# Add more built-in tools as needed
],
name="praisonai-builtin-tools",
transport="stdio" # Use "sse" for web clients
)
@@ -0,0 +1,25 @@
"""Simple MCP Tools Server Example
Minimal example of exposing Python functions as MCP tools.
Usage:
python mcp-tools-server-simple.py
"""
from praisonaiagents.mcp import ToolsMCPServer
def search(query: str) -> dict:
"""Search for information."""
return {"query": query, "results": [f"Result for {query}"]}
def add(a: int, b: int) -> dict:
"""Add two numbers."""
return {"result": a + b}
if __name__ == "__main__":
server = ToolsMCPServer(name="simple-tools")
server.register_tools([search, add])
server.run()
@@ -0,0 +1,65 @@
"""MCP Tools Server with SSE Transport
Expose Python functions as MCP tools via Server-Sent Events (SSE).
This is useful for web-based clients and remote access.
Usage:
python mcp-tools-server-sse.py
# Then connect from another agent:
# agent = Agent(tools=MCP("http://localhost:8080/sse"))
"""
from praisonaiagents.mcp import ToolsMCPServer
def search_web(query: str, max_results: int = 5) -> dict:
"""Search the web for information.
Args:
query: The search query
max_results: Maximum results to return
"""
return {
"query": query,
"results": [
{"title": f"Result {i+1}", "url": f"https://example.com/{i+1}"}
for i in range(max_results)
]
}
def calculate(expression: str) -> dict:
"""Evaluate a mathematical expression.
Args:
expression: Math expression (e.g., "2 + 2 * 3")
"""
try:
allowed = set("0123456789+-*/.() ")
if not all(c in allowed for c in expression):
return {"error": "Invalid characters"}
return {"expression": expression, "result": eval(expression)}
except Exception as e:
return {"error": str(e)}
def get_time() -> dict:
"""Get the current time."""
from datetime import datetime
now = datetime.now()
return {
"time": now.strftime("%H:%M:%S"),
"date": now.strftime("%Y-%m-%d"),
"timestamp": now.isoformat()
}
if __name__ == "__main__":
server = ToolsMCPServer(name="sse-tools-server")
server.register_tools([search_web, calculate, get_time])
print("Starting SSE MCP Server...")
print("Connect with: MCP('http://localhost:8080/sse')")
server.run_sse(host="0.0.0.0", port=8080)
+140
View File
@@ -0,0 +1,140 @@
"""MCP Tools Server Example
Expose Python functions as MCP tools for external clients like Claude Desktop, Cursor, etc.
Usage:
# Run as stdio server (for Claude Desktop)
python mcp-tools-server.py
# Run as SSE server (for web clients)
python mcp-tools-server.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
"""
# 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
allowed_chars = set("0123456789+-*/.() ")
if not all(c in allowed_chars for c in expression):
return {"error": "Invalid characters in expression"}
result = eval(expression) # 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()
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,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")
+1
View File
@@ -86,3 +86,4 @@ build
graph.py
chroma_db/
.qodo
*_COMPARISON_ANALYSIS.md
+405 -13
View File
@@ -84,6 +84,64 @@ agents = PraisonAIAgents(agents=[researcher, writer], tasks=[research_task, writ
result = agents.start()
```
### Planning Mode 🆕
#### Single Agent Planning (Simplest)
Enable planning for any single agent:
```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")
```
#### Multi-Agent Planning
Plan before execution with multiple agents:
```python
from praisonaiagents import Agent, Task, PraisonAIAgents
researcher = Agent(name="Researcher", role="Research Analyst")
writer = Agent(name="Writer", role="Content Writer")
research_task = Task(description="Research AI benefits", agent=researcher)
write_task = Task(description="Write summary", agent=writer)
agents = PraisonAIAgents(
agents=[researcher, writer],
tasks=[research_task, write_task],
planning=True, # Enable planning
planning_tools=[search_web], # Tools for planning research
planning_reasoning=True, # Chain-of-thought reasoning
auto_approve_plan=True # Auto-approve plans
)
result = agents.start()
```
**Features:**
- **Single Agent Planning**: `Agent(planning=True)` - simplest API
- **Multi-Agent Planning**: `PraisonAIAgents(planning=True)` - orchestrated
- **Planning Tools**: Research with tools before creating plans
- **Reasoning Mode**: Chain-of-thought planning with tool use
- **Todo Lists**: Auto-generated from plans with progress tracking
- **Read-Only Mode**: `plan_mode=True` restricts agents to safe tools
- **Approval Flow**: Review and approve plans before execution
- **Persistence**: Plans saved to `.praison/plans/` as markdown
---
## Agent Types
@@ -216,27 +274,307 @@ router = RouterAgent(
## Memory
Persistent memory with short-term, long-term, and graph memory support.
PraisonAI provides a comprehensive memory system.
### Quick Start (Zero Dependencies)
```python
from praisonaiagents import Agent, Memory
from praisonaiagents import Agent
from praisonaiagents.memory import FileMemory
# Simple file-based memory (no extra dependencies)
memory = FileMemory(user_id="user123")
agent = Agent(
name="Memory Agent",
instructions="You are a helpful assistant with memory.",
memory=memory # or memory=True for auto FileMemory
)
response = agent.chat("My name is John and I prefer Python")
# Memory automatically stores this context
```
### Advanced Memory (with providers)
```python
from praisonaiagents.memory import Memory
memory = Memory(
provider="chroma", # or "mongodb", "rag"
provider="chroma", # or "mongodb", "mem0", "rag"
use_short_term=True,
use_long_term=True
)
agent = Agent(name="Memory Agent", memory=memory)
```
### Memory Types
| Type | Description | Storage |
|------|-------------|---------|
| **Short-term** | Rolling buffer of recent context | JSON/SQLite |
| **Long-term** | Persistent facts and knowledge | JSON/SQLite/Vector DB |
| **Entity** | Named entities (people, places, concepts) | JSON/SQLite |
| **Episodic** | Date-based interaction memories | JSON files |
| **Graph** | Relationship-based memory | Mem0/Neo4j |
### Session Save/Resume
```python
from praisonaiagents.memory import FileMemory
memory = FileMemory(user_id="user123")
# Add context during conversation
memory.add_short_term("User is working on ML project")
memory.add_long_term("User prefers Python", importance=0.9)
# Save session for later
memory.save_session("ml_project", conversation_history=[...])
# Resume later
session_data = memory.resume_session("ml_project")
# List all sessions
sessions = memory.list_sessions()
```
### Context Compression
```python
# Auto-compress when memory gets full
memory.auto_compress_if_needed(threshold_percent=0.7)
# Manual compression with LLM summarization
def llm_summarize(prompt):
return agent.chat(prompt)
summary = memory.compress(llm_func=llm_summarize, max_items=10)
```
### Checkpointing
```python
# Create checkpoint before risky operations
checkpoint_id = memory.create_checkpoint("before_refactor", include_files=["main.py"])
# Restore if needed
memory.restore_checkpoint(checkpoint_id, restore_files=True)
# List checkpoints
checkpoints = memory.list_checkpoints()
```
### Memory Slash Commands
```python
# Handle slash commands programmatically
result = memory.handle_command("/memory show")
result = memory.handle_command("/memory add User likes coffee")
result = memory.handle_command("/memory search Python")
result = memory.handle_command("/memory save my_session")
result = memory.handle_command("/memory compress")
result = memory.handle_command("/memory checkpoint")
result = memory.handle_command("/memory help")
```
**Available Commands:**
- `/memory show` - Display stats and recent items
- `/memory add <content>` - Add to long-term memory
- `/memory clear [short|all]` - Clear memory
- `/memory search <query>` - Search memories
- `/memory save <name>` - Save session
- `/memory resume <name>` - Resume session
- `/memory sessions` - List saved sessions
- `/memory compress` - Compress short-term memory
- `/memory checkpoint [name]` - Create checkpoint
- `/memory restore <id>` - Restore checkpoint
- `/memory checkpoints` - List checkpoints
- `/memory refresh` - Reload from disk
---
## Rules & Instructions
PraisonAI automatically discovers and applies rules from multiple sources.
### Supported Instruction Files
| File | Description | Priority |
|------|-------------|----------|
| `PRAISON.md` | PraisonAI native instructions | High (500) |
| `PRAISON.local.md` | Local overrides (gitignored) | Higher (600) |
| `CLAUDE.md` | Claude Code memory file | High (500) |
| `CLAUDE.local.md` | Local overrides (gitignored) | Higher (600) |
| `AGENTS.md` | OpenAI Codex CLI instructions | High (500) |
| `GEMINI.md` | Gemini CLI memory file | High (500) |
| `.cursorrules` | Cursor IDE rules (legacy) | High (500) |
| `.windsurfrules` | Windsurf IDE rules (legacy) | High (500) |
| `.claude/rules/*.md` | Claude Code modular rules | Medium (50) |
| `.windsurf/rules/*.md` | Windsurf modular rules | Medium (50) |
| `.cursor/rules/*.mdc` | Cursor modular rules | Medium (50) |
| `.praison/rules/*.md` | Workspace rules | Medium (0) |
| `~/.praison/rules/*.md` | Global rules | Low (-1000) |
### Auto-Discovery
Rules are automatically loaded when you create an Agent:
```python
from praisonaiagents import Agent
# Agent auto-discovers CLAUDE.md, AGENTS.md, GEMINI.md, etc.
agent = Agent(
name="Memory Agent",
memory=memory
name="Assistant",
instructions="You are helpful."
)
# Rules are injected into system prompt automatically
```
### Rule File Format
Rules support YAML frontmatter for advanced configuration:
```markdown
---
description: Python coding guidelines
globs: ["**/*.py", "**/*.pyx"]
activation: always # always, glob, manual, ai_decision
priority: 10
---
# Python Guidelines
- Use type hints for all functions
- Follow PEP 8 style guide
- Write docstrings for public APIs
```
### Activation Modes
| Mode | Description |
|------|-------------|
| `always` | Always applied (default) |
| `glob` | Applied when file matches glob pattern |
| `manual` | Only when explicitly invoked via @mention |
| `ai_decision` | AI decides when to apply |
### Programmatic Rules Management
```python
from praisonaiagents.memory import RulesManager
rules = RulesManager(workspace_path="/path/to/project")
# Get all active rules
active = rules.get_active_rules()
# Get rules for specific file
python_rules = rules.get_rules_for_file("src/main.py")
# Build context for LLM
context = rules.build_rules_context(file_path="src/main.py")
# Create new rule
rules.create_rule(
name="testing",
content="Always write tests first",
globs=["**/*.test.*"],
activation="glob"
)
# Get stats
stats = rules.get_stats()
# {'total_rules': 5, 'root_rules': 3, 'workspace_rules': 2, ...}
```
### Storage Structure
```
project/
├── CLAUDE.md # Auto-loaded (Claude Code)
├── CLAUDE.local.md # Local overrides (gitignored)
├── AGENTS.md # Auto-loaded (Codex CLI)
├── GEMINI.md # Auto-loaded (Gemini CLI)
├── PRAISON.md # Auto-loaded (PraisonAI)
├── PRAISON.local.md # Local overrides (gitignored)
├── .claude/rules/ # Claude Code modular rules
├── .windsurf/rules/ # Windsurf rules
├── .cursor/rules/ # Cursor rules
├── .praison/
│ ├── rules/ # Workspace rules
│ │ ├── python.md
│ │ └── testing.md
│ ├── workflows/ # Reusable workflows
│ │ └── deploy.md
│ ├── hooks.json # Pre/post operation hooks
│ └── memory/
│ └── {user_id}/
│ ├── short_term.json
│ ├── long_term.json
│ ├── sessions/
│ └── checkpoints/
└── ~/.praison/
└── rules/ # Global rules
└── global.md
```
### @Import Syntax
Include other files in your rules:
```markdown
# CLAUDE.md
See @README for project overview
See @docs/architecture.md for system design
@~/.praison/my-preferences.md
```
### 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"
```
### Workflows
Create reusable multi-step 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"}
)
```
**Memory Types:**
- **Short-term**: Recent conversation context
- **Long-term**: Persistent storage across sessions
- **Graph**: Relationship-based memory with Mem0
### Hooks
```python
from praisonaiagents.memory import HooksManager
hooks = HooksManager()
# Register Python hooks
hooks.register("pre_write_code", lambda ctx: print(f"Writing {ctx['file']}"))
# Or configure in .praison/hooks.json
```
---
@@ -317,12 +655,23 @@ agent = Agent(tools=[my_tool])
## MCP (Model Context Protocol)
Connect to MCP servers for extended capabilities.
Connect to MCP servers for extended capabilities. Supports **Protocol Revision 2025-11-25**.
### MCP Transports
| Transport | URL Pattern | Use Case |
|-----------|-------------|----------|
| **stdio** | Command string | Local NPX/Python servers |
| **Streamable HTTP** | `http(s)://...` | Production deployments |
| **WebSocket** | `ws://` or `wss://` | Real-time bidirectional |
| **SSE (Legacy)** | `http(s)://.../sse` | Backward compatibility |
### MCP Client (Consume MCP Servers)
```python
from praisonaiagents import Agent, MCP
# Local MCP server
# stdio - Local MCP server
agent = Agent(
tools=MCP(
command="npx",
@@ -330,9 +679,19 @@ agent = Agent(
)
)
# Remote MCP server (SSE)
# Streamable HTTP - Production server
agent = Agent(
tools=MCP(url="http://localhost:8080/sse")
tools=MCP("https://api.example.com/mcp")
)
# WebSocket - Real-time bidirectional (NEW)
agent = Agent(
tools=MCP("wss://api.example.com/mcp", auth_token="your-token")
)
# SSE (Legacy) - Backward compatibility
agent = Agent(
tools=MCP("http://localhost:8080/sse")
)
# With environment variables
@@ -345,6 +704,39 @@ agent = Agent(
)
```
### MCP Features (Protocol 2025-11-25)
- **Session Management**: Automatic `Mcp-Session-Id` handling
- **Protocol Versioning**: `Mcp-Protocol-Version` header on all requests
- **Resumability**: SSE stream recovery via `Last-Event-ID`
- **Security**: Origin validation, DNS rebinding prevention
- **WebSocket**: Auto-reconnect with exponential backoff
### MCP Server (Expose Tools as MCP Server)
Expose your Python functions as MCP tools for external clients (Claude Desktop, Cursor, etc.).
```python
from praisonaiagents.mcp import ToolsMCPServer
# Define your tools as regular Python functions
def search_web(query: str, max_results: int = 5) -> dict:
"""Search the web for information."""
return {"results": [f"Result for {query}"]}
def calculate(expression: str) -> dict:
"""Evaluate a mathematical expression."""
return {"result": eval(expression)}
# Create and run MCP server
server = ToolsMCPServer(name="my-tools")
server.register_tools([search_web, calculate])
server.run() # Starts stdio server (for Claude Desktop)
# Or run as SSE server for web clients
# server.run_sse(host="0.0.0.0", port=8080)
```
---
## Native Web Search
@@ -0,0 +1,287 @@
"""
Fast Context Usage Examples.
This file demonstrates how to use the Fast Context feature for rapid code search.
Fast Context is inspired by Windsurf's SWE-grep approach:
- Parallel tool execution (up to 8 concurrent calls)
- Limited turns (max 4) for fast response
- Lightweight model for search operations
- Result caching for repeated queries
"""
# Example 1: Basic FastContext Usage
# ==================================
def example_basic_search():
"""Basic search using FastContext."""
from praisonaiagents.context.fast import FastContext
# Create FastContext instance
fc = FastContext(workspace_path=".")
# Search for a pattern
result = fc.search("def authenticate")
print(f"Found {result.total_files} files")
print(f"Search time: {result.search_time_ms}ms")
# Print results
for file_match in result.files[:5]:
print(f" - {file_match.path}")
for lr in file_match.line_ranges:
print(f" Lines {lr.start}-{lr.end}")
# Example 2: Search with Patterns
# ===============================
def example_search_with_patterns():
"""Search with include/exclude patterns."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=".")
# Search only in Python files
result = fc.search(
"class Agent",
include_patterns=["**/*.py"],
exclude_patterns=["**/tests/**", "**/test_*.py"]
)
print(f"Found {result.total_files} files matching 'class Agent'")
for f in result.files:
print(f" - {f.path}")
# Example 3: Get Context for Agent
# ================================
def example_get_context_for_agent():
"""Get formatted context string for an agent."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=".")
# Get context formatted for injection into an agent
context = fc.get_context_for_agent(
"authentication handlers",
max_files=5,
max_lines_per_file=50
)
print("Context for agent:")
print(context[:500] + "..." if len(context) > 500 else context)
# Example 4: Agent with FastContext Integration
# =============================================
def example_agent_with_fast_context():
"""Use FastContext with an Agent."""
from praisonaiagents import Agent
# Create agent with FastContext enabled
agent = Agent(
name="CodeAssistant",
instructions="You are a helpful code assistant.",
fast_context=True,
fast_context_path=".", # Current directory
fast_context_model="gpt-4o-mini",
fast_context_max_turns=4,
fast_context_parallelism=8
)
# Use delegate_to_fast_context for code search
context = agent.delegate_to_fast_context("find authentication code")
if context:
print("Found relevant code context:")
print(context[:300] + "...")
else:
print("No relevant code found")
# Example 5: File Indexer
# =======================
def example_file_indexer():
"""Use FileIndexer for fast file lookups."""
from praisonaiagents.context.fast.indexer import FileIndexer
# Create and run indexer
indexer = FileIndexer(workspace_path=".")
count = indexer.index()
print(f"Indexed {count} files")
# Find files by pattern
py_files = indexer.find_by_pattern("**/*.py")
print(f"Found {len(py_files)} Python files")
# Find files by extension
js_files = indexer.find_by_extension(".js")
print(f"Found {len(js_files)} JavaScript files")
# Get stats
stats = indexer.get_stats()
print(f"Total size: {stats['total_size_mb']} MB")
# Example 6: Symbol Indexer
# =========================
def example_symbol_indexer():
"""Use SymbolIndexer to find code symbols."""
from praisonaiagents.context.fast.indexer import SymbolIndexer, SymbolType
# Create and run indexer
indexer = SymbolIndexer(workspace_path=".")
count = indexer.index()
print(f"Indexed {count} symbols")
# Find symbols by name
agent_symbols = indexer.find_by_name("Agent")
print(f"Found {len(agent_symbols)} symbols containing 'Agent'")
# Find all classes
classes = indexer.find_by_type(SymbolType.CLASS)
print(f"Found {len(classes)} classes")
# Get stats
stats = indexer.get_stats()
print(f"Symbols by type: {stats['by_type']}")
# Example 7: Context Injector
# ===========================
def example_context_injector():
"""Use ContextInjector to format results for agents."""
from praisonaiagents.context.fast import FastContext
from praisonaiagents.context.fast.context_injector import (
ContextInjector,
InjectionConfig
)
# Search for code
fc = FastContext(workspace_path=".")
result = fc.search("def main")
# Configure injection
config = InjectionConfig(
max_tokens=2000,
max_files=5,
format_style="markdown", # or "xml", "plain"
prioritize_precision=True
)
# Inject into prompts
injector = ContextInjector(config)
formatted = injector.format_context(result)
print("Formatted context:")
print(formatted[:500] + "..." if len(formatted) > 500 else formatted)
# Example 8: Environment Variable Configuration
# =============================================
def example_env_config():
"""Configure FastContext via environment variables."""
import os
# Set environment variables (usually done in shell or .env file)
os.environ["FAST_CONTEXT_MODEL"] = "gpt-4o-mini"
os.environ["FAST_CONTEXT_MAX_TURNS"] = "4"
os.environ["FAST_CONTEXT_PARALLELISM"] = "8"
os.environ["FAST_CONTEXT_TIMEOUT"] = "30.0"
os.environ["FAST_CONTEXT_CACHE"] = "true"
os.environ["FAST_CONTEXT_CACHE_TTL"] = "300"
from praisonaiagents.context.fast import FastContext
# FastContext will use environment variables as defaults
fc = FastContext(workspace_path=".")
print(f"Model: {fc.model}")
print(f"Max turns: {fc.max_turns}")
print(f"Max parallel: {fc.max_parallel}")
print(f"Cache enabled: {fc.cache_enabled}")
# Example 9: Caching
# ==================
def example_caching():
"""Demonstrate result caching."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(
workspace_path=".",
cache_enabled=True,
cache_ttl=300 # 5 minutes
)
# First search - not cached
result1 = fc.search("def main")
print(f"First search: from_cache={result1.from_cache}, time={result1.search_time_ms}ms")
# Second search - cached
result2 = fc.search("def main")
print(f"Second search: from_cache={result2.from_cache}, time={result2.search_time_ms}ms")
# Clear cache
fc.clear_cache()
# Third search - not cached
result3 = fc.search("def main")
print(f"Third search: from_cache={result3.from_cache}, time={result3.search_time_ms}ms")
# Example 10: Context Manager
# ===========================
def example_context_manager():
"""Use FastContext as a context manager."""
from praisonaiagents.context.fast import FastContext
with FastContext(workspace_path=".") as fc:
result = fc.search("import")
print(f"Found {result.total_files} files with imports")
# Resources are automatically cleaned up
if __name__ == "__main__":
print("=" * 60)
print("Fast Context Examples")
print("=" * 60)
print("\n1. Basic Search")
print("-" * 40)
example_basic_search()
print("\n2. Search with Patterns")
print("-" * 40)
example_search_with_patterns()
print("\n3. Get Context for Agent")
print("-" * 40)
example_get_context_for_agent()
print("\n5. File Indexer")
print("-" * 40)
example_file_indexer()
print("\n9. Caching")
print("-" * 40)
example_caching()
print("\n10. Context Manager")
print("-" * 40)
example_context_manager()
print("\n" + "=" * 60)
print("All examples completed!")
@@ -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()
+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()
@@ -55,6 +55,49 @@ except ImportError:
from .session import Session
from .memory.memory import Memory
from .guardrails import GuardrailResult, LLMGuardrail
# Fast Context support (lazy loaded to avoid performance impact)
def __getattr__(name):
"""Lazy load FastContext to avoid impacting package load time."""
if name == "FastContext":
from praisonaiagents.context.fast import FastContext
return FastContext
elif name == "FastContextResult":
from praisonaiagents.context.fast import FastContextResult
return FastContextResult
elif name == "FileMatch":
from praisonaiagents.context.fast import FileMatch
return FileMatch
elif name == "LineRange":
from praisonaiagents.context.fast import LineRange
return LineRange
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
# 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:
@@ -247,3 +290,11 @@ if _mcp_available:
if FlowDisplay is not None:
__all__.extend(['FlowDisplay', 'track_workflow'])
# Add FastContext exports (lazy loaded)
__all__.extend([
'FastContext',
'FastContextResult',
'FileMatch',
'LineRange'
])
@@ -238,7 +238,17 @@ class Agent:
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
claude_memory: Optional[Union[bool, Any]] = None,
plan_mode: bool = False,
planning: bool = False,
planning_tools: Optional[List[Any]] = None,
planning_reasoning: bool = False,
fast_context: bool = False,
fast_context_path: Optional[str] = None,
fast_context_model: str = "gpt-4o-mini",
fast_context_max_turns: int = 4,
fast_context_parallelism: int = 8,
fast_context_timeout: float = 30.0
):
"""Initialize an Agent instance.
@@ -524,11 +534,20 @@ 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
@@ -572,6 +591,15 @@ Your Goal: {self.goal}
self._knowledge_config = knowledge_config
self.knowledge = None # Will be initialized on first use
# Fast Context configuration (lazy loaded)
self.fast_context_enabled = fast_context
self.fast_context_path = fast_context_path or os.getcwd()
self.fast_context_model = fast_context_model
self.fast_context_max_turns = fast_context_max_turns
self.fast_context_parallelism = fast_context_parallelism
self.fast_context_timeout = fast_context_timeout
self._fast_context_instance = None # Lazy loaded
@property
def console(self):
"""Lazily initialize Rich Console only when needed."""
@@ -606,6 +634,37 @@ 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.
@@ -669,6 +728,108 @@ Your Goal: {self.goal}
return supports_prompt_caching(model_name)
@property
def fast_context(self):
"""Lazily initialize FastContext instance when needed.
Returns:
FastContext instance or None if not enabled
"""
if not self.fast_context_enabled:
return None
if self._fast_context_instance is None:
try:
from ..context.fast import FastContext
self._fast_context_instance = FastContext(
workspace_path=self.fast_context_path,
model=self.fast_context_model,
max_turns=self.fast_context_max_turns,
max_parallel=self.fast_context_parallelism,
timeout=self.fast_context_timeout,
verbose=self.verbose
)
except ImportError:
logging.warning("FastContext not available")
return None
return self._fast_context_instance
def delegate_to_fast_context(self, query: str) -> Optional[str]:
"""Delegate a code search query to FastContext subagent.
This method uses the FastContext subagent to rapidly search
the codebase and return relevant context.
Args:
query: Natural language search query
Returns:
Formatted context string or None if FastContext not available
"""
if not self.fast_context_enabled or self.fast_context is None:
return None
try:
result = self.fast_context.search(query)
if result.total_files > 0:
return self.fast_context.get_context_for_agent(query)
return None
except Exception as e:
logging.warning(f"FastContext search failed: {e}")
return None
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.
@@ -1072,6 +1233,12 @@ Your Goal: {self.goal}
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()
@@ -2398,10 +2565,88 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
def run(self, prompt: str, **kwargs):
"""Alias for start() method"""
return self.start(prompt, **kwargs)
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))
@@ -258,11 +258,24 @@ Expanded prompt:"""
prompt_lower = prompt.lower().strip()
return any(prompt_lower.startswith(verb) for verb in task_verbs)
def _is_technical_task(self, prompt: str) -> bool:
"""Check if prompt is a technical/coding task."""
technical_indicators = [
"python", "javascript", "java", "code", "function", "class",
"api", "database", "sql", "html", "css", "react", "node",
"algorithm", "debug", "fix", "error", "bug", "test",
"deploy", "server", "docker", "kubernetes", "aws", "azure",
"typescript", "rust", "go", "c++", "c#", "ruby", "php",
"script", "bash", "shell", "terminal", "command", "cli"
]
prompt_lower = prompt.lower()
return any(ind in prompt_lower for ind in technical_indicators)
def _is_creative_task(self, prompt: str) -> bool:
"""Check if prompt is a creative task."""
"""Check if prompt is a creative task (non-technical)."""
creative_indicators = [
"poem", "story", "script", "song", "art", "creative",
"fiction", "novel", "essay", "blog", "article", "write"
"poem", "story", "song", "art", "creative",
"fiction", "novel", "essay", "blog", "article"
]
prompt_lower = prompt.lower()
return any(ind in prompt_lower for ind in creative_indicators)
@@ -273,6 +286,10 @@ Expanded prompt:"""
if self._is_short_prompt(prompt, threshold=5):
return ExpandStrategy.DETAILED
# Technical/coding tasks get structured expansion (check before creative)
if self._is_technical_task(prompt):
return ExpandStrategy.STRUCTURED
# Creative tasks benefit from creative expansion
if self._is_creative_task(prompt):
return ExpandStrategy.CREATIVE
@@ -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
@@ -0,0 +1,30 @@
"""
Context module for PraisonAI Agents.
This module provides context retrieval capabilities including:
- Fast Context: Rapid parallel code search subagent
"""
# Lazy imports to avoid impacting package load time
def __getattr__(name):
"""Lazy import to avoid impacting package performance."""
if name == "FastContext":
from praisonaiagents.context.fast.fast_context import FastContext
return FastContext
elif name == "FastContextResult":
from praisonaiagents.context.fast.result import FastContextResult
return FastContextResult
elif name == "FileMatch":
from praisonaiagents.context.fast.result import FileMatch
return FileMatch
elif name == "LineRange":
from praisonaiagents.context.fast.result import LineRange
return LineRange
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
"FastContext",
"FastContextResult",
"FileMatch",
"LineRange",
]
@@ -0,0 +1,40 @@
"""
Fast Context module for rapid parallel code search.
This module implements a specialized subagent that retrieves relevant code
from a codebase using parallel tool calls, similar to Windsurf's Fast Context.
Key features:
- Parallel tool execution (up to 8 concurrent calls)
- Limited serial turns (max 4) for fast response
- Restricted tool set (grep, glob, read) for safety
- Returns files + line ranges (no summarization)
"""
# Lazy imports to avoid impacting package load time
def __getattr__(name):
"""Lazy import to avoid impacting package performance."""
if name == "FastContext":
from praisonaiagents.context.fast.fast_context import FastContext
return FastContext
elif name == "FastContextResult":
from praisonaiagents.context.fast.result import FastContextResult
return FastContextResult
elif name == "FileMatch":
from praisonaiagents.context.fast.result import FileMatch
return FileMatch
elif name == "LineRange":
from praisonaiagents.context.fast.result import LineRange
return LineRange
elif name == "FastContextAgent":
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
return FastContextAgent
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
"FastContext",
"FastContextResult",
"FileMatch",
"LineRange",
"FastContextAgent",
]
@@ -0,0 +1,383 @@
"""
Context Injector for Fast Context.
Injects FastContext results into the main agent with:
- Token budget management
- Context formatting for different models
- Precision over recall (avoid context pollution)
"""
import logging
from typing import Optional, List, Dict, Any
from dataclasses import dataclass
from praisonaiagents.context.fast.result import FastContextResult, FileMatch
logger = logging.getLogger(__name__)
@dataclass
class InjectionConfig:
"""Configuration for context injection.
Attributes:
max_tokens: Maximum tokens for injected context
max_files: Maximum files to include
max_lines_per_file: Maximum lines per file
include_line_numbers: Whether to include line numbers
include_file_content: Whether to include file content
format_style: Format style ('markdown', 'xml', 'plain')
prioritize_precision: If True, include fewer but more relevant results
"""
max_tokens: int = 4000
max_files: int = 10
max_lines_per_file: int = 100
include_line_numbers: bool = True
include_file_content: bool = True
format_style: str = "markdown"
prioritize_precision: bool = True
# Approximate tokens per character (conservative estimate)
chars_per_token: float = 4.0
class ContextInjector:
"""Injects FastContext results into agent context.
Handles:
- Token budget management
- Result formatting
- Context pollution prevention
"""
def __init__(self, config: Optional[InjectionConfig] = None):
"""Initialize context injector.
Args:
config: Injection configuration
"""
self.config = config or InjectionConfig()
def estimate_tokens(self, text: str) -> int:
"""Estimate token count for text.
Args:
text: Text to estimate
Returns:
Estimated token count
"""
return int(len(text) / self.config.chars_per_token)
def inject(
self,
result: FastContextResult,
system_prompt: Optional[str] = None,
user_message: Optional[str] = None
) -> Dict[str, str]:
"""Inject FastContext results into agent context.
Args:
result: FastContext search result
system_prompt: Original system prompt
user_message: Original user message
Returns:
Dict with 'system_prompt' and 'user_message' keys
"""
# Format the context
context_str = self.format_context(result)
# Inject into appropriate location
if system_prompt:
# Add to system prompt
injected_system = self._inject_into_system_prompt(system_prompt, context_str)
return {
"system_prompt": injected_system,
"user_message": user_message or ""
}
elif user_message:
# Add to user message
injected_user = self._inject_into_user_message(user_message, context_str)
return {
"system_prompt": "",
"user_message": injected_user
}
else:
return {
"system_prompt": context_str,
"user_message": ""
}
def format_context(self, result: FastContextResult) -> str:
"""Format FastContext result as context string.
Args:
result: FastContext search result
Returns:
Formatted context string
"""
if self.config.format_style == "markdown":
return self._format_markdown(result)
elif self.config.format_style == "xml":
return self._format_xml(result)
else:
return self._format_plain(result)
def _format_markdown(self, result: FastContextResult) -> str:
"""Format result as markdown.
Args:
result: FastContext search result
Returns:
Markdown formatted string
"""
if not result.files:
return ""
lines = ["# Relevant Code Context"]
lines.append(f"*Found {result.total_files} file(s) in {result.search_time_ms}ms*\n")
token_budget = self.config.max_tokens
files_included = 0
# Sort by relevance
sorted_files = sorted(
result.files,
key=lambda f: f.relevance_score,
reverse=True
)
# Apply precision filter if enabled
if self.config.prioritize_precision:
sorted_files = [f for f in sorted_files if f.relevance_score >= 0.5]
for file_match in sorted_files:
if files_included >= self.config.max_files:
break
file_section = self._format_file_markdown(file_match)
section_tokens = self.estimate_tokens(file_section)
if section_tokens > token_budget:
# Try truncated version
file_section = self._format_file_markdown(file_match, truncate=True)
section_tokens = self.estimate_tokens(file_section)
if section_tokens > token_budget:
continue
lines.append(file_section)
token_budget -= section_tokens
files_included += 1
if files_included < result.total_files:
lines.append(f"\n*... and {result.total_files - files_included} more files*")
return "\n".join(lines)
def _format_file_markdown(
self,
file_match: FileMatch,
truncate: bool = False
) -> str:
"""Format a single file match as markdown.
Args:
file_match: File match to format
truncate: Whether to truncate content
Returns:
Markdown formatted string
"""
lines = [f"## `{file_match.path}`"]
if not self.config.include_file_content:
if file_match.line_ranges:
ranges = ", ".join(
f"L{lr.start}-{lr.end}" for lr in file_match.line_ranges[:3]
)
lines.append(f"Relevant lines: {ranges}")
return "\n".join(lines)
max_lines = self.config.max_lines_per_file
if truncate:
max_lines = max_lines // 2
lines_shown = 0
for lr in file_match.line_ranges:
if lines_shown >= max_lines:
break
if self.config.include_line_numbers:
lines.append(f"Lines {lr.start}-{lr.end}:")
if lr.content:
content = lr.content
content_lines = content.split("\n")
if len(content_lines) > max_lines - lines_shown:
content_lines = content_lines[:max_lines - lines_shown]
content = "\n".join(content_lines) + "\n..."
lines.append(f"```\n{content}\n```")
lines_shown += len(content_lines)
return "\n".join(lines)
def _format_xml(self, result: FastContextResult) -> str:
"""Format result as XML.
Args:
result: FastContext search result
Returns:
XML formatted string
"""
if not result.files:
return ""
lines = ["<code_context>"]
lines.append(f" <summary files=\"{result.total_files}\" search_time_ms=\"{result.search_time_ms}\"/>")
token_budget = self.config.max_tokens
files_included = 0
sorted_files = sorted(
result.files,
key=lambda f: f.relevance_score,
reverse=True
)
for file_match in sorted_files[:self.config.max_files]:
if files_included >= self.config.max_files:
break
file_xml = self._format_file_xml(file_match)
section_tokens = self.estimate_tokens(file_xml)
if section_tokens > token_budget:
continue
lines.append(file_xml)
token_budget -= section_tokens
files_included += 1
lines.append("</code_context>")
return "\n".join(lines)
def _format_file_xml(self, file_match: FileMatch) -> str:
"""Format a single file match as XML.
Args:
file_match: File match to format
Returns:
XML formatted string
"""
lines = [f' <file path="{file_match.path}" relevance="{file_match.relevance_score:.2f}">']
for lr in file_match.line_ranges[:3]:
lines.append(f' <lines start="{lr.start}" end="{lr.end}">')
if lr.content and self.config.include_file_content:
# Escape XML special characters
content = lr.content.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
lines.append(f" <![CDATA[{content}]]>")
lines.append(" </lines>")
lines.append(" </file>")
return "\n".join(lines)
def _format_plain(self, result: FastContextResult) -> str:
"""Format result as plain text.
Args:
result: FastContext search result
Returns:
Plain text formatted string
"""
if not result.files:
return ""
lines = [f"Code Context ({result.total_files} files):"]
for file_match in result.files[:self.config.max_files]:
lines.append(f"\nFile: {file_match.path}")
for lr in file_match.line_ranges[:3]:
lines.append(f" Lines {lr.start}-{lr.end}")
if lr.content and self.config.include_file_content:
for line in lr.content.split("\n")[:10]:
lines.append(f" {line}")
return "\n".join(lines)
def _inject_into_system_prompt(
self,
system_prompt: str,
context: str
) -> str:
"""Inject context into system prompt.
Args:
system_prompt: Original system prompt
context: Context to inject
Returns:
Modified system prompt
"""
if not context:
return system_prompt
# Add context section
return f"{system_prompt}\n\n{context}"
def _inject_into_user_message(
self,
user_message: str,
context: str
) -> str:
"""Inject context into user message.
Args:
user_message: Original user message
context: Context to inject
Returns:
Modified user message
"""
if not context:
return user_message
# Prepend context
return f"{context}\n\n---\n\n{user_message}"
def inject_fast_context(
result: FastContextResult,
system_prompt: Optional[str] = None,
user_message: Optional[str] = None,
max_tokens: int = 4000,
format_style: str = "markdown"
) -> Dict[str, str]:
"""Convenience function to inject FastContext results.
Args:
result: FastContext search result
system_prompt: Original system prompt
user_message: Original user message
max_tokens: Maximum tokens for context
format_style: Format style ('markdown', 'xml', 'plain')
Returns:
Dict with 'system_prompt' and 'user_message' keys
"""
config = InjectionConfig(
max_tokens=max_tokens,
format_style=format_style
)
injector = ContextInjector(config)
return injector.inject(result, system_prompt, user_message)
@@ -0,0 +1,384 @@
"""
FastContext - High-level API for fast parallel code search.
This module provides the main FastContext class that can be used
standalone or integrated with the main Agent class.
Example usage:
# Standalone usage
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path="/path/to/project")
result = fc.search("find authentication handlers")
print(result.to_context_string())
# With Agent integration
from praisonaiagents import Agent
agent = Agent(
name="CodeAssistant",
fast_context=True,
fast_context_path="/path/to/project"
)
"""
import os
import hashlib
import json
import logging
from typing import Optional, List, Dict, Any
from praisonaiagents.context.fast.result import FastContextResult, FileMatch, LineRange
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
logger = logging.getLogger(__name__)
class FastContext:
"""High-level API for fast parallel code search.
FastContext provides rapid code search capabilities using:
- Parallel tool execution (up to 8 concurrent calls)
- Limited turns (max 4) for fast response
- Optional LLM-powered intelligent search
- Result caching for repeated queries
Attributes:
workspace_path: Root directory for searches
model: LLM model for intelligent search
max_turns: Maximum search turns
max_parallel: Maximum parallel tool calls
cache_enabled: Whether to cache results
"""
# Environment variable names for configuration
ENV_MODEL = "FAST_CONTEXT_MODEL"
ENV_MAX_TURNS = "FAST_CONTEXT_MAX_TURNS"
ENV_MAX_PARALLEL = "FAST_CONTEXT_PARALLELISM"
ENV_TIMEOUT = "FAST_CONTEXT_TIMEOUT"
ENV_CACHE_ENABLED = "FAST_CONTEXT_CACHE"
ENV_CACHE_TTL = "FAST_CONTEXT_CACHE_TTL"
def __init__(
self,
workspace_path: Optional[str] = None,
model: Optional[str] = None,
max_turns: Optional[int] = None,
max_parallel: Optional[int] = None,
timeout: Optional[float] = None,
cache_enabled: Optional[bool] = None,
cache_ttl: Optional[int] = None,
verbose: bool = False
):
"""Initialize FastContext.
Configuration can be set via parameters or environment variables.
Parameters take precedence over environment variables.
Environment variables:
FAST_CONTEXT_MODEL: LLM model name (default: gpt-4o-mini)
FAST_CONTEXT_MAX_TURNS: Maximum search turns (default: 4)
FAST_CONTEXT_PARALLELISM: Max parallel calls (default: 8)
FAST_CONTEXT_TIMEOUT: Timeout in seconds (default: 30.0)
FAST_CONTEXT_CACHE: Enable caching (default: true)
FAST_CONTEXT_CACHE_TTL: Cache TTL in seconds (default: 300)
Args:
workspace_path: Root directory for searches (defaults to cwd)
model: LLM model for intelligent search
max_turns: Maximum search turns
max_parallel: Maximum parallel tool calls per turn
timeout: Timeout per tool call in seconds
cache_enabled: Whether to cache search results
cache_ttl: Cache time-to-live in seconds
verbose: If True, print debug information
"""
self.workspace_path = os.path.abspath(workspace_path or os.getcwd())
# Apply environment variable overrides with defaults
self.model = model or os.environ.get(self.ENV_MODEL, "gpt-4o-mini")
self.max_turns = max_turns if max_turns is not None else int(os.environ.get(self.ENV_MAX_TURNS, "4"))
self.max_parallel = max_parallel if max_parallel is not None else int(os.environ.get(self.ENV_MAX_PARALLEL, "8"))
self.timeout = timeout if timeout is not None else float(os.environ.get(self.ENV_TIMEOUT, "30.0"))
# Cache settings
cache_env = os.environ.get(self.ENV_CACHE_ENABLED, "true").lower()
self.cache_enabled = cache_enabled if cache_enabled is not None else cache_env in ("true", "1", "yes")
self.cache_ttl = cache_ttl if cache_ttl is not None else int(os.environ.get(self.ENV_CACHE_TTL, "300"))
self.verbose = verbose
# Cache storage
self._cache: Dict[str, Dict[str, Any]] = {}
# Agent instance (lazy loaded)
self._agent: Optional[FastContextAgent] = None
def _get_agent(self) -> FastContextAgent:
"""Get or create FastContextAgent instance."""
if self._agent is None:
self._agent = FastContextAgent(
workspace_path=self.workspace_path,
max_turns=self.max_turns,
max_parallel=self.max_parallel,
model=self.model,
timeout=self.timeout,
verbose=self.verbose
)
return self._agent
def _get_cache_key(self, query: str, **kwargs) -> str:
"""Generate cache key for a query."""
key_data = {
"query": query,
"workspace": self.workspace_path,
**kwargs
}
key_str = json.dumps(key_data, sort_keys=True)
return hashlib.md5(key_str.encode()).hexdigest()
def _get_cached(self, cache_key: str) -> Optional[FastContextResult]:
"""Get cached result if valid."""
if not self.cache_enabled:
return None
import time
cached = self._cache.get(cache_key)
if cached:
if time.time() - cached["timestamp"] < self.cache_ttl:
result = cached["result"]
result.from_cache = True
return result
else:
# Expired
del self._cache[cache_key]
return None
def _set_cached(self, cache_key: str, result: FastContextResult) -> None:
"""Cache a result."""
if not self.cache_enabled:
return
import time
self._cache[cache_key] = {
"result": result,
"timestamp": time.time()
}
def search(
self,
query: str,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None,
use_llm: bool = False
) -> FastContextResult:
"""Search for code relevant to the query.
Args:
query: Search query (pattern or natural language)
include_patterns: Glob patterns to include (e.g., ["*.py"])
exclude_patterns: Glob patterns to exclude
use_llm: If True, use LLM for intelligent search
Returns:
FastContextResult with matching files and lines
"""
# Check cache
cache_key = self._get_cache_key(
query,
include_patterns=include_patterns,
exclude_patterns=exclude_patterns,
use_llm=use_llm
)
cached = self._get_cached(cache_key)
if cached:
if self.verbose:
logger.info(f"Cache hit for query: {query}")
return cached
# Perform search
agent = self._get_agent()
if use_llm:
result = agent.search(query, include_patterns, exclude_patterns)
else:
result = agent.search_simple(query, include_patterns, exclude_patterns)
# Cache result
self._set_cached(cache_key, result)
return result
def search_files(
self,
pattern: str,
max_results: int = 50
) -> FastContextResult:
"""Search for files matching a glob pattern.
Args:
pattern: Glob pattern (e.g., "**/*.py", "src/*.js")
max_results: Maximum results to return
Returns:
FastContextResult with matching files
"""
from praisonaiagents.context.fast.search_tools import glob_search
result = FastContextResult(query=pattern)
result.start_timer()
matches = glob_search(
self.workspace_path,
pattern,
max_results=max_results
)
for match in matches:
result.add_file(FileMatch(
path=match["path"],
relevance_score=1.0
))
result.stop_timer()
result.turns_used = 1
result.total_tool_calls = 1
return result
def read_context(
self,
filepath: str,
start_line: Optional[int] = None,
end_line: Optional[int] = None,
context_lines: int = 5
) -> Optional[str]:
"""Read file content with context.
Args:
filepath: Path to file (relative to workspace)
start_line: Starting line (1-indexed)
end_line: Ending line (1-indexed)
context_lines: Additional context lines
Returns:
File content string or None if error
"""
from praisonaiagents.context.fast.search_tools import read_file
full_path = filepath
if not os.path.isabs(filepath):
full_path = os.path.join(self.workspace_path, filepath)
result = read_file(
full_path,
start_line=start_line,
end_line=end_line,
context_lines=context_lines
)
if result["success"]:
return result["content"]
return None
def get_context_for_agent(
self,
query: str,
max_files: int = 10,
max_lines_per_file: int = 100,
include_content: bool = True
) -> str:
"""Get formatted context string for an agent.
This method searches for relevant code and formats it
as a context string that can be injected into an agent's
system prompt or message.
Args:
query: Search query
max_files: Maximum files to include
max_lines_per_file: Maximum lines per file
include_content: Whether to include file content
Returns:
Formatted context string
"""
result = self.search(query)
if not result.files:
return f"No relevant code found for: {query}"
lines = [f"# Relevant Code Context for: {query}"]
lines.append(f"Found {result.total_files} relevant file(s) in {result.search_time_ms}ms\n")
for i, file_match in enumerate(result.files[:max_files]):
lines.append(f"## {file_match.path}")
if include_content and file_match.line_ranges:
for lr in file_match.line_ranges:
# Limit lines
if lr.line_count > max_lines_per_file:
lr = LineRange(
start=lr.start,
end=lr.start + max_lines_per_file - 1,
content=lr.content,
relevance_score=lr.relevance_score
)
lines.append(f"Lines {lr.start}-{lr.end}:")
# Read content if not already present
if lr.content:
lines.append(f"```\n{lr.content}\n```")
else:
content = self.read_context(
file_match.path,
start_line=lr.start,
end_line=lr.end
)
if content:
lines.append(f"```\n{content}\n```")
lines.append("")
if len(result.files) > max_files:
lines.append(f"... and {len(result.files) - max_files} more files")
return "\n".join(lines)
def clear_cache(self) -> None:
"""Clear the result cache."""
self._cache.clear()
def close(self) -> None:
"""Close resources."""
if self._agent:
self._agent.close()
self._agent = None
self._cache.clear()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
# Convenience function for quick searches
def fast_search(
query: str,
workspace_path: Optional[str] = None,
**kwargs
) -> FastContextResult:
"""Convenience function for quick code search.
Args:
query: Search query
workspace_path: Root directory (defaults to cwd)
**kwargs: Additional arguments for FastContext
Returns:
FastContextResult with matching files
"""
with FastContext(workspace_path=workspace_path, **kwargs) as fc:
return fc.search(query)
@@ -0,0 +1,392 @@
"""
FastContextAgent - Specialized subagent for rapid code search.
This agent is designed for fast context retrieval using:
- Parallel tool execution (up to 8 concurrent calls)
- Limited turns (max 4) for fast response
- Lightweight model (gpt-4o-mini by default)
- Restricted tool set (grep, glob, read, list_dir)
"""
import os
import json
import logging
from typing import List, Dict, Any, Optional
from praisonaiagents.context.fast.result import (
FastContextResult,
FileMatch,
LineRange
)
from praisonaiagents.context.fast.search_tools import (
grep_search,
glob_search,
read_file,
list_directory,
FAST_CONTEXT_TOOLS
)
from praisonaiagents.context.fast.parallel_executor import (
ToolCallBatch,
ParallelSearchCoordinator
)
logger = logging.getLogger(__name__)
class FastContextAgent:
"""Specialized agent for fast parallel code search.
This agent performs rapid code search using parallel tool calls
and a lightweight model optimized for retrieval tasks.
Attributes:
workspace_path: Root directory for searches
max_turns: Maximum search turns (default: 4)
max_parallel: Maximum parallel tool calls (default: 8)
model: LLM model for intelligent search (default: gpt-4o-mini)
timeout: Timeout per tool call in seconds
"""
def __init__(
self,
workspace_path: str,
max_turns: int = 4,
max_parallel: int = 8,
model: str = "gpt-4o-mini",
timeout: float = 30.0,
verbose: bool = False
):
"""Initialize FastContextAgent.
Args:
workspace_path: Root directory for searches
max_turns: Maximum search turns
max_parallel: Maximum parallel tool calls per turn
model: LLM model name for intelligent search
timeout: Timeout per tool call in seconds
verbose: If True, print debug information
"""
self.workspace_path = os.path.abspath(workspace_path)
self.max_turns = max_turns
self.max_parallel = max_parallel
self.model = model
self.timeout = timeout
self.verbose = verbose
# Initialize coordinator
self.coordinator = ParallelSearchCoordinator(
max_parallel=max_parallel,
max_turns=max_turns,
timeout=timeout
)
# LLM instance (lazy loaded)
self._llm = None
def _get_llm(self):
"""Get or create LLM instance (lazy loading)."""
if self._llm is None:
try:
from praisonaiagents.llm import LLM
self._llm = LLM(model=self.model)
except Exception as e:
logger.warning(f"Could not initialize LLM: {e}")
self._llm = None
return self._llm
def get_tools(self) -> List[Dict[str, Any]]:
"""Get tool definitions for LLM function calling.
Returns:
List of tool definition dictionaries
"""
return FAST_CONTEXT_TOOLS
def execute_tool(self, tool_name: str, **kwargs) -> Any:
"""Execute a single tool.
Args:
tool_name: Name of the tool
**kwargs: Tool arguments
Returns:
Tool result
"""
# Prepend workspace path if needed
if tool_name in ("grep_search", "glob_search"):
if "search_path" not in kwargs or kwargs["search_path"] == ".":
kwargs["search_path"] = self.workspace_path
elif not os.path.isabs(kwargs["search_path"]):
kwargs["search_path"] = os.path.join(
self.workspace_path, kwargs["search_path"]
)
elif tool_name == "list_directory":
if "dir_path" not in kwargs or kwargs["dir_path"] == ".":
kwargs["dir_path"] = self.workspace_path
elif not os.path.isabs(kwargs["dir_path"]):
kwargs["dir_path"] = os.path.join(
self.workspace_path, kwargs["dir_path"]
)
elif tool_name == "read_file":
if "filepath" in kwargs and not os.path.isabs(kwargs["filepath"]):
kwargs["filepath"] = os.path.join(
self.workspace_path, kwargs["filepath"]
)
tools = {
"grep_search": grep_search,
"glob_search": glob_search,
"read_file": read_file,
"list_directory": list_directory
}
if tool_name not in tools:
return {"error": f"Unknown tool: {tool_name}"}
try:
return tools[tool_name](**kwargs)
except Exception as e:
return {"error": str(e)}
def search_simple(
self,
query: str,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None,
max_results: int = 50
) -> FastContextResult:
"""Perform a simple search without LLM.
This is a fast, direct search that doesn't use the LLM
for query understanding. Good for simple pattern searches.
Args:
query: Search pattern
include_patterns: Glob patterns to include
exclude_patterns: Glob patterns to exclude
max_results: Maximum results per tool
Returns:
FastContextResult with matching files and lines
"""
result = FastContextResult(query=query)
result.start_timer()
# Execute grep search
grep_results = grep_search(
search_path=self.workspace_path,
pattern=query,
include_patterns=include_patterns,
exclude_patterns=exclude_patterns,
max_results=max_results,
context_lines=2
)
# Convert to FastContextResult format
for match in grep_results:
file_match = FileMatch(
path=match["path"],
relevance_score=1.0,
match_count=1
)
file_match.add_line_range(LineRange(
start=max(1, match["line_number"] - 2),
end=match["line_number"] + 2,
content=match.get("context"),
relevance_score=1.0
))
result.add_file(file_match)
result.stop_timer()
result.turns_used = 1
result.total_tool_calls = 1
return result
def search(
self,
query: str,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None
) -> FastContextResult:
"""Perform intelligent search using LLM.
Uses the LLM to understand the query and plan search strategy,
then executes parallel tool calls to find relevant code.
Args:
query: Natural language search query
include_patterns: Glob patterns to include
exclude_patterns: Glob patterns to exclude
Returns:
FastContextResult with matching files and lines
"""
result = FastContextResult(query=query)
result.start_timer()
llm = self._get_llm()
# If no LLM available, fall back to simple search
if llm is None:
logger.info("LLM not available, falling back to simple search")
return self.search_simple(query, include_patterns, exclude_patterns)
self.coordinator.start()
# System prompt for the search agent
system_prompt = f"""You are a fast code search agent. Your task is to find relevant code in a codebase.
Workspace: {self.workspace_path}
You have access to these tools:
- grep_search: Search for patterns in files
- glob_search: Find files by pattern
- read_file: Read file contents
- list_directory: List directory contents
Strategy:
1. Start with broad searches (glob, list_directory) to understand structure
2. Use grep_search to find specific patterns
3. Use read_file to get detailed context for relevant files
4. Execute up to {self.max_parallel} tool calls in parallel per turn
5. Complete search in {self.max_turns} turns or less
Return files and line ranges that are relevant to the query.
Be efficient - use parallel tool calls to explore multiple paths at once."""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": f"Find code relevant to: {query}"}
]
# Add file pattern constraints if provided
if include_patterns:
messages[1]["content"] += f"\nOnly search in files matching: {include_patterns}"
if exclude_patterns:
messages[1]["content"] += f"\nExclude files matching: {exclude_patterns}"
try:
# Multi-turn search loop
for turn in range(self.max_turns):
if self.verbose:
logger.info(f"Search turn {turn + 1}/{self.max_turns}")
# Get LLM response with tool calls
response = llm.client.chat.completions.create(
model=self.model,
messages=messages,
tools=[{"type": "function", "function": t} for t in self.get_tools()],
tool_choice="auto"
)
message = response.choices[0].message
# Check if we have tool calls
if not message.tool_calls:
# No more tool calls, search complete
break
# Execute tool calls in parallel
batch = ToolCallBatch(max_size=self.max_parallel)
tool_call_map = {} # Map tool call ID to batch index
for tool_call in message.tool_calls[:self.max_parallel]:
try:
args = json.loads(tool_call.function.arguments)
batch.add(tool_call.function.name, **args)
tool_call_map[tool_call.id] = len(batch.tasks) - 1
except json.JSONDecodeError:
continue
# Execute batch
if batch.tasks:
tool_results = self.coordinator.execute_turn_sync(batch)
# Process results
messages.append({"role": "assistant", "content": None, "tool_calls": [
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
}
for tc in message.tool_calls[:self.max_parallel]
]})
# Add tool results to messages
for tool_call in message.tool_calls[:self.max_parallel]:
idx = tool_call_map.get(tool_call.id)
if idx is not None and idx < len(tool_results):
tool_result = tool_results[idx]
# Process grep results into FastContextResult
if tool_call.function.name == "grep_search":
if isinstance(tool_result, list):
for match in tool_result:
file_match = FileMatch(
path=match.get("path", ""),
relevance_score=1.0,
match_count=1
)
line_num = match.get("line_number", 1)
file_match.add_line_range(LineRange(
start=max(1, line_num - 2),
end=line_num + 2,
content=match.get("context"),
relevance_score=1.0
))
result.add_file(file_match)
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": json.dumps(tool_result) if not isinstance(tool_result, str) else tool_result
})
except Exception as e:
logger.error(f"Error during LLM search: {e}")
# Fall back to simple search
return self.search_simple(query, include_patterns, exclude_patterns)
finally:
self.coordinator.stop()
result.stop_timer()
result.turns_used = self.coordinator.turns_used
result.total_tool_calls = self.coordinator.total_tool_calls
result.sort_by_relevance()
return result
async def search_async(
self,
query: str,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None
) -> FastContextResult:
"""Async version of search.
Args:
query: Natural language search query
include_patterns: Glob patterns to include
exclude_patterns: Glob patterns to exclude
Returns:
FastContextResult with matching files and lines
"""
# For now, wrap sync version
# TODO: Implement true async with async LLM calls
return self.search(query, include_patterns, exclude_patterns)
def close(self) -> None:
"""Close resources."""
self.coordinator.close()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
@@ -0,0 +1,23 @@
"""
Indexer module for Fast Context.
Provides file and symbol indexing capabilities for fast code search.
"""
def __getattr__(name):
"""Lazy load indexer classes."""
if name == "FileIndexer":
from .file_indexer import FileIndexer
return FileIndexer
elif name == "SymbolIndexer":
from .symbol_indexer import SymbolIndexer
return SymbolIndexer
elif name == "Symbol":
from .symbol_indexer import Symbol
return Symbol
elif name == "SymbolType":
from .symbol_indexer import SymbolType
return SymbolType
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = ["FileIndexer", "SymbolIndexer", "Symbol", "SymbolType"]
@@ -0,0 +1,445 @@
"""
File Indexer for Fast Context.
Provides file indexing capabilities for faster code search:
- Indexes file paths and metadata
- Supports .gitignore and .praisonignore patterns
- Incremental updates on file changes
- Memory-efficient storage
"""
import os
import fnmatch
import hashlib
import json
import logging
from dataclasses import dataclass
from typing import Dict, List, Set, Optional, Any
from pathlib import Path
import time
logger = logging.getLogger(__name__)
@dataclass
class FileInfo:
"""Information about an indexed file.
Attributes:
path: Relative path from workspace root
absolute_path: Absolute file path
size: File size in bytes
modified: Last modification timestamp
extension: File extension (lowercase)
hash: Content hash for change detection
"""
path: str
absolute_path: str
size: int
modified: float
extension: str
hash: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"path": self.path,
"absolute_path": self.absolute_path,
"size": self.size,
"modified": self.modified,
"extension": self.extension,
"hash": self.hash
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "FileInfo":
"""Create from dictionary."""
return cls(**data)
class FileIndexer:
"""Indexes files in a workspace for fast search.
The indexer maintains an in-memory index of files that can be
used to quickly find files by pattern without filesystem traversal.
Attributes:
workspace_path: Root directory to index
files: Dictionary mapping relative paths to FileInfo
extensions: Set of indexed file extensions
"""
# Default extensions to index
DEFAULT_EXTENSIONS = {
".py", ".js", ".ts", ".jsx", ".tsx", ".java", ".go", ".rs",
".c", ".cpp", ".h", ".hpp", ".cs", ".rb", ".php", ".swift",
".kt", ".scala", ".r", ".m", ".mm", ".sql", ".sh", ".bash",
".yaml", ".yml", ".json", ".xml", ".html", ".css", ".scss",
".md", ".txt", ".rst", ".toml", ".ini", ".cfg", ".conf"
}
# Default patterns to ignore
DEFAULT_IGNORE_PATTERNS = {
"__pycache__", "*.pyc", "*.pyo", ".git", ".svn", ".hg",
"node_modules", "venv", ".venv", "env", ".env",
"dist", "build", ".cache", ".pytest_cache",
"*.egg-info", "*.egg", "*.whl", "*.so", "*.dll", "*.dylib"
}
def __init__(
self,
workspace_path: str,
extensions: Optional[Set[str]] = None,
ignore_patterns: Optional[Set[str]] = None,
respect_gitignore: bool = True,
max_file_size: int = 10 * 1024 * 1024, # 10MB
compute_hashes: bool = False
):
"""Initialize file indexer.
Args:
workspace_path: Root directory to index
extensions: File extensions to index (None = all text files)
ignore_patterns: Patterns to ignore
respect_gitignore: Whether to respect .gitignore
max_file_size: Maximum file size to index
compute_hashes: Whether to compute content hashes
"""
self.workspace_path = os.path.abspath(workspace_path)
self.extensions = extensions or self.DEFAULT_EXTENSIONS
self.ignore_patterns = ignore_patterns or self.DEFAULT_IGNORE_PATTERNS.copy()
self.respect_gitignore = respect_gitignore
self.max_file_size = max_file_size
self.compute_hashes = compute_hashes
# Index storage
self.files: Dict[str, FileInfo] = {}
self.by_extension: Dict[str, List[str]] = {}
# Gitignore patterns
self._gitignore_patterns: Set[str] = set()
# Index metadata
self.last_indexed: Optional[float] = None
self.total_files: int = 0
self.total_size: int = 0
def _load_gitignore(self) -> None:
"""Load patterns from .gitignore and .praisonignore."""
self._gitignore_patterns.clear()
for filename in [".gitignore", ".praisonignore"]:
filepath = os.path.join(self.workspace_path, filename)
if os.path.exists(filepath):
try:
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
for line in f:
line = line.strip()
if line and not line.startswith('#'):
self._gitignore_patterns.add(line)
except Exception:
pass
def _should_ignore(self, path: str, is_dir: bool = False) -> bool:
"""Check if a path should be ignored.
Args:
path: Relative path from workspace
is_dir: Whether the path is a directory
Returns:
True if path should be ignored
"""
name = os.path.basename(path)
parts = Path(path).parts
# Check default ignore patterns
for pattern in self.ignore_patterns:
if fnmatch.fnmatch(name, pattern):
return True
if any(fnmatch.fnmatch(part, pattern) for part in parts):
return True
# Check gitignore patterns
if self.respect_gitignore:
for pattern in self._gitignore_patterns:
# Handle directory patterns
if pattern.endswith('/'):
dir_pattern = pattern.rstrip('/')
if is_dir and fnmatch.fnmatch(name, dir_pattern):
return True
if any(fnmatch.fnmatch(part, dir_pattern) for part in parts):
return True
# Handle ** patterns
elif '**' in pattern:
fnmatch_pattern = pattern.replace('**/', '*').replace('**', '*')
if fnmatch.fnmatch(path, fnmatch_pattern):
return True
# Handle simple patterns
else:
if fnmatch.fnmatch(name, pattern):
return True
if fnmatch.fnmatch(path, pattern):
return True
return False
def _compute_hash(self, filepath: str) -> Optional[str]:
"""Compute MD5 hash of file content.
Args:
filepath: Path to file
Returns:
MD5 hash string or None on error
"""
if not self.compute_hashes:
return None
try:
hasher = hashlib.md5()
with open(filepath, 'rb') as f:
for chunk in iter(lambda: f.read(8192), b''):
hasher.update(chunk)
return hasher.hexdigest()
except Exception:
return None
def index(self) -> int:
"""Index all files in the workspace.
Returns:
Number of files indexed
"""
start_time = time.perf_counter()
# Load gitignore patterns
if self.respect_gitignore:
self._load_gitignore()
# Clear existing index
self.files.clear()
self.by_extension.clear()
self.total_files = 0
self.total_size = 0
# Walk directory tree
for root, dirs, files in os.walk(self.workspace_path):
# Get relative path
rel_root = os.path.relpath(root, self.workspace_path)
if rel_root == '.':
rel_root = ''
# Filter directories in-place
dirs[:] = [d for d in dirs if not self._should_ignore(
os.path.join(rel_root, d) if rel_root else d, is_dir=True
)]
for filename in files:
rel_path = os.path.join(rel_root, filename) if rel_root else filename
abs_path = os.path.join(root, filename)
# Check if should ignore
if self._should_ignore(rel_path):
continue
# Check extension
ext = os.path.splitext(filename)[1].lower()
if self.extensions and ext not in self.extensions:
continue
try:
stat = os.stat(abs_path)
# Check file size
if stat.st_size > self.max_file_size:
continue
# Create file info
file_info = FileInfo(
path=rel_path,
absolute_path=abs_path,
size=stat.st_size,
modified=stat.st_mtime,
extension=ext,
hash=self._compute_hash(abs_path)
)
# Add to index
self.files[rel_path] = file_info
# Add to extension index
if ext not in self.by_extension:
self.by_extension[ext] = []
self.by_extension[ext].append(rel_path)
self.total_files += 1
self.total_size += stat.st_size
except Exception as e:
logger.debug(f"Error indexing {abs_path}: {e}")
continue
self.last_indexed = time.perf_counter()
elapsed = self.last_indexed - start_time
logger.info(f"Indexed {self.total_files} files ({self.total_size / 1024 / 1024:.2f} MB) in {elapsed:.2f}s")
return self.total_files
def find_by_pattern(self, pattern: str) -> List[FileInfo]:
"""Find files matching a glob pattern.
Args:
pattern: Glob pattern (e.g., "**/*.py", "src/*.js")
Returns:
List of matching FileInfo objects
"""
results = []
for rel_path, file_info in self.files.items():
if fnmatch.fnmatch(rel_path, pattern):
results.append(file_info)
elif fnmatch.fnmatch(os.path.basename(rel_path), pattern):
results.append(file_info)
return results
def find_by_extension(self, extension: str) -> List[FileInfo]:
"""Find files by extension.
Args:
extension: File extension (with or without dot)
Returns:
List of matching FileInfo objects
"""
if not extension.startswith('.'):
extension = '.' + extension
extension = extension.lower()
paths = self.by_extension.get(extension, [])
return [self.files[p] for p in paths if p in self.files]
def find_by_name(self, name: str, exact: bool = False) -> List[FileInfo]:
"""Find files by name.
Args:
name: Filename to search for
exact: If True, match exact name; otherwise, partial match
Returns:
List of matching FileInfo objects
"""
results = []
name_lower = name.lower()
for rel_path, file_info in self.files.items():
filename = os.path.basename(rel_path).lower()
if exact:
if filename == name_lower:
results.append(file_info)
else:
if name_lower in filename:
results.append(file_info)
return results
def get_file(self, path: str) -> Optional[FileInfo]:
"""Get file info by path.
Args:
path: Relative path from workspace
Returns:
FileInfo or None if not found
"""
return self.files.get(path)
def is_stale(self, max_age: float = 60.0) -> bool:
"""Check if index is stale.
Args:
max_age: Maximum age in seconds
Returns:
True if index is stale or not indexed
"""
if self.last_indexed is None:
return True
return (time.perf_counter() - self.last_indexed) > max_age
def get_stats(self) -> Dict[str, Any]:
"""Get index statistics.
Returns:
Dictionary with index stats
"""
return {
"workspace_path": self.workspace_path,
"total_files": self.total_files,
"total_size_bytes": self.total_size,
"total_size_mb": round(self.total_size / 1024 / 1024, 2),
"extensions": dict(sorted(
((ext, len(paths)) for ext, paths in self.by_extension.items()),
key=lambda x: x[1],
reverse=True
)),
"last_indexed": self.last_indexed
}
def save_index(self, filepath: str) -> None:
"""Save index to file.
Args:
filepath: Path to save index
"""
data = {
"workspace_path": self.workspace_path,
"last_indexed": self.last_indexed,
"files": {path: info.to_dict() for path, info in self.files.items()}
}
with open(filepath, 'w', encoding='utf-8') as f:
json.dump(data, f, indent=2)
def load_index(self, filepath: str) -> bool:
"""Load index from file.
Args:
filepath: Path to load index from
Returns:
True if loaded successfully
"""
try:
with open(filepath, 'r', encoding='utf-8') as f:
data = json.load(f)
# Verify workspace path matches
if data.get("workspace_path") != self.workspace_path:
return False
self.last_indexed = data.get("last_indexed")
self.files.clear()
self.by_extension.clear()
self.total_files = 0
self.total_size = 0
for path, info_dict in data.get("files", {}).items():
file_info = FileInfo.from_dict(info_dict)
self.files[path] = file_info
ext = file_info.extension
if ext not in self.by_extension:
self.by_extension[ext] = []
self.by_extension[ext].append(path)
self.total_files += 1
self.total_size += file_info.size
return True
except Exception as e:
logger.debug(f"Error loading index: {e}")
return False
@@ -0,0 +1,321 @@
"""
File Watcher for Incremental Indexing.
Monitors file system changes and triggers incremental index updates:
- Watches for file creation, modification, deletion
- Debounces rapid changes
- Supports background indexing
"""
import os
import time
import logging
import threading
from typing import Dict, Set, Callable, Optional, Any
from dataclasses import dataclass
from enum import Enum
logger = logging.getLogger(__name__)
class ChangeType(Enum):
"""Types of file system changes."""
CREATED = "created"
MODIFIED = "modified"
DELETED = "deleted"
@dataclass
class FileChange:
"""Represents a file system change.
Attributes:
path: Path to the changed file
change_type: Type of change
timestamp: When the change was detected
"""
path: str
change_type: ChangeType
timestamp: float
class FileWatcher:
"""Watches for file system changes and triggers callbacks.
Uses polling-based approach for cross-platform compatibility.
For production use, consider using watchdog library.
Attributes:
workspace_path: Root directory to watch
extensions: File extensions to watch
poll_interval: Seconds between polls
"""
def __init__(
self,
workspace_path: str,
extensions: Optional[Set[str]] = None,
poll_interval: float = 1.0,
debounce_delay: float = 0.5,
on_change: Optional[Callable[[FileChange], None]] = None
):
"""Initialize file watcher.
Args:
workspace_path: Root directory to watch
extensions: File extensions to watch (None = all)
poll_interval: Seconds between polls
debounce_delay: Delay before processing changes
on_change: Callback for file changes
"""
self.workspace_path = os.path.abspath(workspace_path)
self.extensions = extensions or {
'.py', '.js', '.ts', '.jsx', '.tsx', '.go', '.rs', '.java',
'.c', '.cpp', '.h', '.hpp', '.cs', '.rb', '.php', '.swift'
}
self.poll_interval = poll_interval
self.debounce_delay = debounce_delay
self.on_change = on_change
# State tracking
self._file_mtimes: Dict[str, float] = {}
self._pending_changes: Dict[str, FileChange] = {}
self._running = False
self._thread: Optional[threading.Thread] = None
self._lock = threading.Lock()
def _scan_files(self) -> Dict[str, float]:
"""Scan workspace and get file modification times.
Returns:
Dict mapping file paths to modification times
"""
mtimes = {}
for root, dirs, files in os.walk(self.workspace_path):
# Skip hidden directories
dirs[:] = [d for d in dirs if not d.startswith('.')]
for filename in files:
ext = os.path.splitext(filename)[1].lower()
if ext in self.extensions:
file_path = os.path.join(root, filename)
try:
mtimes[file_path] = os.path.getmtime(file_path)
except OSError:
pass
return mtimes
def _detect_changes(self, new_mtimes: Dict[str, float]) -> None:
"""Detect changes between old and new file states.
Args:
new_mtimes: New file modification times
"""
old_paths = set(self._file_mtimes.keys())
new_paths = set(new_mtimes.keys())
now = time.time()
# Detect created files
for path in new_paths - old_paths:
self._add_pending_change(FileChange(
path=path,
change_type=ChangeType.CREATED,
timestamp=now
))
# Detect deleted files
for path in old_paths - new_paths:
self._add_pending_change(FileChange(
path=path,
change_type=ChangeType.DELETED,
timestamp=now
))
# Detect modified files
for path in old_paths & new_paths:
if new_mtimes[path] > self._file_mtimes[path]:
self._add_pending_change(FileChange(
path=path,
change_type=ChangeType.MODIFIED,
timestamp=now
))
self._file_mtimes = new_mtimes
def _add_pending_change(self, change: FileChange) -> None:
"""Add a change to pending queue (with debouncing).
Args:
change: File change to add
"""
with self._lock:
self._pending_changes[change.path] = change
def _process_pending_changes(self) -> None:
"""Process pending changes after debounce delay."""
with self._lock:
now = time.time()
to_process = []
for path, change in list(self._pending_changes.items()):
if now - change.timestamp >= self.debounce_delay:
to_process.append(change)
del self._pending_changes[path]
# Process changes outside lock
for change in to_process:
if self.on_change:
try:
self.on_change(change)
except Exception as e:
logger.error(f"Error processing change {change.path}: {e}")
def _poll_loop(self) -> None:
"""Main polling loop."""
# Initial scan
self._file_mtimes = self._scan_files()
while self._running:
try:
# Scan for changes
new_mtimes = self._scan_files()
self._detect_changes(new_mtimes)
# Process pending changes
self._process_pending_changes()
except Exception as e:
logger.error(f"Error in file watcher: {e}")
time.sleep(self.poll_interval)
def start(self) -> None:
"""Start watching for file changes."""
if self._running:
return
self._running = True
self._thread = threading.Thread(target=self._poll_loop, daemon=True)
self._thread.start()
logger.info(f"File watcher started for {self.workspace_path}")
def stop(self) -> None:
"""Stop watching for file changes."""
self._running = False
if self._thread:
self._thread.join(timeout=2.0)
self._thread = None
logger.info("File watcher stopped")
def __enter__(self):
self.start()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.stop()
class IncrementalIndexer:
"""Manages incremental index updates based on file changes.
Combines FileIndexer/SymbolIndexer with FileWatcher for
automatic index updates when files change.
"""
def __init__(
self,
workspace_path: str,
file_indexer: Optional[Any] = None,
symbol_indexer: Optional[Any] = None,
poll_interval: float = 2.0
):
"""Initialize incremental indexer.
Args:
workspace_path: Root directory to index
file_indexer: FileIndexer instance (optional)
symbol_indexer: SymbolIndexer instance (optional)
poll_interval: Seconds between file system polls
"""
self.workspace_path = os.path.abspath(workspace_path)
self.file_indexer = file_indexer
self.symbol_indexer = symbol_indexer
self._watcher = FileWatcher(
workspace_path=workspace_path,
poll_interval=poll_interval,
on_change=self._handle_change
)
# Statistics
self.changes_processed = 0
self.last_change_time: Optional[float] = None
def _handle_change(self, change: FileChange) -> None:
"""Handle a file change.
Args:
change: File change to process
"""
rel_path = os.path.relpath(change.path, self.workspace_path)
logger.debug(f"Processing {change.change_type.value}: {rel_path}")
if change.change_type == ChangeType.DELETED:
# Remove from indexes
if self.file_indexer and hasattr(self.file_indexer, 'files'):
self.file_indexer.files.pop(rel_path, None)
if self.symbol_indexer and hasattr(self.symbol_indexer, 'symbols'):
self.symbol_indexer.symbols.pop(rel_path, None)
else:
# Re-index the file
if self.file_indexer and hasattr(self.file_indexer, 'index_file'):
# FileIndexer doesn't have index_file, so we skip
pass
if self.symbol_indexer and hasattr(self.symbol_indexer, 'index_file'):
self.symbol_indexer.index_file(change.path)
self.changes_processed += 1
self.last_change_time = time.time()
def start(self) -> None:
"""Start incremental indexing."""
# Initial full index
if self.file_indexer:
self.file_indexer.index()
if self.symbol_indexer:
self.symbol_indexer.index()
# Start watching
self._watcher.start()
def stop(self) -> None:
"""Stop incremental indexing."""
self._watcher.stop()
def get_stats(self) -> Dict[str, Any]:
"""Get indexer statistics.
Returns:
Dictionary with stats
"""
stats = {
"changes_processed": self.changes_processed,
"last_change_time": self.last_change_time
}
if self.file_indexer:
stats["file_indexer"] = self.file_indexer.get_stats()
if self.symbol_indexer:
stats["symbol_indexer"] = self.symbol_indexer.get_stats()
return stats
def __enter__(self):
self.start()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.stop()
@@ -0,0 +1,669 @@
"""
Symbol Indexer for Fast Context.
Extracts code symbols from source files:
- Python: functions, classes, imports
- JavaScript/TypeScript: functions, classes, imports, exports
- Go: functions, types, imports
- Rust: functions, structs, impl blocks, use statements
- Java: classes, methods, imports
"""
import os
import re
import logging
from dataclasses import dataclass, field
from typing import Dict, List, Set, Optional, Any
from pathlib import Path
from enum import Enum
logger = logging.getLogger(__name__)
class SymbolType(Enum):
"""Types of code symbols."""
FUNCTION = "function"
CLASS = "class"
METHOD = "method"
IMPORT = "import"
EXPORT = "export"
VARIABLE = "variable"
CONSTANT = "constant"
TYPE = "type"
STRUCT = "struct"
INTERFACE = "interface"
MODULE = "module"
@dataclass
class Symbol:
"""Represents a code symbol.
Attributes:
name: Symbol name
symbol_type: Type of symbol
file_path: File containing the symbol
line_number: Line number (1-indexed)
end_line: End line number (for multi-line symbols)
signature: Full signature (for functions/methods)
parent: Parent symbol name (for methods)
docstring: Documentation string if available
"""
name: str
symbol_type: SymbolType
file_path: str
line_number: int
end_line: Optional[int] = None
signature: Optional[str] = None
parent: Optional[str] = None
docstring: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"name": self.name,
"type": self.symbol_type.value,
"file_path": self.file_path,
"line_number": self.line_number,
"end_line": self.end_line,
"signature": self.signature,
"parent": self.parent,
"docstring": self.docstring
}
class LanguageParser:
"""Base class for language-specific parsers."""
def parse(self, content: str, file_path: str) -> List[Symbol]:
"""Parse content and extract symbols.
Args:
content: File content
file_path: Path to the file
Returns:
List of extracted symbols
"""
raise NotImplementedError
class PythonParser(LanguageParser):
"""Parser for Python files."""
# Regex patterns for Python
FUNCTION_PATTERN = re.compile(
r'^(\s*)def\s+(\w+)\s*\((.*?)\)\s*(?:->.*?)?:',
re.MULTILINE
)
CLASS_PATTERN = re.compile(
r'^(\s*)class\s+(\w+)\s*(?:\((.*?)\))?:',
re.MULTILINE
)
IMPORT_PATTERN = re.compile(
r'^(?:from\s+(\S+)\s+)?import\s+(.+)$',
re.MULTILINE
)
ASYNC_FUNCTION_PATTERN = re.compile(
r'^(\s*)async\s+def\s+(\w+)\s*\((.*?)\)\s*(?:->.*?)?:',
re.MULTILINE
)
def parse(self, content: str, file_path: str) -> List[Symbol]:
symbols = []
lines = content.split('\n')
# Track class context for methods
current_class = None
class_indent = 0
# Parse classes
for match in self.CLASS_PATTERN.finditer(content):
indent = len(match.group(1))
name = match.group(2)
bases = match.group(3) or ""
line_num = content[:match.start()].count('\n') + 1
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.CLASS,
file_path=file_path,
line_number=line_num,
signature=f"class {name}({bases})" if bases else f"class {name}"
))
# Parse functions (including async)
for pattern in [self.FUNCTION_PATTERN, self.ASYNC_FUNCTION_PATTERN]:
for match in pattern.finditer(content):
indent = len(match.group(1))
name = match.group(2)
params = match.group(3)
line_num = content[:match.start()].count('\n') + 1
# Determine if it's a method (inside a class)
is_method = indent > 0
parent = None
if is_method:
# Find parent class
for sym in symbols:
if sym.symbol_type == SymbolType.CLASS:
if sym.line_number < line_num:
parent = sym.name
symbol_type = SymbolType.METHOD if is_method else SymbolType.FUNCTION
is_async = "async" in match.group(0)
prefix = "async def" if is_async else "def"
symbols.append(Symbol(
name=name,
symbol_type=symbol_type,
file_path=file_path,
line_number=line_num,
signature=f"{prefix} {name}({params})",
parent=parent
))
# Parse imports
for match in self.IMPORT_PATTERN.finditer(content):
from_module = match.group(1)
imports = match.group(2)
line_num = content[:match.start()].count('\n') + 1
if from_module:
sig = f"from {from_module} import {imports}"
else:
sig = f"import {imports}"
symbols.append(Symbol(
name=imports.split(',')[0].strip().split(' ')[0],
symbol_type=SymbolType.IMPORT,
file_path=file_path,
line_number=line_num,
signature=sig
))
return symbols
class JavaScriptParser(LanguageParser):
"""Parser for JavaScript/TypeScript files."""
FUNCTION_PATTERN = re.compile(
r'(?:export\s+)?(?:async\s+)?function\s+(\w+)\s*\((.*?)\)',
re.MULTILINE
)
ARROW_FUNCTION_PATTERN = re.compile(
r'(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\((.*?)\)\s*=>',
re.MULTILINE
)
CLASS_PATTERN = re.compile(
r'(?:export\s+)?class\s+(\w+)(?:\s+extends\s+(\w+))?',
re.MULTILINE
)
IMPORT_PATTERN = re.compile(
r'import\s+(?:{([^}]+)}|(\w+))\s+from\s+[\'"]([^\'"]+)[\'"]',
re.MULTILINE
)
EXPORT_PATTERN = re.compile(
r'export\s+(?:default\s+)?(?:const|let|var|function|class)\s+(\w+)',
re.MULTILINE
)
def parse(self, content: str, file_path: str) -> List[Symbol]:
symbols = []
# Parse functions
for match in self.FUNCTION_PATTERN.finditer(content):
name = match.group(1)
params = match.group(2)
line_num = content[:match.start()].count('\n') + 1
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.FUNCTION,
file_path=file_path,
line_number=line_num,
signature=f"function {name}({params})"
))
# Parse arrow functions
for match in self.ARROW_FUNCTION_PATTERN.finditer(content):
name = match.group(1)
params = match.group(2)
line_num = content[:match.start()].count('\n') + 1
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.FUNCTION,
file_path=file_path,
line_number=line_num,
signature=f"const {name} = ({params}) =>"
))
# Parse classes
for match in self.CLASS_PATTERN.finditer(content):
name = match.group(1)
extends = match.group(2)
line_num = content[:match.start()].count('\n') + 1
sig = f"class {name}"
if extends:
sig += f" extends {extends}"
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.CLASS,
file_path=file_path,
line_number=line_num,
signature=sig
))
# Parse imports
for match in self.IMPORT_PATTERN.finditer(content):
named = match.group(1)
default = match.group(2)
module = match.group(3)
line_num = content[:match.start()].count('\n') + 1
name = default or (named.split(',')[0].strip() if named else module)
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.IMPORT,
file_path=file_path,
line_number=line_num,
signature=f"import from '{module}'"
))
# Parse exports
for match in self.EXPORT_PATTERN.finditer(content):
name = match.group(1)
line_num = content[:match.start()].count('\n') + 1
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.EXPORT,
file_path=file_path,
line_number=line_num,
signature=f"export {name}"
))
return symbols
class GoParser(LanguageParser):
"""Parser for Go files."""
FUNCTION_PATTERN = re.compile(
r'^func\s+(?:\((\w+)\s+\*?(\w+)\)\s+)?(\w+)\s*\((.*?)\)',
re.MULTILINE
)
TYPE_PATTERN = re.compile(
r'^type\s+(\w+)\s+(struct|interface)',
re.MULTILINE
)
IMPORT_PATTERN = re.compile(
r'import\s+(?:\(\s*([\s\S]*?)\s*\)|"([^"]+)")',
re.MULTILINE
)
def parse(self, content: str, file_path: str) -> List[Symbol]:
symbols = []
# Parse functions
for match in self.FUNCTION_PATTERN.finditer(content):
receiver_name = match.group(1)
receiver_type = match.group(2)
name = match.group(3)
params = match.group(4)
line_num = content[:match.start()].count('\n') + 1
if receiver_type:
# Method
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.METHOD,
file_path=file_path,
line_number=line_num,
signature=f"func ({receiver_name} {receiver_type}) {name}({params})",
parent=receiver_type
))
else:
# Function
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.FUNCTION,
file_path=file_path,
line_number=line_num,
signature=f"func {name}({params})"
))
# Parse types
for match in self.TYPE_PATTERN.finditer(content):
name = match.group(1)
kind = match.group(2)
line_num = content[:match.start()].count('\n') + 1
symbol_type = SymbolType.STRUCT if kind == "struct" else SymbolType.INTERFACE
symbols.append(Symbol(
name=name,
symbol_type=symbol_type,
file_path=file_path,
line_number=line_num,
signature=f"type {name} {kind}"
))
return symbols
class RustParser(LanguageParser):
"""Parser for Rust files."""
FUNCTION_PATTERN = re.compile(
r'^(\s*)(?:pub\s+)?(?:async\s+)?fn\s+(\w+)\s*(?:<[^>]*>)?\s*\((.*?)\)',
re.MULTILINE
)
STRUCT_PATTERN = re.compile(
r'^(?:pub\s+)?struct\s+(\w+)',
re.MULTILINE
)
IMPL_PATTERN = re.compile(
r'^impl(?:<[^>]*>)?\s+(?:(\w+)\s+for\s+)?(\w+)',
re.MULTILINE
)
USE_PATTERN = re.compile(
r'^use\s+(.+);',
re.MULTILINE
)
def parse(self, content: str, file_path: str) -> List[Symbol]:
symbols = []
# Parse functions
for match in self.FUNCTION_PATTERN.finditer(content):
indent = len(match.group(1))
name = match.group(2)
params = match.group(3)
line_num = content[:match.start()].count('\n') + 1
is_method = indent > 0
symbol_type = SymbolType.METHOD if is_method else SymbolType.FUNCTION
symbols.append(Symbol(
name=name,
symbol_type=symbol_type,
file_path=file_path,
line_number=line_num,
signature=f"fn {name}({params})"
))
# Parse structs
for match in self.STRUCT_PATTERN.finditer(content):
name = match.group(1)
line_num = content[:match.start()].count('\n') + 1
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.STRUCT,
file_path=file_path,
line_number=line_num,
signature=f"struct {name}"
))
# Parse impl blocks
for match in self.IMPL_PATTERN.finditer(content):
trait_name = match.group(1)
type_name = match.group(2)
line_num = content[:match.start()].count('\n') + 1
if trait_name:
sig = f"impl {trait_name} for {type_name}"
else:
sig = f"impl {type_name}"
symbols.append(Symbol(
name=type_name,
symbol_type=SymbolType.TYPE,
file_path=file_path,
line_number=line_num,
signature=sig
))
return symbols
class JavaParser(LanguageParser):
"""Parser for Java files."""
CLASS_PATTERN = re.compile(
r'(?:public\s+)?(?:abstract\s+)?(?:final\s+)?class\s+(\w+)(?:\s+extends\s+(\w+))?(?:\s+implements\s+([^{]+))?',
re.MULTILINE
)
METHOD_PATTERN = re.compile(
r'(?:public|private|protected)?\s*(?:static\s+)?(?:final\s+)?(?:\w+(?:<[^>]+>)?)\s+(\w+)\s*\((.*?)\)',
re.MULTILINE
)
IMPORT_PATTERN = re.compile(
r'^import\s+(?:static\s+)?([^;]+);',
re.MULTILINE
)
def parse(self, content: str, file_path: str) -> List[Symbol]:
symbols = []
# Parse classes
for match in self.CLASS_PATTERN.finditer(content):
name = match.group(1)
extends = match.group(2)
implements = match.group(3)
line_num = content[:match.start()].count('\n') + 1
sig = f"class {name}"
if extends:
sig += f" extends {extends}"
if implements:
sig += f" implements {implements.strip()}"
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.CLASS,
file_path=file_path,
line_number=line_num,
signature=sig
))
# Parse methods
for match in self.METHOD_PATTERN.finditer(content):
name = match.group(1)
params = match.group(2)
line_num = content[:match.start()].count('\n') + 1
# Skip constructors and common keywords
if name in ('if', 'while', 'for', 'switch', 'catch'):
continue
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.METHOD,
file_path=file_path,
line_number=line_num,
signature=f"{name}({params})"
))
# Parse imports
for match in self.IMPORT_PATTERN.finditer(content):
import_path = match.group(1)
line_num = content[:match.start()].count('\n') + 1
# Get the last part of the import
name = import_path.split('.')[-1]
symbols.append(Symbol(
name=name,
symbol_type=SymbolType.IMPORT,
file_path=file_path,
line_number=line_num,
signature=f"import {import_path}"
))
return symbols
class SymbolIndexer:
"""Indexes code symbols from source files.
Supports multiple languages:
- Python (.py)
- JavaScript (.js, .jsx)
- TypeScript (.ts, .tsx)
- Go (.go)
- Rust (.rs)
- Java (.java)
"""
# Map file extensions to parsers
PARSERS = {
'.py': PythonParser,
'.js': JavaScriptParser,
'.jsx': JavaScriptParser,
'.ts': JavaScriptParser,
'.tsx': JavaScriptParser,
'.go': GoParser,
'.rs': RustParser,
'.java': JavaParser,
}
def __init__(self, workspace_path: str):
"""Initialize symbol indexer.
Args:
workspace_path: Root directory to index
"""
self.workspace_path = os.path.abspath(workspace_path)
self.symbols: Dict[str, List[Symbol]] = {} # file_path -> symbols
self.by_name: Dict[str, List[Symbol]] = {} # symbol_name -> symbols
self.by_type: Dict[SymbolType, List[Symbol]] = {} # type -> symbols
# Initialize parsers
self._parsers = {ext: parser() for ext, parser in self.PARSERS.items()}
def index_file(self, file_path: str) -> List[Symbol]:
"""Index symbols from a single file.
Args:
file_path: Path to the file
Returns:
List of extracted symbols
"""
ext = os.path.splitext(file_path)[1].lower()
parser = self._parsers.get(ext)
if not parser:
return []
try:
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
content = f.read()
rel_path = os.path.relpath(file_path, self.workspace_path)
symbols = parser.parse(content, rel_path)
# Store in indexes
self.symbols[rel_path] = symbols
for sym in symbols:
# By name
if sym.name not in self.by_name:
self.by_name[sym.name] = []
self.by_name[sym.name].append(sym)
# By type
if sym.symbol_type not in self.by_type:
self.by_type[sym.symbol_type] = []
self.by_type[sym.symbol_type].append(sym)
return symbols
except Exception as e:
logger.debug(f"Error indexing {file_path}: {e}")
return []
def index(self) -> int:
"""Index all supported files in workspace.
Returns:
Total number of symbols indexed
"""
total_symbols = 0
for root, dirs, files in os.walk(self.workspace_path):
# Skip hidden directories
dirs[:] = [d for d in dirs if not d.startswith('.')]
for filename in files:
ext = os.path.splitext(filename)[1].lower()
if ext in self.PARSERS:
file_path = os.path.join(root, filename)
symbols = self.index_file(file_path)
total_symbols += len(symbols)
logger.info(f"Indexed {total_symbols} symbols from {len(self.symbols)} files")
return total_symbols
def find_by_name(self, name: str, exact: bool = False) -> List[Symbol]:
"""Find symbols by name.
Args:
name: Symbol name to search
exact: If True, match exact name
Returns:
List of matching symbols
"""
if exact:
return self.by_name.get(name, [])
results = []
name_lower = name.lower()
for sym_name, symbols in self.by_name.items():
if name_lower in sym_name.lower():
results.extend(symbols)
return results
def find_by_type(self, symbol_type: SymbolType) -> List[Symbol]:
"""Find symbols by type.
Args:
symbol_type: Type of symbol
Returns:
List of matching symbols
"""
return self.by_type.get(symbol_type, [])
def find_in_file(self, file_path: str) -> List[Symbol]:
"""Find all symbols in a file.
Args:
file_path: Path to the file
Returns:
List of symbols in the file
"""
return self.symbols.get(file_path, [])
def get_stats(self) -> Dict[str, Any]:
"""Get indexer statistics.
Returns:
Dictionary with stats
"""
type_counts = {t.value: len(syms) for t, syms in self.by_type.items()}
return {
"total_files": len(self.symbols),
"total_symbols": sum(len(s) for s in self.symbols.values()),
"by_type": type_counts
}
@@ -0,0 +1,325 @@
"""
Parallel execution engine for Fast Context.
Provides async parallel execution of search tools with:
- Configurable parallelism (default: 8 concurrent calls)
- Timeout handling per tool call
- Result aggregation
- Error isolation (one failure doesn't affect others)
"""
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict, Any, Optional, Callable
import logging
import time
logger = logging.getLogger(__name__)
class ToolCallBatch:
"""Batch of tool calls to execute in parallel.
Attributes:
tasks: List of task dictionaries
max_size: Maximum batch size (default: 8)
"""
def __init__(self, max_size: int = 8):
"""Initialize batch.
Args:
max_size: Maximum number of tasks in batch
"""
self.tasks: List[Dict[str, Any]] = []
self.max_size = max_size
def add(self, tool_name: str, **kwargs) -> bool:
"""Add a tool call to the batch.
Args:
tool_name: Name of the tool to call
**kwargs: Tool arguments
Returns:
True if added, False if batch is full
"""
if self.is_full:
return False
self.tasks.append({
"tool": tool_name,
"args": kwargs
})
return True
@property
def is_full(self) -> bool:
"""Check if batch is at max capacity."""
return len(self.tasks) >= self.max_size
def clear(self) -> None:
"""Clear all tasks from batch."""
self.tasks.clear()
class ParallelExecutor:
"""Executes search tools in parallel.
Uses asyncio for concurrent execution with a thread pool
for CPU-bound operations.
Attributes:
max_parallel: Maximum concurrent executions
timeout: Timeout per tool call in seconds
"""
def __init__(
self,
max_parallel: int = 8,
timeout: float = 30.0
):
"""Initialize executor.
Args:
max_parallel: Maximum concurrent tool calls
timeout: Timeout per tool call in seconds
"""
self.max_parallel = max_parallel
self.timeout = timeout
self._executor: Optional[ThreadPoolExecutor] = None
def _get_executor(self) -> ThreadPoolExecutor:
"""Get or create thread pool executor."""
if self._executor is None:
self._executor = ThreadPoolExecutor(max_workers=self.max_parallel)
return self._executor
def _get_tool_function(self, tool_name: str) -> Optional[Callable]:
"""Get tool function by name.
Args:
tool_name: Name of the tool
Returns:
Tool function or None if not found
"""
# Import here to avoid circular imports
from praisonaiagents.context.fast.search_tools import (
grep_search,
glob_search,
read_file,
list_directory
)
tools = {
"grep_search": grep_search,
"glob_search": glob_search,
"read_file": read_file,
"list_directory": list_directory
}
return tools.get(tool_name)
async def _execute_single(
self,
task: Dict[str, Any],
loop: asyncio.AbstractEventLoop
) -> Any:
"""Execute a single tool call.
Args:
task: Task dictionary with 'tool' and 'args'
loop: Event loop for thread execution
Returns:
Tool result or error dictionary
"""
tool_name = task.get("tool")
args = task.get("args", {})
tool_func = self._get_tool_function(tool_name)
if tool_func is None:
return {"error": f"Unknown tool: {tool_name}"}
try:
# Run in thread pool to avoid blocking
executor = self._get_executor()
result = await asyncio.wait_for(
loop.run_in_executor(executor, lambda: tool_func(**args)),
timeout=self.timeout
)
return result
except asyncio.TimeoutError:
return {"error": f"Timeout executing {tool_name}", "timeout": True}
except Exception as e:
logger.debug(f"Error executing {tool_name}: {e}")
return {"error": str(e)}
async def execute(self, tasks: List[Dict[str, Any]]) -> List[Any]:
"""Execute multiple tool calls in parallel.
Args:
tasks: List of task dictionaries, each with 'tool' and 'args'
Returns:
List of results in same order as tasks
"""
if not tasks:
return []
loop = asyncio.get_event_loop()
# Create semaphore to limit concurrency
semaphore = asyncio.Semaphore(self.max_parallel)
async def bounded_execute(task: Dict[str, Any]) -> Any:
async with semaphore:
return await self._execute_single(task, loop)
# Execute all tasks concurrently
results = await asyncio.gather(
*[bounded_execute(task) for task in tasks],
return_exceptions=True
)
# Convert exceptions to error dicts
processed_results = []
for result in results:
if isinstance(result, Exception):
processed_results.append({"error": str(result)})
else:
processed_results.append(result)
return processed_results
def execute_sync(self, tasks: List[Dict[str, Any]]) -> List[Any]:
"""Synchronous wrapper for execute.
Args:
tasks: List of task dictionaries
Returns:
List of results
"""
try:
loop = asyncio.get_event_loop()
if loop.is_running():
# If already in async context, create new loop
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as pool:
future = pool.submit(asyncio.run, self.execute(tasks))
return future.result()
else:
return loop.run_until_complete(self.execute(tasks))
except RuntimeError:
# No event loop, create one
return asyncio.run(self.execute(tasks))
def close(self) -> None:
"""Close the executor and release resources."""
if self._executor is not None:
self._executor.shutdown(wait=False)
self._executor = None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
class ParallelSearchCoordinator:
"""Coordinates parallel search operations across multiple turns.
Implements the Fast Context search strategy:
- Up to 8 parallel tool calls per turn
- Maximum 4 turns
- Aggregates results across turns
"""
def __init__(
self,
max_parallel: int = 8,
max_turns: int = 4,
timeout: float = 30.0
):
"""Initialize coordinator.
Args:
max_parallel: Max parallel calls per turn
max_turns: Maximum number of turns
timeout: Timeout per tool call
"""
self.max_parallel = max_parallel
self.max_turns = max_turns
self.timeout = timeout
self.executor = ParallelExecutor(max_parallel, timeout)
# Track execution stats
self.total_tool_calls = 0
self.turns_used = 0
self.start_time: Optional[float] = None
self.end_time: Optional[float] = None
def start(self) -> None:
"""Start timing the search."""
self.start_time = time.perf_counter()
self.total_tool_calls = 0
self.turns_used = 0
def stop(self) -> None:
"""Stop timing the search."""
self.end_time = time.perf_counter()
@property
def elapsed_ms(self) -> int:
"""Get elapsed time in milliseconds."""
if self.start_time is None:
return 0
end = self.end_time or time.perf_counter()
return int((end - self.start_time) * 1000)
async def execute_turn(self, batch: ToolCallBatch) -> List[Any]:
"""Execute a single turn of parallel tool calls.
Args:
batch: Batch of tool calls
Returns:
List of results
"""
if self.turns_used >= self.max_turns:
logger.warning("Maximum turns reached")
return []
self.turns_used += 1
self.total_tool_calls += len(batch.tasks)
results = await self.executor.execute(batch.tasks)
return results
def execute_turn_sync(self, batch: ToolCallBatch) -> List[Any]:
"""Synchronous wrapper for execute_turn.
Args:
batch: Batch of tool calls
Returns:
List of results
"""
try:
loop = asyncio.get_event_loop()
if loop.is_running():
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as pool:
future = pool.submit(asyncio.run, self.execute_turn(batch))
return future.result()
else:
return loop.run_until_complete(self.execute_turn(batch))
except RuntimeError:
return asyncio.run(self.execute_turn(batch))
def close(self) -> None:
"""Close resources."""
self.executor.close()
@@ -0,0 +1,210 @@
"""
Result classes for Fast Context.
These dataclasses represent the output of Fast Context searches,
providing structured file and line range information.
"""
from dataclasses import dataclass, field
from typing import List, Optional
import time
@dataclass
class LineRange:
"""Represents a range of lines in a file.
Attributes:
start: Starting line number (1-indexed)
end: Ending line number (1-indexed, inclusive)
content: Optional content of the lines
relevance_score: Score indicating relevance (0.0-1.0)
"""
start: int
end: int
content: Optional[str] = None
relevance_score: float = 1.0
def __post_init__(self):
if self.start < 1:
self.start = 1
if self.end < self.start:
self.end = self.start
@property
def line_count(self) -> int:
"""Number of lines in this range."""
return self.end - self.start + 1
def overlaps(self, other: "LineRange") -> bool:
"""Check if this range overlaps with another."""
return not (self.end < other.start or self.start > other.end)
def merge(self, other: "LineRange") -> "LineRange":
"""Merge with another overlapping range."""
if not self.overlaps(other):
raise ValueError("Cannot merge non-overlapping ranges")
return LineRange(
start=min(self.start, other.start),
end=max(self.end, other.end),
content=None, # Content needs to be re-fetched
relevance_score=max(self.relevance_score, other.relevance_score)
)
@dataclass
class FileMatch:
"""Represents a file match from Fast Context search.
Attributes:
path: Absolute or relative path to the file
line_ranges: List of relevant line ranges in the file
relevance_score: Overall relevance score (0.0-1.0)
match_count: Number of pattern matches found
"""
path: str
line_ranges: List[LineRange] = field(default_factory=list)
relevance_score: float = 1.0
match_count: int = 0
def add_line_range(self, line_range: LineRange) -> None:
"""Add a line range, merging with existing overlapping ranges."""
# Find overlapping ranges
overlapping = []
non_overlapping = []
for existing in self.line_ranges:
if existing.overlaps(line_range):
overlapping.append(existing)
else:
non_overlapping.append(existing)
# Merge all overlapping ranges
merged = line_range
for overlap in overlapping:
merged = merged.merge(overlap)
self.line_ranges = non_overlapping + [merged]
self.line_ranges.sort(key=lambda r: r.start)
@property
def total_lines(self) -> int:
"""Total number of lines across all ranges."""
return sum(r.line_count for r in self.line_ranges)
@dataclass
class FastContextResult:
"""Result of a Fast Context search operation.
Attributes:
files: List of matching files with line ranges
query: The original search query
search_time_ms: Time taken for the search in milliseconds
turns_used: Number of turns used by the subagent
total_tool_calls: Total number of tool calls made
from_cache: Whether result was retrieved from cache
"""
files: List[FileMatch] = field(default_factory=list)
query: str = ""
search_time_ms: int = 0
turns_used: int = 0
total_tool_calls: int = 0
from_cache: bool = False
_start_time: float = field(default=0.0, repr=False)
def start_timer(self) -> None:
"""Start the search timer."""
self._start_time = time.perf_counter()
def stop_timer(self) -> None:
"""Stop the timer and record elapsed time."""
if self._start_time > 0:
elapsed = time.perf_counter() - self._start_time
self.search_time_ms = int(elapsed * 1000)
def add_file(self, file_match: FileMatch) -> None:
"""Add a file match, merging with existing if same path."""
for existing in self.files:
if existing.path == file_match.path:
# Merge line ranges
for lr in file_match.line_ranges:
existing.add_line_range(lr)
existing.relevance_score = max(
existing.relevance_score,
file_match.relevance_score
)
existing.match_count += file_match.match_count
return
self.files.append(file_match)
def sort_by_relevance(self) -> None:
"""Sort files by relevance score (highest first)."""
self.files.sort(key=lambda f: f.relevance_score, reverse=True)
@property
def total_files(self) -> int:
"""Total number of files found."""
return len(self.files)
@property
def total_lines(self) -> int:
"""Total number of lines across all files."""
return sum(f.total_lines for f in self.files)
def to_context_string(self, max_files: int = 10, include_content: bool = True) -> str:
"""Convert result to a context string for the main agent.
Args:
max_files: Maximum number of files to include
include_content: Whether to include line content
Returns:
Formatted context string
"""
if not self.files:
return "No relevant files found."
lines = [f"Found {self.total_files} relevant file(s):"]
for i, file_match in enumerate(self.files[:max_files]):
lines.append(f"\n## {file_match.path}")
if file_match.line_ranges:
for lr in file_match.line_ranges:
lines.append(f" Lines {lr.start}-{lr.end}")
if include_content and lr.content:
# Indent content
content_lines = lr.content.split('\n')
for cl in content_lines:
lines.append(f" {cl}")
if len(self.files) > max_files:
lines.append(f"\n... and {len(self.files) - max_files} more files")
return '\n'.join(lines)
def to_dict(self) -> dict:
"""Convert to dictionary for serialization."""
return {
"files": [
{
"path": f.path,
"relevance_score": f.relevance_score,
"match_count": f.match_count,
"line_ranges": [
{
"start": lr.start,
"end": lr.end,
"relevance_score": lr.relevance_score
}
for lr in f.line_ranges
]
}
for f in self.files
],
"query": self.query,
"search_time_ms": self.search_time_ms,
"turns_used": self.turns_used,
"total_tool_calls": self.total_tool_calls,
"from_cache": self.from_cache
}
@@ -0,0 +1,368 @@
"""
Search Strategy for Fast Context.
Orchestrates multi-turn search with:
- Turn 1-3: Exploration (grep, glob, directory listing)
- Turn 4: Final answer (file list + line ranges)
- Early termination if sufficient context found
"""
import logging
from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional, Set
from enum import Enum
from praisonaiagents.context.fast.result import FastContextResult, FileMatch, LineRange
from praisonaiagents.context.fast.parallel_executor import ToolCallBatch
logger = logging.getLogger(__name__)
class SearchPhase(Enum):
"""Phases of the search strategy."""
DISCOVERY = "discovery" # Turn 1: Discover structure
EXPLORATION = "exploration" # Turn 2-3: Deep search
REFINEMENT = "refinement" # Turn 4: Final answer
@dataclass
class SearchState:
"""Tracks the state of a multi-turn search.
Attributes:
query: Original search query
current_turn: Current turn number (1-indexed)
max_turns: Maximum turns allowed
phase: Current search phase
discovered_files: Files found so far
explored_paths: Paths already explored
result: Accumulated search result
should_terminate: Whether to terminate early
"""
query: str
current_turn: int = 0
max_turns: int = 4
phase: SearchPhase = SearchPhase.DISCOVERY
discovered_files: Set[str] = field(default_factory=set)
explored_paths: Set[str] = field(default_factory=set)
result: FastContextResult = field(default_factory=FastContextResult)
should_terminate: bool = False
# Thresholds for early termination
min_files_for_termination: int = 5
max_files_before_refinement: int = 20
def advance_turn(self) -> None:
"""Advance to the next turn and update phase."""
self.current_turn += 1
if self.current_turn == 1:
self.phase = SearchPhase.DISCOVERY
elif self.current_turn < self.max_turns:
self.phase = SearchPhase.EXPLORATION
else:
self.phase = SearchPhase.REFINEMENT
def check_early_termination(self) -> bool:
"""Check if search should terminate early.
Returns:
True if search should terminate
"""
# Terminate if we have enough high-quality results
if len(self.discovered_files) >= self.min_files_for_termination:
high_relevance = sum(
1 for f in self.result.files
if f.relevance_score >= 0.8
)
if high_relevance >= 3:
self.should_terminate = True
return True
# Terminate if we've explored everything
if self.current_turn >= 2 and not self._has_unexplored_paths():
self.should_terminate = True
return True
return False
def _has_unexplored_paths(self) -> bool:
"""Check if there are unexplored paths."""
# Simple heuristic: if we found files but haven't explored many paths
return len(self.discovered_files) > len(self.explored_paths) * 2
class SearchStrategy:
"""Orchestrates multi-turn search strategy.
Implements the Fast Context search approach:
- Turn 1: Discovery - understand codebase structure
- Turn 2-3: Exploration - deep search for relevant code
- Turn 4: Refinement - finalize results
Supports early termination when sufficient context is found.
"""
def __init__(
self,
workspace_path: str,
max_turns: int = 4,
max_parallel: int = 8
):
"""Initialize search strategy.
Args:
workspace_path: Root directory for searches
max_turns: Maximum search turns
max_parallel: Maximum parallel tool calls per turn
"""
self.workspace_path = workspace_path
self.max_turns = max_turns
self.max_parallel = max_parallel
def create_state(self, query: str) -> SearchState:
"""Create a new search state.
Args:
query: Search query
Returns:
New SearchState instance
"""
state = SearchState(query=query, max_turns=self.max_turns)
state.result.query = query
return state
def plan_discovery_turn(self, state: SearchState) -> ToolCallBatch:
"""Plan the discovery turn (Turn 1).
Focus on understanding codebase structure:
- List root directory
- Find common file patterns
- Initial grep for query terms
Args:
state: Current search state
Returns:
Batch of tool calls to execute
"""
batch = ToolCallBatch(max_size=self.max_parallel)
query = state.query
# List root directory
batch.add("list_directory", dir_path=self.workspace_path, recursive=False)
# Find Python files (most common)
batch.add("glob_search", search_path=self.workspace_path, pattern="**/*.py", max_results=20)
# Initial grep for query terms
query_words = query.lower().split()
for word in query_words[:3]: # First 3 words
if len(word) >= 3: # Skip short words
batch.add("grep_search", search_path=self.workspace_path, pattern=word, max_results=10)
if batch.is_full:
break
return batch
def plan_exploration_turn(self, state: SearchState) -> ToolCallBatch:
"""Plan an exploration turn (Turn 2-3).
Deep search based on discovery results:
- Search in discovered directories
- Read promising files
- Follow references
Args:
state: Current search state
Returns:
Batch of tool calls to execute
"""
batch = ToolCallBatch(max_size=self.max_parallel)
# Read top files found so far
for file_match in state.result.files[:4]:
if file_match.path not in state.explored_paths:
batch.add("read_file", filepath=file_match.path, context_lines=5)
state.explored_paths.add(file_match.path)
if batch.is_full:
break
# Search in unexplored directories
unexplored_dirs = self._find_unexplored_dirs(state)
for dir_path in unexplored_dirs[:2]:
batch.add("list_directory", dir_path=dir_path, recursive=True, max_depth=2)
if batch.is_full:
break
# Additional grep searches with refined patterns
if not batch.is_full:
refined_patterns = self._generate_refined_patterns(state)
for pattern in refined_patterns[:2]:
batch.add("grep_search", search_path=self.workspace_path, pattern=pattern, max_results=15)
if batch.is_full:
break
return batch
def plan_refinement_turn(self, state: SearchState) -> ToolCallBatch:
"""Plan the refinement turn (Turn 4).
Finalize results:
- Read remaining important files
- Get specific line ranges
Args:
state: Current search state
Returns:
Batch of tool calls to execute
"""
batch = ToolCallBatch(max_size=self.max_parallel)
# Read top files that haven't been fully explored
for file_match in state.result.files[:self.max_parallel]:
if file_match.path not in state.explored_paths:
# Read specific line ranges if available
if file_match.line_ranges:
for lr in file_match.line_ranges[:2]:
batch.add(
"read_file",
filepath=file_match.path,
start_line=max(1, lr.start - 5),
end_line=lr.end + 5
)
if batch.is_full:
break
else:
batch.add("read_file", filepath=file_match.path)
state.explored_paths.add(file_match.path)
if batch.is_full:
break
return batch
def plan_next_turn(self, state: SearchState) -> Optional[ToolCallBatch]:
"""Plan the next turn based on current state.
Args:
state: Current search state
Returns:
Batch of tool calls, or None if search should end
"""
state.advance_turn()
# Check for early termination
if state.should_terminate or state.current_turn > state.max_turns:
return None
if state.check_early_termination():
return None
# Plan based on phase
if state.phase == SearchPhase.DISCOVERY:
return self.plan_discovery_turn(state)
elif state.phase == SearchPhase.EXPLORATION:
return self.plan_exploration_turn(state)
else: # REFINEMENT
return self.plan_refinement_turn(state)
def process_results(
self,
state: SearchState,
tool_results: List[Any]
) -> None:
"""Process results from a turn and update state.
Args:
state: Current search state
tool_results: Results from tool execution
"""
for result in tool_results:
if isinstance(result, list):
# Grep or glob results
for item in result:
if isinstance(item, dict):
path = item.get("path", "")
if path:
state.discovered_files.add(path)
# Create or update file match
file_match = FileMatch(
path=path,
relevance_score=0.8,
match_count=1
)
# Add line range if available
line_num = item.get("line_number")
if line_num:
file_match.add_line_range(LineRange(
start=max(1, line_num - 2),
end=line_num + 2,
content=item.get("context"),
relevance_score=0.8
))
state.result.add_file(file_match)
elif isinstance(result, dict):
# Read file or list directory result
if result.get("success"):
path = result.get("path", "")
if path:
state.explored_paths.add(path)
# Handle directory listing
entries = result.get("entries", [])
for entry in entries:
if not entry.get("is_dir"):
state.discovered_files.add(entry.get("path", ""))
def _find_unexplored_dirs(self, state: SearchState) -> List[str]:
"""Find directories that haven't been explored.
Args:
state: Current search state
Returns:
List of unexplored directory paths
"""
import os
# Extract unique directories from discovered files
dirs = set()
for file_path in state.discovered_files:
dir_path = os.path.dirname(file_path)
if dir_path and dir_path not in state.explored_paths:
dirs.add(dir_path)
return list(dirs)[:5] # Limit to 5
def _generate_refined_patterns(self, state: SearchState) -> List[str]:
"""Generate refined search patterns based on results.
Args:
state: Current search state
Returns:
List of refined search patterns
"""
patterns = []
query_words = state.query.lower().split()
# Combine query words
if len(query_words) >= 2:
patterns.append(f"{query_words[0]}.*{query_words[1]}")
# Add common code patterns
for word in query_words[:2]:
if len(word) >= 3:
patterns.append(f"def {word}")
patterns.append(f"class {word}")
return patterns[:4]
@@ -0,0 +1,636 @@
"""
Search tools for Fast Context.
Provides cross-platform compatible search tools:
- grep_search: Pattern search in files
- glob_search: File pattern matching
- read_file: Read file contents with line range support
- list_directory: List directory contents
These tools are designed to be fast, safe, and cross-platform compatible.
"""
import os
import re
import fnmatch
from pathlib import Path
from typing import List, Dict, Any, Optional, Set
import logging
logger = logging.getLogger(__name__)
def _load_gitignore_patterns(root_path: str) -> Set[str]:
"""Load patterns from .gitignore file.
Args:
root_path: Root directory to search for .gitignore
Returns:
Set of gitignore patterns
"""
patterns = set()
gitignore_path = Path(root_path) / ".gitignore"
if gitignore_path.exists():
try:
with open(gitignore_path, 'r', encoding='utf-8', errors='ignore') as f:
for line in f:
line = line.strip()
# Skip empty lines and comments
if line and not line.startswith('#'):
patterns.add(line)
except Exception:
pass
# Also check for .praisonignore
praisonignore_path = Path(root_path) / ".praisonignore"
if praisonignore_path.exists():
try:
with open(praisonignore_path, 'r', encoding='utf-8', errors='ignore') as f:
for line in f:
line = line.strip()
if line and not line.startswith('#'):
patterns.add(line)
except Exception:
pass
return patterns
def _should_ignore(path: str, root_path: str, patterns: Set[str]) -> bool:
"""Check if a path should be ignored based on gitignore patterns.
Args:
path: Path to check
root_path: Root directory
patterns: Set of gitignore patterns
Returns:
True if path should be ignored
"""
if not patterns:
return False
rel_path = os.path.relpath(path, root_path)
path_parts = Path(rel_path).parts
for pattern in patterns:
# Handle directory patterns (ending with /)
if pattern.endswith('/'):
dir_pattern = pattern.rstrip('/')
if any(fnmatch.fnmatch(part, dir_pattern) for part in path_parts):
return True
# Handle ** patterns
elif '**' in pattern:
# Convert gitignore pattern to fnmatch pattern
fnmatch_pattern = pattern.replace('**/', '*').replace('**', '*')
if fnmatch.fnmatch(rel_path, fnmatch_pattern):
return True
# Handle simple patterns
else:
# Check against filename
if fnmatch.fnmatch(os.path.basename(path), pattern):
return True
# Check against relative path
if fnmatch.fnmatch(rel_path, pattern):
return True
# Check if any path component matches
if any(fnmatch.fnmatch(part, pattern) for part in path_parts):
return True
return False
def _is_binary_file(filepath: str) -> bool:
"""Check if a file is binary.
Args:
filepath: Path to the file
Returns:
True if file appears to be binary
"""
try:
with open(filepath, 'rb') as f:
chunk = f.read(8192)
# Check for null bytes (common in binary files)
if b'\x00' in chunk:
return True
# Check for high ratio of non-text characters
text_chars = bytearray({7, 8, 9, 10, 12, 13, 27} | set(range(0x20, 0x100)) - {0x7f})
non_text = sum(1 for byte in chunk if byte not in text_chars)
if len(chunk) > 0 and non_text / len(chunk) > 0.3:
return True
except Exception:
return True
return False
def grep_search(
search_path: str,
pattern: str,
is_regex: bool = False,
case_sensitive: bool = False,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None,
max_results: int = 100,
context_lines: int = 0,
respect_gitignore: bool = True
) -> List[Dict[str, Any]]:
"""Search for pattern in files.
Args:
search_path: Directory or file to search
pattern: Search pattern (string or regex)
is_regex: If True, treat pattern as regex
case_sensitive: If True, search is case sensitive
include_patterns: Glob patterns to include (e.g., ["*.py"])
exclude_patterns: Glob patterns to exclude (e.g., ["**/test_*"])
max_results: Maximum number of results to return
context_lines: Number of context lines before/after match
respect_gitignore: If True, respect .gitignore patterns
Returns:
List of match dictionaries with path, line_number, content, context
"""
results = []
search_path = os.path.abspath(search_path)
# Compile pattern
if is_regex:
flags = 0 if case_sensitive else re.IGNORECASE
try:
regex = re.compile(pattern, flags)
except re.error as e:
logger.warning(f"Invalid regex pattern: {e}")
return []
else:
if case_sensitive:
regex = re.compile(re.escape(pattern))
else:
regex = re.compile(re.escape(pattern), re.IGNORECASE)
# Load gitignore patterns
gitignore_patterns = set()
if respect_gitignore:
gitignore_patterns = _load_gitignore_patterns(search_path)
# Get files to search
if os.path.isfile(search_path):
files_to_search = [search_path]
else:
files_to_search = []
for root, dirs, files in os.walk(search_path):
# Filter directories in-place to skip ignored ones
if respect_gitignore:
dirs[:] = [d for d in dirs if not _should_ignore(
os.path.join(root, d), search_path, gitignore_patterns
)]
for filename in files:
filepath = os.path.join(root, filename)
# Check gitignore
if respect_gitignore and _should_ignore(filepath, search_path, gitignore_patterns):
continue
# Check include patterns
if include_patterns:
rel_path = os.path.relpath(filepath, search_path)
if not any(fnmatch.fnmatch(rel_path, p) or fnmatch.fnmatch(filename, p)
for p in include_patterns):
continue
# Check exclude patterns
if exclude_patterns:
rel_path = os.path.relpath(filepath, search_path)
if any(fnmatch.fnmatch(rel_path, p) or fnmatch.fnmatch(filename, p)
for p in exclude_patterns):
continue
files_to_search.append(filepath)
# Search files
for filepath in files_to_search:
if len(results) >= max_results:
break
# Skip binary files
if _is_binary_file(filepath):
continue
try:
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
lines = f.readlines()
for i, line in enumerate(lines):
if len(results) >= max_results:
break
if regex.search(line):
# Get context
start_ctx = max(0, i - context_lines)
end_ctx = min(len(lines), i + context_lines + 1)
context = ''.join(lines[start_ctx:end_ctx])
results.append({
"path": os.path.relpath(filepath, search_path),
"absolute_path": filepath,
"line_number": i + 1, # 1-indexed
"content": line.rstrip('\n\r'),
"context": context.rstrip('\n\r') if context_lines > 0 else None,
"context_start": start_ctx + 1 if context_lines > 0 else None,
"context_end": end_ctx if context_lines > 0 else None
})
except Exception as e:
logger.debug(f"Error reading {filepath}: {e}")
continue
return results
def glob_search(
search_path: str,
pattern: str,
max_results: int = 100,
include_dirs: bool = False,
respect_gitignore: bool = True
) -> List[Dict[str, Any]]:
"""Search for files matching glob pattern.
Args:
search_path: Directory to search
pattern: Glob pattern (e.g., "**/*.py", "src/*.js")
max_results: Maximum number of results
include_dirs: If True, include directories in results
respect_gitignore: If True, respect .gitignore patterns
Returns:
List of file info dictionaries with path, size, is_dir
"""
results = []
search_path = os.path.abspath(search_path)
# Load gitignore patterns
gitignore_patterns = set()
if respect_gitignore:
gitignore_patterns = _load_gitignore_patterns(search_path)
# Use pathlib for glob
root = Path(search_path)
try:
for match in root.glob(pattern):
if len(results) >= max_results:
break
# Check gitignore
if respect_gitignore and _should_ignore(str(match), search_path, gitignore_patterns):
continue
# Skip directories unless requested
if match.is_dir() and not include_dirs:
continue
# Skip binary files for file results
if match.is_file() and _is_binary_file(str(match)):
continue
try:
stat = match.stat()
results.append({
"path": str(match.relative_to(root)),
"absolute_path": str(match),
"size": stat.st_size if match.is_file() else None,
"is_dir": match.is_dir(),
"modified": stat.st_mtime
})
except Exception:
continue
except Exception as e:
logger.debug(f"Error in glob search: {e}")
return results
def read_file(
filepath: str,
start_line: Optional[int] = None,
end_line: Optional[int] = None,
context_lines: int = 0,
max_lines: int = 500
) -> Dict[str, Any]:
"""Read file contents with optional line range.
Args:
filepath: Path to the file
start_line: Starting line (1-indexed, inclusive)
end_line: Ending line (1-indexed, inclusive)
context_lines: Additional context lines before/after range
max_lines: Maximum lines to return
Returns:
Dictionary with content, line info, and metadata
"""
filepath = os.path.abspath(filepath)
if not os.path.exists(filepath):
return {
"success": False,
"error": f"File not found: {filepath}",
"path": filepath
}
if os.path.isdir(filepath):
return {
"success": False,
"error": f"Path is a directory: {filepath}",
"path": filepath
}
if _is_binary_file(filepath):
return {
"success": False,
"error": f"Binary file: {filepath}",
"path": filepath
}
try:
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
lines = f.readlines()
total_lines = len(lines)
# Determine line range
if start_line is None:
start_line = 1
if end_line is None:
end_line = total_lines
# Apply context
actual_start = max(1, start_line - context_lines)
actual_end = min(total_lines, end_line + context_lines)
# Limit lines
if actual_end - actual_start + 1 > max_lines:
actual_end = actual_start + max_lines - 1
# Extract content (convert to 0-indexed)
content_lines = lines[actual_start - 1:actual_end]
content = ''.join(content_lines)
return {
"success": True,
"path": filepath,
"content": content.rstrip('\n'),
"start_line": actual_start,
"end_line": actual_end,
"total_lines": total_lines,
"lines_returned": len(content_lines)
}
except Exception as e:
return {
"success": False,
"error": str(e),
"path": filepath
}
def list_directory(
dir_path: str,
recursive: bool = False,
max_depth: int = 10,
max_entries: int = 500,
include_hidden: bool = False,
respect_gitignore: bool = True
) -> Dict[str, Any]:
"""List directory contents.
Args:
dir_path: Directory path
recursive: If True, list recursively
max_depth: Maximum recursion depth
max_entries: Maximum entries to return
include_hidden: If True, include hidden files/dirs
respect_gitignore: If True, respect .gitignore patterns
Returns:
Dictionary with entries list and metadata
"""
dir_path = os.path.abspath(dir_path)
if not os.path.exists(dir_path):
return {
"success": False,
"error": f"Directory not found: {dir_path}",
"path": dir_path
}
if not os.path.isdir(dir_path):
return {
"success": False,
"error": f"Not a directory: {dir_path}",
"path": dir_path
}
# Load gitignore patterns
gitignore_patterns = set()
if respect_gitignore:
gitignore_patterns = _load_gitignore_patterns(dir_path)
entries = []
def _list_dir(path: str, depth: int):
if depth > max_depth or len(entries) >= max_entries:
return
try:
for entry in os.scandir(path):
if len(entries) >= max_entries:
break
# Skip hidden files unless requested
if not include_hidden and entry.name.startswith('.'):
continue
# Check gitignore
if respect_gitignore and _should_ignore(entry.path, dir_path, gitignore_patterns):
continue
try:
stat = entry.stat()
entries.append({
"name": entry.name,
"path": os.path.relpath(entry.path, dir_path),
"absolute_path": entry.path,
"is_dir": entry.is_dir(),
"size": stat.st_size if entry.is_file() else None,
"modified": stat.st_mtime
})
except Exception:
continue
# Recurse into directories
if recursive and entry.is_dir():
_list_dir(entry.path, depth + 1)
except PermissionError:
pass
except Exception as e:
logger.debug(f"Error listing {path}: {e}")
_list_dir(dir_path, 0)
# Sort entries: directories first, then by name
entries.sort(key=lambda e: (not e["is_dir"], e["name"].lower()))
return {
"success": True,
"path": dir_path,
"entries": entries,
"total_entries": len(entries),
"truncated": len(entries) >= max_entries
}
# Tool definitions for LLM function calling
FAST_CONTEXT_TOOLS = [
{
"name": "grep_search",
"description": "Search for a pattern in files within a directory. Returns matching lines with file paths and line numbers.",
"parameters": {
"type": "object",
"properties": {
"search_path": {
"type": "string",
"description": "Directory or file path to search"
},
"pattern": {
"type": "string",
"description": "Search pattern (text or regex)"
},
"is_regex": {
"type": "boolean",
"description": "If true, treat pattern as regex",
"default": False
},
"case_sensitive": {
"type": "boolean",
"description": "If true, search is case sensitive",
"default": False
},
"include_patterns": {
"type": "array",
"items": {"type": "string"},
"description": "Glob patterns to include (e.g., ['*.py'])"
},
"max_results": {
"type": "integer",
"description": "Maximum results to return",
"default": 50
}
},
"required": ["search_path", "pattern"]
}
},
{
"name": "glob_search",
"description": "Find files matching a glob pattern. Returns file paths and metadata.",
"parameters": {
"type": "object",
"properties": {
"search_path": {
"type": "string",
"description": "Directory to search"
},
"pattern": {
"type": "string",
"description": "Glob pattern (e.g., '**/*.py', 'src/*.js')"
},
"max_results": {
"type": "integer",
"description": "Maximum results to return",
"default": 50
}
},
"required": ["search_path", "pattern"]
}
},
{
"name": "read_file",
"description": "Read contents of a file, optionally with specific line range.",
"parameters": {
"type": "object",
"properties": {
"filepath": {
"type": "string",
"description": "Path to the file"
},
"start_line": {
"type": "integer",
"description": "Starting line number (1-indexed)"
},
"end_line": {
"type": "integer",
"description": "Ending line number (1-indexed)"
},
"context_lines": {
"type": "integer",
"description": "Additional context lines",
"default": 0
}
},
"required": ["filepath"]
}
},
{
"name": "list_directory",
"description": "List contents of a directory.",
"parameters": {
"type": "object",
"properties": {
"dir_path": {
"type": "string",
"description": "Directory path"
},
"recursive": {
"type": "boolean",
"description": "If true, list recursively",
"default": False
},
"max_depth": {
"type": "integer",
"description": "Maximum recursion depth",
"default": 3
}
},
"required": ["dir_path"]
}
}
]
def execute_tool(tool_name: str, **kwargs) -> Dict[str, Any]:
"""Execute a search tool by name.
Args:
tool_name: Name of the tool to execute
**kwargs: Tool arguments
Returns:
Tool result dictionary
"""
tools = {
"grep_search": grep_search,
"glob_search": glob_search,
"read_file": read_file,
"list_directory": list_directory
}
if tool_name not in tools:
return {"error": f"Unknown tool: {tool_name}"}
try:
return tools[tool_name](**kwargs)
except Exception as e:
return {"error": str(e)}
@@ -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,19 +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 .file_memory import FileMemory, create_memory
from .rules_manager import RulesManager, Rule, create_rules_manager
# Lazy import for Memory to avoid dependency issues
# 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__ = ["FileMemory", "Memory", "create_memory"]
__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)
@@ -837,6 +837,445 @@ class FileMemory:
self._save_summaries()
self._log("Imported memory data")
# -------------------------------------------------------------------------
# Session Management
# -------------------------------------------------------------------------
def save_session(
self,
name: str,
conversation_history: Optional[List[Dict[str, Any]]] = None,
metadata: Optional[Dict[str, Any]] = None
) -> str:
"""
Save current session state for later resumption.
Similar to Gemini CLI's /chat save command.
Args:
name: Session name (used as filename)
conversation_history: Optional conversation messages to save
metadata: Optional additional metadata
Returns:
Path to saved session file
"""
sessions_path = self.user_path / "sessions"
sessions_path.mkdir(parents=True, exist_ok=True)
session_file = sessions_path / f"{name}.json"
session_data = {
"name": name,
"user_id": self.user_id,
"saved_at": time.time(),
"saved_at_iso": datetime.now().isoformat(),
"short_term": [item.to_dict() for item in self._short_term],
"long_term_snapshot": [item.to_dict() for item in self._long_term[-50:]], # Last 50
"entity_ids": list(self._entities.keys()),
"conversation_history": conversation_history or [],
"metadata": metadata or {},
"config": self.config
}
self._write_json(session_file, session_data)
self._log(f"Saved session '{name}' to {session_file}")
return str(session_file)
def resume_session(self, name: str) -> Dict[str, Any]:
"""
Resume a previously saved session.
Similar to Gemini CLI's /chat resume command.
Args:
name: Session name to resume
Returns:
Session data including conversation history
"""
sessions_path = self.user_path / "sessions"
session_file = sessions_path / f"{name}.json"
if not session_file.exists():
raise FileNotFoundError(f"Session '{name}' not found")
session_data = self._read_json(session_file, {})
# Restore short-term memory from session
if "short_term" in session_data:
self._short_term = [MemoryItem.from_dict(item) for item in session_data["short_term"]]
self._save_short_term()
self._log(f"Resumed session '{name}'")
return session_data
def list_sessions(self) -> List[Dict[str, Any]]:
"""
List all saved sessions.
Similar to Gemini CLI's /chat list command.
Returns:
List of session info dicts
"""
sessions_path = self.user_path / "sessions"
if not sessions_path.exists():
return []
sessions = []
for session_file in sessions_path.glob("*.json"):
try:
data = self._read_json(session_file, {})
sessions.append({
"name": data.get("name", session_file.stem),
"saved_at": data.get("saved_at_iso", ""),
"short_term_count": len(data.get("short_term", [])),
"has_conversation": len(data.get("conversation_history", [])) > 0,
"file_path": str(session_file)
})
except Exception:
continue
# Sort by saved_at descending
sessions.sort(key=lambda x: x.get("saved_at", ""), reverse=True)
return sessions
def delete_session(self, name: str) -> bool:
"""Delete a saved session."""
sessions_path = self.user_path / "sessions"
session_file = sessions_path / f"{name}.json"
if session_file.exists():
session_file.unlink()
self._log(f"Deleted session '{name}'")
return True
return False
# -------------------------------------------------------------------------
# Context Compression
# -------------------------------------------------------------------------
def compress(
self,
llm_func: Optional[callable] = None,
max_items: int = 10
) -> str:
"""
Compress short-term memory into a summary.
Similar to Gemini CLI's /compress command.
Args:
llm_func: Optional LLM function for summarization.
Should accept (prompt: str) -> str
max_items: Max items to keep after compression
Returns:
The generated summary
"""
if len(self._short_term) <= max_items:
return "" # No compression needed
# Gather content to compress
items_to_compress = self._short_term[:-max_items]
content_list = [item.content for item in items_to_compress]
# Generate summary
if llm_func:
prompt = f"""Summarize the following conversation context into key points.
Preserve important facts, decisions, and context.
Be concise but comprehensive.
Context to summarize:
{chr(10).join(f'- {c}' for c in content_list)}
Summary:"""
summary = llm_func(prompt)
else:
# Simple concatenation if no LLM
summary = "Compressed context: " + " | ".join(content_list[:5]) + "..."
# Add summary as a high-importance long-term memory
self.add_long_term(
content=f"[Session Summary] {summary}",
metadata={"type": "compression_summary", "items_compressed": len(items_to_compress)},
importance=0.9
)
# Keep only recent items
self._short_term = self._short_term[-max_items:]
self._save_short_term()
self._log(f"Compressed {len(items_to_compress)} items into summary")
return summary
def auto_compress_if_needed(
self,
threshold_percent: float = 0.7,
llm_func: Optional[callable] = None
) -> Optional[str]:
"""
Auto-compress if short-term memory exceeds threshold.
Args:
threshold_percent: Compress when STM is this % full (0.0-1.0)
llm_func: Optional LLM function for summarization
Returns:
Summary if compressed, None otherwise
"""
limit = self.config["short_term_limit"]
threshold = int(limit * threshold_percent)
if len(self._short_term) >= threshold:
return self.compress(llm_func=llm_func, max_items=int(limit * 0.3))
return None
# -------------------------------------------------------------------------
# Checkpointing
# -------------------------------------------------------------------------
def create_checkpoint(
self,
name: Optional[str] = None,
include_files: Optional[List[str]] = None
) -> str:
"""
Create a checkpoint of current memory state.
Similar to Gemini CLI's checkpoint before destructive operations.
Args:
name: Optional checkpoint name (auto-generated if not provided)
include_files: Optional list of file paths to snapshot
Returns:
Checkpoint ID
"""
checkpoints_path = self.user_path / "checkpoints"
checkpoints_path.mkdir(parents=True, exist_ok=True)
checkpoint_id = name or f"checkpoint_{int(time.time())}"
checkpoint_file = checkpoints_path / f"{checkpoint_id}.json"
checkpoint_data = {
"id": checkpoint_id,
"created_at": time.time(),
"created_at_iso": datetime.now().isoformat(),
"memory_export": self.export(),
"file_snapshots": {}
}
# Optionally snapshot files
if include_files:
for file_path in include_files:
try:
path = Path(file_path)
if path.exists() and path.is_file():
checkpoint_data["file_snapshots"][file_path] = path.read_text(encoding="utf-8")
except Exception as e:
self._log(f"Could not snapshot file {file_path}: {e}", logging.WARNING)
self._write_json(checkpoint_file, checkpoint_data)
self._log(f"Created checkpoint '{checkpoint_id}'")
return checkpoint_id
def restore_checkpoint(self, checkpoint_id: str, restore_files: bool = False) -> bool:
"""
Restore memory state from a checkpoint.
Args:
checkpoint_id: Checkpoint ID to restore
restore_files: Whether to restore file snapshots
Returns:
True if restored successfully
"""
checkpoints_path = self.user_path / "checkpoints"
checkpoint_file = checkpoints_path / f"{checkpoint_id}.json"
if not checkpoint_file.exists():
self._log(f"Checkpoint '{checkpoint_id}' not found", logging.ERROR)
return False
checkpoint_data = self._read_json(checkpoint_file, {})
# Restore memory
if "memory_export" in checkpoint_data:
self.import_data(checkpoint_data["memory_export"])
# Optionally restore files
if restore_files and "file_snapshots" in checkpoint_data:
for file_path, content in checkpoint_data["file_snapshots"].items():
try:
Path(file_path).write_text(content, encoding="utf-8")
self._log(f"Restored file: {file_path}")
except Exception as e:
self._log(f"Could not restore file {file_path}: {e}", logging.WARNING)
self._log(f"Restored checkpoint '{checkpoint_id}'")
return True
def list_checkpoints(self) -> List[Dict[str, Any]]:
"""List all checkpoints."""
checkpoints_path = self.user_path / "checkpoints"
if not checkpoints_path.exists():
return []
checkpoints = []
for checkpoint_file in checkpoints_path.glob("*.json"):
try:
data = self._read_json(checkpoint_file, {})
checkpoints.append({
"id": data.get("id", checkpoint_file.stem),
"created_at": data.get("created_at_iso", ""),
"has_file_snapshots": len(data.get("file_snapshots", {})) > 0,
"file_path": str(checkpoint_file)
})
except Exception:
continue
# Sort by created_at descending
checkpoints.sort(key=lambda x: x.get("created_at", ""), reverse=True)
return checkpoints
def delete_checkpoint(self, checkpoint_id: str) -> bool:
"""Delete a checkpoint."""
checkpoints_path = self.user_path / "checkpoints"
checkpoint_file = checkpoints_path / f"{checkpoint_id}.json"
if checkpoint_file.exists():
checkpoint_file.unlink()
self._log(f"Deleted checkpoint '{checkpoint_id}'")
return True
return False
# -------------------------------------------------------------------------
# Memory Commands (Slash Commands)
# -------------------------------------------------------------------------
def handle_command(self, command: str) -> Dict[str, Any]:
"""
Handle slash commands like /memory show, /memory add, etc.
Similar to Gemini CLI's /memory commands.
Args:
command: Full command string (e.g., "/memory show" or "/memory add User likes Python")
Returns:
Command result dict
"""
parts = command.strip().split(maxsplit=2)
if len(parts) < 2:
return {"error": "Invalid command. Use: /memory <action> [args]"}
action = parts[1].lower()
args = parts[2] if len(parts) > 2 else ""
if action == "show":
return {
"action": "show",
"stats": self.get_stats(),
"short_term": [item.content for item in self.get_short_term(limit=10)],
"long_term": [item.content for item in self.get_long_term(limit=10)]
}
elif action == "add":
if not args:
return {"error": "Usage: /memory add <content>"}
mem_id = self.add_long_term(args, importance=0.8)
return {"action": "add", "id": mem_id, "content": args}
elif action == "clear":
target = args.lower() if args else "short"
if target == "all":
self.clear_all()
return {"action": "clear", "target": "all"}
elif target == "short":
self.clear_short_term()
return {"action": "clear", "target": "short_term"}
else:
return {"error": "Usage: /memory clear [short|all]"}
elif action == "search":
if not args:
return {"error": "Usage: /memory search <query>"}
results = self.search(args, limit=10)
return {"action": "search", "query": args, "results": results}
elif action == "save":
if not args:
return {"error": "Usage: /memory save <session_name>"}
path = self.save_session(args)
return {"action": "save", "session": args, "path": path}
elif action == "resume":
if not args:
return {"error": "Usage: /memory resume <session_name>"}
try:
self.resume_session(args)
return {"action": "resume", "session": args, "restored": True}
except FileNotFoundError:
return {"error": f"Session '{args}' not found"}
elif action == "sessions":
sessions = self.list_sessions()
return {"action": "sessions", "sessions": sessions}
elif action == "compress":
summary = self.compress()
return {"action": "compress", "summary": summary}
elif action == "checkpoint":
checkpoint_id = self.create_checkpoint(args if args else None)
return {"action": "checkpoint", "id": checkpoint_id}
elif action == "restore":
if not args:
return {"error": "Usage: /memory restore <checkpoint_id>"}
success = self.restore_checkpoint(args)
return {"action": "restore", "checkpoint": args, "success": success}
elif action == "checkpoints":
checkpoints = self.list_checkpoints()
return {"action": "checkpoints", "checkpoints": checkpoints}
elif action == "refresh":
self._load_all()
return {"action": "refresh", "message": "Memory reloaded from disk"}
elif action == "help":
return {
"action": "help",
"commands": {
"/memory show": "Display memory stats and recent items",
"/memory add <content>": "Add to long-term memory",
"/memory clear [short|all]": "Clear memory",
"/memory search <query>": "Search memories",
"/memory save <name>": "Save session",
"/memory resume <name>": "Resume session",
"/memory sessions": "List saved sessions",
"/memory compress": "Compress short-term memory",
"/memory checkpoint [name]": "Create checkpoint",
"/memory restore <id>": "Restore checkpoint",
"/memory checkpoints": "List checkpoints",
"/memory refresh": "Reload from disk"
}
}
else:
return {"error": f"Unknown action: {action}. Use /memory help for available commands."}
# Convenience function for simple usage
@@ -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)
@@ -176,6 +176,48 @@ TOOL_MAPPINGS = {
# 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
@@ -187,8 +229,8 @@ def __getattr__(name: str) -> Any:
module_path, class_name = TOOL_MAPPINGS[name]
# Return class itself (not instance) for ClaudeMemoryTool
if name == 'ClaudeMemoryTool':
# 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)
@@ -205,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,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,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,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.178"
version = "0.0.191"
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,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!")
Whitespace-only changes.
@@ -0,0 +1,252 @@
"""
Tests for FastContext high-level API.
"""
import pytest
import tempfile
import time
from pathlib import Path
@pytest.fixture
def temp_workspace():
"""Create a temporary workspace with test files."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create directory structure
src_dir = Path(tmpdir) / "src"
src_dir.mkdir()
(src_dir / "main.py").write_text("""
def main():
print("Hello, World!")
authenticate_user()
def authenticate_user():
# Authentication logic here
pass
class UserService:
def login(self, username, password):
return True
""")
(src_dir / "utils.py").write_text("""
import os
def get_config():
return {"debug": True}
def validate_input(data):
if not data:
raise ValueError("Empty data")
return True
""")
auth_dir = src_dir / "auth"
auth_dir.mkdir()
(auth_dir / "handler.py").write_text("""
from typing import Optional
class AuthHandler:
def __init__(self):
self.token = None
def authenticate(self, credentials: dict) -> bool:
return True
def logout(self):
self.token = None
""")
yield tmpdir
class TestFastContextInit:
"""Tests for FastContext initialization."""
def test_default_initialization(self, temp_workspace):
"""Test default initialization."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
assert fc.workspace_path == temp_workspace
assert fc.max_turns == 4
assert fc.max_parallel == 8
assert fc.model == "gpt-4o-mini"
assert fc.cache_enabled is True
def test_custom_initialization(self, temp_workspace):
"""Test custom initialization."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(
workspace_path=temp_workspace,
max_turns=3,
max_parallel=4,
model="claude-3-haiku-20240307",
cache_enabled=False
)
assert fc.max_turns == 3
assert fc.max_parallel == 4
assert fc.model == "claude-3-haiku-20240307"
assert fc.cache_enabled is False
def test_context_manager(self, temp_workspace):
"""Test context manager usage."""
from praisonaiagents.context.fast import FastContext
with FastContext(workspace_path=temp_workspace) as fc:
result = fc.search("def")
assert result.total_files > 0
class TestFastContextSearch:
"""Tests for FastContext search functionality."""
def test_search_basic(self, temp_workspace):
"""Test basic search."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
result = fc.search("authenticate")
assert result.total_files > 0
file_paths = [f.path for f in result.files]
assert any("main.py" in p for p in file_paths)
def test_search_with_patterns(self, temp_workspace):
"""Test search with include patterns."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
result = fc.search("def", include_patterns=["**/auth/*.py"])
for f in result.files:
assert "auth" in f.path
def test_search_no_results(self, temp_workspace):
"""Test search with no matches."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
result = fc.search("nonexistent_xyz_123")
assert result.total_files == 0
class TestFastContextCache:
"""Tests for FastContext caching."""
def test_cache_hit(self, temp_workspace):
"""Test cache hit on repeated query."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace, cache_enabled=True)
# First search
result1 = fc.search("authenticate")
assert result1.from_cache is False
# Second search (should hit cache)
result2 = fc.search("authenticate")
assert result2.from_cache is True
def test_cache_disabled(self, temp_workspace):
"""Test cache disabled."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace, cache_enabled=False)
result1 = fc.search("authenticate")
result2 = fc.search("authenticate")
assert result1.from_cache is False
assert result2.from_cache is False
def test_clear_cache(self, temp_workspace):
"""Test clearing cache."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace, cache_enabled=True)
fc.search("authenticate")
fc.clear_cache()
result = fc.search("authenticate")
assert result.from_cache is False
class TestFastContextFiles:
"""Tests for FastContext file operations."""
def test_search_files(self, temp_workspace):
"""Test searching for files."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
result = fc.search_files("**/*.py")
assert result.total_files >= 3
def test_read_context(self, temp_workspace):
"""Test reading file context."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
content = fc.read_context("src/main.py")
assert content is not None
assert "def main():" in content
def test_read_context_with_lines(self, temp_workspace):
"""Test reading specific lines."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
content = fc.read_context("src/main.py", start_line=1, end_line=5, context_lines=0)
assert content is not None
lines = content.split("\n")
assert len(lines) <= 5
class TestFastContextAgentIntegration:
"""Tests for FastContext agent integration."""
def test_get_context_for_agent(self, temp_workspace):
"""Test getting formatted context for agent."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
context = fc.get_context_for_agent("authenticate")
assert "Relevant Code Context" in context
assert "authenticate" in context.lower()
def test_get_context_for_agent_no_results(self, temp_workspace):
"""Test context for agent with no results."""
from praisonaiagents.context.fast import FastContext
fc = FastContext(workspace_path=temp_workspace)
context = fc.get_context_for_agent("nonexistent_xyz_123")
assert "No relevant code found" in context
class TestFastSearchFunction:
"""Tests for fast_search convenience function."""
def test_fast_search(self, temp_workspace):
"""Test fast_search function."""
from praisonaiagents.context.fast.fast_context import fast_search
result = fast_search("authenticate", workspace_path=temp_workspace)
assert result.total_files > 0
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,205 @@
"""
Tests for FastContextAgent class.
"""
import pytest
import tempfile
from pathlib import Path
@pytest.fixture
def temp_workspace():
"""Create a temporary workspace with test files."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create directory structure
src_dir = Path(tmpdir) / "src"
src_dir.mkdir()
(src_dir / "main.py").write_text("""
def main():
print("Hello, World!")
authenticate_user()
def authenticate_user():
# Authentication logic here
pass
class UserService:
def login(self, username, password):
return True
""")
(src_dir / "utils.py").write_text("""
import os
def get_config():
return {"debug": True}
def validate_input(data):
if not data:
raise ValueError("Empty data")
return True
""")
auth_dir = src_dir / "auth"
auth_dir.mkdir()
(auth_dir / "handler.py").write_text("""
from typing import Optional
class AuthHandler:
def __init__(self):
self.token = None
def authenticate(self, credentials: dict) -> bool:
return True
def logout(self):
self.token = None
""")
yield tmpdir
class TestFastContextAgentInit:
"""Tests for FastContextAgent initialization."""
def test_default_initialization(self, temp_workspace):
"""Test default initialization."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
assert agent.workspace_path == temp_workspace
assert agent.max_turns == 4
assert agent.max_parallel == 8
assert agent.model == "gpt-4o-mini"
def test_custom_initialization(self, temp_workspace):
"""Test custom initialization."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(
workspace_path=temp_workspace,
max_turns=3,
max_parallel=4,
model="claude-3-haiku-20240307"
)
assert agent.max_turns == 3
assert agent.max_parallel == 4
assert agent.model == "claude-3-haiku-20240307"
class TestFastContextAgentSearch:
"""Tests for FastContextAgent search functionality."""
def test_search_returns_result(self, temp_workspace):
"""Test that search returns FastContextResult."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
from praisonaiagents.context.fast.result import FastContextResult
agent = FastContextAgent(workspace_path=temp_workspace)
# Use simple search without LLM
result = agent.search_simple("authenticate")
assert isinstance(result, FastContextResult)
assert result.total_files > 0
def test_search_finds_relevant_files(self, temp_workspace):
"""Test that search finds relevant files."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
result = agent.search_simple("authenticate")
# Should find files containing "authenticate"
file_paths = [f.path for f in result.files]
assert any("main.py" in p for p in file_paths)
assert any("handler.py" in p for p in file_paths)
def test_search_with_file_pattern(self, temp_workspace):
"""Test search with file pattern filter."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
result = agent.search_simple("def", include_patterns=["**/auth/*.py"])
# Should only find in auth directory
for f in result.files:
assert "auth" in f.path
def test_search_records_timing(self, temp_workspace):
"""Test that search records timing information."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
result = agent.search_simple("def")
# search_time_ms >= 0 is valid (can be 0 for very fast searches)
assert result.search_time_ms >= 0
assert result.turns_used == 1
assert result.total_tool_calls == 1
def test_search_no_results(self, temp_workspace):
"""Test search with no matches."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
result = agent.search_simple("nonexistent_pattern_xyz123")
assert result.total_files == 0
class TestFastContextAgentTools:
"""Tests for FastContextAgent tool definitions."""
def test_get_tools(self, temp_workspace):
"""Test getting tool definitions."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
tools = agent.get_tools()
assert len(tools) == 4
tool_names = [t["name"] for t in tools]
assert "grep_search" in tool_names
assert "glob_search" in tool_names
assert "read_file" in tool_names
assert "list_directory" in tool_names
def test_execute_tool(self, temp_workspace):
"""Test executing a tool directly."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
result = agent.execute_tool("glob_search", pattern="**/*.py")
assert len(result) >= 3 # At least 3 .py files
class TestFastContextAgentLLMSearch:
"""Tests for FastContextAgent LLM-powered search."""
def test_search_falls_back_to_simple(self, temp_workspace):
"""Test that search falls back to simple when LLM not available."""
from praisonaiagents.context.fast.fast_context_agent import FastContextAgent
agent = FastContextAgent(workspace_path=temp_workspace)
# search() should fall back to search_simple when LLM is not available
result = agent.search("find authentication code")
# Should still return results (from simple search fallback)
assert result is not None
assert result.query == "find authentication code"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,269 @@
"""
Tests for Fast Context result classes.
"""
import pytest
import time
class TestLineRange:
"""Tests for LineRange dataclass."""
def test_line_range_creation(self):
"""Test basic LineRange creation."""
from praisonaiagents.context.fast.result import LineRange
lr = LineRange(start=10, end=20)
assert lr.start == 10
assert lr.end == 20
assert lr.content is None
assert lr.relevance_score == 1.0
def test_line_range_with_content(self):
"""Test LineRange with content."""
from praisonaiagents.context.fast.result import LineRange
lr = LineRange(start=1, end=3, content="line1\nline2\nline3")
assert lr.content == "line1\nline2\nline3"
def test_line_range_invalid_start(self):
"""Test LineRange corrects invalid start."""
from praisonaiagents.context.fast.result import LineRange
lr = LineRange(start=0, end=5)
assert lr.start == 1 # Corrected to 1
def test_line_range_invalid_end(self):
"""Test LineRange corrects invalid end."""
from praisonaiagents.context.fast.result import LineRange
lr = LineRange(start=10, end=5)
assert lr.end == 10 # Corrected to start
def test_line_count(self):
"""Test line_count property."""
from praisonaiagents.context.fast.result import LineRange
lr = LineRange(start=10, end=20)
assert lr.line_count == 11
def test_overlaps_true(self):
"""Test overlaps returns True for overlapping ranges."""
from praisonaiagents.context.fast.result import LineRange
lr1 = LineRange(start=10, end=20)
lr2 = LineRange(start=15, end=25)
assert lr1.overlaps(lr2) is True
assert lr2.overlaps(lr1) is True
def test_overlaps_false(self):
"""Test overlaps returns False for non-overlapping ranges."""
from praisonaiagents.context.fast.result import LineRange
lr1 = LineRange(start=10, end=20)
lr2 = LineRange(start=25, end=30)
assert lr1.overlaps(lr2) is False
def test_overlaps_adjacent(self):
"""Test overlaps for adjacent ranges."""
from praisonaiagents.context.fast.result import LineRange
lr1 = LineRange(start=10, end=20)
lr2 = LineRange(start=21, end=30)
assert lr1.overlaps(lr2) is False
def test_merge(self):
"""Test merging overlapping ranges."""
from praisonaiagents.context.fast.result import LineRange
lr1 = LineRange(start=10, end=20, relevance_score=0.8)
lr2 = LineRange(start=15, end=25, relevance_score=0.9)
merged = lr1.merge(lr2)
assert merged.start == 10
assert merged.end == 25
assert merged.relevance_score == 0.9 # Max of both
def test_merge_non_overlapping_raises(self):
"""Test merge raises for non-overlapping ranges."""
from praisonaiagents.context.fast.result import LineRange
lr1 = LineRange(start=10, end=20)
lr2 = LineRange(start=25, end=30)
with pytest.raises(ValueError):
lr1.merge(lr2)
class TestFileMatch:
"""Tests for FileMatch dataclass."""
def test_file_match_creation(self):
"""Test basic FileMatch creation."""
from praisonaiagents.context.fast.result import FileMatch
fm = FileMatch(path="/path/to/file.py")
assert fm.path == "/path/to/file.py"
assert fm.line_ranges == []
assert fm.relevance_score == 1.0
assert fm.match_count == 0
def test_add_line_range(self):
"""Test adding line ranges."""
from praisonaiagents.context.fast.result import FileMatch, LineRange
fm = FileMatch(path="/path/to/file.py")
fm.add_line_range(LineRange(start=10, end=20))
assert len(fm.line_ranges) == 1
assert fm.line_ranges[0].start == 10
def test_add_overlapping_line_ranges(self):
"""Test adding overlapping line ranges merges them."""
from praisonaiagents.context.fast.result import FileMatch, LineRange
fm = FileMatch(path="/path/to/file.py")
fm.add_line_range(LineRange(start=10, end=20))
fm.add_line_range(LineRange(start=15, end=25))
assert len(fm.line_ranges) == 1
assert fm.line_ranges[0].start == 10
assert fm.line_ranges[0].end == 25
def test_add_non_overlapping_line_ranges(self):
"""Test adding non-overlapping line ranges keeps them separate."""
from praisonaiagents.context.fast.result import FileMatch, LineRange
fm = FileMatch(path="/path/to/file.py")
fm.add_line_range(LineRange(start=10, end=20))
fm.add_line_range(LineRange(start=30, end=40))
assert len(fm.line_ranges) == 2
def test_total_lines(self):
"""Test total_lines property."""
from praisonaiagents.context.fast.result import FileMatch, LineRange
fm = FileMatch(path="/path/to/file.py")
fm.add_line_range(LineRange(start=10, end=20)) # 11 lines
fm.add_line_range(LineRange(start=30, end=35)) # 6 lines
assert fm.total_lines == 17
class TestFastContextResult:
"""Tests for FastContextResult dataclass."""
def test_result_creation(self):
"""Test basic result creation."""
from praisonaiagents.context.fast.result import FastContextResult
result = FastContextResult(query="find authentication")
assert result.query == "find authentication"
assert result.files == []
assert result.search_time_ms == 0
def test_timer(self):
"""Test timer functionality."""
from praisonaiagents.context.fast.result import FastContextResult
result = FastContextResult()
result.start_timer()
time.sleep(0.1) # 100ms
result.stop_timer()
assert result.search_time_ms >= 100
assert result.search_time_ms < 200 # Should be close to 100ms
def test_add_file(self):
"""Test adding files."""
from praisonaiagents.context.fast.result import FastContextResult, FileMatch
result = FastContextResult()
result.add_file(FileMatch(path="/path/to/file.py"))
assert result.total_files == 1
def test_add_duplicate_file_merges(self):
"""Test adding duplicate file path merges them."""
from praisonaiagents.context.fast.result import (
FastContextResult, FileMatch, LineRange
)
result = FastContextResult()
fm1 = FileMatch(path="/path/to/file.py", relevance_score=0.8)
fm1.add_line_range(LineRange(start=10, end=20))
result.add_file(fm1)
fm2 = FileMatch(path="/path/to/file.py", relevance_score=0.9)
fm2.add_line_range(LineRange(start=30, end=40))
result.add_file(fm2)
assert result.total_files == 1
assert len(result.files[0].line_ranges) == 2
assert result.files[0].relevance_score == 0.9
def test_sort_by_relevance(self):
"""Test sorting files by relevance."""
from praisonaiagents.context.fast.result import FastContextResult, FileMatch
result = FastContextResult()
result.add_file(FileMatch(path="/low.py", relevance_score=0.3))
result.add_file(FileMatch(path="/high.py", relevance_score=0.9))
result.add_file(FileMatch(path="/mid.py", relevance_score=0.6))
result.sort_by_relevance()
assert result.files[0].path == "/high.py"
assert result.files[1].path == "/mid.py"
assert result.files[2].path == "/low.py"
def test_to_context_string(self):
"""Test converting result to context string."""
from praisonaiagents.context.fast.result import (
FastContextResult, FileMatch, LineRange
)
result = FastContextResult()
fm = FileMatch(path="/path/to/file.py")
fm.add_line_range(LineRange(start=10, end=20))
result.add_file(fm)
context = result.to_context_string()
assert "1 relevant file" in context
assert "/path/to/file.py" in context
assert "Lines 10-20" in context
def test_to_context_string_empty(self):
"""Test context string for empty result."""
from praisonaiagents.context.fast.result import FastContextResult
result = FastContextResult()
context = result.to_context_string()
assert "No relevant files found" in context
def test_to_dict(self):
"""Test converting result to dictionary."""
from praisonaiagents.context.fast.result import (
FastContextResult, FileMatch, LineRange
)
result = FastContextResult(query="test", search_time_ms=100, turns_used=2)
fm = FileMatch(path="/file.py", relevance_score=0.8)
fm.add_line_range(LineRange(start=10, end=20))
result.add_file(fm)
d = result.to_dict()
assert d["query"] == "test"
assert d["search_time_ms"] == 100
assert d["turns_used"] == 2
assert len(d["files"]) == 1
assert d["files"][0]["path"] == "/file.py"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,261 @@
"""
Tests for File Indexer.
"""
import pytest
import tempfile
import os
from pathlib import Path
@pytest.fixture
def temp_workspace():
"""Create a temporary workspace with test files."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create directory structure
src_dir = Path(tmpdir) / "src"
src_dir.mkdir()
(src_dir / "main.py").write_text("def main(): pass")
(src_dir / "utils.py").write_text("def util(): pass")
(src_dir / "config.json").write_text('{"key": "value"}')
auth_dir = src_dir / "auth"
auth_dir.mkdir()
(auth_dir / "handler.py").write_text("class Handler: pass")
# Create .gitignore
(Path(tmpdir) / ".gitignore").write_text("__pycache__/\n*.pyc\n")
# Create ignored directory
cache_dir = Path(tmpdir) / "__pycache__"
cache_dir.mkdir()
(cache_dir / "cached.pyc").write_text("cached")
yield tmpdir
class TestFileIndexerInit:
"""Tests for FileIndexer initialization."""
def test_default_initialization(self, temp_workspace):
"""Test default initialization."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
assert indexer.workspace_path == temp_workspace
assert indexer.respect_gitignore is True
assert indexer.total_files == 0
def test_custom_extensions(self, temp_workspace):
"""Test custom extensions."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(
workspace_path=temp_workspace,
extensions={".py"}
)
assert indexer.extensions == {".py"}
class TestFileIndexerIndex:
"""Tests for FileIndexer indexing."""
def test_index_files(self, temp_workspace):
"""Test indexing files."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
count = indexer.index()
assert count >= 3 # At least main.py, utils.py, handler.py
assert indexer.total_files >= 3
def test_index_respects_gitignore(self, temp_workspace):
"""Test that indexing respects .gitignore."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace, respect_gitignore=True)
indexer.index()
# Should not include __pycache__ files
for path in indexer.files:
assert "__pycache__" not in path
assert not path.endswith(".pyc")
def test_index_by_extension(self, temp_workspace):
"""Test indexing organizes by extension."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
assert ".py" in indexer.by_extension
assert len(indexer.by_extension[".py"]) >= 3
def test_index_with_extension_filter(self, temp_workspace):
"""Test indexing with extension filter."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace, extensions={".py"})
indexer.index()
# Should only have .py files
for path in indexer.files:
assert path.endswith(".py")
class TestFileIndexerFind:
"""Tests for FileIndexer find methods."""
def test_find_by_pattern(self, temp_workspace):
"""Test finding files by pattern."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
results = indexer.find_by_pattern("**/*.py")
assert len(results) >= 3
for info in results:
assert info.path.endswith(".py")
def test_find_by_extension(self, temp_workspace):
"""Test finding files by extension."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
results = indexer.find_by_extension(".py")
assert len(results) >= 3
def test_find_by_extension_without_dot(self, temp_workspace):
"""Test finding files by extension without dot."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
results = indexer.find_by_extension("py")
assert len(results) >= 3
def test_find_by_name_exact(self, temp_workspace):
"""Test finding files by exact name."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
results = indexer.find_by_name("main.py", exact=True)
assert len(results) == 1
assert results[0].path.endswith("main.py")
def test_find_by_name_partial(self, temp_workspace):
"""Test finding files by partial name."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
results = indexer.find_by_name("handler")
assert len(results) >= 1
assert any("handler" in r.path for r in results)
def test_get_file(self, temp_workspace):
"""Test getting file by path."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
# Find a file first
results = indexer.find_by_name("main.py", exact=True)
assert len(results) > 0
# Get by path
file_info = indexer.get_file(results[0].path)
assert file_info is not None
assert file_info.path == results[0].path
class TestFileIndexerStats:
"""Tests for FileIndexer statistics."""
def test_get_stats(self, temp_workspace):
"""Test getting index statistics."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
stats = indexer.get_stats()
assert stats["total_files"] >= 3
assert stats["total_size_bytes"] > 0
assert ".py" in stats["extensions"]
def test_is_stale(self, temp_workspace):
"""Test stale check."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
# Before indexing, should be stale
assert indexer.is_stale() is True
indexer.index()
# After indexing, should not be stale
assert indexer.is_stale(max_age=60.0) is False
class TestFileIndexerPersistence:
"""Tests for FileIndexer save/load."""
def test_save_and_load(self, temp_workspace):
"""Test saving and loading index."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
# Save index
index_path = os.path.join(temp_workspace, ".fast_context_index.json")
indexer.save_index(index_path)
assert os.path.exists(index_path)
# Load into new indexer
indexer2 = FileIndexer(workspace_path=temp_workspace)
loaded = indexer2.load_index(index_path)
assert loaded is True
assert indexer2.total_files == indexer.total_files
def test_load_wrong_workspace(self, temp_workspace):
"""Test loading index from wrong workspace."""
from praisonaiagents.context.fast.indexer import FileIndexer
indexer = FileIndexer(workspace_path=temp_workspace)
indexer.index()
index_path = os.path.join(temp_workspace, ".fast_context_index.json")
indexer.save_index(index_path)
# Try to load with different workspace
indexer2 = FileIndexer(workspace_path="/different/path")
loaded = indexer2.load_index(index_path)
assert loaded is False
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,205 @@
"""
Tests for Fast Context parallel executor.
"""
import pytest
import asyncio
import time
import tempfile
from pathlib import Path
@pytest.fixture
def temp_workspace():
"""Create a temporary workspace with test files."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create some test files
for i in range(5):
(Path(tmpdir) / f"file{i}.py").write_text(f"# File {i}\ndef func{i}(): pass")
yield tmpdir
class TestParallelExecutor:
"""Tests for ParallelExecutor class."""
@pytest.mark.asyncio
async def test_execute_single_task(self, temp_workspace):
"""Test executing a single task."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
from praisonaiagents.context.fast.search_tools import glob_search
executor = ParallelExecutor(max_parallel=8)
tasks = [
{"tool": "glob_search", "args": {"search_path": temp_workspace, "pattern": "*.py"}}
]
results = await executor.execute(tasks)
assert len(results) == 1
assert len(results[0]) == 5 # 5 .py files
@pytest.mark.asyncio
async def test_execute_multiple_tasks_parallel(self, temp_workspace):
"""Test executing multiple tasks in parallel."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
executor = ParallelExecutor(max_parallel=8)
tasks = [
{"tool": "glob_search", "args": {"search_path": temp_workspace, "pattern": "*.py"}},
{"tool": "grep_search", "args": {"search_path": temp_workspace, "pattern": "def"}},
{"tool": "list_directory", "args": {"dir_path": temp_workspace}},
]
start = time.perf_counter()
results = await executor.execute(tasks)
elapsed = time.perf_counter() - start
assert len(results) == 3
# All tasks should complete (parallel execution)
assert all(r is not None for r in results)
@pytest.mark.asyncio
async def test_execute_with_timeout(self, temp_workspace):
"""Test execution with timeout."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
executor = ParallelExecutor(max_parallel=8, timeout=0.001) # Very short timeout
# This should timeout for slow operations
tasks = [
{"tool": "glob_search", "args": {"search_path": temp_workspace, "pattern": "**/*.py"}}
]
results = await executor.execute(tasks)
# Should return results (either success or timeout error)
assert len(results) == 1
@pytest.mark.asyncio
async def test_execute_respects_max_parallel(self, temp_workspace):
"""Test that max_parallel limit is respected."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
executor = ParallelExecutor(max_parallel=2)
# Create more tasks than max_parallel
tasks = [
{"tool": "glob_search", "args": {"search_path": temp_workspace, "pattern": "*.py"}}
for _ in range(5)
]
results = await executor.execute(tasks)
assert len(results) == 5
# All should complete successfully
assert all(len(r) == 5 for r in results)
@pytest.mark.asyncio
async def test_execute_handles_errors(self, temp_workspace):
"""Test that errors in one task don't affect others."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
executor = ParallelExecutor(max_parallel=8)
tasks = [
{"tool": "glob_search", "args": {"search_path": temp_workspace, "pattern": "*.py"}},
{"tool": "read_file", "args": {"filepath": "/nonexistent/file.py"}}, # Will fail
{"tool": "list_directory", "args": {"dir_path": temp_workspace}},
]
results = await executor.execute(tasks)
assert len(results) == 3
# First and third should succeed
assert len(results[0]) == 5
assert results[1]["success"] is False # Error result
assert results[2]["success"] is True
@pytest.mark.asyncio
async def test_execute_unknown_tool(self, temp_workspace):
"""Test handling of unknown tool."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
executor = ParallelExecutor(max_parallel=8)
tasks = [
{"tool": "unknown_tool", "args": {}}
]
results = await executor.execute(tasks)
assert len(results) == 1
assert "error" in results[0]
@pytest.mark.asyncio
async def test_execute_empty_tasks(self):
"""Test executing empty task list."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
executor = ParallelExecutor(max_parallel=8)
results = await executor.execute([])
assert results == []
def test_execute_sync(self, temp_workspace):
"""Test synchronous execution wrapper."""
from praisonaiagents.context.fast.parallel_executor import ParallelExecutor
executor = ParallelExecutor(max_parallel=8)
tasks = [
{"tool": "glob_search", "args": {"search_path": temp_workspace, "pattern": "*.py"}}
]
results = executor.execute_sync(tasks)
assert len(results) == 1
assert len(results[0]) == 5
class TestToolCallBatch:
"""Tests for ToolCallBatch class."""
def test_batch_creation(self):
"""Test creating a tool call batch."""
from praisonaiagents.context.fast.parallel_executor import ToolCallBatch
batch = ToolCallBatch()
batch.add("grep_search", search_path="/path", pattern="test")
batch.add("glob_search", search_path="/path", pattern="*.py")
assert len(batch.tasks) == 2
def test_batch_max_size(self):
"""Test batch respects max size."""
from praisonaiagents.context.fast.parallel_executor import ToolCallBatch
batch = ToolCallBatch(max_size=2)
batch.add("grep_search", search_path="/path", pattern="test1")
batch.add("grep_search", search_path="/path", pattern="test2")
# Should not add beyond max_size
added = batch.add("grep_search", search_path="/path", pattern="test3")
assert added is False
assert len(batch.tasks) == 2
def test_batch_is_full(self):
"""Test is_full property."""
from praisonaiagents.context.fast.parallel_executor import ToolCallBatch
batch = ToolCallBatch(max_size=2)
assert batch.is_full is False
batch.add("grep_search", search_path="/path", pattern="test1")
assert batch.is_full is False
batch.add("grep_search", search_path="/path", pattern="test2")
assert batch.is_full is True
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,336 @@
"""
Tests for Fast Context search tools.
"""
import pytest
import tempfile
import os
from pathlib import Path
@pytest.fixture
def temp_workspace():
"""Create a temporary workspace with test files."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create directory structure
src_dir = Path(tmpdir) / "src"
src_dir.mkdir()
# Create Python files
(src_dir / "main.py").write_text("""
def main():
print("Hello, World!")
authenticate_user()
def authenticate_user():
# Authentication logic here
pass
class UserService:
def login(self, username, password):
return True
""")
(src_dir / "utils.py").write_text("""
import os
def get_config():
return {"debug": True}
def validate_input(data):
if not data:
raise ValueError("Empty data")
return True
""")
# Create subdirectory
auth_dir = src_dir / "auth"
auth_dir.mkdir()
(auth_dir / "handler.py").write_text("""
from typing import Optional
class AuthHandler:
def __init__(self):
self.token = None
def authenticate(self, credentials: dict) -> bool:
# Authenticate user
return True
def logout(self):
self.token = None
""")
# Create a .gitignore
(Path(tmpdir) / ".gitignore").write_text("""
__pycache__/
*.pyc
.env
""")
# Create ignored directory
cache_dir = Path(tmpdir) / "__pycache__"
cache_dir.mkdir()
(cache_dir / "cached.pyc").write_text("cached content")
yield tmpdir
class TestGrepSearch:
"""Tests for grep_search function."""
def test_grep_simple_pattern(self, temp_workspace):
"""Test grep with simple pattern."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "authenticate")
assert len(results) > 0
# Should find in main.py and handler.py
paths = [r["path"] for r in results]
assert any("main.py" in p for p in paths)
assert any("handler.py" in p for p in paths)
def test_grep_regex_pattern(self, temp_workspace):
"""Test grep with regex pattern."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, r"def \w+\(", is_regex=True)
assert len(results) > 0
# Should find function definitions
def test_grep_case_insensitive(self, temp_workspace):
"""Test grep case insensitive search."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "AUTHENTICATE", case_sensitive=False)
assert len(results) > 0
def test_grep_case_sensitive(self, temp_workspace):
"""Test grep case sensitive search."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "AUTHENTICATE", case_sensitive=True)
assert len(results) == 0 # No uppercase AUTHENTICATE exists
def test_grep_with_file_pattern(self, temp_workspace):
"""Test grep with file pattern filter."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "def", include_patterns=["*.py"])
assert len(results) > 0
# All results should be .py files
for r in results:
assert r["path"].endswith(".py")
def test_grep_with_exclude_pattern(self, temp_workspace):
"""Test grep with exclude pattern."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "def", exclude_patterns=["**/auth/*"])
# Should not include auth directory files
for r in results:
assert "auth" not in r["path"]
def test_grep_max_results(self, temp_workspace):
"""Test grep with max results limit."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "def", max_results=2)
assert len(results) <= 2
def test_grep_no_matches(self, temp_workspace):
"""Test grep with no matches."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "nonexistent_pattern_xyz")
assert len(results) == 0
def test_grep_respects_gitignore(self, temp_workspace):
"""Test grep respects .gitignore patterns."""
from praisonaiagents.context.fast.search_tools import grep_search
results = grep_search(temp_workspace, "cached", respect_gitignore=True)
# Should not find content in __pycache__
for r in results:
assert "__pycache__" not in r["path"]
class TestGlobSearch:
"""Tests for glob_search function."""
def test_glob_all_python_files(self, temp_workspace):
"""Test glob to find all Python files."""
from praisonaiagents.context.fast.search_tools import glob_search
results = glob_search(temp_workspace, "**/*.py")
assert len(results) >= 3 # main.py, utils.py, handler.py
for r in results:
assert r["path"].endswith(".py")
def test_glob_specific_directory(self, temp_workspace):
"""Test glob in specific directory."""
from praisonaiagents.context.fast.search_tools import glob_search
results = glob_search(temp_workspace, "src/auth/*.py")
assert len(results) == 1
assert "handler.py" in results[0]["path"]
def test_glob_by_name(self, temp_workspace):
"""Test glob by filename."""
from praisonaiagents.context.fast.search_tools import glob_search
results = glob_search(temp_workspace, "**/main.py")
assert len(results) == 1
assert "main.py" in results[0]["path"]
def test_glob_no_matches(self, temp_workspace):
"""Test glob with no matches."""
from praisonaiagents.context.fast.search_tools import glob_search
results = glob_search(temp_workspace, "**/*.java")
assert len(results) == 0
def test_glob_max_results(self, temp_workspace):
"""Test glob with max results."""
from praisonaiagents.context.fast.search_tools import glob_search
results = glob_search(temp_workspace, "**/*.py", max_results=2)
assert len(results) <= 2
def test_glob_respects_gitignore(self, temp_workspace):
"""Test glob respects .gitignore."""
from praisonaiagents.context.fast.search_tools import glob_search
results = glob_search(temp_workspace, "**/*.pyc", respect_gitignore=True)
assert len(results) == 0 # .pyc is in .gitignore
class TestReadFile:
"""Tests for read_file function."""
def test_read_entire_file(self, temp_workspace):
"""Test reading entire file."""
from praisonaiagents.context.fast.search_tools import read_file
result = read_file(os.path.join(temp_workspace, "src", "main.py"))
assert result["success"] is True
assert "def main():" in result["content"]
assert result["total_lines"] > 0
def test_read_file_with_line_range(self, temp_workspace):
"""Test reading specific line range."""
from praisonaiagents.context.fast.search_tools import read_file
result = read_file(
os.path.join(temp_workspace, "src", "main.py"),
start_line=1,
end_line=3
)
assert result["success"] is True
assert result["start_line"] == 1
assert result["end_line"] == 3
lines = result["content"].split("\n")
assert len(lines) <= 3
def test_read_nonexistent_file(self, temp_workspace):
"""Test reading nonexistent file."""
from praisonaiagents.context.fast.search_tools import read_file
result = read_file(os.path.join(temp_workspace, "nonexistent.py"))
assert result["success"] is False
assert "error" in result
def test_read_file_with_context(self, temp_workspace):
"""Test reading with context lines."""
from praisonaiagents.context.fast.search_tools import read_file
result = read_file(
os.path.join(temp_workspace, "src", "main.py"),
start_line=5,
end_line=5,
context_lines=2
)
assert result["success"] is True
# Should include 2 lines before and after line 5
class TestListDirectory:
"""Tests for list_directory function."""
def test_list_directory(self, temp_workspace):
"""Test listing directory contents."""
from praisonaiagents.context.fast.search_tools import list_directory
result = list_directory(temp_workspace)
assert result["success"] is True
assert len(result["entries"]) > 0
# Should contain src directory
names = [e["name"] for e in result["entries"]]
assert "src" in names
def test_list_directory_recursive(self, temp_workspace):
"""Test listing directory recursively."""
from praisonaiagents.context.fast.search_tools import list_directory
result = list_directory(temp_workspace, recursive=True)
assert result["success"] is True
# Should find files in subdirectories
paths = [e["path"] for e in result["entries"]]
assert any("main.py" in p for p in paths)
assert any("handler.py" in p for p in paths)
def test_list_directory_max_depth(self, temp_workspace):
"""Test listing with max depth."""
from praisonaiagents.context.fast.search_tools import list_directory
result = list_directory(temp_workspace, recursive=True, max_depth=1)
assert result["success"] is True
# Should not find files in auth subdirectory (depth 2)
paths = [e["path"] for e in result["entries"]]
assert not any("handler.py" in p for p in paths)
def test_list_nonexistent_directory(self, temp_workspace):
"""Test listing nonexistent directory."""
from praisonaiagents.context.fast.search_tools import list_directory
result = list_directory(os.path.join(temp_workspace, "nonexistent"))
assert result["success"] is False
def test_list_directory_respects_gitignore(self, temp_workspace):
"""Test listing respects .gitignore."""
from praisonaiagents.context.fast.search_tools import list_directory
result = list_directory(temp_workspace, recursive=True, respect_gitignore=True)
# Should not include __pycache__
paths = [e["path"] for e in result["entries"]]
assert not any("__pycache__" in p for p in paths)
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -651,3 +651,179 @@ class TestMemoryIntegrationWithOpenAI:
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,323 @@
"""
Tests for MCP HTTP Stream Transport features.
Tests the HTTPStreamTransport implementation including:
- Protocol version header (Mcp-Protocol-Version)
- Session management (Mcp-Session-Id)
- Resumability (Last-Event-ID, retry delay)
- Session termination (HTTP DELETE)
- SSE event parsing
"""
import pytest
class TestHTTPStreamTransportInit:
"""Test HTTPStreamTransport initialization."""
def test_transport_has_protocol_version_header(self):
"""Test that transport includes Mcp-Protocol-Version header."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
assert 'Mcp-Protocol-Version' in transport.headers
assert transport.headers['Mcp-Protocol-Version'] == '2025-03-26'
def test_transport_default_protocol_version(self):
"""Test default protocol version for backward compatibility."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
assert HTTPStreamTransport.DEFAULT_PROTOCOL_VERSION == '2025-03-26'
def test_transport_custom_protocol_version(self):
"""Test custom protocol version via options."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport(
'https://api.example.com/mcp',
options={'protocol_version': '2025-11-25'}
)
assert transport.headers['Mcp-Protocol-Version'] == '2025-11-25'
def test_transport_has_accept_header(self):
"""Test that transport includes correct Accept header."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
assert 'Accept' in transport.headers
assert 'application/json' in transport.headers['Accept']
assert 'text/event-stream' in transport.headers['Accept']
def test_transport_has_session_id_when_provided(self):
"""Test that transport includes session ID when provided."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport(
'https://api.example.com/mcp',
session_id='test-session-123'
)
assert 'Mcp-Session-Id' in transport.headers
assert transport.headers['Mcp-Session-Id'] == 'test-session-123'
def test_transport_no_session_id_when_not_provided(self):
"""Test that transport doesn't include session ID when not provided."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
assert 'Mcp-Session-Id' not in transport.headers
class TestHTTPStreamTransportResumability:
"""Test HTTPStreamTransport resumability features."""
def test_transport_has_last_event_id_tracking(self):
"""Test that transport tracks last event ID."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
assert hasattr(transport, 'last_event_id')
assert transport.last_event_id is None
def test_transport_has_retry_delay_tracking(self):
"""Test that transport tracks retry delay."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
assert hasattr(transport, '_retry_delay_ms')
assert transport._retry_delay_ms == 3000 # Default 3 seconds
def test_transport_updates_last_event_id(self):
"""Test that last_event_id can be updated."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
transport.last_event_id = 'event-123'
assert transport.last_event_id == 'event-123'
def test_transport_updates_retry_delay(self):
"""Test that retry delay can be updated."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
transport._retry_delay_ms = 5000
assert transport._retry_delay_ms == 5000
class TestHTTPStreamTransportSessionTermination:
"""Test HTTPStreamTransport session termination."""
def test_transport_has_terminate_session_method(self):
"""Test that transport has terminate_session method."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
assert hasattr(transport, 'terminate_session')
assert callable(transport.terminate_session)
@pytest.mark.asyncio
async def test_terminate_session_returns_false_without_session(self):
"""Test that terminate_session returns False without session."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport('https://api.example.com/mcp')
# No session, should return False
result = await transport.terminate_session()
assert result is False
@pytest.mark.asyncio
async def test_terminate_session_returns_false_without_client_session(self):
"""Test that terminate_session returns False without aiohttp session."""
from praisonaiagents.mcp.mcp_http_stream import HTTPStreamTransport
transport = HTTPStreamTransport(
'https://api.example.com/mcp',
session_id='test-session'
)
# _session is None (not initialized)
result = await transport.terminate_session()
assert result is False
class TestHTTPStreamTransportAutoDetection:
"""Test transport auto-detection in MCP class."""
def test_http_url_uses_http_stream_transport(self):
"""Test that HTTP URLs use HTTPStreamTransport."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type('http://localhost:8080/mcp') == 'http_stream'
assert get_transport_type('https://api.example.com/mcp') == 'http_stream'
def test_sse_url_uses_sse_transport(self):
"""Test that SSE URLs use SSE transport."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type('http://localhost:8080/sse') == 'sse'
assert get_transport_type('https://api.example.com/sse') == 'sse'
def test_websocket_url_uses_websocket_transport(self):
"""Test that WebSocket URLs use WebSocket transport."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type('ws://localhost:8080/mcp') == 'websocket'
assert get_transport_type('wss://api.example.com/mcp') == 'websocket'
def test_command_string_uses_stdio_transport(self):
"""Test that command strings use stdio transport."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type('npx @modelcontextprotocol/server-memory') == 'stdio'
assert get_transport_type('python3 server.py') == 'stdio'
class TestSSEEventParsing:
"""Test SSE event parsing for resumability."""
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\": true}"
result = parse_sse_event(event_text)
assert result['id'] == 'event-123'
assert result['data'] == '{"test": true}'
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: {\"test\": true}"
result = parse_sse_event(event_text)
assert result['retry'] == 5000
assert result['data'] == '{"test": true}'
def test_parse_sse_event_with_all_fields(self):
"""Test parsing SSE event with all fields."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "id: event-456\nretry: 3000\ndata: {\"jsonrpc\": \"2.0\"}"
result = parse_sse_event(event_text)
assert result['id'] == 'event-456'
assert result['retry'] == 3000
assert result['data'] == '{"jsonrpc": "2.0"}'
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"
result = parse_sse_event(event_text)
assert result['data'] == 'line1\nline2'
class TestResumabilityManager:
"""Test ResumabilityManager for SSE stream recovery."""
def test_manager_tracks_event_ids(self):
"""Test that manager tracks event IDs per stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id('stream1', 'event-100')
manager.set_last_event_id('stream2', 'event-200')
assert manager.get_last_event_id('stream1') == 'event-100'
assert manager.get_last_event_id('stream2') == 'event-200'
def test_manager_returns_none_for_unknown_stream(self):
"""Test that manager returns None for unknown stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
assert manager.get_last_event_id('unknown') is None
def test_manager_generates_resume_headers(self):
"""Test that manager generates Last-Event-ID header."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id('stream1', 'event-123')
headers = manager.get_resume_headers('stream1')
assert 'Last-Event-ID' in headers
assert headers['Last-Event-ID'] == 'event-123'
def test_manager_tracks_retry_delay(self):
"""Test that manager tracks retry delay."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_retry_delay(5000)
assert manager.get_retry_delay() == 5000
def test_manager_default_retry_delay(self):
"""Test default retry delay."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
# Default should be 3000ms
assert manager.get_retry_delay() == 3000
class TestEventStore:
"""Test EventStore for message replay."""
def test_store_events(self):
"""Test storing events."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore(max_events=100)
store.store('stream1', 'event-1', {'data': 'test1'})
store.store('stream1', 'event-2', {'data': 'test2'})
events = store.get_all_events('stream1')
assert len(events) == 2
def test_get_events_after_id(self):
"""Test getting events after specific ID."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore(max_events=100)
store.store('stream1', 'event-1', {'data': 'test1'})
store.store('stream1', 'event-2', {'data': 'test2'})
store.store('stream1', 'event-3', {'data': 'test3'})
events = store.get_events_after('stream1', 'event-1')
assert len(events) == 2
assert events[0]['data'] == 'test2'
assert events[1]['data'] == 'test3'
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('stream1', f'event-{i}', {'data': f'test{i}'})
events = store.get_all_events('stream1')
assert len(events) == 3
if __name__ == "__main__":
pytest.main([__file__, "-v"])
Loaded 100 of 146 files, more files were not shown because too many files have changed in this diff. Show more