mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-09 03:39:37 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d3aa78241a | ||
|
|
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 | ||
|
|
4a85c135a6 | ||
|
|
cdb42a3f08 | ||
|
|
c8378118e7 | ||
|
|
0213764057 | ||
|
|
2220114809 | ||
|
|
9fe4dbc172 | ||
|
|
180503fcd5 | ||
|
|
0d8f936b19 | ||
|
|
8b53ebf5e3 | ||
|
|
504c5a9759 | ||
|
|
8fef42fc03 | ||
|
|
2f18cc1ea2 | ||
|
|
ff7969cdbf | ||
|
|
5fa84dc3eb | ||
|
|
4bf8b980ed | ||
|
|
53e0fd19ab | ||
|
|
0cb6307d3c | ||
|
|
adce94410b | ||
|
|
42aec561b1 | ||
|
|
acc805768d | ||
|
|
6d9fa167e3 | ||
|
|
9e178b70e6 | ||
|
|
ed406973cd |
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.87" \
|
||||
"praisonai>=2.3.20" \
|
||||
"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.87" \
|
||||
"praisonai>=2.3.20" \
|
||||
"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.87" \
|
||||
"praisonai>=2.3.20" \
|
||||
"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.87" \
|
||||
"praisonai>=2.3.20" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
+2
-2
@@ -121,7 +121,7 @@ healthcheck:
|
||||
## 📦 Package Versions
|
||||
|
||||
All Docker images use consistent, up-to-date versions:
|
||||
- PraisonAI: `>=2.2.87`
|
||||
- PraisonAI: `>=2.2.95`
|
||||
- PraisonAI Agents: `>=0.0.92`
|
||||
- Python: `3.11-slim`
|
||||
|
||||
@@ -218,7 +218,7 @@ docker-compose up -d
|
||||
### Version Pinning
|
||||
To use specific versions, update the Dockerfile:
|
||||
```dockerfile
|
||||
RUN pip install "praisonai==2.2.87" "praisonaiagents==0.0.92"
|
||||
RUN pip install "praisonai==2.2.95" "praisonaiagents==0.0.92"
|
||||
```
|
||||
|
||||
## 🌐 Production Deployment
|
||||
|
||||
@@ -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,481 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "LJRJRQtFTagH"
|
||||
},
|
||||
"source": [
|
||||
"# AI RAG Agent with Web Access using PraisonAI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Ph4RVTBCTcmB"
|
||||
},
|
||||
"source": [
|
||||
"This notebook demonstrates how to build a Retrieval-Augmented Generation (RAG) agent with web access using GPT-4o and PraisonAI. The agent specializes in Thai recipes and can search the web for additional information."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tMxfwPPATy_o"
|
||||
},
|
||||
"source": [
|
||||
"This RAG agent:\n",
|
||||
"* Creates an AI assistant specialized in Thai recipes and cooking\n",
|
||||
"* Has access to a PDF knowledge base about Thai recipes\n",
|
||||
"* Can search the web using DuckDuckGo for additional information\n",
|
||||
"* Provides comprehensive answers about Thai cuisine\n",
|
||||
"* Uses PraisonAI framework for agent management"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "h68oMaCITmqt"
|
||||
},
|
||||
"source": [
|
||||
"[](https://colab.research.google.com/github/Dhivya-Bharathy/PraisonAI/blob/main/examples/cookbooks/rag_examples/agentic_rag/PraisonAI_RecipeAgent.ipynb)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "cAdzoWXJUMfF"
|
||||
},
|
||||
"source": [
|
||||
"# Dependencies"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "9D-2TlOrUNcu"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install praisonaiagents openai requests"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "yY-4-kEIUW40"
|
||||
},
|
||||
"source": [
|
||||
"# Set Up API Key"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "f5bhQ05yUX6J",
|
||||
"outputId": "dcf9d0d9-332e-4136-9289-4c39d898005c"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Enter your OpenAI API key: ··········\n",
|
||||
"✅ API key set successfully!\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"# Set your OpenAI API key here\n",
|
||||
"os.environ['OPENAI_API_KEY'] = 'Enter your api key here'\n",
|
||||
"\n",
|
||||
"# Or use this method to input it securely\n",
|
||||
"from getpass import getpass\n",
|
||||
"api_key = getpass('Enter your OpenAI API key: ')\n",
|
||||
"os.environ['OPENAI_API_KEY'] = api_key\n",
|
||||
"\n",
|
||||
"print(\"✅ API key set successfully!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Eh0j-SegUi13"
|
||||
},
|
||||
"source": [
|
||||
"# Tools"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Z1A9oifIUkQN",
|
||||
"outputId": "9be78226-36fa-43c4-865d-2f7e6e2ba9de"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"✅ Tools imported successfully!\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from praisonaiagents import Agent\n",
|
||||
"import requests\n",
|
||||
"import tempfile\n",
|
||||
"from pathlib import Path\n",
|
||||
"\n",
|
||||
"print(\"✅ Tools imported successfully!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "efnnREHXUnYA"
|
||||
},
|
||||
"source": [
|
||||
"# YAML Prompt"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "kcM14-yQUqG-",
|
||||
"outputId": "c5b6288d-a789-4a93-e20c-1d285370b667"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"✅ Agent instructions configured!\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Agent configuration and instructions\n",
|
||||
"AGENT_INSTRUCTIONS = \"\"\"\n",
|
||||
"You are a helpful AI assistant specialized in Thai recipes and cooking.\n",
|
||||
"\n",
|
||||
"You have access to a PDF knowledge base about Thai recipes from: https://phi-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf\n",
|
||||
"\n",
|
||||
"You can also search the web for additional information about Thai cooking, ingredients, and techniques.\n",
|
||||
"\n",
|
||||
"When answering questions:\n",
|
||||
"1. Use your knowledge about Thai cuisine to provide helpful information\n",
|
||||
"2. If needed, search the web for additional details, current information, or clarification\n",
|
||||
"3. Provide comprehensive, helpful answers about Thai cuisine\n",
|
||||
"4. Always be informative and helpful about Thai cooking!\n",
|
||||
"\n",
|
||||
"You can use the internet_search function to search the web when needed.\n",
|
||||
"\n",
|
||||
"Focus on providing practical cooking advice, ingredient information, and authentic Thai recipe guidance.\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"print(\"✅ Agent instructions configured!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "MCBRcXFJUulY"
|
||||
},
|
||||
"source": [
|
||||
"# Main"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "YaRLbgMmUvXv",
|
||||
"outputId": "8ed8d9f8-85a7-48de-ccc6-f617cd672a85"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"🤖 Thai Recipe RAG Agent created successfully!\n",
|
||||
"The agent is ready to answer questions about Thai cooking.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Create the RAG agent\n",
|
||||
"rag_agent = Agent(\n",
|
||||
" instructions=AGENT_INSTRUCTIONS,\n",
|
||||
" llm=\"gpt-4o\",\n",
|
||||
" markdown=True,\n",
|
||||
" verbose=True\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(\"🤖 Thai Recipe RAG Agent created successfully!\")\n",
|
||||
"print(\"The agent is ready to answer questions about Thai cooking.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QvZM94pGU0IX"
|
||||
},
|
||||
"source": [
|
||||
"# Test the Agent"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 1000,
|
||||
"referenced_widgets": [
|
||||
"47739975c9f04ee5b4e69e8f445cbeda",
|
||||
"08888d583cba4c01b3fdf8f30b06ca4f"
|
||||
]
|
||||
},
|
||||
"id": "QPdrleKxU1VN",
|
||||
"outputId": "a71a86cd-eff8-4309-9b50-173a30066c90"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"🤔 Question: What are the essential ingredients for Thai cooking?\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╭─ </span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8; font-weight: bold\">Agent Info</span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\"> ────────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #ff9b9b; text-decoration-color: #ff9b9b; font-weight: bold\">👤 Agent:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Agent</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #b4b4b3; text-decoration-color: #b4b4b3; font-weight: bold\">Role:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Assistant</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[38;2;210;227;200m╭─\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[1;38;2;210;227;200mAgent Info\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[38;2;210;227;200m───────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\u001b[38;2;210;227;200m─╮\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;255;155;155m👤 Agent:\u001b[0m \u001b[38;2;255;229;229mAgent\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;180;180;179mRole:\u001b[0m \u001b[38;2;255;229;229mAssistant\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "47739975c9f04ee5b4e69e8f445cbeda",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Output()"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
|
||||
],
|
||||
"text/plain": []
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Response generated in 8.4s</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[2mResponse generated in 8.4s\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭───────────────────────────────────────────────────── Task ──────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> What are the essential ingredients for Thai cooking? <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[36m╭─\u001b[0m\u001b[36m────────────────────────────────────────────────────\u001b[0m\u001b[36m Task \u001b[0m\u001b[36m─────────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m What are the essential ingredients for Thai cooking? \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭─────────────────────────────────────────────────── Response ────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> Thai cooking is known for its balance of five fundamental taste senses in each dish or the overall meal: sour, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> sweet, salty, bitter, and spicy. Essential ingredients that help achieve this balance include: <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 1 </span><span style=\"font-weight: bold\">Fish Sauce (Nam Pla)</span>: A salty, savory sauce made from fermented fish. It's a cornerstone of Thai cooking, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>used to add depth of flavor. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 2 </span><span style=\"font-weight: bold\">Lime Juice</span>: Used to add a fresh, tangy citrus flavor to dishes, especially in salads and some curries. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 3 </span><span style=\"font-weight: bold\">Chilies</span>: Thai cuisine often features fresh Thai bird's eye chilies for heat, as well as dried chilies for <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>certain recipes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 4 </span><span style=\"font-weight: bold\">Lemongrass</span>: A citrusy, fragrant herb that is often used in soups, curries, and teas. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 5 </span><span style=\"font-weight: bold\">Galangal</span>: Similar to ginger, galangal has a sharp, peppery flavor and is used in soups and curries. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 6 </span><span style=\"font-weight: bold\">Kaffir Lime Leaves</span>: They add a distinct, zesty citrus aroma and flavor to Thai dishes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 7 </span><span style=\"font-weight: bold\">Coconut Milk</span>: Adds richness and creaminess, commonly used in curries and desserts. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 8 </span><span style=\"font-weight: bold\">Palm Sugar</span>: Used as a sweetener, palm sugar has a nuanced flavor, unlike regular sugar. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 9 </span><span style=\"font-weight: bold\">Tamarind Paste</span>: Provides a sour element to many dishes, particularly pad Thai and some sauces. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 10 </span><span style=\"font-weight: bold\">Thai Basil</span>: Adds a spicy, peppery flavor, different from the sweet basil used in Western cooking. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 11 </span><span style=\"font-weight: bold\">Coriander (Cilantro)</span>: Used as a garnish and for its fresh, citrusy leaves and roots, which add depth to <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>sauces and stocks. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 12 </span><span style=\"font-weight: bold\">Shrimp Paste</span>: A very pungent paste made from fermented shrimp, used sparingly to enrich the flavor of <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>dishes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> These ingredients are commonly used in combination to create the complex and rich flavors characteristic of <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> Thai cuisine. Fresh and high-quality versions of these ingredients contribute significantly to the authenticity <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> of Thai dishes. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[36m╭─\u001b[0m\u001b[36m──────────────────────────────────────────────────\u001b[0m\u001b[36m Response \u001b[0m\u001b[36m───────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m Thai cooking is known for its balance of five fundamental taste senses in each dish or the overall meal: sour, \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m sweet, salty, bitter, and spicy. Essential ingredients that help achieve this balance include: \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 1 \u001b[0m\u001b[1mFish Sauce (Nam Pla)\u001b[0m: A salty, savory sauce made from fermented fish. It's a cornerstone of Thai cooking, \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mused to add depth of flavor. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 2 \u001b[0m\u001b[1mLime Juice\u001b[0m: Used to add a fresh, tangy citrus flavor to dishes, especially in salads and some curries. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 3 \u001b[0m\u001b[1mChilies\u001b[0m: Thai cuisine often features fresh Thai bird's eye chilies for heat, as well as dried chilies for \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mcertain recipes. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 4 \u001b[0m\u001b[1mLemongrass\u001b[0m: A citrusy, fragrant herb that is often used in soups, curries, and teas. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 5 \u001b[0m\u001b[1mGalangal\u001b[0m: Similar to ginger, galangal has a sharp, peppery flavor and is used in soups and curries. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 6 \u001b[0m\u001b[1mKaffir Lime Leaves\u001b[0m: They add a distinct, zesty citrus aroma and flavor to Thai dishes. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 7 \u001b[0m\u001b[1mCoconut Milk\u001b[0m: Adds richness and creaminess, commonly used in curries and desserts. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 8 \u001b[0m\u001b[1mPalm Sugar\u001b[0m: Used as a sweetener, palm sugar has a nuanced flavor, unlike regular sugar. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 9 \u001b[0m\u001b[1mTamarind Paste\u001b[0m: Provides a sour element to many dishes, particularly pad Thai and some sauces. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 10 \u001b[0m\u001b[1mThai Basil\u001b[0m: Adds a spicy, peppery flavor, different from the sweet basil used in Western cooking. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 11 \u001b[0m\u001b[1mCoriander (Cilantro)\u001b[0m: Used as a garnish and for its fresh, citrusy leaves and roots, which add depth to \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0msauces and stocks. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m 12 \u001b[0m\u001b[1mShrimp Paste\u001b[0m: A very pungent paste made from fermented shrimp, used sparingly to enrich the flavor of \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mdishes. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m These ingredients are commonly used in combination to create the complex and rich flavors characteristic of \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m Thai cuisine. Fresh and high-quality versions of these ingredients contribute significantly to the authenticity \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m of Thai dishes. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"�� Answer: Thai cooking is known for its balance of five fundamental taste senses in each dish or the overall meal: sour, sweet, salty, bitter, and spicy. Essential ingredients that help achieve this balance include:\n",
|
||||
"\n",
|
||||
"1. **Fish Sauce (Nam Pla)**: A salty, savory sauce made from fermented fish. It's a cornerstone of Thai cooking, used to add depth of flavor.\n",
|
||||
"\n",
|
||||
"2. **Lime Juice**: Used to add a fresh, tangy citrus flavor to dishes, especially in salads and some curries.\n",
|
||||
"\n",
|
||||
"3. **Chilies**: Thai cuisine often features fresh Thai bird's eye chilies for heat, as well as dried chilies for certain recipes.\n",
|
||||
"\n",
|
||||
"4. **Lemongrass**: A citrusy, fragrant herb that is often used in soups, curries, and teas.\n",
|
||||
"\n",
|
||||
"5. **Galangal**: Similar to ginger, galangal has a sharp, peppery flavor and is used in soups and curries.\n",
|
||||
"\n",
|
||||
"6. **Kaffir Lime Leaves**: They add a distinct, zesty citrus aroma and flavor to Thai dishes.\n",
|
||||
"\n",
|
||||
"7. **Coconut Milk**: Adds richness and creaminess, commonly used in curries and desserts.\n",
|
||||
"\n",
|
||||
"8. **Palm Sugar**: Used as a sweetener, palm sugar has a nuanced flavor, unlike regular sugar.\n",
|
||||
"\n",
|
||||
"9. **Tamarind Paste**: Provides a sour element to many dishes, particularly pad Thai and some sauces.\n",
|
||||
"\n",
|
||||
"10. **Thai Basil**: Adds a spicy, peppery flavor, different from the sweet basil used in Western cooking.\n",
|
||||
"\n",
|
||||
"11. **Coriander (Cilantro)**: Used as a garnish and for its fresh, citrusy leaves and roots, which add depth to sauces and stocks.\n",
|
||||
"\n",
|
||||
"12. **Shrimp Paste**: A very pungent paste made from fermented shrimp, used sparingly to enrich the flavor of dishes.\n",
|
||||
"\n",
|
||||
"These ingredients are commonly used in combination to create the complex and rich flavors characteristic of Thai cuisine. Fresh and high-quality versions of these ingredients contribute significantly to the authenticity of Thai dishes.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Test question 1: Basic Thai cooking\n",
|
||||
"question1 = \"What are the essential ingredients for Thai cooking?\"\n",
|
||||
"print(f\"🤔 Question: {question1}\")\n",
|
||||
"\n",
|
||||
"try:\n",
|
||||
" response1 = rag_agent.start(question1)\n",
|
||||
" print(f\"�� Answer: {response1}\")\n",
|
||||
"except Exception as e:\n",
|
||||
" print(f\"❌ Error: {e}\")"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,556 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EA8l4cSVl3Lx"
|
||||
},
|
||||
"source": [
|
||||
"# Agentic RAG with GPT-5 - Colab Notebook"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "I-8nOxJ2l7nO"
|
||||
},
|
||||
"source": [
|
||||
"This Colab notebook sets up a Streamlit app using PraisonAI Agents framework.\n",
|
||||
"It allows you to interact with an AI agent powered by GPT-5 for Q&A.\n",
|
||||
"Currently, the knowledge base functionality is disabled due to compatibility issues."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "g5I1tMkxmDQe"
|
||||
},
|
||||
"source": [
|
||||
"- GPT-5 Integration (via OpenAI API)\n",
|
||||
"- Streamlit Web App\n",
|
||||
"- Agent Framework: PraisonAI Agents\n",
|
||||
"- Markdown-formatted responses"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "yeEno7U0q2U_"
|
||||
},
|
||||
"source": [
|
||||
"[](https://colab.research.google.com/github/Dhivya-Bharathy/PraisonAI/blob/main/examples/cookbooks/Qwen2.5_InstructionAgent.ipynb)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "WV2GnzZCmGBR"
|
||||
},
|
||||
"source": [
|
||||
"# Dependencies"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "20wAadXemJ0k"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install praisonaiagents python-dotenv streamlit"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "l71uCRarndvy"
|
||||
},
|
||||
"source": [
|
||||
"# Set API Key"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"id": "brfc-ZZKnfre"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from dotenv import load_dotenv\n",
|
||||
"\n",
|
||||
"load_dotenv()\n",
|
||||
"\n",
|
||||
"OPENAI_API_KEY = \"Enter your api key here\"\n",
|
||||
"\n",
|
||||
"os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "d1fGzyQCn1aQ"
|
||||
},
|
||||
"source": [
|
||||
"# Tools and Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"id": "GcNPLkZXntIS"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import streamlit as st\n",
|
||||
"from praisonaiagents import Agent\n",
|
||||
"\n",
|
||||
"# YAML Prompt Configuration\n",
|
||||
"agent_config = {\n",
|
||||
" \"name\": \"Knowledge Agent\",\n",
|
||||
" \"instructions\": [\n",
|
||||
" \"You are a helpful AI assistant. Answer questions based on your general knowledge.\",\n",
|
||||
" \"Provide clear, well-structured answers in markdown format.\",\n",
|
||||
" \"Use proper markdown formatting with headers, lists, and emphasis where appropriate.\",\n",
|
||||
" \"Structure your response with clear sections and bullet points when helpful.\",\n",
|
||||
" ],\n",
|
||||
" \"llm\": \"gpt-5-nano\",\n",
|
||||
" \"markdown\": True,\n",
|
||||
" \"verbose\": True\n",
|
||||
"}\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "jkmrh0gIn_4a"
|
||||
},
|
||||
"source": [
|
||||
"# Initialize Agent"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"id": "qaJvf8RyoAsB"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def create_agent():\n",
|
||||
" \"\"\"Create an agent with reasoning capabilities\"\"\"\n",
|
||||
" return Agent(\n",
|
||||
" name=agent_config[\"name\"],\n",
|
||||
" instructions=agent_config[\"instructions\"],\n",
|
||||
" llm=agent_config[\"llm\"],\n",
|
||||
" markdown=agent_config[\"markdown\"],\n",
|
||||
" verbose=agent_config[\"verbose\"]\n",
|
||||
" )"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "kvGJh3bbqTmh"
|
||||
},
|
||||
"source": [
|
||||
"# Main Function"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 1000,
|
||||
"referenced_widgets": [
|
||||
"ad296d97510246de8eb0138ce8a8a166",
|
||||
"f6e502b3dca442f0ae6c5739f5a6d159"
|
||||
]
|
||||
},
|
||||
"id": "_QyxzqKKqUrW",
|
||||
"outputId": "97be7e02-b034-4004-e9e3-56337e3f2887"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"�� Initializing Agentic RAG with GPT-5...\n",
|
||||
"📚 Loading agent...\n",
|
||||
"✅ Agent loaded successfully!\n",
|
||||
"\n",
|
||||
"🎯 Example queries you can try:\n",
|
||||
"1. What is PraisonAI and how do Agents work?\n",
|
||||
"2. What are Teams in PraisonAI and how do they work?\n",
|
||||
"3. Give me a step-by-step guide to building a RAG system.\n",
|
||||
"4. What are AI Agents?\n",
|
||||
"\n",
|
||||
"============================================================\n",
|
||||
"\n",
|
||||
"🤔 Enter your question (or 'quit' to exit):\n",
|
||||
"> What is PraisonAI and how do Agents work?\n",
|
||||
"🔍 Processing your question...\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╭─ </span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8; font-weight: bold\">Agent Info</span><span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\"> ────────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #ff9b9b; text-decoration-color: #ff9b9b; font-weight: bold\">👤 Agent:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Knowledge Agent</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #b4b4b3; text-decoration-color: #b4b4b3; font-weight: bold\">Role:</span> <span style=\"color: #ffe5e5; text-decoration-color: #ffe5e5\">Assistant</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span> <span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">│</span>\n",
|
||||
"<span style=\"color: #d2e3c8; text-decoration-color: #d2e3c8\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[38;2;210;227;200m╭─\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[1;38;2;210;227;200mAgent Info\u001b[0m\u001b[38;2;210;227;200m \u001b[0m\u001b[38;2;210;227;200m───────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\u001b[38;2;210;227;200m─╮\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;255;155;155m👤 Agent:\u001b[0m \u001b[38;2;255;229;229mKnowledge Agent\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[1;38;2;180;180;179mRole:\u001b[0m \u001b[38;2;255;229;229mAssistant\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m│\u001b[0m \u001b[38;2;210;227;200m│\u001b[0m\n",
|
||||
"\u001b[38;2;210;227;200m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"application/vnd.jupyter.widget-view+json": {
|
||||
"model_id": "ad296d97510246de8eb0138ce8a8a166",
|
||||
"version_major": 2,
|
||||
"version_minor": 0
|
||||
},
|
||||
"text/plain": [
|
||||
"Output()"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
|
||||
],
|
||||
"text/plain": []
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Response generated in 9.4s</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[2mResponse generated in 9.4s\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭───────────────────────────────────────────────────── Task ──────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> What is PraisonAI and how do Agents work? <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[36m╭─\u001b[0m\u001b[36m────────────────────────────────────────────────────\u001b[0m\u001b[36m Task \u001b[0m\u001b[36m─────────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m What is PraisonAI and how do Agents work? \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">╭─────────────────────────────────────────────────── Response ────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┃ <span style=\"font-weight: bold\">What is PraisonAI?</span> ┃ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> PraisonAI is a platform designed to leverage artificial intelligence to enhance personal and business <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> productivity. It typically involves developing AI models that can assist in various tasks, automate processes, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> and provide insights based on data analysis. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">Key Features of PraisonAI</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Task Automation</span>: Automates repetitive tasks, allowing users to focus on more strategic activities. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Data Insights</span>: Analyzes large sets of data to provide actionable insights. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Natural Language Processing</span>: Enables users to interact with the AI through natural language, making it <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>user-friendly. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Customizable Solutions</span>: Tailors AI applications to meet specific user needs across various industries. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┃ <span style=\"font-weight: bold\">How Do Agents Work in PraisonAI?</span> ┃ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> Agents in PraisonAI refer to autonomous or semi-autonomous AI systems that perform designated tasks or <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> functions. Here’s how they typically operate: <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">1. Types of Agents</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Reactive Agents</span>: Respond to the current environment without internal memory. They act solely based on <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>current inputs. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Deliberative Agents</span>: Have a model of their environment, plan their actions, and retain memories of past <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>interactions. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Learning Agents</span>: Improve their performance over time through experience, adapting to new information. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">2. Functionality of Agents</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Task Execution</span>: Agents can carry out specific tasks as programmed, such as scheduling, data entry, or <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>customer support. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Interaction</span>: They often interact with users through chat interfaces or APIs, providing a seamless user <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>experience. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Integration</span>: Can be integrated with other systems or software to enhance functionality, such as CRM systems, <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>databases, or other AI tools. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Continuous Learning</span>: Agents may utilize machine learning algorithms to adapt and evolve based on user <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> </span>interactions and data received. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">3. Benefits of Using Agents</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Increased Efficiency</span>: Automate routine tasks to reduce workload. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">24/7 Availability</span>: Provide consistent support and functionality without downtime. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Scalability</span>: Easily scale operations by deploying more agents as needed. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span><span style=\"font-weight: bold\">Data-Driven Decisions</span>: Use data insights to inform strategy and operations. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"font-weight: bold; text-decoration: underline\">Conclusion</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> PraisonAI represents a robust platform for enhancing productivity through AI, with agents serving as the heart <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> of this solution, performing various tasks and supporting users in achieving their goals efficiently. Whether <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> they are simple reactive agents or complex learning entities, their design focuses on making operations <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">│</span> smoother and more effective. <span style=\"color: #008080; text-decoration-color: #008080\">│</span>\n",
|
||||
"<span style=\"color: #008080; text-decoration-color: #008080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[36m╭─\u001b[0m\u001b[36m──────────────────────────────────────────────────\u001b[0m\u001b[36m Response \u001b[0m\u001b[36m───────────────────────────────────────────────────\u001b[0m\u001b[36m─╮\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m ┃ \u001b[1mWhat is PraisonAI?\u001b[0m ┃ \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m PraisonAI is a platform designed to leverage artificial intelligence to enhance personal and business \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m productivity. It typically involves developing AI models that can assist in various tasks, automate processes, \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m and provide insights based on data analysis. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;4mKey Features of PraisonAI\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mTask Automation\u001b[0m: Automates repetitive tasks, allowing users to focus on more strategic activities. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mData Insights\u001b[0m: Analyzes large sets of data to provide actionable insights. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mNatural Language Processing\u001b[0m: Enables users to interact with the AI through natural language, making it \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0muser-friendly. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mCustomizable Solutions\u001b[0m: Tailors AI applications to meet specific user needs across various industries. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m ┃ \u001b[1mHow Do Agents Work in PraisonAI?\u001b[0m ┃ \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m Agents in PraisonAI refer to autonomous or semi-autonomous AI systems that perform designated tasks or \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m functions. Here’s how they typically operate: \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;4m1. \u001b[0m\u001b[1;4mTypes of Agents\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mReactive Agents\u001b[0m: Respond to the current environment without internal memory. They act solely based on \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mcurrent inputs. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mDeliberative Agents\u001b[0m: Have a model of their environment, plan their actions, and retain memories of past \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0minteractions. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mLearning Agents\u001b[0m: Improve their performance over time through experience, adapting to new information. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;4m2. \u001b[0m\u001b[1;4mFunctionality of Agents\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mTask Execution\u001b[0m: Agents can carry out specific tasks as programmed, such as scheduling, data entry, or \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mcustomer support. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mInteraction\u001b[0m: They often interact with users through chat interfaces or APIs, providing a seamless user \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mexperience. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mIntegration\u001b[0m: Can be integrated with other systems or software to enhance functionality, such as CRM systems, \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0mdatabases, or other AI tools. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mContinuous Learning\u001b[0m: Agents may utilize machine learning algorithms to adapt and evolve based on user \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m \u001b[0minteractions and data received. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;4m3. \u001b[0m\u001b[1;4mBenefits of Using Agents\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mIncreased Efficiency\u001b[0m: Automate routine tasks to reduce workload. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1m24/7 Availability\u001b[0m: Provide consistent support and functionality without downtime. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mScalability\u001b[0m: Easily scale operations by deploying more agents as needed. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;33m • \u001b[0m\u001b[1mData-Driven Decisions\u001b[0m: Use data insights to inform strategy and operations. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[1;4mConclusion\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m PraisonAI represents a robust platform for enhancing productivity through AI, with agents serving as the heart \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m of this solution, performing various tasks and supporting users in achieving their goals efficiently. Whether \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m they are simple reactive agents or complex learning entities, their design focuses on making operations \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m│\u001b[0m smoother and more effective. \u001b[36m│\u001b[0m\n",
|
||||
"\u001b[36m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\n",
|
||||
"💡 Answer:\n",
|
||||
"========================================\n",
|
||||
"# What is PraisonAI?\n",
|
||||
"\n",
|
||||
"PraisonAI is a platform designed to leverage artificial intelligence to enhance personal and business productivity. It typically involves developing AI models that can assist in various tasks, automate processes, and provide insights based on data analysis.\n",
|
||||
"\n",
|
||||
"## Key Features of PraisonAI\n",
|
||||
"\n",
|
||||
"- **Task Automation**: Automates repetitive tasks, allowing users to focus on more strategic activities.\n",
|
||||
"- **Data Insights**: Analyzes large sets of data to provide actionable insights.\n",
|
||||
"- **Natural Language Processing**: Enables users to interact with the AI through natural language, making it user-friendly.\n",
|
||||
"- **Customizable Solutions**: Tailors AI applications to meet specific user needs across various industries.\n",
|
||||
"\n",
|
||||
"# How Do Agents Work in PraisonAI?\n",
|
||||
"\n",
|
||||
"Agents in PraisonAI refer to autonomous or semi-autonomous AI systems that perform designated tasks or functions. Here’s how they typically operate:\n",
|
||||
"\n",
|
||||
"## 1. **Types of Agents**\n",
|
||||
"\n",
|
||||
"- **Reactive Agents**: Respond to the current environment without internal memory. They act solely based on current inputs.\n",
|
||||
"- **Deliberative Agents**: Have a model of their environment, plan their actions, and retain memories of past interactions.\n",
|
||||
"- **Learning Agents**: Improve their performance over time through experience, adapting to new information.\n",
|
||||
"\n",
|
||||
"## 2. **Functionality of Agents**\n",
|
||||
"\n",
|
||||
"- **Task Execution**: Agents can carry out specific tasks as programmed, such as scheduling, data entry, or customer support.\n",
|
||||
"- **Interaction**: They often interact with users through chat interfaces or APIs, providing a seamless user experience.\n",
|
||||
"- **Integration**: Can be integrated with other systems or software to enhance functionality, such as CRM systems, databases, or other AI tools.\n",
|
||||
"- **Continuous Learning**: Agents may utilize machine learning algorithms to adapt and evolve based on user interactions and data received.\n",
|
||||
"\n",
|
||||
"## 3. **Benefits of Using Agents**\n",
|
||||
"\n",
|
||||
"- **Increased Efficiency**: Automate routine tasks to reduce workload.\n",
|
||||
"- **24/7 Availability**: Provide consistent support and functionality without downtime.\n",
|
||||
"- **Scalability**: Easily scale operations by deploying more agents as needed.\n",
|
||||
"- **Data-Driven Decisions**: Use data insights to inform strategy and operations.\n",
|
||||
"\n",
|
||||
"## Conclusion\n",
|
||||
"\n",
|
||||
"PraisonAI represents a robust platform for enhancing productivity through AI, with agents serving as the heart of this solution, performing various tasks and supporting users in achieving their goals efficiently. Whether they are simple reactive agents or complex learning entities, their design focuses on making operations smoother and more effective.\n",
|
||||
"========================================\n",
|
||||
"\n",
|
||||
"🤔 Enter your question (or 'quit' to exit):\n",
|
||||
"> quit\n",
|
||||
"👋 Goodbye!\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"def main():\n",
|
||||
" print(\"�� Initializing Agentic RAG with GPT-5...\")\n",
|
||||
"\n",
|
||||
" # Check if API key is set\n",
|
||||
" if not OPENAI_API_KEY or OPENAI_API_KEY == \"your_openai_api_key_here\":\n",
|
||||
" print(\"❌ Please set your OpenAI API key first!\")\n",
|
||||
" print(\"�� Get your key from: https://platform.openai.com/\")\n",
|
||||
" return\n",
|
||||
"\n",
|
||||
" try:\n",
|
||||
" # Create agent\n",
|
||||
" print(\"📚 Loading agent...\")\n",
|
||||
" agent = create_agent()\n",
|
||||
" print(\"✅ Agent loaded successfully!\")\n",
|
||||
"\n",
|
||||
" # Example queries\n",
|
||||
" example_queries = [\n",
|
||||
" \"What is PraisonAI and how do Agents work?\",\n",
|
||||
" \"What are Teams in PraisonAI and how do they work?\",\n",
|
||||
" \"Give me a step-by-step guide to building a RAG system.\",\n",
|
||||
" \"What are AI Agents?\"\n",
|
||||
" ]\n",
|
||||
"\n",
|
||||
" print(\"\\n🎯 Example queries you can try:\")\n",
|
||||
" for i, query in enumerate(example_queries, 1):\n",
|
||||
" print(f\"{i}. {query}\")\n",
|
||||
"\n",
|
||||
" print(\"\\n\" + \"=\"*60)\n",
|
||||
"\n",
|
||||
" # Interactive query loop\n",
|
||||
" while True:\n",
|
||||
" print(\"\\n🤔 Enter your question (or 'quit' to exit):\")\n",
|
||||
" user_query = input(\"> \")\n",
|
||||
"\n",
|
||||
" if user_query.lower() in ['quit', 'exit', 'q']:\n",
|
||||
" print(\"👋 Goodbye!\")\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
" if user_query.strip():\n",
|
||||
" print(\"🔍 Processing your question...\")\n",
|
||||
" try:\n",
|
||||
" response = agent.start(user_query)\n",
|
||||
" print(\"\\n💡 Answer:\")\n",
|
||||
" print(\"=\"*40)\n",
|
||||
" print(response)\n",
|
||||
" print(\"=\"*40)\n",
|
||||
" except Exception as e:\n",
|
||||
" print(f\"❌ Error: {str(e)}\")\n",
|
||||
" print(f\"🔍 Error type: {type(e).__name__}\")\n",
|
||||
" print(\"💡 This might be due to API issues or model access problems.\")\n",
|
||||
" else:\n",
|
||||
" print(\"❌ Please enter a valid question.\")\n",
|
||||
"\n",
|
||||
" except Exception as e:\n",
|
||||
" print(f\"❌ Failed to initialize agent: {str(e)}\")\n",
|
||||
" print(\"�� Make sure your API key is correct and you have sufficient credits.\")\n",
|
||||
"\n",
|
||||
"# Run the main function\n",
|
||||
"if __name__ == \"__main__\":\n",
|
||||
" main()"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
MCP Backward Compatibility Example
|
||||
|
||||
Demonstrates backward compatibility with older MCP servers
|
||||
using the deprecated HTTP+SSE transport (protocol version 2024-11-05).
|
||||
|
||||
Features:
|
||||
- Legacy SSE endpoint detection (/sse suffix)
|
||||
- Automatic transport negotiation
|
||||
- Protocol version support (2024-11-05 to 2025-11-25)
|
||||
|
||||
Protocol: MCP 2025-11-25 with 2024-11-05 compatibility
|
||||
"""
|
||||
|
||||
# Example code (requires running MCP servers):
|
||||
#
|
||||
# from praisonaiagents import Agent, MCP
|
||||
#
|
||||
# # Legacy SSE server (auto-detected via /sse suffix)
|
||||
# agent_legacy = Agent(
|
||||
# name="Legacy SSE Assistant",
|
||||
# tools=MCP("http://localhost:8080/sse")
|
||||
# )
|
||||
#
|
||||
# # Modern Streamable HTTP server
|
||||
# agent_modern = Agent(
|
||||
# name="Modern HTTP Assistant",
|
||||
# tools=MCP("http://localhost:8080/mcp")
|
||||
# )
|
||||
|
||||
if __name__ == "__main__":
|
||||
from praisonaiagents.mcp.mcp_transport import get_transport_type
|
||||
from praisonaiagents.mcp.mcp_compat import (
|
||||
detect_transport_support,
|
||||
is_legacy_sse_url,
|
||||
TransportNegotiator
|
||||
)
|
||||
from praisonaiagents.mcp.mcp_session import VALID_PROTOCOL_VERSIONS
|
||||
|
||||
print("MCP Backward Compatibility Example")
|
||||
print("=" * 50)
|
||||
|
||||
# URL detection
|
||||
print("\n1. URL-based Transport Detection:")
|
||||
urls = [
|
||||
"http://localhost:8080/sse",
|
||||
"https://api.example.com/sse",
|
||||
"http://localhost:8080/mcp",
|
||||
"https://api.example.com/mcp",
|
||||
]
|
||||
for url in urls:
|
||||
transport = get_transport_type(url)
|
||||
legacy = is_legacy_sse_url(url)
|
||||
print(f" {url}")
|
||||
print(f" Transport: {transport}, Legacy: {legacy}")
|
||||
|
||||
# Transport negotiation
|
||||
print("\n2. Transport Negotiation (HTTP response based):")
|
||||
responses = [
|
||||
(200, "application/json", "Streamable HTTP"),
|
||||
(200, "text/event-stream", "SSE stream"),
|
||||
(400, None, "Fallback to SSE"),
|
||||
(404, None, "Fallback to SSE"),
|
||||
(405, None, "Fallback to SSE"),
|
||||
]
|
||||
for status, content_type, desc in responses:
|
||||
detected = detect_transport_support(status, content_type)
|
||||
print(f" HTTP {status} ({desc}): {detected}")
|
||||
|
||||
# Protocol versions
|
||||
print("\n3. Supported Protocol Versions:")
|
||||
for version in sorted(VALID_PROTOCOL_VERSIONS):
|
||||
print(f" - {version}")
|
||||
|
||||
# Negotiator
|
||||
print("\n4. Transport Negotiator:")
|
||||
negotiator = TransportNegotiator("https://api.example.com/mcp")
|
||||
print(f" Base URL: {negotiator.base_url}")
|
||||
print(f" Try Streamable first: {negotiator.should_try_streamable_first()}")
|
||||
|
||||
negotiator_sse = TransportNegotiator("https://api.example.com/sse")
|
||||
print(f" SSE URL: {negotiator_sse.base_url}")
|
||||
print(f" Try Streamable first: {negotiator_sse.should_try_streamable_first()}")
|
||||
@@ -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,44 @@
|
||||
from praisonaiagents import Agent
|
||||
import os
|
||||
|
||||
# Set OpenAI API key if not already set
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
print("Please set your OpenAI API key: export OPENAI_API_KEY='your-api-key-here'")
|
||||
exit(1)
|
||||
|
||||
# Create the RAG agent with web search capabilities
|
||||
rag_agent = Agent(
|
||||
instructions="""You are a helpful AI assistant specialized in Thai recipes and cooking.
|
||||
|
||||
You have access to a PDF knowledge base about Thai recipes from: https://phi-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf
|
||||
|
||||
You can also search the web for additional information about Thai cooking, ingredients, and techniques.
|
||||
|
||||
When answering questions:
|
||||
1. Use your knowledge about Thai cuisine to provide helpful information
|
||||
2. If needed, search the web for additional details, current information, or clarification
|
||||
3. Provide comprehensive, helpful answers about Thai cuisine
|
||||
4. Always be informative and helpful about Thai cooking!
|
||||
|
||||
You can use the internet_search function to search the web when needed.""",
|
||||
llm="gpt-4o",
|
||||
markdown=True,
|
||||
verbose=True
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("🤖 Thai Recipe RAG Agent is ready!")
|
||||
print("Ask me anything about Thai recipes or cooking!")
|
||||
print("Type 'quit' to exit.\n")
|
||||
|
||||
while True:
|
||||
user_input = input("You: ")
|
||||
if user_input.lower() in ['quit', 'exit', 'bye']:
|
||||
print("👋 Goodbye! Happy cooking!")
|
||||
break
|
||||
|
||||
try:
|
||||
response = rag_agent.start(user_input)
|
||||
print(f"\n🤖 Assistant: {response}\n")
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}\n")
|
||||
@@ -0,0 +1,154 @@
|
||||
# 🧠 Agentic RAG with GPT-5
|
||||
|
||||
An agentic RAG application built with the PraisonAI Agents framework, featuring GPT-5 and built-in vector search for efficient knowledge retrieval and question answering.
|
||||
|
||||
## ✨ Features
|
||||
|
||||
- **🤖 GPT-5**: Latest OpenAI model for intelligent responses
|
||||
- **🗄️ Built-in Vector Search**: ChromaDB integration for fast similarity search
|
||||
- **🔍 Agentic RAG**: Intelligent retrieval augmented generation
|
||||
- **📝 Markdown Formatting**: Beautiful, structured responses
|
||||
- **🌐 Dynamic Knowledge**: Add URLs to expand knowledge base
|
||||
- **⚡ Real-time Responses**: Fast answer generation- **🎯 Clean Interface**: Simplified UI without configuration complexity
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
### Prerequisites
|
||||
|
||||
- Python 3.11+
|
||||
- OpenAI API key with GPT-5 access
|
||||
|
||||
### Installation
|
||||
|
||||
1. **Clone and navigate to the project**
|
||||
```bash
|
||||
cd rag_examples/agentic_rag_gpt5
|
||||
```
|
||||
|
||||
2. **Install dependencies**
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
3. **Set up your OpenAI API key**
|
||||
```bash
|
||||
export OPENAI_API_KEY="your-api-key-here"
|
||||
```
|
||||
Or create a `.env` file:
|
||||
```
|
||||
OPENAI_API_KEY=your-api-key-here
|
||||
```
|
||||
|
||||
4. **Run the application**
|
||||
```bash
|
||||
streamlit run agentic_rag_gpt5.py
|
||||
```
|
||||
|
||||
## 🎯 How to Use
|
||||
|
||||
1. **Enter your OpenAI API key** in the sidebar
|
||||
2. **Add knowledge sources** by entering URLs in the sidebar
|
||||
3. **Ask questions** using the text area or suggested prompts
|
||||
4. **Get answers** with markdown formatting
|
||||
|
||||
### Suggested Questions
|
||||
|
||||
- **"What is PraisonAI?"** - Learn about the PraisonAI Agents framework
|
||||
- **"Teams in PraisonAI"** - Understand how teams work in PraisonAI
|
||||
- **"Build RAG system"** - Get a step-by-step guide to building RAG systems
|
||||
|
||||
## 🏗️ Architecture
|
||||
|
||||
### Core Components
|
||||
|
||||
- **`Agent`**: PraisonAI Agents framework for intelligent Q&A
|
||||
- **`knowledge`**: Built-in knowledge base that handles URLs and documents
|
||||
- **`llm`**: OpenAI GPT-5-nano for generating responses
|
||||
- **Built-in Vector Search**: Automatic similarity search without external setup
|
||||
|
||||
### Data Flow
|
||||
|
||||
1. **Knowledge Loading**: URLs are processed and stored in the built-in vector database
|
||||
2. **Vector Search**: OpenAI embeddings enable semantic search
|
||||
3. **Response Generation**: GPT-5-nano processes information and generates answers
|
||||
4. **Formatted Output**: Markdown-formatted responses
|
||||
|
||||
## 🔧 Configuration
|
||||
|
||||
### Database Settings
|
||||
- **Vector DB**: Built-in vector database with automatic indexing
|
||||
- **Storage**: Local storage managed by PraisonAI Agents
|
||||
- **Search**: Automatic similarity search
|
||||
|
||||
### Model Configuration
|
||||
- **LLM**: OpenAI GPT-5-nano
|
||||
- **Embeddings**: Automatic handling by PraisonAI Agents
|
||||
- **Vector Store**: Built-in with automatic document processing
|
||||
|
||||
## 📚 Knowledge Management
|
||||
|
||||
### Adding Sources
|
||||
- Use the sidebar to add new URLs
|
||||
- Sources are automatically processed and indexed
|
||||
- Current sources are displayed as numbered list
|
||||
|
||||
### Default Knowledge
|
||||
- Starts with PraisonAI documentation: `https://docs.praisonai.com/introduction/agents.md`
|
||||
- Expandable with any web-based documentation
|
||||
|
||||
## 🎨 UI Features
|
||||
|
||||
### Sidebar
|
||||
- **API Key Management**: Secure input for OpenAI credentials
|
||||
- **URL Addition**: Dynamic knowledge base expansion
|
||||
- **Current Sources**: Numbered list of loaded URLs
|
||||
|
||||
### Main Interface
|
||||
- **Suggested Prompts**: Quick access to common questions
|
||||
- **Query Input**: Large text area for custom questions
|
||||
- **Fast Responses**: Quick answer generation
|
||||
- **Markdown Rendering**: Beautiful formatted responses
|
||||
|
||||
## 🛠️ Technical Details
|
||||
|
||||
### Dependencies
|
||||
```
|
||||
streamlit>=1.28.0
|
||||
praisonaiagents>=0.1.0
|
||||
openai>=1.0.0
|
||||
python-dotenv>=1.0.0
|
||||
```
|
||||
|
||||
### Key Features
|
||||
- **Built-in Knowledge Base**: Automatic document processing and indexing
|
||||
- **Vector Search**: Efficient similarity search with built-in database
|
||||
- **Caching**: Efficient resource loading with Streamlit caching
|
||||
- **Error Handling**: Graceful handling of API and processing errors
|
||||
|
||||
## 🔍 Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
**Knowledge base not loading**
|
||||
- Check OpenAI API key is valid
|
||||
- Ensure URLs are accessible
|
||||
- Verify internet connection
|
||||
|
||||
**Agent initialization errors**
|
||||
- Check if PraisonAI Agents is properly installed
|
||||
- Verify OpenAI API key has sufficient credits
|
||||
- Ensure Python version is 3.11+
|
||||
|
||||
### Performance Tips
|
||||
- **Cache Resources**: Knowledge base and agent are cached for efficiency
|
||||
- **Built-in Vector Search**: Fast similarity search without external setup
|
||||
- **Local Storage**: Optimized local storage for optimal performance
|
||||
|
||||
## 🎯 Use Cases
|
||||
|
||||
- **Documentation Q&A**: Ask questions about technical documentation
|
||||
- **Research Assistant**: Get answers from multiple knowledge sources
|
||||
- **Learning Tool**: Interactive exploration of complex topics
|
||||
- **Content Discovery**: Find relevant information across multiple sources
|
||||
|
||||
**Built with ❤️ using PraisonAI Agents and GPT-5**
|
||||
@@ -0,0 +1,201 @@
|
||||
import streamlit as st
|
||||
import os
|
||||
from praisonaiagents import Agent
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Page configuration
|
||||
st.set_page_config(
|
||||
page_title="Agentic RAG with GPT-5",
|
||||
page_icon="🧠",
|
||||
layout="wide"
|
||||
)
|
||||
|
||||
# Main title and description
|
||||
st.title("🧠 Agentic RAG with GPT-5")
|
||||
st.markdown("""
|
||||
This app demonstrates an intelligent AI agent that:
|
||||
1. **Answers** your questions clearly and concisely using GPT-5
|
||||
|
||||
⚠️ **Note**: Knowledge base functionality is temporarily disabled due to a compatibility issue with the current version of PraisonAI Agents.
|
||||
|
||||
Enter your OpenAI API key in the sidebar to get started!
|
||||
""")
|
||||
|
||||
# Sidebar for API key and settings
|
||||
with st.sidebar:
|
||||
st.header("🔧 Configuration")
|
||||
|
||||
# OpenAI API Key
|
||||
openai_key = st.text_input(
|
||||
"OpenAI API Key",
|
||||
type="password",
|
||||
value=os.getenv("OPENAI_API_KEY", ""),
|
||||
help="Get your key from https://platform.openai.com/"
|
||||
)
|
||||
|
||||
# Add URLs to knowledge base
|
||||
st.subheader("🌐 Add Knowledge Sources")
|
||||
st.info("⚠️ Knowledge base functionality is temporarily disabled due to compatibility issues.")
|
||||
new_url = st.text_input(
|
||||
"Add URL",
|
||||
placeholder="https://docs.praisonai.com/introduction",
|
||||
help="Enter a URL to add to the knowledge base (currently disabled)"
|
||||
)
|
||||
|
||||
if st.button("➕ Add URL", type="primary", disabled=True):
|
||||
st.info("Knowledge base functionality is temporarily disabled.")
|
||||
|
||||
# Check if API key is provided
|
||||
if openai_key:
|
||||
# Initialize knowledge base (cached to avoid reloading)
|
||||
@st.cache_resource(show_spinner="📚 Loading knowledge base...")
|
||||
def load_knowledge() -> list:
|
||||
"""Load and initialize the knowledge base with default URL"""
|
||||
return ["https://docs.praisonai.com/introduction/agents.md"] # Default URL
|
||||
|
||||
# Initialize agent (cached to avoid reloading)
|
||||
@st.cache_resource(show_spinner="🤖 Loading agent...")
|
||||
def load_agent(_knowledge: list) -> Agent:
|
||||
"""Create an agent with reasoning capabilities"""
|
||||
# Note: Temporarily removed knowledge parameter to avoid rerank error
|
||||
# TODO: Re-enable when PraisonAI Agents knowledge issue is resolved
|
||||
return Agent(
|
||||
name="Knowledge Agent",
|
||||
instructions=[
|
||||
"You are a helpful AI assistant. Answer questions based on your general knowledge.",
|
||||
"Provide clear, well-structured answers in markdown format.",
|
||||
"Use proper markdown formatting with headers, lists, and emphasis where appropriate.",
|
||||
"Structure your response with clear sections and bullet points when helpful.",
|
||||
],
|
||||
llm="gpt-5-nano
|
||||
markdown=True,
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Load knowledge and agent
|
||||
knowledge = load_knowledge()
|
||||
agent = load_agent(knowledge)
|
||||
|
||||
# Display current URLs in knowledge base
|
||||
if knowledge:
|
||||
st.sidebar.subheader("📚 Current Knowledge Sources")
|
||||
for i, url in enumerate(knowledge, 1):
|
||||
st.sidebar.markdown(f"{i}. {url}")
|
||||
|
||||
# Handle URL additions
|
||||
if hasattr(st.session_state, 'urls_to_add') and st.session_state.urls_to_add:
|
||||
with st.spinner("📥 Loading new documents..."):
|
||||
knowledge.append(st.session_state.urls_to_add)
|
||||
# Reinitialize agent with new knowledge
|
||||
agent = load_agent(knowledge)
|
||||
st.success(f"✅ Added: {st.session_state.urls_to_add}")
|
||||
del st.session_state.urls_to_add
|
||||
st.rerun()
|
||||
|
||||
# Main query section
|
||||
st.divider()
|
||||
st.subheader("🤔 Ask a Question")
|
||||
|
||||
# Suggested prompts
|
||||
st.markdown("**Try these prompts:**")
|
||||
col1, col2, col3 = st.columns(3)
|
||||
with col1:
|
||||
if st.button("What is PraisonAI?", use_container_width=True):
|
||||
st.session_state.query = "What is PraisonAI and how do Agents work?"
|
||||
with col2:
|
||||
if st.button("Teams in PraisonAI", use_container_width=True):
|
||||
st.session_state.query = "What are Teams in PraisonAI and how do they work?"
|
||||
with col3:
|
||||
if st.button("Build RAG system", use_container_width=True):
|
||||
st.session_state.query = "Give me a step-by-step guide to building a RAG system."
|
||||
|
||||
# Query input
|
||||
query = st.text_area(
|
||||
"Your question:",
|
||||
value=st.session_state.get("query", "What are AI Agents?"),
|
||||
height=100,
|
||||
help="Ask anything about the loaded knowledge sources"
|
||||
)
|
||||
|
||||
# Run button
|
||||
if st.button("🚀 Get Answer", type="primary"):
|
||||
if query:
|
||||
# Create container for answer
|
||||
st.markdown("### 💡 Answer")
|
||||
answer_container = st.container()
|
||||
answer_placeholder = answer_container.empty()
|
||||
|
||||
# Get the agent's response
|
||||
with st.spinner("🔍 Searching and generating answer..."):
|
||||
try:
|
||||
st.info("🤖 Agent is processing your question...")
|
||||
response = agent.start(query)
|
||||
answer_placeholder.markdown(
|
||||
response,
|
||||
unsafe_allow_html=True
|
||||
)
|
||||
except Exception as e:
|
||||
st.error(f"Error getting response: {str(e)}")
|
||||
st.error(f"Error type: {type(e).__name__}")
|
||||
st.error(f"Full error details: {repr(e)}")
|
||||
# Try to provide a helpful response
|
||||
answer_placeholder.markdown("""
|
||||
**⚠️ Error occurred while processing your question.**
|
||||
|
||||
This might be due to:
|
||||
- Knowledge base configuration issues
|
||||
- Model access problems
|
||||
- Network connectivity issues
|
||||
|
||||
Please try again or check your OpenAI API key.
|
||||
""")
|
||||
else:
|
||||
st.error("Please enter a question")
|
||||
|
||||
else:
|
||||
# Show instructions if API key is missing
|
||||
st.info("""
|
||||
👋 **Welcome! To use this app, you need:**
|
||||
|
||||
- **OpenAI API Key** (set it in the sidebar)
|
||||
- Sign up at [platform.openai.com](https://platform.openai.com/)
|
||||
- Generate a new API key
|
||||
|
||||
Once you enter the key, the app will load the knowledge base and agent.
|
||||
""")
|
||||
|
||||
# Footer with explanation
|
||||
st.divider()
|
||||
with st.expander("📖 How This Works"):
|
||||
st.markdown("""
|
||||
**This app uses the PraisonAI Agents framework to create an intelligent Q&A system:**
|
||||
|
||||
⚠️ **Current Status**: Knowledge base functionality is temporarily disabled due to compatibility issues with the current version of PraisonAI Agents (v0.0.157).
|
||||
|
||||
**What Works Now:**
|
||||
1. **GPT-5 Integration**: OpenAI's GPT-5 model for generating intelligent responses
|
||||
2. **Markdown Formatting**: Beautiful, structured responses
|
||||
3. **Agent Framework**: PraisonAI Agents for orchestration
|
||||
|
||||
**What's Temporarily Disabled:**
|
||||
- URL-based knowledge base
|
||||
- Vector search and retrieval
|
||||
- Document processing
|
||||
|
||||
**Key Components:**
|
||||
- `Agent`: PraisonAI Agents framework for creating intelligent agents
|
||||
- `llm`: OpenAI GPT-5-nano for generating responses
|
||||
|
||||
**Why PraisonAI Agents?**
|
||||
- Easy-to-use agent framework
|
||||
- Built-in error handling
|
||||
- Clean API design
|
||||
- Perfect for prototyping and production applications
|
||||
|
||||
**Next Steps:**
|
||||
- Monitor PraisonAI Agents updates for knowledge base fixes
|
||||
- Re-enable knowledge functionality when compatible version is available
|
||||
""")
|
||||
@@ -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,55 @@
|
||||
"""
|
||||
Workflow Loop with CSV Example
|
||||
|
||||
Demonstrates iterating over a CSV file, processing each row.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, WorkflowContext, StepResult
|
||||
from praisonaiagents.memory.workflows import loop
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
# Create sample CSV file
|
||||
csv_content = """name,email,task
|
||||
Alice,[email protected],Write documentation
|
||||
Bob,[email protected],Review pull requests
|
||||
Charlie,[email protected],Fix critical bug
|
||||
Diana,[email protected],Deploy to production"""
|
||||
|
||||
# Processor for each CSV row
|
||||
def process_task(ctx: WorkflowContext) -> StepResult:
|
||||
"""Process a single task from CSV."""
|
||||
item = ctx.variables.get("item", {})
|
||||
name = item.get("name", "Unknown")
|
||||
task = item.get("task", "No task")
|
||||
index = ctx.variables.get("loop_index", 0)
|
||||
|
||||
return StepResult(
|
||||
output=f"[{index+1}] {name}: {task} ✅",
|
||||
variables={"last_processed": name}
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create temp CSV file
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f:
|
||||
f.write(csv_content)
|
||||
csv_path = f.name
|
||||
|
||||
try:
|
||||
print("=== Testing Loop with CSV ===\n")
|
||||
print(f"CSV file: {csv_path}\n")
|
||||
|
||||
# Create workflow
|
||||
workflow = Workflow(
|
||||
name="CSV Task Processor",
|
||||
steps=[loop(process_task, from_csv=csv_path)]
|
||||
)
|
||||
|
||||
result = workflow.start("Process all tasks", verbose=True)
|
||||
|
||||
print(f"\nAll outputs:\n{result['output']}")
|
||||
print(f"\nProcessed {len(result['variables']['loop_outputs'])} items")
|
||||
print(f"Last processed: {result['variables'].get('last_processed')}")
|
||||
|
||||
finally:
|
||||
os.unlink(csv_path)
|
||||
@@ -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,53 @@
|
||||
"""
|
||||
Workflow Parallel Execution Example
|
||||
|
||||
Demonstrates running multiple steps concurrently and
|
||||
combining their results.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, WorkflowContext, StepResult
|
||||
from praisonaiagents.memory.workflows import parallel
|
||||
import time
|
||||
|
||||
# Parallel workers - each does independent work
|
||||
def research_market(ctx: WorkflowContext) -> StepResult:
|
||||
"""Research market trends."""
|
||||
time.sleep(0.1) # Simulate work
|
||||
return StepResult(output="📊 Market: Growth expected 15% YoY")
|
||||
|
||||
def research_competitors(ctx: WorkflowContext) -> StepResult:
|
||||
"""Research competitor analysis."""
|
||||
time.sleep(0.1) # Simulate work
|
||||
return StepResult(output="🏢 Competitors: 3 major players identified")
|
||||
|
||||
def research_customers(ctx: WorkflowContext) -> StepResult:
|
||||
"""Research customer feedback."""
|
||||
time.sleep(0.1) # Simulate work
|
||||
return StepResult(output="👥 Customers: 85% satisfaction rate")
|
||||
|
||||
# Aggregator - combines parallel results
|
||||
def summarize_research(ctx: WorkflowContext) -> StepResult:
|
||||
"""Summarize all research findings."""
|
||||
outputs = ctx.variables.get("parallel_outputs", [])
|
||||
summary = "📋 RESEARCH SUMMARY:\n" + "\n".join(f" • {o}" for o in outputs)
|
||||
return StepResult(output=summary)
|
||||
|
||||
# Create workflow with parallel execution
|
||||
workflow = Workflow(
|
||||
name="Parallel Research",
|
||||
steps=[
|
||||
parallel([research_market, research_competitors, research_customers]),
|
||||
summarize_research
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=== Testing Parallel Workflow ===\n")
|
||||
|
||||
start_time = time.time()
|
||||
result = workflow.start("Analyze the business landscape", verbose=True)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
print(f"\nFinal Output:\n{result['output']}")
|
||||
print(f"\nCompleted in {elapsed:.2f}s (parallel execution)")
|
||||
print(f"Individual outputs: {len(result['variables']['parallel_outputs'])}")
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
Workflow Repeat (Evaluator-Optimizer) Example
|
||||
|
||||
Demonstrates repeating a step until a condition is met.
|
||||
This is useful for iterative improvement patterns.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, WorkflowContext, StepResult
|
||||
from praisonaiagents.memory.workflows import repeat
|
||||
|
||||
# Simulated content generator that improves each iteration
|
||||
class ContentGenerator:
|
||||
def __init__(self):
|
||||
self.iteration = 0
|
||||
self.content = []
|
||||
|
||||
def generate(self, ctx: WorkflowContext) -> StepResult:
|
||||
"""Generate content, adding more each iteration."""
|
||||
self.iteration += 1
|
||||
|
||||
# Add new content each iteration
|
||||
new_items = [
|
||||
f"Point {self.iteration}.1: Key insight about the topic",
|
||||
f"Point {self.iteration}.2: Supporting evidence"
|
||||
]
|
||||
self.content.extend(new_items)
|
||||
|
||||
output = f"Iteration {self.iteration}: Generated {len(self.content)} points total"
|
||||
|
||||
return StepResult(
|
||||
output=output,
|
||||
variables={
|
||||
"total_points": len(self.content),
|
||||
"content": self.content.copy()
|
||||
}
|
||||
)
|
||||
|
||||
# Condition: stop when we have enough content
|
||||
def has_enough_content(ctx: WorkflowContext) -> bool:
|
||||
"""Check if we have generated enough content."""
|
||||
total = ctx.variables.get("total_points", 0)
|
||||
return total >= 6 # Stop when we have 6+ points
|
||||
|
||||
# Create generator instance
|
||||
generator = ContentGenerator()
|
||||
|
||||
# Create workflow with repeat pattern
|
||||
workflow = Workflow(
|
||||
name="Content Generator",
|
||||
steps=[
|
||||
repeat(
|
||||
generator.generate,
|
||||
until=has_enough_content,
|
||||
max_iterations=10
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=== Testing Repeat (Evaluator-Optimizer) Pattern ===\n")
|
||||
|
||||
result = workflow.start("Generate comprehensive content", verbose=True)
|
||||
|
||||
print(f"\nFinal output: {result['output']}")
|
||||
print(f"Total iterations: {result['variables']['repeat_iterations']}")
|
||||
print(f"Total points generated: {result['variables']['total_points']}")
|
||||
print("\nGenerated content:")
|
||||
for point in result['variables']['content']:
|
||||
print(f" • {point}")
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
Workflow Routing Example
|
||||
|
||||
Demonstrates decision-based routing where the workflow
|
||||
takes different paths based on the output of a decision step.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Workflow, WorkflowContext, StepResult
|
||||
from praisonaiagents.memory.workflows import route
|
||||
|
||||
# Decision maker - determines which route to take
|
||||
def classify_request(ctx: WorkflowContext) -> StepResult:
|
||||
"""Classify the input and decide the route."""
|
||||
input_lower = ctx.input.lower()
|
||||
|
||||
if "urgent" in input_lower or "emergency" in input_lower:
|
||||
return StepResult(output="priority: high")
|
||||
elif "question" in input_lower or "help" in input_lower:
|
||||
return StepResult(output="priority: support")
|
||||
else:
|
||||
return StepResult(output="priority: normal")
|
||||
|
||||
# Route handlers
|
||||
def handle_high_priority(ctx: WorkflowContext) -> StepResult:
|
||||
return StepResult(output="🚨 HIGH PRIORITY: Escalating to senior team immediately!")
|
||||
|
||||
def handle_support(ctx: WorkflowContext) -> StepResult:
|
||||
return StepResult(output="💬 SUPPORT: Routing to help desk for assistance.")
|
||||
|
||||
def handle_normal(ctx: WorkflowContext) -> StepResult:
|
||||
return StepResult(output="📋 NORMAL: Added to standard processing queue.")
|
||||
|
||||
# Create workflow with routing
|
||||
workflow = Workflow(
|
||||
name="Request Router",
|
||||
steps=[
|
||||
classify_request,
|
||||
route({
|
||||
"high": [handle_high_priority],
|
||||
"support": [handle_support],
|
||||
"normal": [handle_normal],
|
||||
"default": [handle_normal]
|
||||
})
|
||||
]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test different inputs
|
||||
print("=== Testing Routing Workflow ===\n")
|
||||
|
||||
# Test 1: Urgent request
|
||||
result = workflow.start("This is an URGENT matter!", verbose=True)
|
||||
print(f"Result: {result['output']}\n")
|
||||
|
||||
# Test 2: Support request
|
||||
result = workflow.start("I have a question about my account", verbose=True)
|
||||
print(f"Result: {result['output']}\n")
|
||||
|
||||
# Test 3: Normal request
|
||||
result = workflow.start("Please process my order", verbose=True)
|
||||
print(f"Result: {result['output']}")
|
||||
@@ -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
|
||||
+1117
-54
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
Gemini Deep Research Agent Example
|
||||
|
||||
Uses Google's deep-research-pro model via the Interactions API
|
||||
for automated research with web search.
|
||||
|
||||
Supports two modes:
|
||||
1. Polling (default): Checks status periodically, shows progress logs
|
||||
2. Streaming: Real-time output with thinking summaries (recommended)
|
||||
|
||||
Requirements:
|
||||
- GEMINI_API_KEY or GOOGLE_API_KEY environment variable
|
||||
- google-genai >= 1.55.0 (pip install google-genai)
|
||||
"""
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
from praisonaiagents import DeepResearchAgent
|
||||
|
||||
agent = DeepResearchAgent(
|
||||
name="Gemini Researcher",
|
||||
instructions="""
|
||||
You are a professional researcher. Focus on:
|
||||
- Comprehensive analysis with multiple perspectives
|
||||
- Data-backed insights and trends
|
||||
- Clear structure with executive summary
|
||||
""",
|
||||
model="deep-research-pro", # Auto-detected as Gemini
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
print(f"Agent: {agent}")
|
||||
print(f"Provider: {agent.provider}")
|
||||
|
||||
# Streaming is enabled by default for real-time progress updates
|
||||
# This shows thinking summaries and streams the report as it's generated
|
||||
result = agent.research(
|
||||
"What is the current price of Ethereum?"
|
||||
# stream=True is the default, set stream=False to disable
|
||||
)
|
||||
|
||||
# The report is already printed during streaming when verbose=True
|
||||
# But we can also access it programmatically:
|
||||
print("\n" + "=" * 60)
|
||||
print(f"Report length: {len(result.report)} characters")
|
||||
print(f"Interaction ID: {result.interaction_id}")
|
||||
print(f"Reasoning steps captured: {len(result.reasoning_steps)}")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,47 @@
|
||||
"""
|
||||
OpenAI Deep Research Agent Example
|
||||
|
||||
Uses OpenAI's o3-deep-research or o4-mini-deep-research models
|
||||
via the Responses API for automated research with web search.
|
||||
|
||||
Supports two modes:
|
||||
1. Non-streaming (default): Waits for completion, returns full report
|
||||
2. Streaming: Real-time progress with reasoning summaries and web searches
|
||||
|
||||
Requirements:
|
||||
- OPENAI_API_KEY environment variable
|
||||
"""
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
from praisonaiagents import DeepResearchAgent
|
||||
|
||||
agent = DeepResearchAgent(
|
||||
name="Research Assistant",
|
||||
instructions="""
|
||||
You are a professional researcher. Focus on:
|
||||
- Data-rich insights with specific figures
|
||||
- Reliable sources and citations
|
||||
- Clear, structured responses
|
||||
""",
|
||||
model="o4-mini-deep-research", # or "o3-deep-research" for higher quality
|
||||
verbose=True
|
||||
)
|
||||
|
||||
print(f"Agent: {agent}")
|
||||
print(f"Provider: {agent.provider}")
|
||||
|
||||
# Streaming is enabled by default for real-time progress updates
|
||||
# Shows: reasoning summaries, web search queries, and report text as generated
|
||||
result = agent.research(
|
||||
"What is the current price of Bitcoin?"
|
||||
# stream=True is the default, set stream=False to disable
|
||||
)
|
||||
|
||||
# The report is already printed during streaming when verbose=True
|
||||
# But we can also access it programmatically:
|
||||
print("\n" + "=" * 60)
|
||||
print(f"Report length: {len(result.report)} characters")
|
||||
print(f"Citations: {len(result.citations)}")
|
||||
print(f"Web searches: {len(result.web_searches)}")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,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,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant",
|
||||
llm="gpt-5-nano"
|
||||
)
|
||||
agent.start("Why sky is Blue?")
|
||||
agent.start("What was my previous question?")
|
||||
@@ -0,0 +1,14 @@
|
||||
from praisonaiagents import Agent
|
||||
import asyncio
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant",
|
||||
llm="gpt-5-nano"
|
||||
)
|
||||
agent.run("Why sea is Blue?")
|
||||
|
||||
async def main():
|
||||
await agent.arun("Why sky is Blue?")
|
||||
await agent.arun("What was my previous question?")
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -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()
|
||||
@@ -15,9 +15,33 @@ _logging.configure_root_logger()
|
||||
from .agent.agent import Agent
|
||||
from .agent.image_agent import ImageAgent
|
||||
from .agent.context_agent import ContextAgent, create_context_agent
|
||||
from .agent.deep_research_agent import (
|
||||
DeepResearchAgent,
|
||||
DeepResearchResponse,
|
||||
Citation,
|
||||
ReasoningStep,
|
||||
WebSearchCall,
|
||||
CodeExecutionStep,
|
||||
MCPCall,
|
||||
FileSearchCall,
|
||||
Provider
|
||||
)
|
||||
from .agent.query_rewriter_agent import (
|
||||
QueryRewriterAgent,
|
||||
RewriteStrategy,
|
||||
RewriteResult
|
||||
)
|
||||
from .agent.prompt_expander_agent import (
|
||||
PromptExpanderAgent,
|
||||
ExpandStrategy,
|
||||
ExpandResult
|
||||
)
|
||||
from .agents.agents import PraisonAIAgents
|
||||
from .task.task import Task
|
||||
from .tools.tools import Tools
|
||||
from .tools.base import BaseTool, ToolResult, ToolValidationError, validate_tool
|
||||
from .tools.decorator import tool, FunctionTool
|
||||
from .tools.registry import get_registry, register_tool, get_tool, ToolRegistry
|
||||
from .agents.autoagents import AutoAgents
|
||||
from .knowledge.knowledge import Knowledge
|
||||
from .knowledge.chunking import Chunking
|
||||
@@ -30,7 +54,55 @@ except ImportError:
|
||||
MCP = None
|
||||
from .session import Session
|
||||
from .memory.memory import Memory
|
||||
from .memory.workflows import (
|
||||
Workflow, WorkflowStep, WorkflowContext, StepResult,
|
||||
Route, Parallel, Loop, Repeat,
|
||||
route, parallel, loop, repeat
|
||||
)
|
||||
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:
|
||||
@@ -154,6 +226,17 @@ __all__ = [
|
||||
'Agents',
|
||||
'Tools',
|
||||
'Task',
|
||||
# Plugin system
|
||||
'BaseTool',
|
||||
'ToolResult',
|
||||
'ToolValidationError',
|
||||
'validate_tool',
|
||||
'tool',
|
||||
'FunctionTool',
|
||||
'ToolRegistry',
|
||||
'get_registry',
|
||||
'register_tool',
|
||||
'get_tool',
|
||||
'TaskOutput',
|
||||
'ReflectionOutput',
|
||||
'AutoAgents',
|
||||
@@ -186,7 +269,22 @@ __all__ = [
|
||||
'disable_performance_mode',
|
||||
'cleanup_telemetry_resources',
|
||||
'MinimalTelemetry',
|
||||
'TelemetryCollector'
|
||||
'TelemetryCollector',
|
||||
'DeepResearchAgent',
|
||||
'DeepResearchResponse',
|
||||
'Citation',
|
||||
'ReasoningStep',
|
||||
'WebSearchCall',
|
||||
'CodeExecutionStep',
|
||||
'MCPCall',
|
||||
'FileSearchCall',
|
||||
'Provider',
|
||||
'QueryRewriterAgent',
|
||||
'RewriteStrategy',
|
||||
'RewriteResult',
|
||||
'PromptExpanderAgent',
|
||||
'ExpandStrategy',
|
||||
'ExpandResult'
|
||||
]
|
||||
|
||||
# Add MCP to __all__ if available
|
||||
@@ -197,3 +295,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'
|
||||
])
|
||||
|
||||
@@ -4,5 +4,52 @@ from .image_agent import ImageAgent
|
||||
from .context_agent import ContextAgent, create_context_agent
|
||||
from .handoff import Handoff, handoff, handoff_filters, RECOMMENDED_PROMPT_PREFIX, prompt_with_handoff_instructions
|
||||
from .router_agent import RouterAgent
|
||||
from .deep_research_agent import (
|
||||
DeepResearchAgent,
|
||||
DeepResearchResponse,
|
||||
Citation,
|
||||
ReasoningStep,
|
||||
WebSearchCall,
|
||||
CodeExecutionStep,
|
||||
MCPCall,
|
||||
FileSearchCall,
|
||||
Provider
|
||||
)
|
||||
from .query_rewriter_agent import (
|
||||
QueryRewriterAgent,
|
||||
RewriteStrategy,
|
||||
RewriteResult
|
||||
)
|
||||
from .prompt_expander_agent import (
|
||||
PromptExpanderAgent,
|
||||
ExpandStrategy,
|
||||
ExpandResult
|
||||
)
|
||||
|
||||
__all__ = ['Agent', 'ImageAgent', 'ContextAgent', 'create_context_agent', 'Handoff', 'handoff', 'handoff_filters', 'RECOMMENDED_PROMPT_PREFIX', 'prompt_with_handoff_instructions', 'RouterAgent']
|
||||
__all__ = [
|
||||
'Agent',
|
||||
'ImageAgent',
|
||||
'ContextAgent',
|
||||
'create_context_agent',
|
||||
'Handoff',
|
||||
'handoff',
|
||||
'handoff_filters',
|
||||
'RECOMMENDED_PROMPT_PREFIX',
|
||||
'prompt_with_handoff_instructions',
|
||||
'RouterAgent',
|
||||
'DeepResearchAgent',
|
||||
'DeepResearchResponse',
|
||||
'Citation',
|
||||
'ReasoningStep',
|
||||
'WebSearchCall',
|
||||
'CodeExecutionStep',
|
||||
'MCPCall',
|
||||
'FileSearchCall',
|
||||
'Provider',
|
||||
'QueryRewriterAgent',
|
||||
'RewriteStrategy',
|
||||
'RewriteResult',
|
||||
'PromptExpanderAgent',
|
||||
'ExpandStrategy',
|
||||
'ExpandResult'
|
||||
]
|
||||
@@ -234,7 +234,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()
|
||||
@@ -2082,27 +2561,199 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
display_error(f"Error in _achat_completion: {e}")
|
||||
return None
|
||||
|
||||
async def arun(self, prompt: str, **kwargs):
|
||||
"""Async alias for astart() method"""
|
||||
return await self.astart(prompt, **kwargs)
|
||||
|
||||
async def astart(self, prompt: str, **kwargs):
|
||||
"""Async version of start method"""
|
||||
return await self.achat(prompt, **kwargs)
|
||||
|
||||
def run(self):
|
||||
def run(self, prompt: str, **kwargs):
|
||||
"""Alias for start() method"""
|
||||
return self.start()
|
||||
return self.start(prompt, **kwargs)
|
||||
|
||||
def _get_planning_agent(self):
|
||||
"""Lazy load PlanningAgent for planning mode."""
|
||||
if self._planning_agent is None and self.planning:
|
||||
from ..planning import PlanningAgent
|
||||
self._planning_agent = PlanningAgent(
|
||||
llm=self.llm if hasattr(self, 'llm') else (self.llm_instance.model if hasattr(self, 'llm_instance') else "gpt-4o-mini"),
|
||||
tools=self.planning_tools,
|
||||
reasoning=self.planning_reasoning,
|
||||
verbose=1 if self.verbose else 0
|
||||
)
|
||||
return self._planning_agent
|
||||
|
||||
def _start_with_planning(self, prompt: str, **kwargs):
|
||||
"""Execute with planning mode - creates plan then executes each step."""
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.markdown import Markdown
|
||||
|
||||
console = Console()
|
||||
|
||||
# Step 1: Create the plan
|
||||
console.print("\n[bold blue]📋 PLANNING PHASE[/bold blue]")
|
||||
console.print("[dim]Creating implementation plan...[/dim]\n")
|
||||
|
||||
planner = self._get_planning_agent()
|
||||
plan = planner.create_plan_sync(request=prompt, agents=[self])
|
||||
|
||||
if not plan or not plan.steps:
|
||||
console.print("[yellow]⚠️ Planning failed, falling back to direct execution[/yellow]")
|
||||
return self.chat(prompt, **kwargs)
|
||||
|
||||
# Display the plan
|
||||
console.print(Panel(
|
||||
Markdown(plan.to_markdown()),
|
||||
title="[bold green]Generated Plan[/bold green]",
|
||||
border_style="green"
|
||||
))
|
||||
|
||||
# Step 2: Execute each step
|
||||
console.print("\n[bold blue]🚀 EXECUTION PHASE[/bold blue]\n")
|
||||
|
||||
results = []
|
||||
context = ""
|
||||
|
||||
for i, step in enumerate(plan.steps):
|
||||
progress = (i + 1) / len(plan.steps)
|
||||
bar_length = 30
|
||||
filled = int(bar_length * progress)
|
||||
bar = "█" * filled + "░" * (bar_length - filled)
|
||||
|
||||
console.print(f"[dim]Progress: [{bar}] {progress * 100:.0f}%[/dim]")
|
||||
console.print(f"\n[bold]📌 Step {i + 1}/{len(plan.steps)}:[/bold] {step.description[:60]}...")
|
||||
|
||||
# Build prompt with context from previous steps
|
||||
step_prompt = step.description
|
||||
if context:
|
||||
step_prompt = f"{step.description}\n\nContext from previous steps:\n{context}"
|
||||
|
||||
# Execute the step
|
||||
result = self.chat(step_prompt, **kwargs)
|
||||
results.append({"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."""
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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)
|
||||
@@ -0,0 +1,632 @@
|
||||
"""
|
||||
Query Rewriter Agent Module
|
||||
|
||||
This module provides the QueryRewriterAgent class for transforming user queries
|
||||
to improve retrieval quality in RAG applications.
|
||||
|
||||
Supported Rewriting Strategies:
|
||||
- **BASIC**: Simple rephrasing for clarity and keyword optimization
|
||||
- **HYDE**: Hypothetical Document Embeddings - generates a hypothetical answer
|
||||
- **STEP_BACK**: Generates higher-level concept questions for complex queries
|
||||
- **SUB_QUERIES**: Decomposes multi-part questions into focused sub-queries
|
||||
- **MULTI_QUERY**: Generates multiple paraphrased versions for ensemble retrieval
|
||||
- **CONTEXTUAL**: Uses conversation history to resolve references and context
|
||||
|
||||
Example:
|
||||
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
|
||||
|
||||
agent = QueryRewriterAgent(verbose=True)
|
||||
|
||||
# Basic rewriting
|
||||
result = agent.rewrite("AI trends")
|
||||
print(result.rewritten_queries)
|
||||
|
||||
# HyDE for better semantic matching
|
||||
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
|
||||
print(result.hypothetical_document)
|
||||
|
||||
# Contextual with chat history
|
||||
chat_history = [
|
||||
{"role": "user", "content": "Tell me about Python"},
|
||||
{"role": "assistant", "content": "Python is a programming language..."}
|
||||
]
|
||||
result = agent.rewrite("What about its performance?", chat_history=chat_history)
|
||||
|
||||
# With search tools for context-aware rewriting
|
||||
from praisonaiagents.tools import internet_search
|
||||
agent = QueryRewriterAgent(tools=[internet_search], verbose=True)
|
||||
result = agent.rewrite("latest AI developments") # Searches first, then rewrites
|
||||
"""
|
||||
|
||||
import logging
|
||||
import json
|
||||
from typing import List, Optional, Any, Dict
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class RewriteStrategy(Enum):
|
||||
"""Enumeration of available query rewriting strategies."""
|
||||
BASIC = "basic"
|
||||
HYDE = "hyde"
|
||||
STEP_BACK = "step_back"
|
||||
SUB_QUERIES = "sub_queries"
|
||||
MULTI_QUERY = "multi_query"
|
||||
CONTEXTUAL = "contextual"
|
||||
AUTO = "auto"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RewriteResult:
|
||||
"""Result of a query rewriting operation."""
|
||||
original_query: str
|
||||
rewritten_queries: List[str]
|
||||
strategy_used: RewriteStrategy
|
||||
hypothetical_document: Optional[str] = None
|
||||
step_back_question: Optional[str] = None
|
||||
sub_queries: Optional[List[str]] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __repr__(self):
|
||||
return f"RewriteResult(strategy={self.strategy_used.value}, queries={len(self.rewritten_queries)})"
|
||||
|
||||
@property
|
||||
def primary_query(self) -> str:
|
||||
"""Returns the primary rewritten query."""
|
||||
return self.rewritten_queries[0] if self.rewritten_queries else self.original_query
|
||||
|
||||
@property
|
||||
def all_queries(self) -> List[str]:
|
||||
"""Returns all queries including original and rewritten."""
|
||||
queries = [self.original_query]
|
||||
queries.extend(self.rewritten_queries)
|
||||
if self.step_back_question:
|
||||
queries.append(self.step_back_question)
|
||||
if self.sub_queries:
|
||||
queries.extend(self.sub_queries)
|
||||
return list(set(queries)) # Remove duplicates
|
||||
|
||||
|
||||
class QueryRewriterAgent:
|
||||
"""
|
||||
Agent for rewriting queries to improve retrieval quality in RAG applications.
|
||||
|
||||
This agent transforms user queries using various strategies to bridge the gap
|
||||
between how users ask questions and how information is stored in knowledge bases.
|
||||
|
||||
Attributes:
|
||||
name: Name of the agent
|
||||
model: LLM model to use for rewriting
|
||||
verbose: Whether to print detailed logs
|
||||
max_queries: Maximum number of queries to generate for multi-query strategy
|
||||
abbreviations: Dictionary of common abbreviations to expand
|
||||
|
||||
Example:
|
||||
agent = QueryRewriterAgent(model="gpt-4o-mini", verbose=True)
|
||||
result = agent.rewrite("ML best practices", strategy=RewriteStrategy.MULTI_QUERY)
|
||||
for query in result.rewritten_queries:
|
||||
print(query)
|
||||
"""
|
||||
|
||||
# Default prompts for each strategy
|
||||
PROMPTS = {
|
||||
"basic": """You are an expert at improving prompts and queries.
|
||||
|
||||
Given the following input, improve it while preserving its original intent and type.
|
||||
|
||||
Guidelines:
|
||||
- If it's a TASK (e.g., "write X", "create Y", "generate Z"), keep it as a task - do NOT convert to a question
|
||||
- If it's a QUESTION (e.g., "what is X?", "how does Y work?"), keep it as a question
|
||||
- Expand abbreviations and acronyms
|
||||
- Fix typos and grammatical errors
|
||||
- Add relevant context if too short
|
||||
- Keep the original intent and action intact
|
||||
|
||||
IMPORTANT: If the input is a command/task to DO something, the output must also be a command/task to DO that thing.
|
||||
|
||||
Input: {query}
|
||||
|
||||
Improved version:""",
|
||||
|
||||
"hyde": """You are an expert at generating hypothetical documents.
|
||||
|
||||
Given the following query, write a detailed hypothetical document that would perfectly answer this query. This document will be used to find similar real documents in a knowledge base.
|
||||
|
||||
Guidelines:
|
||||
- Write as if you are answering the query directly
|
||||
- Include specific details, facts, and terminology
|
||||
- Use the same style and vocabulary as the target documents
|
||||
- Be comprehensive but focused on the query topic
|
||||
|
||||
Query: {query}
|
||||
|
||||
Hypothetical document:""",
|
||||
|
||||
"step_back": """You are an expert at generating step-back questions.
|
||||
|
||||
Given the following query, generate a more general, higher-level question that captures the broader concept. This helps retrieve background information needed to answer the specific query.
|
||||
|
||||
Guidelines:
|
||||
- Identify the underlying concept or principle
|
||||
- Create a question about fundamentals or general knowledge
|
||||
- The step-back question should provide context for the original query
|
||||
|
||||
Query: {query}
|
||||
|
||||
Step-back question:""",
|
||||
|
||||
"sub_queries": """You are an expert at decomposing complex queries.
|
||||
|
||||
Given the following query, break it down into simpler, focused sub-queries. Each sub-query should address one specific aspect of the original question.
|
||||
|
||||
Guidelines:
|
||||
- Identify distinct aspects or components of the query
|
||||
- Create independent, self-contained sub-queries
|
||||
- Each sub-query should be answerable on its own
|
||||
- Return as a JSON array of strings
|
||||
|
||||
Query: {query}
|
||||
|
||||
Sub-queries (JSON array):""",
|
||||
|
||||
"multi_query": """You are an expert at query expansion and paraphrasing.
|
||||
|
||||
Given the following query, generate {num_queries} different versions that capture the same intent but use different words, phrases, or perspectives.
|
||||
|
||||
Guidelines:
|
||||
- Use synonyms and alternative phrasings
|
||||
- Try different question formats (what, how, why, etc.)
|
||||
- Include both formal and informal versions
|
||||
- Add relevant context or specificity
|
||||
- Return as a JSON array of strings
|
||||
|
||||
Query: {query}
|
||||
|
||||
Alternative queries (JSON array):""",
|
||||
|
||||
"contextual": """You are an expert at understanding conversational context.
|
||||
|
||||
Given the conversation history and a follow-up query, rewrite the query as a standalone question that captures all relevant context.
|
||||
|
||||
Guidelines:
|
||||
- Resolve pronouns and references (it, this, that, etc.)
|
||||
- Include relevant context from the conversation
|
||||
- Keep entity names and important details
|
||||
- Make the query self-contained and clear
|
||||
|
||||
Conversation history:
|
||||
{chat_history}
|
||||
|
||||
Follow-up query: {query}
|
||||
|
||||
Standalone question:""",
|
||||
|
||||
"auto": """You are an expert at analyzing and improving prompts and queries.
|
||||
|
||||
Analyze the following input and improve it while preserving its original intent and type.
|
||||
|
||||
Input types to consider:
|
||||
1. TASK/COMMAND (write, create, generate, build) → Keep as a task, make it clearer
|
||||
2. QUESTION (what, how, why, when) → Keep as a question, make it more specific
|
||||
3. Short/keyword inputs → Expand with context but preserve type
|
||||
4. Complex multi-part inputs → Clarify structure
|
||||
5. Well-formed inputs → Minor optimization only
|
||||
|
||||
IMPORTANT: Never convert a task into a question. "Write X" should remain "Write X", not "How to write X?"
|
||||
|
||||
Input: {query}
|
||||
{context}
|
||||
|
||||
Improved version:"""
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "QueryRewriterAgent",
|
||||
model: str = "gpt-4o-mini",
|
||||
instructions: Optional[str] = None,
|
||||
verbose: bool = False,
|
||||
max_queries: int = 5,
|
||||
abbreviations: Optional[Dict[str, str]] = None,
|
||||
temperature: float = 0.3,
|
||||
max_tokens: int = 500,
|
||||
tools: Optional[List[Any]] = None
|
||||
):
|
||||
"""
|
||||
Initialize the QueryRewriterAgent.
|
||||
|
||||
Args:
|
||||
name: Name of the agent
|
||||
model: LLM model to use (default: gpt-4o-mini)
|
||||
instructions: Custom instructions for the agent
|
||||
verbose: Whether to print detailed logs
|
||||
max_queries: Maximum queries for multi-query strategy
|
||||
abbreviations: Custom abbreviation expansions
|
||||
temperature: Temperature for LLM generation
|
||||
max_tokens: Maximum tokens for LLM response
|
||||
tools: Optional list of tools (agent decides when to use them)
|
||||
"""
|
||||
self.name = name
|
||||
self.model = model
|
||||
self.instructions = instructions
|
||||
self.verbose = verbose
|
||||
self.max_queries = max_queries
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
self.tools = tools or []
|
||||
self._agent = None # Lazy initialized
|
||||
|
||||
# Default abbreviations (can be extended)
|
||||
self.abbreviations = abbreviations or {
|
||||
"AI": "Artificial Intelligence",
|
||||
"ML": "Machine Learning",
|
||||
"NLP": "Natural Language Processing",
|
||||
"LLM": "Large Language Model",
|
||||
"RAG": "Retrieval Augmented Generation",
|
||||
"API": "Application Programming Interface",
|
||||
"DB": "Database",
|
||||
"UI": "User Interface",
|
||||
"UX": "User Experience",
|
||||
"SDK": "Software Development Kit",
|
||||
"CLI": "Command Line Interface",
|
||||
}
|
||||
|
||||
# Setup logging
|
||||
self.logger = logging.getLogger(__name__)
|
||||
if self.verbose:
|
||||
self.logger.setLevel(logging.DEBUG)
|
||||
|
||||
@property
|
||||
def agent(self):
|
||||
"""Lazy initialization of internal Agent."""
|
||||
if self._agent is None:
|
||||
from .agent import Agent
|
||||
self._agent = Agent(
|
||||
name=self.name,
|
||||
role="Query Rewriting Specialist",
|
||||
goal="Rewrite queries for better search and retrieval results. Use available tools when needed to gather context.",
|
||||
backstory="You are an expert at understanding user intent and transforming queries for optimal information retrieval.",
|
||||
tools=self.tools,
|
||||
llm=self.model,
|
||||
verbose=self.verbose,
|
||||
markdown=False
|
||||
)
|
||||
return self._agent
|
||||
|
||||
def _call_agent(self, prompt: str) -> str:
|
||||
"""Call the internal Agent with the given prompt."""
|
||||
try:
|
||||
response = self.agent.chat(prompt)
|
||||
return response if response else ""
|
||||
except Exception as e:
|
||||
self.logger.error(f"Agent call failed: {e}")
|
||||
raise
|
||||
|
||||
def _expand_abbreviations(self, query: str) -> str:
|
||||
"""Expand known abbreviations in the query."""
|
||||
words = query.split()
|
||||
expanded = []
|
||||
for word in words:
|
||||
clean_word = word.strip(".,!?;:")
|
||||
if clean_word.upper() in self.abbreviations:
|
||||
expanded.append(f"{clean_word} ({self.abbreviations[clean_word.upper()]})")
|
||||
else:
|
||||
expanded.append(word)
|
||||
return " ".join(expanded)
|
||||
|
||||
def _is_short_query(self, query: str, threshold: int = 5) -> bool:
|
||||
"""Check if query is considered short (keyword-style)."""
|
||||
return len(query.split()) < threshold
|
||||
|
||||
def _parse_json_array(self, text: str) -> List[str]:
|
||||
"""Parse a JSON array from LLM response."""
|
||||
try:
|
||||
# Try direct JSON parsing
|
||||
result = json.loads(text)
|
||||
if isinstance(result, list):
|
||||
return [str(item) for item in result]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Try to extract JSON array from text
|
||||
import re
|
||||
match = re.search(r'\[.*?\]', text, re.DOTALL)
|
||||
if match:
|
||||
try:
|
||||
result = json.loads(match.group())
|
||||
if isinstance(result, list):
|
||||
return [str(item) for item in result]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Fallback: split by newlines and clean up
|
||||
lines = [line.strip() for line in text.split('\n') if line.strip()]
|
||||
lines = [line.lstrip('0123456789.-) ').strip('"\'') for line in lines]
|
||||
return [line for line in lines if line and len(line) > 3]
|
||||
|
||||
def rewrite(
|
||||
self,
|
||||
query: str,
|
||||
strategy: RewriteStrategy = RewriteStrategy.AUTO,
|
||||
chat_history: Optional[List[Dict[str, str]]] = None,
|
||||
context: Optional[str] = None,
|
||||
num_queries: int = None
|
||||
) -> RewriteResult:
|
||||
"""
|
||||
Rewrite a query using the specified strategy.
|
||||
|
||||
All operations go through an internal Agent which automatically handles
|
||||
tool calling when tools are provided. The agent decides when to use tools
|
||||
based on the query context.
|
||||
|
||||
Args:
|
||||
query: The original user query
|
||||
strategy: Rewriting strategy to use
|
||||
chat_history: Previous conversation messages for contextual rewriting
|
||||
context: Additional context about the knowledge base
|
||||
num_queries: Number of queries for multi-query strategy
|
||||
|
||||
Returns:
|
||||
RewriteResult containing rewritten queries
|
||||
"""
|
||||
if self.verbose:
|
||||
self.logger.debug(f"Rewriting query: '{query}' with strategy: {strategy.value}")
|
||||
if self.tools:
|
||||
print(f"[cyan]Agent has {len(self.tools)} tools available (will use if needed)...[/cyan]")
|
||||
|
||||
# Auto-detect strategy if needed
|
||||
if strategy == RewriteStrategy.AUTO:
|
||||
strategy = self._detect_strategy(query, chat_history)
|
||||
if self.verbose:
|
||||
self.logger.debug(f"Auto-detected strategy: {strategy.value}")
|
||||
|
||||
# All calls go through the internal Agent which handles tool calling
|
||||
combined_context = context or ""
|
||||
|
||||
# Dispatch to appropriate method
|
||||
if strategy == RewriteStrategy.BASIC:
|
||||
return self._rewrite_basic(query, combined_context)
|
||||
elif strategy == RewriteStrategy.HYDE:
|
||||
return self._rewrite_hyde(query, combined_context)
|
||||
elif strategy == RewriteStrategy.STEP_BACK:
|
||||
return self._rewrite_step_back(query, combined_context)
|
||||
elif strategy == RewriteStrategy.SUB_QUERIES:
|
||||
return self._rewrite_sub_queries(query, combined_context)
|
||||
elif strategy == RewriteStrategy.MULTI_QUERY:
|
||||
return self._rewrite_multi_query(query, num_queries or self.max_queries, combined_context)
|
||||
elif strategy == RewriteStrategy.CONTEXTUAL:
|
||||
return self._rewrite_contextual(query, chat_history or [], combined_context)
|
||||
else:
|
||||
# Default to basic
|
||||
return self._rewrite_basic(query, combined_context)
|
||||
|
||||
def _detect_strategy(
|
||||
self,
|
||||
query: str,
|
||||
chat_history: Optional[List[Dict[str, str]]] = None
|
||||
) -> RewriteStrategy:
|
||||
"""Auto-detect the best rewriting strategy for a query."""
|
||||
# If chat history exists and query seems like a follow-up
|
||||
if chat_history and self._is_follow_up(query):
|
||||
return RewriteStrategy.CONTEXTUAL
|
||||
|
||||
# Short keyword queries
|
||||
if self._is_short_query(query):
|
||||
return RewriteStrategy.BASIC
|
||||
|
||||
# Complex multi-part queries
|
||||
if self._is_complex_query(query):
|
||||
return RewriteStrategy.SUB_QUERIES
|
||||
|
||||
# Questions that might benefit from step-back
|
||||
if self._needs_step_back(query):
|
||||
return RewriteStrategy.STEP_BACK
|
||||
|
||||
# Default to basic
|
||||
return RewriteStrategy.BASIC
|
||||
|
||||
def _is_follow_up(self, query: str) -> bool:
|
||||
"""Check if query appears to be a follow-up question."""
|
||||
follow_up_indicators = [
|
||||
"it", "this", "that", "they", "them", "their",
|
||||
"what about", "how about", "and", "also", "more",
|
||||
"same", "similar", "related"
|
||||
]
|
||||
query_lower = query.lower()
|
||||
return any(indicator in query_lower for indicator in follow_up_indicators)
|
||||
|
||||
def _is_complex_query(self, query: str) -> bool:
|
||||
"""Check if query contains multiple distinct questions."""
|
||||
# Check for multiple question marks or conjunctions
|
||||
if query.count('?') > 1:
|
||||
return True
|
||||
complex_indicators = [" and ", " or ", " also ", " as well as ", " plus "]
|
||||
return any(ind in query.lower() for ind in complex_indicators)
|
||||
|
||||
def _needs_step_back(self, query: str) -> bool:
|
||||
"""Check if query would benefit from a step-back question."""
|
||||
# Queries about specific details often benefit from broader context
|
||||
specific_indicators = [
|
||||
"specific", "exactly", "precisely", "particular",
|
||||
"difference between", "compare", "versus", "vs"
|
||||
]
|
||||
return any(ind in query.lower() for ind in specific_indicators)
|
||||
|
||||
def _rewrite_basic(self, query: str, context: str = "") -> RewriteResult:
|
||||
"""Basic query rewriting for clarity and keyword optimization."""
|
||||
# First expand abbreviations
|
||||
expanded_query = self._expand_abbreviations(query)
|
||||
|
||||
# Build prompt with optional context
|
||||
prompt = self.PROMPTS["basic"].format(query=expanded_query)
|
||||
if context:
|
||||
prompt += f"\n\nAdditional context to help with rewriting:\n{context}"
|
||||
|
||||
rewritten = self._call_agent(prompt).strip()
|
||||
|
||||
return RewriteResult(
|
||||
original_query=query,
|
||||
rewritten_queries=[rewritten],
|
||||
strategy_used=RewriteStrategy.BASIC,
|
||||
metadata={"expanded_abbreviations": expanded_query != query, "has_context": bool(context)}
|
||||
)
|
||||
|
||||
def _rewrite_hyde(self, query: str, context: str = "") -> RewriteResult:
|
||||
"""HyDE: Generate hypothetical document for better semantic matching."""
|
||||
prompt = self.PROMPTS["hyde"].format(query=query)
|
||||
if context:
|
||||
prompt += f"\n\nUse this context to make the hypothetical document more accurate:\n{context}"
|
||||
|
||||
hypothetical_doc = self._call_agent(prompt).strip()
|
||||
|
||||
# The hypothetical document itself becomes the search query
|
||||
return RewriteResult(
|
||||
original_query=query,
|
||||
rewritten_queries=[hypothetical_doc[:500]], # Truncate for embedding
|
||||
strategy_used=RewriteStrategy.HYDE,
|
||||
hypothetical_document=hypothetical_doc,
|
||||
metadata={"full_document_length": len(hypothetical_doc), "has_context": bool(context)}
|
||||
)
|
||||
|
||||
def _rewrite_step_back(self, query: str, context: str = "") -> RewriteResult:
|
||||
"""Generate step-back question for broader context retrieval."""
|
||||
# Generate step-back question
|
||||
prompt = self.PROMPTS["step_back"].format(query=query)
|
||||
if context:
|
||||
prompt += f"\n\nContext:\n{context}"
|
||||
step_back = self._call_agent(prompt).strip()
|
||||
|
||||
# Also do basic rewriting of original
|
||||
basic_result = self._rewrite_basic(query, context)
|
||||
|
||||
return RewriteResult(
|
||||
original_query=query,
|
||||
rewritten_queries=basic_result.rewritten_queries,
|
||||
strategy_used=RewriteStrategy.STEP_BACK,
|
||||
step_back_question=step_back,
|
||||
metadata={"includes_step_back": True, "has_context": bool(context)}
|
||||
)
|
||||
|
||||
def _rewrite_sub_queries(self, query: str, context: str = "") -> RewriteResult:
|
||||
"""Decompose complex query into focused sub-queries."""
|
||||
prompt = self.PROMPTS["sub_queries"].format(query=query)
|
||||
if context:
|
||||
prompt += f"\n\nContext to help with decomposition:\n{context}"
|
||||
response = self._call_agent(prompt).strip()
|
||||
|
||||
sub_queries = self._parse_json_array(response)
|
||||
|
||||
# Ensure we have at least one sub-query
|
||||
if not sub_queries:
|
||||
sub_queries = [query]
|
||||
|
||||
return RewriteResult(
|
||||
original_query=query,
|
||||
rewritten_queries=sub_queries,
|
||||
strategy_used=RewriteStrategy.SUB_QUERIES,
|
||||
sub_queries=sub_queries,
|
||||
metadata={"num_sub_queries": len(sub_queries), "has_context": bool(context)}
|
||||
)
|
||||
|
||||
def _rewrite_multi_query(self, query: str, num_queries: int, context: str = "") -> RewriteResult:
|
||||
"""Generate multiple paraphrased versions for ensemble retrieval."""
|
||||
prompt = self.PROMPTS["multi_query"].format(
|
||||
query=query,
|
||||
num_queries=num_queries
|
||||
)
|
||||
if context:
|
||||
prompt += f"\n\nContext to help with paraphrasing:\n{context}"
|
||||
response = self._call_agent(prompt).strip()
|
||||
|
||||
queries = self._parse_json_array(response)
|
||||
|
||||
# Ensure we have queries and limit to max
|
||||
if not queries:
|
||||
queries = [query]
|
||||
queries = queries[:num_queries]
|
||||
|
||||
return RewriteResult(
|
||||
original_query=query,
|
||||
rewritten_queries=queries,
|
||||
strategy_used=RewriteStrategy.MULTI_QUERY,
|
||||
metadata={"num_generated": len(queries), "has_context": bool(context)}
|
||||
)
|
||||
|
||||
def _rewrite_contextual(
|
||||
self,
|
||||
query: str,
|
||||
chat_history: List[Dict[str, str]],
|
||||
context: str = ""
|
||||
) -> RewriteResult:
|
||||
"""Rewrite query using conversation context."""
|
||||
# Format chat history
|
||||
history_str = ""
|
||||
for msg in chat_history[-10:]: # Last 10 messages
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
history_str += f"{role.capitalize()}: {content}\n"
|
||||
|
||||
prompt = self.PROMPTS["contextual"].format(
|
||||
chat_history=history_str,
|
||||
query=query
|
||||
)
|
||||
if context:
|
||||
prompt += f"\n\nAdditional context:\n{context}"
|
||||
rewritten = self._call_agent(prompt).strip()
|
||||
|
||||
return RewriteResult(
|
||||
original_query=query,
|
||||
rewritten_queries=[rewritten],
|
||||
strategy_used=RewriteStrategy.CONTEXTUAL,
|
||||
metadata={
|
||||
"history_length": len(chat_history),
|
||||
"history_used": min(10, len(chat_history)),
|
||||
"has_context": bool(context)
|
||||
}
|
||||
)
|
||||
|
||||
# Convenience methods for direct strategy access
|
||||
def rewrite_basic(self, query: str) -> RewriteResult:
|
||||
"""Convenience method for basic rewriting."""
|
||||
return self.rewrite(query, strategy=RewriteStrategy.BASIC)
|
||||
|
||||
def rewrite_hyde(self, query: str) -> RewriteResult:
|
||||
"""Convenience method for HyDE rewriting."""
|
||||
return self.rewrite(query, strategy=RewriteStrategy.HYDE)
|
||||
|
||||
def rewrite_step_back(self, query: str) -> RewriteResult:
|
||||
"""Convenience method for step-back rewriting."""
|
||||
return self.rewrite(query, strategy=RewriteStrategy.STEP_BACK)
|
||||
|
||||
def rewrite_sub_queries(self, query: str) -> RewriteResult:
|
||||
"""Convenience method for sub-query decomposition."""
|
||||
return self.rewrite(query, strategy=RewriteStrategy.SUB_QUERIES)
|
||||
|
||||
def rewrite_multi_query(self, query: str, num_queries: int = None) -> RewriteResult:
|
||||
"""Convenience method for multi-query generation."""
|
||||
return self.rewrite(
|
||||
query,
|
||||
strategy=RewriteStrategy.MULTI_QUERY,
|
||||
num_queries=num_queries
|
||||
)
|
||||
|
||||
def rewrite_contextual(
|
||||
self,
|
||||
query: str,
|
||||
chat_history: List[Dict[str, str]]
|
||||
) -> RewriteResult:
|
||||
"""Convenience method for contextual rewriting."""
|
||||
return self.rewrite(
|
||||
query,
|
||||
strategy=RewriteStrategy.CONTEXTUAL,
|
||||
chat_history=chat_history
|
||||
)
|
||||
|
||||
def add_abbreviation(self, abbrev: str, expansion: str) -> None:
|
||||
"""Add a custom abbreviation expansion."""
|
||||
self.abbreviations[abbrev.upper()] = expansion
|
||||
|
||||
def add_abbreviations(self, abbreviations: Dict[str, str]) -> None:
|
||||
"""Add multiple custom abbreviation expansions."""
|
||||
for abbrev, expansion in abbreviations.items():
|
||||
self.abbreviations[abbrev.upper()] = expansion
|
||||
@@ -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
|
||||
Loaded 100 of 218 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user