mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 17:33:33 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dfb77c39f9 | ||
|
|
7e8f9b0574 | ||
|
|
19bcdcd000 | ||
|
|
8f5160f827 | ||
|
|
3574a4d2c0 | ||
|
|
b552843fa9 | ||
|
|
0461a133df | ||
|
|
b33857de33 | ||
|
|
433c643f6c | ||
|
|
214119a063 | ||
|
|
989a9f6fc5 | ||
|
|
cf276a5640 | ||
|
|
b5134ffb2f | ||
|
|
5664036712 | ||
|
|
6c2ca64cbe | ||
|
|
08c742287c | ||
|
|
4ee32b8eb2 | ||
|
|
35f07e5653 | ||
|
|
6481aec376 | ||
|
|
98c846ef78 | ||
|
|
6e4a9faaf8 | ||
|
|
05d62db355 | ||
|
|
375c8c9826 | ||
|
|
c6977c8221 | ||
|
|
d407c83cf3 | ||
|
|
214add68d9 | ||
|
|
6d2cb9676d | ||
|
|
c519cb1379 | ||
|
|
356dc4909a | ||
|
|
dc278ad07b | ||
|
|
5cc53eb4bb | ||
|
|
b72576d086 | ||
|
|
30314b1573 | ||
|
|
f721bd5e9a | ||
|
|
d00ea3e3c3 | ||
|
|
67dbd1007a | ||
|
|
cf2a5a135f | ||
|
|
bc23d3a104 | ||
|
|
bb4e1d8b9e | ||
|
|
69740baa7c | ||
|
|
5035978a29 | ||
|
|
4ddb043324 | ||
|
|
337014f4db | ||
|
|
35c371b25e | ||
|
|
bba75d9fdd | ||
|
|
bf8cd488da | ||
|
|
28905db46d | ||
|
|
3e7aa44dfa | ||
|
|
2f7a65ab92 | ||
|
|
e6e2379b02 | ||
|
|
f31f9a04aa | ||
|
|
5a8adf295e | ||
|
|
8dc243fd4f | ||
|
|
548b4afb88 | ||
|
|
ffa938b3df | ||
|
|
8a9beed264 | ||
|
|
daca5a1762 | ||
|
|
b183ac5cbd | ||
|
|
5f57e1d2f1 | ||
|
|
5114165944 | ||
|
|
389442ea04 | ||
|
|
6d7cc9c3ba | ||
|
|
690e4497cc | ||
|
|
041f5f0b08 | ||
|
|
bc0d9a6213 | ||
|
|
a4427fa135 | ||
|
|
51e8c65c8c | ||
|
|
285ed672c5 | ||
|
|
c03e51e337 | ||
|
|
f3ba05f6aa | ||
|
|
9cea3edcd6 | ||
|
|
02a3bd3597 | ||
|
|
128c77926d | ||
|
|
3aa1b825e4 | ||
|
|
b4ee30d837 | ||
|
|
825648aead | ||
|
|
1532c31e72 | ||
|
|
b313fefd86 | ||
|
|
435e1afb55 | ||
|
|
4cb3c0565a | ||
|
|
a7634b66dc | ||
|
|
3f0eaccbcb | ||
|
|
47713d933a | ||
|
|
c82bfccd1d | ||
|
|
ef7d1264e2 | ||
|
|
2799304270 |
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -27,6 +27,9 @@ jobs:
|
||||
comprehensive-test:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
OPENAI_MODEL_NAME: gpt-4o-mini
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ["3.11"]
|
||||
@@ -55,11 +58,13 @@ jobs:
|
||||
# Install knowledge dependencies from praisonai-agents
|
||||
uv pip install --system "praisonaiagents[knowledge]"
|
||||
|
||||
- name: Set environment variables
|
||||
- name: Verify environment variables
|
||||
run: |
|
||||
echo "OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY || 'sk-test-key-for-github-actions-testing-only-not-real' }}" >> $GITHUB_ENV
|
||||
echo "OPENAI_API_BASE=${{ secrets.OPENAI_API_BASE || 'https://api.openai.com/v1' }}" >> $GITHUB_ENV
|
||||
echo "OPENAI_MODEL_NAME=${{ secrets.OPENAI_MODEL_NAME || 'gpt-5-nano' }}" >> $GITHUB_ENV
|
||||
if [ -n "$OPENAI_API_KEY" ]; then
|
||||
echo "✅ OPENAI_API_KEY is set (length: ${#OPENAI_API_KEY})"
|
||||
else
|
||||
echo "⚠️ OPENAI_API_KEY is NOT set"
|
||||
fi
|
||||
|
||||
- name: Run Comprehensive Test Suite
|
||||
run: |
|
||||
|
||||
@@ -34,7 +34,7 @@ jobs:
|
||||
cd src/praisonai
|
||||
uv pip install --system ."[ui,gradio,api,agentops,google,openai,anthropic,cohere,chat,code,realtime,call,crewai,autogen]"
|
||||
uv pip install --system duckduckgo_search
|
||||
uv pip install --system pytest pytest-asyncio pytest-cov
|
||||
uv pip install --system pytest pytest-asyncio pytest-cov pytest-timeout
|
||||
# Install knowledge dependencies from praisonai-agents
|
||||
uv pip install --system "praisonaiagents[knowledge]"
|
||||
|
||||
@@ -262,7 +262,7 @@ jobs:
|
||||
export OPENAI_API_KEY="$OPENAI_API_KEY"
|
||||
export OPENAI_API_BASE="$OPENAI_API_BASE"
|
||||
export OPENAI_MODEL_NAME="$OPENAI_MODEL_NAME"
|
||||
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings
|
||||
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings -m "not real" --timeout=60
|
||||
|
||||
- name: Upload coverage reports to Codecov
|
||||
if: matrix.python-version == '3.11'
|
||||
|
||||
@@ -27,7 +27,7 @@ jobs:
|
||||
cd src/praisonai
|
||||
uv pip install --system ."[ui,gradio,api,agentops,google,openai,anthropic,cohere,chat,code,realtime,call,crewai,autogen]"
|
||||
uv pip install --system duckduckgo_search
|
||||
uv pip install --system pytest pytest-asyncio pytest-cov
|
||||
uv pip install --system pytest pytest-asyncio pytest-cov pytest-timeout
|
||||
# Install knowledge dependencies from praisonai-agents
|
||||
uv pip install --system "praisonaiagents[knowledge]"
|
||||
|
||||
@@ -489,7 +489,7 @@ jobs:
|
||||
|
||||
- name: Run Legacy Example Tests
|
||||
run: |
|
||||
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings
|
||||
cd src/praisonai && python -m pytest tests/test.py -v --tb=short --disable-warnings -m "not real" --timeout=60
|
||||
continue-on-error: true
|
||||
|
||||
- name: Upload coverage reports to Codecov
|
||||
|
||||
+1
-1
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
flask \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.3.19" \
|
||||
"praisonai[api]" \
|
||||
gunicorn \
|
||||
markdown
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.3.19" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.3.19" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.3.19" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -29,7 +29,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
@@ -71,7 +71,6 @@
|
||||
"import pandas as pd\n",
|
||||
"\n",
|
||||
"print(\"📦 Dependencies installed!\")\n",
|
||||
"print(\"🔄 Converting from AGNO → PraisonAI\")\n",
|
||||
"print(\"🎯 Building Job Hunting Agent System\")"
|
||||
]
|
||||
},
|
||||
|
||||
@@ -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()}")
|
||||
@@ -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")
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
@@ -0,0 +1,55 @@
|
||||
"""
|
||||
MCP WebSocket Transport Example
|
||||
|
||||
Demonstrates connecting to MCP servers via WebSocket (ws:// or wss://).
|
||||
WebSocket transport provides bidirectional real-time communication,
|
||||
ideal for cloud deployments and long-lived connections.
|
||||
|
||||
Features:
|
||||
- Auto-detection of ws:// and wss:// URLs
|
||||
- Authentication token support
|
||||
- Auto-reconnect with exponential backoff
|
||||
- Ping/pong keepalive
|
||||
|
||||
Protocol: SEP-1288 (WebSocket Transport)
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, MCP
|
||||
|
||||
# Basic WebSocket connection
|
||||
agent = Agent(
|
||||
name="WebSocket Assistant",
|
||||
instructions="You are a helpful assistant connected via WebSocket.",
|
||||
tools=MCP("ws://localhost:8080/mcp")
|
||||
)
|
||||
|
||||
# Secure WebSocket with authentication
|
||||
agent_secure = Agent(
|
||||
name="Secure WebSocket Assistant",
|
||||
instructions="You are a helpful assistant with secure WebSocket connection.",
|
||||
tools=MCP(
|
||||
"wss://api.example.com/mcp",
|
||||
auth_token="Bearer your-secret-token",
|
||||
timeout=60
|
||||
)
|
||||
)
|
||||
|
||||
# WebSocket with custom options
|
||||
agent_custom = Agent(
|
||||
name="Custom WebSocket Assistant",
|
||||
tools=MCP(
|
||||
"wss://api.example.com/mcp",
|
||||
auth_token="your-token",
|
||||
timeout=120,
|
||||
debug=True
|
||||
)
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Note: Requires a running WebSocket MCP server
|
||||
# Example server: https://github.com/modelcontextprotocol/servers
|
||||
|
||||
print("WebSocket MCP Transport Example")
|
||||
print("=" * 40)
|
||||
print("This example requires a WebSocket MCP server running.")
|
||||
print("URL patterns: ws://host:port/path or wss://host:port/path")
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Simple Workflow Example
|
||||
|
||||
The easiest way to create and run workflows in PraisonAI.
|
||||
Just pass functions or agents as steps - no complex configuration needed!
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, WorkflowContext, StepResult
|
||||
|
||||
# Define simple handler functions
|
||||
def validate(ctx: WorkflowContext) -> StepResult:
|
||||
"""Validate the input data."""
|
||||
if not ctx.input:
|
||||
return StepResult(output="No input provided", stop_workflow=True)
|
||||
return StepResult(output=f"Validated: {ctx.input}")
|
||||
|
||||
def process(ctx: WorkflowContext) -> StepResult:
|
||||
"""Process the validated data."""
|
||||
return StepResult(output=f"Processed: {ctx.previous_result}")
|
||||
|
||||
def format_output(ctx: WorkflowContext) -> StepResult:
|
||||
"""Format the final output."""
|
||||
return StepResult(output=f"✅ Final: {ctx.previous_result}")
|
||||
|
||||
# Create workflow - just list your functions!
|
||||
workflow = Workflow(
|
||||
name="Simple Pipeline",
|
||||
steps=[validate, process, format_output]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the workflow - PraisonAI style uses .start()
|
||||
result = workflow.start("Hello, World!", verbose=True)
|
||||
|
||||
print(f"\nFinal output: {result['output']}")
|
||||
print(f"Steps completed: {len(result['steps'])}")
|
||||
@@ -0,0 +1,57 @@
|
||||
"""
|
||||
Task Callbacks Example
|
||||
|
||||
Demonstrates using on_task_start and on_task_complete callbacks
|
||||
with PraisonAIAgents for workflow monitoring and logging.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
# Create an agent
|
||||
researcher = Agent(
|
||||
name="Researcher",
|
||||
role="Research Assistant",
|
||||
goal="Find and summarize information",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create tasks with variables
|
||||
task1 = Task(
|
||||
name="research",
|
||||
description="Research the topic: {{topic}}",
|
||||
expected_output="A brief summary of the topic",
|
||||
agent=researcher,
|
||||
variables={"topic": "artificial intelligence"}
|
||||
)
|
||||
|
||||
task2 = Task(
|
||||
name="summarize",
|
||||
description="Create a one-paragraph summary of the research.",
|
||||
expected_output="A concise summary",
|
||||
agent=researcher
|
||||
)
|
||||
|
||||
# Define callbacks
|
||||
def on_start(task, task_id):
|
||||
"""Called before each task starts."""
|
||||
print(f"🚀 Starting task: {task.name} (ID: {task_id})")
|
||||
|
||||
def on_complete(task, output):
|
||||
"""Called after each task completes."""
|
||||
print(f"✅ Completed task: {task.name}")
|
||||
print(f" Output preview: {str(output)[:100]}...")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create agents with callbacks
|
||||
agents = PraisonAIAgents(
|
||||
agents=[researcher],
|
||||
tasks=[task1, task2],
|
||||
process="workflow",
|
||||
on_task_start=on_start,
|
||||
on_task_complete=on_complete,
|
||||
variables={"global_var": "shared_value"} # Global variables
|
||||
)
|
||||
|
||||
# Run the workflow
|
||||
result = agents.start()
|
||||
print("\n🎉 Workflow completed!")
|
||||
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
Workflow Branching Example
|
||||
|
||||
Demonstrates conditional branching in workflows where steps can route
|
||||
to different next steps based on output content.
|
||||
"""
|
||||
|
||||
from praisonaiagents.memory.workflows import WorkflowManager, WorkflowStep, Workflow
|
||||
|
||||
# Create a workflow with branching
|
||||
workflow = Workflow(
|
||||
name="Decision Workflow",
|
||||
description="A workflow that branches based on validation result",
|
||||
steps=[
|
||||
WorkflowStep(
|
||||
name="validate",
|
||||
action="Check if the number 42 is positive. Reply with 'valid' or 'invalid'.",
|
||||
next_steps=["success_handler", "error_handler"],
|
||||
branch_condition={
|
||||
"valid": ["success_handler"],
|
||||
"invalid": ["error_handler"]
|
||||
}
|
||||
),
|
||||
WorkflowStep(
|
||||
name="success_handler",
|
||||
action="The validation passed! Generate a success message."
|
||||
),
|
||||
WorkflowStep(
|
||||
name="error_handler",
|
||||
action="The validation failed. Generate an error message."
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create manager and register workflow
|
||||
manager = WorkflowManager()
|
||||
manager.workflows["Decision Workflow"] = workflow
|
||||
|
||||
# Execute - will branch to success_handler since 42 is positive
|
||||
result = manager.execute(
|
||||
"Decision Workflow",
|
||||
default_llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
print("Workflow completed!")
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['status']}")
|
||||
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
Workflow Checkpoints Example
|
||||
|
||||
Demonstrates saving and resuming workflow execution using checkpoints.
|
||||
Useful for long-running workflows that may be interrupted.
|
||||
"""
|
||||
|
||||
from praisonaiagents.memory.workflows import WorkflowManager, WorkflowStep, Workflow
|
||||
|
||||
# Create a multi-step workflow
|
||||
workflow = Workflow(
|
||||
name="Long Process",
|
||||
description="A workflow with checkpoints for resumability",
|
||||
steps=[
|
||||
WorkflowStep(
|
||||
name="step1",
|
||||
action="Initialize the process and prepare data."
|
||||
),
|
||||
WorkflowStep(
|
||||
name="step2",
|
||||
action="Process the first batch of data."
|
||||
),
|
||||
WorkflowStep(
|
||||
name="step3",
|
||||
action="Process the second batch of data."
|
||||
),
|
||||
WorkflowStep(
|
||||
name="step4",
|
||||
action="Finalize and generate report."
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
manager = WorkflowManager()
|
||||
manager.workflows["Long Process"] = workflow
|
||||
|
||||
# Execute with checkpoint - saves after each step
|
||||
print("=== Starting workflow with checkpoint ===")
|
||||
result = manager.execute(
|
||||
"Long Process",
|
||||
default_llm="gpt-4o-mini",
|
||||
checkpoint="my-checkpoint" # Saves progress after each step
|
||||
)
|
||||
|
||||
print(f"Completed {len(result['results'])} steps")
|
||||
|
||||
# List available checkpoints
|
||||
print("\n=== Available checkpoints ===")
|
||||
checkpoints = manager.list_checkpoints()
|
||||
for cp in checkpoints:
|
||||
print(f" {cp['name']}: {cp['completed_steps']} steps completed")
|
||||
|
||||
# Resume from checkpoint (if workflow was interrupted)
|
||||
# result = manager.execute(
|
||||
# "Long Process",
|
||||
# default_llm="gpt-4o-mini",
|
||||
# resume="my-checkpoint" # Resumes from last saved state
|
||||
# )
|
||||
|
||||
# Clean up checkpoint
|
||||
manager.delete_checkpoint("my-checkpoint")
|
||||
print("\n✅ Checkpoint deleted")
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
Workflow Conditional Execution Example
|
||||
|
||||
Demonstrates using should_run to conditionally execute steps
|
||||
based on input or previous results.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, WorkflowStep, WorkflowContext, StepResult
|
||||
|
||||
# Condition functions - return True to run the step, False to skip
|
||||
def is_sensitive_content(ctx: WorkflowContext) -> bool:
|
||||
"""Only run compliance check for sensitive topics."""
|
||||
sensitive_keywords = ["legal", "medical", "financial", "security"]
|
||||
return any(keyword in ctx.input.lower() for keyword in sensitive_keywords)
|
||||
|
||||
def needs_translation(ctx: WorkflowContext) -> bool:
|
||||
"""Only translate if content is not in English."""
|
||||
return ctx.variables.get("language", "en") != "en"
|
||||
|
||||
# Handler functions
|
||||
def process_content(ctx: WorkflowContext) -> StepResult:
|
||||
return StepResult(output=f"Processed: {ctx.input}")
|
||||
|
||||
def compliance_check(ctx: WorkflowContext) -> StepResult:
|
||||
return StepResult(output=f"✅ Compliance verified for: {ctx.previous_result}")
|
||||
|
||||
def translate(ctx: WorkflowContext) -> StepResult:
|
||||
lang = ctx.variables.get("language", "unknown")
|
||||
return StepResult(output=f"🌐 Translated to {lang}: {ctx.previous_result}")
|
||||
|
||||
def finalize(ctx: WorkflowContext) -> StepResult:
|
||||
return StepResult(output=f"📄 Final: {ctx.previous_result}")
|
||||
|
||||
# Create workflow with conditional steps
|
||||
workflow = Workflow(
|
||||
name="Conditional Pipeline",
|
||||
steps=[
|
||||
WorkflowStep(name="process", handler=process_content),
|
||||
WorkflowStep(
|
||||
name="compliance",
|
||||
handler=compliance_check,
|
||||
should_run=is_sensitive_content # Only runs for sensitive content
|
||||
),
|
||||
WorkflowStep(
|
||||
name="translate",
|
||||
handler=translate,
|
||||
should_run=needs_translation # Only runs if language != "en"
|
||||
),
|
||||
WorkflowStep(name="finalize", handler=finalize)
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test 1: Normal content (skips compliance)
|
||||
print("=== Test 1: Normal content ===")
|
||||
result = workflow.start("Hello world", verbose=True)
|
||||
print(f"Output: {result['output']}\n")
|
||||
|
||||
# Test 2: Sensitive content (runs compliance)
|
||||
print("=== Test 2: Legal content ===")
|
||||
result = workflow.start("Review this legal document", verbose=True)
|
||||
print(f"Output: {result['output']}\n")
|
||||
|
||||
# Test 3: Non-English content (runs translation)
|
||||
print("=== Test 3: Non-English content ===")
|
||||
workflow.variables = {"language": "es"}
|
||||
result = workflow.start("Hola mundo", verbose=True)
|
||||
print(f"Output: {result['output']}")
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
Workflow Early Stop Example
|
||||
|
||||
Demonstrates how to stop a workflow early using custom handler
|
||||
functions that return StepResult with stop_workflow=True.
|
||||
"""
|
||||
|
||||
from praisonaiagents.memory.workflows import (
|
||||
WorkflowManager, WorkflowStep, Workflow,
|
||||
WorkflowContext, StepResult
|
||||
)
|
||||
|
||||
# Custom validator that can stop the workflow
|
||||
def validate_data(context: WorkflowContext) -> StepResult:
|
||||
"""Check if data is valid and stop workflow if not."""
|
||||
data = context.variables.get("data", {})
|
||||
|
||||
if data.get("value", 0) < 0:
|
||||
return StepResult(
|
||||
output="❌ Validation failed: negative value detected. Stopping workflow.",
|
||||
stop_workflow=True # This stops the entire workflow
|
||||
)
|
||||
|
||||
return StepResult(
|
||||
output="✅ Validation passed: data is valid.",
|
||||
stop_workflow=False # Continue to next step
|
||||
)
|
||||
|
||||
# Create workflow with early stop capability
|
||||
workflow = Workflow(
|
||||
name="Data Processing",
|
||||
description="Process data with validation gate",
|
||||
variables={
|
||||
"data": {"value": -5} # Invalid data - will trigger early stop
|
||||
},
|
||||
steps=[
|
||||
WorkflowStep(
|
||||
name="validate",
|
||||
handler=validate_data # Custom function
|
||||
),
|
||||
WorkflowStep(
|
||||
name="process",
|
||||
action="Process the validated data." # Won't run if validation fails
|
||||
),
|
||||
WorkflowStep(
|
||||
name="report",
|
||||
action="Generate final report." # Won't run if validation fails
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
manager = WorkflowManager()
|
||||
manager.workflows["Data Processing"] = workflow
|
||||
|
||||
# Test with invalid data (will stop early)
|
||||
print("=== Testing with invalid data ===")
|
||||
result = manager.execute("Data Processing", default_llm="gpt-4o-mini")
|
||||
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['output']}")
|
||||
|
||||
# Test with valid data (will complete)
|
||||
print("\n=== Testing with valid data ===")
|
||||
workflow.variables["data"] = {"value": 42}
|
||||
result = manager.execute("Data Processing", default_llm="gpt-4o-mini")
|
||||
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['status']}")
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Workflow Loops Example
|
||||
|
||||
Demonstrates iterating over a list of items in a workflow step.
|
||||
Each item is processed individually with the loop variable.
|
||||
"""
|
||||
|
||||
from praisonaiagents.memory.workflows import WorkflowManager, WorkflowStep, Workflow
|
||||
|
||||
# Create a workflow with loop
|
||||
workflow = Workflow(
|
||||
name="Process Items",
|
||||
description="Process a list of items using loop",
|
||||
variables={
|
||||
"items": ["apple", "banana", "cherry"] # List to iterate over
|
||||
},
|
||||
steps=[
|
||||
WorkflowStep(
|
||||
name="process_each",
|
||||
action="Describe the fruit: {{item}}",
|
||||
loop_over="items", # Variable name containing the list
|
||||
loop_var="item" # Variable name for current item
|
||||
),
|
||||
WorkflowStep(
|
||||
name="summarize",
|
||||
action="Summarize all the fruits that were processed."
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create manager and register workflow
|
||||
manager = WorkflowManager()
|
||||
manager.workflows["Process Items"] = workflow
|
||||
|
||||
# Execute - will process each fruit in the list
|
||||
result = manager.execute(
|
||||
"Process Items",
|
||||
default_llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
print("Workflow completed!")
|
||||
for step_result in result["results"]:
|
||||
print(f" {step_result['step']}: {step_result['status']}")
|
||||
if isinstance(step_result.get("output"), list):
|
||||
print(f" Loop processed {len(step_result['output'])} items")
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
Mixed Steps Workflow Example
|
||||
|
||||
Demonstrates combining different step types in a single workflow:
|
||||
- Agent instances
|
||||
- Handler functions
|
||||
- WorkflowStep with action strings
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, WorkflowStep, WorkflowContext, StepResult, Agent
|
||||
|
||||
# Create an agent
|
||||
analyzer = Agent(
|
||||
name="Analyzer",
|
||||
role="Data Analyst",
|
||||
goal="Analyze and summarize data",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Custom handler functions
|
||||
def validate_input(ctx: WorkflowContext) -> StepResult:
|
||||
"""Validate the input before processing."""
|
||||
if len(ctx.input) < 5:
|
||||
return StepResult(
|
||||
output="Input too short. Minimum 5 characters required.",
|
||||
stop_workflow=True
|
||||
)
|
||||
return StepResult(output=f"✅ Valid input: {ctx.input}")
|
||||
|
||||
def format_output(ctx: WorkflowContext) -> StepResult:
|
||||
"""Format the final output."""
|
||||
return StepResult(
|
||||
output=f"📊 Report:\n{ctx.previous_result}",
|
||||
variables={"report_generated": True}
|
||||
)
|
||||
|
||||
# Create workflow with mixed step types
|
||||
workflow = Workflow(
|
||||
name="Mixed Pipeline",
|
||||
steps=[
|
||||
# Step 1: Function - validates input
|
||||
validate_input,
|
||||
|
||||
# Step 2: Agent - analyzes content
|
||||
analyzer,
|
||||
|
||||
# Step 3: WorkflowStep with action string
|
||||
WorkflowStep(
|
||||
name="enhance",
|
||||
action="Add a conclusion to this analysis: {{previous_output}}"
|
||||
),
|
||||
|
||||
# Step 4: Function - formats output
|
||||
format_output
|
||||
],
|
||||
default_llm="gpt-4o-mini" # Used for action-based steps
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test with valid input
|
||||
print("=== Running mixed workflow ===")
|
||||
result = workflow.start(
|
||||
"Analyze the trends in AI adoption for 2024",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
print(f"\nFinal output:\n{result['output']}")
|
||||
print(f"\nVariables: {result['variables']}")
|
||||
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
Workflow with Agents Example
|
||||
|
||||
Demonstrates using Agent instances directly as workflow steps.
|
||||
Agents are automatically wrapped and executed in sequence.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, Agent
|
||||
|
||||
# Create agents
|
||||
researcher = Agent(
|
||||
name="Researcher",
|
||||
role="Research Specialist",
|
||||
goal="Find accurate information",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
writer = Agent(
|
||||
name="Writer",
|
||||
role="Content Writer",
|
||||
goal="Write engaging content",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
editor = Agent(
|
||||
name="Editor",
|
||||
role="Content Editor",
|
||||
goal="Polish and improve content",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create workflow with agents as steps
|
||||
workflow = Workflow(
|
||||
name="Content Pipeline",
|
||||
steps=[researcher, writer, editor]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the workflow
|
||||
result = workflow.start(
|
||||
"Write a short paragraph about artificial intelligence",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
print(f"\nFinal output:\n{result['output']}")
|
||||
@@ -86,3 +86,4 @@ build
|
||||
graph.py
|
||||
chroma_db/
|
||||
.qodo
|
||||
*_COMPARISON_ANALYSIS.md
|
||||
+559
-13
@@ -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,193 @@ 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
|
||||
|
||||
Enable real-time web search capabilities using LiteLLM's native web search.
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Enable web search with a supported model
|
||||
agent = Agent(
|
||||
name="Researcher",
|
||||
instructions="You are a helpful research assistant.",
|
||||
llm="openai/gpt-4o-search-preview",
|
||||
web_search=True
|
||||
)
|
||||
|
||||
result = agent.start("What are the latest developments in AI today?")
|
||||
```
|
||||
|
||||
**Supported Providers:**
|
||||
- **OpenAI**: `gpt-4o-search-preview`, `gpt-4o-mini-search-preview`
|
||||
- **Anthropic**: `claude-3-5-sonnet-latest`, `claude-sonnet-4`
|
||||
- **Google/Gemini**: `gemini-2.0-flash`, `gemini-2.5-pro`
|
||||
- **xAI**: `grok-3`
|
||||
- **Perplexity**: All models
|
||||
|
||||
---
|
||||
|
||||
## Web Fetch
|
||||
|
||||
Retrieve full content from specific URLs (Anthropic models only).
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Enable web fetch with an Anthropic model
|
||||
agent = Agent(
|
||||
name="Content Analyzer",
|
||||
instructions="You can fetch and analyze web content.",
|
||||
llm="anthropic/claude-sonnet-4-20250514",
|
||||
web_fetch=True
|
||||
)
|
||||
|
||||
result = agent.start("Fetch and summarize https://example.com")
|
||||
```
|
||||
|
||||
**Web Fetch Options:**
|
||||
```python
|
||||
agent = Agent(
|
||||
llm="anthropic/claude-sonnet-4-20250514",
|
||||
web_fetch={
|
||||
"max_uses": 10,
|
||||
"allowed_domains": ["example.com"],
|
||||
"citations": {"enabled": True}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Prompt Caching
|
||||
|
||||
Reduce costs and latency by caching parts of prompts.
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Enable prompt caching for Anthropic models
|
||||
agent = Agent(
|
||||
name="Legal Analyst",
|
||||
instructions="You are an AI assistant analyzing legal documents." * 50, # Long prompt
|
||||
llm="anthropic/claude-3-5-sonnet-latest",
|
||||
prompt_caching=True
|
||||
)
|
||||
|
||||
# First call creates cache, subsequent calls use cached tokens
|
||||
result = agent.start("Summarize the key terms")
|
||||
```
|
||||
|
||||
**Supported Providers:**
|
||||
- **OpenAI** (`openai/`) - Automatic caching for prompts ≥1024 tokens
|
||||
- **Anthropic** (`anthropic/`) - Manual caching with `cache_control`
|
||||
- **Bedrock** (`bedrock/`) - All models supporting prompt caching
|
||||
- **Deepseek** (`deepseek/`) - Works like OpenAI
|
||||
|
||||
---
|
||||
|
||||
## Claude Memory Tool (Beta)
|
||||
|
||||
Enable Claude to store and retrieve information across conversations.
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Enable Claude memory (Anthropic models only)
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
llm="anthropic/claude-sonnet-4-20250514",
|
||||
claude_memory=True
|
||||
)
|
||||
|
||||
# Claude will automatically:
|
||||
# 1. Check /memories directory before tasks
|
||||
# 2. Store progress/learnings in files
|
||||
# 3. Reference memories in future conversations
|
||||
result = agent.start("Research AI trends and remember key findings")
|
||||
```
|
||||
|
||||
**Features:**
|
||||
- File-based persistent memory in `.praison/claude_memory/`
|
||||
- Claude autonomously decides what to store/retrieve
|
||||
- Supports: view, create, str_replace, insert, delete, rename commands
|
||||
- Cross-conversation learning
|
||||
|
||||
**Supported Models:** Claude Sonnet 4, Claude Opus 4, Claude Haiku 4.5
|
||||
|
||||
---
|
||||
|
||||
## Memory (Zero Dependencies)
|
||||
|
||||
Enable persistent memory for agents without any extra packages.
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Simple: Enable file-based memory (no extra dependencies!)
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
memory=True # Uses FileMemory by default
|
||||
)
|
||||
|
||||
# Store memories
|
||||
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)
|
||||
|
||||
# Get memory context for prompts
|
||||
context = agent.get_memory_context(query="What does the user prefer?")
|
||||
```
|
||||
|
||||
**Memory Types:**
|
||||
- **Short-term**: Rolling buffer of recent context (auto-expires)
|
||||
- **Long-term**: Persistent important facts
|
||||
- **Entity**: People, places, organizations
|
||||
- **Episodic**: Date-based interaction history
|
||||
|
||||
**Storage Providers:**
|
||||
| Provider | Dependencies | Use Case |
|
||||
|----------|-------------|----------|
|
||||
| `memory=True` or `"file"` | None | Default, zero-config |
|
||||
| `"sqlite"` | Built-in | Search, indexing |
|
||||
| `"chromadb"` | chromadb | Semantic/vector search |
|
||||
| `"mem0"` | mem0ai | Graph memory, cloud |
|
||||
|
||||
---
|
||||
|
||||
## Guardrails
|
||||
|
||||
@@ -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()
|
||||
@@ -0,0 +1,94 @@
|
||||
# PraisonAI Example Plugin
|
||||
|
||||
This is a template for creating tool plugins for PraisonAI Agents.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
# Install in development mode
|
||||
pip install -e .
|
||||
|
||||
# Or install from PyPI (after publishing)
|
||||
pip install praisonai-example-plugin
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
After installation, tools are automatically discovered:
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Use tools by name - they're auto-discovered!
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
tools=["example_greet", "example_math", "example_reverse"]
|
||||
)
|
||||
|
||||
# Or import and use directly
|
||||
from praisonai_example_plugin import GreetTool, MathTool, reverse_text
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
tools=[GreetTool(), MathTool(), reverse_text]
|
||||
)
|
||||
```
|
||||
|
||||
## Creating Your Own Plugin
|
||||
|
||||
1. **Copy this template**
|
||||
2. **Rename the package** in `pyproject.toml` and `src/` directory
|
||||
3. **Add your tools** following the patterns in `__init__.py`
|
||||
4. **Update entry_points** in `pyproject.toml` to register your tools
|
||||
|
||||
### Tool Patterns
|
||||
|
||||
**Class-based tool:**
|
||||
```python
|
||||
from praisonaiagents import BaseTool
|
||||
|
||||
class MyTool(BaseTool):
|
||||
name = "my_tool"
|
||||
description = "What this tool does"
|
||||
|
||||
def run(self, param: str) -> str:
|
||||
return f"Result: {param}"
|
||||
```
|
||||
|
||||
**Decorator-based tool:**
|
||||
```python
|
||||
from praisonaiagents import tool
|
||||
|
||||
@tool
|
||||
def my_function(param: str) -> str:
|
||||
"""What this function does."""
|
||||
return f"Result: {param}"
|
||||
```
|
||||
|
||||
### Registering Tools
|
||||
|
||||
In `pyproject.toml`:
|
||||
```toml
|
||||
[project.entry-points."praisonaiagents.tools"]
|
||||
my_tool = "my_package:MyTool"
|
||||
my_function = "my_package:my_function"
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
pip install -e ".[dev]"
|
||||
pytest
|
||||
```
|
||||
|
||||
## Publishing
|
||||
|
||||
```bash
|
||||
pip install build twine
|
||||
python -m build
|
||||
twine upload dist/*
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,49 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonai-example-plugin"
|
||||
version = "1.0.0"
|
||||
description = "Example plugin for PraisonAI Agents - demonstrates tool plugin development"
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
requires-python = ">=3.10"
|
||||
authors = [
|
||||
{name = "Your Name", email = "[email protected]"}
|
||||
]
|
||||
keywords = ["praisonai", "agents", "tools", "plugin"]
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
]
|
||||
|
||||
dependencies = [
|
||||
"praisonaiagents>=0.0.168",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=7.0.0",
|
||||
"pytest-mock>=3.10.0",
|
||||
]
|
||||
|
||||
# This is the key part - register your tools as entry points
|
||||
# The format is: tool_name = "module.path:ClassName" or "module.path:function_name"
|
||||
[project.entry-points."praisonaiagents.tools"]
|
||||
example_greet = "praisonai_example_plugin:GreetTool"
|
||||
example_math = "praisonai_example_plugin:MathTool"
|
||||
example_reverse = "praisonai_example_plugin:reverse_text"
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/yourusername/praisonai-example-plugin"
|
||||
Documentation = "https://github.com/yourusername/praisonai-example-plugin#readme"
|
||||
Repository = "https://github.com/yourusername/praisonai-example-plugin"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
+119
@@ -0,0 +1,119 @@
|
||||
"""PraisonAI Example Plugin - Demonstrates how to create tool plugins.
|
||||
|
||||
This package shows how external developers can create pip-installable
|
||||
tool plugins for PraisonAI Agents.
|
||||
|
||||
Usage after installation:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Tools are auto-discovered via entry_points!
|
||||
agent = Agent(tools=["example_greet", "example_math"])
|
||||
|
||||
# Or import directly
|
||||
from praisonai_example_plugin import GreetTool, MathTool
|
||||
agent = Agent(tools=[GreetTool(), MathTool()])
|
||||
"""
|
||||
|
||||
from typing import Dict, Any
|
||||
from praisonaiagents import BaseTool, tool
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Class-Based Tools
|
||||
# =============================================================================
|
||||
|
||||
class GreetTool(BaseTool):
|
||||
"""A tool that generates personalized greetings.
|
||||
|
||||
This demonstrates a simple class-based tool with multiple parameters.
|
||||
"""
|
||||
name = "example_greet"
|
||||
description = "Generate a personalized greeting message"
|
||||
version = "1.0.0"
|
||||
|
||||
def run(self, name: str, greeting: str = "Hello", punctuation: str = "!") -> str:
|
||||
"""Generate a greeting.
|
||||
|
||||
Args:
|
||||
name: The name of the person to greet
|
||||
greeting: The greeting word (default: "Hello")
|
||||
punctuation: End punctuation (default: "!")
|
||||
|
||||
Returns:
|
||||
A formatted greeting string
|
||||
"""
|
||||
return f"{greeting}, {name}{punctuation}"
|
||||
|
||||
|
||||
class MathTool(BaseTool):
|
||||
"""A tool for basic mathematical operations.
|
||||
|
||||
Demonstrates returning structured data from a tool.
|
||||
"""
|
||||
name = "example_math"
|
||||
description = "Perform basic math operations: add, subtract, multiply, divide"
|
||||
version = "1.0.0"
|
||||
|
||||
def run(
|
||||
self,
|
||||
a: float,
|
||||
b: float,
|
||||
operation: str = "add"
|
||||
) -> Dict[str, Any]:
|
||||
"""Perform a math operation.
|
||||
|
||||
Args:
|
||||
a: First number
|
||||
b: Second number
|
||||
operation: One of 'add', 'subtract', 'multiply', 'divide'
|
||||
|
||||
Returns:
|
||||
Dict containing result and operation details
|
||||
"""
|
||||
ops = {
|
||||
"add": ("+", lambda x, y: x + y),
|
||||
"subtract": ("-", lambda x, y: x - y),
|
||||
"multiply": ("*", lambda x, y: x * y),
|
||||
"divide": ("/", lambda x, y: x / y if y != 0 else float('inf')),
|
||||
}
|
||||
|
||||
if operation not in ops:
|
||||
return {"error": f"Unknown operation: {operation}. Use: {list(ops.keys())}"}
|
||||
|
||||
symbol, func = ops[operation]
|
||||
result = func(a, b)
|
||||
|
||||
return {
|
||||
"result": result,
|
||||
"expression": f"{a} {symbol} {b} = {result}",
|
||||
"operation": operation
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Decorator-Based Tools
|
||||
# =============================================================================
|
||||
|
||||
@tool(name="example_reverse", description="Reverse a string of text")
|
||||
def reverse_text(text: str) -> str:
|
||||
"""Reverse the input text.
|
||||
|
||||
Args:
|
||||
text: The text to reverse
|
||||
|
||||
Returns:
|
||||
The reversed text
|
||||
"""
|
||||
return text[::-1]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Package Exports
|
||||
# =============================================================================
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__all__ = [
|
||||
"GreetTool",
|
||||
"MathTool",
|
||||
"reverse_text",
|
||||
]
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -31,9 +31,17 @@ from .agent.query_rewriter_agent import (
|
||||
RewriteStrategy,
|
||||
RewriteResult
|
||||
)
|
||||
from .agent.prompt_expander_agent import (
|
||||
PromptExpanderAgent,
|
||||
ExpandStrategy,
|
||||
ExpandResult
|
||||
)
|
||||
from .agents.agents import PraisonAIAgents
|
||||
from .task.task import Task
|
||||
from .tools.tools import Tools
|
||||
from .tools.base import BaseTool, ToolResult, ToolValidationError, validate_tool
|
||||
from .tools.decorator import tool, FunctionTool
|
||||
from .tools.registry import get_registry, register_tool, get_tool, ToolRegistry
|
||||
from .agents.autoagents import AutoAgents
|
||||
from .knowledge.knowledge import Knowledge
|
||||
from .knowledge.chunking import Chunking
|
||||
@@ -46,7 +54,51 @@ except ImportError:
|
||||
MCP = None
|
||||
from .session import Session
|
||||
from .memory.memory import Memory
|
||||
from .memory.workflows import Workflow, WorkflowStep, WorkflowContext, StepResult
|
||||
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:
|
||||
@@ -170,6 +222,17 @@ __all__ = [
|
||||
'Agents',
|
||||
'Tools',
|
||||
'Task',
|
||||
# Plugin system
|
||||
'BaseTool',
|
||||
'ToolResult',
|
||||
'ToolValidationError',
|
||||
'validate_tool',
|
||||
'tool',
|
||||
'FunctionTool',
|
||||
'ToolRegistry',
|
||||
'get_registry',
|
||||
'register_tool',
|
||||
'get_tool',
|
||||
'TaskOutput',
|
||||
'ReflectionOutput',
|
||||
'AutoAgents',
|
||||
@@ -214,7 +277,10 @@ __all__ = [
|
||||
'Provider',
|
||||
'QueryRewriterAgent',
|
||||
'RewriteStrategy',
|
||||
'RewriteResult'
|
||||
'RewriteResult',
|
||||
'PromptExpanderAgent',
|
||||
'ExpandStrategy',
|
||||
'ExpandResult'
|
||||
]
|
||||
|
||||
# Add MCP to __all__ if available
|
||||
@@ -225,3 +291,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'
|
||||
])
|
||||
|
||||
@@ -20,6 +20,11 @@ from .query_rewriter_agent import (
|
||||
RewriteStrategy,
|
||||
RewriteResult
|
||||
)
|
||||
from .prompt_expander_agent import (
|
||||
PromptExpanderAgent,
|
||||
ExpandStrategy,
|
||||
ExpandResult
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
'Agent',
|
||||
@@ -43,5 +48,8 @@ __all__ = [
|
||||
'Provider',
|
||||
'QueryRewriterAgent',
|
||||
'RewriteStrategy',
|
||||
'RewriteResult'
|
||||
'RewriteResult',
|
||||
'PromptExpanderAgent',
|
||||
'ExpandStrategy',
|
||||
'ExpandResult'
|
||||
]
|
||||
@@ -234,7 +234,23 @@ class Agent:
|
||||
max_guardrail_retries: int = 3,
|
||||
handoffs: Optional[List[Union['Agent', 'Handoff']]] = None,
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None
|
||||
api_key: Optional[str] = None,
|
||||
web_search: Optional[Union[bool, Dict[str, Any]]] = None,
|
||||
web_fetch: Optional[Union[bool, Dict[str, Any]]] = None,
|
||||
prompt_caching: Optional[bool] = None,
|
||||
claude_memory: Optional[Union[bool, Any]] = None,
|
||||
plan_mode: bool = False,
|
||||
planning: bool = False,
|
||||
planning_tools: Optional[List[Any]] = None,
|
||||
planning_reasoning: bool = False,
|
||||
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,
|
||||
history_in_context: Optional[int] = None,
|
||||
auto_save: Optional[str] = None
|
||||
):
|
||||
"""Initialize an Agent instance.
|
||||
|
||||
@@ -400,7 +416,11 @@ class Agent:
|
||||
model=model_name,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
metrics=metrics
|
||||
metrics=metrics,
|
||||
web_search=web_search,
|
||||
web_fetch=web_fetch,
|
||||
prompt_caching=prompt_caching,
|
||||
claude_memory=claude_memory
|
||||
)
|
||||
self._using_custom_llm = True
|
||||
except ImportError as e:
|
||||
@@ -435,6 +455,10 @@ class Agent:
|
||||
if api_key:
|
||||
llm_params['api_key'] = api_key
|
||||
llm_params['metrics'] = metrics
|
||||
llm_params['web_search'] = web_search
|
||||
llm_params['web_fetch'] = web_fetch
|
||||
llm_params['prompt_caching'] = prompt_caching
|
||||
llm_params['claude_memory'] = claude_memory
|
||||
self.llm_instance = LLM(**llm_params)
|
||||
self._using_custom_llm = True
|
||||
|
||||
@@ -455,6 +479,15 @@ class Agent:
|
||||
if callable(tools):
|
||||
# If a single function/callable is passed, wrap it in a list
|
||||
self.tools = [tools]
|
||||
elif isinstance(tools, str):
|
||||
# Single tool name string - resolve from registry
|
||||
self.tools = self._resolve_tool_names([tools])
|
||||
elif isinstance(tools, (list, tuple)):
|
||||
# Check if list contains strings (tool names) that need resolution
|
||||
if tools and all(isinstance(t, str) for t in tools):
|
||||
self.tools = self._resolve_tool_names(tools)
|
||||
else:
|
||||
self.tools = list(tools)
|
||||
else:
|
||||
# Handle all falsy values (None, False, 0, "", etc.) by defaulting to empty list
|
||||
self.tools = tools or []
|
||||
@@ -462,7 +495,8 @@ class Agent:
|
||||
self.max_iter = max_iter
|
||||
self.max_rpm = max_rpm
|
||||
self.max_execution_time = max_execution_time
|
||||
self.memory = memory
|
||||
self._memory_instance = None
|
||||
self._init_memory(memory, user_id)
|
||||
self.verbose = verbose
|
||||
self.allow_delegation = allow_delegation
|
||||
self.step_callback = step_callback
|
||||
@@ -502,6 +536,36 @@ 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
|
||||
|
||||
# Session management
|
||||
self.history_in_context = history_in_context # Number of past sessions to include
|
||||
self.auto_save = auto_save # Session name for auto-saving
|
||||
|
||||
# Initialize rules manager for persistent context (like Cursor/Windsurf)
|
||||
self._rules_manager = None
|
||||
self._init_rules_manager()
|
||||
|
||||
# Handle web_search fallback: inject DuckDuckGo tool for unsupported models
|
||||
if web_search and not self._model_supports_web_search():
|
||||
from ..tools.duckduckgo_tools import internet_search
|
||||
# Check if internet_search is not already in tools
|
||||
tool_names = [getattr(t, '__name__', str(t)) for t in self.tools]
|
||||
if 'internet_search' not in tool_names and 'duckduckgo' not in tool_names:
|
||||
self.tools.append(internet_search)
|
||||
logging.info("Model does not support native web search. Added DuckDuckGo fallback tool.")
|
||||
|
||||
# Log warning if web_fetch is enabled but model doesn't support it
|
||||
if web_fetch and not self._model_supports_web_fetch():
|
||||
logging.warning(f"Model '{self.llm}' does not support native web fetch. Web fetch will be ignored.")
|
||||
|
||||
# Initialize guardrail settings
|
||||
self.guardrail = guardrail
|
||||
@@ -533,6 +597,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."""
|
||||
@@ -567,6 +640,345 @@ Your Goal: {self.goal}
|
||||
self._agent_id = str(uuid.uuid4())
|
||||
return self._agent_id
|
||||
|
||||
def get_available_tools(self) -> List[Any]:
|
||||
"""
|
||||
Get tools available to this agent, filtered by plan_mode if enabled.
|
||||
|
||||
In plan_mode, only read-only tools are available to prevent
|
||||
modifications during the planning phase.
|
||||
|
||||
Returns:
|
||||
List of available tools
|
||||
"""
|
||||
if not self.plan_mode:
|
||||
return self.tools
|
||||
|
||||
# Filter to read-only tools only
|
||||
from ..planning import READ_ONLY_TOOLS, RESTRICTED_TOOLS
|
||||
|
||||
filtered_tools = []
|
||||
for tool in self.tools:
|
||||
tool_name = getattr(tool, '__name__', str(tool)).lower()
|
||||
|
||||
# Check if tool is in restricted list
|
||||
is_restricted = any(
|
||||
restricted.lower() in tool_name
|
||||
for restricted in RESTRICTED_TOOLS
|
||||
)
|
||||
|
||||
if not is_restricted:
|
||||
filtered_tools.append(tool)
|
||||
|
||||
return filtered_tools
|
||||
|
||||
def _model_supports_web_search(self) -> bool:
|
||||
"""
|
||||
Check if the agent's model supports native web search via LiteLLM.
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports native web search, False otherwise
|
||||
"""
|
||||
from ..llm.model_capabilities import supports_web_search
|
||||
|
||||
# Get the model name
|
||||
if hasattr(self, 'llm_instance') and self.llm_instance:
|
||||
model_name = self.llm_instance.model
|
||||
elif hasattr(self, 'llm') and self.llm:
|
||||
model_name = self.llm
|
||||
else:
|
||||
model_name = "gpt-5-nano"
|
||||
|
||||
return supports_web_search(model_name)
|
||||
|
||||
def _model_supports_web_fetch(self) -> bool:
|
||||
"""
|
||||
Check if the agent's model supports web fetch via LiteLLM.
|
||||
|
||||
Web fetch allows the model to retrieve full content from specific URLs.
|
||||
Currently only supported by Anthropic Claude models.
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports web fetch, False otherwise
|
||||
"""
|
||||
from ..llm.model_capabilities import supports_web_fetch
|
||||
|
||||
# Get the model name
|
||||
if hasattr(self, 'llm_instance') and self.llm_instance:
|
||||
model_name = self.llm_instance.model
|
||||
elif hasattr(self, 'llm') and self.llm:
|
||||
model_name = self.llm
|
||||
else:
|
||||
model_name = "gpt-5-nano"
|
||||
|
||||
return supports_web_fetch(model_name)
|
||||
|
||||
def _model_supports_prompt_caching(self) -> bool:
|
||||
"""
|
||||
Check if the agent's model supports prompt caching via LiteLLM.
|
||||
|
||||
Prompt caching allows caching parts of prompts to reduce costs and latency.
|
||||
Supported by OpenAI, Anthropic, Bedrock, and Deepseek.
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports prompt caching, False otherwise
|
||||
"""
|
||||
from ..llm.model_capabilities import supports_prompt_caching
|
||||
|
||||
# Get the model name
|
||||
if hasattr(self, 'llm_instance') and self.llm_instance:
|
||||
model_name = self.llm_instance.model
|
||||
elif hasattr(self, 'llm') and self.llm:
|
||||
model_name = self.llm
|
||||
else:
|
||||
model_name = "gpt-4o-mini"
|
||||
|
||||
return supports_prompt_caching(model_name)
|
||||
|
||||
@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.
|
||||
|
||||
Args:
|
||||
memory: Can be:
|
||||
- True: Use FileMemory with default settings
|
||||
- False/None: No memory
|
||||
- "file": Use FileMemory
|
||||
- "sqlite": Use existing Memory class with SQLite
|
||||
- dict: Configuration for memory
|
||||
- Memory/FileMemory instance: Use directly
|
||||
user_id: User identifier for memory isolation
|
||||
"""
|
||||
self.memory = memory
|
||||
|
||||
if memory is None or memory is False:
|
||||
self._memory_instance = None
|
||||
return
|
||||
|
||||
# Determine user_id
|
||||
mem_user_id = user_id or getattr(self, 'user_id', None) or "default"
|
||||
|
||||
if memory is True or memory == "file":
|
||||
# Use FileMemory (zero dependencies)
|
||||
from ..memory.file_memory import FileMemory
|
||||
self._memory_instance = FileMemory(
|
||||
user_id=mem_user_id,
|
||||
verbose=1 if getattr(self, 'verbose', False) else 0
|
||||
)
|
||||
elif isinstance(memory, str) and memory in ("sqlite", "chromadb", "mem0", "mongodb"):
|
||||
# Use full Memory class with specific provider
|
||||
try:
|
||||
from ..memory.memory import Memory
|
||||
config = {"provider": memory if memory != "sqlite" else "rag"}
|
||||
self._memory_instance = Memory(config)
|
||||
except ImportError:
|
||||
logging.warning(f"Memory provider '{memory}' requires additional dependencies. Falling back to FileMemory.")
|
||||
from ..memory.file_memory import FileMemory
|
||||
self._memory_instance = FileMemory(user_id=mem_user_id)
|
||||
elif isinstance(memory, dict):
|
||||
# Configuration dict
|
||||
provider = memory.get("provider", "file")
|
||||
if provider == "file":
|
||||
from ..memory.file_memory import FileMemory
|
||||
self._memory_instance = FileMemory(
|
||||
user_id=memory.get("user_id", mem_user_id),
|
||||
config=memory
|
||||
)
|
||||
else:
|
||||
try:
|
||||
from ..memory.memory import Memory
|
||||
self._memory_instance = Memory(memory)
|
||||
except ImportError:
|
||||
logging.warning("Full Memory class requires additional dependencies. Falling back to FileMemory.")
|
||||
from ..memory.file_memory import FileMemory
|
||||
self._memory_instance = FileMemory(user_id=mem_user_id)
|
||||
else:
|
||||
# Assume it's already a memory instance
|
||||
self._memory_instance = memory
|
||||
|
||||
def get_memory_context(self, query: Optional[str] = None) -> str:
|
||||
"""
|
||||
Get memory context for the current conversation.
|
||||
|
||||
Args:
|
||||
query: Optional query to focus the context
|
||||
|
||||
Returns:
|
||||
Formatted memory context string
|
||||
"""
|
||||
if not self._memory_instance:
|
||||
return ""
|
||||
|
||||
if hasattr(self._memory_instance, 'get_context'):
|
||||
return self._memory_instance.get_context(query=query)
|
||||
|
||||
return ""
|
||||
|
||||
def store_memory(self, content: str, memory_type: str = "short_term", **kwargs):
|
||||
"""
|
||||
Store content in memory.
|
||||
|
||||
Args:
|
||||
content: Content to store
|
||||
memory_type: Type of memory (short_term, long_term, entity, episodic)
|
||||
**kwargs: Additional arguments for the memory method
|
||||
"""
|
||||
if not self._memory_instance:
|
||||
return
|
||||
|
||||
if memory_type == "short_term" and hasattr(self._memory_instance, 'add_short_term'):
|
||||
self._memory_instance.add_short_term(content, **kwargs)
|
||||
elif memory_type == "long_term" and hasattr(self._memory_instance, 'add_long_term'):
|
||||
self._memory_instance.add_long_term(content, **kwargs)
|
||||
elif memory_type == "entity" and hasattr(self._memory_instance, 'add_entity'):
|
||||
self._memory_instance.add_entity(content, **kwargs)
|
||||
elif memory_type == "episodic" and hasattr(self._memory_instance, 'add_episodic'):
|
||||
self._memory_instance.add_episodic(content, **kwargs)
|
||||
|
||||
def _display_memory_info(self):
|
||||
"""Display memory information to user in a friendly format."""
|
||||
if not self._memory_instance:
|
||||
return
|
||||
|
||||
# Only display once per chat session
|
||||
if hasattr(self, '_memory_displayed') and self._memory_displayed:
|
||||
return
|
||||
self._memory_displayed = True
|
||||
|
||||
stats = self._memory_instance.get_stats()
|
||||
short_count = stats.get('short_term_count', 0)
|
||||
long_count = stats.get('long_term_count', 0)
|
||||
entity_count = stats.get('entity_count', 0)
|
||||
storage_path = stats.get('storage_path', '')
|
||||
|
||||
total_memories = short_count + long_count + entity_count
|
||||
|
||||
if total_memories > 0:
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
# Build memory info text
|
||||
info_parts = []
|
||||
if long_count > 0:
|
||||
info_parts.append(f"💾 {long_count} long-term")
|
||||
if short_count > 0:
|
||||
info_parts.append(f"⚡ {short_count} short-term")
|
||||
if entity_count > 0:
|
||||
info_parts.append(f"👤 {entity_count} entities")
|
||||
|
||||
memory_text = Text()
|
||||
memory_text.append("🧠 Memory loaded: ", style="bold cyan")
|
||||
memory_text.append(" | ".join(info_parts))
|
||||
memory_text.append(f"\n📁 Storage: {storage_path}", style="dim")
|
||||
|
||||
self.console.print(Panel(
|
||||
memory_text,
|
||||
title="[bold]Agent Memory[/bold]",
|
||||
border_style="cyan",
|
||||
expand=False
|
||||
))
|
||||
|
||||
@property
|
||||
def llm_model(self):
|
||||
"""Unified property to get the LLM model regardless of configuration type.
|
||||
@@ -813,17 +1225,35 @@ Your Goal: {self.goal}
|
||||
if not self.use_system_prompt:
|
||||
return None
|
||||
|
||||
# Check cache first
|
||||
tools_key = self._get_tools_cache_key(tools)
|
||||
cache_key = f"{self.role}:{self.goal}:{tools_key}"
|
||||
|
||||
if cache_key in self._system_prompt_cache:
|
||||
return self._system_prompt_cache[cache_key]
|
||||
# Check cache first (skip cache if memory is enabled since context is dynamic)
|
||||
if not self._memory_instance:
|
||||
tools_key = self._get_tools_cache_key(tools)
|
||||
cache_key = f"{self.role}:{self.goal}:{tools_key}"
|
||||
|
||||
if cache_key in self._system_prompt_cache:
|
||||
return self._system_prompt_cache[cache_key]
|
||||
else:
|
||||
cache_key = None # Don't cache when memory is enabled
|
||||
|
||||
system_prompt = f"""{self.backstory}\n
|
||||
Your Role: {self.role}\n
|
||||
Your Goal: {self.goal}"""
|
||||
|
||||
# Add rules context if rules manager is enabled
|
||||
if self._rules_manager:
|
||||
rules_context = self.get_rules_context()
|
||||
if rules_context:
|
||||
system_prompt += f"\n\n## Rules (Guidelines you must follow)\n{rules_context}"
|
||||
|
||||
# Add memory context if memory is enabled
|
||||
if self._memory_instance:
|
||||
memory_context = self.get_memory_context()
|
||||
if memory_context:
|
||||
system_prompt += f"\n\n## Memory (Information you remember about the user)\n{memory_context}"
|
||||
# Display memory info to user if verbose
|
||||
if self.verbose:
|
||||
self._display_memory_info()
|
||||
|
||||
# Add tool usage instructions if tools are available
|
||||
# Use provided tools or fall back to self.tools
|
||||
tools_to_use = tools if tools is not None else self.tools
|
||||
@@ -853,9 +1283,9 @@ Your Goal: {self.goal}"""
|
||||
if tool_names:
|
||||
system_prompt += f"\n\nYou have access to the following tools: {', '.join(tool_names)}. Use these tools when appropriate to help complete your tasks. Always use tools when they can help provide accurate information or perform actions."
|
||||
|
||||
# Cache the generated system prompt
|
||||
# Cache the generated system prompt (only if cache_key is set, i.e., memory not enabled)
|
||||
# Simple cache size limit to prevent unbounded growth
|
||||
if len(self._system_prompt_cache) < self._max_cache_size:
|
||||
if cache_key and len(self._system_prompt_cache) < self._max_cache_size:
|
||||
self._system_prompt_cache[cache_key] = system_prompt
|
||||
return system_prompt
|
||||
|
||||
@@ -1008,6 +1438,31 @@ Your Goal: {self.goal}"""
|
||||
tools=self.tools
|
||||
)
|
||||
|
||||
def _resolve_tool_names(self, tool_names):
|
||||
"""Resolve tool names to actual tool instances from registry.
|
||||
|
||||
Args:
|
||||
tool_names: List of tool name strings
|
||||
|
||||
Returns:
|
||||
List of resolved tool instances
|
||||
"""
|
||||
resolved = []
|
||||
try:
|
||||
from ..tools.registry import get_registry
|
||||
registry = get_registry()
|
||||
|
||||
for name in tool_names:
|
||||
tool = registry.get(name)
|
||||
if tool is not None:
|
||||
resolved.append(tool)
|
||||
else:
|
||||
logging.warning(f"Tool '{name}' not found in registry")
|
||||
except ImportError:
|
||||
logging.warning("Tool registry not available, cannot resolve tool names")
|
||||
|
||||
return resolved
|
||||
|
||||
def _cast_arguments(self, func, arguments):
|
||||
"""Cast arguments to their expected types based on function signature."""
|
||||
if not callable(func) or not arguments:
|
||||
@@ -1105,13 +1560,31 @@ Your Goal: {self.goal}"""
|
||||
# Try to find the function in the agent's tools list first
|
||||
func = None
|
||||
for tool in self.tools if isinstance(self.tools, (list, tuple)) else []:
|
||||
# Check for BaseTool instances (plugin system)
|
||||
from ..tools.base import BaseTool
|
||||
if isinstance(tool, BaseTool) and tool.name == function_name:
|
||||
func = tool
|
||||
break
|
||||
# Check for FunctionTool (decorated functions)
|
||||
if hasattr(tool, 'name') and getattr(tool, 'name', None) == function_name:
|
||||
func = tool
|
||||
break
|
||||
if (callable(tool) and getattr(tool, '__name__', '') == function_name) or \
|
||||
(inspect.isclass(tool) and tool.__name__ == function_name):
|
||||
func = tool
|
||||
break
|
||||
|
||||
if func is None:
|
||||
# If not found in tools, try globals and main
|
||||
# Check the global tool registry for plugins
|
||||
try:
|
||||
from ..tools.registry import get_registry
|
||||
registry = get_registry()
|
||||
func = registry.get(function_name)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if func is None:
|
||||
# If not found in tools or registry, try globals and main
|
||||
func = globals().get(function_name)
|
||||
if not func:
|
||||
import __main__
|
||||
@@ -1119,6 +1592,12 @@ Your Goal: {self.goal}"""
|
||||
|
||||
if func:
|
||||
try:
|
||||
# BaseTool instances (plugin system) - call run() method
|
||||
from ..tools.base import BaseTool
|
||||
if isinstance(func, BaseTool):
|
||||
casted_arguments = self._cast_arguments(func.run, arguments)
|
||||
return func.run(**casted_arguments)
|
||||
|
||||
# Langchain: If it's a class with run but not _run, instantiate and call run
|
||||
if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
|
||||
instance = func()
|
||||
@@ -2092,21 +2571,189 @@ 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({"step": i + 1, "description": step.description, "result": result})
|
||||
|
||||
# Update context for next step (use full result, not truncated)
|
||||
context += f"\n\nStep {i + 1} result: {result 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")
|
||||
|
||||
# Compile all results into a comprehensive final output
|
||||
if len(results) > 1:
|
||||
# Create a compilation prompt
|
||||
all_results_text = "\n\n".join([
|
||||
f"## Step {r['step']}: {r['description']}\n\n{r['result']}"
|
||||
for r in results
|
||||
])
|
||||
|
||||
compilation_prompt = f"""You are tasked with compiling a comprehensive, detailed report from the following research steps.
|
||||
|
||||
IMPORTANT: Write a DETAILED and COMPREHENSIVE report. Do NOT summarize or compress the information.
|
||||
Include ALL relevant details, data, statistics, and findings from each step.
|
||||
Organize the information logically with clear sections and subsections.
|
||||
|
||||
## Research Results to Compile:
|
||||
|
||||
{all_results_text}
|
||||
|
||||
## Instructions:
|
||||
1. Combine all the information into a single, well-organized document
|
||||
2. Preserve ALL details, numbers, statistics, and specific findings
|
||||
3. Use clear headings and subheadings
|
||||
4. Do not omit any important information
|
||||
5. Make it comprehensive and detailed
|
||||
|
||||
Write the complete compiled report:"""
|
||||
|
||||
console.print("\n[bold blue]📝 COMPILING FINAL REPORT[/bold blue]")
|
||||
console.print("[dim]Creating comprehensive output from all steps...[/dim]\n")
|
||||
|
||||
final_result = self.chat(compilation_prompt, **kwargs)
|
||||
return final_result
|
||||
|
||||
# Return the single result if only one step
|
||||
return results[0]["result"] 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 streaming is enabled (either from kwargs or agent's stream attribute)
|
||||
stream_enabled = kwargs.get('stream', getattr(self, 'stream', False))
|
||||
# Load history from past sessions if history_in_context is set
|
||||
self._load_history_context()
|
||||
|
||||
if stream_enabled:
|
||||
# Check if planning mode is enabled
|
||||
if self.planning:
|
||||
result = self._start_with_planning(prompt, **kwargs)
|
||||
elif kwargs.get('stream', getattr(self, 'stream', False)):
|
||||
# Return a generator for streaming response
|
||||
return self._start_stream(prompt, **kwargs)
|
||||
result = self._start_stream(prompt, **kwargs)
|
||||
else:
|
||||
# Return regular chat response for backward compatibility
|
||||
# Explicitly pass the resolved stream parameter to avoid chat() method default
|
||||
kwargs['stream'] = stream_enabled
|
||||
return self.chat(prompt, **kwargs)
|
||||
kwargs['stream'] = False
|
||||
result = self.chat(prompt, **kwargs)
|
||||
|
||||
# Auto-save session if enabled
|
||||
self._auto_save_session()
|
||||
|
||||
return result
|
||||
|
||||
def _load_history_context(self):
|
||||
"""Load history from past sessions into context if history_in_context is set."""
|
||||
if not self.history_in_context or not self._memory_instance:
|
||||
return
|
||||
|
||||
try:
|
||||
sessions = self._memory_instance.list_sessions()
|
||||
if not sessions:
|
||||
return
|
||||
|
||||
# Get the last N sessions
|
||||
sessions_to_load = sessions[:self.history_in_context]
|
||||
|
||||
for session_info in reversed(sessions_to_load):
|
||||
try:
|
||||
session_data = self._memory_instance.resume_session(session_info["name"])
|
||||
history = session_data.get("conversation_history", [])
|
||||
|
||||
# Add history to chat_history with a marker
|
||||
for msg in history:
|
||||
if msg not in self.chat_history:
|
||||
# Mark as from history to avoid duplication
|
||||
msg_copy = msg.copy()
|
||||
msg_copy["_from_history"] = True
|
||||
self.chat_history.append(msg_copy)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if sessions_to_load:
|
||||
logging.debug(f"Loaded history from {len(sessions_to_load)} past session(s)")
|
||||
except Exception as e:
|
||||
logging.debug(f"Error loading history context: {e}")
|
||||
|
||||
def _auto_save_session(self):
|
||||
"""Auto-save session if auto_save is enabled."""
|
||||
if not self.auto_save or not self._memory_instance:
|
||||
return
|
||||
|
||||
try:
|
||||
# Filter out history markers before saving
|
||||
clean_history = [
|
||||
{k: v for k, v in msg.items() if k != "_from_history"}
|
||||
for msg in self.chat_history
|
||||
]
|
||||
|
||||
self._memory_instance.save_session(
|
||||
name=self.auto_save,
|
||||
conversation_history=clean_history,
|
||||
metadata={"agent_name": self.name, "user_id": self.user_id}
|
||||
)
|
||||
logging.debug(f"Auto-saved session: {self.auto_save}")
|
||||
except Exception as e:
|
||||
logging.debug(f"Error auto-saving session: {e}")
|
||||
|
||||
def _start_stream(self, prompt: str, **kwargs) -> Generator[str, None, None]:
|
||||
"""Stream generator for real-time response chunks."""
|
||||
|
||||
@@ -0,0 +1,415 @@
|
||||
"""
|
||||
Prompt Expander Agent Module
|
||||
|
||||
This module provides the PromptExpanderAgent class for expanding short prompts
|
||||
into detailed, comprehensive prompts for better task execution.
|
||||
|
||||
Unlike QueryRewriterAgent (which optimizes queries for search/retrieval),
|
||||
PromptExpanderAgent focuses on enriching prompts for task execution.
|
||||
|
||||
Supported Expansion Strategies:
|
||||
- **BASIC**: Simple expansion with clarity improvements
|
||||
- **DETAILED**: Rich expansion with context, constraints, and examples
|
||||
- **STRUCTURED**: Expansion with clear structure (task, format, requirements)
|
||||
- **CREATIVE**: Expansion with creative flair and vivid language
|
||||
- **AUTO**: Automatically selects the best strategy based on prompt analysis
|
||||
|
||||
Example:
|
||||
from praisonaiagents import PromptExpanderAgent, ExpandStrategy
|
||||
|
||||
agent = PromptExpanderAgent(verbose=True)
|
||||
|
||||
# Basic expansion
|
||||
result = agent.expand("write a movie script in 3 lines")
|
||||
print(result.expanded_prompt)
|
||||
|
||||
# Detailed expansion
|
||||
result = agent.expand("blog about AI", strategy=ExpandStrategy.DETAILED)
|
||||
print(result.expanded_prompt)
|
||||
|
||||
# With tools for context gathering
|
||||
from praisonaiagents.tools import internet_search
|
||||
agent = PromptExpanderAgent(tools=[internet_search], verbose=True)
|
||||
result = agent.expand("latest AI developments")
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional, Any, Dict
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from datetime import datetime
|
||||
|
||||
class ExpandStrategy(Enum):
|
||||
"""Enumeration of available prompt expansion strategies."""
|
||||
BASIC = "basic"
|
||||
DETAILED = "detailed"
|
||||
STRUCTURED = "structured"
|
||||
CREATIVE = "creative"
|
||||
AUTO = "auto"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExpandResult:
|
||||
"""Result of a prompt expansion operation."""
|
||||
original_prompt: str
|
||||
expanded_prompt: str
|
||||
strategy_used: ExpandStrategy
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __repr__(self):
|
||||
return f"ExpandResult(strategy={self.strategy_used.value}, original_len={len(self.original_prompt)}, expanded_len={len(self.expanded_prompt)})"
|
||||
|
||||
|
||||
class PromptExpanderAgent:
|
||||
"""
|
||||
Agent for expanding short prompts into detailed, comprehensive prompts.
|
||||
|
||||
This agent transforms brief task descriptions into rich, detailed prompts
|
||||
that provide better context and guidance for task execution.
|
||||
|
||||
Key difference from QueryRewriterAgent:
|
||||
- QueryRewriterAgent: Optimizes queries for search/retrieval (RAG)
|
||||
- PromptExpanderAgent: Expands prompts for detailed task execution
|
||||
|
||||
Attributes:
|
||||
name: Name of the agent
|
||||
model: LLM model to use for expansion
|
||||
verbose: Whether to print detailed logs
|
||||
temperature: Temperature for LLM generation (higher = more creative)
|
||||
tools: Optional tools for context gathering
|
||||
|
||||
Example:
|
||||
agent = PromptExpanderAgent(model="gpt-4o-mini", verbose=True)
|
||||
result = agent.expand("write a poem", strategy=ExpandStrategy.CREATIVE)
|
||||
print(result.expanded_prompt)
|
||||
"""
|
||||
|
||||
# Default prompts for each strategy
|
||||
PROMPTS = {
|
||||
"basic": """You are an expert at improving task prompts.
|
||||
|
||||
Current date: {current_date}
|
||||
|
||||
Given the following task prompt, expand it to be clearer and more actionable while preserving the original intent.
|
||||
|
||||
Guidelines:
|
||||
- Keep it as a TASK/COMMAND (do NOT convert to a question)
|
||||
- Add clarity without changing the core request
|
||||
- Fix any ambiguity
|
||||
- Keep it concise but complete
|
||||
- Use the current date for any time-sensitive context
|
||||
|
||||
Original prompt: {prompt}
|
||||
|
||||
Expanded prompt:""",
|
||||
|
||||
"detailed": """You are an expert at creating comprehensive task prompts.
|
||||
|
||||
Current date: {current_date}
|
||||
|
||||
Given the following brief task prompt, expand it into a detailed, actionable prompt that provides rich context and clear guidance.
|
||||
|
||||
Guidelines:
|
||||
- Keep it as a TASK/COMMAND (do NOT convert to a question)
|
||||
- Add specific details about what to include
|
||||
- Specify quality expectations
|
||||
- Add relevant constraints or requirements
|
||||
- Include format guidance if applicable
|
||||
- Mention style or tone if relevant
|
||||
- Use the current date for any time-sensitive context
|
||||
|
||||
Original prompt: {prompt}
|
||||
|
||||
Detailed expanded prompt:""",
|
||||
|
||||
"structured": """You are an expert at creating well-structured task prompts.
|
||||
|
||||
Current date: {current_date}
|
||||
|
||||
Given the following task prompt, expand it into a structured format with clear sections.
|
||||
|
||||
Guidelines:
|
||||
- Keep it as a TASK/COMMAND (do NOT convert to a question)
|
||||
- Use a clear structure:
|
||||
* Task: What to do
|
||||
* Format: Expected output format
|
||||
* Requirements: Specific requirements
|
||||
* Style: Tone or style guidance
|
||||
* Constraints: Any limitations
|
||||
- Be specific and actionable
|
||||
- Use the current date for any time-sensitive context
|
||||
|
||||
Original prompt: {prompt}
|
||||
|
||||
Structured expanded prompt:""",
|
||||
|
||||
"creative": """You are an expert at crafting inspiring, creative task prompts.
|
||||
|
||||
Current date: {current_date}
|
||||
|
||||
Given the following task prompt, expand it with vivid language and creative direction while keeping it actionable.
|
||||
|
||||
Guidelines:
|
||||
- Keep it as a TASK/COMMAND (do NOT convert to a question)
|
||||
- Use evocative, inspiring language
|
||||
- Paint a picture of the desired outcome
|
||||
- Add creative direction and artistic guidance
|
||||
- Make it exciting and engaging
|
||||
- Preserve the core task intent
|
||||
- Use the current date for any time-sensitive context
|
||||
|
||||
Original prompt: {prompt}
|
||||
|
||||
Creative expanded prompt:""",
|
||||
|
||||
"auto": """You are an expert at analyzing and expanding task prompts.
|
||||
|
||||
Current date: {current_date}
|
||||
|
||||
Analyze the following prompt and expand it appropriately based on its nature.
|
||||
|
||||
For short/vague prompts: Add significant detail and context
|
||||
For task prompts: Preserve the action-oriented nature
|
||||
For creative tasks: Add artistic direction
|
||||
For technical tasks: Add precision and requirements
|
||||
|
||||
IMPORTANT: Keep it as a TASK/COMMAND. Never convert to a question.
|
||||
Use the current date for any time-sensitive context.
|
||||
|
||||
Original prompt: {prompt}
|
||||
{context}
|
||||
|
||||
Expanded prompt:"""
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "PromptExpanderAgent",
|
||||
model: str = "gpt-4o-mini",
|
||||
instructions: Optional[str] = None,
|
||||
verbose: bool = False,
|
||||
temperature: float = 0.7, # Higher for creativity
|
||||
max_tokens: int = 1000,
|
||||
tools: Optional[List[Any]] = None
|
||||
):
|
||||
"""
|
||||
Initialize the PromptExpanderAgent.
|
||||
|
||||
Args:
|
||||
name: Name of the agent
|
||||
model: LLM model to use (default: gpt-4o-mini)
|
||||
instructions: Custom instructions for the agent
|
||||
verbose: Whether to print detailed logs
|
||||
temperature: Temperature for LLM generation (default: 0.7 for creativity)
|
||||
max_tokens: Maximum tokens for LLM response
|
||||
tools: Optional list of tools for context gathering
|
||||
"""
|
||||
self.name = name
|
||||
self.model = model
|
||||
self.instructions = instructions
|
||||
self.verbose = verbose
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
self.tools = tools or []
|
||||
self._agent = None # Lazy initialized
|
||||
|
||||
# Setup logging
|
||||
self.logger = logging.getLogger(__name__)
|
||||
if self.verbose:
|
||||
self.logger.setLevel(logging.DEBUG)
|
||||
|
||||
@property
|
||||
def agent(self):
|
||||
"""Lazy initialization of internal Agent."""
|
||||
if self._agent is None:
|
||||
from .agent import Agent
|
||||
self._agent = Agent(
|
||||
name=self.name,
|
||||
role="Prompt Expansion Specialist",
|
||||
goal="Expand brief prompts into detailed, actionable prompts for better task execution. Use available tools when needed to gather context.",
|
||||
backstory="You are an expert at understanding user intent and transforming brief requests into comprehensive, detailed prompts.",
|
||||
tools=self.tools,
|
||||
llm=self.model,
|
||||
verbose=self.verbose,
|
||||
markdown=False
|
||||
)
|
||||
return self._agent
|
||||
|
||||
def _call_agent(self, prompt: str) -> str:
|
||||
"""Call the internal Agent with the given prompt."""
|
||||
try:
|
||||
response = self.agent.chat(prompt)
|
||||
return response if response else ""
|
||||
except Exception as e:
|
||||
self.logger.error(f"Agent call failed: {e}")
|
||||
raise
|
||||
|
||||
def _is_short_prompt(self, prompt: str, threshold: int = 10) -> bool:
|
||||
"""Check if prompt is considered short."""
|
||||
return len(prompt.split()) < threshold
|
||||
|
||||
def _is_task_prompt(self, prompt: str) -> bool:
|
||||
"""Check if prompt is a task/command (vs a question)."""
|
||||
task_verbs = [
|
||||
"write", "create", "generate", "build", "make", "design",
|
||||
"develop", "implement", "compose", "draft", "produce",
|
||||
"craft", "construct", "prepare", "draw", "paint", "code"
|
||||
]
|
||||
prompt_lower = prompt.lower().strip()
|
||||
return any(prompt_lower.startswith(verb) for verb in task_verbs)
|
||||
|
||||
def _is_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 (non-technical)."""
|
||||
creative_indicators = [
|
||||
"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)
|
||||
|
||||
def _detect_strategy(self, prompt: str) -> ExpandStrategy:
|
||||
"""Auto-detect the best expansion strategy for a prompt."""
|
||||
# Very short prompts need detailed expansion
|
||||
if self._is_short_prompt(prompt, threshold=5):
|
||||
return ExpandStrategy.DETAILED
|
||||
|
||||
# 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
|
||||
|
||||
# Task prompts get structured expansion
|
||||
if self._is_task_prompt(prompt):
|
||||
return ExpandStrategy.STRUCTURED
|
||||
|
||||
# Default to basic
|
||||
return ExpandStrategy.BASIC
|
||||
|
||||
def expand(
|
||||
self,
|
||||
prompt: str,
|
||||
strategy: ExpandStrategy = ExpandStrategy.AUTO,
|
||||
context: Optional[str] = None
|
||||
) -> ExpandResult:
|
||||
"""
|
||||
Expand a prompt using the specified strategy.
|
||||
|
||||
All operations go through an internal Agent which automatically handles
|
||||
tool calling when tools are provided. The agent decides when to use tools
|
||||
based on the prompt context.
|
||||
|
||||
Args:
|
||||
prompt: The original user prompt
|
||||
strategy: Expansion strategy to use
|
||||
context: Additional context to help with expansion
|
||||
|
||||
Returns:
|
||||
ExpandResult containing the expanded prompt
|
||||
"""
|
||||
if self.verbose:
|
||||
self.logger.debug(f"Expanding prompt: '{prompt}' with strategy: {strategy.value}")
|
||||
if self.tools:
|
||||
print(f"[cyan]Agent has {len(self.tools)} tools available (will use if needed)...[/cyan]")
|
||||
|
||||
# Auto-detect strategy if needed
|
||||
if strategy == ExpandStrategy.AUTO:
|
||||
strategy = self._detect_strategy(prompt)
|
||||
if self.verbose:
|
||||
self.logger.debug(f"Auto-detected strategy: {strategy.value}")
|
||||
|
||||
# Build the expansion prompt
|
||||
if strategy == ExpandStrategy.AUTO:
|
||||
# Fallback if auto-detect returns AUTO
|
||||
strategy = ExpandStrategy.DETAILED
|
||||
|
||||
# Get current date for time-sensitive context
|
||||
current_date = datetime.now().strftime("%B %d, %Y")
|
||||
|
||||
expansion_prompt = self.PROMPTS[strategy.value].format(
|
||||
prompt=prompt,
|
||||
context=context or "",
|
||||
current_date=current_date
|
||||
)
|
||||
|
||||
# Add tool usage instructions if tools are available
|
||||
if self.tools:
|
||||
# Build dynamic tool descriptions from docstrings
|
||||
tool_descriptions = []
|
||||
for tool in self.tools:
|
||||
tool_name = getattr(tool, '__name__', str(tool))
|
||||
tool_doc = getattr(tool, '__doc__', '')
|
||||
# Get first line of docstring as description
|
||||
if tool_doc:
|
||||
first_line = tool_doc.strip().split('\n')[0].strip()
|
||||
tool_descriptions.append(f"- {tool_name}: {first_line}")
|
||||
else:
|
||||
tool_descriptions.append(f"- {tool_name}")
|
||||
|
||||
tools_list = '\n'.join(tool_descriptions)
|
||||
|
||||
tool_instruction = f"""
|
||||
|
||||
IMPORTANT: You have the following tools available:
|
||||
{tools_list}
|
||||
|
||||
Before expanding this prompt, FIRST use the available tools to gather relevant current information about the topic. This will help you create a more informed and accurate expanded prompt.
|
||||
|
||||
Steps:
|
||||
1. Use the appropriate tool(s) to gather relevant information about the topic
|
||||
2. Use the gathered information to inform your expansion
|
||||
3. Create an expanded prompt that incorporates the latest context
|
||||
|
||||
"""
|
||||
expansion_prompt = tool_instruction + expansion_prompt
|
||||
|
||||
# Call the agent
|
||||
expanded = self._call_agent(expansion_prompt).strip()
|
||||
|
||||
# If expansion failed, return original
|
||||
if not expanded:
|
||||
expanded = prompt
|
||||
|
||||
return ExpandResult(
|
||||
original_prompt=prompt,
|
||||
expanded_prompt=expanded,
|
||||
strategy_used=strategy,
|
||||
metadata={
|
||||
"original_length": len(prompt),
|
||||
"expanded_length": len(expanded),
|
||||
"has_context": bool(context),
|
||||
"has_tools": bool(self.tools)
|
||||
}
|
||||
)
|
||||
|
||||
# Convenience methods for direct strategy access
|
||||
def expand_basic(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
|
||||
"""Convenience method for basic expansion."""
|
||||
return self.expand(prompt, strategy=ExpandStrategy.BASIC, context=context)
|
||||
|
||||
def expand_detailed(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
|
||||
"""Convenience method for detailed expansion."""
|
||||
return self.expand(prompt, strategy=ExpandStrategy.DETAILED, context=context)
|
||||
|
||||
def expand_structured(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
|
||||
"""Convenience method for structured expansion."""
|
||||
return self.expand(prompt, strategy=ExpandStrategy.STRUCTURED, context=context)
|
||||
|
||||
def expand_creative(self, prompt: str, context: Optional[str] = None) -> ExpandResult:
|
||||
"""Convenience method for creative expansion."""
|
||||
return self.expand(prompt, strategy=ExpandStrategy.CREATIVE, context=context)
|
||||
@@ -110,20 +110,23 @@ class QueryRewriterAgent:
|
||||
|
||||
# Default prompts for each strategy
|
||||
PROMPTS = {
|
||||
"basic": """You are an expert at rewriting search queries for better retrieval.
|
||||
"basic": """You are an expert at improving prompts and queries.
|
||||
|
||||
Given the following query, rewrite it to be clearer, more specific, and better suited for searching a knowledge base.
|
||||
Given the following input, improve it while preserving its original intent and type.
|
||||
|
||||
Guidelines:
|
||||
- If it's a TASK (e.g., "write X", "create Y", "generate Z"), keep it as a task - do NOT convert to a question
|
||||
- If it's a QUESTION (e.g., "what is X?", "how does Y work?"), keep it as a question
|
||||
- Expand abbreviations and acronyms
|
||||
- Fix typos and grammatical errors
|
||||
- Add relevant context if the query is too short
|
||||
- Transform keyword searches into natural questions
|
||||
- Keep the original intent intact
|
||||
- Add relevant context if too short
|
||||
- Keep the original intent and action intact
|
||||
|
||||
Query: {query}
|
||||
IMPORTANT: If the input is a command/task to DO something, the output must also be a command/task to DO that thing.
|
||||
|
||||
Rewritten query:""",
|
||||
Input: {query}
|
||||
|
||||
Improved version:""",
|
||||
|
||||
"hyde": """You are an expert at generating hypothetical documents.
|
||||
|
||||
@@ -198,21 +201,23 @@ Follow-up query: {query}
|
||||
|
||||
Standalone question:""",
|
||||
|
||||
"auto": """You are an expert at analyzing and rewriting queries.
|
||||
"auto": """You are an expert at analyzing and improving prompts and queries.
|
||||
|
||||
Analyze the following query and determine the best rewriting strategy, then apply it.
|
||||
Analyze the following input and improve it while preserving its original intent and type.
|
||||
|
||||
Query types to consider:
|
||||
1. Short/keyword queries → Expand into full questions
|
||||
2. Ambiguous queries → Clarify with context
|
||||
3. Complex multi-part queries → Break into sub-queries
|
||||
4. Follow-up queries → Resolve references from context
|
||||
5. Well-formed queries → Minor optimization only
|
||||
Input types to consider:
|
||||
1. TASK/COMMAND (write, create, generate, build) → Keep as a task, make it clearer
|
||||
2. QUESTION (what, how, why, when) → Keep as a question, make it more specific
|
||||
3. Short/keyword inputs → Expand with context but preserve type
|
||||
4. Complex multi-part inputs → Clarify structure
|
||||
5. Well-formed inputs → Minor optimization only
|
||||
|
||||
Query: {query}
|
||||
IMPORTANT: Never convert a task into a question. "Write X" should remain "Write X", not "How to write X?"
|
||||
|
||||
Input: {query}
|
||||
{context}
|
||||
|
||||
Analysis and rewritten query:"""
|
||||
Improved version:"""
|
||||
}
|
||||
|
||||
def __init__(
|
||||
|
||||
@@ -2,7 +2,7 @@ import os
|
||||
import time
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict, Optional, List
|
||||
from typing import Any, Dict, Optional, List, Callable
|
||||
from pydantic import BaseModel
|
||||
from rich.text import Text
|
||||
from rich.panel import Panel
|
||||
@@ -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, on_task_start: Optional[Callable] = None, on_task_complete: Optional[Callable] = None, variables: Optional[Dict[str, Any]] = None):
|
||||
# Add check at the start if memory is requested
|
||||
if memory:
|
||||
try:
|
||||
@@ -162,6 +162,11 @@ class PraisonAIAgents:
|
||||
self.stream = stream
|
||||
self.name = name # Store the name for the Agents collection
|
||||
|
||||
# Callbacks for workflow execution
|
||||
self.on_task_start = on_task_start # Called before each task
|
||||
self.on_task_complete = on_task_complete # Called after each task
|
||||
self.variables = variables if variables else {} # Global variables for substitution
|
||||
|
||||
# Check for manager_llm in environment variable if not provided
|
||||
self.manager_llm = manager_llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-5-nano')
|
||||
|
||||
@@ -213,6 +218,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:
|
||||
@@ -651,13 +666,24 @@ Context:
|
||||
|
||||
executor_agent = task.agent
|
||||
|
||||
# Create agent from agent_config if provided and no agent assigned
|
||||
if executor_agent is None and getattr(task, 'agent_config', None):
|
||||
executor_agent = self._create_agent_from_config(task.agent_config)
|
||||
task.agent = executor_agent
|
||||
|
||||
# Set current agent for token tracking
|
||||
llm = getattr(executor_agent, 'llm', None) or getattr(executor_agent, 'llm_instance', None)
|
||||
if llm and hasattr(llm, 'set_current_agent'):
|
||||
llm.set_current_agent(executor_agent.name)
|
||||
|
||||
# Substitute variables in task description if provided
|
||||
task_description = task.description
|
||||
if getattr(task, 'variables', None):
|
||||
for key, value in task.variables.items():
|
||||
task_description = task_description.replace(f"{{{{{key}}}}}", str(value))
|
||||
|
||||
task_prompt = f"""
|
||||
You need to do the following task: {task.description}.
|
||||
You need to do the following task: {task_description}.
|
||||
Expected Output: {task.expected_output}.
|
||||
"""
|
||||
if task.context:
|
||||
@@ -815,6 +841,17 @@ Context:
|
||||
logger.info(f"Task with ID {task_id} is already completed")
|
||||
return
|
||||
|
||||
# Call on_task_start callback if provided
|
||||
if self.on_task_start:
|
||||
try:
|
||||
self.on_task_start(task, task_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_task_start callback: {e}")
|
||||
|
||||
# Apply global variables to task if not already set
|
||||
if self.variables and not getattr(task, 'variables', None):
|
||||
task.variables = self.variables
|
||||
|
||||
retries = 0
|
||||
while task.status != "completed" and retries < self.max_retries:
|
||||
logger.debug(f"Attempt {retries+1} for task {task_id}")
|
||||
@@ -847,6 +884,14 @@ Context:
|
||||
logger.exception(e)
|
||||
|
||||
self.save_output_to_file(task, task_output)
|
||||
|
||||
# Call on_task_complete callback if provided
|
||||
if self.on_task_complete:
|
||||
try:
|
||||
self.on_task_complete(task, task_output)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_task_complete callback: {e}")
|
||||
|
||||
if self.verbose >= 1:
|
||||
logger.info(f"Task {task_id} completed successfully.")
|
||||
else:
|
||||
@@ -929,9 +974,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 +1581,377 @@ 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 _create_agent_from_config(self, agent_config: dict) -> 'Agent':
|
||||
"""
|
||||
Create an Agent from a configuration dictionary.
|
||||
|
||||
Args:
|
||||
agent_config: Dict with keys like 'name', 'role', 'goal', 'backstory', 'llm', 'tools'
|
||||
|
||||
Returns:
|
||||
Agent instance
|
||||
"""
|
||||
from ..agent.agent import Agent
|
||||
|
||||
return Agent(
|
||||
name=agent_config.get('name', 'TaskAgent'),
|
||||
role=agent_config.get('role'),
|
||||
goal=agent_config.get('goal'),
|
||||
backstory=agent_config.get('backstory'),
|
||||
llm=agent_config.get('llm'),
|
||||
tools=agent_config.get('tools'),
|
||||
verbose=self.verbose >= 1,
|
||||
memory=self.memory
|
||||
)
|
||||
|
||||
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("&", "&").replace("<", "<").replace(">", ">")
|
||||
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)}
|
||||
@@ -239,6 +239,10 @@ class LLM:
|
||||
api_key: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
events: List[Any] = [],
|
||||
web_search: Optional[Union[bool, Dict[str, Any]]] = None,
|
||||
web_fetch: Optional[Union[bool, Dict[str, Any]]] = None,
|
||||
prompt_caching: Optional[bool] = None,
|
||||
claude_memory: Optional[Union[bool, Any]] = None,
|
||||
**extra_settings
|
||||
):
|
||||
# Configure logging only once at the class level
|
||||
@@ -273,6 +277,11 @@ class LLM:
|
||||
self.base_url = base_url
|
||||
self.events = events
|
||||
self.extra_settings = extra_settings
|
||||
self.web_search = web_search
|
||||
self.web_fetch = web_fetch
|
||||
self.prompt_caching = prompt_caching
|
||||
self.claude_memory = claude_memory
|
||||
self._claude_memory_tool = None # Lazy initialized
|
||||
self._console = None # Lazy load console when needed
|
||||
self.chat_history = []
|
||||
self.verbose = extra_settings.get('verbose', True)
|
||||
@@ -381,6 +390,109 @@ class LLM:
|
||||
else:
|
||||
return False
|
||||
|
||||
def _supports_web_search(self) -> bool:
|
||||
"""
|
||||
Check if the current model supports native web search via LiteLLM.
|
||||
|
||||
Native web search allows the model to search the web in real-time
|
||||
without requiring external tools like DuckDuckGo.
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports native web search, False otherwise
|
||||
"""
|
||||
from .model_capabilities import supports_web_search
|
||||
return supports_web_search(self.model)
|
||||
|
||||
def _supports_web_fetch(self) -> bool:
|
||||
"""
|
||||
Check if the current model supports web fetch via LiteLLM.
|
||||
|
||||
Web fetch allows the model to retrieve full content from specific URLs
|
||||
(web pages and PDF documents). Currently only supported by Anthropic Claude models.
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports web fetch, False otherwise
|
||||
"""
|
||||
from .model_capabilities import supports_web_fetch
|
||||
return supports_web_fetch(self.model)
|
||||
|
||||
def _supports_prompt_caching(self) -> bool:
|
||||
"""
|
||||
Check if the current model supports prompt caching via LiteLLM.
|
||||
|
||||
Prompt caching allows caching parts of prompts to reduce costs and latency.
|
||||
Supported by OpenAI, Anthropic, Bedrock, and Deepseek.
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports prompt caching, False otherwise
|
||||
"""
|
||||
from .model_capabilities import supports_prompt_caching
|
||||
return supports_prompt_caching(self.model)
|
||||
|
||||
def _get_claude_memory_tool(self):
|
||||
"""
|
||||
Get or initialize the Claude Memory Tool instance.
|
||||
|
||||
Returns:
|
||||
ClaudeMemoryTool instance if claude_memory is enabled, None otherwise
|
||||
"""
|
||||
if not self.claude_memory:
|
||||
return None
|
||||
|
||||
if self._claude_memory_tool is None:
|
||||
from ..tools.claude_memory_tool import ClaudeMemoryTool
|
||||
|
||||
if isinstance(self.claude_memory, bool):
|
||||
# Create default memory tool
|
||||
self._claude_memory_tool = ClaudeMemoryTool()
|
||||
else:
|
||||
# Use provided memory tool instance
|
||||
self._claude_memory_tool = self.claude_memory
|
||||
|
||||
return self._claude_memory_tool
|
||||
|
||||
def _supports_claude_memory(self) -> bool:
|
||||
"""
|
||||
Check if the current model supports Claude Memory Tool.
|
||||
|
||||
Claude Memory Tool is only supported by Anthropic Claude models.
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports Claude Memory Tool
|
||||
"""
|
||||
if not self.claude_memory:
|
||||
return False
|
||||
|
||||
return self._is_anthropic_model()
|
||||
|
||||
def _is_memory_tool_call(self, function_name: str) -> bool:
|
||||
"""
|
||||
Check if a tool call is for the Claude Memory Tool.
|
||||
|
||||
Args:
|
||||
function_name: Name of the function being called
|
||||
|
||||
Returns:
|
||||
bool: True if this is a memory tool call
|
||||
"""
|
||||
return function_name == "memory" and self._supports_claude_memory()
|
||||
|
||||
def _execute_memory_tool_call(self, arguments: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Execute a Claude Memory Tool call.
|
||||
|
||||
Args:
|
||||
arguments: The arguments from Claude's tool call
|
||||
|
||||
Returns:
|
||||
Result string from the memory operation
|
||||
"""
|
||||
memory_tool = self._get_claude_memory_tool()
|
||||
if not memory_tool:
|
||||
return "Error: Memory tool not initialized"
|
||||
|
||||
return memory_tool.process_tool_call(arguments)
|
||||
|
||||
def _generate_ollama_tool_summary(self, tool_results: List[Any], response_text: str) -> Optional[str]:
|
||||
"""
|
||||
Generate a summary from tool results for Ollama to prevent infinite loops.
|
||||
@@ -724,7 +836,21 @@ class LLM:
|
||||
|
||||
# Skip system messages for legacy o1 models as they don't support them
|
||||
if not self._needs_system_message_skip():
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
# Apply prompt caching for Anthropic models if enabled
|
||||
if self.prompt_caching and self._supports_prompt_caching() and self._is_anthropic_model():
|
||||
# Anthropic requires cache_control in content array format
|
||||
messages.append({
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": system_prompt,
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
]
|
||||
})
|
||||
else:
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
|
||||
# Add chat history if provided
|
||||
if chat_history:
|
||||
@@ -2295,6 +2421,12 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
return False
|
||||
return any(prefix in self.model.lower() for prefix in ['gemini', 'gemini/', 'google/gemini'])
|
||||
|
||||
def _is_anthropic_model(self) -> bool:
|
||||
"""Check if the model is an Anthropic Claude model."""
|
||||
if not self.model:
|
||||
return False
|
||||
return any(prefix in self.model.lower() for prefix in ['claude', 'anthropic/', 'anthropic/claude'])
|
||||
|
||||
def _is_streaming_error_recoverable(self, error: Exception) -> bool:
|
||||
"""Check if a streaming error is recoverable (e.g., malformed chunk vs connection error)."""
|
||||
error_msg = str(error).lower()
|
||||
@@ -3171,6 +3303,79 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
logging.debug(f"Could not verify function calling support: {e}. Setting tool_choice anyway.")
|
||||
params['tool_choice'] = 'auto'
|
||||
|
||||
# Add web_search_options if web_search is enabled and model supports it
|
||||
if self.web_search and self._supports_web_search():
|
||||
if isinstance(self.web_search, dict):
|
||||
# Use custom web_search_options provided by user
|
||||
params['web_search_options'] = self.web_search
|
||||
else:
|
||||
# Use default web_search_options
|
||||
params['web_search_options'] = {"search_context_size": "medium"}
|
||||
logging.debug(f"Web search enabled with options: {params['web_search_options']}")
|
||||
|
||||
# Some web search models (like gpt-4o-search-preview) don't support certain parameters
|
||||
# Remove incompatible parameters to avoid BadRequestError
|
||||
incompatible_params = ['temperature', 'top_p', 'presence_penalty', 'frequency_penalty', 'logit_bias']
|
||||
for param in incompatible_params:
|
||||
if param in params:
|
||||
del params[param]
|
||||
logging.debug(f"Removed incompatible parameter '{param}' for web search model")
|
||||
|
||||
# Add web_fetch tool if web_fetch is enabled and model supports it
|
||||
# Web fetch is passed as a tool (unlike web_search which uses web_search_options)
|
||||
if self.web_fetch and self._supports_web_fetch():
|
||||
# Build the web_fetch tool definition
|
||||
web_fetch_tool = {
|
||||
"type": "web_fetch_20250910",
|
||||
"name": "web_fetch",
|
||||
}
|
||||
|
||||
# Add optional parameters from web_fetch dict if provided
|
||||
if isinstance(self.web_fetch, dict):
|
||||
if 'max_uses' in self.web_fetch:
|
||||
web_fetch_tool['max_uses'] = self.web_fetch['max_uses']
|
||||
if 'allowed_domains' in self.web_fetch:
|
||||
web_fetch_tool['allowed_domains'] = self.web_fetch['allowed_domains']
|
||||
if 'blocked_domains' in self.web_fetch:
|
||||
web_fetch_tool['blocked_domains'] = self.web_fetch['blocked_domains']
|
||||
if 'citations' in self.web_fetch:
|
||||
web_fetch_tool['citations'] = self.web_fetch['citations']
|
||||
if 'max_content_tokens' in self.web_fetch:
|
||||
web_fetch_tool['max_content_tokens'] = self.web_fetch['max_content_tokens']
|
||||
else:
|
||||
# Default max_uses if web_fetch=True
|
||||
web_fetch_tool['max_uses'] = 5
|
||||
|
||||
# Add web_fetch tool to existing tools or create tools list
|
||||
if 'tools' in params and params['tools']:
|
||||
params['tools'].append(web_fetch_tool)
|
||||
else:
|
||||
params['tools'] = [web_fetch_tool]
|
||||
|
||||
logging.debug(f"Web fetch enabled with tool: {web_fetch_tool}")
|
||||
|
||||
# Add Claude Memory Tool if enabled and model supports it
|
||||
if self._supports_claude_memory():
|
||||
memory_tool = self._get_claude_memory_tool()
|
||||
if memory_tool:
|
||||
# Add the memory tool definition
|
||||
memory_tool_def = memory_tool.get_tool_definition()
|
||||
|
||||
# Add memory tool to existing tools or create tools list
|
||||
if 'tools' in params and params['tools']:
|
||||
params['tools'].append(memory_tool_def)
|
||||
else:
|
||||
params['tools'] = [memory_tool_def]
|
||||
|
||||
# Add the beta header for Anthropic
|
||||
beta_header = memory_tool.get_beta_header()
|
||||
if 'extra_headers' in params:
|
||||
params['extra_headers']['anthropic-beta'] = beta_header
|
||||
else:
|
||||
params['extra_headers'] = {'anthropic-beta': beta_header}
|
||||
|
||||
logging.debug(f"Claude memory tool enabled with beta header: {beta_header}")
|
||||
|
||||
return params
|
||||
|
||||
def _prepare_response_logging(self, temperature: float, stream: bool, verbose: bool, markdown: bool, **kwargs) -> Optional[Dict[str, Any]]:
|
||||
|
||||
@@ -1,62 +1,30 @@
|
||||
"""
|
||||
Model capabilities configuration for different LLM providers.
|
||||
This module defines which models support specific features like structured outputs.
|
||||
|
||||
This module uses LiteLLM's helper functions as the primary source for model capability detection.
|
||||
LiteLLM is maintained by many contributors and is more accurate and up-to-date.
|
||||
|
||||
LiteLLM Helper Functions:
|
||||
- litellm.supports_web_search(model=) - Check web search support
|
||||
- litellm.supports_function_calling(model=) - Check function calling support
|
||||
- litellm.supports_parallel_function_calling(model=) - Check parallel function calling
|
||||
- litellm.supports_response_schema(model=) - Check JSON schema/structured outputs support
|
||||
- litellm.utils.supports_prompt_caching(model=) - Check prompt caching support
|
||||
- litellm.get_supported_openai_params(model=) - Get all supported params
|
||||
|
||||
Sources:
|
||||
- https://docs.litellm.ai/docs/completion/web_search
|
||||
- https://docs.litellm.ai/docs/completion/json_mode
|
||||
- https://docs.litellm.ai/docs/completion/function_call
|
||||
- https://docs.litellm.ai/docs/completion/prompt_caching
|
||||
"""
|
||||
|
||||
# Models that support OpenAI-style structured outputs (response_format with Pydantic models)
|
||||
MODELS_SUPPORTING_STRUCTURED_OUTPUTS = {
|
||||
# OpenAI models
|
||||
"gpt-5-nano",
|
||||
"gpt-5-nano",
|
||||
"gpt-4-turbo",
|
||||
"gpt-4-turbo-preview",
|
||||
"gpt-4-turbo-2024-04-09",
|
||||
"gpt-4-1106-preview",
|
||||
"gpt-4-0125-preview",
|
||||
"gpt-3.5-turbo",
|
||||
"gpt-3.5-turbo-1106",
|
||||
"gpt-3.5-turbo-0125",
|
||||
|
||||
# New/Future OpenAI models (as mentioned by user)
|
||||
"codex-mini",
|
||||
"o3-pro",
|
||||
"gpt-4.5-preview",
|
||||
"o3-mini",
|
||||
"o1",
|
||||
"o1-preview",
|
||||
"o1-mini",
|
||||
"gpt-4.1",
|
||||
"gpt-4.1-nano",
|
||||
"gpt-4.1-mini",
|
||||
"o4-mini",
|
||||
"o3",
|
||||
|
||||
# Gemini models that support structured outputs
|
||||
"gemini-2.0-flash",
|
||||
"gemini-2.0-flash-exp",
|
||||
"gemini-1.5-pro",
|
||||
"gemini-1.5-pro-latest",
|
||||
"gemini-1.5-flash",
|
||||
"gemini-1.5-flash-latest",
|
||||
"gemini-1.5-flash-8b",
|
||||
"gemini-1.5-flash-8b-latest",
|
||||
}
|
||||
|
||||
# Models that explicitly DON'T support structured outputs
|
||||
MODELS_NOT_SUPPORTING_STRUCTURED_OUTPUTS = {
|
||||
# Audio preview models
|
||||
"gpt-4o-audio-preview",
|
||||
"gpt-5-nano-audio-preview",
|
||||
|
||||
# Legacy o1 models (don't support system messages either)
|
||||
"o1-preview-2024-09-12",
|
||||
"o1-mini-2024-09-12",
|
||||
}
|
||||
|
||||
|
||||
def supports_structured_outputs(model_name: str) -> bool:
|
||||
"""
|
||||
Check if a model supports OpenAI-style structured outputs.
|
||||
Check if a model supports structured outputs (JSON schema).
|
||||
|
||||
Uses LiteLLM's supports_response_schema() as the primary check.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check
|
||||
@@ -67,34 +35,72 @@ def supports_structured_outputs(model_name: str) -> bool:
|
||||
if not model_name:
|
||||
return False
|
||||
|
||||
# Strip provider prefixes (e.g., 'google/', 'openai/', etc.)
|
||||
model_without_provider = model_name
|
||||
for prefix in ['google/', 'openai/', 'anthropic/', 'gemini/', 'mistral/', 'deepseek/', 'groq/']:
|
||||
if model_name.startswith(prefix):
|
||||
model_without_provider = model_name[len(prefix):]
|
||||
break
|
||||
try:
|
||||
import litellm
|
||||
# Use LiteLLM's built-in check - most accurate and up-to-date
|
||||
if hasattr(litellm, 'supports_response_schema'):
|
||||
return litellm.supports_response_schema(model=model_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# First check if it's explicitly in the NOT supporting list
|
||||
if model_without_provider in MODELS_NOT_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return False
|
||||
|
||||
|
||||
def supports_function_calling(model_name: str) -> bool:
|
||||
"""
|
||||
Check if a model supports function calling.
|
||||
|
||||
Uses LiteLLM's supports_function_calling() as the primary check.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports function calling, False otherwise
|
||||
"""
|
||||
if not model_name:
|
||||
return False
|
||||
|
||||
# Then check if it's in the supporting list
|
||||
if model_without_provider in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return True
|
||||
try:
|
||||
import litellm
|
||||
# Use LiteLLM's built-in check - most accurate and up-to-date
|
||||
if hasattr(litellm, 'supports_function_calling'):
|
||||
return litellm.supports_function_calling(model=model_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# For models with version suffixes, check the base model name
|
||||
base_model = model_without_provider.split('-2024-')[0].split('-2025-')[0]
|
||||
if base_model in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def supports_parallel_function_calling(model_name: str) -> bool:
|
||||
"""
|
||||
Check if a model supports parallel function calling.
|
||||
|
||||
Uses LiteLLM's supports_parallel_function_calling() as the primary check.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports parallel function calling, False otherwise
|
||||
"""
|
||||
if not model_name:
|
||||
return False
|
||||
|
||||
try:
|
||||
import litellm
|
||||
# Use LiteLLM's built-in check - most accurate and up-to-date
|
||||
if hasattr(litellm, 'supports_parallel_function_calling'):
|
||||
return litellm.supports_parallel_function_calling(model=model_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Default to False for unknown models
|
||||
return False
|
||||
|
||||
|
||||
def supports_streaming_with_tools(model_name: str) -> bool:
|
||||
"""
|
||||
Check if a model supports streaming when tools are provided.
|
||||
Most models that support structured outputs also support streaming with tools.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check
|
||||
@@ -102,15 +108,153 @@ def supports_streaming_with_tools(model_name: str) -> bool:
|
||||
Returns:
|
||||
bool: True if the model supports streaming with tools, False otherwise
|
||||
"""
|
||||
# For now, use the same logic as structured outputs
|
||||
# In the future, this could be a separate list if needed
|
||||
return supports_structured_outputs(model_name)
|
||||
# Models that support function calling generally support streaming with tools
|
||||
return supports_function_calling(model_name)
|
||||
|
||||
|
||||
# Supported Gemini internal tools
|
||||
GEMINI_INTERNAL_TOOLS = {'googleSearch', 'urlContext', 'codeExecution'}
|
||||
|
||||
|
||||
def supports_web_search(model_name: str) -> bool:
|
||||
"""
|
||||
Check if a model supports native web search via LiteLLM.
|
||||
|
||||
Uses LiteLLM's supports_web_search() as the primary check.
|
||||
|
||||
Native web search allows the model to search the web in real-time
|
||||
without requiring external tools like DuckDuckGo.
|
||||
|
||||
Supported providers:
|
||||
- OpenAI (gpt-4o-search-preview, gpt-4o-mini-search-preview)
|
||||
- xAI (grok-3)
|
||||
- Anthropic (claude-3-5-sonnet-latest, claude-sonnet-4, etc.)
|
||||
- Google/Vertex AI (gemini-2.0-flash, gemini-2.5-*, etc.)
|
||||
- Perplexity (all models)
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check (with or without provider prefix)
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports native web search, False otherwise
|
||||
"""
|
||||
if not model_name:
|
||||
return False
|
||||
|
||||
try:
|
||||
import litellm
|
||||
# Use LiteLLM's built-in check - most accurate and up-to-date
|
||||
if hasattr(litellm, 'supports_web_search'):
|
||||
return litellm.supports_web_search(model=model_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def supports_prompt_caching(model_name: str) -> bool:
|
||||
"""
|
||||
Check if a model supports prompt caching via LiteLLM.
|
||||
|
||||
Uses LiteLLM's supports_prompt_caching() as the primary check.
|
||||
|
||||
Prompt caching allows caching parts of prompts to reduce costs and latency
|
||||
on subsequent requests with similar prompts.
|
||||
|
||||
Supported providers:
|
||||
- OpenAI (openai/) - Automatic caching for prompts ≥1024 tokens
|
||||
- Anthropic (anthropic/) - Manual caching with cache_control
|
||||
- Bedrock (bedrock/) - All models that support prompt caching
|
||||
- Deepseek (deepseek/) - Works like OpenAI
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check (with or without provider prefix)
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports prompt caching, False otherwise
|
||||
"""
|
||||
if not model_name:
|
||||
return False
|
||||
|
||||
try:
|
||||
from litellm.utils import supports_prompt_caching as litellm_supports_prompt_caching
|
||||
return litellm_supports_prompt_caching(model=model_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
# Models that support web fetch via LiteLLM (Anthropic only)
|
||||
# Web fetch retrieves full content from specific URLs (web pages and PDFs)
|
||||
# Source: https://docs.litellm.ai/docs/completion/web_fetch
|
||||
# Note: LiteLLM doesn't have a supports_web_fetch() helper yet, so we maintain a static list
|
||||
MODELS_SUPPORTING_WEB_FETCH = {
|
||||
# Anthropic Claude models with web fetch support
|
||||
"claude-opus-4-1-20250805",
|
||||
"claude-opus-4-1",
|
||||
"claude-opus-4-20250514",
|
||||
"claude-opus-4",
|
||||
"claude-sonnet-4-20250514",
|
||||
"claude-sonnet-4",
|
||||
"claude-3-7-sonnet-20250219",
|
||||
"claude-3-7-sonnet-latest",
|
||||
"claude-3-5-sonnet-latest",
|
||||
"claude-3-5-sonnet-20241022",
|
||||
"claude-3-5-haiku-latest",
|
||||
"claude-3-5-haiku-20241022",
|
||||
}
|
||||
|
||||
|
||||
def supports_web_fetch(model_name: str) -> bool:
|
||||
"""
|
||||
Check if a model supports web fetch via LiteLLM.
|
||||
|
||||
Web fetch allows the model to retrieve full content from specific URLs
|
||||
(web pages and PDF documents). Currently only supported by Anthropic Claude models.
|
||||
|
||||
Note: LiteLLM doesn't have a supports_web_fetch() helper yet, so we use a static list
|
||||
with auto-detection for Claude 4+ models.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check (with or without provider prefix)
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports web fetch, False otherwise
|
||||
"""
|
||||
if not model_name:
|
||||
return False
|
||||
|
||||
# Strip provider prefixes
|
||||
model_without_provider = model_name
|
||||
for prefix in ['anthropic/', 'bedrock/', 'vertex_ai/']:
|
||||
if model_name.startswith(prefix):
|
||||
model_without_provider = model_name[len(prefix):]
|
||||
break
|
||||
|
||||
# Check our static list
|
||||
if model_without_provider in MODELS_SUPPORTING_WEB_FETCH:
|
||||
return True
|
||||
|
||||
# Check base model name (without version suffix)
|
||||
base_model = model_without_provider.split('-2024-')[0].split('-2025-')[0]
|
||||
if base_model in MODELS_SUPPORTING_WEB_FETCH:
|
||||
return True
|
||||
|
||||
# Auto-support for Claude 4+ models (Anthropic only)
|
||||
model_lower = model_without_provider.lower()
|
||||
if 'claude' in model_lower:
|
||||
import re
|
||||
# Match patterns like claude-4, claude-5, claude-sonnet-4, claude-opus-4, etc.
|
||||
version_match = re.search(r'claude-(?:sonnet-|opus-|haiku-)?(\d+)', model_lower)
|
||||
if version_match:
|
||||
version = int(version_match.group(1))
|
||||
if version >= 4: # Claude 4 and later
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def is_gemini_internal_tool(tool) -> bool:
|
||||
"""
|
||||
Check if a tool is a Gemini internal tool and should be included in formatted tools.
|
||||
|
||||
@@ -2,8 +2,65 @@
|
||||
Model Context Protocol (MCP) integration for PraisonAI Agents.
|
||||
|
||||
This package provides classes and utilities for connecting to MCP servers
|
||||
using different transport methods (stdio, SSE, etc.).
|
||||
using different transport methods:
|
||||
- stdio: Standard input/output (subprocess)
|
||||
- sse: Server-Sent Events (legacy HTTP+SSE)
|
||||
- http_stream: Streamable HTTP (current standard)
|
||||
- websocket: WebSocket (SEP-1288)
|
||||
|
||||
Protocol Revision: 2025-11-25
|
||||
"""
|
||||
from .mcp import MCP
|
||||
|
||||
__all__ = ["MCP"]
|
||||
# Lazy imports for optional components
|
||||
def __getattr__(name):
|
||||
"""Lazy import for optional modules."""
|
||||
if name == "WebSocketMCPClient":
|
||||
from .mcp_websocket import WebSocketMCPClient
|
||||
return WebSocketMCPClient
|
||||
elif name == "SessionManager":
|
||||
from .mcp_session import SessionManager
|
||||
return SessionManager
|
||||
elif name == "SecurityConfig":
|
||||
from .mcp_security import SecurityConfig
|
||||
return SecurityConfig
|
||||
elif name == "BaseTransport":
|
||||
from .mcp_transport import BaseTransport
|
||||
return BaseTransport
|
||||
elif name == "TransportConfig":
|
||||
from .mcp_transport import TransportConfig
|
||||
return TransportConfig
|
||||
elif name == "ToolsMCPServer":
|
||||
from .mcp_server import ToolsMCPServer
|
||||
return ToolsMCPServer
|
||||
elif name == "launch_tools_mcp_server":
|
||||
from .mcp_server import launch_tools_mcp_server
|
||||
return launch_tools_mcp_server
|
||||
elif name == "function_to_mcp_schema":
|
||||
from .mcp_utils import function_to_mcp_schema
|
||||
return function_to_mcp_schema
|
||||
elif name == "get_tool_metadata":
|
||||
from .mcp_utils import get_tool_metadata
|
||||
return get_tool_metadata
|
||||
elif name == "python_type_to_json_schema":
|
||||
from .mcp_utils import python_type_to_json_schema
|
||||
return python_type_to_json_schema
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
__all__ = [
|
||||
# Client
|
||||
"MCP",
|
||||
# Server (lazy loaded)
|
||||
"ToolsMCPServer",
|
||||
"launch_tools_mcp_server",
|
||||
# Utilities (lazy loaded)
|
||||
"function_to_mcp_schema",
|
||||
"get_tool_metadata",
|
||||
"python_type_to_json_schema",
|
||||
# Transport components (lazy loaded)
|
||||
"WebSocketMCPClient",
|
||||
"SessionManager",
|
||||
"SecurityConfig",
|
||||
"BaseTransport",
|
||||
"TransportConfig",
|
||||
]
|
||||
@@ -227,6 +227,26 @@ class MCP:
|
||||
self.timeout = timeout
|
||||
self.debug = debug
|
||||
|
||||
# Check if this is a WebSocket URL (ws:// or wss://)
|
||||
if isinstance(command_or_string, str) and re.match(r'^wss?://', command_or_string):
|
||||
from .mcp_websocket import WebSocketMCPClient
|
||||
# Extract auth token if provided
|
||||
auth_token = kwargs.pop('auth_token', None)
|
||||
|
||||
self.websocket_client = WebSocketMCPClient(
|
||||
command_or_string,
|
||||
debug=debug,
|
||||
timeout=timeout,
|
||||
auth_token=auth_token,
|
||||
options=kwargs
|
||||
)
|
||||
self._tools = list(self.websocket_client.tools)
|
||||
self.is_sse = False
|
||||
self.is_http_stream = False
|
||||
self.is_websocket = True
|
||||
self.is_npx = False
|
||||
return
|
||||
|
||||
# Check if this is an HTTP URL
|
||||
if isinstance(command_or_string, str) and re.match(r'^https?://', command_or_string):
|
||||
# Determine transport type based on URL or kwargs
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
"""
|
||||
MCP Backward Compatibility Module.
|
||||
|
||||
This module provides utilities for maintaining backward compatibility
|
||||
with older MCP servers and the deprecated HTTP+SSE transport from
|
||||
protocol version 2024-11-05.
|
||||
|
||||
Per MCP specification, clients wanting to support older servers should:
|
||||
1. Attempt POST to server URL with InitializeRequest
|
||||
2. If it fails with 400, 404, or 405, fall back to legacy SSE
|
||||
3. Issue GET to server URL expecting SSE stream with endpoint event
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def detect_transport_support(
|
||||
response_status: int,
|
||||
content_type: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Detect which transport a server supports based on response.
|
||||
|
||||
Per MCP spec backward compatibility guide:
|
||||
- Success (200) with JSON indicates Streamable HTTP
|
||||
- 400, 404, 405 indicates need for legacy SSE fallback
|
||||
|
||||
Args:
|
||||
response_status: HTTP response status code
|
||||
content_type: Response Content-Type header
|
||||
|
||||
Returns:
|
||||
Transport type: "http_stream" or "sse"
|
||||
"""
|
||||
# Check for fallback indicators
|
||||
if response_status in (400, 404, 405):
|
||||
return "sse"
|
||||
|
||||
# Successful response indicates Streamable HTTP support
|
||||
if response_status == 200:
|
||||
return "http_stream"
|
||||
|
||||
# Default to SSE for unknown cases
|
||||
return "sse"
|
||||
|
||||
|
||||
def is_legacy_sse_url(url: str) -> bool:
|
||||
"""
|
||||
Check if URL indicates legacy SSE transport.
|
||||
|
||||
Legacy SSE URLs typically end with /sse.
|
||||
|
||||
Args:
|
||||
url: Server URL
|
||||
|
||||
Returns:
|
||||
True if URL indicates legacy SSE
|
||||
"""
|
||||
return url.rstrip('/').endswith('/sse')
|
||||
|
||||
|
||||
def get_legacy_sse_endpoint(base_url: str) -> str:
|
||||
"""
|
||||
Get the legacy SSE endpoint URL.
|
||||
|
||||
Args:
|
||||
base_url: Base server URL
|
||||
|
||||
Returns:
|
||||
SSE endpoint URL
|
||||
"""
|
||||
base = base_url.rstrip('/')
|
||||
if not base.endswith('/sse'):
|
||||
return f"{base}/sse"
|
||||
return base
|
||||
|
||||
|
||||
def get_streamable_http_endpoint(base_url: str) -> str:
|
||||
"""
|
||||
Get the Streamable HTTP endpoint URL.
|
||||
|
||||
Args:
|
||||
base_url: Base server URL
|
||||
|
||||
Returns:
|
||||
MCP endpoint URL
|
||||
"""
|
||||
base = base_url.rstrip('/')
|
||||
# Remove /sse suffix if present
|
||||
if base.endswith('/sse'):
|
||||
base = base[:-4]
|
||||
if not base.endswith('/mcp'):
|
||||
return f"{base}/mcp"
|
||||
return base
|
||||
|
||||
|
||||
class TransportNegotiator:
|
||||
"""
|
||||
Handles automatic transport negotiation with MCP servers.
|
||||
|
||||
This class implements the backward compatibility algorithm
|
||||
from the MCP specification for detecting server capabilities.
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: str):
|
||||
"""
|
||||
Initialize negotiator.
|
||||
|
||||
Args:
|
||||
base_url: Server base URL
|
||||
"""
|
||||
self.base_url = base_url
|
||||
self._detected_transport: Optional[str] = None
|
||||
|
||||
@property
|
||||
def detected_transport(self) -> Optional[str]:
|
||||
"""Get the detected transport type."""
|
||||
return self._detected_transport
|
||||
|
||||
def set_detected_transport(self, transport: str) -> None:
|
||||
"""
|
||||
Set the detected transport type.
|
||||
|
||||
Args:
|
||||
transport: Transport type ("http_stream" or "sse")
|
||||
"""
|
||||
self._detected_transport = transport
|
||||
|
||||
def get_endpoint_url(self) -> str:
|
||||
"""
|
||||
Get the appropriate endpoint URL based on detected transport.
|
||||
|
||||
Returns:
|
||||
Endpoint URL for the detected transport
|
||||
"""
|
||||
if self._detected_transport == "sse":
|
||||
return get_legacy_sse_endpoint(self.base_url)
|
||||
return get_streamable_http_endpoint(self.base_url)
|
||||
|
||||
def should_try_streamable_first(self) -> bool:
|
||||
"""
|
||||
Determine if Streamable HTTP should be tried first.
|
||||
|
||||
Per MCP spec, clients should try Streamable HTTP first
|
||||
unless the URL explicitly indicates legacy SSE.
|
||||
|
||||
Returns:
|
||||
True if should try Streamable HTTP first
|
||||
"""
|
||||
return not is_legacy_sse_url(self.base_url)
|
||||
@@ -200,16 +200,31 @@ class HTTPStreamTransport:
|
||||
|
||||
This transport provides a single endpoint for all MCP communication,
|
||||
supporting both batch (JSON) and streaming (SSE) response modes.
|
||||
|
||||
Per MCP Protocol Revision 2025-11-25:
|
||||
- Includes MCP-Protocol-Version header on all requests
|
||||
- Handles session management via Mcp-Session-Id header
|
||||
- Supports SSE resumability via Last-Event-ID
|
||||
"""
|
||||
|
||||
# Default protocol version for backward compatibility
|
||||
DEFAULT_PROTOCOL_VERSION = '2025-03-26'
|
||||
|
||||
def __init__(self, base_url: str, session_id: Optional[str] = None, options: Optional[Dict[str, Any]] = None):
|
||||
self.base_url = base_url
|
||||
self.session_id = session_id
|
||||
self.options = options or {}
|
||||
self.response_mode = self.options.get('responseMode', 'batch')
|
||||
self.protocol_version = self.options.get('protocol_version', self.DEFAULT_PROTOCOL_VERSION)
|
||||
|
||||
# Track last event ID for resumability
|
||||
self.last_event_id: Optional[str] = None
|
||||
self._retry_delay_ms: int = 3000 # Default retry delay
|
||||
|
||||
self.headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Accept': 'application/json, text/event-stream'
|
||||
'Accept': 'application/json, text/event-stream',
|
||||
'Mcp-Protocol-Version': self.protocol_version # Required per spec
|
||||
}
|
||||
if session_id:
|
||||
self.headers['Mcp-Session-Id'] = session_id
|
||||
@@ -281,7 +296,7 @@ class HTTPStreamTransport:
|
||||
raise
|
||||
|
||||
async def _process_sse_response(self, response):
|
||||
"""Process SSE response stream."""
|
||||
"""Process SSE response stream with resumability support."""
|
||||
buffer = ""
|
||||
async for chunk in response.content:
|
||||
buffer += chunk.decode('utf-8')
|
||||
@@ -291,12 +306,34 @@ class HTTPStreamTransport:
|
||||
event, buffer = buffer.split("\n\n", 1)
|
||||
lines = event.strip().split("\n")
|
||||
|
||||
# Parse SSE event
|
||||
data = None
|
||||
for line in lines:
|
||||
if line.startswith("data: "):
|
||||
data = line[6:] # Remove "data: " prefix
|
||||
# Parse SSE event fields per spec
|
||||
event_id = None
|
||||
data_lines = []
|
||||
retry_value = None
|
||||
|
||||
for line in lines:
|
||||
if line.startswith("id:"):
|
||||
# Track event ID for resumability
|
||||
event_id = line[3:].strip()
|
||||
elif line.startswith("data:"):
|
||||
data_lines.append(line[5:].strip())
|
||||
elif line.startswith("retry:"):
|
||||
# Handle retry field per SSE spec
|
||||
try:
|
||||
retry_value = int(line[6:].strip())
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Update last event ID for resumability
|
||||
if event_id:
|
||||
self.last_event_id = event_id
|
||||
|
||||
# Update retry delay if provided
|
||||
if retry_value is not None:
|
||||
self._retry_delay_ms = retry_value
|
||||
|
||||
# Process data if present
|
||||
data = '\n'.join(data_lines) if data_lines else None
|
||||
if data:
|
||||
try:
|
||||
message = json.loads(data)
|
||||
@@ -310,13 +347,49 @@ class HTTPStreamTransport:
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"Failed to parse SSE event: {data}")
|
||||
|
||||
async def terminate_session(self) -> bool:
|
||||
"""
|
||||
Terminate the session via HTTP DELETE.
|
||||
|
||||
Per MCP spec: Clients that no longer need a particular session
|
||||
SHOULD send an HTTP DELETE to the MCP endpoint with the
|
||||
Mcp-Session-Id header to explicitly terminate the session.
|
||||
|
||||
Returns:
|
||||
True if session was terminated, False if server doesn't support it
|
||||
"""
|
||||
if not self._session or not self.session_id:
|
||||
return False
|
||||
|
||||
try:
|
||||
headers = {
|
||||
'Mcp-Session-Id': self.session_id,
|
||||
'Mcp-Protocol-Version': self.protocol_version
|
||||
}
|
||||
async with self._session.delete(self.base_url, headers=headers) as response:
|
||||
if response.status == 405:
|
||||
# Server doesn't allow client-initiated session termination
|
||||
logger.debug("Server does not allow session termination via DELETE")
|
||||
return False
|
||||
elif response.status in (200, 202, 204):
|
||||
# Session terminated successfully
|
||||
self.session_id = None
|
||||
self.headers.pop('Mcp-Session-Id', None)
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"Unexpected response to DELETE: {response.status}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Error terminating session: {e}")
|
||||
return False
|
||||
|
||||
async def start_sse_listener(self):
|
||||
"""Start listening for SSE events from the server."""
|
||||
if self._sse_task is None or self._sse_task.done():
|
||||
self._sse_task = asyncio.create_task(self._sse_listener())
|
||||
|
||||
async def _sse_listener(self):
|
||||
"""Background task to listen for SSE events."""
|
||||
"""Background task to listen for SSE events with resumability support."""
|
||||
while True:
|
||||
try:
|
||||
# Check if we should stop
|
||||
@@ -324,18 +397,28 @@ class HTTPStreamTransport:
|
||||
break
|
||||
|
||||
url = self.base_url
|
||||
if self.session_id:
|
||||
# Add session as query parameter for SSE connection
|
||||
url = f"{url}?session={self.session_id}"
|
||||
|
||||
# Build headers per MCP spec
|
||||
headers = {
|
||||
'Accept': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache'
|
||||
'Cache-Control': 'no-cache',
|
||||
'Mcp-Protocol-Version': self.protocol_version # Required per spec
|
||||
}
|
||||
if self.session_id:
|
||||
headers['Mcp-Session-Id'] = self.session_id
|
||||
|
||||
# Include Last-Event-ID for resumability per spec
|
||||
if self.last_event_id:
|
||||
headers['Last-Event-ID'] = self.last_event_id
|
||||
|
||||
async with self._session.get(url, headers=headers) as response:
|
||||
# Handle session expiration (HTTP 404)
|
||||
if response.status == 404:
|
||||
logger.warning("Session expired (HTTP 404), need to reinitialize")
|
||||
self.session_id = None
|
||||
self.headers.pop('Mcp-Session-Id', None)
|
||||
break
|
||||
|
||||
buffer = ""
|
||||
async for chunk in response.content:
|
||||
# Check if we should stop
|
||||
@@ -349,12 +432,32 @@ class HTTPStreamTransport:
|
||||
event, buffer = buffer.split("\n\n", 1)
|
||||
lines = event.strip().split("\n")
|
||||
|
||||
# Parse SSE event
|
||||
data = None
|
||||
for line in lines:
|
||||
if line.startswith("data: "):
|
||||
data = line[6:] # Remove "data: " prefix
|
||||
# Parse SSE event fields
|
||||
event_id = None
|
||||
data_lines = []
|
||||
retry_value = None
|
||||
|
||||
for line in lines:
|
||||
if line.startswith("id:"):
|
||||
event_id = line[3:].strip()
|
||||
elif line.startswith("data:"):
|
||||
data_lines.append(line[5:].strip())
|
||||
elif line.startswith("retry:"):
|
||||
try:
|
||||
retry_value = int(line[6:].strip())
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Update last event ID for resumability
|
||||
if event_id:
|
||||
self.last_event_id = event_id
|
||||
|
||||
# Update retry delay if provided
|
||||
if retry_value is not None:
|
||||
self._retry_delay_ms = retry_value
|
||||
|
||||
# Process data
|
||||
data = '\n'.join(data_lines) if data_lines else None
|
||||
if data:
|
||||
try:
|
||||
message = json.loads(data)
|
||||
@@ -368,7 +471,8 @@ class HTTPStreamTransport:
|
||||
except Exception as e:
|
||||
if not (hasattr(self, '_closing') and self._closing):
|
||||
logger.error(f"SSE listener error: {e}")
|
||||
await asyncio.sleep(1) # Reconnect after 1 second
|
||||
# Use retry delay from server or default
|
||||
await asyncio.sleep(self._retry_delay_ms / 1000.0)
|
||||
else:
|
||||
break
|
||||
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
"""
|
||||
MCP Security Module.
|
||||
|
||||
This module provides security utilities for MCP transports,
|
||||
implementing security best practices from MCP Protocol
|
||||
Revision 2025-11-25.
|
||||
|
||||
Security features:
|
||||
- Origin header validation (DNS rebinding prevention)
|
||||
- Localhost binding recommendations
|
||||
- Authentication header support
|
||||
- Secure session ID generation
|
||||
"""
|
||||
|
||||
import secrets
|
||||
import base64
|
||||
from typing import Optional, Dict, List, Set
|
||||
from urllib.parse import urlparse
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
# Localhost addresses
|
||||
LOCALHOST_ADDRESSES: Set[str] = {
|
||||
'localhost',
|
||||
'127.0.0.1',
|
||||
'::1', # IPv6 localhost
|
||||
}
|
||||
|
||||
|
||||
def is_valid_origin(
|
||||
origin: Optional[str],
|
||||
allowed_origins: List[str],
|
||||
allow_missing: bool = False
|
||||
) -> bool:
|
||||
"""
|
||||
Validate Origin header against allowed origins.
|
||||
|
||||
Per MCP spec, servers MUST validate the Origin header on all
|
||||
incoming connections to prevent DNS rebinding attacks.
|
||||
|
||||
Args:
|
||||
origin: The Origin header value
|
||||
allowed_origins: List of allowed origin hosts
|
||||
allow_missing: Whether to allow missing Origin header
|
||||
|
||||
Returns:
|
||||
True if origin is valid, False otherwise
|
||||
"""
|
||||
if not origin:
|
||||
return allow_missing
|
||||
|
||||
# Extract host from origin
|
||||
host = extract_origin_host(origin)
|
||||
if not host:
|
||||
return False
|
||||
|
||||
# Check against allowed origins
|
||||
return host in allowed_origins
|
||||
|
||||
|
||||
def extract_origin_host(origin: str) -> Optional[str]:
|
||||
"""
|
||||
Extract the host from an Origin header value.
|
||||
|
||||
Args:
|
||||
origin: Origin header value (e.g., "https://example.com:443")
|
||||
|
||||
Returns:
|
||||
Host portion of the origin, or None if invalid
|
||||
"""
|
||||
if not origin:
|
||||
return None
|
||||
|
||||
try:
|
||||
parsed = urlparse(origin)
|
||||
return parsed.hostname
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def is_localhost_address(address: str) -> bool:
|
||||
"""
|
||||
Check if an address is a localhost address.
|
||||
|
||||
Args:
|
||||
address: IP address or hostname
|
||||
|
||||
Returns:
|
||||
True if localhost, False otherwise
|
||||
"""
|
||||
return address.lower() in LOCALHOST_ADDRESSES
|
||||
|
||||
|
||||
def should_bind_localhost_only(is_local: bool = True) -> bool:
|
||||
"""
|
||||
Determine if server should bind to localhost only.
|
||||
|
||||
Per MCP spec, when running locally, servers SHOULD bind only
|
||||
to localhost (127.0.0.1) rather than all network interfaces.
|
||||
|
||||
Args:
|
||||
is_local: Whether this is a local development server
|
||||
|
||||
Returns:
|
||||
True if should bind to localhost only
|
||||
"""
|
||||
return is_local
|
||||
|
||||
|
||||
def is_potential_dns_rebinding(origin: str, server_host: str) -> bool:
|
||||
"""
|
||||
Detect potential DNS rebinding attack.
|
||||
|
||||
DNS rebinding occurs when an external origin tries to access
|
||||
a local server through DNS manipulation.
|
||||
|
||||
Args:
|
||||
origin: The Origin header value
|
||||
server_host: The host the server is bound to
|
||||
|
||||
Returns:
|
||||
True if potential DNS rebinding attempt
|
||||
"""
|
||||
origin_host = extract_origin_host(origin)
|
||||
if not origin_host:
|
||||
return True # Treat invalid origin as suspicious
|
||||
|
||||
# If server is on localhost, only localhost origins are safe
|
||||
if is_localhost_address(server_host):
|
||||
return not is_localhost_address(origin_host)
|
||||
|
||||
# For non-localhost servers, check if origin matches server
|
||||
return origin_host != server_host
|
||||
|
||||
|
||||
def create_auth_header(
|
||||
credentials: str,
|
||||
auth_type: str = "bearer",
|
||||
header_name: str = "Authorization"
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
Create an authentication header.
|
||||
|
||||
Args:
|
||||
credentials: The credentials/token
|
||||
auth_type: Type of auth ("bearer", "basic", "custom")
|
||||
header_name: Custom header name (for auth_type="custom")
|
||||
|
||||
Returns:
|
||||
Dictionary with authentication header
|
||||
"""
|
||||
if auth_type.lower() == "bearer":
|
||||
return {"Authorization": f"Bearer {credentials}"}
|
||||
|
||||
elif auth_type.lower() == "basic":
|
||||
# Base64 encode for Basic auth
|
||||
encoded = base64.b64encode(credentials.encode()).decode()
|
||||
return {"Authorization": f"Basic {encoded}"}
|
||||
|
||||
elif auth_type.lower() == "custom":
|
||||
return {header_name: credentials}
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown auth_type: {auth_type}")
|
||||
|
||||
|
||||
def validate_auth_header(headers: Dict[str, str]) -> bool:
|
||||
"""
|
||||
Validate that authentication header is present.
|
||||
|
||||
Args:
|
||||
headers: Request headers dictionary
|
||||
|
||||
Returns:
|
||||
True if Authorization header is present
|
||||
"""
|
||||
# Check for Authorization header (case-insensitive)
|
||||
for key in headers:
|
||||
if key.lower() == 'authorization':
|
||||
return bool(headers[key])
|
||||
return False
|
||||
|
||||
|
||||
def generate_secure_session_id(length: int = 32) -> str:
|
||||
"""
|
||||
Generate a cryptographically secure session ID.
|
||||
|
||||
Per MCP spec, session IDs SHOULD be globally unique and
|
||||
cryptographically secure (e.g., securely generated UUID,
|
||||
JWT, or cryptographic hash).
|
||||
|
||||
The ID MUST only contain visible ASCII characters (0x21-0x7E).
|
||||
|
||||
Args:
|
||||
length: Desired length of session ID (minimum 32)
|
||||
|
||||
Returns:
|
||||
Secure session ID string
|
||||
"""
|
||||
# Ensure minimum length for security
|
||||
length = max(length, 32)
|
||||
|
||||
# Use URL-safe base64 which only uses visible ASCII
|
||||
# Generate extra bytes to account for base64 expansion
|
||||
num_bytes = (length * 3) // 4 + 1
|
||||
random_bytes = secrets.token_bytes(num_bytes)
|
||||
|
||||
# Encode to URL-safe base64 (uses A-Z, a-z, 0-9, -, _)
|
||||
session_id = base64.urlsafe_b64encode(random_bytes).decode()
|
||||
|
||||
# Trim to desired length
|
||||
return session_id[:length]
|
||||
|
||||
|
||||
@dataclass
|
||||
class SecurityConfig:
|
||||
"""
|
||||
Security configuration for MCP transports.
|
||||
|
||||
This class holds security settings that can be applied
|
||||
to MCP server implementations.
|
||||
|
||||
Attributes:
|
||||
validate_origin: Whether to validate Origin header
|
||||
allow_missing_origin: Whether to allow missing Origin
|
||||
allowed_origins: List of allowed origin hosts
|
||||
require_auth: Whether authentication is required
|
||||
bind_localhost_only: Whether to bind to localhost only
|
||||
"""
|
||||
validate_origin: bool = True
|
||||
allow_missing_origin: bool = False
|
||||
allowed_origins: List[str] = field(default_factory=lambda: ["localhost", "127.0.0.1"])
|
||||
require_auth: bool = False
|
||||
bind_localhost_only: bool = True
|
||||
|
||||
def is_origin_allowed(self, origin: Optional[str]) -> bool:
|
||||
"""Check if an origin is allowed by this config."""
|
||||
if not self.validate_origin:
|
||||
return True
|
||||
return is_valid_origin(
|
||||
origin,
|
||||
self.allowed_origins,
|
||||
self.allow_missing_origin
|
||||
)
|
||||
|
||||
def get_bind_address(self) -> str:
|
||||
"""Get recommended bind address based on config."""
|
||||
return "127.0.0.1" if self.bind_localhost_only else "0.0.0.0"
|
||||
@@ -0,0 +1,356 @@
|
||||
"""MCP Server wrapper for exposing PraisonAI tools as MCP servers.
|
||||
|
||||
This module provides a simple way to expose Python functions as MCP tools
|
||||
that can be consumed by any MCP client (Claude Desktop, Cursor, etc.).
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.mcp import ToolsMCPServer
|
||||
|
||||
def search(query: str) -> str:
|
||||
'''Search the web.'''
|
||||
return f"Results for {query}"
|
||||
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
server.register_tool(search)
|
||||
server.run() # Starts MCP server (stdio by default)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from .mcp_utils import function_to_mcp_schema
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ToolsMCPServer:
|
||||
"""MCP Server that exposes Python functions as MCP tools.
|
||||
|
||||
This class wraps Python functions and exposes them as MCP tools
|
||||
using the FastMCP library. It supports both sync and async functions.
|
||||
|
||||
Example:
|
||||
server = ToolsMCPServer(name="my-tools")
|
||||
|
||||
def search(query: str) -> str:
|
||||
'''Search the web.'''
|
||||
return f"Results for {query}"
|
||||
|
||||
server.register_tool(search)
|
||||
server.run() # Starts stdio server
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "praisonai-tools",
|
||||
tools: Optional[List[Callable]] = None,
|
||||
debug: bool = False
|
||||
):
|
||||
"""Initialize the MCP server.
|
||||
|
||||
Args:
|
||||
name: Name of the MCP server
|
||||
tools: Optional list of tools to register on init
|
||||
debug: Enable debug logging
|
||||
"""
|
||||
self.name = name
|
||||
self.tools: List[Callable] = []
|
||||
self._tool_map: Dict[str, Callable] = {}
|
||||
self._fastmcp = None
|
||||
self._debug = debug
|
||||
|
||||
if debug:
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
# Register initial tools if provided
|
||||
if tools:
|
||||
self.register_tools(tools)
|
||||
|
||||
def register_tool(self, func: Callable) -> None:
|
||||
"""Register a single tool function.
|
||||
|
||||
Args:
|
||||
func: The function to register as an MCP tool
|
||||
"""
|
||||
if not callable(func):
|
||||
raise TypeError(f"Expected callable, got {type(func)}")
|
||||
|
||||
self.tools.append(func)
|
||||
tool_name = getattr(func, '__mcp_name__', None) or func.__name__
|
||||
self._tool_map[tool_name] = func
|
||||
|
||||
logger.debug(f"Registered tool: {tool_name}")
|
||||
|
||||
def register_tools(self, funcs: List[Callable]) -> None:
|
||||
"""Register multiple tool functions.
|
||||
|
||||
Args:
|
||||
funcs: List of functions to register as MCP tools
|
||||
"""
|
||||
for func in funcs:
|
||||
self.register_tool(func)
|
||||
|
||||
def get_tool_schemas(self) -> List[Dict[str, Any]]:
|
||||
"""Get MCP schemas for all registered tools.
|
||||
|
||||
Returns:
|
||||
List of MCP tool schema dictionaries
|
||||
"""
|
||||
return [function_to_mcp_schema(tool) for tool in self.tools]
|
||||
|
||||
def get_tool_names(self) -> List[str]:
|
||||
"""Get names of all registered tools.
|
||||
|
||||
Returns:
|
||||
List of tool names
|
||||
"""
|
||||
return list(self._tool_map.keys())
|
||||
|
||||
def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
|
||||
"""Execute a registered tool synchronously.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool to execute
|
||||
arguments: Arguments to pass to the tool
|
||||
|
||||
Returns:
|
||||
Result from the tool execution
|
||||
|
||||
Raises:
|
||||
ValueError: If tool is not found
|
||||
"""
|
||||
if tool_name not in self._tool_map:
|
||||
raise ValueError(f"Tool '{tool_name}' not found")
|
||||
|
||||
func = self._tool_map[tool_name]
|
||||
|
||||
# Handle async functions
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
return loop.run_until_complete(func(**arguments))
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
return func(**arguments)
|
||||
|
||||
async def execute_tool_async(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
|
||||
"""Execute a registered tool asynchronously.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool to execute
|
||||
arguments: Arguments to pass to the tool
|
||||
|
||||
Returns:
|
||||
Result from the tool execution
|
||||
|
||||
Raises:
|
||||
ValueError: If tool is not found
|
||||
"""
|
||||
if tool_name not in self._tool_map:
|
||||
raise ValueError(f"Tool '{tool_name}' not found")
|
||||
|
||||
func = self._tool_map[tool_name]
|
||||
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
return await func(**arguments)
|
||||
|
||||
# Run sync function in executor
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, lambda: func(**arguments))
|
||||
|
||||
def get_fastmcp(self):
|
||||
"""Get or create FastMCP instance with registered tools.
|
||||
|
||||
Returns:
|
||||
FastMCP instance with all tools registered
|
||||
"""
|
||||
if self._fastmcp is None:
|
||||
self._fastmcp = self._create_fastmcp()
|
||||
return self._fastmcp
|
||||
|
||||
def _create_fastmcp(self):
|
||||
"""Create FastMCP instance and register all tools.
|
||||
|
||||
Returns:
|
||||
Configured FastMCP instance
|
||||
"""
|
||||
try:
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"FastMCP not available. Install with: pip install mcp"
|
||||
)
|
||||
|
||||
mcp = FastMCP(self.name)
|
||||
|
||||
# Register each tool with FastMCP
|
||||
for func in self.tools:
|
||||
self._register_tool_with_fastmcp(mcp, func)
|
||||
|
||||
return mcp
|
||||
|
||||
def _register_tool_with_fastmcp(self, mcp, func: Callable) -> None:
|
||||
"""Register a single function with FastMCP.
|
||||
|
||||
Args:
|
||||
mcp: FastMCP instance
|
||||
func: Function to register
|
||||
"""
|
||||
# Get tool metadata
|
||||
tool_name = getattr(func, '__mcp_name__', None) or func.__name__
|
||||
description = getattr(func, '__mcp_description__', None)
|
||||
if description is None:
|
||||
docstring = inspect.getdoc(func)
|
||||
description = docstring.split('\n')[0] if docstring else tool_name
|
||||
|
||||
# Create wrapper that preserves signature
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
@mcp.tool(name=tool_name, description=description)
|
||||
async def tool_wrapper(**kwargs):
|
||||
return await func(**kwargs)
|
||||
else:
|
||||
@mcp.tool(name=tool_name, description=description)
|
||||
def tool_wrapper(**kwargs):
|
||||
return func(**kwargs)
|
||||
|
||||
# Copy signature for proper schema generation
|
||||
tool_wrapper.__signature__ = inspect.signature(func)
|
||||
tool_wrapper.__annotations__ = getattr(func, '__annotations__', {})
|
||||
|
||||
def run(self, transport: str = "stdio") -> None:
|
||||
"""Run the MCP server.
|
||||
|
||||
Args:
|
||||
transport: Transport type - "stdio" or "sse"
|
||||
"""
|
||||
if transport == "stdio":
|
||||
self.run_stdio()
|
||||
elif transport == "sse":
|
||||
self.run_sse()
|
||||
else:
|
||||
raise ValueError(f"Unknown transport: {transport}. Use 'stdio' or 'sse'")
|
||||
|
||||
def run_stdio(self) -> None:
|
||||
"""Run the MCP server using stdio transport."""
|
||||
mcp = self.get_fastmcp()
|
||||
|
||||
print(f"🚀 Starting MCP server '{self.name}' with stdio transport")
|
||||
print(f"🛠️ Available tools: {', '.join(self.get_tool_names())}")
|
||||
|
||||
mcp.run()
|
||||
|
||||
def run_sse(self, host: str = "0.0.0.0", port: int = 8080) -> None:
|
||||
"""Run the MCP server using SSE transport.
|
||||
|
||||
Args:
|
||||
host: Host to bind to
|
||||
port: Port to listen on
|
||||
"""
|
||||
try:
|
||||
import uvicorn
|
||||
from mcp.server.sse import SseServerTransport
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
from starlette.routing import Mount, Route
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
f"SSE dependencies not available: {e}. "
|
||||
"Install with: pip install uvicorn starlette"
|
||||
)
|
||||
|
||||
mcp = self.get_fastmcp()
|
||||
|
||||
# Set up SSE transport
|
||||
sse_path = "/sse"
|
||||
messages_path = "/messages/"
|
||||
|
||||
sse_transport = SseServerTransport(messages_path)
|
||||
|
||||
async def handle_sse(request: Request):
|
||||
async with sse_transport.connect_sse(
|
||||
request.scope, request.receive, request._send
|
||||
) as (read_stream, write_stream):
|
||||
await mcp._mcp_server.run(
|
||||
read_stream,
|
||||
write_stream,
|
||||
mcp._mcp_server.create_initialization_options()
|
||||
)
|
||||
|
||||
app = Starlette(
|
||||
debug=self._debug,
|
||||
routes=[
|
||||
Route(sse_path, endpoint=handle_sse),
|
||||
Mount(messages_path, app=sse_transport.handle_post_message),
|
||||
]
|
||||
)
|
||||
|
||||
print(f"🚀 Starting MCP server '{self.name}' with SSE transport")
|
||||
print(f"📡 SSE endpoint: http://{host}:{port}{sse_path}")
|
||||
print(f"🛠️ Available tools: {', '.join(self.get_tool_names())}")
|
||||
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
|
||||
|
||||
def launch_tools_mcp_server(
|
||||
tools: Optional[List[Callable]] = None,
|
||||
tool_names: Optional[List[str]] = None,
|
||||
name: str = "praisonai-tools",
|
||||
transport: str = "stdio",
|
||||
host: str = "0.0.0.0",
|
||||
port: int = 8080,
|
||||
debug: bool = False
|
||||
) -> None:
|
||||
"""Convenience function to launch an MCP server with tools.
|
||||
|
||||
Args:
|
||||
tools: List of tool functions to expose
|
||||
tool_names: List of tool names to load from praisonaiagents.tools
|
||||
name: Name of the MCP server
|
||||
transport: Transport type - "stdio" or "sse"
|
||||
host: Host for SSE transport
|
||||
port: Port for SSE transport
|
||||
debug: Enable debug logging
|
||||
|
||||
Example:
|
||||
# With custom tools
|
||||
def my_search(query: str) -> str:
|
||||
return f"Results for {query}"
|
||||
|
||||
launch_tools_mcp_server(tools=[my_search])
|
||||
|
||||
# With built-in tools
|
||||
launch_tools_mcp_server(tool_names=["tavily_search", "exa_search"])
|
||||
"""
|
||||
server = ToolsMCPServer(name=name, debug=debug)
|
||||
|
||||
# Register custom tools
|
||||
if tools:
|
||||
server.register_tools(tools)
|
||||
|
||||
# Load and register tools by name
|
||||
if tool_names:
|
||||
from praisonaiagents import tools as praison_tools
|
||||
|
||||
for tool_name in tool_names:
|
||||
if hasattr(praison_tools, tool_name):
|
||||
tool_func = getattr(praison_tools, tool_name)
|
||||
if callable(tool_func):
|
||||
server.register_tool(tool_func)
|
||||
else:
|
||||
logger.warning(f"Tool '{tool_name}' is not callable")
|
||||
else:
|
||||
logger.warning(f"Tool '{tool_name}' not found in praisonaiagents.tools")
|
||||
|
||||
if not server.tools:
|
||||
raise ValueError("No tools registered. Provide 'tools' or 'tool_names'.")
|
||||
|
||||
# Run the server
|
||||
if transport == "stdio":
|
||||
server.run_stdio()
|
||||
elif transport == "sse":
|
||||
server.run_sse(host=host, port=port)
|
||||
else:
|
||||
raise ValueError(f"Unknown transport: {transport}")
|
||||
@@ -0,0 +1,457 @@
|
||||
"""
|
||||
MCP Session Management Module.
|
||||
|
||||
This module provides session management utilities for MCP transports,
|
||||
implementing the session management requirements from MCP Protocol
|
||||
Revision 2025-11-25.
|
||||
|
||||
Features:
|
||||
- Session ID validation (visible ASCII 0x21-0x7E)
|
||||
- Session lifecycle management
|
||||
- Protocol version header handling
|
||||
- Session expiration detection
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
|
||||
# Valid protocol versions
|
||||
VALID_PROTOCOL_VERSIONS = {
|
||||
'2024-11-05', # Original HTTP+SSE
|
||||
'2025-03-26', # Streamable HTTP introduction
|
||||
'2025-06-18', # Updates
|
||||
'2025-11-25', # Current revision
|
||||
}
|
||||
|
||||
# Default protocol version for backward compatibility
|
||||
DEFAULT_PROTOCOL_VERSION = '2025-03-26'
|
||||
|
||||
|
||||
def is_valid_session_id(session_id: Optional[str]) -> bool:
|
||||
"""
|
||||
Validate MCP session ID per specification.
|
||||
|
||||
Session IDs MUST only contain visible ASCII characters (0x21 to 0x7E).
|
||||
This excludes space (0x20), control characters (0x00-0x1F), and DEL (0x7F).
|
||||
|
||||
Args:
|
||||
session_id: The session ID to validate
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
"""
|
||||
if session_id is None or session_id == "":
|
||||
return False
|
||||
|
||||
# Check each character is in visible ASCII range (0x21-0x7E)
|
||||
for char in session_id:
|
||||
code = ord(char)
|
||||
if code < 0x21 or code > 0x7E:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_valid_protocol_version(version: Optional[str]) -> bool:
|
||||
"""
|
||||
Validate MCP protocol version string.
|
||||
|
||||
Args:
|
||||
version: Protocol version string (e.g., '2025-11-25')
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
"""
|
||||
if not version:
|
||||
return False
|
||||
|
||||
# Check format: YYYY-MM-DD
|
||||
if not re.match(r'^\d{4}-\d{2}-\d{2}$', version):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_session_expired(status_code: int) -> bool:
|
||||
"""
|
||||
Check if HTTP status code indicates session expiration.
|
||||
|
||||
Per MCP spec, HTTP 404 indicates the session has been terminated.
|
||||
|
||||
Args:
|
||||
status_code: HTTP response status code
|
||||
|
||||
Returns:
|
||||
True if session is expired, False otherwise
|
||||
"""
|
||||
return status_code == 404
|
||||
|
||||
|
||||
def extract_session_id(headers: Dict[str, str]) -> Optional[str]:
|
||||
"""
|
||||
Extract session ID from HTTP response headers.
|
||||
|
||||
The header name is case-insensitive per HTTP spec.
|
||||
|
||||
Args:
|
||||
headers: HTTP response headers dictionary
|
||||
|
||||
Returns:
|
||||
Session ID if present, None otherwise
|
||||
"""
|
||||
# Check various case combinations
|
||||
for key in headers:
|
||||
if key.lower() == 'mcp-session-id':
|
||||
return headers[key]
|
||||
return None
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""
|
||||
Manages MCP session state for HTTP transports.
|
||||
|
||||
This class handles:
|
||||
- Session ID storage and validation
|
||||
- Protocol version negotiation
|
||||
- Session lifecycle (active, expired, terminated)
|
||||
- Header generation for requests
|
||||
|
||||
Example:
|
||||
```python
|
||||
manager = SessionManager()
|
||||
|
||||
# After receiving InitializeResult with session ID
|
||||
manager.set_session_id("abc123")
|
||||
|
||||
# Get headers for subsequent requests
|
||||
headers = manager.get_headers()
|
||||
# {'Mcp-Session-Id': 'abc123', 'Mcp-Protocol-Version': '2025-03-26'}
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, protocol_version: Optional[str] = None):
|
||||
"""
|
||||
Initialize session manager.
|
||||
|
||||
Args:
|
||||
protocol_version: Initial protocol version (defaults to 2025-03-26)
|
||||
"""
|
||||
self._session_id: Optional[str] = None
|
||||
self._protocol_version: str = protocol_version or DEFAULT_PROTOCOL_VERSION
|
||||
self._is_expired: bool = False
|
||||
self._metadata: Dict[str, Any] = {}
|
||||
|
||||
@property
|
||||
def session_id(self) -> Optional[str]:
|
||||
"""Get current session ID."""
|
||||
return self._session_id
|
||||
|
||||
@property
|
||||
def protocol_version(self) -> str:
|
||||
"""Get current protocol version."""
|
||||
return self._protocol_version
|
||||
|
||||
@property
|
||||
def is_active(self) -> bool:
|
||||
"""Check if session is active (has ID and not expired)."""
|
||||
return self._session_id is not None and not self._is_expired
|
||||
|
||||
@property
|
||||
def is_expired(self) -> bool:
|
||||
"""Check if session has expired."""
|
||||
return self._is_expired
|
||||
|
||||
def set_session_id(self, session_id: str) -> None:
|
||||
"""
|
||||
Set the session ID.
|
||||
|
||||
Args:
|
||||
session_id: The session ID from server
|
||||
|
||||
Raises:
|
||||
ValueError: If session ID is invalid
|
||||
"""
|
||||
if not is_valid_session_id(session_id):
|
||||
raise ValueError("Invalid session ID: must contain only visible ASCII (0x21-0x7E)")
|
||||
|
||||
self._session_id = session_id
|
||||
self._is_expired = False
|
||||
|
||||
def set_protocol_version(self, version: str) -> None:
|
||||
"""
|
||||
Set the negotiated protocol version.
|
||||
|
||||
Args:
|
||||
version: Protocol version string
|
||||
|
||||
Raises:
|
||||
ValueError: If version format is invalid
|
||||
"""
|
||||
if not is_valid_protocol_version(version):
|
||||
raise ValueError(f"Invalid protocol version format: {version}")
|
||||
|
||||
self._protocol_version = version
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear the session (logout/disconnect)."""
|
||||
self._session_id = None
|
||||
self._is_expired = False
|
||||
self._metadata = {}
|
||||
|
||||
def mark_expired(self) -> None:
|
||||
"""Mark the session as expired (received HTTP 404)."""
|
||||
self._is_expired = True
|
||||
|
||||
def get_headers(self) -> Dict[str, str]:
|
||||
"""
|
||||
Get HTTP headers for MCP requests.
|
||||
|
||||
Returns:
|
||||
Dictionary of headers to include in requests
|
||||
"""
|
||||
headers = {
|
||||
'Mcp-Protocol-Version': self._protocol_version
|
||||
}
|
||||
|
||||
if self._session_id:
|
||||
headers['Mcp-Session-Id'] = self._session_id
|
||||
|
||||
return headers
|
||||
|
||||
def update_from_response(self, headers: Dict[str, str], status_code: int = 200) -> None:
|
||||
"""
|
||||
Update session state from HTTP response.
|
||||
|
||||
Args:
|
||||
headers: Response headers
|
||||
status_code: HTTP status code
|
||||
"""
|
||||
# Check for session expiration
|
||||
if is_session_expired(status_code):
|
||||
self.mark_expired()
|
||||
return
|
||||
|
||||
# Extract and update session ID if present
|
||||
new_session_id = extract_session_id(headers)
|
||||
if new_session_id:
|
||||
self.set_session_id(new_session_id)
|
||||
|
||||
def set_metadata(self, key: str, value: Any) -> None:
|
||||
"""Store arbitrary metadata with the session."""
|
||||
self._metadata[key] = value
|
||||
|
||||
def get_metadata(self, key: str, default: Any = None) -> Any:
|
||||
"""Retrieve metadata from the session."""
|
||||
return self._metadata.get(key, default)
|
||||
|
||||
|
||||
class ResumabilityManager:
|
||||
"""
|
||||
Manages SSE stream resumability for MCP transports.
|
||||
|
||||
This class handles:
|
||||
- Event ID tracking per stream
|
||||
- Last-Event-ID header generation
|
||||
- Retry delay handling
|
||||
|
||||
Per MCP spec, event IDs should be assigned per-stream to act as
|
||||
a cursor within that particular stream.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize resumability manager."""
|
||||
self._last_event_ids: Dict[str, str] = {} # stream_id -> last_event_id
|
||||
self._retry_delay: int = 3000 # Default retry delay in ms
|
||||
|
||||
def set_last_event_id(self, stream_id: str, event_id: str) -> None:
|
||||
"""
|
||||
Store the last received event ID for a stream.
|
||||
|
||||
Args:
|
||||
stream_id: Identifier for the SSE stream
|
||||
event_id: The event ID from the SSE event
|
||||
"""
|
||||
self._last_event_ids[stream_id] = event_id
|
||||
|
||||
def get_last_event_id(self, stream_id: str) -> Optional[str]:
|
||||
"""
|
||||
Get the last received event ID for a stream.
|
||||
|
||||
Args:
|
||||
stream_id: Identifier for the SSE stream
|
||||
|
||||
Returns:
|
||||
Last event ID if available, None otherwise
|
||||
"""
|
||||
return self._last_event_ids.get(stream_id)
|
||||
|
||||
def set_retry_delay(self, delay_ms: int) -> None:
|
||||
"""
|
||||
Set the retry delay from SSE retry field.
|
||||
|
||||
Args:
|
||||
delay_ms: Delay in milliseconds
|
||||
"""
|
||||
self._retry_delay = delay_ms
|
||||
|
||||
def get_retry_delay(self) -> int:
|
||||
"""Get current retry delay in milliseconds."""
|
||||
return self._retry_delay
|
||||
|
||||
def get_resume_headers(self, stream_id: str) -> Dict[str, str]:
|
||||
"""
|
||||
Get headers for resuming a stream.
|
||||
|
||||
Args:
|
||||
stream_id: Identifier for the SSE stream
|
||||
|
||||
Returns:
|
||||
Headers dict with Last-Event-ID if available
|
||||
"""
|
||||
headers = {}
|
||||
last_event_id = self.get_last_event_id(stream_id)
|
||||
if last_event_id:
|
||||
headers['Last-Event-ID'] = last_event_id
|
||||
return headers
|
||||
|
||||
def clear_stream(self, stream_id: str) -> None:
|
||||
"""Clear tracking for a specific stream."""
|
||||
self._last_event_ids.pop(stream_id, None)
|
||||
|
||||
def clear_all(self) -> None:
|
||||
"""Clear all stream tracking."""
|
||||
self._last_event_ids.clear()
|
||||
|
||||
|
||||
def parse_sse_event(event_text: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Parse an SSE event into its components.
|
||||
|
||||
SSE events have the format:
|
||||
id: <event-id>
|
||||
retry: <milliseconds>
|
||||
data: <data-line>
|
||||
data: <data-line-2>
|
||||
|
||||
(blank line ends event)
|
||||
|
||||
Args:
|
||||
event_text: Raw SSE event text
|
||||
|
||||
Returns:
|
||||
Dictionary with 'id', 'retry', 'data', 'event' fields as available
|
||||
"""
|
||||
result: Dict[str, Any] = {}
|
||||
data_lines = []
|
||||
|
||||
for line in event_text.strip().split('\n'):
|
||||
if line.startswith('id:'):
|
||||
result['id'] = line[3:].strip()
|
||||
elif line.startswith('retry:'):
|
||||
try:
|
||||
result['retry'] = int(line[6:].strip())
|
||||
except ValueError:
|
||||
pass
|
||||
elif line.startswith('data:'):
|
||||
data_lines.append(line[5:].strip())
|
||||
elif line.startswith('event:'):
|
||||
result['event'] = line[6:].strip()
|
||||
|
||||
# Join multiple data lines with newlines
|
||||
if data_lines:
|
||||
result['data'] = '\n'.join(data_lines)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class EventStore:
|
||||
"""
|
||||
Optional event store for message redelivery support.
|
||||
|
||||
This class stores events per stream to enable replay when
|
||||
clients reconnect with Last-Event-ID header.
|
||||
|
||||
Note: This is an in-memory store suitable for single-server
|
||||
deployments. For multi-node deployments, use an external
|
||||
store (Redis, database, etc.).
|
||||
"""
|
||||
|
||||
def __init__(self, max_events: int = 1000):
|
||||
"""
|
||||
Initialize event store.
|
||||
|
||||
Args:
|
||||
max_events: Maximum events to store per stream
|
||||
"""
|
||||
self.max_events = max_events
|
||||
self._streams: Dict[str, list] = {} # stream_id -> [(event_id, data), ...]
|
||||
|
||||
def store(self, stream_id: str, event_id: str, data: Any) -> None:
|
||||
"""
|
||||
Store an event for potential replay.
|
||||
|
||||
Args:
|
||||
stream_id: Identifier for the SSE stream
|
||||
event_id: The event ID
|
||||
data: The event data
|
||||
"""
|
||||
if stream_id not in self._streams:
|
||||
self._streams[stream_id] = []
|
||||
|
||||
self._streams[stream_id].append((event_id, data))
|
||||
|
||||
# Trim to max_events
|
||||
if len(self._streams[stream_id]) > self.max_events:
|
||||
self._streams[stream_id] = self._streams[stream_id][-self.max_events:]
|
||||
|
||||
def get_events_after(self, stream_id: str, after_event_id: str) -> list:
|
||||
"""
|
||||
Get events after a specific event ID.
|
||||
|
||||
Args:
|
||||
stream_id: Identifier for the SSE stream
|
||||
after_event_id: The Last-Event-ID from client
|
||||
|
||||
Returns:
|
||||
List of event data after the specified event
|
||||
"""
|
||||
if stream_id not in self._streams:
|
||||
return []
|
||||
|
||||
events = self._streams[stream_id]
|
||||
|
||||
# Find the index of the after_event_id
|
||||
start_index = None
|
||||
for i, (event_id, _) in enumerate(events):
|
||||
if event_id == after_event_id:
|
||||
start_index = i + 1
|
||||
break
|
||||
|
||||
if start_index is None:
|
||||
return []
|
||||
|
||||
return [data for _, data in events[start_index:]]
|
||||
|
||||
def get_all_events(self, stream_id: str) -> list:
|
||||
"""
|
||||
Get all stored events for a stream.
|
||||
|
||||
Args:
|
||||
stream_id: Identifier for the SSE stream
|
||||
|
||||
Returns:
|
||||
List of all event data
|
||||
"""
|
||||
if stream_id not in self._streams:
|
||||
return []
|
||||
|
||||
return [data for _, data in self._streams[stream_id]]
|
||||
|
||||
def clear_stream(self, stream_id: str) -> None:
|
||||
"""Clear events for a specific stream."""
|
||||
self._streams.pop(stream_id, None)
|
||||
|
||||
def clear_all(self) -> None:
|
||||
"""Clear all stored events."""
|
||||
self._streams.clear()
|
||||
@@ -0,0 +1,298 @@
|
||||
"""
|
||||
MCP Transport Abstraction Layer.
|
||||
|
||||
This module provides a common interface for all MCP transports,
|
||||
enabling easy switching between transport types and custom
|
||||
transport implementations.
|
||||
|
||||
Transport types:
|
||||
- stdio: Standard input/output (subprocess)
|
||||
- sse: Server-Sent Events (legacy HTTP+SSE)
|
||||
- http_stream: Streamable HTTP (current standard)
|
||||
- websocket: WebSocket (SEP-1288)
|
||||
"""
|
||||
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, Optional, List, Type
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransportConfig:
|
||||
"""
|
||||
Configuration for MCP transports.
|
||||
|
||||
Attributes:
|
||||
timeout: Connection/operation timeout in seconds
|
||||
debug: Enable debug logging
|
||||
retry_count: Number of retry attempts
|
||||
retry_delay: Base delay between retries in seconds
|
||||
auth_token: Optional authentication token
|
||||
headers: Additional headers for HTTP transports
|
||||
"""
|
||||
timeout: int = 60
|
||||
debug: bool = False
|
||||
retry_count: int = 3
|
||||
retry_delay: float = 1.0
|
||||
auth_token: Optional[str] = None
|
||||
headers: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
class BaseTransport(ABC):
|
||||
"""
|
||||
Abstract base class for MCP transports.
|
||||
|
||||
All MCP transports must implement this interface to ensure
|
||||
consistent behavior across transport types.
|
||||
|
||||
The transport is responsible for:
|
||||
- Establishing and maintaining connections
|
||||
- Sending JSON-RPC messages
|
||||
- Receiving JSON-RPC messages
|
||||
- Proper cleanup on close
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def connect(self) -> None:
|
||||
"""
|
||||
Establish connection to the MCP server.
|
||||
|
||||
Raises:
|
||||
ConnectionError: If connection fails
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def send(self, message: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Send a JSON-RPC message to the server.
|
||||
|
||||
Args:
|
||||
message: JSON-RPC message dictionary
|
||||
|
||||
Raises:
|
||||
RuntimeError: If not connected
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def receive(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Receive a JSON-RPC message from the server.
|
||||
|
||||
Returns:
|
||||
JSON-RPC message dictionary
|
||||
|
||||
Raises:
|
||||
RuntimeError: If not connected
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
"""
|
||||
Close the connection and cleanup resources.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_connected(self) -> bool:
|
||||
"""Check if transport is currently connected."""
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
"""Async context manager entry."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
|
||||
|
||||
def get_transport_type(target: str) -> str:
|
||||
"""
|
||||
Determine the appropriate transport type for a target.
|
||||
|
||||
Args:
|
||||
target: URL or command string
|
||||
|
||||
Returns:
|
||||
Transport type string: "stdio", "sse", "http_stream", or "websocket"
|
||||
"""
|
||||
if not isinstance(target, str):
|
||||
return "stdio"
|
||||
|
||||
# WebSocket URLs
|
||||
if re.match(r'^wss?://', target, re.IGNORECASE):
|
||||
return "websocket"
|
||||
|
||||
# HTTP URLs
|
||||
if re.match(r'^https?://', target, re.IGNORECASE):
|
||||
# Legacy SSE endpoint
|
||||
if target.endswith('/sse'):
|
||||
return "sse"
|
||||
# Default to Streamable HTTP
|
||||
return "http_stream"
|
||||
|
||||
# Everything else is stdio (command)
|
||||
return "stdio"
|
||||
|
||||
|
||||
class TransportRegistry:
|
||||
"""
|
||||
Registry for MCP transport implementations.
|
||||
|
||||
This allows registration of custom transports and provides
|
||||
a central place to look up transport classes by name.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize empty registry."""
|
||||
self._transports: Dict[str, Type[BaseTransport]] = {}
|
||||
|
||||
def register(self, name: str, transport_class: Type[BaseTransport]) -> None:
|
||||
"""
|
||||
Register a transport class.
|
||||
|
||||
Args:
|
||||
name: Transport name (e.g., "websocket")
|
||||
transport_class: Transport class (must extend BaseTransport)
|
||||
"""
|
||||
if not issubclass(transport_class, BaseTransport):
|
||||
raise TypeError(f"{transport_class} must extend BaseTransport")
|
||||
self._transports[name] = transport_class
|
||||
|
||||
def get(self, name: str) -> Optional[Type[BaseTransport]]:
|
||||
"""
|
||||
Get a registered transport class.
|
||||
|
||||
Args:
|
||||
name: Transport name
|
||||
|
||||
Returns:
|
||||
Transport class or None if not found
|
||||
"""
|
||||
return self._transports.get(name)
|
||||
|
||||
def list_transports(self) -> List[str]:
|
||||
"""
|
||||
List all registered transport names.
|
||||
|
||||
Returns:
|
||||
List of transport names
|
||||
"""
|
||||
return list(self._transports.keys())
|
||||
|
||||
def unregister(self, name: str) -> None:
|
||||
"""
|
||||
Unregister a transport.
|
||||
|
||||
Args:
|
||||
name: Transport name to remove
|
||||
"""
|
||||
self._transports.pop(name, None)
|
||||
|
||||
|
||||
# Placeholder transport classes for registry
|
||||
# These are minimal implementations to satisfy the registry
|
||||
|
||||
class StdioTransportPlaceholder(BaseTransport):
|
||||
"""Placeholder for stdio transport."""
|
||||
async def connect(self): pass
|
||||
async def send(self, message): pass
|
||||
async def receive(self): return {}
|
||||
async def close(self): pass
|
||||
@property
|
||||
def is_connected(self): return False
|
||||
|
||||
|
||||
class SSETransportPlaceholder(BaseTransport):
|
||||
"""Placeholder for SSE transport."""
|
||||
async def connect(self): pass
|
||||
async def send(self, message): pass
|
||||
async def receive(self): return {}
|
||||
async def close(self): pass
|
||||
@property
|
||||
def is_connected(self): return False
|
||||
|
||||
|
||||
class HTTPStreamTransportPlaceholder(BaseTransport):
|
||||
"""Placeholder for HTTP Stream transport."""
|
||||
async def connect(self): pass
|
||||
async def send(self, message): pass
|
||||
async def receive(self): return {}
|
||||
async def close(self): pass
|
||||
@property
|
||||
def is_connected(self): return False
|
||||
|
||||
|
||||
class WebSocketTransportPlaceholder(BaseTransport):
|
||||
"""Placeholder for WebSocket transport."""
|
||||
async def connect(self): pass
|
||||
async def send(self, message): pass
|
||||
async def receive(self): return {}
|
||||
async def close(self): pass
|
||||
@property
|
||||
def is_connected(self): return False
|
||||
|
||||
|
||||
# Default registry instance
|
||||
_default_registry: Optional[TransportRegistry] = None
|
||||
|
||||
|
||||
def get_default_registry() -> TransportRegistry:
|
||||
"""
|
||||
Get the default transport registry with built-in transports.
|
||||
|
||||
Returns:
|
||||
TransportRegistry with default transports registered
|
||||
"""
|
||||
global _default_registry
|
||||
|
||||
if _default_registry is None:
|
||||
_default_registry = TransportRegistry()
|
||||
|
||||
# Register default transports
|
||||
_default_registry.register("stdio", StdioTransportPlaceholder)
|
||||
_default_registry.register("sse", SSETransportPlaceholder)
|
||||
_default_registry.register("http_stream", HTTPStreamTransportPlaceholder)
|
||||
_default_registry.register("websocket", WebSocketTransportPlaceholder)
|
||||
|
||||
return _default_registry
|
||||
|
||||
|
||||
def create_transport(
|
||||
target: str,
|
||||
config: Optional[TransportConfig] = None,
|
||||
registry: Optional[TransportRegistry] = None
|
||||
) -> BaseTransport:
|
||||
"""
|
||||
Create a transport instance for the given target.
|
||||
|
||||
Args:
|
||||
target: URL or command string
|
||||
config: Optional transport configuration
|
||||
registry: Optional custom registry (uses default if None)
|
||||
|
||||
Returns:
|
||||
Transport instance
|
||||
|
||||
Raises:
|
||||
ValueError: If no suitable transport found
|
||||
"""
|
||||
if registry is None:
|
||||
registry = get_default_registry()
|
||||
|
||||
if config is None:
|
||||
config = TransportConfig()
|
||||
|
||||
transport_type = get_transport_type(target)
|
||||
transport_class = registry.get(transport_type)
|
||||
|
||||
if transport_class is None:
|
||||
raise ValueError(f"No transport registered for type: {transport_type}")
|
||||
|
||||
return transport_class()
|
||||
@@ -0,0 +1,238 @@
|
||||
"""MCP utility functions for converting Python functions to MCP tool schemas.
|
||||
|
||||
This module provides utilities to convert Python functions with type hints
|
||||
into MCP-compatible tool schemas for exposing as MCP server tools.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.mcp.mcp_utils import function_to_mcp_schema
|
||||
|
||||
def my_tool(query: str, max_results: int = 10) -> str:
|
||||
'''Search for information.'''
|
||||
return "results"
|
||||
|
||||
schema = function_to_mcp_schema(my_tool)
|
||||
# Returns MCP-compatible tool schema
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from typing import Any, Callable, Dict, List, Union, get_type_hints, get_origin, get_args
|
||||
|
||||
|
||||
def python_type_to_json_schema(python_type: Any) -> Dict[str, Any]:
|
||||
"""Convert a Python type annotation to JSON Schema.
|
||||
|
||||
Args:
|
||||
python_type: A Python type annotation (e.g., str, int, List[str], Optional[int])
|
||||
|
||||
Returns:
|
||||
JSON Schema dictionary representing the type
|
||||
"""
|
||||
# Handle None type
|
||||
if python_type is None or python_type is type(None):
|
||||
return {"type": "null"}
|
||||
|
||||
# Handle Any type
|
||||
if python_type is Any:
|
||||
return {}
|
||||
|
||||
# Handle basic types
|
||||
type_mapping = {
|
||||
str: {"type": "string"},
|
||||
int: {"type": "integer"},
|
||||
float: {"type": "number"},
|
||||
bool: {"type": "boolean"},
|
||||
bytes: {"type": "string"},
|
||||
}
|
||||
|
||||
if python_type in type_mapping:
|
||||
return type_mapping[python_type]
|
||||
|
||||
# Handle generic types (List, Dict, Optional, Union)
|
||||
origin = get_origin(python_type)
|
||||
args = get_args(python_type)
|
||||
|
||||
# Handle List[T]
|
||||
if origin is list:
|
||||
if args:
|
||||
return {
|
||||
"type": "array",
|
||||
"items": python_type_to_json_schema(args[0])
|
||||
}
|
||||
return {"type": "array", "items": {}}
|
||||
|
||||
# Handle Dict[K, V]
|
||||
if origin is dict:
|
||||
return {"type": "object"}
|
||||
|
||||
# Handle Optional[T] (Union[T, None])
|
||||
if origin is Union:
|
||||
# Filter out None type
|
||||
non_none_args = [arg for arg in args if arg is not type(None)]
|
||||
if len(non_none_args) == 1:
|
||||
# Optional[T] case
|
||||
return python_type_to_json_schema(non_none_args[0])
|
||||
elif len(non_none_args) > 1:
|
||||
# Union of multiple types
|
||||
return {"oneOf": [python_type_to_json_schema(arg) for arg in non_none_args]}
|
||||
return {}
|
||||
|
||||
# Handle tuple
|
||||
if origin is tuple:
|
||||
if args:
|
||||
return {
|
||||
"type": "array",
|
||||
"items": [python_type_to_json_schema(arg) for arg in args]
|
||||
}
|
||||
return {"type": "array"}
|
||||
|
||||
# Handle set
|
||||
if origin is set:
|
||||
if args:
|
||||
return {
|
||||
"type": "array",
|
||||
"items": python_type_to_json_schema(args[0]),
|
||||
"uniqueItems": True
|
||||
}
|
||||
return {"type": "array", "uniqueItems": True}
|
||||
|
||||
# Default to object for unknown types
|
||||
return {"type": "object"}
|
||||
|
||||
|
||||
def get_tool_metadata(func: Callable) -> Dict[str, str]:
|
||||
"""Extract metadata from a function for MCP tool registration.
|
||||
|
||||
Checks for custom attributes (__mcp_name__, __mcp_description__) first,
|
||||
then falls back to function name and docstring.
|
||||
|
||||
Args:
|
||||
func: The function to extract metadata from
|
||||
|
||||
Returns:
|
||||
Dictionary with 'name' and 'description' keys
|
||||
"""
|
||||
# Get name - check for custom attribute first
|
||||
name = getattr(func, '__mcp_name__', None) or func.__name__
|
||||
|
||||
# Get description - check for custom attribute first
|
||||
description = getattr(func, '__mcp_description__', None)
|
||||
|
||||
if description is None:
|
||||
# Extract from docstring
|
||||
docstring = inspect.getdoc(func)
|
||||
if docstring:
|
||||
# Use first line of docstring as description
|
||||
description = docstring.split('\n')[0].strip()
|
||||
else:
|
||||
# Fall back to function name
|
||||
description = name
|
||||
|
||||
return {
|
||||
"name": name,
|
||||
"description": description
|
||||
}
|
||||
|
||||
|
||||
def function_to_mcp_schema(func: Callable) -> Dict[str, Any]:
|
||||
"""Convert a Python function to an MCP tool schema.
|
||||
|
||||
This function inspects the function's signature, type hints, and docstring
|
||||
to generate a complete MCP-compatible tool schema.
|
||||
|
||||
Args:
|
||||
func: The function to convert
|
||||
|
||||
Returns:
|
||||
MCP tool schema dictionary with:
|
||||
- name: Tool name
|
||||
- description: Tool description
|
||||
- inputSchema: JSON Schema for parameters
|
||||
|
||||
Example:
|
||||
def search(query: str, max_results: int = 10) -> str:
|
||||
'''Search the web for information.'''
|
||||
return "results"
|
||||
|
||||
schema = function_to_mcp_schema(search)
|
||||
# {
|
||||
# "name": "search",
|
||||
# "description": "Search the web for information.",
|
||||
# "inputSchema": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "query": {"type": "string"},
|
||||
# "max_results": {"type": "integer"}
|
||||
# },
|
||||
# "required": ["query"]
|
||||
# }
|
||||
# }
|
||||
"""
|
||||
# Get metadata
|
||||
metadata = get_tool_metadata(func)
|
||||
|
||||
# Get function signature
|
||||
sig = inspect.signature(func)
|
||||
|
||||
# Try to get type hints
|
||||
try:
|
||||
type_hints = get_type_hints(func)
|
||||
except Exception:
|
||||
type_hints = {}
|
||||
|
||||
# Build properties and required list
|
||||
properties = {}
|
||||
required = []
|
||||
|
||||
for param_name, param in sig.parameters.items():
|
||||
# Skip self, cls, and special parameters
|
||||
if param_name in ('self', 'cls'):
|
||||
continue
|
||||
|
||||
# Skip *args and **kwargs
|
||||
if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
|
||||
continue
|
||||
|
||||
# Get type annotation
|
||||
param_type = type_hints.get(param_name, param.annotation)
|
||||
|
||||
if param_type is inspect.Parameter.empty:
|
||||
# No type annotation - default to string
|
||||
properties[param_name] = {"type": "string"}
|
||||
else:
|
||||
properties[param_name] = python_type_to_json_schema(param_type)
|
||||
|
||||
# Check if parameter is required (no default value)
|
||||
if param.default is inspect.Parameter.empty:
|
||||
# Check if it's Optional type
|
||||
origin = get_origin(param_type) if param_type is not inspect.Parameter.empty else None
|
||||
args = get_args(param_type) if param_type is not inspect.Parameter.empty else ()
|
||||
|
||||
is_optional = origin is Union and type(None) in args
|
||||
|
||||
if not is_optional:
|
||||
required.append(param_name)
|
||||
|
||||
# Build the schema
|
||||
schema = {
|
||||
"name": metadata["name"],
|
||||
"description": metadata["description"],
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": properties,
|
||||
"required": required
|
||||
}
|
||||
}
|
||||
|
||||
return schema
|
||||
|
||||
|
||||
def functions_to_mcp_tools(functions: List[Callable]) -> List[Dict[str, Any]]:
|
||||
"""Convert multiple functions to MCP tool schemas.
|
||||
|
||||
Args:
|
||||
functions: List of functions to convert
|
||||
|
||||
Returns:
|
||||
List of MCP tool schemas
|
||||
"""
|
||||
return [function_to_mcp_schema(func) for func in functions]
|
||||
@@ -0,0 +1,596 @@
|
||||
"""
|
||||
WebSocket Transport for MCP (Model Context Protocol).
|
||||
|
||||
This module provides WebSocket-based transport for connecting to MCP servers.
|
||||
WebSocket transport offers:
|
||||
- Long-lived bidirectional communication
|
||||
- Better cloud provider support (hibernatable connections)
|
||||
- Lower protocol overhead compared to HTTP
|
||||
- Session persistence across network interruptions
|
||||
|
||||
This is an optional transport - the 'websockets' package is only imported
|
||||
when ws:// or wss:// URLs are used, ensuring no performance impact on
|
||||
existing stdio/SSE/HTTP transports.
|
||||
|
||||
Based on SEP-1288: WebSocket Transport for MCP (in review).
|
||||
Protocol Revision: 2025-11-25
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import inspect
|
||||
import re
|
||||
from typing import Any, Dict, Optional, List
|
||||
|
||||
logger = logging.getLogger("mcp-websocket")
|
||||
|
||||
|
||||
def is_websocket_url(url: str) -> bool:
|
||||
"""
|
||||
Check if a URL is a WebSocket URL.
|
||||
|
||||
Args:
|
||||
url: The URL string to check
|
||||
|
||||
Returns:
|
||||
True if the URL starts with ws:// or wss://, False otherwise
|
||||
"""
|
||||
if not isinstance(url, str):
|
||||
return False
|
||||
return bool(re.match(r'^wss?://', url, re.IGNORECASE))
|
||||
|
||||
|
||||
def calculate_backoff(attempt: int, base_delay: float = 1.0, max_delay: float = 60.0) -> float:
|
||||
"""
|
||||
Calculate exponential backoff delay for reconnection.
|
||||
|
||||
Args:
|
||||
attempt: The retry attempt number (0-indexed)
|
||||
base_delay: Base delay in seconds
|
||||
max_delay: Maximum delay cap in seconds
|
||||
|
||||
Returns:
|
||||
Delay in seconds for this attempt
|
||||
"""
|
||||
delay = base_delay * (2 ** attempt)
|
||||
return min(delay, max_delay)
|
||||
|
||||
|
||||
# Global event loop for WebSocket operations (similar to other MCP transports)
|
||||
_event_loop = None
|
||||
_loop_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_event_loop():
|
||||
"""Get or create a global event loop for WebSocket operations."""
|
||||
global _event_loop
|
||||
with _loop_lock:
|
||||
if _event_loop is None or _event_loop.is_closed():
|
||||
_event_loop = asyncio.new_event_loop()
|
||||
return _event_loop
|
||||
|
||||
|
||||
class WebSocketTransport:
|
||||
"""
|
||||
Low-level WebSocket transport for MCP communication.
|
||||
|
||||
This class handles the WebSocket connection and JSON-RPC message framing.
|
||||
It supports:
|
||||
- Secure (wss://) and insecure (ws://) connections
|
||||
- Session ID management
|
||||
- Authentication via subprotocol or headers
|
||||
- Automatic reconnection with exponential backoff
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
session_id: Optional[str] = None,
|
||||
auth_token: Optional[str] = None,
|
||||
max_retries: int = 3,
|
||||
retry_delay: float = 1.0,
|
||||
timeout: float = 30.0,
|
||||
ping_interval: float = 30.0,
|
||||
ping_timeout: float = 10.0
|
||||
):
|
||||
"""
|
||||
Initialize WebSocket transport.
|
||||
|
||||
Args:
|
||||
url: WebSocket URL (ws:// or wss://)
|
||||
session_id: Optional MCP session ID
|
||||
auth_token: Optional authentication token
|
||||
max_retries: Maximum reconnection attempts
|
||||
retry_delay: Base delay between retries (exponential backoff)
|
||||
timeout: Connection timeout in seconds
|
||||
ping_interval: Interval between ping messages (keepalive)
|
||||
ping_timeout: Timeout for ping response
|
||||
"""
|
||||
self.url = url
|
||||
self.session_id = session_id
|
||||
self.auth_token = auth_token
|
||||
self.max_retries = max_retries
|
||||
self.retry_delay = retry_delay
|
||||
self.timeout = timeout
|
||||
self.ping_interval = ping_interval
|
||||
self.ping_timeout = ping_timeout
|
||||
|
||||
self._ws = None
|
||||
self._closed = False
|
||||
self._message_queue: asyncio.Queue = asyncio.Queue()
|
||||
self._receive_task = None
|
||||
|
||||
async def __aenter__(self):
|
||||
"""Async context manager entry - establish connection."""
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Async context manager exit - close connection."""
|
||||
await self.close()
|
||||
|
||||
async def connect(self):
|
||||
"""Establish WebSocket connection."""
|
||||
try:
|
||||
# Lazy import websockets
|
||||
import websockets
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"websockets package is required for WebSocket transport. "
|
||||
"Install it with: pip install websockets"
|
||||
)
|
||||
|
||||
# Build connection options
|
||||
extra_headers = {}
|
||||
subprotocols = []
|
||||
|
||||
# Add session ID if present
|
||||
if self.session_id:
|
||||
extra_headers['Mcp-Session-Id'] = self.session_id
|
||||
|
||||
# Add auth token via subprotocol (SEP-1288 recommendation)
|
||||
if self.auth_token:
|
||||
# Smuggle auth token in subprotocol list
|
||||
# Format: "mcp-auth-<base64-encoded-token>"
|
||||
import base64
|
||||
encoded_token = base64.b64encode(self.auth_token.encode()).decode()
|
||||
subprotocols.append(f"mcp-auth-{encoded_token}")
|
||||
|
||||
# Always include MCP subprotocol
|
||||
subprotocols.append("mcp")
|
||||
|
||||
# Connect with retry logic
|
||||
last_error = None
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
self._ws = await asyncio.wait_for(
|
||||
websockets.connect(
|
||||
self.url,
|
||||
extra_headers=extra_headers if extra_headers else None,
|
||||
subprotocols=subprotocols if subprotocols else None,
|
||||
ping_interval=self.ping_interval,
|
||||
ping_timeout=self.ping_timeout,
|
||||
close_timeout=10
|
||||
),
|
||||
timeout=self.timeout
|
||||
)
|
||||
logger.debug(f"WebSocket connected to {self.url}")
|
||||
|
||||
# Start background receive task
|
||||
self._receive_task = asyncio.create_task(self._receive_loop())
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
if attempt < self.max_retries:
|
||||
delay = calculate_backoff(attempt, self.retry_delay)
|
||||
logger.warning(f"WebSocket connection failed, retrying in {delay}s: {e}")
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
raise ConnectionError(f"Failed to connect after {self.max_retries + 1} attempts: {last_error}")
|
||||
|
||||
async def _receive_loop(self):
|
||||
"""Background task to receive messages and queue them."""
|
||||
try:
|
||||
async for message in self._ws:
|
||||
try:
|
||||
data = json.loads(message)
|
||||
await self._message_queue.put(data)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to parse WebSocket message: {e}")
|
||||
except Exception as e:
|
||||
if not self._closed:
|
||||
logger.error(f"WebSocket receive error: {e}")
|
||||
await self._message_queue.put({"error": str(e)})
|
||||
|
||||
async def send(self, message: Dict[str, Any]):
|
||||
"""
|
||||
Send a JSON-RPC message over WebSocket.
|
||||
|
||||
Args:
|
||||
message: JSON-RPC message dictionary
|
||||
"""
|
||||
if self._ws is None:
|
||||
raise RuntimeError("WebSocket not connected")
|
||||
|
||||
data = json.dumps(message)
|
||||
await self._ws.send(data)
|
||||
logger.debug(f"Sent: {message.get('method', message.get('id', 'response'))}")
|
||||
|
||||
async def receive(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Receive a JSON-RPC message from WebSocket.
|
||||
|
||||
Returns:
|
||||
Parsed JSON-RPC message dictionary
|
||||
"""
|
||||
message = await self._message_queue.get()
|
||||
if "error" in message and len(message) == 1:
|
||||
raise ConnectionError(message["error"])
|
||||
return message
|
||||
|
||||
async def close(self):
|
||||
"""Close WebSocket connection."""
|
||||
self._closed = True
|
||||
|
||||
if self._receive_task:
|
||||
self._receive_task.cancel()
|
||||
try:
|
||||
await self._receive_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self._ws:
|
||||
await self._ws.close()
|
||||
self._ws = None
|
||||
|
||||
logger.debug("WebSocket connection closed")
|
||||
|
||||
|
||||
class WebSocketMCPTool:
|
||||
"""
|
||||
A wrapper for an MCP tool that can be used with praisonaiagents.
|
||||
|
||||
This class wraps a tool from a WebSocket MCP server and makes it
|
||||
compatible with the Agent class tool interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
description: str,
|
||||
session: Any, # ClientSession from mcp package
|
||||
input_schema: Optional[Dict[str, Any]] = None,
|
||||
timeout: int = 60
|
||||
):
|
||||
"""
|
||||
Initialize WebSocket MCP tool wrapper.
|
||||
|
||||
Args:
|
||||
name: Tool name
|
||||
description: Tool description
|
||||
session: MCP ClientSession for making calls
|
||||
input_schema: JSON Schema for tool input
|
||||
timeout: Timeout for tool calls in seconds
|
||||
"""
|
||||
self.name = name
|
||||
self.__name__ = name
|
||||
self.__qualname__ = name
|
||||
self.__doc__ = description
|
||||
self.description = description
|
||||
self.session = session
|
||||
self.input_schema = input_schema or {}
|
||||
self.timeout = timeout
|
||||
|
||||
# Build function signature from input schema
|
||||
self.__signature__ = self._build_signature()
|
||||
|
||||
def _build_signature(self) -> inspect.Signature:
|
||||
"""Build function signature from input schema."""
|
||||
params = []
|
||||
|
||||
if self.input_schema and 'properties' in self.input_schema:
|
||||
required = self.input_schema.get('required', [])
|
||||
|
||||
for param_name, prop_schema in self.input_schema['properties'].items():
|
||||
# Determine type annotation
|
||||
prop_type = prop_schema.get('type', 'string') if isinstance(prop_schema, dict) else 'string'
|
||||
annotation = self._json_type_to_python(prop_type)
|
||||
|
||||
# Determine default value
|
||||
default = inspect.Parameter.empty if param_name in required else None
|
||||
|
||||
params.append(
|
||||
inspect.Parameter(
|
||||
name=param_name,
|
||||
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
default=default,
|
||||
annotation=annotation
|
||||
)
|
||||
)
|
||||
|
||||
return inspect.Signature(params)
|
||||
|
||||
def _json_type_to_python(self, json_type: str) -> type:
|
||||
"""Convert JSON Schema type to Python type."""
|
||||
type_map = {
|
||||
'string': str,
|
||||
'integer': int,
|
||||
'number': float,
|
||||
'boolean': bool,
|
||||
'array': list,
|
||||
'object': dict
|
||||
}
|
||||
return type_map.get(json_type, Any)
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
"""
|
||||
Call the tool synchronously.
|
||||
|
||||
This wraps the async call for synchronous usage.
|
||||
"""
|
||||
logger.debug(f"Tool {self.name} called with args: {kwargs}")
|
||||
|
||||
loop = get_event_loop()
|
||||
future = asyncio.run_coroutine_threadsafe(self._async_call(**kwargs), loop)
|
||||
|
||||
try:
|
||||
return future.result(timeout=self.timeout)
|
||||
except Exception as e:
|
||||
logger.error(f"Error calling tool {self.name}: {e}")
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
async def _async_call(self, **kwargs):
|
||||
"""Call the tool asynchronously."""
|
||||
logger.debug(f"Async calling tool {self.name} with args: {kwargs}")
|
||||
|
||||
try:
|
||||
result = await self.session.call_tool(self.name, kwargs)
|
||||
|
||||
# Extract text from result
|
||||
if hasattr(result, 'content') and result.content:
|
||||
if hasattr(result.content[0], 'text'):
|
||||
return result.content[0].text
|
||||
return str(result.content[0])
|
||||
return str(result)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in _async_call for {self.name}: {e}")
|
||||
raise
|
||||
|
||||
def _fix_array_schemas(self, schema: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Fix array schemas by adding missing 'items' attribute.
|
||||
|
||||
OpenAI requires array types to specify item types.
|
||||
"""
|
||||
if not isinstance(schema, dict):
|
||||
return schema
|
||||
|
||||
fixed_schema = schema.copy()
|
||||
|
||||
# Fix array types at current level
|
||||
if fixed_schema.get("type") == "array" and "items" not in fixed_schema:
|
||||
fixed_schema["items"] = {"type": "string"}
|
||||
|
||||
# Recursively fix nested schemas
|
||||
if "properties" in fixed_schema:
|
||||
fixed_properties = {}
|
||||
for prop_name, prop_schema in fixed_schema["properties"].items():
|
||||
fixed_properties[prop_name] = self._fix_array_schemas(prop_schema)
|
||||
fixed_schema["properties"] = fixed_properties
|
||||
|
||||
if "items" in fixed_schema:
|
||||
fixed_schema["items"] = self._fix_array_schemas(fixed_schema["items"])
|
||||
|
||||
return fixed_schema
|
||||
|
||||
def to_openai_tool(self) -> Dict[str, Any]:
|
||||
"""Convert the tool to OpenAI function calling format."""
|
||||
fixed_schema = self._fix_array_schemas(self.input_schema)
|
||||
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"parameters": fixed_schema
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class WebSocketMCPClient:
|
||||
"""
|
||||
High-level client for connecting to MCP servers over WebSocket.
|
||||
|
||||
This client handles:
|
||||
- Connection establishment and lifecycle
|
||||
- Tool discovery and wrapping
|
||||
- Session management
|
||||
- Automatic reconnection
|
||||
|
||||
Example:
|
||||
```python
|
||||
from praisonaiagents.mcp.mcp_websocket import WebSocketMCPClient
|
||||
|
||||
client = WebSocketMCPClient("ws://localhost:8080/mcp")
|
||||
for tool in client.tools:
|
||||
print(tool.name)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_url: str,
|
||||
debug: bool = False,
|
||||
timeout: int = 60,
|
||||
auth_token: Optional[str] = None,
|
||||
options: Optional[Dict[str, Any]] = None
|
||||
):
|
||||
"""
|
||||
Initialize WebSocket MCP client.
|
||||
|
||||
Args:
|
||||
server_url: WebSocket URL (ws:// or wss://)
|
||||
debug: Enable debug logging
|
||||
timeout: Timeout for operations in seconds
|
||||
auth_token: Optional authentication token
|
||||
options: Additional configuration options
|
||||
"""
|
||||
self.server_url = server_url
|
||||
self.debug = debug
|
||||
self.timeout = timeout
|
||||
self.auth_token = auth_token
|
||||
self.options = options or {}
|
||||
|
||||
self.session = None
|
||||
self.transport = None
|
||||
self.tools: List[WebSocketMCPTool] = []
|
||||
self._closed = False
|
||||
|
||||
# Configure logging
|
||||
if debug:
|
||||
logger.setLevel(logging.DEBUG)
|
||||
else:
|
||||
logger.setLevel(logging.WARNING)
|
||||
|
||||
# Initialize connection
|
||||
self._initialize()
|
||||
|
||||
def _initialize(self):
|
||||
"""Initialize the connection and discover tools."""
|
||||
loop = get_event_loop()
|
||||
|
||||
# Start event loop in background thread if not running
|
||||
def run_event_loop():
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_forever()
|
||||
|
||||
self.loop_thread = threading.Thread(target=run_event_loop, daemon=True)
|
||||
self.loop_thread.start()
|
||||
|
||||
# Run initialization
|
||||
future = asyncio.run_coroutine_threadsafe(self._async_initialize(), loop)
|
||||
self.tools = future.result(timeout=self.timeout)
|
||||
|
||||
async def _async_initialize(self):
|
||||
"""Asynchronously initialize connection and discover tools."""
|
||||
logger.debug(f"Connecting to MCP server at {self.server_url}")
|
||||
|
||||
# Lazy import mcp package
|
||||
try:
|
||||
from mcp import ClientSession
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"mcp package is required. Install it with: pip install mcp"
|
||||
)
|
||||
|
||||
# Create transport
|
||||
self.transport = WebSocketTransport(
|
||||
self.server_url,
|
||||
auth_token=self.auth_token,
|
||||
timeout=self.timeout
|
||||
)
|
||||
await self.transport.connect()
|
||||
|
||||
# Create read/write streams for ClientSession
|
||||
# The ClientSession expects async iterators for read/write
|
||||
read_stream = self._create_read_stream()
|
||||
write_stream = self._create_write_stream()
|
||||
|
||||
# Create and initialize session
|
||||
self._session_context = ClientSession(read_stream, write_stream)
|
||||
self.session = await self._session_context.__aenter__()
|
||||
|
||||
await self.session.initialize()
|
||||
|
||||
# Discover tools
|
||||
logger.debug("Listing tools...")
|
||||
response = await self.session.list_tools()
|
||||
tools_data = response.tools
|
||||
logger.debug(f"Found {len(tools_data)} tools: {[t.name for t in tools_data]}")
|
||||
|
||||
# Create tool wrappers
|
||||
tools = []
|
||||
for tool in tools_data:
|
||||
input_schema = tool.inputSchema if hasattr(tool, 'inputSchema') else None
|
||||
wrapper = WebSocketMCPTool(
|
||||
name=tool.name,
|
||||
description=tool.description if hasattr(tool, 'description') else f"Call the {tool.name} tool",
|
||||
session=self.session,
|
||||
input_schema=input_schema,
|
||||
timeout=self.timeout
|
||||
)
|
||||
tools.append(wrapper)
|
||||
|
||||
return tools
|
||||
|
||||
def _create_read_stream(self):
|
||||
"""Create async iterator for reading from transport."""
|
||||
async def read():
|
||||
while not self._closed:
|
||||
try:
|
||||
message = await self.transport.receive()
|
||||
yield message
|
||||
except Exception as e:
|
||||
if not self._closed:
|
||||
logger.error(f"Read error: {e}")
|
||||
break
|
||||
return read()
|
||||
|
||||
def _create_write_stream(self):
|
||||
"""Create write function for transport."""
|
||||
async def write(message):
|
||||
if hasattr(message, 'to_dict'):
|
||||
message = message.to_dict()
|
||||
await self.transport.send(message)
|
||||
return write
|
||||
|
||||
def __iter__(self):
|
||||
"""Return iterator over tools."""
|
||||
return iter(self.tools)
|
||||
|
||||
def to_openai_tools(self) -> List[Dict[str, Any]]:
|
||||
"""Convert all tools to OpenAI format."""
|
||||
return [tool.to_openai_tool() for tool in self.tools]
|
||||
|
||||
async def aclose(self):
|
||||
"""Async cleanup."""
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
self._closed = True
|
||||
|
||||
try:
|
||||
if hasattr(self, '_session_context') and self._session_context:
|
||||
await self._session_context.__aexit__(None, None, None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
if self.transport:
|
||||
await self.transport.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
"""Synchronous cleanup."""
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
try:
|
||||
loop = get_event_loop()
|
||||
if not loop.is_closed():
|
||||
future = asyncio.run_coroutine_threadsafe(self.aclose(), loop)
|
||||
future.result(timeout=5)
|
||||
except Exception:
|
||||
self._closed = True
|
||||
|
||||
def __del__(self):
|
||||
"""Cleanup on garbage collection."""
|
||||
try:
|
||||
if not self._closed:
|
||||
self.close()
|
||||
except Exception:
|
||||
pass
|
||||
@@ -8,8 +8,111 @@ 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)
|
||||
- Docs management (like Cursor docs)
|
||||
- MCP config management (like Cursor .cursor/mcp/)
|
||||
|
||||
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
|
||||
- DocsManager: Documentation context management
|
||||
- MCPConfigManager: MCP server configuration management
|
||||
"""
|
||||
|
||||
from .memory import Memory
|
||||
from .file_memory import FileMemory, create_memory
|
||||
from .rules_manager import RulesManager, Rule, create_rules_manager
|
||||
from .docs_manager import DocsManager, Doc
|
||||
from .mcp_config import MCPConfigManager, MCPConfig
|
||||
|
||||
__all__ = ["Memory"]
|
||||
# Lazy imports for optional modules to avoid dependency issues and improve startup time
|
||||
def __getattr__(name):
|
||||
if name == "Memory":
|
||||
from .memory import Memory
|
||||
return Memory
|
||||
if name == "AutoMemory":
|
||||
from .auto_memory import AutoMemory
|
||||
return AutoMemory
|
||||
if name == "AutoMemoryExtractor":
|
||||
from .auto_memory import AutoMemoryExtractor
|
||||
return AutoMemoryExtractor
|
||||
if name == "create_auto_memory":
|
||||
from .auto_memory import create_auto_memory
|
||||
return create_auto_memory
|
||||
if name == "WorkflowManager":
|
||||
from .workflows import WorkflowManager
|
||||
return WorkflowManager
|
||||
if name == "Workflow":
|
||||
from .workflows import Workflow
|
||||
return Workflow
|
||||
if name == "WorkflowStep":
|
||||
from .workflows import WorkflowStep
|
||||
return WorkflowStep
|
||||
if name == "WorkflowContext":
|
||||
from .workflows import WorkflowContext
|
||||
return WorkflowContext
|
||||
if name == "StepResult":
|
||||
from .workflows import StepResult
|
||||
return StepResult
|
||||
# Backward compatibility aliases
|
||||
if name == "StepInput":
|
||||
from .workflows import WorkflowContext
|
||||
return WorkflowContext
|
||||
if name == "StepOutput":
|
||||
from .workflows import StepResult
|
||||
return StepResult
|
||||
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
|
||||
# DocsManager and MCPConfigManager are already imported at module level
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
__all__ = [
|
||||
# Core memory
|
||||
"FileMemory",
|
||||
"Memory",
|
||||
"create_memory",
|
||||
# Rules management
|
||||
"RulesManager",
|
||||
"Rule",
|
||||
"create_rules_manager",
|
||||
# Auto memory
|
||||
"AutoMemory",
|
||||
"AutoMemoryExtractor",
|
||||
"create_auto_memory",
|
||||
# Workflows
|
||||
"WorkflowManager",
|
||||
"Workflow",
|
||||
"WorkflowStep",
|
||||
"WorkflowContext",
|
||||
"StepResult",
|
||||
"create_workflow_manager",
|
||||
# Hooks
|
||||
"HooksManager",
|
||||
"HookResult",
|
||||
"create_hooks_manager",
|
||||
# Docs management
|
||||
"DocsManager",
|
||||
"Doc",
|
||||
# MCP config management
|
||||
"MCPConfigManager",
|
||||
"MCPConfig",
|
||||
]
|
||||
@@ -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)
|
||||
@@ -0,0 +1,431 @@
|
||||
"""
|
||||
Docs Manager for PraisonAI Agents.
|
||||
|
||||
Provides reusable documentation context support similar to:
|
||||
- Cursor Notepads (deprecated in favor of docs)
|
||||
- Cursor .cursor/docs/
|
||||
- Project documentation for AI context
|
||||
|
||||
Features:
|
||||
- Auto-discovery of docs from .praison/docs/
|
||||
- Support for markdown, text, and code files
|
||||
- @doc mentions for including specific docs
|
||||
- Priority ordering for docs
|
||||
- Character limits for context window management
|
||||
|
||||
Storage Structure:
|
||||
.praison/docs/
|
||||
├── project-overview.md # Project overview
|
||||
├── architecture.md # Architecture decisions
|
||||
├── api-reference.md # API documentation
|
||||
└── coding-standards.md # Coding guidelines
|
||||
|
||||
Doc File Format (YAML frontmatter + Markdown):
|
||||
---
|
||||
description: Project architecture overview
|
||||
priority: 10
|
||||
tags: [architecture, design]
|
||||
---
|
||||
|
||||
# Architecture Overview
|
||||
This project uses...
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Literal
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Doc:
|
||||
"""A single documentation file with metadata and content."""
|
||||
name: str
|
||||
content: str
|
||||
description: str = ""
|
||||
priority: int = 0
|
||||
tags: List[str] = field(default_factory=list)
|
||||
file_path: Optional[str] = None
|
||||
|
||||
def matches_tag(self, tag: str) -> bool:
|
||||
"""Check if this doc has a specific tag."""
|
||||
return tag.lower() in [t.lower() for t in self.tags]
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"content": self.content,
|
||||
"description": self.description,
|
||||
"priority": self.priority,
|
||||
"tags": self.tags,
|
||||
"file_path": self.file_path
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> 'Doc':
|
||||
return cls(**data)
|
||||
|
||||
|
||||
class DocsManager:
|
||||
"""
|
||||
Manages documentation context for AI agents.
|
||||
|
||||
Provides:
|
||||
- Auto-discovery of docs from .praison/docs/
|
||||
- Global docs from ~/.praison/docs/
|
||||
- @doc mentions for including specific docs
|
||||
- Priority-based ordering
|
||||
"""
|
||||
|
||||
DOCS_DIR_NAME = ".praison/docs"
|
||||
SUPPORTED_EXTENSIONS = [".md", ".txt", ".rst"]
|
||||
|
||||
# Maximum characters per doc (context window management)
|
||||
MAX_DOC_CHARS = 12000
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace_path: Optional[str] = None,
|
||||
global_docs_path: Optional[str] = None,
|
||||
verbose: int = 0
|
||||
):
|
||||
"""
|
||||
Initialize DocsManager.
|
||||
|
||||
Args:
|
||||
workspace_path: Path to workspace/project root
|
||||
global_docs_path: Path to global docs (default: ~/.praison/docs)
|
||||
verbose: Verbosity level
|
||||
"""
|
||||
self.workspace_path = Path(workspace_path) if workspace_path else Path.cwd()
|
||||
self.global_docs_path = Path(global_docs_path) if global_docs_path else Path.home() / ".praison" / "docs"
|
||||
self.verbose = verbose
|
||||
|
||||
self._docs: Dict[str, Doc] = {}
|
||||
self._load_all_docs()
|
||||
|
||||
def _log(self, msg: str, level: int = logging.INFO):
|
||||
"""Log message if verbose."""
|
||||
if self.verbose >= 1:
|
||||
logger.log(level, msg)
|
||||
|
||||
def _parse_frontmatter(self, content: str) -> tuple:
|
||||
"""
|
||||
Parse YAML frontmatter from content.
|
||||
|
||||
Returns:
|
||||
Tuple of (frontmatter_dict, remaining_content)
|
||||
"""
|
||||
frontmatter = {}
|
||||
remaining = content
|
||||
|
||||
# Check for YAML frontmatter
|
||||
if content.startswith("---"):
|
||||
parts = content.split("---", 2)
|
||||
if len(parts) >= 3:
|
||||
try:
|
||||
import yaml
|
||||
frontmatter = yaml.safe_load(parts[1]) or {}
|
||||
remaining = parts[2].strip()
|
||||
except Exception:
|
||||
# If YAML parsing fails, use content as-is
|
||||
pass
|
||||
|
||||
return frontmatter, remaining
|
||||
|
||||
def _load_doc_file(self, file_path: Path, scope: str = "workspace") -> Optional[Doc]:
|
||||
"""Load a single doc file."""
|
||||
try:
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
frontmatter, doc_content = self._parse_frontmatter(content)
|
||||
|
||||
name = file_path.stem
|
||||
|
||||
# Determine priority based on scope
|
||||
base_priority = 100 if scope == "workspace" else 0
|
||||
|
||||
doc = Doc(
|
||||
name=name,
|
||||
content=doc_content,
|
||||
description=frontmatter.get("description", ""),
|
||||
priority=frontmatter.get("priority", base_priority),
|
||||
tags=frontmatter.get("tags", []),
|
||||
file_path=str(file_path)
|
||||
)
|
||||
|
||||
return doc
|
||||
except Exception as e:
|
||||
self._log(f"Failed to load doc {file_path}: {e}", logging.WARNING)
|
||||
return None
|
||||
|
||||
def _load_docs_from_dir(self, docs_dir: Path, scope: str = "workspace"):
|
||||
"""Load all docs from a directory."""
|
||||
if not docs_dir.exists():
|
||||
return
|
||||
|
||||
for ext in self.SUPPORTED_EXTENSIONS:
|
||||
for file_path in docs_dir.glob(f"*{ext}"):
|
||||
if file_path.is_file():
|
||||
doc = self._load_doc_file(file_path, scope)
|
||||
if doc:
|
||||
key = f"{scope}:{doc.name}"
|
||||
self._docs[key] = doc
|
||||
self._log(f"Loaded doc: {key}")
|
||||
|
||||
def _load_all_docs(self):
|
||||
"""Load all docs from global and workspace directories."""
|
||||
self._docs.clear()
|
||||
|
||||
# 1. Load global docs (lowest priority)
|
||||
if self.global_docs_path.exists():
|
||||
self._load_docs_from_dir(self.global_docs_path, "global")
|
||||
for key, doc in self._docs.items():
|
||||
if key.startswith("global:"):
|
||||
doc.priority = doc.priority - 1000 # Lower priority for global
|
||||
|
||||
# 2. Load workspace docs
|
||||
workspace_docs_dir = self.workspace_path / self.DOCS_DIR_NAME.replace("/", os.sep)
|
||||
if workspace_docs_dir.exists():
|
||||
self._load_docs_from_dir(workspace_docs_dir, "workspace")
|
||||
|
||||
self._log(f"Loaded {len(self._docs)} docs total")
|
||||
|
||||
def reload(self):
|
||||
"""Reload all docs from disk."""
|
||||
self._load_all_docs()
|
||||
|
||||
def get_all_docs(self) -> List[Doc]:
|
||||
"""Get all loaded docs, sorted by priority."""
|
||||
docs = list(self._docs.values())
|
||||
docs.sort(key=lambda d: d.priority, reverse=True)
|
||||
return docs
|
||||
|
||||
def get_doc(self, name: str) -> Optional[Doc]:
|
||||
"""Get a specific doc by name."""
|
||||
# Try workspace first, then global
|
||||
for scope in ["workspace", "global"]:
|
||||
key = f"{scope}:{name}"
|
||||
if key in self._docs:
|
||||
return self._docs[key]
|
||||
return None
|
||||
|
||||
def get_docs_by_tag(self, tag: str) -> List[Doc]:
|
||||
"""Get all docs with a specific tag."""
|
||||
docs = [d for d in self._docs.values() if d.matches_tag(tag)]
|
||||
docs.sort(key=lambda d: d.priority, reverse=True)
|
||||
return docs
|
||||
|
||||
def get_docs_for_context(
|
||||
self,
|
||||
include_docs: Optional[List[str]] = None,
|
||||
include_tags: Optional[List[str]] = None,
|
||||
max_chars: int = 50000
|
||||
) -> List[Doc]:
|
||||
"""
|
||||
Get docs for inclusion in agent context.
|
||||
|
||||
Args:
|
||||
include_docs: Specific doc names to include (via @doc mentions)
|
||||
include_tags: Tags to filter by
|
||||
max_chars: Maximum total characters
|
||||
|
||||
Returns:
|
||||
List of docs sorted by priority
|
||||
"""
|
||||
active = []
|
||||
|
||||
# Include specific docs by name
|
||||
if include_docs:
|
||||
for name in include_docs:
|
||||
doc = self.get_doc(name)
|
||||
if doc and doc not in active:
|
||||
active.append(doc)
|
||||
|
||||
# Include docs by tag
|
||||
if include_tags:
|
||||
for tag in include_tags:
|
||||
for doc in self.get_docs_by_tag(tag):
|
||||
if doc not in active:
|
||||
active.append(doc)
|
||||
|
||||
# If no specific docs requested, include all high-priority docs
|
||||
if not include_docs and not include_tags:
|
||||
for doc in self.get_all_docs():
|
||||
if doc.priority >= 100: # Only high-priority docs by default
|
||||
active.append(doc)
|
||||
|
||||
# Sort by priority
|
||||
active.sort(key=lambda d: d.priority, reverse=True)
|
||||
|
||||
return active
|
||||
|
||||
def format_docs_for_prompt(
|
||||
self,
|
||||
include_docs: Optional[List[str]] = None,
|
||||
include_tags: Optional[List[str]] = None,
|
||||
max_chars: int = 50000
|
||||
) -> str:
|
||||
"""
|
||||
Format docs as a string for inclusion in prompts.
|
||||
|
||||
Args:
|
||||
include_docs: Specific doc names to include
|
||||
include_tags: Tags to filter by
|
||||
max_chars: Maximum total characters
|
||||
|
||||
Returns:
|
||||
Formatted docs string
|
||||
"""
|
||||
active_docs = self.get_docs_for_context(include_docs, include_tags, max_chars)
|
||||
|
||||
if not active_docs:
|
||||
return ""
|
||||
|
||||
parts = ["# Project Documentation\n"]
|
||||
total_chars = len(parts[0])
|
||||
|
||||
for doc in active_docs:
|
||||
if total_chars >= max_chars:
|
||||
break
|
||||
|
||||
doc_text = doc.content.strip()
|
||||
if not doc_text:
|
||||
continue
|
||||
|
||||
# Add doc with header
|
||||
if doc.description:
|
||||
header = f"## {doc.name}: {doc.description}"
|
||||
else:
|
||||
header = f"## {doc.name}"
|
||||
|
||||
section = f"{header}\n{doc_text}\n"
|
||||
|
||||
if total_chars + len(section) <= max_chars:
|
||||
parts.append(section)
|
||||
total_chars += len(section)
|
||||
else:
|
||||
# Truncate last doc
|
||||
remaining = max_chars - total_chars
|
||||
if remaining > 100:
|
||||
parts.append(section[:remaining] + "\n... (truncated)")
|
||||
break
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
def create_doc(
|
||||
self,
|
||||
name: str,
|
||||
content: str,
|
||||
description: str = "",
|
||||
priority: int = 0,
|
||||
tags: Optional[List[str]] = None,
|
||||
scope: Literal["global", "workspace"] = "workspace"
|
||||
) -> Doc:
|
||||
"""
|
||||
Create a new doc file.
|
||||
|
||||
Args:
|
||||
name: Doc name (used as filename)
|
||||
content: Doc content (markdown)
|
||||
description: Short description
|
||||
priority: Priority (higher = applied first)
|
||||
tags: Tags for categorization
|
||||
scope: Where to save (global or workspace)
|
||||
|
||||
Returns:
|
||||
Created Doc object
|
||||
"""
|
||||
# Determine save path
|
||||
if scope == "global":
|
||||
docs_dir = self.global_docs_path
|
||||
else:
|
||||
docs_dir = self.workspace_path / self.DOCS_DIR_NAME.replace("/", os.sep)
|
||||
|
||||
docs_dir.mkdir(parents=True, exist_ok=True)
|
||||
file_path = docs_dir / f"{name}.md"
|
||||
|
||||
# Build frontmatter
|
||||
frontmatter_lines = ["---"]
|
||||
if description:
|
||||
frontmatter_lines.append(f'description: "{description}"')
|
||||
if priority != 0:
|
||||
frontmatter_lines.append(f"priority: {priority}")
|
||||
if tags:
|
||||
tags_str = ", ".join(f'"{t}"' for t in tags)
|
||||
frontmatter_lines.append(f"tags: [{tags_str}]")
|
||||
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 doc
|
||||
doc = Doc(
|
||||
name=name,
|
||||
content=content,
|
||||
description=description,
|
||||
priority=priority,
|
||||
tags=tags or [],
|
||||
file_path=str(file_path)
|
||||
)
|
||||
|
||||
self._docs[f"{scope}:{name}"] = doc
|
||||
self._log(f"Created doc '{name}' at {file_path}")
|
||||
|
||||
return doc
|
||||
|
||||
def delete_doc(self, name: str, scope: Optional[str] = None) -> bool:
|
||||
"""
|
||||
Delete a doc file.
|
||||
|
||||
Args:
|
||||
name: Doc name to delete
|
||||
scope: Scope to delete from (None = try both)
|
||||
|
||||
Returns:
|
||||
True if deleted, False if not found
|
||||
"""
|
||||
scopes = [scope] if scope else ["workspace", "global"]
|
||||
|
||||
for s in scopes:
|
||||
key = f"{s}:{name}"
|
||||
if key in self._docs:
|
||||
doc = self._docs[key]
|
||||
if doc.file_path:
|
||||
try:
|
||||
Path(doc.file_path).unlink()
|
||||
del self._docs[key]
|
||||
self._log(f"Deleted doc '{name}' from {s}")
|
||||
return True
|
||||
except Exception as e:
|
||||
self._log(f"Failed to delete doc file: {e}", logging.ERROR)
|
||||
|
||||
return False
|
||||
|
||||
def list_docs(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all docs with metadata.
|
||||
|
||||
Returns:
|
||||
List of doc info dicts
|
||||
"""
|
||||
docs = []
|
||||
for key, doc in self._docs.items():
|
||||
scope = key.split(":")[0]
|
||||
docs.append({
|
||||
"name": doc.name,
|
||||
"scope": scope,
|
||||
"description": doc.description,
|
||||
"priority": doc.priority,
|
||||
"tags": doc.tags,
|
||||
"file_path": doc.file_path,
|
||||
"content_length": len(doc.content)
|
||||
})
|
||||
|
||||
# Sort by priority
|
||||
docs.sort(key=lambda d: d["priority"], reverse=True)
|
||||
return docs
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,417 @@
|
||||
"""
|
||||
Cascade Hooks System for PraisonAI Agents.
|
||||
|
||||
Provides event hooks similar to Windsurf's Cascade Hooks,
|
||||
allowing custom actions before/after agent operations.
|
||||
|
||||
Features:
|
||||
- Pre/post hooks for read, write, command, prompt operations
|
||||
- JSON configuration (.praison/hooks.json)
|
||||
- Script execution with timeout
|
||||
- Exit code handling
|
||||
- Lazy loading for performance
|
||||
|
||||
Configuration (.praison/hooks.json):
|
||||
{
|
||||
"hooks": {
|
||||
"pre_write_code": "./scripts/lint.sh",
|
||||
"post_write_code": "./scripts/format.sh",
|
||||
"pre_run_command": "./scripts/validate.sh",
|
||||
"post_user_prompt": "./scripts/log.sh"
|
||||
},
|
||||
"timeout": 30,
|
||||
"enabled": true
|
||||
}
|
||||
|
||||
Hook Events:
|
||||
- pre_read_code: Before reading a file
|
||||
- post_read_code: After reading a file
|
||||
- pre_write_code: Before writing to a file
|
||||
- post_write_code: After writing to a file
|
||||
- pre_run_command: Before running a terminal command
|
||||
- post_run_command: After running a terminal command
|
||||
- pre_user_prompt: Before processing user prompt
|
||||
- post_user_prompt: After processing user prompt
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Callable, Literal
|
||||
from dataclasses import dataclass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Hook event types
|
||||
HookEvent = Literal[
|
||||
"pre_read_code",
|
||||
"post_read_code",
|
||||
"pre_write_code",
|
||||
"post_write_code",
|
||||
"pre_run_command",
|
||||
"post_run_command",
|
||||
"pre_user_prompt",
|
||||
"post_user_prompt",
|
||||
"pre_mcp_tool_use",
|
||||
"post_mcp_tool_use"
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookResult:
|
||||
"""Result of a hook execution."""
|
||||
success: bool
|
||||
exit_code: int = 0
|
||||
stdout: str = ""
|
||||
stderr: str = ""
|
||||
blocked: bool = False # If True, the operation should be blocked
|
||||
modified_input: Optional[str] = None # Modified input if applicable
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookConfig:
|
||||
"""Configuration for a single hook."""
|
||||
event: HookEvent
|
||||
command: str
|
||||
timeout: int = 30
|
||||
enabled: bool = True
|
||||
block_on_failure: bool = False
|
||||
pass_input: bool = True
|
||||
|
||||
|
||||
class HooksManager:
|
||||
"""
|
||||
Manages hook discovery, loading, and execution.
|
||||
|
||||
Hooks are configured in .praison/hooks.json and can execute
|
||||
scripts before/after agent operations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
hooks = HooksManager(workspace_path="/path/to/project")
|
||||
|
||||
# Execute pre-write hook
|
||||
result = hooks.execute(
|
||||
"pre_write_code",
|
||||
context={"file_path": "main.py", "content": "..."}
|
||||
)
|
||||
|
||||
if result.blocked:
|
||||
print("Write blocked by hook")
|
||||
```
|
||||
"""
|
||||
|
||||
CONFIG_FILE = ".praison/hooks.json"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace_path: Optional[str] = None,
|
||||
verbose: int = 0
|
||||
):
|
||||
"""
|
||||
Initialize HooksManager.
|
||||
|
||||
Args:
|
||||
workspace_path: Path to workspace/project root
|
||||
verbose: Verbosity level
|
||||
"""
|
||||
self.workspace_path = Path(workspace_path) if workspace_path else Path.cwd()
|
||||
self.verbose = verbose
|
||||
|
||||
self._hooks: Dict[HookEvent, List[HookConfig]] = {}
|
||||
self._global_timeout = 30
|
||||
self._enabled = True
|
||||
self._loaded = False
|
||||
|
||||
# Python callable hooks (for programmatic use)
|
||||
self._callable_hooks: Dict[HookEvent, List[Callable]] = {}
|
||||
|
||||
def _log(self, msg: str, level: int = logging.INFO):
|
||||
"""Log message if verbose."""
|
||||
if self.verbose >= 1:
|
||||
logger.log(level, msg)
|
||||
|
||||
def _ensure_loaded(self):
|
||||
"""Lazy load hooks on first access."""
|
||||
if not self._loaded:
|
||||
self._load_config()
|
||||
self._loaded = True
|
||||
|
||||
def _load_config(self):
|
||||
"""Load hooks configuration from file."""
|
||||
config_path = self.workspace_path / self.CONFIG_FILE.replace("/", os.sep)
|
||||
|
||||
if not config_path.exists():
|
||||
return
|
||||
|
||||
try:
|
||||
config = json.loads(config_path.read_text(encoding="utf-8"))
|
||||
|
||||
self._enabled = config.get("enabled", True)
|
||||
self._global_timeout = config.get("timeout", 30)
|
||||
|
||||
hooks_config = config.get("hooks", {})
|
||||
|
||||
for event, command in hooks_config.items():
|
||||
if event not in self._hooks:
|
||||
self._hooks[event] = []
|
||||
|
||||
if isinstance(command, str):
|
||||
# Simple command string
|
||||
self._hooks[event].append(HookConfig(
|
||||
event=event,
|
||||
command=command,
|
||||
timeout=self._global_timeout
|
||||
))
|
||||
elif isinstance(command, dict):
|
||||
# Detailed config
|
||||
self._hooks[event].append(HookConfig(
|
||||
event=event,
|
||||
command=command.get("command", ""),
|
||||
timeout=command.get("timeout", self._global_timeout),
|
||||
enabled=command.get("enabled", True),
|
||||
block_on_failure=command.get("block_on_failure", False),
|
||||
pass_input=command.get("pass_input", True)
|
||||
))
|
||||
elif isinstance(command, list):
|
||||
# Multiple commands for same event
|
||||
for cmd in command:
|
||||
if isinstance(cmd, str):
|
||||
self._hooks[event].append(HookConfig(
|
||||
event=event,
|
||||
command=cmd,
|
||||
timeout=self._global_timeout
|
||||
))
|
||||
|
||||
self._log(f"Loaded {sum(len(h) for h in self._hooks.values())} hooks")
|
||||
|
||||
except Exception as e:
|
||||
self._log(f"Error loading hooks config: {e}", logging.WARNING)
|
||||
|
||||
def reload(self):
|
||||
"""Reload hooks configuration."""
|
||||
self._hooks = {}
|
||||
self._loaded = False
|
||||
self._ensure_loaded()
|
||||
|
||||
def register(
|
||||
self,
|
||||
event: HookEvent,
|
||||
handler: Callable[[Dict[str, Any]], Optional[Dict[str, Any]]]
|
||||
):
|
||||
"""
|
||||
Register a Python callable as a hook.
|
||||
|
||||
Args:
|
||||
event: Hook event type
|
||||
handler: Callable that receives context dict and optionally returns modified context
|
||||
"""
|
||||
if event not in self._callable_hooks:
|
||||
self._callable_hooks[event] = []
|
||||
self._callable_hooks[event].append(handler)
|
||||
|
||||
def unregister(self, event: HookEvent, handler: Callable):
|
||||
"""Unregister a callable hook."""
|
||||
if event in self._callable_hooks:
|
||||
self._callable_hooks[event] = [
|
||||
h for h in self._callable_hooks[event] if h != handler
|
||||
]
|
||||
|
||||
def execute(
|
||||
self,
|
||||
event: HookEvent,
|
||||
context: Optional[Dict[str, Any]] = None
|
||||
) -> HookResult:
|
||||
"""
|
||||
Execute hooks for an event.
|
||||
|
||||
Args:
|
||||
event: Hook event type
|
||||
context: Context data to pass to hooks
|
||||
|
||||
Returns:
|
||||
HookResult with execution status
|
||||
"""
|
||||
self._ensure_loaded()
|
||||
|
||||
if not self._enabled:
|
||||
return HookResult(success=True)
|
||||
|
||||
context = context or {}
|
||||
combined_result = HookResult(success=True)
|
||||
|
||||
# Execute callable hooks first
|
||||
if event in self._callable_hooks:
|
||||
for handler in self._callable_hooks[event]:
|
||||
try:
|
||||
result = handler(context)
|
||||
if isinstance(result, dict):
|
||||
context.update(result)
|
||||
except Exception as e:
|
||||
self._log(f"Callable hook error for {event}: {e}", logging.WARNING)
|
||||
combined_result.success = False
|
||||
|
||||
# Execute script hooks
|
||||
if event in self._hooks:
|
||||
for hook in self._hooks[event]:
|
||||
if not hook.enabled:
|
||||
continue
|
||||
|
||||
result = self._execute_script(hook, context)
|
||||
|
||||
if not result.success:
|
||||
combined_result.success = False
|
||||
combined_result.stderr += result.stderr
|
||||
|
||||
if hook.block_on_failure:
|
||||
combined_result.blocked = True
|
||||
break
|
||||
|
||||
combined_result.stdout += result.stdout
|
||||
|
||||
if result.modified_input:
|
||||
combined_result.modified_input = result.modified_input
|
||||
|
||||
return combined_result
|
||||
|
||||
def _execute_script(
|
||||
self,
|
||||
hook: HookConfig,
|
||||
context: Dict[str, Any]
|
||||
) -> HookResult:
|
||||
"""Execute a script hook."""
|
||||
try:
|
||||
# Prepare environment
|
||||
env = os.environ.copy()
|
||||
|
||||
# Pass context as environment variables
|
||||
if hook.pass_input:
|
||||
for key, value in context.items():
|
||||
if isinstance(value, (str, int, float, bool)):
|
||||
env[f"PRAISON_HOOK_{key.upper()}"] = str(value)
|
||||
|
||||
# Also pass as JSON
|
||||
env["PRAISON_HOOK_CONTEXT"] = json.dumps(context)
|
||||
|
||||
# Resolve command path
|
||||
command = hook.command
|
||||
if not os.path.isabs(command) and not command.startswith("./"):
|
||||
command = str(self.workspace_path / command)
|
||||
|
||||
# Execute
|
||||
result = subprocess.run(
|
||||
command,
|
||||
shell=True,
|
||||
cwd=str(self.workspace_path),
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=hook.timeout
|
||||
)
|
||||
|
||||
success = result.returncode == 0
|
||||
|
||||
# Check for special exit codes
|
||||
blocked = result.returncode == 1 and hook.block_on_failure
|
||||
|
||||
# Check stdout for modified input
|
||||
modified_input = None
|
||||
if result.stdout.startswith("MODIFIED:"):
|
||||
modified_input = result.stdout[9:].strip()
|
||||
|
||||
self._log(f"Hook {hook.event} executed: exit={result.returncode}")
|
||||
|
||||
return HookResult(
|
||||
success=success,
|
||||
exit_code=result.returncode,
|
||||
stdout=result.stdout,
|
||||
stderr=result.stderr,
|
||||
blocked=blocked,
|
||||
modified_input=modified_input
|
||||
)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
self._log(f"Hook {hook.event} timed out after {hook.timeout}s", logging.WARNING)
|
||||
return HookResult(
|
||||
success=False,
|
||||
exit_code=-1,
|
||||
stderr=f"Timeout after {hook.timeout}s"
|
||||
)
|
||||
except Exception as e:
|
||||
self._log(f"Hook {hook.event} error: {e}", logging.WARNING)
|
||||
return HookResult(
|
||||
success=False,
|
||||
exit_code=-1,
|
||||
stderr=str(e)
|
||||
)
|
||||
|
||||
def has_hooks(self, event: HookEvent) -> bool:
|
||||
"""Check if any hooks are registered for an event."""
|
||||
self._ensure_loaded()
|
||||
return (
|
||||
event in self._hooks and len(self._hooks[event]) > 0 or
|
||||
event in self._callable_hooks and len(self._callable_hooks[event]) > 0
|
||||
)
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get hooks statistics."""
|
||||
self._ensure_loaded()
|
||||
|
||||
script_hooks = sum(len(h) for h in self._hooks.values())
|
||||
callable_hooks = sum(len(h) for h in self._callable_hooks.values())
|
||||
|
||||
return {
|
||||
"enabled": self._enabled,
|
||||
"script_hooks": script_hooks,
|
||||
"callable_hooks": callable_hooks,
|
||||
"total_hooks": script_hooks + callable_hooks,
|
||||
"events": list(set(list(self._hooks.keys()) + list(self._callable_hooks.keys())))
|
||||
}
|
||||
|
||||
def create_config(
|
||||
self,
|
||||
hooks: Optional[Dict[str, str]] = None,
|
||||
timeout: int = 30,
|
||||
enabled: bool = True
|
||||
):
|
||||
"""
|
||||
Create a hooks configuration file.
|
||||
|
||||
Args:
|
||||
hooks: Dict mapping event names to commands
|
||||
timeout: Global timeout in seconds
|
||||
enabled: Whether hooks are enabled
|
||||
"""
|
||||
config_dir = self.workspace_path / ".praison"
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
config = {
|
||||
"enabled": enabled,
|
||||
"timeout": timeout,
|
||||
"hooks": hooks or {}
|
||||
}
|
||||
|
||||
config_path = config_dir / "hooks.json"
|
||||
config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
|
||||
|
||||
self._log(f"Created hooks config at {config_path}")
|
||||
self.reload()
|
||||
|
||||
|
||||
def create_hooks_manager(
|
||||
workspace_path: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> HooksManager:
|
||||
"""
|
||||
Create a HooksManager instance.
|
||||
|
||||
Args:
|
||||
workspace_path: Path to workspace
|
||||
**kwargs: Additional configuration
|
||||
|
||||
Returns:
|
||||
HooksManager instance
|
||||
"""
|
||||
return HooksManager(workspace_path=workspace_path, **kwargs)
|
||||
@@ -0,0 +1,369 @@
|
||||
"""
|
||||
MCP Config Manager for PraisonAI Agents.
|
||||
|
||||
Provides JSON-based MCP server configuration support similar to:
|
||||
- Cursor .cursor/mcp/*.mcpjson
|
||||
- Windsurf MCP configuration
|
||||
|
||||
Features:
|
||||
- Auto-discovery of MCP configs from .praison/mcp/
|
||||
- JSON config files for MCP servers
|
||||
- Environment variable support
|
||||
- Easy server management
|
||||
|
||||
Storage Structure:
|
||||
.praison/mcp/
|
||||
├── filesystem.json # Filesystem MCP server
|
||||
├── brave-search.json # Brave search MCP server
|
||||
└── postgres.json # PostgreSQL MCP server
|
||||
|
||||
Config File Format (JSON):
|
||||
{
|
||||
"name": "filesystem",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-filesystem", "."],
|
||||
"env": {
|
||||
"SOME_VAR": "value"
|
||||
},
|
||||
"enabled": true,
|
||||
"description": "Filesystem access for the agent"
|
||||
}
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPConfig:
|
||||
"""A single MCP server configuration."""
|
||||
name: str
|
||||
command: str
|
||||
args: List[str] = field(default_factory=list)
|
||||
env: Dict[str, str] = field(default_factory=dict)
|
||||
enabled: bool = True
|
||||
description: str = ""
|
||||
file_path: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"command": self.command,
|
||||
"args": self.args,
|
||||
"env": self.env,
|
||||
"enabled": self.enabled,
|
||||
"description": self.description
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any], file_path: Optional[str] = None) -> 'MCPConfig':
|
||||
return cls(
|
||||
name=data.get("name", ""),
|
||||
command=data.get("command", ""),
|
||||
args=data.get("args", []),
|
||||
env=data.get("env", {}),
|
||||
enabled=data.get("enabled", True),
|
||||
description=data.get("description", ""),
|
||||
file_path=file_path
|
||||
)
|
||||
|
||||
def to_mcp_instance(self):
|
||||
"""
|
||||
Convert this config to an MCP instance.
|
||||
|
||||
Returns:
|
||||
MCP instance ready for use with agents
|
||||
"""
|
||||
try:
|
||||
from praisonaiagents.tools.mcp import MCP
|
||||
|
||||
# Resolve environment variables
|
||||
resolved_env = {}
|
||||
for key, value in self.env.items():
|
||||
if value.startswith("$"):
|
||||
# Environment variable reference
|
||||
env_var = value[1:]
|
||||
resolved_env[key] = os.environ.get(env_var, value)
|
||||
else:
|
||||
resolved_env[key] = value
|
||||
|
||||
return MCP(
|
||||
command=self.command,
|
||||
args=self.args,
|
||||
env=resolved_env if resolved_env else None
|
||||
)
|
||||
except ImportError:
|
||||
logger.warning("MCP tools not available")
|
||||
return None
|
||||
|
||||
|
||||
class MCPConfigManager:
|
||||
"""
|
||||
Manages MCP server configurations for AI agents.
|
||||
|
||||
Provides:
|
||||
- Auto-discovery of MCP configs from .praison/mcp/
|
||||
- Global configs from ~/.praison/mcp/
|
||||
- Easy server management
|
||||
"""
|
||||
|
||||
MCP_DIR_NAME = ".praison/mcp"
|
||||
SUPPORTED_EXTENSIONS = [".json", ".mcpjson"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace_path: Optional[str] = None,
|
||||
global_mcp_path: Optional[str] = None,
|
||||
verbose: int = 0
|
||||
):
|
||||
"""
|
||||
Initialize MCPConfigManager.
|
||||
|
||||
Args:
|
||||
workspace_path: Path to workspace/project root
|
||||
global_mcp_path: Path to global MCP configs (default: ~/.praison/mcp)
|
||||
verbose: Verbosity level
|
||||
"""
|
||||
self.workspace_path = Path(workspace_path) if workspace_path else Path.cwd()
|
||||
self.global_mcp_path = Path(global_mcp_path) if global_mcp_path else Path.home() / ".praison" / "mcp"
|
||||
self.verbose = verbose
|
||||
|
||||
self._configs: Dict[str, MCPConfig] = {}
|
||||
self._load_all_configs()
|
||||
|
||||
def _log(self, msg: str, level: int = logging.INFO):
|
||||
"""Log message if verbose."""
|
||||
if self.verbose >= 1:
|
||||
logger.log(level, msg)
|
||||
|
||||
def _load_config_file(self, file_path: Path, scope: str = "workspace") -> Optional[MCPConfig]:
|
||||
"""Load a single MCP config file."""
|
||||
try:
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
data = json.loads(content)
|
||||
|
||||
# Use filename as name if not specified
|
||||
if not data.get("name"):
|
||||
data["name"] = file_path.stem
|
||||
|
||||
config = MCPConfig.from_dict(data, str(file_path))
|
||||
return config
|
||||
except json.JSONDecodeError as e:
|
||||
self._log(f"Invalid JSON in {file_path}: {e}", logging.WARNING)
|
||||
return None
|
||||
except Exception as e:
|
||||
self._log(f"Failed to load MCP config {file_path}: {e}", logging.WARNING)
|
||||
return None
|
||||
|
||||
def _load_configs_from_dir(self, mcp_dir: Path, scope: str = "workspace"):
|
||||
"""Load all MCP configs from a directory."""
|
||||
if not mcp_dir.exists():
|
||||
return
|
||||
|
||||
for ext in self.SUPPORTED_EXTENSIONS:
|
||||
for file_path in mcp_dir.glob(f"*{ext}"):
|
||||
if file_path.is_file():
|
||||
config = self._load_config_file(file_path, scope)
|
||||
if config:
|
||||
key = f"{scope}:{config.name}"
|
||||
self._configs[key] = config
|
||||
self._log(f"Loaded MCP config: {key}")
|
||||
|
||||
def _load_all_configs(self):
|
||||
"""Load all MCP configs from global and workspace directories."""
|
||||
self._configs.clear()
|
||||
|
||||
# 1. Load global configs
|
||||
if self.global_mcp_path.exists():
|
||||
self._load_configs_from_dir(self.global_mcp_path, "global")
|
||||
|
||||
# 2. Load workspace configs (override global)
|
||||
workspace_mcp_dir = self.workspace_path / self.MCP_DIR_NAME.replace("/", os.sep)
|
||||
if workspace_mcp_dir.exists():
|
||||
self._load_configs_from_dir(workspace_mcp_dir, "workspace")
|
||||
|
||||
self._log(f"Loaded {len(self._configs)} MCP configs total")
|
||||
|
||||
def reload(self):
|
||||
"""Reload all configs from disk."""
|
||||
self._load_all_configs()
|
||||
|
||||
def get_all_configs(self) -> List[MCPConfig]:
|
||||
"""Get all loaded MCP configs."""
|
||||
return list(self._configs.values())
|
||||
|
||||
def get_enabled_configs(self) -> List[MCPConfig]:
|
||||
"""Get all enabled MCP configs."""
|
||||
return [c for c in self._configs.values() if c.enabled]
|
||||
|
||||
def get_config(self, name: str) -> Optional[MCPConfig]:
|
||||
"""Get a specific MCP config by name."""
|
||||
# Try workspace first, then global
|
||||
for scope in ["workspace", "global"]:
|
||||
key = f"{scope}:{name}"
|
||||
if key in self._configs:
|
||||
return self._configs[key]
|
||||
return None
|
||||
|
||||
def get_mcp_tools(self, names: Optional[List[str]] = None) -> List:
|
||||
"""
|
||||
Get MCP instances for use with agents.
|
||||
|
||||
Args:
|
||||
names: Specific config names to load (None = all enabled)
|
||||
|
||||
Returns:
|
||||
List of MCP instances
|
||||
"""
|
||||
tools = []
|
||||
|
||||
if names:
|
||||
configs = [self.get_config(name) for name in names]
|
||||
configs = [c for c in configs if c and c.enabled]
|
||||
else:
|
||||
configs = self.get_enabled_configs()
|
||||
|
||||
for config in configs:
|
||||
mcp = config.to_mcp_instance()
|
||||
if mcp:
|
||||
tools.append(mcp)
|
||||
|
||||
return tools
|
||||
|
||||
def create_config(
|
||||
self,
|
||||
name: str,
|
||||
command: str,
|
||||
args: Optional[List[str]] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
enabled: bool = True,
|
||||
description: str = "",
|
||||
scope: str = "workspace"
|
||||
) -> MCPConfig:
|
||||
"""
|
||||
Create a new MCP config file.
|
||||
|
||||
Args:
|
||||
name: Config name (used as filename)
|
||||
command: Command to run (e.g., "npx")
|
||||
args: Command arguments
|
||||
env: Environment variables
|
||||
enabled: Whether the config is enabled
|
||||
description: Description of the MCP server
|
||||
scope: Where to save (global or workspace)
|
||||
|
||||
Returns:
|
||||
Created MCPConfig object
|
||||
"""
|
||||
# Determine save path
|
||||
if scope == "global":
|
||||
mcp_dir = self.global_mcp_path
|
||||
else:
|
||||
mcp_dir = self.workspace_path / self.MCP_DIR_NAME.replace("/", os.sep)
|
||||
|
||||
mcp_dir.mkdir(parents=True, exist_ok=True)
|
||||
file_path = mcp_dir / f"{name}.json"
|
||||
|
||||
# Build config
|
||||
config = MCPConfig(
|
||||
name=name,
|
||||
command=command,
|
||||
args=args or [],
|
||||
env=env or {},
|
||||
enabled=enabled,
|
||||
description=description,
|
||||
file_path=str(file_path)
|
||||
)
|
||||
|
||||
# Write file
|
||||
file_path.write_text(
|
||||
json.dumps(config.to_dict(), indent=2),
|
||||
encoding="utf-8"
|
||||
)
|
||||
|
||||
self._configs[f"{scope}:{name}"] = config
|
||||
self._log(f"Created MCP config '{name}' at {file_path}")
|
||||
|
||||
return config
|
||||
|
||||
def delete_config(self, name: str, scope: Optional[str] = None) -> bool:
|
||||
"""
|
||||
Delete an MCP config file.
|
||||
|
||||
Args:
|
||||
name: Config name to delete
|
||||
scope: Scope to delete from (None = try both)
|
||||
|
||||
Returns:
|
||||
True if deleted, False if not found
|
||||
"""
|
||||
scopes = [scope] if scope else ["workspace", "global"]
|
||||
|
||||
for s in scopes:
|
||||
key = f"{s}:{name}"
|
||||
if key in self._configs:
|
||||
config = self._configs[key]
|
||||
if config.file_path:
|
||||
try:
|
||||
Path(config.file_path).unlink()
|
||||
del self._configs[key]
|
||||
self._log(f"Deleted MCP config '{name}' from {s}")
|
||||
return True
|
||||
except Exception as e:
|
||||
self._log(f"Failed to delete config file: {e}", logging.ERROR)
|
||||
|
||||
return False
|
||||
|
||||
def enable_config(self, name: str) -> bool:
|
||||
"""Enable an MCP config."""
|
||||
config = self.get_config(name)
|
||||
if config:
|
||||
config.enabled = True
|
||||
self._save_config(config)
|
||||
return True
|
||||
return False
|
||||
|
||||
def disable_config(self, name: str) -> bool:
|
||||
"""Disable an MCP config."""
|
||||
config = self.get_config(name)
|
||||
if config:
|
||||
config.enabled = False
|
||||
self._save_config(config)
|
||||
return True
|
||||
return False
|
||||
|
||||
def _save_config(self, config: MCPConfig):
|
||||
"""Save a config back to its file."""
|
||||
if config.file_path:
|
||||
Path(config.file_path).write_text(
|
||||
json.dumps(config.to_dict(), indent=2),
|
||||
encoding="utf-8"
|
||||
)
|
||||
|
||||
def list_configs(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all MCP configs with metadata.
|
||||
|
||||
Returns:
|
||||
List of config info dicts
|
||||
"""
|
||||
configs = []
|
||||
for key, config in self._configs.items():
|
||||
scope = key.split(":")[0]
|
||||
configs.append({
|
||||
"name": config.name,
|
||||
"scope": scope,
|
||||
"command": config.command,
|
||||
"args": config.args,
|
||||
"enabled": config.enabled,
|
||||
"description": config.description,
|
||||
"file_path": config.file_path
|
||||
})
|
||||
|
||||
return configs
|
||||
@@ -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)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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)
|
||||
@@ -44,7 +44,9 @@ class Task:
|
||||
retain_full_context: bool = False, # By default, only use previous task output, not all previous tasks
|
||||
guardrail: Optional[Union[Callable[[TaskOutput], Tuple[bool, Any]], str]] = None,
|
||||
max_retries: int = 3,
|
||||
retry_count: int = 0
|
||||
retry_count: int = 0,
|
||||
agent_config: Optional[Dict[str, Any]] = None,
|
||||
variables: Optional[Dict[str, Any]] = None
|
||||
):
|
||||
# Add check if memory config is provided
|
||||
if memory is not None or (config and config.get('memory_config')):
|
||||
@@ -89,6 +91,8 @@ class Task:
|
||||
self.retry_count = retry_count
|
||||
self._guardrail_fn = None
|
||||
self.validation_feedback = None # Store validation failure feedback for retry attempts
|
||||
self.agent_config = agent_config # Per-task agent configuration {role, goal, backstory, llm}
|
||||
self.variables = variables if variables else {} # Variables for substitution in description
|
||||
|
||||
# Set logger level based on config verbose level
|
||||
verbose = self.config.get("verbose", 0)
|
||||
|
||||
Loaded 100 of 196 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user