mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 17:33:33 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7b41444aad | ||
|
|
7c8b77063f | ||
|
|
59eb8d6356 | ||
|
|
89bddc31c6 | ||
|
|
e6e908298b | ||
|
|
2b3509c1d1 | ||
|
|
6e8029055a | ||
|
|
d2c9e8825a | ||
|
|
cd9816e218 | ||
|
|
d400149075 | ||
|
|
73b08e6b76 | ||
|
|
19c3a1f6a7 | ||
|
|
71a7968318 | ||
|
|
1b258229c1 | ||
|
|
d0e09f51f3 | ||
|
|
5adc862ead | ||
|
|
856db54334 | ||
|
|
f94e5bb3d6 | ||
|
|
1cece603ec | ||
|
|
9e2bbb9bbb | ||
|
|
35b4cff7f7 | ||
|
|
a094d1ff30 | ||
|
|
000bdf3418 | ||
|
|
826ff56fc5 | ||
|
|
4bd0195d5d | ||
|
|
812abec2eb | ||
|
|
e58f6deb4b | ||
|
|
f23680c121 | ||
|
|
e569400260 | ||
|
|
75c4d8432c | ||
|
|
6249d14d3f | ||
|
|
32537a227e | ||
|
|
f83ffcdafe | ||
|
|
519aebd271 | ||
|
|
fb4ecf1893 | ||
|
|
5a4cbf8aeb | ||
|
|
885a73f697 | ||
|
|
9336f15e67 | ||
|
|
6d55cfa87e | ||
|
|
0b431b1fe1 | ||
|
|
e252f86ac9 | ||
|
|
ed395de6f3 | ||
|
|
633469d12e | ||
|
|
c5c57dff26 | ||
|
|
573b9efcc1 | ||
|
|
dc4e51ac81 | ||
|
|
08e211c693 | ||
|
|
abd8021925 | ||
|
|
8a6c4eec7f | ||
|
|
4e77a6a019 | ||
|
|
41e33c09e0 | ||
|
|
96f00e2172 | ||
|
|
e858cebed9 | ||
|
|
cca79bd9fc | ||
|
|
4ffbc1b52a | ||
|
|
b948151a14 | ||
|
|
11eee70d9a | ||
|
|
da6a2c1c04 | ||
|
|
731268df5c | ||
|
|
d60021cb40 | ||
|
|
75894f26c1 | ||
|
|
c53e09480a | ||
|
|
2ef506fa8a | ||
|
|
7ce783e59b | ||
|
|
95ffd0c53f | ||
|
|
e1c90bf569 | ||
|
|
5fbd8aa783 | ||
|
|
381c86de79 | ||
|
|
35250e429e | ||
|
|
3027937057 | ||
|
|
d79443868f | ||
|
|
6705e753e8 | ||
|
|
b4d1532512 | ||
|
|
4b5d3b2eda | ||
|
|
a1d4f3ee43 | ||
|
|
b5e930ff2e | ||
|
|
1a6a14f8f2 | ||
|
|
ca69de4c6a | ||
|
|
9b3d36e6d7 | ||
|
|
a9359a0785 | ||
|
|
4cb016c3ca | ||
|
|
d2e4776884 | ||
|
|
193b82aeba | ||
|
|
029de283c4 | ||
|
|
26376a77af | ||
|
|
7cb580bb12 | ||
|
|
06b458cc54 | ||
|
|
e29c2fbe90 | ||
|
|
76032be618 | ||
|
|
80f5ba940b | ||
|
|
65791781e2 | ||
|
|
ee66154b73 | ||
|
|
0dc51060d8 | ||
|
|
2daae1f464 | ||
|
|
81bcda0440 | ||
|
|
e269cc7bf2 | ||
|
|
66e5dfcb41 | ||
|
|
ae3e69a286 | ||
|
|
40ea6fe1aa | ||
|
|
e30a7564be | ||
|
|
fbd0a736f0 | ||
|
|
eea9f757af | ||
|
|
02a8f075a6 | ||
|
|
57c39653d6 | ||
|
|
2f21d0f907 | ||
|
|
e0caaf8421 | ||
|
|
4bc35c3b6e | ||
|
|
4fd9b863a6 | ||
|
|
fc02392ddb | ||
|
|
4564639f18 | ||
|
|
16598104ab | ||
|
|
b5d13a2b7f | ||
|
|
4622bcda43 | ||
|
|
31bae3a231 | ||
|
|
466bcfedd7 | ||
|
|
c24bf80314 | ||
|
|
c36b3aee6e | ||
|
|
08512ba6d3 | ||
|
|
63ebfc0b54 | ||
|
|
b29269071d | ||
|
|
b6a86bf348 | ||
|
|
02ba71be3d | ||
|
|
e60f1c1aa1 | ||
|
|
450f237805 | ||
|
|
66e53fc82e | ||
|
|
7e6ee5f9f2 | ||
|
|
b669b9815d | ||
|
|
6d89cd603e | ||
|
|
6e955b5941 | ||
|
|
87a67b5f5d | ||
|
|
baec9cd0a4 | ||
|
|
e144dc497d | ||
|
|
8a77ab0ab5 | ||
|
|
0034b5c964 | ||
|
|
543c627374 | ||
|
|
93154adf9e | ||
|
|
3ab2462d96 | ||
|
|
aecf65a130 | ||
|
|
0a453d30a2 | ||
|
|
354a2c717b | ||
|
|
62514fbef4 | ||
|
|
f3c77a2e8b | ||
|
|
50dde4fa02 | ||
|
|
bbc4530329 | ||
|
|
aab975bec7 | ||
|
|
50c50fa567 | ||
|
|
68f280634a | ||
|
|
346488fbec | ||
|
|
57b5a1fd8d | ||
|
|
d75a5bb728 | ||
|
|
519ecc025a | ||
|
|
8290b95132 | ||
|
|
3d95e1ff0e | ||
|
|
74405d1bbe | ||
|
|
d307c3857c | ||
|
|
fbe5dddaae | ||
|
|
7e4108918d | ||
|
|
da7e38100e | ||
|
|
848459b816 | ||
|
|
cd63845398 | ||
|
|
7827427dc0 | ||
|
|
0e378d0754 | ||
|
|
206022eb9c | ||
|
|
901ba8074d | ||
|
|
953ee6e7bd | ||
|
|
181d7e1ede | ||
|
|
4039878422 | ||
|
|
461eba18d5 | ||
|
|
331a1de610 | ||
|
|
91632eb5e4 | ||
|
|
ad46225483 | ||
|
|
9af6758eb6 | ||
|
|
aac67a6572 | ||
|
|
2ceff9bfae | ||
|
|
3b7af02bea | ||
|
|
1ee7cfa2ef | ||
|
|
b01bd41574 | ||
|
|
6570e445b8 | ||
|
|
b421bcf925 | ||
|
|
41d0998b76 | ||
|
|
393e0ee33d | ||
|
|
0dce4de4a4 | ||
|
|
7fe4b48c6e | ||
|
|
5286912f34 | ||
|
|
d6cb1c0165 | ||
|
|
08ed3f94dc | ||
|
|
d3bfbc4b92 | ||
|
|
af7081d294 | ||
|
|
3b2f26699e | ||
|
|
7aa61b9155 | ||
|
|
7e37c940d4 | ||
|
|
2199a57dfe | ||
|
|
d49f046322 | ||
|
|
c16c2e38af | ||
|
|
b5df9f323c | ||
|
|
6e4abeef52 | ||
|
|
5656570172 | ||
|
|
c08c7cfedf | ||
|
|
fbbe80fd4d | ||
|
|
cabddd9a22 | ||
|
|
20b838f36a | ||
|
|
cf149a3c2d | ||
|
|
30f6f5706c | ||
|
|
0423c9c4af | ||
|
|
be82abfac3 | ||
|
|
4af6228f89 | ||
|
|
33eaf34ffb | ||
|
|
e2512b38c0 | ||
|
|
b02cda252d | ||
|
|
2a4b369e15 | ||
|
|
c0a9e55f6a | ||
|
|
7c4b28885d | ||
|
|
b0ec8b9a22 | ||
|
|
1b5cf2baea | ||
|
|
b1cc2c31ce | ||
|
|
2a40c57033 | ||
|
|
350d6e6488 | ||
|
|
20a49ffd70 | ||
|
|
4bf9916c5f | ||
|
|
9764483f36 | ||
|
|
ff405be834 | ||
|
|
6d384a96ca | ||
|
|
6770996c8b | ||
|
|
93c7731ccb | ||
|
|
2057c68a20 | ||
|
|
1ca928a2e4 | ||
|
|
66ea9d7148 | ||
|
|
74a71807a6 | ||
|
|
9e92e85a72 | ||
|
|
6d827a45c8 | ||
|
|
aabaaa6318 | ||
|
|
f518bd301a | ||
|
|
834f710cf6 | ||
|
|
ede71a7fbe | ||
|
|
7c751c88a3 | ||
|
|
17f3aab1e1 | ||
|
|
028fa8158d | ||
|
|
c62685adc5 | ||
|
|
60d6178310 | ||
|
|
b198154b71 | ||
|
|
205de2ba0b | ||
|
|
ee668b2711 | ||
|
|
a5bbb65017 | ||
|
|
6296755ebb | ||
|
|
4b6c847260 | ||
|
|
2fc459fe05 | ||
|
|
1c3c2941cc | ||
|
|
6e38a88cd3 | ||
|
|
f0f22bb13d | ||
|
|
a8fa70c455 | ||
|
|
1b4b4a10b5 | ||
|
|
1efdaac18f | ||
|
|
1b2182968f | ||
|
|
5e5a10589e | ||
|
|
6c5e42762f | ||
|
|
afb7e47080 | ||
|
|
e674a2ae84 | ||
|
|
d471d9754a | ||
|
|
9d25b1ed3d | ||
|
|
d0d04bc17e | ||
|
|
1f4e2b1c5a | ||
|
|
88d4241121 | ||
|
|
9e081e9413 | ||
|
|
bc0449b8de | ||
|
|
c262942ba7 | ||
|
|
209c2a42ae | ||
|
|
75c856b948 | ||
|
|
85fefac429 | ||
|
|
5bd3e45be4 | ||
|
|
4c62731e81 | ||
|
|
7ddbf98794 | ||
|
|
99506e80cf |
No files matched your search
@@ -0,0 +1,81 @@
|
||||
name: 'Gemini CLI Action'
|
||||
description: 'Uses the Gemini CLI to automate software development tasks.'
|
||||
author: 'Gemini'
|
||||
|
||||
inputs:
|
||||
prompt:
|
||||
description: 'A specific prompt to guide Gemini.'
|
||||
required: false
|
||||
default: 'You are a helpful assistant.'
|
||||
GEMINI_API_KEY:
|
||||
description: 'Your Gemini API key.'
|
||||
required: true
|
||||
OTLP_GCP_WIF_PROVIDER:
|
||||
description: 'The workload identity provider for GCP authentication.'
|
||||
required: false
|
||||
OTLP_GOOGLE_CLOUD_PROJECT:
|
||||
description: 'The Google Cloud project for telemetry.'
|
||||
required: false
|
||||
settings_json:
|
||||
description: |
|
||||
A JSON string to configure the Gemini CLI. This will be written to .gemini/settings.json.
|
||||
To enable observability with Google Cloud, include:
|
||||
{"telemetry": {"enabled": true, "target": "gcp"}, "sandbox": false}
|
||||
required: false
|
||||
version:
|
||||
description: |
|
||||
The version of @google/gemini-cli to execute.
|
||||
Can be a specific version from npm (e.g., '0.1.0', 'latest'), a branch name (e.g., 'main'), or a commit hash.
|
||||
required: false
|
||||
default: 'latest'
|
||||
|
||||
runs:
|
||||
using: "composite"
|
||||
steps:
|
||||
- name: Configure Gemini CLI
|
||||
if: inputs.settings_json != ''
|
||||
run: |
|
||||
mkdir -p .gemini
|
||||
echo "$SETTINGS_JSON" > .gemini/settings.json
|
||||
shell: bash
|
||||
env:
|
||||
SETTINGS_JSON: ${{ inputs.settings_json }}
|
||||
|
||||
- name: Authenticate to Google Cloud for Telemetry
|
||||
if: inputs.OTLP_GCP_WIF_PROVIDER != ''
|
||||
id: auth
|
||||
uses: 'google-github-actions/auth@v2'
|
||||
with:
|
||||
project_id: ${{ inputs.OTLP_GOOGLE_CLOUD_PROJECT }}
|
||||
workload_identity_provider: ${{ inputs.OTLP_GCP_WIF_PROVIDER }}
|
||||
|
||||
- name: Install Gemini CLI
|
||||
id: install
|
||||
run: |
|
||||
if [[ "${{ inputs.version }}" == "latest" || "${{ inputs.version }}" =~ ^[0-9]+\.[0-9]+\.[0-9]+(-[a-zA-Z0-9\.-]+)?(\+[a-zA-Z0-9\.-]+)?$ ]]; then
|
||||
echo "Installing Gemini CLI from npm: @google/gemini-cli@${{ inputs.version }}"
|
||||
npm install -g @google/gemini-cli@${{ inputs.version }}
|
||||
else
|
||||
echo "Installing Gemini CLI from GitHub: github:google-gemini/gemini-cli#${{ inputs.version }}"
|
||||
npm install -g github:google-gemini/gemini-cli#${{ inputs.version }}
|
||||
fi
|
||||
echo "Verifying installation:"
|
||||
if command -v gemini >/dev/null 2>&1; then
|
||||
gemini --version || echo "Gemini CLI installed successfully (version command not available)"
|
||||
else
|
||||
echo "Error: Gemini CLI not found in PATH"
|
||||
exit 1
|
||||
fi
|
||||
shell: bash
|
||||
|
||||
- name: Run Gemini CLI
|
||||
run: |
|
||||
gemini --yolo --prompt "$PROMPT"
|
||||
shell: bash
|
||||
env:
|
||||
GEMINI_API_KEY: ${{ inputs.GEMINI_API_KEY }}
|
||||
PROMPT: ${{ inputs.prompt }}
|
||||
|
||||
branding:
|
||||
icon: 'terminal'
|
||||
color: 'blue'
|
||||
@@ -1,4 +1,4 @@
|
||||
name: Auto Issue Comment
|
||||
name: Auto Issue Comment with PR Creation
|
||||
|
||||
on:
|
||||
issues:
|
||||
@@ -9,14 +9,18 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
issues: write
|
||||
contents: write
|
||||
pull-requests: write
|
||||
actions: read
|
||||
steps:
|
||||
- name: Add automated comment
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
github-token: ${{ secrets.GH_TOKEN }}
|
||||
script: |
|
||||
github.rest.issues.createComment({
|
||||
issue_number: context.issue.number,
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
body: '@claude review this issue with multiple parallel running agents and do a detailed analysis and fix this if the existing code doesn\'t have the solution implemented. Making sure it has backward compatibility, no existing features removed. After making those changes, again use multi agents to review the applied changes'
|
||||
})
|
||||
body: '@claude review this issue and do a detailed analysis and fix this if the existing code doesn\'t have the solution implemented. Making sure it has backward compatibility, no existing features removed. After making those changes, again review the applied changes. Use @web to search if you dont know any information or to find the latest documentation or to find the latest version. Run the code if you think you need to run it to test it. Minimal code change to start with if required any changes. **Please create a PR with your changes**.'
|
||||
})
|
||||
@@ -13,10 +13,11 @@ jobs:
|
||||
- name: Add automated comment
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
github-token: ${{ secrets.GH_TOKEN }}
|
||||
script: |
|
||||
github.rest.issues.createComment({
|
||||
issue_number: context.issue.number,
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
body: '@claude review this pull request with multiple parallel running agents and do a detailed analysis and fix this if the existing code doesn\'t have the solution implemented. Making sure it has backward compatibility, no existing features removed. After making those changes, again use multi agents to review the applied changes'
|
||||
body: '@claude review this pull request and do a detailed analysis and fix this if the existing code doesn\'t have the solution implemented. Making sure it has backward compatibility, no existing features removed. After making those changes, again review the applied changes. Use @web to search if you dont know any information or to find the latest documentation or to find the latest version. Run the code if you think you need to run it to test it. Minimal code change to start with if required any changes.'
|
||||
})
|
||||
+132
-20
@@ -1,4 +1,4 @@
|
||||
name: Claude Code
|
||||
name: Claude Assistant
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
@@ -6,34 +6,146 @@ on:
|
||||
pull_request_review_comment:
|
||||
types: [created]
|
||||
issues:
|
||||
types: [opened, assigned]
|
||||
types: [opened, assigned, labeled]
|
||||
pull_request_review:
|
||||
types: [submitted]
|
||||
|
||||
jobs:
|
||||
claude:
|
||||
claude-response:
|
||||
# Prevent bots from triggering this action
|
||||
if: |
|
||||
(github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
|
||||
(github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
|
||||
github.actor != 'dependabot[bot]' &&
|
||||
github.actor != 'github-actions[bot]' &&
|
||||
github.actor != 'cursor[bot]' &&
|
||||
!contains(github.actor, '[bot]') &&
|
||||
github.actor != 'renovate[bot]'
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
issues: read
|
||||
id-token: write
|
||||
contents: write
|
||||
pull-requests: write
|
||||
issues: write
|
||||
actions: read
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Run Claude Code (OAuth)
|
||||
id: claude
|
||||
uses: grll/claude-code-action@beta
|
||||
fetch-depth: 0
|
||||
|
||||
- uses: anthropics/claude-code-action@beta
|
||||
with:
|
||||
use_oauth: true
|
||||
claude_access_token: ${{ secrets.CLAUDE_ACCESS_TOKEN }}
|
||||
claude_refresh_token: ${{ secrets.CLAUDE_REFRESH_TOKEN }}
|
||||
claude_expires_at: ${{ secrets.CLAUDE_EXPIRES_AT }}
|
||||
claude_code_oauth_token: ${{ secrets.CLAUDE_ACCESS_TOKEN }}
|
||||
github_token: ${{ secrets.GITHUB_TOKEN }}
|
||||
trigger_phrase: "@claude"
|
||||
allowed_tools: |
|
||||
Bash(git:*)
|
||||
Bash(python:*)
|
||||
Bash(pip:*)
|
||||
Bash(conda:*)
|
||||
Bash(pytest:*)
|
||||
Bash(python -m pytest:*)
|
||||
Bash(python -m pip:*)
|
||||
Bash(poetry:*)
|
||||
View
|
||||
GlobTool
|
||||
GrepTool
|
||||
BatchTool
|
||||
Edit
|
||||
Replace
|
||||
mcp__github__get_issue
|
||||
mcp__github__get_issue_comments
|
||||
mcp__github__update_issue
|
||||
timeout_minutes: 30
|
||||
|
||||
post-review-comment:
|
||||
needs: claude-response
|
||||
if: success() && needs.claude-response.result == 'success'
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
issues: write
|
||||
steps:
|
||||
- name: Check for existing review comment
|
||||
id: check-comment
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
github-token: ${{ secrets.GH_TOKEN }}
|
||||
script: |
|
||||
const reviewText = '@claude review if the changes made are as per the requirements';
|
||||
let comments = [];
|
||||
let issueNumber;
|
||||
|
||||
try {
|
||||
// Determine issue/PR number based on event type
|
||||
if (context.eventName === 'pull_request_review_comment' || context.eventName === 'pull_request_review') {
|
||||
issueNumber = context.payload.pull_request?.number || context.issue.number;
|
||||
} else if (context.eventName === 'issue_comment' && context.payload.issue.pull_request) {
|
||||
// This is a comment on a PR (not an issue)
|
||||
issueNumber = context.payload.issue.number;
|
||||
} else {
|
||||
// This is a regular issue
|
||||
issueNumber = context.issue.number;
|
||||
}
|
||||
|
||||
comments = await github.rest.issues.listComments({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
issue_number: issueNumber
|
||||
});
|
||||
|
||||
// Look for Claude's completed review comments
|
||||
const claudeReviewComments = comments.data.filter(comment =>
|
||||
(comment.user.login === 'claude[bot]' || comment.body.includes('✅') || comment.body.includes('## Summary')) &&
|
||||
!comment.body.includes('⏳') &&
|
||||
!comment.body.includes('🔄') &&
|
||||
!comment.body.includes('Loading...')
|
||||
);
|
||||
|
||||
// Check if our review request comment already exists
|
||||
const existingReviewRequest = comments.data.find(comment =>
|
||||
comment.body.includes(reviewText)
|
||||
);
|
||||
|
||||
// Post if Claude has completed a review AND we haven't posted the review request yet
|
||||
const shouldPost = claudeReviewComments.length > 0 && !existingReviewRequest;
|
||||
core.setOutput('should-post', shouldPost);
|
||||
console.log(`Claude review completed: ${claudeReviewComments.length > 0}`);
|
||||
console.log(`Review request exists: ${!!existingReviewRequest}`);
|
||||
console.log(`Should post comment: ${shouldPost}`);
|
||||
|
||||
} catch (error) {
|
||||
console.log('Error checking comments:', error);
|
||||
core.setOutput('should-post', false); // If we can't check, don't post
|
||||
}
|
||||
|
||||
- name: Wait 30 seconds before posting comment
|
||||
if: steps.check-comment.outputs.should-post == 'true'
|
||||
run: sleep 30
|
||||
|
||||
- name: Post review comment
|
||||
if: steps.check-comment.outputs.should-post == 'true'
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
github-token: ${{ secrets.GH_TOKEN }}
|
||||
script: |
|
||||
let issueNumber;
|
||||
|
||||
// Determine issue/PR number based on event type
|
||||
if (context.eventName === 'pull_request_review_comment' || context.eventName === 'pull_request_review') {
|
||||
issueNumber = context.payload.pull_request?.number || context.issue.number;
|
||||
} else if (context.eventName === 'issue_comment' && context.payload.issue.pull_request) {
|
||||
// This is a comment on a PR (not an issue)
|
||||
issueNumber = context.payload.issue.number;
|
||||
} else {
|
||||
// This is a regular issue
|
||||
issueNumber = context.issue.number;
|
||||
}
|
||||
|
||||
console.log(`Posting review comment on issue/PR #${issueNumber}`);
|
||||
|
||||
await github.rest.issues.createComment({
|
||||
issue_number: issueNumber,
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
body: '@claude review if the changes made are as per the requirements and make changes if issues still exist.'
|
||||
});
|
||||
@@ -40,5 +40,6 @@ jobs:
|
||||
uses: docker/build-push-action@v2
|
||||
with:
|
||||
context: .
|
||||
file: ./docker/Dockerfile
|
||||
push: true
|
||||
tags: ghcr.io/mervinpraison/praisonai:latest
|
||||
@@ -0,0 +1,76 @@
|
||||
name: Gemini Automated Issue Triage
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
jobs:
|
||||
triage-issue:
|
||||
if: >
|
||||
github.event_name == 'issue_comment' &&
|
||||
!github.event.issue.pull_request &&
|
||||
contains(github.event.comment.body, '@gemini triage') &&
|
||||
(github.event.comment.author_association == 'OWNER' ||
|
||||
github.event.comment.author_association == 'MEMBER' ||
|
||||
github.event.comment.author_association == 'COLLABORATOR')
|
||||
timeout-minutes: 5
|
||||
permissions:
|
||||
issues: write
|
||||
contents: read
|
||||
id-token: write
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.issue.number }}
|
||||
cancel-in-progress: true
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Generate GitHub App Token
|
||||
id: generate_token
|
||||
uses: actions/create-github-app-token@v1
|
||||
with:
|
||||
app-id: ${{ secrets.APP_ID }}
|
||||
private-key: ${{ secrets.PRIVATE_KEY }}
|
||||
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
token: ${{ steps.generate_token.outputs.token }}
|
||||
|
||||
- name: Run Gemini Issue Triage
|
||||
uses: google-gemini/gemini-cli-action@main
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ steps.generate_token.outputs.token }}
|
||||
ISSUE_TITLE: ${{ github.event.issue.title }}
|
||||
ISSUE_BODY: ${{ github.event.issue.body }}
|
||||
ISSUE_NUMBER: ${{ github.event.issue.number }}
|
||||
REPOSITORY: ${{ github.repository }}
|
||||
with:
|
||||
version: 0.1.8-rc.0
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
OTLP_GCP_WIF_PROVIDER: ${{ secrets.OTLP_GCP_WIF_PROVIDER }}
|
||||
OTLP_GOOGLE_CLOUD_PROJECT: ${{ secrets.OTLP_GOOGLE_CLOUD_PROJECT }}
|
||||
settings_json: |
|
||||
{
|
||||
"coreTools": [
|
||||
"run_shell_command(gh label list)",
|
||||
"run_shell_command(gh issue edit)"
|
||||
],
|
||||
"telemetry": {
|
||||
"enabled": true,
|
||||
"target": "gcp"
|
||||
},
|
||||
"sandbox": false
|
||||
}
|
||||
prompt: |
|
||||
You are an issue triage assistant. Analyze the current GitHub issue and apply the most appropriate existing labels.
|
||||
|
||||
Steps:
|
||||
1. Run: `gh label list --limit 100` to get all available labels.
|
||||
2. Review the issue title and body provided in the environment variables.
|
||||
3. Select the most relevant labels from the existing labels, focusing on kind/*, area/*, and priority/*.
|
||||
4. Apply the selected labels to this issue using: `gh issue edit ISSUE_NUMBER --add-label "label1,label2"`
|
||||
|
||||
Guidelines:
|
||||
- Only use labels that already exist in the repository.
|
||||
- Do not add comments or modify the issue content.
|
||||
- Triage only the current issue.
|
||||
- Assign all applicable kind/*, area/*, and priority/* labels based on the issue content.
|
||||
@@ -0,0 +1,198 @@
|
||||
name: Gemini Issue Review
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
issues: write
|
||||
pull-requests: read
|
||||
|
||||
jobs:
|
||||
review-issue:
|
||||
if: >
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
(github.event_name == 'issue_comment' &&
|
||||
!github.event.issue.pull_request &&
|
||||
contains(github.event.comment.body, '@gemini') &&
|
||||
!contains(github.event.comment.body, '@gemini triage') &&
|
||||
(github.event.comment.author_association == 'OWNER' ||
|
||||
github.event.comment.author_association == 'MEMBER' ||
|
||||
github.event.comment.author_association == 'COLLABORATOR'))
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
steps:
|
||||
- name: Generate token
|
||||
id: generate_token
|
||||
uses: actions/create-github-app-token@v1
|
||||
with:
|
||||
app-id: ${{ secrets.APP_ID }}
|
||||
private-key: ${{ secrets.PRIVATE_KEY }}
|
||||
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
token: ${{ steps.generate_token.outputs.token }}
|
||||
|
||||
- name: Set up Git
|
||||
run: |
|
||||
git config --local user.email "[email protected]"
|
||||
git config --local user.name "Gemini Bot"
|
||||
|
||||
- name: Get issue data
|
||||
id: issue_data
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.generate_token.outputs.token }}
|
||||
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 }}"
|
||||
else
|
||||
echo "Unable to determine issue number"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Get issue details
|
||||
gh issue view $ISSUE_NUMBER --json title,body,labels,assignees > issue_details.json
|
||||
|
||||
# Extract additional instructions from comment if triggered by comment
|
||||
ADDITIONAL_INSTRUCTIONS=""
|
||||
if [ "${{ github.event_name }}" = "issue_comment" ]; then
|
||||
COMMENT_BODY="${{ github.event.comment.body }}"
|
||||
ADDITIONAL_INSTRUCTIONS=$(echo "$COMMENT_BODY" | sed -n 's/.*@gemini[[:space:]]*\(.*\)/\1/p' | head -1)
|
||||
fi
|
||||
|
||||
echo "ISSUE_NUMBER=$ISSUE_NUMBER" >> $GITHUB_OUTPUT
|
||||
echo "ADDITIONAL_INSTRUCTIONS=$ADDITIONAL_INSTRUCTIONS" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Run Gemini CLI
|
||||
uses: ./.github/actions/gemini
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ steps.generate_token.outputs.token }}
|
||||
ISSUE_NUMBER: ${{ steps.issue_data.outputs.ISSUE_NUMBER }}
|
||||
ADDITIONAL_INSTRUCTIONS: ${{ steps.issue_data.outputs.ADDITIONAL_INSTRUCTIONS }}
|
||||
REPOSITORY: ${{ github.repository }}
|
||||
with:
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
OTLP_GCP_WIF_PROVIDER: ${{ secrets.OTLP_GCP_WIF_PROVIDER }}
|
||||
OTLP_GOOGLE_CLOUD_PROJECT: ${{ secrets.OTLP_GOOGLE_CLOUD_PROJECT }}
|
||||
prompt: |
|
||||
You are a helpful GitHub issue assistant. Your task is to analyze and help with GitHub issues.
|
||||
|
||||
**Issue Details:**
|
||||
- Issue Number: #${ISSUE_NUMBER}
|
||||
- Repository: ${REPOSITORY}
|
||||
- Issue URL: ${{ github.event.issue.html_url }}
|
||||
|
||||
**Additional Instructions:** ${ADDITIONAL_INSTRUCTIONS}
|
||||
|
||||
**CRITICAL: When making code modifications, follow this EXACT workflow:**
|
||||
|
||||
**STEP 1: Analyze the Issue**
|
||||
- Use: `gh issue view ${ISSUE_NUMBER} --json title,body,labels,assignees`
|
||||
- Analyze issue content thoroughly
|
||||
- Search codebase for related files using `glob`, `search_file_content`, `list_directory`
|
||||
|
||||
**STEP 2: If Code Changes Are Needed**
|
||||
Create a new branch with 'gemini/' prefix:
|
||||
- Use: `git checkout -b gemini/issue-${ISSUE_NUMBER}-description`
|
||||
- Example: `git checkout -b gemini/issue-123-fix-sequential-tools`
|
||||
|
||||
**STEP 3: Make Code Modifications**
|
||||
- Use `read_file` to examine existing code
|
||||
- Use `write_file` or `replace` to make changes
|
||||
- Use `search_file_content` to find related code
|
||||
- Make focused, targeted changes to fix the issue
|
||||
|
||||
**STEP 4: Test Changes (if applicable)**
|
||||
- Run relevant tests if test files exist
|
||||
- Use `run_shell_command` to execute test commands
|
||||
|
||||
**STEP 5: Commit and Push Changes**
|
||||
- Use: `git add .` (or specific files)
|
||||
- Use: `git commit -m "Gemini: [Issue #${ISSUE_NUMBER}] Brief description of fix"`
|
||||
- Use: `git push -u origin gemini/issue-${ISSUE_NUMBER}-description`
|
||||
|
||||
**STEP 6: Create Pull Request**
|
||||
- Use: `gh pr create --title "Fix #${ISSUE_NUMBER}: Issue description" --body "Fixes #${ISSUE_NUMBER}\n\nSummary of changes:\n- Change 1\n- Change 2" --head gemini/issue-${ISSUE_NUMBER}-description --base main`
|
||||
- This will return a PR URL like: https://github.com/repo/pull/123
|
||||
|
||||
**STEP 7: Comment on Issue with Results**
|
||||
Use: `gh issue comment ${ISSUE_NUMBER} -b "MESSAGE"`
|
||||
|
||||
**Your final comment MUST include:**
|
||||
- Summary of what was analyzed
|
||||
- What changes were made (if any)
|
||||
- Link to the created PR (if applicable)
|
||||
- Brief explanation of the solution
|
||||
|
||||
**Example final comment:**
|
||||
```
|
||||
Hi! I've analyzed and fixed this issue.
|
||||
|
||||
**Problem:** Sequential tool calling was getting stuck after first tool call.
|
||||
|
||||
**Solution:** Added proper prompt handling after each tool call in the Ollama model integration.
|
||||
|
||||
**Changes Made:**
|
||||
- Modified `src/praisonai-agents/ollama-sequential.py`
|
||||
- Added continuation prompts for multi-step tool execution
|
||||
|
||||
**Pull Request:** https://github.com/MervinPraison/PraisonAI/pull/856
|
||||
|
||||
The fix is ready for review. Please test and let me know if you encounter any issues!
|
||||
```
|
||||
|
||||
**IMPORTANT NOTES:**
|
||||
- Always create branches with 'gemini/' prefix
|
||||
- Always include issue number in branch name and commit message
|
||||
- Always create PR that references the issue
|
||||
- Always provide PR link in your final response
|
||||
- Use the GITHUB_TOKEN environment variable for authentication
|
||||
|
||||
Please proceed with helping resolve this issue following this exact workflow.
|
||||
|
||||
settings_json: |
|
||||
{
|
||||
"coreTools": [
|
||||
"run_shell_command(echo)",
|
||||
"run_shell_command(gh issue view)",
|
||||
"run_shell_command(gh issue comment)",
|
||||
"run_shell_command(gh issue list)",
|
||||
"run_shell_command(gh issue edit)",
|
||||
"run_shell_command(gh pr create)",
|
||||
"run_shell_command(gh pr view)",
|
||||
"run_shell_command(gh pr list)",
|
||||
"run_shell_command(cat)",
|
||||
"run_shell_command(head)",
|
||||
"run_shell_command(tail)",
|
||||
"run_shell_command(grep)",
|
||||
"run_shell_command(find)",
|
||||
"run_shell_command(git checkout)",
|
||||
"run_shell_command(git add)",
|
||||
"run_shell_command(git commit)",
|
||||
"run_shell_command(git push)",
|
||||
"run_shell_command(git status)",
|
||||
"run_shell_command(git branch)",
|
||||
"list_directory",
|
||||
"read_file",
|
||||
"read_many_files",
|
||||
"write_file",
|
||||
"replace",
|
||||
"glob",
|
||||
"search_file_content",
|
||||
"web_fetch",
|
||||
"google_web_search",
|
||||
"save_memory"
|
||||
],
|
||||
"telemetry": {
|
||||
"enabled": false,
|
||||
"otlpGcpWifProvider": "${{ secrets.OTLP_GCP_WIF_PROVIDER }}",
|
||||
"googleCloudProject": "${{ secrets.OTLP_GOOGLE_CLOUD_PROJECT }}"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
name: Gemini PR Review
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
pr_number:
|
||||
description: 'PR number to review'
|
||||
required: true
|
||||
type: number
|
||||
|
||||
jobs:
|
||||
review-pr:
|
||||
if: >
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
(github.event_name == 'issue_comment' &&
|
||||
github.event.issue.pull_request &&
|
||||
contains(github.event.comment.body, '@gemini') &&
|
||||
(github.event.comment.author_association == 'OWNER' ||
|
||||
github.event.comment.author_association == 'MEMBER' ||
|
||||
github.event.comment.author_association == 'COLLABORATOR'))
|
||||
timeout-minutes: 15
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
issues: write
|
||||
steps:
|
||||
- name: Generate GitHub App Token
|
||||
id: generate_token
|
||||
uses: actions/create-github-app-token@v1
|
||||
with:
|
||||
app-id: ${{ secrets.APP_ID }}
|
||||
private-key: ${{ secrets.PRIVATE_KEY }}
|
||||
|
||||
- name: Determine checkout ref
|
||||
id: checkout_ref
|
||||
run: |
|
||||
if [ "${{ github.event_name }}" = "issue_comment" ]; then
|
||||
echo "ref=refs/pull/${{ github.event.issue.number }}/head" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "ref=${{ github.event.pull_request.head.sha }}" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Checkout PR code
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
token: ${{ steps.generate_token.outputs.token }}
|
||||
ref: ${{ steps.checkout_ref.outputs.ref }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Get PR details
|
||||
id: get_pr
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ steps.generate_token.outputs.token }}
|
||||
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 }}
|
||||
else
|
||||
PR_NUMBER=${{ github.event.pull_request.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 }}"
|
||||
ADDITIONAL_INSTRUCTIONS=$(echo "$COMMENT_BODY" | sed 's/.*@gemini//' | xargs)
|
||||
fi
|
||||
echo "additional_instructions=$ADDITIONAL_INSTRUCTIONS" >> "$GITHUB_OUTPUT"
|
||||
|
||||
# Get PR details
|
||||
PR_DATA=$(gh pr view $PR_NUMBER --json title,body,additions,deletions,changedFiles,baseRefName,headRefName)
|
||||
echo "pr_data=$PR_DATA" >> "$GITHUB_OUTPUT"
|
||||
|
||||
# Get file changes
|
||||
CHANGED_FILES=$(gh pr diff $PR_NUMBER --name-only)
|
||||
echo "changed_files<<EOF" >> "$GITHUB_OUTPUT"
|
||||
echo "$CHANGED_FILES" >> "$GITHUB_OUTPUT"
|
||||
echo "EOF" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Run Gemini PR Review
|
||||
uses: ./.github/actions/gemini
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ steps.generate_token.outputs.token }}
|
||||
PR_NUMBER: ${{ steps.get_pr.outputs.pr_number }}
|
||||
PR_DATA: ${{ steps.get_pr.outputs.pr_data }}
|
||||
CHANGED_FILES: ${{ steps.get_pr.outputs.changed_files }}
|
||||
ADDITIONAL_INSTRUCTIONS: ${{ steps.get_pr.outputs.additional_instructions }}
|
||||
REPOSITORY: ${{ github.repository }}
|
||||
with:
|
||||
version: 0.1.8-rc.0
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
OTLP_GCP_WIF_PROVIDER: ${{ secrets.OTLP_GCP_WIF_PROVIDER }}
|
||||
OTLP_GOOGLE_CLOUD_PROJECT: ${{ secrets.OTLP_GOOGLE_CLOUD_PROJECT }}
|
||||
settings_json: |
|
||||
{
|
||||
"coreTools": [
|
||||
"run_shell_command(echo)",
|
||||
"run_shell_command(gh pr view)",
|
||||
"run_shell_command(gh pr diff)",
|
||||
"run_shell_command(gh pr comment)",
|
||||
"run_shell_command(gh pr list)",
|
||||
"run_shell_command(cat)",
|
||||
"run_shell_command(head)",
|
||||
"run_shell_command(tail)",
|
||||
"run_shell_command(grep)",
|
||||
"run_shell_command(find)",
|
||||
"run_shell_command(git checkout)",
|
||||
"run_shell_command(git add)",
|
||||
"run_shell_command(git commit)",
|
||||
"run_shell_command(git push)",
|
||||
"run_shell_command(git status)",
|
||||
"run_shell_command(git branch)",
|
||||
"run_shell_command(git fetch)",
|
||||
"list_directory",
|
||||
"read_file",
|
||||
"read_many_files",
|
||||
"write_file",
|
||||
"replace",
|
||||
"glob",
|
||||
"search_file_content",
|
||||
"web_fetch",
|
||||
"google_web_search",
|
||||
"save_memory"
|
||||
],
|
||||
"telemetry": {
|
||||
"enabled": false,
|
||||
"target": "gcp"
|
||||
},
|
||||
"sandbox": false
|
||||
}
|
||||
prompt: |
|
||||
You are an expert code reviewer and AI assistant for Pull Requests. You can review code AND make modifications to improve the PR.
|
||||
|
||||
**PR Details:**
|
||||
- PR Number: #$PR_NUMBER
|
||||
- Repository: $REPOSITORY
|
||||
- Additional Instructions: $ADDITIONAL_INSTRUCTIONS
|
||||
|
||||
**CRITICAL: When making code modifications, follow this EXACT workflow:**
|
||||
|
||||
**STEP 1: Analyze the PR**
|
||||
- Use: `echo "$PR_DATA"` to get PR details (JSON format)
|
||||
- Use: `echo "$CHANGED_FILES"` to get changed files list
|
||||
- Use: `gh pr diff ${PR_NUMBER}` to see the full diff
|
||||
- Use: `gh pr view ${PR_NUMBER} --json title,body,headRefName` to get PR info
|
||||
- Analyze the code changes thoroughly
|
||||
|
||||
**STEP 2: If Code Improvements Are Needed**
|
||||
**IMPORTANT:** DO NOT create a new branch! Work on the EXISTING PR branch.
|
||||
- Use: `git fetch origin` to get latest changes
|
||||
- Get PR branch name: `gh pr view ${PR_NUMBER} --json headRefName`
|
||||
- Switch to PR branch: `git checkout BRANCH_NAME` (use actual branch name from PR)
|
||||
- Use: `git pull origin BRANCH_NAME` to get latest PR changes
|
||||
|
||||
**STEP 3: Make Code Modifications**
|
||||
- Use `read_file` to examine existing code
|
||||
- Use `write_file` or `replace` to make improvements
|
||||
- Use `search_file_content` to find related code
|
||||
- Make focused, targeted improvements to the PR
|
||||
|
||||
**STEP 4: Test Changes (if applicable)**
|
||||
- Run relevant tests if test files exist
|
||||
- Use `run_shell_command` to execute test commands
|
||||
|
||||
**STEP 5: Commit and Push Changes to EXISTING PR Branch**
|
||||
- Use: `git add .` (or specific files)
|
||||
- Use: `git commit -m "Gemini: [PR #${PR_NUMBER}] Brief description of improvements"`
|
||||
- Use: `git push origin BRANCH_NAME` (push to EXISTING PR branch)
|
||||
|
||||
**STEP 6: Review and Comment on PR**
|
||||
Write comprehensive review and post comment with your analysis and any improvements made.
|
||||
Use: `gh pr comment ${PR_NUMBER} -b "MESSAGE"`
|
||||
|
||||
Review Guidelines:
|
||||
- Focus on code quality, security, performance, and maintainability
|
||||
- Check for common issues: potential bugs, security vulnerabilities, performance bottlenecks
|
||||
- Verify error handling and edge cases
|
||||
- Look for code style and best practices
|
||||
- Comment on architecture and design decisions if significant
|
||||
- Be constructive and specific in feedback
|
||||
- Highlight both issues and positive aspects
|
||||
- Suggest improvements with examples when possible
|
||||
|
||||
Review Areas:
|
||||
- **Security**: Authentication, authorization, input validation, data sanitization
|
||||
- **Performance**: Algorithms, database queries, caching, resource usage
|
||||
- **Reliability**: Error handling, logging, testing coverage, edge cases
|
||||
- **Maintainability**: Code structure, documentation, naming conventions
|
||||
- **Functionality**: Logic correctness, requirements fulfillment
|
||||
|
||||
Output Format:
|
||||
Structure your review using this exact format with markdown:
|
||||
|
||||
## 📋 Review Summary
|
||||
Provide a brief 2-3 sentence overview of the PR and overall assessment.
|
||||
|
||||
## 🔍 General Feedback
|
||||
- List general observations about code quality
|
||||
- Mention overall patterns or architectural decisions
|
||||
- Highlight positive aspects of the implementation
|
||||
- Note any recurring themes across files
|
||||
|
||||
## 🎯 Specific Feedback
|
||||
Only include sections below that have actual issues. If there are no issues in a priority category, omit that entire section.
|
||||
|
||||
### 🔴 Critical
|
||||
(Only include this section if there are critical issues)
|
||||
Issues that must be addressed before merging (security vulnerabilities, breaking changes, major bugs):
|
||||
- **File: `filename:line`** - Description of critical issue with specific recommendation
|
||||
|
||||
### 🟡 High
|
||||
(Only include this section if there are high priority issues)
|
||||
Important issues that should be addressed (performance problems, design flaws, significant bugs):
|
||||
- **File: `filename:line`** - Description of high priority issue with suggested fix
|
||||
|
||||
### 🟢 Medium
|
||||
(Only include this section if there are medium priority issues)
|
||||
Improvements that would enhance code quality (style issues, minor optimizations, better practices):
|
||||
- **File: `filename:line`** - Description of medium priority improvement
|
||||
|
||||
### 🔵 Low
|
||||
(Only include this section if there are suggestions)
|
||||
Nice-to-have improvements and suggestions (documentation, naming, minor refactoring):
|
||||
- **File: `filename:line`** - Description of suggestion or enhancement
|
||||
|
||||
**Note**: If no specific issues are found in any category, simply state "No specific issues identified in this review."
|
||||
|
||||
## ✅ Highlights
|
||||
(Only include this section if there are positive aspects to highlight)
|
||||
- Mention specific good practices or implementations
|
||||
- Acknowledge well-written code sections
|
||||
- Note improvements from previous versions
|
||||
|
||||
---
|
||||
*Review completed by Gemini CLI*
|
||||
@@ -151,7 +151,7 @@ jobs:
|
||||
return 'INVALID_TYPE'
|
||||
|
||||
print('🔧 Direct PraisonAI API Key Check:')
|
||||
env_api_key = os.environ.get(\\"OPENAI_API_KEY\\", \\"NOT_SET\\")
|
||||
env_api_key = os.environ.get(\"OPENAI_API_KEY\", \"NOT_SET\")
|
||||
print(f'Environment OPENAI_API_KEY: {get_key_display_value(env_api_key)}...')
|
||||
|
||||
from praisonai import PraisonAI
|
||||
@@ -163,7 +163,7 @@ jobs:
|
||||
|
||||
from praisonai.inc.models import PraisonAIModel
|
||||
|
||||
print('\\\\n🧪 Testing PraisonAIModel with explicit API key (CrewAI method):')
|
||||
print('\\n🧪 Testing PraisonAIModel with explicit API key (CrewAI method):')
|
||||
model_with_explicit_key = PraisonAIModel(
|
||||
model='openai/gpt-4o-mini',
|
||||
base_url=praisonai.config_list[0].get('base_url'),
|
||||
|
||||
+3
-1
@@ -75,4 +75,6 @@ agents/praisonaiagents/praisonaiagents.egg-info
|
||||
# Local Netlify folder
|
||||
.netlify
|
||||
.qodo
|
||||
CopilotKit*
|
||||
CopilotKit*
|
||||
.gemini
|
||||
issue_details.json
|
||||
+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.48" \
|
||||
"praisonai>=2.2.58" \
|
||||
"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.48" \
|
||||
"praisonai>=2.2.58" \
|
||||
"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.48" \
|
||||
"praisonai>=2.2.58" \
|
||||
"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.48" \
|
||||
"praisonai>=2.2.58" \
|
||||
"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.48`
|
||||
- PraisonAI: `>=2.2.58`
|
||||
- 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.48" "praisonaiagents==0.0.92"
|
||||
RUN pip install "praisonai==2.2.58" "praisonaiagents==0.0.92"
|
||||
```
|
||||
|
||||
## 🌐 Production Deployment
|
||||
|
||||
@@ -0,0 +1,293 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZOWgx2fYqW8T"
|
||||
},
|
||||
"source": [
|
||||
"# GitHub Repository Insights with PraisonAIAgents"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "XgJdLRiIqZ2L"
|
||||
},
|
||||
"source": [
|
||||
"This notebook demonstrates how to explore and analyze GitHub repositories using PraisonAIAgents in Google Colab. You can query repository issues, pull requests, and activity patterns with natural language, leveraging the power of agentic tools—all without leaving your notebook."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Y_koMT6qzxne"
|
||||
},
|
||||
"source": [
|
||||
"[](https://colab.research.google.com/github/Dhivya-Bharathy/PraisonAI/blob/main/examples/cookbooks/Programming_Code_Analysis_Agents/GitHub_Repository_Insights.ipynb)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "2ey7zWAJrIPx"
|
||||
},
|
||||
"source": [
|
||||
"# Dependencies"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 24,
|
||||
"metadata": {
|
||||
"id": "_xFH-iLTxzd8"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install praisonaiagents openai requests --quiet"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Kzty0zxzsbZ_"
|
||||
},
|
||||
"source": [
|
||||
"# Set Up Key"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 25,
|
||||
"metadata": {
|
||||
"id": "HCVP5fmvx2hH"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"# Set your GitHub token here\n",
|
||||
"GITHUB_TOKEN = \"your-github-token\" # <-- Replace with your actual token\n",
|
||||
"\n",
|
||||
"# Set your OpenAI API key here (required by PraisonAIAgents, even if not used for LLM)\n",
|
||||
"OPENAI_API_KEY = \"your-openai-api-key\" # <-- Replace with your actual OpenAI API key\n",
|
||||
"\n",
|
||||
"# Set environment variable for OpenAI\n",
|
||||
"os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "10BTLbF4s_f-"
|
||||
},
|
||||
"source": [
|
||||
"# Import Libraries"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 26,
|
||||
"metadata": {
|
||||
"id": "_sLahvfXyEb0"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from praisonaiagents import Agent\n",
|
||||
"import requests\n",
|
||||
"from IPython.display import display, Markdown"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "GRftdghxtUDs"
|
||||
},
|
||||
"source": [
|
||||
"# Custom Tool"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 27,
|
||||
"metadata": {
|
||||
"id": "nT11EFNNyIHl"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def github_query_tool(query: str, repo: str, github_token: str) -> str:\n",
|
||||
" \"\"\"\n",
|
||||
" Custom tool to query GitHub repository data using the GitHub API.\n",
|
||||
" Supports basic queries for issues, pull requests, and repository info.\n",
|
||||
" \"\"\"\n",
|
||||
" headers = {\n",
|
||||
" \"Authorization\": f\"token {github_token}\",\n",
|
||||
" \"Accept\": \"application/vnd.github.v3+json\"\n",
|
||||
" }\n",
|
||||
" base_url = f\"https://api.github.com/repos/{repo}\"\n",
|
||||
" result = \"\"\n",
|
||||
"\n",
|
||||
" if \"issue\" in query.lower():\n",
|
||||
" url = f\"{base_url}/issues\"\n",
|
||||
" response = requests.get(url, headers=headers)\n",
|
||||
" if response.ok:\n",
|
||||
" issues = response.json()\n",
|
||||
" result = f\"Found {len(issues)} issues in {repo}.\\n\"\n",
|
||||
" for issue in issues[:5]:\n",
|
||||
" result += f\"- #{issue['number']}: {issue['title']} (State: {issue['state']})\\n\"\n",
|
||||
" else:\n",
|
||||
" result = f\"Failed to fetch issues: {response.text}\"\n",
|
||||
"\n",
|
||||
" elif \"pull request\" in query.lower() or \"pr\" in query.lower():\n",
|
||||
" url = f\"{base_url}/pulls\"\n",
|
||||
" response = requests.get(url, headers=headers)\n",
|
||||
" if response.ok:\n",
|
||||
" pulls = response.json()\n",
|
||||
" result = f\"Found {len(pulls)} pull requests in {repo}.\\n\"\n",
|
||||
" for pr in pulls[:5]:\n",
|
||||
" result += f\"- #{pr['number']}: {pr['title']} (State: {pr['state']})\\n\"\n",
|
||||
" else:\n",
|
||||
" result = f\"Failed to fetch pull requests: {response.text}\"\n",
|
||||
"\n",
|
||||
" elif \"activity\" in query.lower() or \"stats\" in query.lower():\n",
|
||||
" url = f\"{base_url}/stats/contributors\"\n",
|
||||
" response = requests.get(url, headers=headers)\n",
|
||||
" if response.ok:\n",
|
||||
" stats = response.json()\n",
|
||||
" result = f\"Repository activity for {repo}:\\n\"\n",
|
||||
" for contributor in stats[:5]:\n",
|
||||
" result += f\"- {contributor['author']['login']}: {contributor['total']} commits\\n\"\n",
|
||||
" else:\n",
|
||||
" result = f\"Failed to fetch activity stats: {response.text}\"\n",
|
||||
"\n",
|
||||
" else:\n",
|
||||
" # Default: basic repo info\n",
|
||||
" response = requests.get(base_url, headers=headers)\n",
|
||||
" if response.ok:\n",
|
||||
" info = response.json()\n",
|
||||
" result = (\n",
|
||||
" f\"Repository: {info['full_name']}\\n\"\n",
|
||||
" f\"Description: {info.get('description', 'No description')}\\n\"\n",
|
||||
" f\"Stars: {info['stargazers_count']}\\n\"\n",
|
||||
" f\"Forks: {info['forks_count']}\\n\"\n",
|
||||
" f\"Open Issues: {info['open_issues_count']}\\n\"\n",
|
||||
" )\n",
|
||||
" else:\n",
|
||||
" result = f\"Failed to fetch repository info: {response.text}\"\n",
|
||||
"\n",
|
||||
" return result"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "mAFg5SW-tacD"
|
||||
},
|
||||
"source": [
|
||||
"# Register Tool and Create Agent"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 28,
|
||||
"metadata": {
|
||||
"id": "x_MmUc_AyMaB"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def github_tool_wrapper(inputs):\n",
|
||||
" query = inputs.get(\"query\", \"\")\n",
|
||||
" repo = inputs.get(\"repo\", \"\")\n",
|
||||
" github_token = inputs.get(\"github_token\", \"\")\n",
|
||||
" return github_query_tool(query, repo, github_token)\n",
|
||||
"\n",
|
||||
"agent = Agent(\n",
|
||||
" instructions=\"\"\"\n",
|
||||
" You are a GitHub assistant. Help users explore repositories and their activity.\n",
|
||||
" - Provide organized, concise insights about the repository.\n",
|
||||
" - Focus on facts and data from the GitHub API.\n",
|
||||
" - Use markdown formatting for better readability.\n",
|
||||
" - Present numerical data in tables when appropriate.\n",
|
||||
" - Include links to relevant GitHub pages when helpful.\n",
|
||||
" \"\"\",\n",
|
||||
" tools={\n",
|
||||
" \"github_query_tool\": github_tool_wrapper\n",
|
||||
" }\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "lPstz7bbtr0k"
|
||||
},
|
||||
"source": [
|
||||
"# User Input and Run Agent"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 30,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 118
|
||||
},
|
||||
"id": "MymtRmlnyQqd",
|
||||
"outputId": "33a23810-4a67-4615-f9e9-146b3030fcbf"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Enter the GitHub repository (owner/repo): Dhivya-Bharathy/PraisonAI\n",
|
||||
"Enter your question (e.g., Show me open issues): Show me open issues\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"### Results\n",
|
||||
"Found 0 issues in Dhivya-Bharathy/PraisonAI.\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# User input\n",
|
||||
"repo = input(\"Enter the GitHub repository (owner/repo): \")\n",
|
||||
"query = input(\"Enter your question (e.g., Show me open issues): \")\n",
|
||||
"\n",
|
||||
"# Call the tool directly (since agent.chat expects an LLM, but we want tool output)\n",
|
||||
"response = agent.execute_tool(\"github_query_tool\", {\n",
|
||||
" \"query\": query,\n",
|
||||
" \"repo\": repo,\n",
|
||||
" \"github_token\": GITHUB_TOKEN\n",
|
||||
"})\n",
|
||||
"\n",
|
||||
"# Display the result\n",
|
||||
"display(Markdown(f\"### Results\\n{response}\"))"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vmQmHoy-5cAh"
|
||||
},
|
||||
"source": [
|
||||
"# Notion Workspace Assistant with PraisonAIAgents"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "WOow52NM5dmh"
|
||||
},
|
||||
"source": [
|
||||
"This notebook demonstrates how to interact with your Notion workspace using PraisonAIAgents in Google Colab. You can read, search, and update Notion pages with natural language queries, leveraging agentic tools for seamless document management—all within a notebook environment."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "TlnFAv9P_lGL"
|
||||
},
|
||||
"source": [
|
||||
"[](https://colab.research.google.com/github/Dhivya-Bharathy/PraisonAI/blob/main/examples/cookbooks/Programming_Code_Analysis_Agents/Notion_Workspace_Assistant.ipynb)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "zMn7hUGZ5hn5"
|
||||
},
|
||||
"source": [
|
||||
"# Dependencies"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "U6X3lKAu6iU3"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install praisonaiagents openai requests --quiet"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4_96-D7z6lrC"
|
||||
},
|
||||
"source": [
|
||||
"# Set Up Key"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"id": "kltUFVIW6rEO"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"# Set your Notion and OpenAI API keys here\n",
|
||||
"NOTION_API_KEY = \"ntn_..\" # <-- Replace with your actual Notion integration token\n",
|
||||
"OPENAI_API_KEY = \"sk-proj-_..\" # <-- Replace with your actual OpenAI API key\n",
|
||||
"\n",
|
||||
"# Set environment variables for downstream libraries\n",
|
||||
"os.environ[\"NOTION_API_KEY\"] = NOTION_API_KEY\n",
|
||||
"os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "FOLTBEXu635I"
|
||||
},
|
||||
"source": [
|
||||
"# Tools"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"id": "fLWxSASM5jlU"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from praisonaiagents import Agent\n",
|
||||
"import requests\n",
|
||||
"from IPython.display import display, Markdown"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KN_Sxha37EQ0"
|
||||
},
|
||||
"source": [
|
||||
"# Custom Tool"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"id": "WWxcfO9L7Fsu"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def notion_read_page_tool(page_id: str, notion_token: str) -> str:\n",
|
||||
" \"\"\"\n",
|
||||
" Custom tool to read the content of a Notion page using the Notion API.\n",
|
||||
" \"\"\"\n",
|
||||
" url = f\"https://api.notion.com/v1/blocks/{page_id}/children\"\n",
|
||||
" headers = {\n",
|
||||
" \"Authorization\": f\"Bearer {notion_token}\",\n",
|
||||
" \"Notion-Version\": \"2022-06-28\"\n",
|
||||
" }\n",
|
||||
" response = requests.get(url, headers=headers)\n",
|
||||
" if response.ok:\n",
|
||||
" data = response.json()\n",
|
||||
" blocks = data.get(\"results\", [])\n",
|
||||
" content = []\n",
|
||||
" for block in blocks:\n",
|
||||
" block_type = block.get(\"type\", \"\")\n",
|
||||
" text = \"\"\n",
|
||||
" if block_type == \"paragraph\":\n",
|
||||
" text = block.get(\"paragraph\", {}).get(\"text\", [])\n",
|
||||
" if text:\n",
|
||||
" text = text[0].get(\"plain_text\", \"\")\n",
|
||||
" elif block_type == \"heading_1\":\n",
|
||||
" text = block.get(\"heading_1\", {}).get(\"text\", [])\n",
|
||||
" if text:\n",
|
||||
" text = \"# \" + text[0].get(\"plain_text\", \"\")\n",
|
||||
" elif block_type == \"heading_2\":\n",
|
||||
" text = block.get(\"heading_2\", {}).get(\"text\", [])\n",
|
||||
" if text:\n",
|
||||
" text = \"## \" + text[0].get(\"plain_text\", \"\")\n",
|
||||
" elif block_type == \"heading_3\":\n",
|
||||
" text = block.get(\"heading_3\", {}).get(\"text\", [])\n",
|
||||
" if text:\n",
|
||||
" text = \"### \" + text[0].get(\"plain_text\", \"\")\n",
|
||||
" if text:\n",
|
||||
" content.append(text)\n",
|
||||
" return \"\\n\".join(content) if content else \"No readable content found on this page.\"\n",
|
||||
" else:\n",
|
||||
" return f\"Failed to fetch page content: {response.text}\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ebuhw9z27Nae"
|
||||
},
|
||||
"source": [
|
||||
"# YAML Prompt"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "eMjDDjzO7Pf3",
|
||||
"outputId": "b935f96d-f147-44de-80ea-7f75d23f9e7b"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\n",
|
||||
"agent_instructions: |\n",
|
||||
" You are a Notion workspace assistant. Help users read and manage their Notion pages.\n",
|
||||
" - Provide clear, organized summaries of page content.\n",
|
||||
" - Use the Notion API to fetch and display information.\n",
|
||||
" - If a page cannot be read, explain the issue clearly.\n",
|
||||
"\n",
|
||||
"tools:\n",
|
||||
" - name: notion_read_page_tool\n",
|
||||
" description: >\n",
|
||||
" Reads the content of a Notion page using its page ID and the user's Notion integration token.\n",
|
||||
" parameters:\n",
|
||||
" - name: page_id\n",
|
||||
" type: string\n",
|
||||
" description: The Notion page ID to read.\n",
|
||||
" - name: notion_token\n",
|
||||
" type: string\n",
|
||||
" description: The user's Notion integration token.\n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"yaml_prompt = \"\"\"\n",
|
||||
"agent_instructions: |\n",
|
||||
" You are a Notion workspace assistant. Help users read and manage their Notion pages.\n",
|
||||
" - Provide clear, organized summaries of page content.\n",
|
||||
" - Use the Notion API to fetch and display information.\n",
|
||||
" - If a page cannot be read, explain the issue clearly.\n",
|
||||
"\n",
|
||||
"tools:\n",
|
||||
" - name: notion_read_page_tool\n",
|
||||
" description: >\n",
|
||||
" Reads the content of a Notion page using its page ID and the user's Notion integration token.\n",
|
||||
" parameters:\n",
|
||||
" - name: page_id\n",
|
||||
" type: string\n",
|
||||
" description: The Notion page ID to read.\n",
|
||||
" - name: notion_token\n",
|
||||
" type: string\n",
|
||||
" description: The user's Notion integration token.\n",
|
||||
"\"\"\"\n",
|
||||
"print(yaml_prompt)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1ozm6Z467WQK"
|
||||
},
|
||||
"source": [
|
||||
"# Main"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 15,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 83
|
||||
},
|
||||
"id": "eNg-jwpx7X5s",
|
||||
"outputId": "4440b52a-46b3-4683-9183-13cd21596e25"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"### Notion Page Content\n",
|
||||
"*No content found or page is not accessible. Please check your page ID and sharing settings.*"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Define a wrapper to match the agent's expected tool signature\n",
|
||||
"def notion_tool_wrapper(inputs):\n",
|
||||
" page_id = inputs.get(\"page_id\", \"\")\n",
|
||||
" notion_token = inputs.get(\"notion_token\", \"\")\n",
|
||||
" return notion_read_page_tool(page_id, notion_token)\n",
|
||||
"\n",
|
||||
"# Create the agent and register the tool\n",
|
||||
"agent = Agent(\n",
|
||||
" instructions=\"\"\"\n",
|
||||
" You are a Notion workspace assistant. Help users read and manage their Notion pages.\n",
|
||||
" - Provide clear, organized summaries of page content.\n",
|
||||
" - Use the Notion API to fetch and display information.\n",
|
||||
" - If a page cannot be read, explain the issue clearly.\n",
|
||||
" \"\"\",\n",
|
||||
" tools={\n",
|
||||
" \"notion_read_page_tool\": notion_tool_wrapper\n",
|
||||
" }\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Call the tool directly (since agent.chat expects an LLM, but we want tool output)\n",
|
||||
"response = agent.execute_tool(\"notion_read_page_tool\", {\n",
|
||||
" \"page_id\": page_id,\n",
|
||||
" \"notion_token\": NOTION_API_KEY\n",
|
||||
"})\n",
|
||||
"\n",
|
||||
"# Check for error pattern and display a friendly message\n",
|
||||
"if response.startswith(\"Failed to fetch page content:\"):\n",
|
||||
" display(Markdown(\"### Notion Page Content\\n*No content found or page is not accessible. Please check your page ID and sharing settings.*\"))\n",
|
||||
"else:\n",
|
||||
" display(Markdown(f\"### Notion Page Content\\n{response}\"))"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Fixed example demonstrating proper guardrail usage with PraisonAI Agents.
|
||||
This addresses the issues reported in issue #875.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, Task, GuardrailResult, PraisonAIAgents, TaskOutput
|
||||
from typing import Tuple, Any
|
||||
import trafilatura
|
||||
|
||||
# Example 1: Using GuardrailResult return type (now supported!)
|
||||
def validate_length_guardrailresult(output: TaskOutput) -> GuardrailResult:
|
||||
"""Ensure output is between 100-500 characters using GuardrailResult"""
|
||||
# Extract the raw text from the TaskOutput object
|
||||
text = output.raw if hasattr(output, 'raw') else str(output)
|
||||
length = len(text)
|
||||
|
||||
if 100 <= length <= 500:
|
||||
return GuardrailResult(
|
||||
success=True,
|
||||
result=output, # Pass through the original output
|
||||
error=""
|
||||
)
|
||||
else:
|
||||
return GuardrailResult(
|
||||
success=False,
|
||||
result=None,
|
||||
error=f"Output must be 100-500 chars, got {length}"
|
||||
)
|
||||
|
||||
# Example 2: Using Tuple[bool, Any] return type (original method)
|
||||
def validate_length_tuple(output: TaskOutput) -> Tuple[bool, Any]:
|
||||
"""Ensure output is between 100-500 characters using tuple"""
|
||||
text = output.raw if hasattr(output, 'raw') else str(output)
|
||||
length = len(text)
|
||||
|
||||
if 100 <= length <= 500:
|
||||
return True, output
|
||||
else:
|
||||
return False, f"Output must be 100-500 chars, got {length}"
|
||||
|
||||
# Tool function
|
||||
def get_url_context(url):
|
||||
"""Fetch and extract content from a URL"""
|
||||
downloaded = trafilatura.fetch_url(url)
|
||||
if not downloaded:
|
||||
return "Sorry, I couldn't fetch the content from that URL."
|
||||
|
||||
extracted = trafilatura.extract(
|
||||
downloaded,
|
||||
include_comments=False,
|
||||
include_links=True,
|
||||
output_format='json',
|
||||
with_metadata=True,
|
||||
url=url
|
||||
)
|
||||
|
||||
if not extracted:
|
||||
return "Sorry, I couldn't extract readable content from that page."
|
||||
|
||||
return extracted # returns JSON string
|
||||
|
||||
# Create agent with FIXED tools parameter (must be a list!)
|
||||
agent = Agent(
|
||||
name="Content Summarizer",
|
||||
role="Content Analysis Expert",
|
||||
goal="Summarize web content concisely",
|
||||
instructions="You are a helpful assistant that summarizes web content",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
self_reflect=False,
|
||||
verbose=True,
|
||||
tools=[get_url_context] # FIX: tools must be a list, not a single function
|
||||
)
|
||||
|
||||
# Create task with GuardrailResult guardrail
|
||||
task_with_guardrailresult = Task(
|
||||
name="summarise article with GuardrailResult",
|
||||
description="get the context of this url: https://blog.google/technology/ai/dolphingemma/ and produce a summary below 500 characters",
|
||||
agent=agent,
|
||||
guardrail=validate_length_guardrailresult, # Using GuardrailResult
|
||||
expected_output="summary of the article below 500 characters",
|
||||
max_retries=3 # Will retry up to 3 times if guardrail fails
|
||||
)
|
||||
|
||||
# Alternative: Create task with tuple guardrail
|
||||
task_with_tuple = Task(
|
||||
name="summarise article with tuple",
|
||||
description="get the context of this url: https://blog.google/technology/ai/dolphingemma/ and produce a summary below 500 characters",
|
||||
agent=agent,
|
||||
guardrail=validate_length_tuple, # Using Tuple[bool, Any]
|
||||
expected_output="summary of the article below 500 characters",
|
||||
max_retries=3
|
||||
)
|
||||
|
||||
# Example with string-based LLM guardrail
|
||||
task_with_llm_guardrail = Task(
|
||||
name="summarise with LLM guardrail",
|
||||
description="get the context of this url: https://blog.google/technology/ai/dolphingemma/ and produce a summary",
|
||||
agent=agent,
|
||||
guardrail="Ensure the summary is professional, factual, and between 100-500 characters",
|
||||
expected_output="professional summary of the article"
|
||||
)
|
||||
|
||||
# Run with GuardrailResult example
|
||||
print("=== Running with GuardrailResult guardrail ===")
|
||||
agents_gr = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task_with_guardrailresult]
|
||||
)
|
||||
|
||||
# Uncomment to run:
|
||||
# result_gr = agents_gr.start()
|
||||
|
||||
# Run with Tuple example
|
||||
print("\n=== Running with Tuple[bool, Any] guardrail ===")
|
||||
agents_tuple = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task_with_tuple]
|
||||
)
|
||||
|
||||
# Uncomment to run:
|
||||
# result_tuple = agents_tuple.start()
|
||||
|
||||
# Run with LLM guardrail example
|
||||
print("\n=== Running with LLM-based guardrail ===")
|
||||
agents_llm = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task_with_llm_guardrail]
|
||||
)
|
||||
|
||||
# Uncomment to run:
|
||||
# result_llm = agents_llm.start()
|
||||
|
||||
print("""
|
||||
Key fixes applied:
|
||||
1. GuardrailResult is now accepted as a valid return type annotation
|
||||
2. tools parameter must be a list: tools=[get_url_context] not tools=get_url_context
|
||||
3. Both GuardrailResult and Tuple[bool, Any] return types are supported
|
||||
4. String-based LLM guardrails are also supported
|
||||
|
||||
The guardrail will automatically retry (up to max_retries times) if validation fails.
|
||||
""")
|
||||
@@ -0,0 +1,259 @@
|
||||
# Multi-Provider/Multi-Model Support
|
||||
|
||||
PraisonAI now supports intelligent multi-provider and multi-model capabilities, allowing agents to automatically select the most appropriate AI model based on task requirements, cost considerations, and performance needs.
|
||||
|
||||
## 🚀 Key Features
|
||||
|
||||
### 1. **Automatic Model Selection**
|
||||
- Analyzes task complexity and selects the best model
|
||||
- Considers cost vs. performance trade-offs
|
||||
- Supports fallback mechanisms for reliability
|
||||
|
||||
### 2. **Cost Optimization**
|
||||
- Routes simple tasks to cheaper models (e.g., GPT-4o-mini, Gemini Flash)
|
||||
- Reserves expensive models for complex tasks
|
||||
- Tracks usage and provides cost estimates
|
||||
|
||||
### 3. **Multi-Provider Support**
|
||||
- Works with OpenAI, Anthropic, Google, Groq, and more
|
||||
- Seamless switching between providers
|
||||
- Provider preference settings
|
||||
|
||||
### 4. **Flexible Routing Strategies**
|
||||
- `auto`: Automatic selection based on task analysis
|
||||
- `cost-optimized`: Prioritize cheaper models
|
||||
- `performance-optimized`: Prioritize capability
|
||||
- `manual`: Use specified model
|
||||
|
||||
## 📦 Installation
|
||||
|
||||
The multi-provider support is included in the standard installation:
|
||||
|
||||
```bash
|
||||
pip install praisonaiagents
|
||||
```
|
||||
|
||||
## 🔧 Basic Usage
|
||||
|
||||
### Simple Multi-Model Agent
|
||||
|
||||
```python
|
||||
from praisonaiagents import Task, PraisonAIAgents
|
||||
from praisonaiagents.agent import RouterAgent
|
||||
|
||||
# Create a multi-model agent
|
||||
agent = RouterAgent(
|
||||
name="Smart Assistant",
|
||||
role="Adaptive AI Assistant",
|
||||
goal="Complete tasks using the most appropriate model",
|
||||
models=["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet-20241022"],
|
||||
routing_strategy="auto" # Automatic model selection
|
||||
)
|
||||
|
||||
# Create tasks
|
||||
simple_task = Task(
|
||||
name="calculate",
|
||||
description="What is 15% of 250?",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
complex_task = Task(
|
||||
name="analyze",
|
||||
description="Write a Python implementation of the A* pathfinding algorithm",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
# Run tasks - agent will automatically select appropriate models
|
||||
agents = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[simple_task, complex_task]
|
||||
)
|
||||
|
||||
results = agents.start()
|
||||
```
|
||||
|
||||
### Cost-Optimized Workflow
|
||||
|
||||
```python
|
||||
from praisonaiagents.llm import ModelRouter
|
||||
|
||||
# Create custom router with cost constraints
|
||||
router = ModelRouter(
|
||||
cost_threshold=0.005, # Max $0.005 per 1k tokens
|
||||
preferred_providers=["google", "openai"]
|
||||
)
|
||||
|
||||
# Create cost-conscious agent
|
||||
analyzer = RouterAgent(
|
||||
name="Budget Analyzer",
|
||||
role="Data Analyst",
|
||||
goal="Analyze data efficiently",
|
||||
models={
|
||||
"gemini/gemini-1.5-flash": {},
|
||||
"gpt-4o-mini": {},
|
||||
"claude-3-haiku-20240307": {}
|
||||
},
|
||||
model_router=router,
|
||||
routing_strategy="cost-optimized"
|
||||
)
|
||||
```
|
||||
|
||||
## 🎯 Routing Strategies
|
||||
|
||||
### Auto Routing (Default)
|
||||
Automatically selects models based on:
|
||||
- Task complexity analysis
|
||||
- Required capabilities (tools, vision, etc.)
|
||||
- Context size requirements
|
||||
- Cost/performance balance
|
||||
|
||||
### Cost-Optimized
|
||||
Prioritizes cheaper models while meeting task requirements:
|
||||
- Uses lightweight models for simple tasks
|
||||
- Only escalates to expensive models when necessary
|
||||
- Ideal for high-volume operations
|
||||
|
||||
### Performance-Optimized
|
||||
Prioritizes model capability:
|
||||
- Uses the best available model for the task
|
||||
- Ideal for quality-critical applications
|
||||
- Less concern for cost
|
||||
|
||||
### Manual
|
||||
Uses the specified model without routing:
|
||||
- Direct control over model selection
|
||||
- Useful for testing or specific requirements
|
||||
|
||||
## 📊 Model Profiles
|
||||
|
||||
The system includes pre-configured profiles for popular models:
|
||||
|
||||
| Model | Provider | Best For | Cost/1k tokens |
|
||||
|-------|----------|----------|----------------|
|
||||
| gpt-4o-mini | OpenAI | Simple tasks, speed | $0.00075 |
|
||||
| gemini-1.5-flash | Google | Cost-effective, multimodal | $0.000125 |
|
||||
| claude-3-haiku | Anthropic | Fast responses | $0.0008 |
|
||||
| gpt-4o | OpenAI | General purpose | $0.0075 |
|
||||
| claude-3-5-sonnet-20241022 | Anthropic | Complex reasoning | $0.009 |
|
||||
| deepseek-chat | DeepSeek | Code & math | $0.0014 |
|
||||
|
||||
## 🔍 Task Complexity Analysis
|
||||
|
||||
The system analyzes tasks to determine complexity:
|
||||
|
||||
- **Simple**: Basic calculations, definitions, yes/no questions
|
||||
- **Moderate**: Summarization, basic analysis, classification
|
||||
- **Complex**: Code generation, algorithm implementation
|
||||
- **Very Complex**: Multi-step reasoning, system design
|
||||
|
||||
## 📈 Usage Tracking
|
||||
|
||||
Track model usage and costs:
|
||||
|
||||
```python
|
||||
# Get usage report
|
||||
report = agent.get_usage_report()
|
||||
print(report)
|
||||
|
||||
# Output:
|
||||
{
|
||||
'agent_name': 'Smart Assistant',
|
||||
'routing_strategy': 'auto',
|
||||
'model_usage': {
|
||||
'gpt-4o-mini': {'calls': 5, 'tokens': 1500, 'cost': 0.0011},
|
||||
'gpt-4o': {'calls': 2, 'tokens': 3000, 'cost': 0.0225}
|
||||
},
|
||||
'total_cost_estimate': 0.0236,
|
||||
'total_calls': 7
|
||||
}
|
||||
```
|
||||
|
||||
## 🛠️ Custom Model Configuration
|
||||
|
||||
Add custom model profiles:
|
||||
|
||||
```python
|
||||
from praisonaiagents.llm import ModelProfile, TaskComplexity
|
||||
|
||||
custom_model = ModelProfile(
|
||||
name="custom-model",
|
||||
provider="custom",
|
||||
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE),
|
||||
cost_per_1k_tokens=0.002,
|
||||
strengths=["domain-specific"],
|
||||
capabilities=["text"],
|
||||
context_window=32000
|
||||
)
|
||||
|
||||
router = ModelRouter(models=[custom_model])
|
||||
```
|
||||
|
||||
## 🔗 Integration with AutoAgents
|
||||
|
||||
Works seamlessly with AutoAgents:
|
||||
|
||||
```python
|
||||
from praisonaiagents.agents import AutoAgents
|
||||
|
||||
# Create auto agents
|
||||
auto = AutoAgents(
|
||||
instructions="Analyze market trends and create a report",
|
||||
max_agents=3
|
||||
)
|
||||
|
||||
# Convert to multi-model agents
|
||||
for i, agent in enumerate(auto.agents):
|
||||
auto.agents[i] = RouterAgent(
|
||||
name=agent.name,
|
||||
role=agent.role,
|
||||
goal=agent.goal,
|
||||
models=["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet"],
|
||||
routing_strategy="auto"
|
||||
)
|
||||
|
||||
results = auto.start()
|
||||
```
|
||||
|
||||
## 🌟 Best Practices
|
||||
|
||||
1. **Start with Auto Routing**: Let the system learn your needs
|
||||
2. **Monitor Costs**: Use usage reports to optimize
|
||||
3. **Set Cost Thresholds**: Prevent unexpected expenses
|
||||
4. **Use Appropriate Models**: Don't use GPT-4 for simple math
|
||||
5. **Leverage Fallbacks**: Configure fallback models for reliability
|
||||
|
||||
## 🔐 Environment Variables
|
||||
|
||||
Set API keys for each provider:
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY="sk-..."
|
||||
export ANTHROPIC_API_KEY="sk-ant-..."
|
||||
export GEMINI_API_KEY="..."
|
||||
export GROQ_API_KEY="..."
|
||||
```
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
The multi-provider system is extensible. To add new models:
|
||||
|
||||
1. Add model profile to `ModelRouter.DEFAULT_MODELS`
|
||||
2. Ensure the provider is supported by LiteLLM
|
||||
3. Test with various task complexities
|
||||
|
||||
## 📚 Examples
|
||||
|
||||
See the complete example in `examples/python/agents/multi-provider-agent.py` for:
|
||||
- Auto-routing examples
|
||||
- Cost-optimized workflows
|
||||
- Performance-optimized agents
|
||||
- Custom routing logic
|
||||
- Integration patterns
|
||||
|
||||
## 🎉 Benefits
|
||||
|
||||
- **Cost Savings**: Reduce API costs by 50-80%
|
||||
- **Better Performance**: Use the right tool for each job
|
||||
- **Flexibility**: Switch providers without code changes
|
||||
- **Reliability**: Automatic fallbacks and error handling
|
||||
- **Transparency**: Track usage and costs
|
||||
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
Multi-Provider/Multi-Model Agent Example
|
||||
|
||||
This example demonstrates how to use multiple AI providers/models with intelligent
|
||||
agent-based selection for cost optimization and performance.
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
from praisonaiagents.agent import RouterAgent
|
||||
from praisonaiagents.llm.model_router import ModelRouter, ModelProfile, TaskComplexity
|
||||
|
||||
# Example 1: Simple Multi-Model Agent with Auto-Routing
|
||||
def example_auto_routing():
|
||||
"""Example of automatic model routing based on task complexity"""
|
||||
print("\n=== Example 1: Auto-Routing Multi-Model Agent ===\n")
|
||||
|
||||
# Create a router agent that automatically selects models
|
||||
research_agent = RouterAgent(
|
||||
name="Smart Researcher",
|
||||
role="Adaptive Research Assistant",
|
||||
goal="Research topics using the most appropriate AI model",
|
||||
backstory="I analyze task complexity and route to the best model",
|
||||
models=["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet-20241022"],
|
||||
routing_strategy="auto", # Automatic model selection
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Create tasks with different complexity levels
|
||||
simple_task = Task(
|
||||
name="simple_calculation",
|
||||
description="Calculate the sum of 1542 and 2873",
|
||||
expected_output="The numerical result",
|
||||
agent=research_agent
|
||||
)
|
||||
|
||||
moderate_task = Task(
|
||||
name="summarize_article",
|
||||
description="Summarize the key points about renewable energy trends",
|
||||
expected_output="A concise summary with main points",
|
||||
agent=research_agent
|
||||
)
|
||||
|
||||
complex_task = Task(
|
||||
name="analyze_code",
|
||||
description="Implement a Python function to find the longest palindromic substring in a given string with O(n^2) complexity",
|
||||
expected_output="Complete Python implementation with explanation",
|
||||
agent=research_agent
|
||||
)
|
||||
|
||||
# Run the tasks
|
||||
agents = PraisonAIAgents(
|
||||
agents=[research_agent],
|
||||
tasks=[simple_task, moderate_task, complex_task],
|
||||
process="sequential",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
results = agents.start()
|
||||
|
||||
# Show usage report
|
||||
print("\n=== Model Usage Report ===")
|
||||
print(research_agent.get_usage_report())
|
||||
|
||||
|
||||
# Example 2: Cost-Optimized Workflow with Routing
|
||||
def example_cost_optimized_workflow():
|
||||
"""Example of cost-optimized workflow with model routing"""
|
||||
print("\n=== Example 2: Cost-Optimized Workflow ===\n")
|
||||
|
||||
# Create custom model router with cost threshold
|
||||
cost_router = ModelRouter(
|
||||
cost_threshold=0.01, # Max $0.01 per 1k tokens
|
||||
preferred_providers=["google", "openai", "anthropic"] # Provider preference
|
||||
)
|
||||
|
||||
# Create specialized agents with different model strategies
|
||||
analyzer = RouterAgent(
|
||||
name="Cost-Conscious Analyzer",
|
||||
role="Data Analyzer",
|
||||
goal="Analyze data efficiently while minimizing costs",
|
||||
models={
|
||||
"gemini/gemini-1.5-flash": {},
|
||||
"gpt-4o-mini": {},
|
||||
"deepseek-chat": {}
|
||||
},
|
||||
model_router=cost_router,
|
||||
routing_strategy="cost-optimized",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
writer = RouterAgent(
|
||||
name="Quality Writer",
|
||||
role="Content Writer",
|
||||
goal="Create high-quality content",
|
||||
models={
|
||||
"claude-3-5-sonnet-20241022": {},
|
||||
"gpt-4o": {}
|
||||
},
|
||||
routing_strategy="performance-optimized", # Prefer better models
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Create workflow tasks
|
||||
analysis_task = Task(
|
||||
name="analyze_data",
|
||||
description="Analyze this dataset: [1, 2, 3, 4, 5]. Find mean, median, and mode.",
|
||||
expected_output="Statistical analysis results",
|
||||
agent=analyzer
|
||||
)
|
||||
|
||||
writing_task = Task(
|
||||
name="write_report",
|
||||
description="Write a professional report based on the analysis",
|
||||
expected_output="Professional report with insights",
|
||||
agent=writer,
|
||||
context=[analysis_task] # Use results from analysis
|
||||
)
|
||||
|
||||
# Run workflow
|
||||
workflow = PraisonAIAgents(
|
||||
agents=[analyzer, writer],
|
||||
tasks=[analysis_task, writing_task],
|
||||
process="sequential",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
results = workflow.start()
|
||||
|
||||
# Show cost comparison
|
||||
print("\n=== Cost Analysis ===")
|
||||
print(f"Analyzer costs: {analyzer.get_usage_report()}")
|
||||
print(f"Writer costs: {writer.get_usage_report()}")
|
||||
|
||||
|
||||
# Example 3: Auto Agent Mode with Multi-Provider Support
|
||||
def example_auto_agents_multi_provider():
|
||||
"""Example using AutoAgents with multi-provider support"""
|
||||
print("\n=== Example 3: AutoAgents with Multi-Provider Support ===\n")
|
||||
|
||||
from praisonaiagents.agents import AutoAgents
|
||||
|
||||
# Create AutoAgents that will automatically assign appropriate models
|
||||
auto_agents = AutoAgents(
|
||||
instructions="Create a market research report on electric vehicles. Include data analysis, competitor analysis, and future projections.",
|
||||
max_agents=3,
|
||||
llm="gpt-4o-mini", # Default model for agent generation
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# After agents are created, upgrade them to multi-model agents
|
||||
multi_model_agents = []
|
||||
for agent in auto_agents.agents:
|
||||
# Convert regular agents to multi-model agents
|
||||
multi_agent = RouterAgent(
|
||||
name=agent.name,
|
||||
role=agent.role,
|
||||
goal=agent.goal,
|
||||
backstory=agent.backstory,
|
||||
tools=agent.tools,
|
||||
models=["gpt-4o-mini", "gemini/gemini-1.5-flash", "claude-3-haiku-20240307", "gpt-4o"],
|
||||
routing_strategy="auto",
|
||||
verbose=True
|
||||
)
|
||||
multi_model_agents.append(multi_agent)
|
||||
|
||||
# Update the agents in the AutoAgents instance
|
||||
auto_agents.agents = multi_model_agents
|
||||
|
||||
# Run the auto-generated workflow
|
||||
results = auto_agents.start()
|
||||
|
||||
# Show model usage across all agents
|
||||
print("\n=== Multi-Provider Usage Summary ===")
|
||||
total_calls = 0
|
||||
for agent in multi_model_agents:
|
||||
report = agent.get_usage_report()
|
||||
print(f"\nAgent: {agent.name}")
|
||||
for model, stats in report['model_usage'].items():
|
||||
if stats['calls'] > 0:
|
||||
print(f" - {model}: {stats['calls']} calls")
|
||||
total_calls += stats['calls']
|
||||
|
||||
print(f"\nTotal API calls: {total_calls}")
|
||||
|
||||
|
||||
# Example 4: Custom Routing Logic
|
||||
def example_custom_routing():
|
||||
"""Example with custom routing logic based on specific requirements"""
|
||||
print("\n=== Example 4: Custom Routing Logic ===\n")
|
||||
|
||||
# Create custom model profiles for specific use cases
|
||||
custom_models = [
|
||||
ModelProfile(
|
||||
name="gpt-4o",
|
||||
provider="openai",
|
||||
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.VERY_COMPLEX),
|
||||
cost_per_1k_tokens=0.0075,
|
||||
strengths=["code-generation", "debugging"],
|
||||
capabilities=["text", "function-calling"],
|
||||
context_window=128000
|
||||
),
|
||||
ModelProfile(
|
||||
name="deepseek-chat",
|
||||
provider="deepseek",
|
||||
complexity_range=(TaskComplexity.MODERATE, TaskComplexity.VERY_COMPLEX),
|
||||
cost_per_1k_tokens=0.001,
|
||||
strengths=["mathematics", "algorithms"],
|
||||
capabilities=["text", "function-calling"],
|
||||
context_window=128000
|
||||
)
|
||||
]
|
||||
|
||||
# Create router with custom models
|
||||
custom_router = ModelRouter(models=custom_models)
|
||||
|
||||
# Tool for web search (example)
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information"""
|
||||
return f"Search results for: {query}"
|
||||
|
||||
# Create specialized agent
|
||||
coder = RouterAgent(
|
||||
name="Adaptive Coder",
|
||||
role="Software Developer",
|
||||
goal="Write and debug code using the best model for each task",
|
||||
backstory="I'm an expert coder who knows when to use different AI models",
|
||||
model_router=custom_router,
|
||||
routing_strategy="auto",
|
||||
tools=[search_web],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Create coding tasks
|
||||
tasks = [
|
||||
Task(
|
||||
name="fix_bug",
|
||||
description="Fix this Python bug: 'list' object has no attribute 'appendx'",
|
||||
expected_output="Corrected code with explanation",
|
||||
agent=coder
|
||||
),
|
||||
Task(
|
||||
name="implement_algorithm",
|
||||
description="Implement Dijkstra's shortest path algorithm in Python",
|
||||
expected_output="Complete implementation with comments",
|
||||
agent=coder,
|
||||
tools=[search_web]
|
||||
)
|
||||
]
|
||||
|
||||
# Run tasks
|
||||
agents = PraisonAIAgents(
|
||||
agents=[coder],
|
||||
tasks=tasks,
|
||||
process="sequential",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
results = agents.start()
|
||||
|
||||
print("\n=== Custom Routing Results ===")
|
||||
print(coder.get_usage_report())
|
||||
|
||||
|
||||
# Run examples
|
||||
if __name__ == "__main__":
|
||||
# Make sure to set your API keys
|
||||
# os.environ["OPENAI_API_KEY"] = "your-openai-key"
|
||||
# os.environ["ANTHROPIC_API_KEY"] = "your-anthropic-key"
|
||||
# os.environ["GEMINI_API_KEY"] = "your-gemini-key"
|
||||
|
||||
# Run examples
|
||||
example_auto_routing()
|
||||
example_cost_optimized_workflow()
|
||||
example_auto_agents_multi_provider()
|
||||
example_custom_routing()
|
||||
|
||||
print("\n=== All Examples Completed ===")
|
||||
print("\nKey Features Demonstrated:")
|
||||
print("1. Automatic model selection based on task complexity")
|
||||
print("2. Cost-optimized routing for budget-conscious operations")
|
||||
print("3. Performance-optimized routing for quality-critical tasks")
|
||||
print("4. Multi-provider support with fallback mechanisms")
|
||||
print("5. Integration with AutoAgents for automatic workflow generation")
|
||||
print("6. Custom routing logic for specialized use cases")
|
||||
@@ -0,0 +1,197 @@
|
||||
"""
|
||||
Hosted MCP Server implementation for PraisonAI Agents.
|
||||
|
||||
This module provides a base class for creating hosted MCP servers
|
||||
that can handle requests and integrate with the MCP protocol.
|
||||
|
||||
Note: This is an example implementation. To use it, ensure you have installed:
|
||||
pip install praisonaiagents[mcp] starlette>=0.27.0
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any, Optional, List, Callable
|
||||
import json
|
||||
|
||||
try:
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server import Server
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
from starlette.routing import Mount, Route
|
||||
from mcp.server.sse import SseServerTransport
|
||||
import uvicorn
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"MCP server dependencies not installed. "
|
||||
"Please install with: pip install praisonaiagents[mcp] starlette>=0.27.0"
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HostedMCPServer:
|
||||
"""
|
||||
Base class for creating hosted MCP servers.
|
||||
|
||||
This class provides a foundation for building MCP servers that can:
|
||||
- Handle incoming requests
|
||||
- Define custom tools
|
||||
- Support SSE transport
|
||||
- Be extended with custom functionality like latency tracking
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "hosted-mcp-server", host: str = "localhost", port: int = 8080):
|
||||
"""
|
||||
Initialize the hosted MCP server.
|
||||
|
||||
Args:
|
||||
name: Server name for identification
|
||||
host: Host to bind to (default: localhost)
|
||||
port: Port to listen on (default: 8080)
|
||||
"""
|
||||
self.name = name
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.mcp = FastMCP(name)
|
||||
self._tools: Dict[str, Callable] = {}
|
||||
self._server: Optional[Server] = None
|
||||
self._app: Optional[Starlette] = None
|
||||
|
||||
def handle_request(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Handle incoming MCP requests.
|
||||
|
||||
This method can be overridden in subclasses to add custom request handling,
|
||||
such as latency tracking, authentication, or request modification.
|
||||
|
||||
Args:
|
||||
request_data: The incoming request data
|
||||
|
||||
Returns:
|
||||
Response data
|
||||
"""
|
||||
# Default implementation - can be overridden
|
||||
method = request_data.get('method', '')
|
||||
request_id = request_data.get('id', 'unknown')
|
||||
|
||||
logger.debug(f"Handling request {request_id}: {method}")
|
||||
|
||||
# Basic response structure
|
||||
response = {
|
||||
'id': request_id,
|
||||
'jsonrpc': '2.0',
|
||||
'result': {}
|
||||
}
|
||||
|
||||
return response
|
||||
|
||||
def add_tool(self, func: Callable, name: Optional[str] = None, description: Optional[str] = None):
|
||||
"""
|
||||
Add a tool to the MCP server.
|
||||
|
||||
Args:
|
||||
func: The function to expose as a tool
|
||||
name: Optional name for the tool (defaults to function name)
|
||||
description: Optional description for the tool
|
||||
"""
|
||||
tool_name = name or func.__name__
|
||||
|
||||
# Register with FastMCP
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
# Already async
|
||||
self.mcp.tool(name=tool_name)(func)
|
||||
else:
|
||||
# Wrap sync function in async
|
||||
async def async_wrapper(*args, **kwargs):
|
||||
return func(*args, **kwargs)
|
||||
async_wrapper.__name__ = func.__name__
|
||||
async_wrapper.__doc__ = description or func.__doc__
|
||||
self.mcp.tool(name=tool_name)(async_wrapper)
|
||||
|
||||
self._tools[tool_name] = func
|
||||
logger.info(f"Added tool: {tool_name}")
|
||||
|
||||
def create_app(self, debug: bool = False) -> Starlette:
|
||||
"""
|
||||
Create a Starlette application for serving the MCP server.
|
||||
|
||||
Args:
|
||||
debug: Enable debug mode
|
||||
|
||||
Returns:
|
||||
Starlette application instance
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the MCP server is not properly initialized
|
||||
"""
|
||||
if not self._server:
|
||||
if not hasattr(self.mcp, '_mcp_server'):
|
||||
raise RuntimeError("MCP server not properly initialized. Ensure FastMCP is correctly set up.")
|
||||
self._server = self.mcp._mcp_server
|
||||
|
||||
sse = SseServerTransport("/messages/")
|
||||
|
||||
async def handle_sse(request: Request) -> None:
|
||||
logger.debug(f"SSE connection from {request.client}")
|
||||
async with sse.connect_sse(
|
||||
request.scope,
|
||||
request.receive,
|
||||
request._send,
|
||||
) as (read_stream, write_stream):
|
||||
await self._server.run(
|
||||
read_stream,
|
||||
write_stream,
|
||||
self._server.create_initialization_options(),
|
||||
)
|
||||
|
||||
self._app = Starlette(
|
||||
debug=debug,
|
||||
routes=[
|
||||
Route("/sse", endpoint=handle_sse),
|
||||
Mount("/messages/", app=sse.handle_post_message),
|
||||
],
|
||||
)
|
||||
|
||||
return self._app
|
||||
|
||||
def start(self, debug: bool = False, **uvicorn_kwargs):
|
||||
"""
|
||||
Start the MCP server.
|
||||
|
||||
Args:
|
||||
debug: Enable debug mode
|
||||
**uvicorn_kwargs: Additional arguments to pass to uvicorn
|
||||
"""
|
||||
app = self.create_app(debug=debug)
|
||||
|
||||
print(f"Starting {self.name} MCP server on {self.host}:{self.port}")
|
||||
print(f"Available tools: {', '.join(self._tools.keys())}")
|
||||
print(f"SSE endpoint: http://{self.host}:{self.port}/sse")
|
||||
|
||||
uvicorn.run(app, host=self.host, port=self.port, **uvicorn_kwargs)
|
||||
|
||||
async def start_async(self, debug: bool = False):
|
||||
"""
|
||||
Start the MCP server asynchronously.
|
||||
|
||||
Args:
|
||||
debug: Enable debug mode
|
||||
"""
|
||||
app = self.create_app(debug=debug)
|
||||
|
||||
config = uvicorn.Config(app, host=self.host, port=self.port)
|
||||
server = uvicorn.Server(config)
|
||||
|
||||
print(f"Starting {self.name} MCP server on {self.host}:{self.port}")
|
||||
print(f"Available tools: {', '.join(self._tools.keys())}")
|
||||
|
||||
await server.serve()
|
||||
|
||||
def get_tools(self) -> List[str]:
|
||||
"""Get list of available tool names."""
|
||||
return list(self._tools.keys())
|
||||
|
||||
def get_endpoint(self) -> str:
|
||||
"""Get the SSE endpoint URL."""
|
||||
return f"http://{self.host}:{self.port}/sse"
|
||||
@@ -6,7 +6,7 @@ without modifying PraisonAI core files.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, PraisonAIAgents
|
||||
from praisonaiagents.mcp import HostedMCPServer
|
||||
from hosted_server import HostedMCPServer # Import from local file
|
||||
from latency_tracker_tool import tracker, get_latency_metrics
|
||||
import json
|
||||
|
||||
|
||||
@@ -72,7 +72,7 @@ workflow = PraisonAIAgents(
|
||||
)
|
||||
|
||||
# Run the workflow
|
||||
results = workflow.start()
|
||||
results = workflow.start(return_dict=True)
|
||||
|
||||
# Print results
|
||||
print("\nWorkflow Results:")
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example: Using Gemini Embedding Models with PraisonAI Memory
|
||||
|
||||
This example demonstrates how to configure PraisonAI to use Google's Gemini
|
||||
embedding models for memory storage and retrieval, as requested in issue #870.
|
||||
|
||||
Supported Gemini embedding models:
|
||||
- gemini-embedding-exp-03-07 (experimental)
|
||||
- text-embedding-004 (stable)
|
||||
|
||||
Prerequisites:
|
||||
1. Install praisonaiagents with memory support:
|
||||
pip install "praisonaiagents[memory]"
|
||||
|
||||
2. Set your Google API key:
|
||||
export GOOGLE_API_KEY='your-api-key'
|
||||
|
||||
3. Ensure litellm is installed (comes with praisonaiagents)
|
||||
"""
|
||||
|
||||
import os
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
def main():
|
||||
"""Demonstrate Gemini embedding usage"""
|
||||
|
||||
# Verify Google API key
|
||||
if not os.environ.get("GOOGLE_API_KEY"):
|
||||
print("Error: GOOGLE_API_KEY environment variable not set")
|
||||
print("Please set it with: export GOOGLE_API_KEY='your-api-key'")
|
||||
return
|
||||
|
||||
print("=== Gemini Embedding Example ===\n")
|
||||
|
||||
# Create agents for different tasks
|
||||
researcher = Agent(
|
||||
name="Research Agent",
|
||||
role="Information Researcher",
|
||||
goal="Research and store important information about topics",
|
||||
backstory="Expert at analyzing and documenting information with semantic understanding",
|
||||
llm="gpt-4o-mini" # Using GPT for chat, Gemini for embeddings
|
||||
)
|
||||
|
||||
retriever = Agent(
|
||||
name="Retrieval Agent",
|
||||
role="Information Retriever",
|
||||
goal="Retrieve relevant information from memory using semantic search",
|
||||
backstory="Specialist in finding and presenting stored information",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Task 1: Store information with semantic meaning
|
||||
store_task = Task(
|
||||
description="""
|
||||
Research and store the following information:
|
||||
1. Gemini is Google's family of multimodal AI models
|
||||
2. Gemini models support text, image, video, and audio understanding
|
||||
3. The embedding models provide semantic representations of text
|
||||
4. Gemini Ultra is the largest and most capable model
|
||||
5. Gemini was announced in December 2023
|
||||
|
||||
Store each fact with its semantic meaning preserved.
|
||||
""",
|
||||
expected_output="Confirmation that all facts have been stored in memory",
|
||||
agent=researcher
|
||||
)
|
||||
|
||||
# Task 2: Semantic retrieval test
|
||||
retrieve_task = Task(
|
||||
description="""
|
||||
Using semantic search, find information about:
|
||||
1. Google's AI capabilities
|
||||
2. Multimodal understanding features
|
||||
3. When was Google's latest AI announced
|
||||
|
||||
Note: Use semantic similarity, not exact keyword matching.
|
||||
""",
|
||||
expected_output="Retrieved information based on semantic meaning",
|
||||
agent=retriever
|
||||
)
|
||||
|
||||
# Configure with Gemini embeddings
|
||||
# Option 1: Using stable text-embedding-004
|
||||
agents = PraisonAIAgents(
|
||||
agents=[researcher, retriever],
|
||||
tasks=[store_task, retrieve_task],
|
||||
verbose=True,
|
||||
memory=True,
|
||||
embedder={
|
||||
"provider": "gemini",
|
||||
"config": {
|
||||
"model": "text-embedding-004",
|
||||
# Optional: specify task type for better results
|
||||
# "task_type": "RETRIEVAL_DOCUMENT"
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Alternative: Using experimental model with advanced features
|
||||
# agents = PraisonAIAgents(
|
||||
# agents=[researcher, retriever],
|
||||
# tasks=[store_task, retrieve_task],
|
||||
# verbose=True,
|
||||
# memory=True,
|
||||
# embedder={
|
||||
# "provider": "gemini",
|
||||
# "config": {
|
||||
# "model": "gemini-embedding-exp-03-07",
|
||||
# # Task types: SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING,
|
||||
# # RETRIEVAL_DOCUMENT, RETRIEVAL_QUERY, QUESTION_ANSWERING,
|
||||
# # FACT_VERIFICATION
|
||||
# "task_type": "RETRIEVAL_DOCUMENT"
|
||||
# }
|
||||
# }
|
||||
# )
|
||||
|
||||
# Execute the workflow
|
||||
print("Starting agents with Gemini embeddings...\n")
|
||||
result = agents.start()
|
||||
|
||||
# Demonstrate direct memory access with Gemini embeddings
|
||||
print("\n=== Direct Memory Search Example ===\n")
|
||||
|
||||
if agents.shared_memory:
|
||||
# Search queries that test semantic understanding
|
||||
queries = [
|
||||
"What are Google's latest AI innovations?",
|
||||
"Explain multimodal AI capabilities",
|
||||
"Large language model features"
|
||||
]
|
||||
|
||||
for query in queries:
|
||||
print(f"\nSearching for: '{query}'")
|
||||
results = agents.shared_memory.search_long_term(query, limit=3)
|
||||
|
||||
if results:
|
||||
print(f"Found {len(results)} semantically similar results:")
|
||||
for i, result in enumerate(results, 1):
|
||||
print(f"\n{i}. Relevance Score: {result.get('score', 0):.3f}")
|
||||
print(f" Content: {result.get('text', '')[:150]}...")
|
||||
else:
|
||||
print("No results found")
|
||||
|
||||
print("\n=== Example Complete ===")
|
||||
print("\nKey Points:")
|
||||
print("1. Gemini embeddings provide semantic understanding of text")
|
||||
print("2. Results are based on meaning, not just keyword matching")
|
||||
print("3. Different task types can optimize embedding quality")
|
||||
print("4. Compatible with LiteLLM for easy integration")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,165 @@
|
||||
# Gemini Embedding Support in PraisonAI
|
||||
|
||||
This document describes how to use Google's Gemini embedding models with PraisonAI's memory system.
|
||||
|
||||
## Overview
|
||||
|
||||
PraisonAI now supports configurable embedding models through LiteLLM, including Google's Gemini embedding models. This allows you to use models like `text-embedding-004` or `gemini-embedding-exp-03-07` for semantic search and memory storage.
|
||||
|
||||
## Configuration
|
||||
|
||||
### Default Configuration (OpenAI)
|
||||
|
||||
By default, PraisonAI uses OpenAI's `text-embedding-3-small` model:
|
||||
|
||||
```python
|
||||
from praisonaiagents import PraisonAIAgents
|
||||
|
||||
agents = PraisonAIAgents(
|
||||
agents=[...],
|
||||
tasks=[...],
|
||||
memory=True # Uses default OpenAI embeddings
|
||||
)
|
||||
```
|
||||
|
||||
### Using Gemini Embeddings
|
||||
|
||||
To use Gemini embeddings, configure the `embedder` parameter:
|
||||
|
||||
```python
|
||||
from praisonaiagents import PraisonAIAgents
|
||||
|
||||
agents = PraisonAIAgents(
|
||||
agents=[...],
|
||||
tasks=[...],
|
||||
memory=True,
|
||||
embedder={
|
||||
"provider": "gemini",
|
||||
"config": {
|
||||
"model": "text-embedding-004" # or "gemini-embedding-exp-03-07"
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
## Prerequisites
|
||||
|
||||
1. **API Key**: Set your Google API key:
|
||||
```bash
|
||||
export GOOGLE_API_KEY='your-api-key'
|
||||
```
|
||||
|
||||
2. **Dependencies**: Ensure you have the memory dependencies installed:
|
||||
```bash
|
||||
pip install "praisonaiagents[memory]"
|
||||
```
|
||||
|
||||
## Supported Models
|
||||
|
||||
### Gemini Embedding Models
|
||||
|
||||
- `text-embedding-004` - Stable embedding model
|
||||
- `gemini-embedding-exp-03-07` - Experimental model with advanced features
|
||||
|
||||
### Task Types (Optional)
|
||||
|
||||
Gemini models support task-specific optimizations:
|
||||
- `SEMANTIC_SIMILARITY`
|
||||
- `CLASSIFICATION`
|
||||
- `CLUSTERING`
|
||||
- `RETRIEVAL_DOCUMENT`
|
||||
- `RETRIEVAL_QUERY`
|
||||
- `QUESTION_ANSWERING`
|
||||
- `FACT_VERIFICATION`
|
||||
|
||||
Example with task type:
|
||||
```python
|
||||
embedder={
|
||||
"provider": "gemini",
|
||||
"config": {
|
||||
"model": "text-embedding-004",
|
||||
"task_type": "RETRIEVAL_DOCUMENT"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Complete Example
|
||||
|
||||
```python
|
||||
import os
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
# Ensure Google API key is set
|
||||
os.environ["GOOGLE_API_KEY"] = "your-api-key"
|
||||
|
||||
# Create agents
|
||||
researcher = Agent(
|
||||
name="Researcher",
|
||||
role="Information Researcher",
|
||||
goal="Store information with semantic understanding",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create task
|
||||
task = Task(
|
||||
description="Research and store information about AI advancements",
|
||||
expected_output="Confirmation of stored information",
|
||||
agent=researcher
|
||||
)
|
||||
|
||||
# Initialize with Gemini embeddings
|
||||
agents = PraisonAIAgents(
|
||||
agents=[researcher],
|
||||
tasks=[task],
|
||||
memory=True,
|
||||
embedder={
|
||||
"provider": "gemini",
|
||||
"config": {
|
||||
"model": "text-embedding-004"
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Run the agents
|
||||
agents.start()
|
||||
```
|
||||
|
||||
## Backward Compatibility
|
||||
|
||||
The implementation maintains full backward compatibility:
|
||||
|
||||
1. **No embedder config**: Uses default OpenAI `text-embedding-3-small`
|
||||
2. **Existing configs**: Continue to work without modification
|
||||
3. **Other providers**: Supports any embedding model available through LiteLLM
|
||||
|
||||
## Memory Configuration with Mem0
|
||||
|
||||
When using Mem0 as the memory provider, the embedder configuration is passed through:
|
||||
|
||||
```python
|
||||
memory_config = {
|
||||
"provider": "mem0",
|
||||
"config": {
|
||||
"api_key": "your-mem0-key",
|
||||
"embedder": {
|
||||
"provider": "gemini",
|
||||
"config": {
|
||||
"model": "text-embedding-004"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
1. **API Key Error**: Ensure `GOOGLE_API_KEY` is set in your environment
|
||||
2. **Model Not Found**: Verify the model name is correct
|
||||
3. **LiteLLM Error**: Ensure litellm is installed and up to date
|
||||
4. **Embedding Size Mismatch**: When switching models, you may need to clear the vector database as different models produce different embedding dimensions
|
||||
|
||||
## See Also
|
||||
|
||||
- [Gemini API Documentation](https://ai.google.dev/gemini-api/docs/embeddings)
|
||||
- [LiteLLM Embedding Models](https://docs.litellm.ai/docs/embedding/supported_embedding)
|
||||
- Example: `/examples/python/memory/gemini-embedding-example.py`
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a competitor intelligence AI agent. "
|
||||
"Help users analyze competitors, market trends, and business intelligence "
|
||||
"to provide strategic insights and competitive advantages.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need to analyze my main competitor's pricing strategy and market positioning. "
|
||||
"Can you help me gather insights?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a finance AI agent. "
|
||||
"Help users with financial analysis, investment strategies, "
|
||||
"budget planning, and market trend analysis for informed decision making.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need to analyze my investment portfolio and create a diversification strategy. "
|
||||
"Can you help me optimize my financial planning?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a game design AI agent. "
|
||||
"Help users with game concept development, mechanics design, "
|
||||
"and creative storytelling for various gaming platforms and genres.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I want to create a mobile puzzle game concept. "
|
||||
"Can you help me design engaging gameplay mechanics and story elements?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a legal AI agent. "
|
||||
"Help users with legal document analysis, contract review, "
|
||||
"and legal advice while ensuring compliance and risk assessment.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need to review a software licensing agreement. "
|
||||
"Can you help me identify potential risks and key terms?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a recruitment AI agent. "
|
||||
"Help users with candidate screening, job posting optimization, "
|
||||
"and recruitment process automation to find the best talent.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need to create a job posting for a senior software engineer position. "
|
||||
"Can you help me write an attractive and comprehensive job description?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are an AI services agency agent. "
|
||||
"Help users with comprehensive business services including marketing, "
|
||||
"content creation, strategy, and operational optimization using CrewAI framework.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need help with my startup's marketing strategy and content creation. "
|
||||
"Can you provide a comprehensive plan?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are an AI teaching agent. "
|
||||
"Help users with educational content creation, lesson planning, "
|
||||
"and personalized learning experiences across various subjects.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need to create a lesson plan for teaching Python programming to beginners. "
|
||||
"Can you help me design an engaging curriculum?"
|
||||
)
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a data analysis AI agent. Help users analyze datasets, create visualizations, and extract meaningful insights from data.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I have sales data for the last quarter. Can you help me identify trends and patterns?")
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a deep research AI agent. Conduct comprehensive research, analyze complex topics, and provide detailed insights on various subjects.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I need to research the impact of artificial intelligence on job markets. Can you provide a comprehensive analysis?")
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a health and fitness AI agent. Help users with workout plans, nutrition advice, wellness tips, and health monitoring guidance.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I want to lose 20 pounds and build muscle. What's the best workout and nutrition plan for me?")
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are an investment advisor AI agent. Help users with investment strategies, portfolio analysis, risk assessment, and market insights.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I have $50,000 to invest. What would be a good diversified portfolio for someone in their 30s?")
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a journalist AI agent. Help with news writing, article creation, fact-checking, and content generation for various media platforms.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I need to write an article about the latest developments in renewable energy. Can you help me structure it?")
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a medical imaging AI agent. Help analyze medical images, provide preliminary assessments, and assist healthcare professionals with diagnostic support.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I have an X-ray image that needs analysis. Can you help identify any abnormalities?")
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a music generation AI agent. Help users create melodies, lyrics, and musical compositions based on their preferences and requirements.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I need a catchy jingle for my coffee shop. Can you help me create one?")
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a personal finance AI agent. Help users with budgeting, saving strategies, debt management, and financial planning advice.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I want to save $10,000 in the next year. What's the best strategy for me?")
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a startup trend analysis AI agent. Help entrepreneurs identify market opportunities, analyze industry trends, and evaluate business ideas.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I'm thinking of starting a tech startup. What are the current trending opportunities in the market?")
|
||||
@@ -0,0 +1,8 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a travel planning AI agent. Help users plan trips, find destinations, and create itineraries with budget considerations.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start("I want to plan a 7-day trip to Europe with a budget of $3000. What are my best options?")
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a multi MCP AI agent. "
|
||||
"Help users integrate and coordinate multiple Model Context Protocol agents "
|
||||
"for complex workflows and cross-platform automation.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need to coordinate multiple tools and services for my project. "
|
||||
"Can you help me set up an integrated workflow using MCP agents?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a multimodal coding AI agent. "
|
||||
"Help users with code generation, review, and debugging "
|
||||
"by analyzing both text and visual inputs like screenshots and diagrams.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I have a screenshot of a UI design and need to generate the corresponding HTML/CSS code. "
|
||||
"Can you help me create the frontend code?"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a multimodal design AI agent. "
|
||||
"Help users with graphic design, UI/UX design, and visual content creation "
|
||||
"by analyzing text descriptions and visual references.",
|
||||
llm="xai/grok-4"
|
||||
)
|
||||
|
||||
response = agent.start(
|
||||
"I need to create a logo for my tech startup. "
|
||||
"Can you help me design a modern and professional logo concept?"
|
||||
)
|
||||
+1
-1
@@ -196,4 +196,4 @@ def main():
|
||||
print("Make sure SearxNG is running at http://localhost:32768")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
@@ -0,0 +1,420 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Custom Trafilatura Tools for High-Quality URL Content Extraction
|
||||
|
||||
This is an example of creating custom tools for PraisonAI agents.
|
||||
This tool is not part of the core package - it's a standalone custom tool
|
||||
that demonstrates how to implement web content extraction using Trafilatura.
|
||||
|
||||
Install dependencies:
|
||||
pip install trafilatura newspaper3k
|
||||
|
||||
Usage:
|
||||
from custom_trafilatura_tools import TrafilaturaTools
|
||||
|
||||
tools = TrafilaturaTools()
|
||||
content = tools.extract_content("https://example.com/article")
|
||||
metadata = tools.extract_metadata("https://example.com/article")
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, Any, Optional, Union, List
|
||||
from importlib import util
|
||||
import json
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
class TrafilaturaTools:
|
||||
"""Custom tools for extracting high-quality content from URLs using Trafilatura."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize TrafilaturaTools and check for required packages."""
|
||||
self.trafilatura = self._check_and_import_trafilatura()
|
||||
|
||||
def _check_and_import_trafilatura(self):
|
||||
"""Check if trafilatura package is installed and import it."""
|
||||
if util.find_spec("trafilatura") is None:
|
||||
raise ImportError(
|
||||
"trafilatura package is not available. Please install it using: pip install trafilatura"
|
||||
)
|
||||
import trafilatura
|
||||
return trafilatura
|
||||
|
||||
def _validate_url(self, url: str) -> bool:
|
||||
"""
|
||||
Validate URL to prevent SSRF attacks.
|
||||
|
||||
Args:
|
||||
url: URL to validate
|
||||
|
||||
Returns:
|
||||
bool: True if URL is safe, False otherwise
|
||||
"""
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
|
||||
# Only allow http/https protocols
|
||||
if parsed.scheme not in ['http', 'https']:
|
||||
return False
|
||||
|
||||
# Reject URLs with no hostname
|
||||
if not parsed.hostname:
|
||||
return False
|
||||
|
||||
# Reject local/internal addresses
|
||||
hostname = parsed.hostname.lower()
|
||||
|
||||
# Block localhost and loopback
|
||||
if hostname in ['localhost', '127.0.0.1', '0.0.0.0', '::1']:
|
||||
return False
|
||||
|
||||
# Block private IP ranges
|
||||
import ipaddress
|
||||
try:
|
||||
ip = ipaddress.ip_address(hostname)
|
||||
if ip.is_private or ip.is_reserved or ip.is_loopback or ip.is_link_local:
|
||||
return False
|
||||
except ValueError:
|
||||
# Not an IP address, continue with domain validation
|
||||
pass
|
||||
|
||||
# Block common internal domains
|
||||
if any(hostname.endswith(domain) for domain in ['.local', '.internal', '.localdomain']):
|
||||
return False
|
||||
|
||||
# Block metadata service endpoints
|
||||
return hostname not in ['169.254.169.254', 'metadata.google.internal']
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def extract_content(
|
||||
self,
|
||||
url: str,
|
||||
include_comments: bool = False,
|
||||
include_links: bool = True,
|
||||
output_format: str = 'json',
|
||||
include_metadata: bool = True,
|
||||
target_language: Optional[str] = None
|
||||
) -> Union[Dict[str, Any], str]:
|
||||
"""
|
||||
Extract high-quality content from a URL using Trafilatura.
|
||||
|
||||
Args:
|
||||
url: URL to extract content from
|
||||
include_comments: Include comments in extraction
|
||||
include_links: Include links in extraction
|
||||
output_format: Output format ('json', 'text', 'xml')
|
||||
include_metadata: Include metadata in extraction
|
||||
target_language: Target language for extraction (None for auto-detect)
|
||||
|
||||
Returns:
|
||||
Dict or str: Extracted content (dict for json format, str for text/xml)
|
||||
"""
|
||||
# Validate URL
|
||||
if not self._validate_url(url):
|
||||
error_msg = f"Invalid or unsafe URL: {url}"
|
||||
logging.error(error_msg)
|
||||
return {"error": error_msg} if output_format == 'json' else error_msg
|
||||
|
||||
try:
|
||||
# Fetch the URL content
|
||||
downloaded = self.trafilatura.fetch_url(url)
|
||||
if not downloaded:
|
||||
error_msg = f"Could not fetch content from URL: {url}"
|
||||
logging.error(error_msg)
|
||||
return {"error": error_msg} if output_format == 'json' else error_msg
|
||||
|
||||
# Extract content
|
||||
extracted = self.trafilatura.extract(
|
||||
downloaded,
|
||||
include_comments=include_comments,
|
||||
include_links=include_links,
|
||||
output_format=output_format,
|
||||
with_metadata=include_metadata,
|
||||
url=url,
|
||||
target_language=target_language
|
||||
)
|
||||
|
||||
if not extracted:
|
||||
error_msg = f"Could not extract readable content from: {url}"
|
||||
logging.warning(error_msg)
|
||||
return {"error": error_msg} if output_format == 'json' else error_msg
|
||||
|
||||
# Parse JSON output if requested
|
||||
if output_format == 'json':
|
||||
return json.loads(extracted)
|
||||
else:
|
||||
return extracted
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Error extracting content from {url}: {str(e)}"
|
||||
logging.error(error_msg)
|
||||
return {"error": error_msg} if output_format == 'json' else error_msg
|
||||
|
||||
def extract_metadata(self, url: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Extract metadata from a URL.
|
||||
|
||||
Args:
|
||||
url: URL to extract metadata from
|
||||
|
||||
Returns:
|
||||
Dict: Metadata including title, author, date, etc.
|
||||
"""
|
||||
# Validate URL
|
||||
if not self._validate_url(url):
|
||||
error_msg = f"Invalid or unsafe URL: {url}"
|
||||
logging.error(error_msg)
|
||||
return {"error": error_msg}
|
||||
|
||||
try:
|
||||
# Fetch the URL content
|
||||
downloaded = self.trafilatura.fetch_url(url)
|
||||
if not downloaded:
|
||||
error_msg = f"Could not fetch content from URL: {url}"
|
||||
logging.error(error_msg)
|
||||
return {"error": error_msg}
|
||||
|
||||
# Extract metadata
|
||||
metadata = self.trafilatura.extract_metadata(downloaded, url=url)
|
||||
|
||||
if not metadata:
|
||||
return {
|
||||
"url": url,
|
||||
"error": "No metadata found"
|
||||
}
|
||||
|
||||
# Convert metadata object to dict
|
||||
result = {
|
||||
"url": url,
|
||||
"title": metadata.title,
|
||||
"author": metadata.author,
|
||||
"date": metadata.date,
|
||||
"description": metadata.description,
|
||||
"categories": metadata.categories,
|
||||
"tags": metadata.tags,
|
||||
"language": metadata.language,
|
||||
"image": metadata.image,
|
||||
"pagetype": metadata.pagetype,
|
||||
"license": metadata.license
|
||||
}
|
||||
|
||||
# Remove None values
|
||||
result = {k: v for k, v in result.items() if v is not None}
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Error extracting metadata from {url}: {str(e)}"
|
||||
logging.error(error_msg)
|
||||
return {"error": error_msg}
|
||||
|
||||
def extract_text_only(
|
||||
self,
|
||||
url: str,
|
||||
include_comments: bool = False
|
||||
) -> str:
|
||||
"""
|
||||
Extract only the main text content from a URL.
|
||||
|
||||
Args:
|
||||
url: URL to extract text from
|
||||
include_comments: Include comments in extraction
|
||||
|
||||
Returns:
|
||||
str: Extracted text content
|
||||
"""
|
||||
result = self.extract_content(
|
||||
url=url,
|
||||
include_comments=include_comments,
|
||||
include_links=False,
|
||||
output_format='text',
|
||||
include_metadata=False
|
||||
)
|
||||
return result
|
||||
|
||||
def compare_extraction(
|
||||
self,
|
||||
url: str,
|
||||
include_newspaper: bool = True,
|
||||
include_spider: bool = True
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Compare Trafilatura extraction with other tools for the same URL.
|
||||
|
||||
Args:
|
||||
url: URL to extract content from
|
||||
include_newspaper: Include newspaper3k comparison
|
||||
include_spider: Include spider_tools comparison (if available)
|
||||
|
||||
Returns:
|
||||
Dict: Comparison of extraction results
|
||||
"""
|
||||
comparison = {
|
||||
"url": url,
|
||||
"trafilatura": None,
|
||||
"newspaper": None,
|
||||
"spider": None
|
||||
}
|
||||
|
||||
# Get Trafilatura extraction
|
||||
try:
|
||||
comparison["trafilatura"] = self.extract_content(url)
|
||||
except Exception as e:
|
||||
comparison["trafilatura"] = {"error": str(e)}
|
||||
|
||||
# Get Newspaper extraction if requested
|
||||
if include_newspaper:
|
||||
try:
|
||||
# Try to import and use newspaper3k
|
||||
if util.find_spec("newspaper") is not None:
|
||||
import newspaper
|
||||
article = newspaper.Article(url)
|
||||
article.download()
|
||||
article.parse()
|
||||
comparison["newspaper"] = {
|
||||
"title": article.title,
|
||||
"text": article.text,
|
||||
"authors": article.authors,
|
||||
"publish_date": str(article.publish_date) if article.publish_date else None,
|
||||
"top_image": article.top_image,
|
||||
"summary": article.summary if hasattr(article, 'summary') else None
|
||||
}
|
||||
else:
|
||||
comparison["newspaper"] = {"error": "newspaper3k not installed"}
|
||||
except Exception as e:
|
||||
comparison["newspaper"] = {"error": str(e)}
|
||||
|
||||
# Get Spider extraction if requested (placeholder for external spider tools)
|
||||
if include_spider:
|
||||
try:
|
||||
# This is a placeholder - you could integrate with other scraping tools here
|
||||
# For now, we'll use a simple requests + BeautifulSoup approach
|
||||
if util.find_spec("requests") and util.find_spec("bs4"):
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
response = requests.get(url, timeout=10)
|
||||
soup = BeautifulSoup(response.content, 'html.parser')
|
||||
|
||||
# Extract basic content
|
||||
title = soup.find('title')
|
||||
title_text = title.get_text().strip() if title else None
|
||||
|
||||
# Remove script and style elements
|
||||
for script in soup(["script", "style"]):
|
||||
script.decompose()
|
||||
|
||||
# Get text
|
||||
text = soup.get_text()
|
||||
# Clean up whitespace
|
||||
lines = (line.strip() for line in text.splitlines())
|
||||
chunks = (phrase.strip() for line in lines for phrase in line.split(" "))
|
||||
text = ' '.join(chunk for chunk in chunks if chunk)
|
||||
|
||||
comparison["spider"] = {
|
||||
"title": title_text,
|
||||
"text": text[:2000], # Limit to first 2000 chars
|
||||
"method": "requests+beautifulsoup"
|
||||
}
|
||||
else:
|
||||
comparison["spider"] = {"error": "requests or beautifulsoup4 not installed"}
|
||||
except Exception as e:
|
||||
comparison["spider"] = {"error": str(e)}
|
||||
|
||||
return comparison
|
||||
|
||||
def get_available_methods(self) -> List[str]:
|
||||
"""
|
||||
Get list of available methods in this tool.
|
||||
|
||||
Returns:
|
||||
List[str]: List of method names
|
||||
"""
|
||||
return [
|
||||
'extract_content',
|
||||
'extract_metadata',
|
||||
'extract_text_only',
|
||||
'compare_extraction'
|
||||
]
|
||||
|
||||
|
||||
# Convenience functions for easier usage
|
||||
def create_trafilatura_tools() -> TrafilaturaTools:
|
||||
"""Create and return a TrafilaturaTools instance."""
|
||||
return TrafilaturaTools()
|
||||
|
||||
|
||||
# Test the tool
|
||||
if __name__ == "__main__":
|
||||
# Test URL
|
||||
test_url = "https://www.python.org/about/"
|
||||
|
||||
print("Testing Custom Trafilatura Tools")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
# Create tools instance
|
||||
tools = TrafilaturaTools()
|
||||
|
||||
# Test 1: Extract content
|
||||
print("\n1. Testing extract_content():")
|
||||
print("-" * 50)
|
||||
content = tools.extract_content(test_url)
|
||||
if isinstance(content, dict) and 'error' not in content:
|
||||
print(f"✓ Title: {content.get('title', 'N/A')}")
|
||||
print(f"✓ Language: {content.get('language', 'N/A')}")
|
||||
print(f"✓ Text length: {len(content.get('text', ''))} characters")
|
||||
print(f"✓ Author: {content.get('author', 'N/A')}")
|
||||
print(f"✓ Date: {content.get('date', 'N/A')}")
|
||||
else:
|
||||
print(f"✗ Error: {content}")
|
||||
|
||||
# Test 2: Extract metadata
|
||||
print("\n2. Testing extract_metadata():")
|
||||
print("-" * 50)
|
||||
metadata = tools.extract_metadata(test_url)
|
||||
if isinstance(metadata, dict) and 'error' not in metadata:
|
||||
print("✓ Metadata extracted successfully:")
|
||||
for key, value in metadata.items():
|
||||
if key != 'url':
|
||||
print(f" - {key}: {value}")
|
||||
else:
|
||||
print(f"✗ Error: {metadata}")
|
||||
|
||||
# Test 3: Extract text only
|
||||
print("\n3. Testing extract_text_only():")
|
||||
print("-" * 50)
|
||||
text = tools.extract_text_only(test_url)
|
||||
if isinstance(text, str) and text and "error" not in text.lower():
|
||||
print(f"✓ Text extracted: {len(text)} characters")
|
||||
print(f"✓ Preview: {text[:150]}...")
|
||||
else:
|
||||
print(f"✗ Error: {text}")
|
||||
|
||||
# Test 4: Test URL validation
|
||||
print("\n4. Testing URL validation (security):")
|
||||
print("-" * 50)
|
||||
unsafe_urls = ["http://localhost/test", "http://127.0.0.1/test"]
|
||||
for unsafe_url in unsafe_urls:
|
||||
result = tools.extract_content(unsafe_url)
|
||||
if isinstance(result, dict) and 'error' in result and 'Invalid or unsafe URL' in result['error']:
|
||||
print(f"✓ Correctly rejected: {unsafe_url}")
|
||||
else:
|
||||
print(f"✗ Failed to reject: {unsafe_url}")
|
||||
|
||||
# Test 5: Available methods
|
||||
print("\n5. Available methods:")
|
||||
print("-" * 50)
|
||||
methods = tools.get_available_methods()
|
||||
for method in methods:
|
||||
print(f"✓ {method}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to initialize tools: {e}")
|
||||
print("\nMake sure to install dependencies:")
|
||||
print("pip install trafilatura newspaper3k requests beautifulsoup4")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("Custom Trafilatura Tools testing complete!")
|
||||
@@ -0,0 +1,266 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example: Using Custom Trafilatura Tools for High-Quality URL Content Extraction
|
||||
|
||||
This example demonstrates how to create and use custom Trafilatura tools in PraisonAI
|
||||
for extracting high-quality content from web pages.
|
||||
|
||||
This shows how to create custom tools that are not part of the core package.
|
||||
|
||||
Install dependencies:
|
||||
pip install trafilatura newspaper3k requests beautifulsoup4
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
from custom_trafilatura_tools import TrafilaturaTools, create_trafilatura_tools
|
||||
|
||||
# Create tools instance (global for reuse across examples)
|
||||
tools = None
|
||||
|
||||
def get_tools():
|
||||
"""Get or create tools instance"""
|
||||
global tools
|
||||
if tools is None:
|
||||
tools = create_trafilatura_tools()
|
||||
return tools
|
||||
|
||||
# Example 1: Basic content extraction
|
||||
async def basic_extraction_example():
|
||||
"""Extract content from a URL using custom Trafilatura tools"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
# Extract content from a URL
|
||||
url = "https://www.python.org/about/"
|
||||
content = await asyncio.to_thread(tools_instance.extract_content, url)
|
||||
|
||||
if 'error' not in content:
|
||||
print(f"Title: {content.get('title', 'N/A')}")
|
||||
print(f"Author: {content.get('author', 'N/A')}")
|
||||
print(f"Date: {content.get('date', 'N/A')}")
|
||||
print(f"Text preview: {content.get('text', '')[:200]}...")
|
||||
else:
|
||||
print(f"Error: {content['error']}")
|
||||
|
||||
# Example 2: Extract text only
|
||||
async def text_only_example():
|
||||
"""Extract only the main text content"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
url = "https://www.python.org/about/"
|
||||
text = await asyncio.to_thread(tools_instance.extract_text_only, url)
|
||||
|
||||
print(f"Extracted text ({len(text)} characters):")
|
||||
print(f"{text[:300]}...")
|
||||
|
||||
# Example 3: Extract metadata
|
||||
async def metadata_example():
|
||||
"""Extract metadata from a web page"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
url = "https://www.python.org/about/"
|
||||
metadata = await asyncio.to_thread(tools_instance.extract_metadata, url)
|
||||
|
||||
print("Page metadata:")
|
||||
for key, value in metadata.items():
|
||||
if key != 'url':
|
||||
print(f" {key}: {value}")
|
||||
|
||||
# Example 4: Compare extraction methods
|
||||
async def comparison_example():
|
||||
"""Compare Trafilatura with other extraction tools"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
url = "https://www.python.org/about/"
|
||||
comparison = await asyncio.to_thread(
|
||||
tools_instance.compare_extraction,
|
||||
url,
|
||||
include_newspaper=True, # Compare with newspaper3k
|
||||
include_spider=True # Compare with basic scraping
|
||||
)
|
||||
|
||||
print("Extraction comparison:")
|
||||
for tool, result in comparison.items():
|
||||
if tool != 'url' and result:
|
||||
if isinstance(result, dict) and 'error' not in result:
|
||||
print(f"\n{tool.upper()}:")
|
||||
if 'title' in result:
|
||||
print(f" Title: {result.get('title', 'N/A')}")
|
||||
if 'text' in result:
|
||||
text_len = len(result.get('text', ''))
|
||||
print(f" Text length: {text_len} characters")
|
||||
else:
|
||||
print(f"\n{tool.upper()}: {result}")
|
||||
|
||||
# Example 5: Using Custom Trafilatura tools with an Agent
|
||||
async def agent_example():
|
||||
"""Use custom Trafilatura tools with an AI agent"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
# Create wrapper functions for the agent to use
|
||||
def extract_content_tool(url: str) -> dict:
|
||||
"""Extract content from URL using Trafilatura"""
|
||||
return tools_instance.extract_content(url)
|
||||
|
||||
def extract_metadata_tool(url: str) -> dict:
|
||||
"""Extract metadata from URL using Trafilatura"""
|
||||
return tools_instance.extract_metadata(url)
|
||||
|
||||
# Create a research agent with custom Trafilatura tools
|
||||
research_agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Web Content Analyzer",
|
||||
goal="Extract and analyze content from web pages using custom tools",
|
||||
backstory="You are an expert at extracting and analyzing web content using custom Trafilatura tools",
|
||||
tools=[extract_content_tool, extract_metadata_tool],
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create a task for content analysis
|
||||
analysis_task = Task(
|
||||
description="""
|
||||
Extract content from this URL: https://www.python.org/about/
|
||||
|
||||
Use the custom Trafilatura tools to provide a summary of:
|
||||
1. The main content and purpose of the page
|
||||
2. Key metadata (author, date, language)
|
||||
3. The overall quality and structure of the content
|
||||
""",
|
||||
expected_output="A comprehensive analysis of the web page content using custom tools",
|
||||
agent=research_agent
|
||||
)
|
||||
|
||||
# Run the analysis
|
||||
agents = PraisonAIAgents(
|
||||
agents=[research_agent],
|
||||
tasks=[analysis_task]
|
||||
)
|
||||
|
||||
result = await agents.astart()
|
||||
print("\nAgent Analysis Result:")
|
||||
print(result)
|
||||
|
||||
# Example 6: Content extraction with language specification
|
||||
async def language_example():
|
||||
"""Extract content with specific language targeting"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
# Extract content targeting specific language
|
||||
url = "https://www.python.org/about/"
|
||||
content = await asyncio.to_thread(
|
||||
tools_instance.extract_content,
|
||||
url,
|
||||
target_language="en", # Target English content
|
||||
include_links=True, # Include links in extraction
|
||||
include_comments=False # Exclude comments
|
||||
)
|
||||
|
||||
if 'error' not in content:
|
||||
print(f"Language: {content.get('language', 'N/A')}")
|
||||
print(f"Number of links: {len(content.get('links', []))}")
|
||||
|
||||
# Example 7: Output format options
|
||||
async def format_example():
|
||||
"""Demonstrate different output formats"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
url = "https://www.python.org/about/"
|
||||
|
||||
# JSON format (default)
|
||||
json_content = await asyncio.to_thread(tools_instance.extract_content, url, output_format='json')
|
||||
print("JSON format - type:", type(json_content))
|
||||
|
||||
# Text format
|
||||
text_content = await asyncio.to_thread(tools_instance.extract_content, url, output_format='text')
|
||||
print("\nText format - first 200 chars:")
|
||||
print(text_content[:200] + "...")
|
||||
|
||||
# XML format
|
||||
xml_content = await asyncio.to_thread(tools_instance.extract_content, url, output_format='xml')
|
||||
print("\nXML format - first 200 chars:")
|
||||
print(xml_content[:200] + "...")
|
||||
|
||||
# Example 8: Show available methods
|
||||
async def available_methods_example():
|
||||
"""Show what methods are available in the custom tool"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
methods = tools_instance.get_available_methods()
|
||||
print("Available methods in custom Trafilatura tools:")
|
||||
for method in methods:
|
||||
print(f" - {method}")
|
||||
|
||||
# Example 9: Security testing
|
||||
async def security_example():
|
||||
"""Test URL validation security features"""
|
||||
tools_instance = get_tools()
|
||||
|
||||
print("Testing URL security validation:")
|
||||
unsafe_urls = [
|
||||
"http://localhost/test",
|
||||
"http://127.0.0.1/test",
|
||||
"http://169.254.169.254/metadata"
|
||||
]
|
||||
|
||||
for unsafe_url in unsafe_urls:
|
||||
result = await asyncio.to_thread(tools_instance.extract_content, unsafe_url)
|
||||
if isinstance(result, dict) and 'error' in result:
|
||||
print(f"✓ Correctly rejected: {unsafe_url}")
|
||||
else:
|
||||
print(f"✗ Failed to reject: {unsafe_url}")
|
||||
|
||||
async def main():
|
||||
"""Run all examples"""
|
||||
print("=== Custom Trafilatura Tools Examples ===\n")
|
||||
|
||||
try:
|
||||
# First check if we can create the tools
|
||||
tools_instance = get_tools()
|
||||
print("✓ Successfully created custom Trafilatura tools\n")
|
||||
|
||||
print("1. Available Methods")
|
||||
print("-" * 40)
|
||||
await available_methods_example()
|
||||
|
||||
print("\n2. Basic Content Extraction")
|
||||
print("-" * 40)
|
||||
await basic_extraction_example()
|
||||
|
||||
print("\n3. Text-Only Extraction")
|
||||
print("-" * 40)
|
||||
await text_only_example()
|
||||
|
||||
print("\n4. Metadata Extraction")
|
||||
print("-" * 40)
|
||||
await metadata_example()
|
||||
|
||||
print("\n5. Tool Comparison")
|
||||
print("-" * 40)
|
||||
await comparison_example()
|
||||
|
||||
print("\n6. Language-Specific Extraction")
|
||||
print("-" * 40)
|
||||
await language_example()
|
||||
|
||||
print("\n7. Output Format Options")
|
||||
print("-" * 40)
|
||||
await format_example()
|
||||
|
||||
print("\n8. Security Testing")
|
||||
print("-" * 40)
|
||||
await security_example()
|
||||
|
||||
print("\n9. Agent-Based Analysis")
|
||||
print("-" * 40)
|
||||
await agent_example()
|
||||
|
||||
except ImportError as e:
|
||||
print(f"✗ Error: {e}")
|
||||
print("\nPlease install required dependencies:")
|
||||
print("pip install trafilatura newspaper3k requests beautifulsoup4")
|
||||
except Exception as e:
|
||||
print(f"✗ Unexpected error: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the examples
|
||||
asyncio.run(main())
|
||||
@@ -139,10 +139,32 @@ task = Task(
|
||||
#### Task-Level Guardrails
|
||||
```python
|
||||
from typing import Tuple, Any
|
||||
from praisonaiagents import GuardrailResult
|
||||
|
||||
# Function-based guardrail
|
||||
def validate_output(task_output: TaskOutput) -> Tuple[bool, Any]:
|
||||
"""Custom validation function."""
|
||||
# Function-based guardrail (Option 1: Using GuardrailResult)
|
||||
def validate_output(task_output: TaskOutput) -> GuardrailResult:
|
||||
"""Custom validation function returning GuardrailResult."""
|
||||
if "error" in task_output.raw.lower():
|
||||
return GuardrailResult(
|
||||
success=False,
|
||||
result=None,
|
||||
error="Output contains errors"
|
||||
)
|
||||
if len(task_output.raw) < 10:
|
||||
return GuardrailResult(
|
||||
success=False,
|
||||
result=None,
|
||||
error="Output is too short"
|
||||
)
|
||||
return GuardrailResult(
|
||||
success=True,
|
||||
result=task_output,
|
||||
error=""
|
||||
)
|
||||
|
||||
# Function-based guardrail (Option 2: Using Tuple[bool, Any])
|
||||
def validate_output_tuple(task_output: TaskOutput) -> Tuple[bool, Any]:
|
||||
"""Custom validation function returning tuple."""
|
||||
if "error" in task_output.raw.lower():
|
||||
return False, "Output contains errors"
|
||||
if len(task_output.raw) < 10:
|
||||
@@ -170,7 +192,27 @@ task = Task(
|
||||
#### Agent-Level Guardrails
|
||||
```python
|
||||
# Agent guardrails apply to ALL outputs from that agent
|
||||
def validate_professional_tone(task_output: TaskOutput) -> Tuple[bool, Any]:
|
||||
|
||||
# Option 1: Using GuardrailResult
|
||||
def validate_professional_tone(task_output: TaskOutput) -> GuardrailResult:
|
||||
"""Ensure professional tone in all agent responses."""
|
||||
content = task_output.raw.lower()
|
||||
casual_words = ['yo', 'dude', 'awesome', 'cool']
|
||||
for word in casual_words:
|
||||
if word in content:
|
||||
return GuardrailResult(
|
||||
success=False,
|
||||
result=None,
|
||||
error=f"Unprofessional language detected: {word}"
|
||||
)
|
||||
return GuardrailResult(
|
||||
success=True,
|
||||
result=task_output,
|
||||
error=""
|
||||
)
|
||||
|
||||
# Option 2: Using Tuple[bool, Any]
|
||||
def validate_professional_tone_tuple(task_output: TaskOutput) -> Tuple[bool, Any]:
|
||||
"""Ensure professional tone in all agent responses."""
|
||||
content = task_output.raw.lower()
|
||||
casual_words = ['yo', 'dude', 'awesome', 'cool']
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def get_weather(city: str) -> str:
|
||||
return f"The weather in {city} is sunny"
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant",
|
||||
llm="gemini/gemini-1.5-flash-latest",
|
||||
tools=[get_weather]
|
||||
)
|
||||
|
||||
agent.start("What is the weather in Tokyo?")
|
||||
@@ -0,0 +1,30 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def get_stock_price(company_name: str) -> str:
|
||||
"""
|
||||
Get the stock price of a company
|
||||
|
||||
Args:
|
||||
company_name (str): The name of the company
|
||||
|
||||
Returns:
|
||||
str: The stock price of the company
|
||||
"""
|
||||
return f"The stock price of {company_name} is 100"
|
||||
|
||||
def multiply(a: int, b: int) -> int:
|
||||
"""
|
||||
Multiply two numbers
|
||||
"""
|
||||
return a * b
|
||||
|
||||
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant. You can use the tools provided to you to help the user.",
|
||||
llm="gemini/gemini-2.5-pro",
|
||||
tools=[get_stock_price, multiply]
|
||||
)
|
||||
|
||||
result = agent.start("what is the stock price of Google? multiply the Google stock price with 2")
|
||||
print(result)
|
||||
@@ -0,0 +1,25 @@
|
||||
from pydantic import BaseModel
|
||||
from praisonaiagents import Agent, Task, Agents
|
||||
|
||||
class Recipe(BaseModel):
|
||||
recipe_name: str
|
||||
ingredients: list[str]
|
||||
|
||||
agent = Agent(
|
||||
name="Chef",
|
||||
role="Recipe Creator",
|
||||
llm="gemini/gemini-2.5-flash"
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description="Create a cookie recipe",
|
||||
agent=agent,
|
||||
output_pydantic=Recipe # Will use Gemini's native structured output!
|
||||
)
|
||||
|
||||
agents = Agents(
|
||||
agents=[agent],
|
||||
tasks=[task]
|
||||
)
|
||||
|
||||
agents.start()
|
||||
@@ -0,0 +1,179 @@
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.agent.handoff import handoff, handoff_filters
|
||||
from pydantic import BaseModel
|
||||
from typing import Dict, Any
|
||||
import logging
|
||||
|
||||
# Mock data for testing
|
||||
MOCK_ORDERS = {
|
||||
"12345": {
|
||||
"id": "12345",
|
||||
"status": "shipped",
|
||||
"items": ["Widget A", "Widget B"],
|
||||
"total": 99.99,
|
||||
"tracking_number": "TRK123456789",
|
||||
"shipping_address": "123 Main St, City, ST 12345"
|
||||
},
|
||||
"67890": {
|
||||
"id": "67890",
|
||||
"status": "processing",
|
||||
"items": ["Gadget X"],
|
||||
"total": 149.99,
|
||||
"tracking_number": None,
|
||||
"shipping_address": "456 Oak Ave, Town, ST 67890"
|
||||
}
|
||||
}
|
||||
|
||||
MOCK_SHIPMENTS = {
|
||||
"TRK123456789": {
|
||||
"tracking_number": "TRK123456789",
|
||||
"status": "in_transit",
|
||||
"location": "Distribution Center - Chicago",
|
||||
"estimated_delivery": "2024-01-15",
|
||||
"updates": [
|
||||
{"date": "2024-01-10", "status": "picked_up", "location": "Warehouse"},
|
||||
{"date": "2024-01-12", "status": "in_transit", "location": "Chicago"}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
MOCK_TICKETS = []
|
||||
|
||||
# Mock tool functions
|
||||
def check_order_status(order_id: str) -> Dict[str, Any]:
|
||||
"""Check the status of an order"""
|
||||
if order_id in MOCK_ORDERS:
|
||||
return {"success": True, "order": MOCK_ORDERS[order_id]}
|
||||
return {"success": False, "error": f"Order {order_id} not found"}
|
||||
|
||||
def track_shipment(tracking_number: str) -> Dict[str, Any]:
|
||||
"""Track a shipment by tracking number"""
|
||||
if tracking_number in MOCK_SHIPMENTS:
|
||||
return {"success": True, "shipment": MOCK_SHIPMENTS[tracking_number]}
|
||||
return {"success": False, "error": f"Tracking number {tracking_number} not found"}
|
||||
|
||||
def process_refund(order_id: str, reason: str = "customer_request") -> Dict[str, Any]:
|
||||
"""Process a refund for an order"""
|
||||
if order_id in MOCK_ORDERS:
|
||||
order = MOCK_ORDERS[order_id]
|
||||
return {
|
||||
"success": True,
|
||||
"refund_id": f"REF{order_id}",
|
||||
"amount": order["total"],
|
||||
"status": "processing",
|
||||
"estimated_days": 3-5
|
||||
}
|
||||
return {"success": False, "error": f"Order {order_id} not found"}
|
||||
|
||||
def check_refund_policy(item_type: str = "general") -> Dict[str, Any]:
|
||||
"""Check refund policy for different item types"""
|
||||
policies = {
|
||||
"general": {"days": 30, "condition": "unopened", "shipping": "customer_pays"},
|
||||
"electronics": {"days": 15, "condition": "original_packaging", "shipping": "free"},
|
||||
"clothing": {"days": 60, "condition": "tags_attached", "shipping": "customer_pays"}
|
||||
}
|
||||
return {"success": True, "policy": policies.get(item_type, policies["general"])}
|
||||
|
||||
def diagnose_issue(description: str) -> Dict[str, Any]:
|
||||
"""Diagnose a technical issue"""
|
||||
common_solutions = {
|
||||
"login": "Clear cookies and cache, try password reset",
|
||||
"payment": "Check card details, try different payment method",
|
||||
"loading": "Check internet connection, try different browser",
|
||||
"mobile": "Update app, restart device, check storage space"
|
||||
}
|
||||
|
||||
# Simple keyword matching for demo
|
||||
for issue_type, solution in common_solutions.items():
|
||||
if issue_type in description.lower():
|
||||
return {"success": True, "diagnosis": solution, "severity": "low"}
|
||||
|
||||
return {"success": True, "diagnosis": "Please provide more details", "severity": "medium"}
|
||||
|
||||
def create_ticket(issue_description: str, priority: str = "medium") -> Dict[str, Any]:
|
||||
"""Create a support ticket"""
|
||||
ticket_id = f"TICK{len(MOCK_TICKETS) + 1001}"
|
||||
ticket = {
|
||||
"id": ticket_id,
|
||||
"description": issue_description,
|
||||
"priority": priority,
|
||||
"status": "open",
|
||||
"created_at": "2024-01-13T10:00:00Z"
|
||||
}
|
||||
MOCK_TICKETS.append(ticket)
|
||||
return {"success": True, "ticket": ticket}
|
||||
|
||||
def log_escalation(source_agent: str, escalation_data: Dict[str, Any]) -> None:
|
||||
"""Log escalation for audit purposes"""
|
||||
logging.info(f"Escalation from {source_agent}: {escalation_data}")
|
||||
|
||||
# Escalation with structured input
|
||||
class EscalationRequest(BaseModel):
|
||||
issue_type: str
|
||||
severity: str
|
||||
description: str
|
||||
|
||||
# Specialist agents
|
||||
order_agent = Agent(
|
||||
name="Order Management",
|
||||
instructions="You help with order status and tracking.",
|
||||
tools=[check_order_status, track_shipment]
|
||||
)
|
||||
|
||||
refund_agent = Agent(
|
||||
name="Refund Specialist",
|
||||
instructions="You process refunds and handle return requests.",
|
||||
tools=[process_refund, check_refund_policy]
|
||||
)
|
||||
|
||||
technical_agent = Agent(
|
||||
name="Technical Support",
|
||||
instructions="You solve technical issues and provide guidance.",
|
||||
tools=[diagnose_issue, create_ticket]
|
||||
)
|
||||
|
||||
escalation_agent = Agent(
|
||||
name="Senior Manager",
|
||||
instructions="You handle escalated issues requiring management attention."
|
||||
)
|
||||
|
||||
# Main triage agent
|
||||
triage_agent = Agent(
|
||||
name="Customer Service",
|
||||
instructions="""You are the first point of contact. Route customers to the right specialist:
|
||||
- Order issues → Order Management
|
||||
- Refunds → Refund Specialist
|
||||
- Technical problems → Technical Support
|
||||
- Complex issues → Escalate to management
|
||||
""",
|
||||
handoffs=[
|
||||
order_agent,
|
||||
refund_agent,
|
||||
technical_agent,
|
||||
handoff(
|
||||
escalation_agent,
|
||||
tool_name_override="escalate_to_management",
|
||||
input_type=EscalationRequest,
|
||||
on_handoff=lambda src, data: log_escalation(src.name, data)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Start the service
|
||||
if __name__ == "__main__":
|
||||
print("Customer Service Handoff System Demo")
|
||||
print("=" * 40)
|
||||
|
||||
# Test different scenarios
|
||||
test_queries = [
|
||||
"I need a refund for order #12345",
|
||||
"What's the status of order #67890?",
|
||||
"I'm having trouble logging into my account",
|
||||
"This is urgent - I need to speak to a manager about a serious billing issue"
|
||||
]
|
||||
|
||||
for query in test_queries:
|
||||
print(f"\nCustomer: {query}")
|
||||
response = triage_agent.chat(query)
|
||||
print(f"Response: {response}")
|
||||
print("-" * 40)
|
||||
@@ -0,0 +1,30 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def get_stock_price(company_name: str) -> str:
|
||||
"""
|
||||
Get the stock price of a company
|
||||
|
||||
Args:
|
||||
company_name (str): The name of the company
|
||||
|
||||
Returns:
|
||||
str: The stock price of the company
|
||||
"""
|
||||
return f"The stock price of {company_name} is 100"
|
||||
|
||||
def multiply(a: int, b: int) -> int:
|
||||
"""
|
||||
Multiply two numbers
|
||||
"""
|
||||
return a * b
|
||||
|
||||
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant. You can use the tools provided to you to help the user.",
|
||||
llm="ollama/llama3.2",
|
||||
tools=[get_stock_price, multiply]
|
||||
)
|
||||
|
||||
result = agent.start("what is the stock price of Google? multiply the Google stock price with 2")
|
||||
print(result)
|
||||
@@ -0,0 +1,30 @@
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def get_stock_price(company_name: str) -> str:
|
||||
"""
|
||||
Get the stock price of a company
|
||||
|
||||
Args:
|
||||
company_name (str): The name of the company
|
||||
|
||||
Returns:
|
||||
str: The stock price of the company
|
||||
"""
|
||||
return f"The stock price of {company_name} is 100"
|
||||
|
||||
def multiply(a: int, b: int) -> int:
|
||||
"""
|
||||
Multiply two numbers
|
||||
"""
|
||||
return a * b
|
||||
|
||||
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant. You can use the tools provided to you to help the user.",
|
||||
llm="gpt-4o",
|
||||
tools=[get_stock_price, multiply]
|
||||
)
|
||||
|
||||
result = agent.start("what is the stock price of Google? multiply the Google stock price with 2")
|
||||
print(result)
|
||||
@@ -2,6 +2,30 @@
|
||||
Praison AI Agents - A package for hierarchical AI agent task execution
|
||||
"""
|
||||
|
||||
# Configure logging before any other imports
|
||||
import os
|
||||
import logging
|
||||
from rich.logging import RichHandler
|
||||
|
||||
# Get log level from environment variable
|
||||
LOGLEVEL = os.environ.get('LOGLEVEL', 'INFO').upper()
|
||||
|
||||
# Configure root logger
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, LOGLEVEL, logging.INFO),
|
||||
format="%(asctime)s %(filename)s:%(lineno)d %(levelname)s %(message)s",
|
||||
datefmt="[%X]",
|
||||
handlers=[RichHandler(rich_tracebacks=True)]
|
||||
)
|
||||
|
||||
# Suppress specific noisy loggers
|
||||
logging.getLogger("litellm").setLevel(logging.WARNING)
|
||||
logging.getLogger("litellm.utils").setLevel(logging.WARNING)
|
||||
logging.getLogger("markdown_it").setLevel(logging.WARNING)
|
||||
logging.getLogger("rich.markdown").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
||||
|
||||
from .agent.agent import Agent
|
||||
from .agent.image_agent import ImageAgent
|
||||
from .agents.agents import PraisonAIAgents
|
||||
|
||||
@@ -2,5 +2,6 @@
|
||||
from .agent import Agent
|
||||
from .image_agent import ImageAgent
|
||||
from .handoff import Handoff, handoff, handoff_filters, RECOMMENDED_PROMPT_PREFIX, prompt_with_handoff_instructions
|
||||
from .router_agent import RouterAgent
|
||||
|
||||
__all__ = ['Agent', 'ImageAgent', 'Handoff', 'handoff', 'handoff_filters', 'RECOMMENDED_PROMPT_PREFIX', 'prompt_with_handoff_instructions']
|
||||
__all__ = ['Agent', 'ImageAgent', 'Handoff', 'handoff', 'handoff_filters', 'RECOMMENDED_PROMPT_PREFIX', 'prompt_with_handoff_instructions', 'RouterAgent']
|
||||
@@ -27,7 +27,8 @@ from ..main import (
|
||||
display_self_reflection,
|
||||
ReflectionOutput,
|
||||
adisplay_instruction,
|
||||
approval_callback
|
||||
approval_callback,
|
||||
execute_sync_callback
|
||||
)
|
||||
import inspect
|
||||
import uuid
|
||||
@@ -424,7 +425,13 @@ class Agent:
|
||||
# Otherwise, fall back to OpenAI environment/name
|
||||
else:
|
||||
self.llm = llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-4o')
|
||||
self.tools = tools if tools else [] # Store original tools
|
||||
# Handle tools parameter - ensure it's always a list
|
||||
if callable(tools):
|
||||
# If a single function/callable is passed, wrap it in a list
|
||||
self.tools = [tools]
|
||||
else:
|
||||
# Handle all falsy values (None, False, 0, "", etc.) by defaulting to empty list
|
||||
self.tools = tools or []
|
||||
self.function_calling_llm = function_calling_llm
|
||||
self.max_iter = max_iter
|
||||
self.max_rpm = max_rpm
|
||||
@@ -444,6 +451,8 @@ class Agent:
|
||||
self.embedder_config = embedder_config
|
||||
self.knowledge = knowledge
|
||||
self.use_system_prompt = use_system_prompt
|
||||
# NOTE: chat_history is not thread-safe. If concurrent access is needed,
|
||||
# consider using threading.Lock or other synchronization mechanisms
|
||||
self.chat_history = []
|
||||
self.markdown = markdown
|
||||
self.stream = stream
|
||||
@@ -711,8 +720,55 @@ Your Goal: {self.goal}
|
||||
)
|
||||
|
||||
return current_response
|
||||
|
||||
def _build_system_prompt(self, tools=None):
|
||||
"""Build the system prompt with tool information.
|
||||
|
||||
Args:
|
||||
tools: Optional list of tools to use (defaults to self.tools)
|
||||
|
||||
Returns:
|
||||
str: The system prompt or None if use_system_prompt is False
|
||||
"""
|
||||
if not self.use_system_prompt:
|
||||
return None
|
||||
|
||||
system_prompt = f"""{self.backstory}\n
|
||||
Your Role: {self.role}\n
|
||||
Your Goal: {self.goal}"""
|
||||
|
||||
# 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
|
||||
if tools_to_use:
|
||||
tool_names = []
|
||||
for tool in tools_to_use:
|
||||
try:
|
||||
if callable(tool) and hasattr(tool, '__name__'):
|
||||
tool_names.append(tool.__name__)
|
||||
elif isinstance(tool, dict) and isinstance(tool.get('function'), dict) and 'name' in tool['function']:
|
||||
tool_names.append(tool['function']['name'])
|
||||
elif isinstance(tool, str):
|
||||
tool_names.append(tool)
|
||||
elif hasattr(tool, "to_openai_tool"):
|
||||
# Handle MCP tools
|
||||
openai_tools = tool.to_openai_tool()
|
||||
if isinstance(openai_tools, list):
|
||||
for t in openai_tools:
|
||||
if isinstance(t, dict) and 'function' in t and 'name' in t['function']:
|
||||
tool_names.append(t['function']['name'])
|
||||
elif isinstance(openai_tools, dict) and 'function' in openai_tools:
|
||||
tool_names.append(openai_tools['function']['name'])
|
||||
except (AttributeError, KeyError, TypeError) as e:
|
||||
logging.warning(f"Could not extract tool name from {tool}: {e}")
|
||||
continue
|
||||
|
||||
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."
|
||||
|
||||
return system_prompt
|
||||
|
||||
def _build_messages(self, prompt, temperature=0.2, output_json=None, output_pydantic=None):
|
||||
def _build_messages(self, prompt, temperature=0.2, output_json=None, output_pydantic=None, tools=None):
|
||||
"""Build messages list for chat completion.
|
||||
|
||||
Args:
|
||||
@@ -720,17 +776,13 @@ Your Goal: {self.goal}
|
||||
temperature: Temperature for the chat
|
||||
output_json: Optional Pydantic model for JSON output
|
||||
output_pydantic: Optional Pydantic model for JSON output (alias)
|
||||
tools: Optional list of tools to use (defaults to self.tools)
|
||||
|
||||
Returns:
|
||||
tuple: (messages list, original prompt)
|
||||
"""
|
||||
# Build system prompt if enabled
|
||||
system_prompt = None
|
||||
if self.use_system_prompt:
|
||||
system_prompt = f"""{self.backstory}\n
|
||||
Your Role: {self.role}\n
|
||||
Your Goal: {self.goal}
|
||||
"""
|
||||
# Build system prompt using the helper method
|
||||
system_prompt = self._build_system_prompt(tools)
|
||||
|
||||
# Use openai_client's build_messages method if available
|
||||
if self._openai_client is not None:
|
||||
@@ -1048,6 +1100,10 @@ Your Goal: {self.goal}
|
||||
execute_tool_fn=self.execute_tool,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id,
|
||||
reasoning_steps=reasoning_steps
|
||||
)
|
||||
else:
|
||||
@@ -1064,6 +1120,10 @@ Your Goal: {self.goal}
|
||||
execute_tool_fn=self.execute_tool,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id,
|
||||
reasoning_steps=reasoning_steps
|
||||
)
|
||||
else:
|
||||
@@ -1097,8 +1157,39 @@ Your Goal: {self.goal}
|
||||
except Exception as e:
|
||||
display_error(f"Error in chat completion: {e}")
|
||||
return None
|
||||
|
||||
def _execute_callback_and_display(self, prompt: str, response: str, generation_time: float):
|
||||
"""Helper method to execute callbacks and display interaction.
|
||||
|
||||
This centralizes the logic for callback execution and display to avoid duplication.
|
||||
"""
|
||||
# Always execute callbacks regardless of verbose setting (only when not using custom LLM)
|
||||
if not self._using_custom_llm:
|
||||
execute_sync_callback(
|
||||
'interaction',
|
||||
message=prompt,
|
||||
response=response,
|
||||
markdown=self.markdown,
|
||||
generation_time=generation_time,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
|
||||
task_name=None, # Not available in this context
|
||||
task_description=None, # Not available in this context
|
||||
task_id=None # Not available in this context
|
||||
)
|
||||
# Only display interaction if not using custom LLM (to avoid double output) and verbose is True
|
||||
if self.verbose and not self._using_custom_llm:
|
||||
display_interaction(prompt, response, markdown=self.markdown,
|
||||
generation_time=generation_time, console=self.console,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
|
||||
task_name=None, # Not available in this context
|
||||
task_description=None, # Not available in this context
|
||||
task_id=None) # Not available in this context
|
||||
|
||||
def chat(self, prompt, temperature=0.2, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=True):
|
||||
def chat(self, prompt, temperature=0.2, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=True, task_name=None, task_description=None, task_id=None):
|
||||
# Log all parameter values when in debug mode
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
param_info = {
|
||||
@@ -1154,10 +1245,27 @@ Your Goal: {self.goal}
|
||||
tool_param = [openai_tool]
|
||||
logging.debug(f"Converted MCP tool: {tool_param}")
|
||||
|
||||
# Pass everything to LLM class
|
||||
response_text = self.llm_instance.get_response(
|
||||
# Store chat history length for potential rollback
|
||||
chat_history_length = len(self.chat_history)
|
||||
|
||||
# Normalize prompt content for consistent chat history storage
|
||||
normalized_content = prompt
|
||||
if isinstance(prompt, list):
|
||||
# Extract text from multimodal prompts
|
||||
normalized_content = next((item["text"] for item in prompt if item.get("type") == "text"), "")
|
||||
|
||||
# Prevent duplicate messages
|
||||
if not (self.chat_history and
|
||||
self.chat_history[-1].get("role") == "user" and
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
|
||||
try:
|
||||
# Pass everything to LLM class
|
||||
response_text = self.llm_instance.get_response(
|
||||
prompt=prompt,
|
||||
system_prompt=f"{self.backstory}\n\nYour Role: {self.role}\n\nYour Goal: {self.goal}" if self.use_system_prompt else None,
|
||||
system_prompt=self._build_system_prompt(tools),
|
||||
chat_history=self.chat_history,
|
||||
temperature=temperature,
|
||||
tools=tool_param,
|
||||
@@ -1172,24 +1280,34 @@ Your Goal: {self.goal}
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools if tools is not None else self.tools)],
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id,
|
||||
execute_tool_fn=self.execute_tool, # Pass tool execution function
|
||||
reasoning_steps=reasoning_steps
|
||||
)
|
||||
reasoning_steps=reasoning_steps,
|
||||
stream=stream # Pass the stream parameter from chat method
|
||||
)
|
||||
|
||||
self.chat_history.append({"role": "user", "content": prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
|
||||
# Log completion time if in debug mode
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
total_time = time.time() - start_time
|
||||
logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")
|
||||
# Log completion time if in debug mode
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
total_time = time.time() - start_time
|
||||
logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")
|
||||
|
||||
# Apply guardrail validation for custom LLM response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
|
||||
return validated_response
|
||||
# Apply guardrail validation for custom LLM response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
|
||||
# Rollback chat history if LLM call fails
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
display_error(f"Error in LLM chat: {e}")
|
||||
return None
|
||||
except Exception as e:
|
||||
display_error(f"Error in LLM chat: {e}")
|
||||
@@ -1197,191 +1315,214 @@ Your Goal: {self.goal}
|
||||
else:
|
||||
# Use the new _build_messages helper method
|
||||
messages, original_prompt = self._build_messages(prompt, temperature, output_json, output_pydantic)
|
||||
|
||||
# Store chat history length for potential rollback
|
||||
chat_history_length = len(self.chat_history)
|
||||
|
||||
# Normalize original_prompt for consistent chat history storage
|
||||
normalized_content = original_prompt
|
||||
if isinstance(original_prompt, list):
|
||||
# Extract text from multimodal prompts
|
||||
normalized_content = next((item["text"] for item in original_prompt if item.get("type") == "text"), "")
|
||||
|
||||
# Prevent duplicate messages
|
||||
if not (self.chat_history and
|
||||
self.chat_history[-1].get("role") == "user" and
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
|
||||
final_response_text = None
|
||||
reflection_count = 0
|
||||
start_time = time.time()
|
||||
|
||||
while True:
|
||||
try:
|
||||
if self.verbose:
|
||||
# Handle both string and list prompts for instruction display
|
||||
display_text = prompt
|
||||
if isinstance(prompt, list):
|
||||
# Extract text content from multimodal prompt
|
||||
display_text = next((item["text"] for item in prompt if item["type"] == "text"), "")
|
||||
|
||||
if display_text and str(display_text).strip():
|
||||
# Pass agent information to display_instruction
|
||||
agent_tools = [t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools]
|
||||
display_instruction(
|
||||
f"Agent {self.name} is processing prompt: {display_text}",
|
||||
console=self.console,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=agent_tools
|
||||
)
|
||||
|
||||
response = self._chat_completion(messages, temperature=temperature, tools=tools if tools else None, reasoning_steps=reasoning_steps, stream=self.stream)
|
||||
if not response:
|
||||
return None
|
||||
|
||||
response_text = response.choices[0].message.content.strip()
|
||||
|
||||
# Handle output_json or output_pydantic if specified
|
||||
if output_json or output_pydantic:
|
||||
# Add to chat history and return raw response
|
||||
self.chat_history.append({"role": "user", "content": original_prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Only display interaction if not using custom LLM (to avoid double output) and verbose is True
|
||||
if self.verbose and not self._using_custom_llm:
|
||||
display_interaction(original_prompt, response_text, markdown=self.markdown,
|
||||
generation_time=time.time() - start_time, console=self.console)
|
||||
return response_text
|
||||
|
||||
if not self.self_reflect:
|
||||
self.chat_history.append({"role": "user", "content": original_prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
|
||||
# Wrap entire while loop in try-except for rollback on any failure
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
if self.verbose:
|
||||
logging.debug(f"Agent {self.name} final response: {response_text}")
|
||||
# Only display interaction if not using custom LLM (to avoid double output) and verbose is True
|
||||
if self.verbose and not self._using_custom_llm:
|
||||
display_interaction(original_prompt, response_text, markdown=self.markdown, generation_time=time.time() - start_time, console=self.console)
|
||||
# Return only reasoning content if reasoning_steps is True
|
||||
if reasoning_steps and hasattr(response.choices[0].message, 'reasoning_content'):
|
||||
# Apply guardrail to reasoning content
|
||||
try:
|
||||
validated_reasoning = self._apply_guardrail_with_retry(response.choices[0].message.reasoning_content, original_prompt, temperature, tools)
|
||||
return validated_reasoning
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for reasoning content: {e}")
|
||||
return None
|
||||
# Apply guardrail to regular response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
|
||||
# Handle both string and list prompts for instruction display
|
||||
display_text = prompt
|
||||
if isinstance(prompt, list):
|
||||
# Extract text content from multimodal prompt
|
||||
display_text = next((item["text"] for item in prompt if item["type"] == "text"), "")
|
||||
|
||||
if display_text and str(display_text).strip():
|
||||
# Pass agent information to display_instruction
|
||||
agent_tools = [t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools]
|
||||
display_instruction(
|
||||
f"Agent {self.name} is processing prompt: {display_text}",
|
||||
console=self.console,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=agent_tools
|
||||
)
|
||||
|
||||
response = self._chat_completion(messages, temperature=temperature, tools=tools if tools else None, reasoning_steps=reasoning_steps, stream=self.stream)
|
||||
if not response:
|
||||
# Rollback chat history on response failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
|
||||
reflection_prompt = f"""
|
||||
response_text = response.choices[0].message.content.strip()
|
||||
|
||||
# Handle output_json or output_pydantic if specified
|
||||
if output_json or output_pydantic:
|
||||
# Add to chat history and return raw response
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Apply guardrail validation even for JSON output
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for JSON output: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
|
||||
if not self.self_reflect:
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
if self.verbose:
|
||||
logging.debug(f"Agent {self.name} final response: {response_text}")
|
||||
# Return only reasoning content if reasoning_steps is True
|
||||
if reasoning_steps and hasattr(response.choices[0].message, 'reasoning_content'):
|
||||
# Apply guardrail to reasoning content
|
||||
try:
|
||||
validated_reasoning = self._apply_guardrail_with_retry(response.choices[0].message.reasoning_content, original_prompt, temperature, tools)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_reasoning, time.time() - start_time)
|
||||
return validated_reasoning
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for reasoning content: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
# Apply guardrail to regular response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
|
||||
reflection_prompt = f"""
|
||||
Reflect on your previous response: '{response_text}'.
|
||||
{self.reflect_prompt if self.reflect_prompt else "Identify any flaws, improvements, or actions."}
|
||||
Provide a "satisfactory" status ('yes' or 'no').
|
||||
Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
"""
|
||||
logging.debug(f"{self.name} reflection attempt {reflection_count+1}, sending prompt: {reflection_prompt}")
|
||||
messages.append({"role": "user", "content": reflection_prompt})
|
||||
"""
|
||||
logging.debug(f"{self.name} reflection attempt {reflection_count+1}, sending prompt: {reflection_prompt}")
|
||||
messages.append({"role": "user", "content": reflection_prompt})
|
||||
|
||||
try:
|
||||
# Check if we're using a custom LLM (like Gemini)
|
||||
if self._using_custom_llm or self._openai_client is None:
|
||||
# For custom LLMs, we need to handle reflection differently
|
||||
# Use non-streaming to get complete JSON response
|
||||
reflection_response = self._chat_completion(messages, temperature=temperature, tools=None, stream=False, reasoning_steps=False)
|
||||
|
||||
if not reflection_response or not reflection_response.choices:
|
||||
raise Exception("No response from reflection request")
|
||||
|
||||
reflection_text = reflection_response.choices[0].message.content.strip()
|
||||
|
||||
# Clean the JSON output
|
||||
cleaned_json = self.clean_json_output(reflection_text)
|
||||
|
||||
# Parse the JSON manually
|
||||
reflection_data = json.loads(cleaned_json)
|
||||
|
||||
# Create a reflection output object manually
|
||||
class CustomReflectionOutput:
|
||||
def __init__(self, data):
|
||||
self.reflection = data.get('reflection', '')
|
||||
self.satisfactory = data.get('satisfactory', 'no').lower()
|
||||
|
||||
reflection_output = CustomReflectionOutput(reflection_data)
|
||||
else:
|
||||
# Use OpenAI's structured output for OpenAI models
|
||||
reflection_response = self._openai_client.sync_client.beta.chat.completions.parse(
|
||||
model=self.reflect_llm if self.reflect_llm else self.llm,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
response_format=ReflectionOutput
|
||||
)
|
||||
try:
|
||||
# Check if we're using a custom LLM (like Gemini)
|
||||
if self._using_custom_llm or self._openai_client is None:
|
||||
# For custom LLMs, we need to handle reflection differently
|
||||
# Use non-streaming to get complete JSON response
|
||||
reflection_response = self._chat_completion(messages, temperature=temperature, tools=None, stream=False, reasoning_steps=False)
|
||||
|
||||
if not reflection_response or not reflection_response.choices:
|
||||
raise Exception("No response from reflection request")
|
||||
|
||||
reflection_text = reflection_response.choices[0].message.content.strip()
|
||||
|
||||
# Clean the JSON output
|
||||
cleaned_json = self.clean_json_output(reflection_text)
|
||||
|
||||
# Parse the JSON manually
|
||||
reflection_data = json.loads(cleaned_json)
|
||||
|
||||
# Create a reflection output object manually
|
||||
class CustomReflectionOutput:
|
||||
def __init__(self, data):
|
||||
self.reflection = data.get('reflection', '')
|
||||
self.satisfactory = data.get('satisfactory', 'no').lower()
|
||||
|
||||
reflection_output = CustomReflectionOutput(reflection_data)
|
||||
else:
|
||||
# Use OpenAI's structured output for OpenAI models
|
||||
reflection_response = self._openai_client.sync_client.beta.chat.completions.parse(
|
||||
model=self.reflect_llm if self.reflect_llm else self.llm,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
response_format=ReflectionOutput
|
||||
)
|
||||
|
||||
reflection_output = reflection_response.choices[0].message.parsed
|
||||
reflection_output = reflection_response.choices[0].message.parsed
|
||||
|
||||
if self.verbose:
|
||||
display_self_reflection(f"Agent {self.name} self reflection (using {self.reflect_llm if self.reflect_llm else self.llm}): reflection='{reflection_output.reflection}' satisfactory='{reflection_output.satisfactory}'", console=self.console)
|
||||
|
||||
messages.append({"role": "assistant", "content": f"Self Reflection: {reflection_output.reflection} Satisfactory?: {reflection_output.satisfactory}"})
|
||||
|
||||
# Only consider satisfactory after minimum reflections
|
||||
if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
|
||||
if self.verbose:
|
||||
display_self_reflection("Agent marked the response as satisfactory after meeting minimum reflections", console=self.console)
|
||||
self.chat_history.append({"role": "user", "content": prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Only display interaction if not using custom LLM (to avoid double output) and verbose is True
|
||||
if self.verbose and not self._using_custom_llm:
|
||||
display_interaction(prompt, response_text, markdown=self.markdown, generation_time=time.time() - start_time, console=self.console)
|
||||
# Apply guardrail validation after satisfactory reflection
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
|
||||
return None
|
||||
display_self_reflection(f"Agent {self.name} self reflection (using {self.reflect_llm if self.reflect_llm else self.llm}): reflection='{reflection_output.reflection}' satisfactory='{reflection_output.satisfactory}'", console=self.console)
|
||||
|
||||
# Check if we've hit max reflections
|
||||
if reflection_count >= self.max_reflect - 1:
|
||||
if self.verbose:
|
||||
display_self_reflection("Maximum reflection count reached, returning current response", console=self.console)
|
||||
self.chat_history.append({"role": "user", "content": prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Only display interaction if not using custom LLM (to avoid double output) and verbose is True
|
||||
if self.verbose and not self._using_custom_llm:
|
||||
display_interaction(prompt, response_text, markdown=self.markdown, generation_time=time.time() - start_time, console=self.console)
|
||||
# Apply guardrail validation after max reflections
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after max reflections: {e}")
|
||||
return None
|
||||
messages.append({"role": "assistant", "content": f"Self Reflection: {reflection_output.reflection} Satisfactory?: {reflection_output.satisfactory}"})
|
||||
|
||||
logging.debug(f"{self.name} reflection count {reflection_count + 1}, continuing reflection process")
|
||||
messages.append({"role": "user", "content": "Now regenerate your response using the reflection you made"})
|
||||
# For custom LLMs during reflection, always use non-streaming to ensure complete responses
|
||||
use_stream = self.stream if not self._using_custom_llm else False
|
||||
response = self._chat_completion(messages, temperature=temperature, tools=None, stream=use_stream)
|
||||
response_text = response.choices[0].message.content.strip()
|
||||
reflection_count += 1
|
||||
continue # Continue the loop for more reflections
|
||||
# Only consider satisfactory after minimum reflections
|
||||
if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
|
||||
if self.verbose:
|
||||
display_self_reflection("Agent marked the response as satisfactory after meeting minimum reflections", console=self.console)
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Apply guardrail validation after satisfactory reflection
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
display_error(f"Error in parsing self-reflection json {e}. Retrying", console=self.console)
|
||||
logging.error("Reflection parsing failed.", exc_info=True)
|
||||
messages.append({"role": "assistant", "content": f"Self Reflection failed."})
|
||||
reflection_count += 1
|
||||
continue # Continue even after error to try again
|
||||
|
||||
except Exception as e:
|
||||
display_error(f"Error in chat: {e}", console=self.console)
|
||||
return None
|
||||
# Check if we've hit max reflections
|
||||
if reflection_count >= self.max_reflect - 1:
|
||||
if self.verbose:
|
||||
display_self_reflection("Maximum reflection count reached, returning current response", console=self.console)
|
||||
# User message already added before LLM call via _build_messages
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
# Apply guardrail validation after max reflections
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after max reflections: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
|
||||
# If not satisfactory and not at max reflections, continue with regeneration
|
||||
logging.debug(f"{self.name} reflection count {reflection_count + 1}, continuing reflection process")
|
||||
messages.append({"role": "user", "content": "Now regenerate your response using the reflection you made"})
|
||||
# For custom LLMs during reflection, always use non-streaming to ensure complete responses
|
||||
use_stream = self.stream if not self._using_custom_llm else False
|
||||
response = self._chat_completion(messages, temperature=temperature, tools=None, stream=use_stream)
|
||||
response_text = response.choices[0].message.content.strip()
|
||||
reflection_count += 1
|
||||
continue # Continue the loop for more reflections
|
||||
|
||||
# Log completion time if in debug mode
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
total_time = time.time() - start_time
|
||||
logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")
|
||||
|
||||
# Apply guardrail validation before returning
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
|
||||
if self.verbose:
|
||||
display_error(f"Guardrail validation failed: {e}", console=self.console)
|
||||
return None
|
||||
except Exception as e:
|
||||
display_error(f"Error in parsing self-reflection json {e}. Retrying", console=self.console)
|
||||
logging.error("Reflection parsing failed.", exc_info=True)
|
||||
messages.append({"role": "assistant", "content": "Self Reflection failed."})
|
||||
reflection_count += 1
|
||||
continue # Continue even after error to try again
|
||||
except Exception:
|
||||
# Catch any exception from the inner try block and re-raise to outer handler
|
||||
raise
|
||||
except Exception as e:
|
||||
# Catch any exceptions that escape the while loop
|
||||
display_error(f"Unexpected error in chat: {e}", console=self.console)
|
||||
# Rollback chat history
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
|
||||
def clean_json_output(self, output: str) -> str:
|
||||
"""Clean and extract JSON from response text."""
|
||||
@@ -1430,10 +1571,26 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
prompt = f"{prompt}\n\nKnowledge: {knowledge_content}"
|
||||
|
||||
if self._using_custom_llm:
|
||||
# Store chat history length for potential rollback
|
||||
chat_history_length = len(self.chat_history)
|
||||
|
||||
# Normalize prompt content for consistent chat history storage
|
||||
normalized_content = prompt
|
||||
if isinstance(prompt, list):
|
||||
# Extract text from multimodal prompts
|
||||
normalized_content = next((item["text"] for item in prompt if item.get("type") == "text"), "")
|
||||
|
||||
# Prevent duplicate messages
|
||||
if not (self.chat_history and
|
||||
self.chat_history[-1].get("role") == "user" and
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
|
||||
try:
|
||||
response_text = await self.llm_instance.get_response_async(
|
||||
prompt=prompt,
|
||||
system_prompt=f"{self.backstory}\n\nYour Role: {self.role}\n\nYour Goal: {self.goal}" if self.use_system_prompt else None,
|
||||
system_prompt=self._build_system_prompt(tools),
|
||||
chat_history=self.chat_history,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
@@ -1447,19 +1604,32 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
console=self.console,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools],
|
||||
agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools if tools is not None else self.tools)],
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id,
|
||||
execute_tool_fn=self.execute_tool_async,
|
||||
reasoning_steps=reasoning_steps
|
||||
)
|
||||
|
||||
self.chat_history.append({"role": "user", "content": prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
total_time = time.time() - start_time
|
||||
logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
|
||||
return response_text
|
||||
|
||||
# Apply guardrail validation for custom LLM response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
except Exception as e:
|
||||
# Rollback chat history if LLM call fails
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
display_error(f"Error in LLM chat: {e}")
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
total_time = time.time() - start_time
|
||||
@@ -1469,6 +1639,22 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# For OpenAI client
|
||||
# Use the new _build_messages helper method
|
||||
messages, original_prompt = self._build_messages(prompt, temperature, output_json, output_pydantic)
|
||||
|
||||
# Store chat history length for potential rollback
|
||||
chat_history_length = len(self.chat_history)
|
||||
|
||||
# Normalize original_prompt for consistent chat history storage
|
||||
normalized_content = original_prompt
|
||||
if isinstance(original_prompt, list):
|
||||
# Extract text from multimodal prompts
|
||||
normalized_content = next((item["text"] for item in original_prompt if item.get("type") == "text"), "")
|
||||
|
||||
# Prevent duplicate messages
|
||||
if not (self.chat_history and
|
||||
self.chat_history[-1].get("role") == "user" and
|
||||
self.chat_history[-1].get("content") == normalized_content):
|
||||
# Add user message to chat history BEFORE LLM call so handoffs can access it
|
||||
self.chat_history.append({"role": "user", "content": normalized_content})
|
||||
|
||||
reflection_count = 0
|
||||
start_time = time.time()
|
||||
@@ -1615,7 +1801,16 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
total_time = time.time() - start_time
|
||||
logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
|
||||
return response_text
|
||||
|
||||
# Apply guardrail validation for OpenAI client response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools)
|
||||
return validated_response
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for OpenAI client: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self.chat_history = self.chat_history[:chat_history_length]
|
||||
return None
|
||||
except Exception as e:
|
||||
display_error(f"Error in chat completion: {e}")
|
||||
if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
|
||||
@@ -1731,6 +1926,26 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
"""Start the agent with a prompt. This is a convenience method that wraps chat()."""
|
||||
return self.chat(prompt, **kwargs)
|
||||
|
||||
def execute(self, task, context=None):
|
||||
"""Execute a task synchronously - backward compatibility method"""
|
||||
if hasattr(task, 'description'):
|
||||
prompt = task.description
|
||||
elif isinstance(task, str):
|
||||
prompt = task
|
||||
else:
|
||||
prompt = str(task)
|
||||
return self.chat(prompt)
|
||||
|
||||
async def aexecute(self, task, context=None):
|
||||
"""Execute a task asynchronously - backward compatibility method"""
|
||||
if hasattr(task, 'description'):
|
||||
prompt = task.description
|
||||
elif isinstance(task, str):
|
||||
prompt = task
|
||||
else:
|
||||
prompt = str(task)
|
||||
return await self.achat(prompt)
|
||||
|
||||
async def execute_tool_async(self, function_name: str, arguments: Dict[str, Any]) -> Any:
|
||||
"""Async version of execute_tool"""
|
||||
try:
|
||||
|
||||
@@ -144,8 +144,17 @@ class ImageAgent(Agent):
|
||||
config = self.image_config.dict(exclude_none=True)
|
||||
config.update(kwargs)
|
||||
|
||||
# Use llm parameter as the model
|
||||
config['model'] = self.llm
|
||||
# Get the model name robustly from the parent Agent's property
|
||||
model_info = self.llm_model
|
||||
model_name = model_info.model if hasattr(model_info, 'model') else str(model_info)
|
||||
|
||||
# Use the model name in config
|
||||
config['model'] = model_name
|
||||
|
||||
# Check if we're using a Gemini model and remove unsupported parameters
|
||||
if 'gemini' in model_name.lower():
|
||||
# Gemini models don't support response_format parameter
|
||||
config.pop('response_format', None)
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
@@ -154,7 +163,7 @@ class ImageAgent(Agent):
|
||||
) as progress:
|
||||
try:
|
||||
# Add a task for image generation
|
||||
task = progress.add_task(f"[cyan]Generating image with {self.llm}...", total=None)
|
||||
task = progress.add_task(f"[cyan]Generating image with {model_name}...", total=None)
|
||||
|
||||
# Use litellm's image generation
|
||||
response = self.litellm(
|
||||
|
||||
@@ -0,0 +1,334 @@
|
||||
"""
|
||||
Router Agent that can use different LLM models based on task requirements.
|
||||
|
||||
This module extends the base Agent class to support multiple models and intelligent
|
||||
model selection based on task characteristics.
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Any, Union
|
||||
from .agent import Agent
|
||||
from ..llm.model_router import ModelRouter
|
||||
from ..llm import LLM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RouterAgent(Agent):
|
||||
"""
|
||||
An enhanced agent that can dynamically select and use different LLM models
|
||||
based on task requirements, optimizing for cost and performance.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
models: Optional[Union[List[str], Dict[str, Any]]] = None,
|
||||
model_router: Optional[ModelRouter] = None,
|
||||
routing_strategy: str = "auto", # "auto", "manual", "cost-optimized", "performance-optimized"
|
||||
primary_model: Optional[str] = None,
|
||||
fallback_model: Optional[str] = None,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
Initialize a RouterAgent.
|
||||
|
||||
Args:
|
||||
models: List of model names or dict mapping model names to configurations
|
||||
model_router: Custom ModelRouter instance for model selection
|
||||
routing_strategy: Strategy for model selection
|
||||
primary_model: Primary model to use (overrides routing for simple tasks)
|
||||
fallback_model: Fallback model if selected model fails
|
||||
**kwargs: Additional arguments passed to parent Agent class
|
||||
"""
|
||||
# Initialize model router
|
||||
self.model_router = model_router or ModelRouter()
|
||||
self.routing_strategy = routing_strategy
|
||||
self.fallback_model = fallback_model or os.getenv('OPENAI_MODEL_NAME', 'gpt-4o-mini')
|
||||
|
||||
# Process models configuration
|
||||
self.available_models = self._process_models_config(models)
|
||||
|
||||
# Set primary model for parent class initialization
|
||||
if primary_model:
|
||||
kwargs['llm'] = primary_model
|
||||
elif self.available_models:
|
||||
# Use the most cost-effective model as default
|
||||
cheapest_model = min(
|
||||
self.available_models.keys(),
|
||||
key=lambda m: self.model_router.get_model_info(m).cost_per_1k_tokens
|
||||
if self.model_router.get_model_info(m) else float('inf')
|
||||
)
|
||||
kwargs['llm'] = cheapest_model
|
||||
|
||||
# Store the original llm parameter for later use
|
||||
self._llm_config = kwargs.get('llm')
|
||||
|
||||
# Store api_key and base_url for LLM initialization
|
||||
self._base_url = kwargs.get('base_url')
|
||||
self._api_key = kwargs.get('api_key')
|
||||
|
||||
# Initialize parent Agent class
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Initialize LLM instances for each model
|
||||
self._llm_instances: Dict[str, LLM] = {}
|
||||
self._initialize_llm_instances()
|
||||
|
||||
# Track usage statistics
|
||||
self.model_usage_stats = {model: {'calls': 0, 'tokens': 0, 'cost': 0.0}
|
||||
for model in self.available_models}
|
||||
|
||||
def _process_models_config(self, models: Optional[Union[List[str], Dict[str, Any]]]) -> Dict[str, Any]:
|
||||
"""Process the models configuration into a standardized format."""
|
||||
if not models:
|
||||
# Use default models from router
|
||||
return {m.name: {} for m in self.model_router.models}
|
||||
|
||||
if isinstance(models, list):
|
||||
# Simple list of model names
|
||||
return {model: {} for model in models}
|
||||
|
||||
# Already a dict with model configurations
|
||||
return models
|
||||
|
||||
def _initialize_llm_instances(self):
|
||||
"""Initialize LLM instances for each available model."""
|
||||
base_url = self._base_url
|
||||
api_key = self._api_key
|
||||
|
||||
for model_name, config in self.available_models.items():
|
||||
try:
|
||||
# Merge base configuration with model-specific config
|
||||
llm_config = {
|
||||
'model': model_name,
|
||||
'base_url': config.get('base_url', base_url),
|
||||
'api_key': config.get('api_key', api_key),
|
||||
'verbose': self.verbose,
|
||||
'markdown': self.markdown,
|
||||
'stream': self.stream
|
||||
}
|
||||
|
||||
# Add any model-specific parameters
|
||||
llm_config.update(config)
|
||||
|
||||
# Create LLM instance
|
||||
self._llm_instances[model_name] = LLM(**llm_config)
|
||||
logger.debug(f"Initialized LLM instance for model: {model_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize LLM for model {model_name}: {e}")
|
||||
|
||||
def _select_model_for_task(
|
||||
self,
|
||||
task_description: str,
|
||||
tools: Optional[List[Any]] = None,
|
||||
context_size: Optional[int] = None
|
||||
) -> str:
|
||||
"""
|
||||
Select the most appropriate model for a given task.
|
||||
|
||||
Args:
|
||||
task_description: Description of the task
|
||||
tools: Tools required for the task
|
||||
context_size: Estimated context size
|
||||
|
||||
Returns:
|
||||
Selected model name
|
||||
"""
|
||||
if self.routing_strategy == "manual":
|
||||
# Use the configured primary model from llm_model property
|
||||
llm_model = self.llm_model
|
||||
if hasattr(llm_model, 'model'):
|
||||
# If it's an LLM instance, get the model name
|
||||
return llm_model.model
|
||||
elif isinstance(llm_model, str):
|
||||
# If it's a string, use it directly
|
||||
return llm_model
|
||||
# Fallback if no model is configured
|
||||
return self.fallback_model
|
||||
|
||||
# Determine required capabilities
|
||||
required_capabilities = []
|
||||
if tools:
|
||||
required_capabilities.append("function-calling")
|
||||
|
||||
# Determine budget consciousness based on strategy
|
||||
budget_conscious = self.routing_strategy in ["auto", "cost-optimized"]
|
||||
|
||||
# Get tool names for analysis
|
||||
tool_names = []
|
||||
if tools:
|
||||
tool_names = [t.__name__ if hasattr(t, '__name__') else str(t) for t in tools]
|
||||
|
||||
# Use router to select model
|
||||
selected_model = self.model_router.select_model(
|
||||
task_description=task_description,
|
||||
required_capabilities=required_capabilities,
|
||||
tools_required=tool_names,
|
||||
context_size=context_size,
|
||||
budget_conscious=budget_conscious
|
||||
)
|
||||
|
||||
# Ensure selected model is available
|
||||
if selected_model not in self._llm_instances:
|
||||
logger.warning(f"Selected model {selected_model} not available, using fallback")
|
||||
return self.fallback_model
|
||||
|
||||
return selected_model
|
||||
|
||||
def _execute_with_model(
|
||||
self,
|
||||
model_name: str,
|
||||
prompt: str,
|
||||
context: Optional[str] = None,
|
||||
tools: Optional[List[Any]] = None,
|
||||
**kwargs
|
||||
) -> str:
|
||||
"""
|
||||
Execute a task with a specific model.
|
||||
|
||||
Args:
|
||||
model_name: Name of the model to use
|
||||
prompt: The prompt to send to the model
|
||||
context: Additional context
|
||||
tools: Tools to make available
|
||||
**kwargs: Additional arguments for the LLM
|
||||
|
||||
Returns:
|
||||
Model response
|
||||
"""
|
||||
llm_instance = self._llm_instances.get(model_name)
|
||||
if not llm_instance:
|
||||
logger.error(f"Model {model_name} not initialized, using fallback")
|
||||
llm_instance = self._llm_instances.get(self.fallback_model)
|
||||
model_name = self.fallback_model
|
||||
|
||||
if not llm_instance:
|
||||
raise ValueError("No LLM instance available for execution")
|
||||
|
||||
# Prepare the full prompt
|
||||
full_prompt = prompt
|
||||
if context:
|
||||
full_prompt = f"{context}\n\n{prompt}"
|
||||
|
||||
try:
|
||||
# Execute with the selected model
|
||||
response = llm_instance.get_response(
|
||||
prompt=full_prompt,
|
||||
system_prompt=self._build_system_prompt(),
|
||||
tools=tools,
|
||||
verbose=self.verbose,
|
||||
markdown=self.markdown,
|
||||
stream=self.stream,
|
||||
agent_name=self.name,
|
||||
agent_role=self.role,
|
||||
agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools or [])],
|
||||
execute_tool_fn=self.execute_tool if tools else None,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
# Update usage statistics
|
||||
self.model_usage_stats[model_name]['calls'] += 1
|
||||
|
||||
# TODO: Implement token tracking when LLM.get_response() is updated to return token usage
|
||||
# The LLM response currently returns only text, but litellm provides usage info in:
|
||||
# response.get("usage") with prompt_tokens, completion_tokens, and total_tokens
|
||||
# This would require modifying the LLM class to return both text and metadata
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing with model {model_name}: {e}")
|
||||
|
||||
# Try fallback model if different
|
||||
if model_name != self.fallback_model and self.fallback_model in self._llm_instances:
|
||||
logger.info(f"Attempting with fallback model: {self.fallback_model}")
|
||||
return self._execute_with_model(
|
||||
self.fallback_model, prompt, context, tools, **kwargs
|
||||
)
|
||||
|
||||
raise
|
||||
|
||||
def execute(
|
||||
self,
|
||||
task_description: str,
|
||||
context: Optional[str] = None,
|
||||
tools: Optional[List[Any]] = None,
|
||||
**kwargs
|
||||
) -> str:
|
||||
"""
|
||||
Execute a task with automatic model selection.
|
||||
|
||||
This method overrides the parent Agent's execute method to add
|
||||
intelligent model selection.
|
||||
|
||||
Args:
|
||||
task_description: Description of the task to execute
|
||||
context: Optional context for the task
|
||||
tools: Optional tools to use
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
Task execution result
|
||||
"""
|
||||
# Estimate context size in tokens (rough estimate: ~4 chars per token)
|
||||
# This is a simplified heuristic; actual tokenization varies by model
|
||||
text_length = len(task_description) + (len(context) if context else 0)
|
||||
context_size = text_length // 4 # Approximate token count
|
||||
|
||||
# Select the best model for this task
|
||||
selected_model = self._select_model_for_task(
|
||||
task_description=task_description,
|
||||
tools=tools,
|
||||
context_size=context_size
|
||||
)
|
||||
|
||||
logger.info(f"RouterAgent '{self.name}' selected model: {selected_model} for task")
|
||||
|
||||
# Execute with the selected model
|
||||
return self._execute_with_model(
|
||||
model_name=selected_model,
|
||||
prompt=task_description,
|
||||
context=context,
|
||||
tools=tools,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
def get_usage_report(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get a report of model usage statistics.
|
||||
|
||||
Returns:
|
||||
Dictionary containing usage statistics and cost estimates
|
||||
"""
|
||||
total_cost = 0.0
|
||||
report = {
|
||||
'agent_name': self.name,
|
||||
'routing_strategy': self.routing_strategy,
|
||||
'model_usage': {}
|
||||
}
|
||||
|
||||
for model, stats in self.model_usage_stats.items():
|
||||
model_info = self.model_router.get_model_info(model)
|
||||
if model_info and stats['tokens'] > 0:
|
||||
cost = self.model_router.estimate_cost(model, stats['tokens'])
|
||||
stats['cost'] = cost
|
||||
total_cost += cost
|
||||
|
||||
report['model_usage'][model] = stats
|
||||
|
||||
report['total_cost_estimate'] = total_cost
|
||||
report['total_calls'] = sum(s['calls'] for s in self.model_usage_stats.values())
|
||||
|
||||
return report
|
||||
|
||||
def _build_system_prompt(self) -> str:
|
||||
"""Build system prompt (inherited from parent but can be customized)."""
|
||||
base_prompt = super()._build_system_prompt()
|
||||
|
||||
# Add multi-model context if needed
|
||||
if self.routing_strategy == "auto":
|
||||
base_prompt += "\n\nNote: You are part of a multi-model system. Focus on your specific task."
|
||||
|
||||
return base_prompt
|
||||
@@ -480,31 +480,46 @@ Context:
|
||||
)
|
||||
|
||||
if self.process == "workflow":
|
||||
# Collect all tasks that should run in parallel
|
||||
parallel_tasks = []
|
||||
tasks_to_run = []
|
||||
async for task_id in process.aworkflow():
|
||||
if self.tasks[task_id].async_execution and self.tasks[task_id].is_start:
|
||||
parallel_tasks.append(task_id)
|
||||
elif parallel_tasks:
|
||||
# Execute collected parallel tasks
|
||||
await asyncio.gather(*[self.arun_task(t) for t in parallel_tasks])
|
||||
parallel_tasks = []
|
||||
# Run the current non-parallel task
|
||||
if self.tasks[task_id].async_execution:
|
||||
await self.arun_task(task_id)
|
||||
else:
|
||||
self.run_task(task_id)
|
||||
|
||||
# Execute any remaining parallel tasks
|
||||
if parallel_tasks:
|
||||
await asyncio.gather(*[self.arun_task(t) for t in parallel_tasks])
|
||||
if self.tasks[task_id].async_execution:
|
||||
tasks_to_run.append(self.arun_task(task_id))
|
||||
else:
|
||||
# If we encounter a sync task, we must wait for the previous async tasks to finish.
|
||||
if tasks_to_run:
|
||||
await asyncio.gather(*tasks_to_run)
|
||||
tasks_to_run = []
|
||||
|
||||
# Run sync task in an executor to avoid blocking the event loop
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self.run_task, task_id)
|
||||
|
||||
if tasks_to_run:
|
||||
await asyncio.gather(*tasks_to_run)
|
||||
|
||||
elif self.process == "sequential":
|
||||
async_tasks_to_run = []
|
||||
|
||||
async def flush_async_tasks():
|
||||
"""Execute all pending async tasks"""
|
||||
nonlocal async_tasks_to_run
|
||||
if async_tasks_to_run:
|
||||
await asyncio.gather(*async_tasks_to_run)
|
||||
async_tasks_to_run = []
|
||||
|
||||
async for task_id in process.asequential():
|
||||
if self.tasks[task_id].async_execution:
|
||||
await self.arun_task(task_id)
|
||||
# Collect async tasks to run in parallel
|
||||
async_tasks_to_run.append(self.arun_task(task_id))
|
||||
else:
|
||||
self.run_task(task_id)
|
||||
# Before running a sync task, execute all pending async tasks
|
||||
await flush_async_tasks()
|
||||
# Run sync task in an executor to avoid blocking the event loop
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self.run_task, task_id)
|
||||
|
||||
# Execute any remaining async tasks at the end
|
||||
await flush_async_tasks()
|
||||
elif self.process == "hierarchical":
|
||||
async for task_id in process.ahierarchical():
|
||||
if isinstance(task_id, Task):
|
||||
@@ -512,7 +527,9 @@ Context:
|
||||
if self.tasks[task_id].async_execution:
|
||||
await self.arun_task(task_id)
|
||||
else:
|
||||
self.run_task(task_id)
|
||||
# Run sync task in an executor to avoid blocking the event loop
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self.run_task, task_id)
|
||||
|
||||
async def astart(self, content=None, return_dict=False, **kwargs):
|
||||
"""Async version of start method
|
||||
@@ -672,7 +689,10 @@ Context:
|
||||
_get_multimodal_message(task_prompt, task.images),
|
||||
tools=task.tools,
|
||||
output_json=task.output_json,
|
||||
output_pydantic=task.output_pydantic
|
||||
output_pydantic=task.output_pydantic,
|
||||
task_name=task.name,
|
||||
task_description=task.description,
|
||||
task_id=task_id
|
||||
)
|
||||
else:
|
||||
agent_output = executor_agent.chat(
|
||||
@@ -681,6 +701,9 @@ Context:
|
||||
output_json=task.output_json,
|
||||
output_pydantic=task.output_pydantic,
|
||||
stream=self.stream,
|
||||
task_name=task.name,
|
||||
task_description=task.description,
|
||||
task_id=task_id
|
||||
)
|
||||
|
||||
if agent_output:
|
||||
@@ -1118,7 +1141,7 @@ Context:
|
||||
response = await agent_instance.achat(current_input)
|
||||
else:
|
||||
# Run sync function in a thread to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
loop = asyncio.get_running_loop()
|
||||
# Correctly pass current_input to the lambda for closure
|
||||
response = await loop.run_in_executor(None, lambda ci=current_input: agent_instance.chat(ci))
|
||||
|
||||
@@ -1273,7 +1296,7 @@ Context:
|
||||
if hasattr(agent_instance, 'achat') and asyncio.iscoroutinefunction(agent_instance.achat):
|
||||
response = await agent_instance.achat(current_input, tools=agent_instance.tools)
|
||||
elif hasattr(agent_instance, 'chat'): # Fallback to sync chat if achat not suitable
|
||||
loop = asyncio.get_event_loop()
|
||||
loop = asyncio.get_running_loop()
|
||||
response = await loop.run_in_executor(None, lambda ci=current_input: agent_instance.chat(ci, tools=agent_instance.tools))
|
||||
else:
|
||||
logging.warning(f"Agent {agent_instance.name} has no suitable chat or achat method.")
|
||||
|
||||
@@ -136,7 +136,7 @@ class AutoAgents(PraisonAIAgents):
|
||||
completion_checker=completion_checker,
|
||||
max_retries=max_retries,
|
||||
process=process,
|
||||
manager_llm=manager_llm
|
||||
manager_llm=manager_llm or self.llm
|
||||
)
|
||||
|
||||
def _display_agents_and_tasks(self, agents: List[Agent], tasks: List[Task]):
|
||||
@@ -296,6 +296,10 @@ IMPORTANT: Each task MUST be an object with name, description, expected_output,
|
||||
last_error = None
|
||||
|
||||
for attempt in range(max_retries):
|
||||
# Initialize variables for this attempt
|
||||
use_openai_structured = False
|
||||
client = None
|
||||
|
||||
# Prepare prompt for this attempt
|
||||
if attempt > 0 and last_response and last_error:
|
||||
# On retry, include the previous response and error
|
||||
@@ -322,23 +326,16 @@ DO NOT use strings for tasks. Each task MUST be a complete object with all four
|
||||
prompt = base_prompt
|
||||
|
||||
try:
|
||||
# Try to use OpenAI's structured output if available
|
||||
use_openai_structured = False
|
||||
client = None
|
||||
|
||||
try:
|
||||
# Check if we have OpenAI API and the model supports structured output
|
||||
if self.llm and (self.llm.startswith('gpt-') or self.llm.startswith('o1-') or self.llm.startswith('o3-')):
|
||||
# Create a new client instance if custom parameters are provided
|
||||
if self.api_key or self.base_url:
|
||||
client = OpenAIClient(api_key=self.api_key, base_url=self.base_url)
|
||||
else:
|
||||
client = get_openai_client()
|
||||
use_openai_structured = True
|
||||
except:
|
||||
# If OpenAI client is not available, we'll use the LLM class
|
||||
pass
|
||||
|
||||
# Check if we have OpenAI API and the model supports structured output
|
||||
from ..llm import supports_structured_outputs
|
||||
if self.llm and supports_structured_outputs(self.llm):
|
||||
client = get_openai_client()
|
||||
use_openai_structured = True
|
||||
except:
|
||||
# If OpenAI client is not available, we'll use the LLM class
|
||||
pass
|
||||
|
||||
try:
|
||||
if use_openai_structured and client:
|
||||
# Use OpenAI's structured output for OpenAI models (backward compatibility)
|
||||
config = client.parse_structured_output(
|
||||
@@ -359,7 +356,7 @@ DO NOT use strings for tasks. Each task MUST be a complete object with all four
|
||||
api_key=self.api_key
|
||||
)
|
||||
|
||||
response_text = llm_instance.response(
|
||||
response_text = llm_instance.get_response(
|
||||
prompt=prompt,
|
||||
system_prompt="You are a helpful assistant designed to generate AI agent configurations.",
|
||||
output_pydantic=AutoAgentsConfig,
|
||||
@@ -503,4 +500,4 @@ DO NOT use strings for tasks. Each task MUST be a complete object with all four
|
||||
Creates tasks based on the instructions, then starts execution.
|
||||
Returns the task status and results dictionary.
|
||||
"""
|
||||
return super().start()
|
||||
return super().start()
|
||||
@@ -19,12 +19,54 @@ class LLMGuardrail:
|
||||
|
||||
Args:
|
||||
description: Natural language description of what to validate
|
||||
llm: The LLM instance to use for validation
|
||||
llm: The LLM instance to use for validation (can be string or LLM instance)
|
||||
"""
|
||||
self.description = description
|
||||
self.llm = llm
|
||||
self.llm = self._initialize_llm(llm)
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
def _initialize_llm(self, llm: Any) -> Any:
|
||||
"""Initialize the LLM instance from string identifier or existing instance.
|
||||
|
||||
Args:
|
||||
llm: String identifier, LLM instance, or None
|
||||
|
||||
Returns:
|
||||
LLM instance or None
|
||||
"""
|
||||
# Local import to avoid circular dependencies
|
||||
def _get_llm_class():
|
||||
from ..llm.llm import LLM
|
||||
return LLM
|
||||
|
||||
if llm is None:
|
||||
return None
|
||||
|
||||
# If it's already an LLM instance, return as-is
|
||||
if hasattr(llm, 'chat') or hasattr(llm, 'get_response') or callable(llm):
|
||||
return llm
|
||||
|
||||
# If it's a string, convert to LLM instance
|
||||
if isinstance(llm, str):
|
||||
try:
|
||||
# Handle string identifiers (both provider/model and simple names)
|
||||
return _get_llm_class()(model=llm)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to initialize LLM from string '{llm}': {str(e)}")
|
||||
return None
|
||||
|
||||
# If it's a dict, pass parameters to LLM
|
||||
if isinstance(llm, dict) and "model" in llm:
|
||||
try:
|
||||
return _get_llm_class()(**llm)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to initialize LLM from dict: {str(e)}")
|
||||
return None
|
||||
|
||||
# Unknown type
|
||||
self.logger.warning(f"Unknown LLM type: {type(llm)}, treating as-is")
|
||||
return llm
|
||||
|
||||
def __call__(self, task_output: TaskOutput) -> Tuple[bool, Union[str, TaskOutput]]:
|
||||
"""Validate the task output using the LLM.
|
||||
|
||||
|
||||
@@ -67,9 +67,6 @@ class Knowledge:
|
||||
|
||||
# Disable OpenAI API request logging
|
||||
logging.getLogger('openai').setLevel(logging.WARNING)
|
||||
|
||||
# Set root logger to warning to catch any uncategorized logs
|
||||
logging.getLogger().setLevel(logging.WARNING)
|
||||
|
||||
@cached_property
|
||||
def _deps(self):
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
# Ollama Handling in LLM.py
|
||||
|
||||
The `ollama_handled` flag is used in `llm.py` (not in `agent.py`) to track whether special Ollama model handling has been applied for tool results. Here's what I found:
|
||||
|
||||
## Where it's used in llm.py:
|
||||
- **Line 863**: `ollama_handled = False` - initialized at the beginning of the tool execution loop
|
||||
- **Line 864**: Special handling logic via `_handle_ollama_model()` method
|
||||
- **Line 907**: `ollama_handled = True` - set when Ollama special handling is applied
|
||||
- **Line 928 & 959**: Checks `if not ollama_handled` to avoid duplicate processing
|
||||
|
||||
## Why it's needed:
|
||||
The flag prevents duplicate processing when Ollama models require special handling. Some Ollama models return only the tool call JSON without processing the results, requiring a follow-up prompt to get the final answer. The `ollama_handled` flag ensures this special handling happens only once.
|
||||
|
||||
## Does this contradict the requirement?
|
||||
**No, it doesn't.** The issue states:
|
||||
|
||||
> "ollama/xxxx model is already handled via llm.py via litellm, so no special condition or handling required for ollama inside agent.py"
|
||||
|
||||
✅ **Current implementation is correct:**
|
||||
|
||||
- `agent.py` has NO ollama-specific conditions (only one documentation mention as an example)
|
||||
- All ollama handling is contained within `llm.py` as intended
|
||||
- The `ollama_handled` flag is an internal implementation detail within `llm.py`
|
||||
|
||||
The codebase already follows the desired pattern where ollama models are handled through the generic LLM interface via litellm without special conditions in agent.py.
|
||||
@@ -5,19 +5,16 @@ import os
|
||||
# Disable litellm telemetry before any imports
|
||||
os.environ["LITELLM_TELEMETRY"] = "False"
|
||||
|
||||
# Suppress all relevant logs at module level
|
||||
logging.getLogger("litellm").setLevel(logging.ERROR)
|
||||
logging.getLogger("openai").setLevel(logging.ERROR)
|
||||
logging.getLogger("httpx").setLevel(logging.ERROR)
|
||||
logging.getLogger("httpcore").setLevel(logging.ERROR)
|
||||
logging.getLogger("pydantic").setLevel(logging.ERROR)
|
||||
# Suppress all relevant logs at module level - consistent with main __init__.py
|
||||
logging.getLogger("litellm").setLevel(logging.WARNING)
|
||||
logging.getLogger("openai").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
||||
logging.getLogger("pydantic").setLevel(logging.WARNING)
|
||||
|
||||
# Suppress pydantic warnings
|
||||
warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")
|
||||
|
||||
# Configure logging to suppress all INFO messages
|
||||
logging.basicConfig(level=logging.WARNING)
|
||||
|
||||
# Import after suppressing warnings
|
||||
from .llm import LLM, LLMContextLengthExceededException
|
||||
from .openai_client import (
|
||||
@@ -32,6 +29,16 @@ from .openai_client import (
|
||||
ToolCall,
|
||||
process_stream_chunks
|
||||
)
|
||||
from .model_capabilities import (
|
||||
supports_structured_outputs,
|
||||
supports_streaming_with_tools
|
||||
)
|
||||
from .model_router import (
|
||||
ModelRouter,
|
||||
ModelProfile,
|
||||
TaskComplexity,
|
||||
create_routing_agent
|
||||
)
|
||||
|
||||
# Ensure telemetry is disabled after import as well
|
||||
try:
|
||||
@@ -52,5 +59,11 @@ __all__ = [
|
||||
"CompletionUsage",
|
||||
"ChatCompletion",
|
||||
"ToolCall",
|
||||
"process_stream_chunks"
|
||||
"process_stream_chunks",
|
||||
"supports_structured_outputs",
|
||||
"supports_streaming_with_tools",
|
||||
"ModelRouter",
|
||||
"ModelProfile",
|
||||
"TaskComplexity",
|
||||
"create_routing_agent"
|
||||
]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
Model capabilities configuration for different LLM providers.
|
||||
This module defines which models support specific features like structured outputs.
|
||||
"""
|
||||
|
||||
# Models that support OpenAI-style structured outputs (response_format with Pydantic models)
|
||||
MODELS_SUPPORTING_STRUCTURED_OUTPUTS = {
|
||||
# OpenAI models
|
||||
"gpt-4o",
|
||||
"gpt-4o-mini",
|
||||
"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-4o-mini-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.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to check
|
||||
|
||||
Returns:
|
||||
bool: True if the model supports structured outputs, False otherwise
|
||||
"""
|
||||
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
|
||||
|
||||
# First check if it's explicitly in the NOT supporting list
|
||||
if model_without_provider in MODELS_NOT_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return False
|
||||
|
||||
# Then check if it's in the supporting list
|
||||
if model_without_provider in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return True
|
||||
|
||||
# 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
|
||||
|
||||
# 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
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,348 @@
|
||||
"""
|
||||
Model Router for intelligent model selection based on task characteristics.
|
||||
|
||||
This module provides functionality to automatically select the most appropriate
|
||||
LLM model/provider based on task complexity, cost considerations, and model capabilities.
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from enum import IntEnum
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TaskComplexity(IntEnum):
|
||||
"""Enum for task complexity levels"""
|
||||
SIMPLE = 1 # Basic queries, math, factual questions
|
||||
MODERATE = 2 # Summarization, basic analysis
|
||||
COMPLEX = 3 # Code generation, deep reasoning
|
||||
VERY_COMPLEX = 4 # Multi-step reasoning, complex analysis
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelProfile:
|
||||
"""Profile for an LLM model with its characteristics"""
|
||||
name: str
|
||||
provider: str
|
||||
complexity_range: Tuple[TaskComplexity, TaskComplexity]
|
||||
cost_per_1k_tokens: float # Average of input/output costs
|
||||
strengths: List[str]
|
||||
capabilities: List[str]
|
||||
context_window: int
|
||||
supports_tools: bool = True
|
||||
supports_streaming: bool = True
|
||||
|
||||
|
||||
class ModelRouter:
|
||||
"""
|
||||
Intelligent model router that selects the best model based on task requirements.
|
||||
|
||||
This router implements a strategy pattern for model selection, considering:
|
||||
- Task complexity
|
||||
- Cost optimization
|
||||
- Model capabilities
|
||||
- Specific strengths for different task types
|
||||
"""
|
||||
|
||||
# Default model profiles - can be customized via configuration
|
||||
DEFAULT_MODELS = [
|
||||
# Lightweight/cheap models for simple tasks
|
||||
ModelProfile(
|
||||
name="gpt-4o-mini",
|
||||
provider="openai",
|
||||
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE),
|
||||
cost_per_1k_tokens=0.00075, # Average of $0.00015 input, $0.0006 output
|
||||
strengths=["speed", "cost-effective", "basic-reasoning"],
|
||||
capabilities=["text", "function-calling"],
|
||||
context_window=128000
|
||||
),
|
||||
ModelProfile(
|
||||
name="gemini/gemini-1.5-flash",
|
||||
provider="google",
|
||||
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE),
|
||||
cost_per_1k_tokens=0.000125, # Very cost-effective
|
||||
strengths=["speed", "cost-effective", "multimodal"],
|
||||
capabilities=["text", "vision", "function-calling"],
|
||||
context_window=1048576
|
||||
),
|
||||
ModelProfile(
|
||||
name="claude-3-haiku-20240307",
|
||||
provider="anthropic",
|
||||
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE),
|
||||
cost_per_1k_tokens=0.0008, # Average of $0.00025 input, $0.00125 output
|
||||
strengths=["speed", "instruction-following"],
|
||||
capabilities=["text", "function-calling"],
|
||||
context_window=200000
|
||||
),
|
||||
|
||||
# Mid-tier models for moderate complexity
|
||||
ModelProfile(
|
||||
name="gpt-4o",
|
||||
provider="openai",
|
||||
complexity_range=(TaskComplexity.MODERATE, TaskComplexity.COMPLEX),
|
||||
cost_per_1k_tokens=0.0075, # Average of $0.0025 input, $0.01 output
|
||||
strengths=["reasoning", "code-generation", "general-purpose"],
|
||||
capabilities=["text", "vision", "function-calling"],
|
||||
context_window=128000
|
||||
),
|
||||
ModelProfile(
|
||||
name="claude-3-5-sonnet-20241022",
|
||||
provider="anthropic",
|
||||
complexity_range=(TaskComplexity.MODERATE, TaskComplexity.VERY_COMPLEX),
|
||||
cost_per_1k_tokens=0.009, # Average of $0.003 input, $0.015 output
|
||||
strengths=["reasoning", "code-generation", "analysis", "writing"],
|
||||
capabilities=["text", "vision", "function-calling"],
|
||||
context_window=200000
|
||||
),
|
||||
|
||||
# High-end models for complex tasks
|
||||
ModelProfile(
|
||||
name="gemini/gemini-1.5-pro",
|
||||
provider="google",
|
||||
complexity_range=(TaskComplexity.COMPLEX, TaskComplexity.VERY_COMPLEX),
|
||||
cost_per_1k_tokens=0.00625, # Average of $0.00125 input, $0.005 output
|
||||
strengths=["reasoning", "long-context", "multimodal"],
|
||||
capabilities=["text", "vision", "function-calling"],
|
||||
context_window=2097152 # 2M context
|
||||
),
|
||||
ModelProfile(
|
||||
name="claude-3-opus-20240229",
|
||||
provider="anthropic",
|
||||
complexity_range=(TaskComplexity.COMPLEX, TaskComplexity.VERY_COMPLEX),
|
||||
cost_per_1k_tokens=0.045, # Average of $0.015 input, $0.075 output
|
||||
strengths=["deep-reasoning", "complex-analysis", "creative-writing"],
|
||||
capabilities=["text", "vision", "function-calling"],
|
||||
context_window=200000
|
||||
),
|
||||
ModelProfile(
|
||||
name="deepseek-chat",
|
||||
provider="deepseek",
|
||||
complexity_range=(TaskComplexity.COMPLEX, TaskComplexity.VERY_COMPLEX),
|
||||
cost_per_1k_tokens=0.0014, # Very cost-effective for capability
|
||||
strengths=["reasoning", "code-generation", "mathematics"],
|
||||
capabilities=["text", "function-calling"],
|
||||
context_window=128000
|
||||
),
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
models: Optional[List[ModelProfile]] = None,
|
||||
default_model: Optional[str] = None,
|
||||
cost_threshold: Optional[float] = None,
|
||||
preferred_providers: Optional[List[str]] = None
|
||||
):
|
||||
"""
|
||||
Initialize the ModelRouter.
|
||||
|
||||
Args:
|
||||
models: Custom list of model profiles to use
|
||||
default_model: Default model to use if no suitable model found
|
||||
cost_threshold: Maximum cost per 1k tokens to consider
|
||||
preferred_providers: List of preferred providers in order
|
||||
"""
|
||||
self.models = models or self.DEFAULT_MODELS
|
||||
self.default_model = default_model or os.getenv('OPENAI_MODEL_NAME', 'gpt-4o')
|
||||
self.cost_threshold = cost_threshold
|
||||
self.preferred_providers = preferred_providers or []
|
||||
|
||||
# Build lookup indices for efficient access
|
||||
self._model_by_name = {m.name: m for m in self.models}
|
||||
self._models_by_complexity = self._build_complexity_index()
|
||||
|
||||
def _build_complexity_index(self) -> Dict[TaskComplexity, List[ModelProfile]]:
|
||||
"""Build an index of models by complexity level"""
|
||||
index = {level: [] for level in TaskComplexity}
|
||||
|
||||
for model in self.models:
|
||||
min_complexity, max_complexity = model.complexity_range
|
||||
for level in TaskComplexity:
|
||||
if min_complexity.value <= level.value <= max_complexity.value:
|
||||
index[level].append(model)
|
||||
|
||||
return index
|
||||
|
||||
def analyze_task_complexity(
|
||||
self,
|
||||
task_description: str,
|
||||
tools_required: Optional[List[str]] = None,
|
||||
context_size: Optional[int] = None
|
||||
) -> TaskComplexity:
|
||||
"""
|
||||
Analyze task description to determine complexity level.
|
||||
|
||||
This is a simple heuristic-based approach. In production, this could be
|
||||
replaced with a more sophisticated ML-based classifier.
|
||||
"""
|
||||
description_lower = task_description.lower()
|
||||
|
||||
# Keywords indicating different complexity levels
|
||||
simple_keywords = [
|
||||
"calculate", "compute", "what is", "define", "list", "count",
|
||||
"simple", "basic", "check", "verify", "yes or no", "true or false"
|
||||
]
|
||||
|
||||
moderate_keywords = [
|
||||
"summarize", "explain", "compare", "describe", "analyze briefly",
|
||||
"find", "search", "extract", "classify", "categorize"
|
||||
]
|
||||
|
||||
complex_keywords = [
|
||||
"implement", "code", "develop", "design", "create algorithm",
|
||||
"optimize", "debug", "refactor", "architect", "solve"
|
||||
]
|
||||
|
||||
very_complex_keywords = [
|
||||
"multi-step", "comprehensive analysis", "deep dive", "research",
|
||||
"strategic", "framework", "system design", "proof", "theorem"
|
||||
]
|
||||
|
||||
# Check for keyword matches
|
||||
if any(keyword in description_lower for keyword in very_complex_keywords):
|
||||
return TaskComplexity.VERY_COMPLEX
|
||||
elif any(keyword in description_lower for keyword in complex_keywords):
|
||||
return TaskComplexity.COMPLEX
|
||||
elif any(keyword in description_lower for keyword in moderate_keywords):
|
||||
return TaskComplexity.MODERATE
|
||||
elif any(keyword in description_lower for keyword in simple_keywords):
|
||||
return TaskComplexity.SIMPLE
|
||||
|
||||
# Consider tool requirements
|
||||
if tools_required and len(tools_required) > 3:
|
||||
return TaskComplexity.COMPLEX
|
||||
|
||||
# Consider context size requirements
|
||||
if context_size and context_size > 50000:
|
||||
return TaskComplexity.COMPLEX
|
||||
|
||||
# Default to moderate
|
||||
return TaskComplexity.MODERATE
|
||||
|
||||
def select_model(
|
||||
self,
|
||||
task_description: str,
|
||||
required_capabilities: Optional[List[str]] = None,
|
||||
tools_required: Optional[List[str]] = None,
|
||||
context_size: Optional[int] = None,
|
||||
budget_conscious: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Select the most appropriate model for a given task.
|
||||
|
||||
Args:
|
||||
task_description: Description of the task to perform
|
||||
required_capabilities: List of required capabilities (e.g., ["vision", "function-calling"])
|
||||
tools_required: List of tools that will be used
|
||||
context_size: Estimated context size needed
|
||||
budget_conscious: Whether to optimize for cost
|
||||
|
||||
Returns:
|
||||
Model name string to use
|
||||
"""
|
||||
# Analyze task complexity
|
||||
complexity = self.analyze_task_complexity(task_description, tools_required, context_size)
|
||||
|
||||
# Get candidate models for this complexity level
|
||||
candidates = self._models_by_complexity.get(complexity, [])
|
||||
|
||||
if not candidates:
|
||||
logger.warning(f"No models found for complexity {complexity}, using default")
|
||||
return self.default_model
|
||||
|
||||
# Filter by required capabilities
|
||||
if required_capabilities:
|
||||
candidates = [
|
||||
m for m in candidates
|
||||
if all(cap in m.capabilities for cap in required_capabilities)
|
||||
]
|
||||
|
||||
# Filter by tool support if needed
|
||||
if tools_required:
|
||||
candidates = [m for m in candidates if m.supports_tools]
|
||||
|
||||
# Filter by context window if specified
|
||||
if context_size:
|
||||
candidates = [m for m in candidates if m.context_window >= context_size]
|
||||
|
||||
# Filter by cost threshold if specified
|
||||
if self.cost_threshold:
|
||||
candidates = [m for m in candidates if m.cost_per_1k_tokens <= self.cost_threshold]
|
||||
|
||||
if not candidates:
|
||||
logger.warning("No models meet all criteria, using default")
|
||||
return self.default_model
|
||||
|
||||
# Sort by selection criteria
|
||||
if budget_conscious:
|
||||
# Sort by cost (ascending)
|
||||
candidates.sort(key=lambda m: m.cost_per_1k_tokens)
|
||||
else:
|
||||
# Sort by capability (descending complexity)
|
||||
candidates.sort(key=lambda m: m.complexity_range[1].value, reverse=True)
|
||||
|
||||
# Apply provider preferences
|
||||
if self.preferred_providers:
|
||||
for provider in self.preferred_providers:
|
||||
for model in candidates:
|
||||
if model.provider == provider:
|
||||
logger.info(f"Selected model: {model.name} (complexity: {complexity}, cost: ${model.cost_per_1k_tokens}/1k tokens)")
|
||||
return model.name
|
||||
|
||||
# Return the best candidate
|
||||
selected = candidates[0]
|
||||
logger.info(f"Selected model: {selected.name} (complexity: {complexity}, cost: ${selected.cost_per_1k_tokens}/1k tokens)")
|
||||
return selected.name
|
||||
|
||||
def get_model_info(self, model_name: str) -> Optional[ModelProfile]:
|
||||
"""Get profile information for a specific model"""
|
||||
return self._model_by_name.get(model_name)
|
||||
|
||||
def estimate_cost(self, model_name: str, estimated_tokens: int) -> float:
|
||||
"""Estimate the cost for a given model and token count"""
|
||||
model = self._model_by_name.get(model_name)
|
||||
if not model:
|
||||
return 0.0
|
||||
return (model.cost_per_1k_tokens * estimated_tokens) / 1000
|
||||
|
||||
|
||||
def create_routing_agent(
|
||||
models: Optional[List[str]] = None,
|
||||
router: Optional[ModelRouter] = None,
|
||||
**agent_kwargs
|
||||
) -> 'Agent':
|
||||
"""
|
||||
Create a specialized routing agent that can select models dynamically.
|
||||
|
||||
Args:
|
||||
models: List of model names to route between
|
||||
router: Custom ModelRouter instance
|
||||
**agent_kwargs: Additional arguments to pass to Agent constructor
|
||||
|
||||
Returns:
|
||||
Agent configured for model routing
|
||||
"""
|
||||
from ..agent import Agent
|
||||
|
||||
if not router:
|
||||
router = ModelRouter()
|
||||
|
||||
routing_agent = Agent(
|
||||
name=agent_kwargs.pop('name', 'ModelRouter'),
|
||||
role=agent_kwargs.pop('role', 'Intelligent Model Router'),
|
||||
goal=agent_kwargs.pop('goal', 'Select the most appropriate AI model based on task requirements'),
|
||||
backstory=agent_kwargs.pop('backstory',
|
||||
'I analyze tasks and route them to the most suitable AI model, '
|
||||
'optimizing for performance, cost, and capability requirements.'
|
||||
),
|
||||
**agent_kwargs
|
||||
)
|
||||
|
||||
# TODO: Consider creating a proper RoutingAgent subclass instead of setting private attributes
|
||||
# For now, store the router on the agent for use in execution
|
||||
routing_agent._model_router = router
|
||||
routing_agent._available_models = models or [m.name for m in router.models]
|
||||
|
||||
return routing_agent
|
||||
@@ -890,45 +890,8 @@ class OpenAIClient:
|
||||
"content": results_str
|
||||
})
|
||||
|
||||
# Check if we should continue (for tools like sequential thinking)
|
||||
should_continue = False
|
||||
for tool_call in tool_calls:
|
||||
# Handle both ToolCall dataclass and OpenAI object
|
||||
if isinstance(tool_call, ToolCall):
|
||||
function_name = tool_call.function["name"]
|
||||
arguments = json.loads(tool_call.function["arguments"])
|
||||
else:
|
||||
function_name = tool_call.function.name
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
|
||||
# For sequential thinking tool, check if nextThoughtNeeded is True
|
||||
if function_name == "sequentialthinking" and arguments.get("nextThoughtNeeded", False):
|
||||
should_continue = True
|
||||
break
|
||||
|
||||
if not should_continue:
|
||||
# Get final response after tool calls
|
||||
if stream:
|
||||
final_response = self.process_stream_response(
|
||||
messages=messages,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
tools=formatted_tools,
|
||||
start_time=start_time,
|
||||
console=console,
|
||||
display_fn=display_fn,
|
||||
reasoning_steps=reasoning_steps,
|
||||
**kwargs
|
||||
)
|
||||
else:
|
||||
final_response = self.create_completion(
|
||||
messages=messages,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
stream=False,
|
||||
**kwargs
|
||||
)
|
||||
break
|
||||
# Continue the loop to allow more tool calls
|
||||
# The model will see tool results and can make additional tool calls
|
||||
|
||||
iteration_count += 1
|
||||
else:
|
||||
@@ -1067,45 +1030,8 @@ class OpenAIClient:
|
||||
"content": results_str
|
||||
})
|
||||
|
||||
# Check if we should continue (for tools like sequential thinking)
|
||||
should_continue = False
|
||||
for tool_call in tool_calls:
|
||||
# Handle both ToolCall dataclass and OpenAI object
|
||||
if isinstance(tool_call, ToolCall):
|
||||
function_name = tool_call.function["name"]
|
||||
arguments = json.loads(tool_call.function["arguments"])
|
||||
else:
|
||||
function_name = tool_call.function.name
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
|
||||
# For sequential thinking tool, check if nextThoughtNeeded is True
|
||||
if function_name == "sequentialthinking" and arguments.get("nextThoughtNeeded", False):
|
||||
should_continue = True
|
||||
break
|
||||
|
||||
if not should_continue:
|
||||
# Get final response after tool calls
|
||||
if stream:
|
||||
final_response = await self.process_stream_response_async(
|
||||
messages=messages,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
tools=formatted_tools,
|
||||
start_time=start_time,
|
||||
console=console,
|
||||
display_fn=display_fn,
|
||||
reasoning_steps=reasoning_steps,
|
||||
**kwargs
|
||||
)
|
||||
else:
|
||||
final_response = await self.acreate_completion(
|
||||
messages=messages,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
stream=False,
|
||||
**kwargs
|
||||
)
|
||||
break
|
||||
# Continue the loop to allow more tool calls
|
||||
# The model will see tool results and can make additional tool calls
|
||||
|
||||
iteration_count += 1
|
||||
else:
|
||||
|
||||
@@ -9,32 +9,15 @@ from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
from rich.markdown import Markdown
|
||||
from rich.logging import RichHandler
|
||||
from rich.live import Live
|
||||
import asyncio
|
||||
|
||||
# # Configure root logger
|
||||
# logging.basicConfig(level=logging.WARNING)
|
||||
|
||||
# Suppress litellm logs
|
||||
# Logging is already configured in __init__.py, just clean up handlers for litellm
|
||||
logging.getLogger("litellm").handlers = []
|
||||
logging.getLogger("litellm.utils").handlers = []
|
||||
logging.getLogger("litellm").propagate = False
|
||||
logging.getLogger("litellm.utils").propagate = False
|
||||
|
||||
LOGLEVEL = os.environ.get('LOGLEVEL', 'INFO').upper()
|
||||
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, LOGLEVEL, logging.INFO),
|
||||
format="%(asctime)s %(filename)s:%(lineno)d %(levelname)s %(message)s",
|
||||
datefmt="[%X]",
|
||||
handlers=[RichHandler(rich_tracebacks=True)]
|
||||
)
|
||||
|
||||
# Add these lines to suppress markdown parser debug logs
|
||||
logging.getLogger('markdown_it').setLevel(logging.WARNING)
|
||||
logging.getLogger('rich.markdown').setLevel(logging.WARNING)
|
||||
|
||||
# Global list to store error logs
|
||||
error_logs = []
|
||||
|
||||
@@ -79,6 +62,30 @@ def register_approval_callback(callback_fn):
|
||||
global approval_callback
|
||||
approval_callback = callback_fn
|
||||
|
||||
def execute_sync_callback(display_type: str, **kwargs):
|
||||
"""Execute synchronous callback for a given display type without displaying anything.
|
||||
|
||||
This function is used to trigger callbacks even when verbose=False.
|
||||
|
||||
Args:
|
||||
display_type (str): Type of display event
|
||||
**kwargs: Arguments to pass to the callback function
|
||||
"""
|
||||
if display_type in sync_display_callbacks:
|
||||
callback = sync_display_callbacks[display_type]
|
||||
import inspect
|
||||
sig = inspect.signature(callback)
|
||||
|
||||
# Filter kwargs to what the callback accepts to maintain backward compatibility
|
||||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
# Callback accepts **kwargs, so pass all arguments
|
||||
supported_kwargs = kwargs
|
||||
else:
|
||||
# Only pass arguments that the callback signature supports
|
||||
supported_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
callback(**supported_kwargs)
|
||||
|
||||
async def execute_callback(display_type: str, **kwargs):
|
||||
"""Execute both sync and async callbacks for a given display type.
|
||||
|
||||
@@ -86,16 +93,38 @@ async def execute_callback(display_type: str, **kwargs):
|
||||
display_type (str): Type of display event
|
||||
**kwargs: Arguments to pass to the callback functions
|
||||
"""
|
||||
import inspect
|
||||
|
||||
# Execute synchronous callback if registered
|
||||
if display_type in sync_display_callbacks:
|
||||
callback = sync_display_callbacks[display_type]
|
||||
sig = inspect.signature(callback)
|
||||
|
||||
# Filter kwargs to what the callback accepts to maintain backward compatibility
|
||||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
# Callback accepts **kwargs, so pass all arguments
|
||||
supported_kwargs = kwargs
|
||||
else:
|
||||
# Only pass arguments that the callback signature supports
|
||||
supported_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, lambda: callback(**kwargs))
|
||||
await loop.run_in_executor(None, lambda: callback(**supported_kwargs))
|
||||
|
||||
# Execute asynchronous callback if registered
|
||||
if display_type in async_display_callbacks:
|
||||
callback = async_display_callbacks[display_type]
|
||||
await callback(**kwargs)
|
||||
sig = inspect.signature(callback)
|
||||
|
||||
# Filter kwargs to what the callback accepts to maintain backward compatibility
|
||||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
# Callback accepts **kwargs, so pass all arguments
|
||||
supported_kwargs = kwargs
|
||||
else:
|
||||
# Only pass arguments that the callback signature supports
|
||||
supported_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
await callback(**supported_kwargs)
|
||||
|
||||
def _clean_display_content(content: str, max_length: int = 20000) -> str:
|
||||
"""Helper function to clean and truncate content for display."""
|
||||
@@ -118,7 +147,7 @@ def _clean_display_content(content: str, max_length: int = 20000) -> str:
|
||||
|
||||
return content.strip()
|
||||
|
||||
def display_interaction(message, response, markdown=True, generation_time=None, console=None):
|
||||
def display_interaction(message, response, markdown=True, generation_time=None, console=None, agent_name=None, agent_role=None, agent_tools=None, task_name=None, task_description=None, task_id=None):
|
||||
"""Synchronous version of display_interaction."""
|
||||
if console is None:
|
||||
console = Console()
|
||||
@@ -130,15 +159,20 @@ def display_interaction(message, response, markdown=True, generation_time=None,
|
||||
message = _clean_display_content(str(message))
|
||||
response = _clean_display_content(str(response))
|
||||
|
||||
# Execute synchronous callback if registered
|
||||
if 'interaction' in sync_display_callbacks:
|
||||
sync_display_callbacks['interaction'](
|
||||
message=message,
|
||||
response=response,
|
||||
markdown=markdown,
|
||||
generation_time=generation_time
|
||||
)
|
||||
|
||||
# Execute synchronous callbacks
|
||||
execute_sync_callback(
|
||||
'interaction',
|
||||
message=message,
|
||||
response=response,
|
||||
markdown=markdown,
|
||||
generation_time=generation_time,
|
||||
agent_name=agent_name,
|
||||
agent_role=agent_role,
|
||||
agent_tools=agent_tools,
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id
|
||||
)
|
||||
# Rest of the display logic...
|
||||
if generation_time:
|
||||
console.print(Text(f"Response generated in {generation_time:.1f}s", style="dim"))
|
||||
@@ -157,9 +191,8 @@ def display_self_reflection(message: str, console=None):
|
||||
console = Console()
|
||||
message = _clean_display_content(str(message))
|
||||
|
||||
# Execute callback if registered
|
||||
if 'self_reflection' in sync_display_callbacks:
|
||||
sync_display_callbacks['self_reflection'](message=message)
|
||||
# Execute synchronous callbacks
|
||||
execute_sync_callback('self_reflection', message=message)
|
||||
|
||||
console.print(Panel.fit(Text(message, style="bold yellow"), title="Self Reflection", border_style="magenta"))
|
||||
|
||||
@@ -170,9 +203,8 @@ def display_instruction(message: str, console=None, agent_name: str = None, agen
|
||||
console = Console()
|
||||
message = _clean_display_content(str(message))
|
||||
|
||||
# Execute callback if registered
|
||||
if 'instruction' in sync_display_callbacks:
|
||||
sync_display_callbacks['instruction'](message=message)
|
||||
# Execute synchronous callbacks
|
||||
execute_sync_callback('instruction', message=message, agent_name=agent_name, agent_role=agent_role, agent_tools=agent_tools)
|
||||
|
||||
# Display agent info if available
|
||||
if agent_name:
|
||||
@@ -198,9 +230,8 @@ def display_tool_call(message: str, console=None):
|
||||
message = _clean_display_content(str(message))
|
||||
logging.debug(f"Cleaned message in display_tool_call: {repr(message)}")
|
||||
|
||||
# Execute callback if registered
|
||||
if 'tool_call' in sync_display_callbacks:
|
||||
sync_display_callbacks['tool_call'](message=message)
|
||||
# Execute synchronous callbacks
|
||||
execute_sync_callback('tool_call', message=message)
|
||||
|
||||
console.print(Panel.fit(Text(message, style="bold cyan"), title="Tool Call", border_style="green"))
|
||||
|
||||
@@ -211,9 +242,8 @@ def display_error(message: str, console=None):
|
||||
console = Console()
|
||||
message = _clean_display_content(str(message))
|
||||
|
||||
# Execute callback if registered
|
||||
if 'error' in sync_display_callbacks:
|
||||
sync_display_callbacks['error'](message=message)
|
||||
# Execute synchronous callbacks
|
||||
execute_sync_callback('error', message=message)
|
||||
|
||||
console.print(Panel.fit(Text(message, style="bold red"), title="Error", border_style="red"))
|
||||
error_logs.append(message)
|
||||
@@ -230,17 +260,13 @@ def display_generating(content: str = "", start_time: Optional[float] = None):
|
||||
|
||||
content = _clean_display_content(str(content))
|
||||
|
||||
# Execute callback if registered
|
||||
if 'generating' in sync_display_callbacks:
|
||||
sync_display_callbacks['generating'](
|
||||
content=content,
|
||||
elapsed_time=elapsed_str.strip() if elapsed_str else None
|
||||
)
|
||||
# Execute synchronous callbacks
|
||||
execute_sync_callback('generating', content=content, elapsed_time=elapsed_str.strip() if elapsed_str else None)
|
||||
|
||||
return Panel(Markdown(content), title=f"Generating...{elapsed_str}", border_style="green")
|
||||
|
||||
# Async versions with 'a' prefix
|
||||
async def adisplay_interaction(message, response, markdown=True, generation_time=None, console=None):
|
||||
async def adisplay_interaction(message, response, markdown=True, generation_time=None, console=None, agent_name=None, agent_role=None, agent_tools=None, task_name=None, task_description=None, task_id=None):
|
||||
"""Async version of display_interaction."""
|
||||
if console is None:
|
||||
console = Console()
|
||||
@@ -258,7 +284,13 @@ async def adisplay_interaction(message, response, markdown=True, generation_time
|
||||
message=message,
|
||||
response=response,
|
||||
markdown=markdown,
|
||||
generation_time=generation_time
|
||||
generation_time=generation_time,
|
||||
agent_name=agent_name,
|
||||
agent_role=agent_role,
|
||||
agent_tools=agent_tools,
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id
|
||||
)
|
||||
|
||||
# Rest of the display logic...
|
||||
@@ -280,8 +312,8 @@ async def adisplay_self_reflection(message: str, console=None):
|
||||
console = Console()
|
||||
message = _clean_display_content(str(message))
|
||||
|
||||
if 'self_reflection' in async_display_callbacks:
|
||||
await async_display_callbacks['self_reflection'](message=message)
|
||||
# Execute callbacks
|
||||
await execute_callback('self_reflection', message=message)
|
||||
|
||||
console.print(Panel.fit(Text(message, style="bold yellow"), title="Self Reflection", border_style="magenta"))
|
||||
|
||||
@@ -293,8 +325,8 @@ async def adisplay_instruction(message: str, console=None, agent_name: str = Non
|
||||
console = Console()
|
||||
message = _clean_display_content(str(message))
|
||||
|
||||
if 'instruction' in async_display_callbacks:
|
||||
await async_display_callbacks['instruction'](message=message)
|
||||
# Execute callbacks
|
||||
await execute_callback('instruction', message=message, agent_name=agent_name, agent_role=agent_role, agent_tools=agent_tools)
|
||||
|
||||
# Display agent info if available
|
||||
if agent_name:
|
||||
@@ -321,8 +353,8 @@ async def adisplay_tool_call(message: str, console=None):
|
||||
message = _clean_display_content(str(message))
|
||||
logging.debug(f"Cleaned message in adisplay_tool_call: {repr(message)}")
|
||||
|
||||
if 'tool_call' in async_display_callbacks:
|
||||
await async_display_callbacks['tool_call'](message=message)
|
||||
# Execute callbacks
|
||||
await execute_callback('tool_call', message=message)
|
||||
|
||||
console.print(Panel.fit(Text(message, style="bold cyan"), title="Tool Call", border_style="green"))
|
||||
|
||||
@@ -334,8 +366,8 @@ async def adisplay_error(message: str, console=None):
|
||||
console = Console()
|
||||
message = _clean_display_content(str(message))
|
||||
|
||||
if 'error' in async_display_callbacks:
|
||||
await async_display_callbacks['error'](message=message)
|
||||
# Execute callbacks
|
||||
await execute_callback('error', message=message)
|
||||
|
||||
console.print(Panel.fit(Text(message, style="bold red"), title="Error", border_style="red"))
|
||||
error_logs.append(message)
|
||||
@@ -353,11 +385,8 @@ async def adisplay_generating(content: str = "", start_time: Optional[float] = N
|
||||
|
||||
content = _clean_display_content(str(content))
|
||||
|
||||
if 'generating' in async_display_callbacks:
|
||||
await async_display_callbacks['generating'](
|
||||
content=content,
|
||||
elapsed_time=elapsed_str.strip() if elapsed_str else None
|
||||
)
|
||||
# Execute callbacks
|
||||
await execute_callback('generating', content=content, elapsed_time=elapsed_str.strip() if elapsed_str else None)
|
||||
|
||||
return Panel(Markdown(content), title=f"Generating...{elapsed_str}", border_style="green")
|
||||
|
||||
|
||||
@@ -187,15 +187,45 @@ class MCP:
|
||||
self.timeout = timeout
|
||||
self.debug = debug
|
||||
|
||||
# Check if this is an SSE URL
|
||||
# Check if this is an HTTP URL
|
||||
if isinstance(command_or_string, str) and re.match(r'^https?://', command_or_string):
|
||||
# Import the SSE client implementation
|
||||
from .mcp_sse import SSEMCPClient
|
||||
self.sse_client = SSEMCPClient(command_or_string, debug=debug, timeout=timeout)
|
||||
self._tools = list(self.sse_client.tools)
|
||||
self.is_sse = True
|
||||
self.is_npx = False
|
||||
return
|
||||
# Determine transport type based on URL or kwargs
|
||||
if command_or_string.endswith('/sse') and 'transport_type' not in kwargs:
|
||||
# Legacy SSE URL - use SSE transport for backward compatibility
|
||||
from .mcp_sse import SSEMCPClient
|
||||
self.sse_client = SSEMCPClient(command_or_string, debug=debug, timeout=timeout)
|
||||
self._tools = list(self.sse_client.tools)
|
||||
self.is_sse = True
|
||||
self.is_http_stream = False
|
||||
self.is_npx = False
|
||||
return
|
||||
else:
|
||||
# Use HTTP Stream transport for all other HTTP URLs
|
||||
from .mcp_http_stream import HTTPStreamMCPClient
|
||||
# Extract transport options from kwargs
|
||||
transport_options = {}
|
||||
if 'responseMode' in kwargs:
|
||||
transport_options['responseMode'] = kwargs.pop('responseMode')
|
||||
if 'headers' in kwargs:
|
||||
transport_options['headers'] = kwargs.pop('headers')
|
||||
if 'cors' in kwargs:
|
||||
transport_options['cors'] = kwargs.pop('cors')
|
||||
if 'session' in kwargs:
|
||||
transport_options['session'] = kwargs.pop('session')
|
||||
if 'resumability' in kwargs:
|
||||
transport_options['resumability'] = kwargs.pop('resumability')
|
||||
|
||||
self.http_stream_client = HTTPStreamMCPClient(
|
||||
command_or_string,
|
||||
debug=debug,
|
||||
timeout=timeout,
|
||||
options=transport_options
|
||||
)
|
||||
self._tools = list(self.http_stream_client.tools)
|
||||
self.is_sse = False
|
||||
self.is_http_stream = True
|
||||
self.is_npx = False
|
||||
return
|
||||
|
||||
# Handle the single string format for stdio client
|
||||
if isinstance(command_or_string, str) and args is None:
|
||||
@@ -219,6 +249,7 @@ class MCP:
|
||||
|
||||
# Set up stdio client
|
||||
self.is_sse = False
|
||||
self.is_http_stream = False
|
||||
|
||||
# Ensure UTF-8 encoding in environment for Docker compatibility
|
||||
env = kwargs.get('env', {})
|
||||
@@ -275,6 +306,9 @@ class MCP:
|
||||
"""
|
||||
if self.is_sse:
|
||||
return list(self.sse_client.tools)
|
||||
|
||||
if self.is_http_stream:
|
||||
return list(self.http_stream_client.tools)
|
||||
|
||||
tool_functions = []
|
||||
|
||||
@@ -448,6 +482,10 @@ class MCP:
|
||||
if self.is_sse and hasattr(self, 'sse_client') and self.sse_client.tools:
|
||||
# Return all tools from SSE client
|
||||
return self.sse_client.to_openai_tools()
|
||||
|
||||
if self.is_http_stream and hasattr(self, 'http_stream_client') and self.http_stream_client.tools:
|
||||
# Return all tools from HTTP Stream client
|
||||
return self.http_stream_client.to_openai_tools()
|
||||
|
||||
# For simplicity, we'll convert the first tool only if multiple exist
|
||||
# More complex implementations could handle multiple tools
|
||||
|
||||
@@ -0,0 +1,466 @@
|
||||
"""
|
||||
HTTP Stream client implementation for MCP (Model Context Protocol).
|
||||
This module provides the necessary classes and functions to connect to an MCP server
|
||||
over HTTP Stream transport, implementing the Streamable HTTP transport protocol.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import inspect
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from typing import List, Dict, Any, Optional, Callable, Iterable, Union
|
||||
from urllib.parse import urlparse, urljoin
|
||||
|
||||
from mcp import ClientSession
|
||||
try:
|
||||
import aiohttp
|
||||
except ImportError:
|
||||
raise ImportError("aiohttp is required for HTTP Stream transport. Install with: pip install praisonaiagents[mcp]")
|
||||
|
||||
logger = logging.getLogger("mcp-http-stream")
|
||||
|
||||
# Global event loop for async operations
|
||||
_event_loop = None
|
||||
|
||||
def get_event_loop():
|
||||
"""Get or create a global event loop."""
|
||||
global _event_loop
|
||||
if _event_loop is None or _event_loop.is_closed():
|
||||
_event_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(_event_loop)
|
||||
return _event_loop
|
||||
|
||||
|
||||
class HTTPStreamMCPTool:
|
||||
"""A wrapper for an MCP tool that can be used with praisonaiagents."""
|
||||
|
||||
def __init__(self, name: str, description: str, session: ClientSession, input_schema: Optional[Dict[str, Any]] = None, timeout: int = 60):
|
||||
self.name = name
|
||||
self.__name__ = name # Required for Agent to recognize it as a tool
|
||||
self.__qualname__ = name # Required for Agent to recognize it as a tool
|
||||
self.__doc__ = description # Required for Agent to recognize it as a tool
|
||||
self.description = description
|
||||
self.session = session
|
||||
self.input_schema = input_schema or {}
|
||||
self.timeout = timeout
|
||||
|
||||
# Create a signature based on input schema
|
||||
params = []
|
||||
if input_schema and 'properties' in input_schema:
|
||||
for param_name, prop_schema in input_schema['properties'].items():
|
||||
# Determine type annotation based on schema
|
||||
prop_type = prop_schema.get('type', 'string') if isinstance(prop_schema, dict) else 'string'
|
||||
if prop_type == 'string':
|
||||
annotation = str
|
||||
elif prop_type == 'integer':
|
||||
annotation = int
|
||||
elif prop_type == 'number':
|
||||
annotation = float
|
||||
elif prop_type == 'boolean':
|
||||
annotation = bool
|
||||
elif prop_type == 'array':
|
||||
annotation = list
|
||||
elif prop_type == 'object':
|
||||
annotation = dict
|
||||
else:
|
||||
annotation = Any
|
||||
|
||||
params.append(
|
||||
inspect.Parameter(
|
||||
name=param_name,
|
||||
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
default=inspect.Parameter.empty if param_name in input_schema.get('required', []) else None,
|
||||
annotation=annotation
|
||||
)
|
||||
)
|
||||
|
||||
self.__signature__ = inspect.Signature(params)
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
"""Synchronous wrapper for the async call."""
|
||||
logger.debug(f"Tool {self.name} called with args: {kwargs}")
|
||||
|
||||
# Use the global event loop
|
||||
loop = get_event_loop()
|
||||
|
||||
# Run the async call in the event loop
|
||||
future = asyncio.run_coroutine_threadsafe(self._async_call(**kwargs), loop)
|
||||
try:
|
||||
# Wait for the result with a timeout
|
||||
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 with the provided arguments."""
|
||||
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):
|
||||
"""
|
||||
Fix array schemas by adding missing 'items' attribute required by OpenAI.
|
||||
|
||||
This ensures compatibility with OpenAI's function calling format which
|
||||
requires array types to specify the type of items they contain.
|
||||
|
||||
Args:
|
||||
schema: The schema dictionary to fix
|
||||
|
||||
Returns:
|
||||
dict: The fixed schema
|
||||
"""
|
||||
if not isinstance(schema, dict):
|
||||
return schema
|
||||
|
||||
# Create a copy to avoid modifying the original
|
||||
fixed_schema = schema.copy()
|
||||
|
||||
# Fix array types at the current level
|
||||
if fixed_schema.get("type") == "array" and "items" not in fixed_schema:
|
||||
# Add a default items schema for arrays without it
|
||||
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
|
||||
|
||||
# Fix items schema if it exists
|
||||
if "items" in fixed_schema:
|
||||
fixed_schema["items"] = self._fix_array_schemas(fixed_schema["items"])
|
||||
|
||||
return fixed_schema
|
||||
|
||||
def to_openai_tool(self):
|
||||
"""Convert the tool to OpenAI format."""
|
||||
# Fix array schemas to include 'items' attribute
|
||||
fixed_schema = self._fix_array_schemas(self.input_schema)
|
||||
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"parameters": fixed_schema
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class HTTPStreamTransport:
|
||||
"""
|
||||
HTTP Stream Transport implementation for MCP.
|
||||
|
||||
This transport provides a single endpoint for all MCP communication,
|
||||
supporting both batch (JSON) and streaming (SSE) response modes.
|
||||
"""
|
||||
|
||||
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.headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Accept': 'application/json, text/event-stream'
|
||||
}
|
||||
if session_id:
|
||||
self.headers['Mcp-Session-Id'] = session_id
|
||||
|
||||
# Add custom headers if provided
|
||||
if 'headers' in self.options:
|
||||
self.headers.update(self.options['headers'])
|
||||
|
||||
self._session = None
|
||||
self._sse_task = None
|
||||
self._message_queue = asyncio.Queue()
|
||||
self._pending_requests = {}
|
||||
self._closing = False
|
||||
|
||||
async def __aenter__(self):
|
||||
self._session = aiohttp.ClientSession()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
# Set closing flag to stop listener gracefully
|
||||
self._closing = True
|
||||
|
||||
if self._sse_task:
|
||||
self._sse_task.cancel()
|
||||
try:
|
||||
await self._sse_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
if self._session:
|
||||
await self._session.close()
|
||||
|
||||
async def send_request(self, request: Dict[str, Any]) -> Union[Dict[str, Any], None]:
|
||||
"""Send a request to the HTTP Stream endpoint."""
|
||||
if not self._session:
|
||||
raise RuntimeError("Transport not initialized. Use async context manager.")
|
||||
|
||||
try:
|
||||
async with self._session.post(self.base_url, json=request, headers=self.headers) as response:
|
||||
# Update session ID if provided in response
|
||||
if 'Mcp-Session-Id' in response.headers:
|
||||
self.session_id = response.headers['Mcp-Session-Id']
|
||||
self.headers['Mcp-Session-Id'] = self.session_id
|
||||
|
||||
# Handle different response types
|
||||
content_type = response.headers.get('Content-Type', '')
|
||||
|
||||
if 'text/event-stream' in content_type:
|
||||
# Stream mode - process SSE events
|
||||
return await self._process_sse_response(response)
|
||||
else:
|
||||
# Batch mode - return JSON response
|
||||
return await response.json()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending request: {e}")
|
||||
raise
|
||||
|
||||
async def _process_sse_response(self, response):
|
||||
"""Process SSE response stream."""
|
||||
buffer = ""
|
||||
async for chunk in response.content:
|
||||
buffer += chunk.decode('utf-8')
|
||||
|
||||
# Process complete SSE events
|
||||
while "\n\n" in buffer:
|
||||
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
|
||||
|
||||
if data:
|
||||
try:
|
||||
message = json.loads(data)
|
||||
# Process the message
|
||||
if 'id' in message and message['id'] in self._pending_requests:
|
||||
# This is a response to a pending request
|
||||
self._pending_requests[message['id']].set_result(message)
|
||||
else:
|
||||
# This is a server-initiated message
|
||||
await self._message_queue.put(message)
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"Failed to parse SSE event: {data}")
|
||||
|
||||
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."""
|
||||
while True:
|
||||
try:
|
||||
# Check if we should stop
|
||||
if hasattr(self, '_closing') and self._closing:
|
||||
break
|
||||
|
||||
url = self.base_url
|
||||
if self.session_id:
|
||||
# Add session as query parameter for SSE connection
|
||||
url = f"{url}?session={self.session_id}"
|
||||
|
||||
headers = {
|
||||
'Accept': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache'
|
||||
}
|
||||
if self.session_id:
|
||||
headers['Mcp-Session-Id'] = self.session_id
|
||||
|
||||
async with self._session.get(url, headers=headers) as response:
|
||||
buffer = ""
|
||||
async for chunk in response.content:
|
||||
# Check if we should stop
|
||||
if hasattr(self, '_closing') and self._closing:
|
||||
break
|
||||
|
||||
buffer += chunk.decode('utf-8')
|
||||
|
||||
# Process complete SSE events
|
||||
while "\n\n" in buffer:
|
||||
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
|
||||
|
||||
if data:
|
||||
try:
|
||||
message = json.loads(data)
|
||||
await self._message_queue.put(message)
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"Failed to parse SSE event: {data}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# Proper shutdown
|
||||
break
|
||||
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
|
||||
else:
|
||||
break
|
||||
|
||||
def read_stream(self):
|
||||
"""Create a read stream for the ClientSession."""
|
||||
async def _read():
|
||||
while True:
|
||||
message = await self._message_queue.get()
|
||||
yield message
|
||||
return _read()
|
||||
|
||||
def write_stream(self):
|
||||
"""Create a write stream for the ClientSession."""
|
||||
async def _write(message):
|
||||
if hasattr(message, 'to_dict'):
|
||||
message = message.to_dict()
|
||||
response = await self.send_request(message)
|
||||
return response
|
||||
return _write
|
||||
|
||||
|
||||
class HTTPStreamMCPClient:
|
||||
"""A client for connecting to an MCP server over HTTP Stream transport."""
|
||||
|
||||
def __init__(self, server_url: str, debug: bool = False, timeout: int = 60, options: Optional[Dict[str, Any]] = None):
|
||||
"""
|
||||
Initialize an HTTP Stream MCP client.
|
||||
|
||||
Args:
|
||||
server_url: The URL of the HTTP Stream MCP server
|
||||
debug: Whether to enable debug logging
|
||||
timeout: Timeout in seconds for operations (default: 60)
|
||||
options: Additional configuration options for the transport
|
||||
"""
|
||||
# Parse URL to extract base URL and endpoint
|
||||
parsed = urlparse(server_url)
|
||||
|
||||
# If the URL already has a path, use it; otherwise use default /mcp endpoint
|
||||
if parsed.path and parsed.path != '/':
|
||||
self.base_url = server_url
|
||||
else:
|
||||
# Default endpoint is /mcp
|
||||
self.base_url = urljoin(server_url, '/mcp')
|
||||
|
||||
self.debug = debug
|
||||
self.timeout = timeout
|
||||
self.options = options or {}
|
||||
self.session = None
|
||||
self.tools = []
|
||||
self.transport = None
|
||||
|
||||
# Set up logging
|
||||
if debug:
|
||||
logger.setLevel(logging.DEBUG)
|
||||
else:
|
||||
# Set to WARNING by default to hide INFO messages
|
||||
logger.setLevel(logging.WARNING)
|
||||
|
||||
self._initialize()
|
||||
|
||||
def _initialize(self):
|
||||
"""Initialize the connection and tools."""
|
||||
# Use the global event loop
|
||||
loop = get_event_loop()
|
||||
|
||||
# Start a background thread to run the event loop
|
||||
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 the initialization in the event loop
|
||||
future = asyncio.run_coroutine_threadsafe(self._async_initialize(), loop)
|
||||
self.tools = future.result(timeout=self.timeout)
|
||||
|
||||
async def _async_initialize(self):
|
||||
"""Asynchronously initialize the connection and tools."""
|
||||
logger.debug(f"Connecting to MCP server at {self.base_url}")
|
||||
|
||||
# Create HTTP Stream transport
|
||||
self.transport = HTTPStreamTransport(self.base_url, options=self.options)
|
||||
await self.transport.__aenter__()
|
||||
|
||||
# Create read and write streams
|
||||
read_stream = self.transport.read_stream()
|
||||
write_stream = self.transport.write_stream()
|
||||
|
||||
# Start SSE listener if in stream mode
|
||||
if self.options.get('responseMode', 'batch') == 'stream':
|
||||
await self.transport.start_sse_listener()
|
||||
|
||||
# Create client session
|
||||
self._session_context = ClientSession(read_stream, write_stream)
|
||||
self.session = await self._session_context.__aenter__()
|
||||
|
||||
# Initialize
|
||||
await self.session.initialize()
|
||||
|
||||
# List available tools
|
||||
logger.debug("Listing tools...")
|
||||
response = await self.session.list_tools()
|
||||
tools_data = response.tools
|
||||
logger.debug(f"Found {len(tools_data)} tools: {[tool.name for tool in tools_data]}")
|
||||
|
||||
# Create tool wrappers
|
||||
tools = []
|
||||
for tool in tools_data:
|
||||
input_schema = tool.inputSchema if hasattr(tool, 'inputSchema') else None
|
||||
wrapper = HTTPStreamMCPTool(
|
||||
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 __iter__(self):
|
||||
"""Return an iterator over the tools."""
|
||||
return iter(self.tools)
|
||||
|
||||
def to_openai_tools(self):
|
||||
"""Convert all tools to OpenAI format."""
|
||||
return [tool.to_openai_tool() for tool in self.tools]
|
||||
|
||||
async def __aenter__(self):
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Async context manager exit."""
|
||||
if self.transport:
|
||||
await self.transport.__aexit__(exc_type, exc_val, exc_tb)
|
||||
if hasattr(self, '_session_context') and self._session_context:
|
||||
await self._session_context.__aexit__(exc_type, exc_val, exc_tb)
|
||||
@@ -44,13 +44,30 @@ class SSEMCPTool:
|
||||
# Create a signature based on input schema
|
||||
params = []
|
||||
if input_schema and 'properties' in input_schema:
|
||||
for param_name in input_schema['properties']:
|
||||
for param_name, prop_schema in input_schema['properties'].items():
|
||||
# Determine type annotation based on schema
|
||||
prop_type = prop_schema.get('type', 'string') if isinstance(prop_schema, dict) else 'string'
|
||||
if prop_type == 'string':
|
||||
annotation = str
|
||||
elif prop_type == 'integer':
|
||||
annotation = int
|
||||
elif prop_type == 'number':
|
||||
annotation = float
|
||||
elif prop_type == 'boolean':
|
||||
annotation = bool
|
||||
elif prop_type == 'array':
|
||||
annotation = list
|
||||
elif prop_type == 'object':
|
||||
annotation = dict
|
||||
else:
|
||||
annotation = str # Default to string for SSE
|
||||
|
||||
params.append(
|
||||
inspect.Parameter(
|
||||
name=param_name,
|
||||
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
default=inspect.Parameter.empty if param_name in input_schema.get('required', []) else None,
|
||||
annotation=str # Default to string
|
||||
annotation=annotation
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -116,6 +116,16 @@ class Memory:
|
||||
self.use_mem0 = (self.provider.lower() == "mem0") and MEM0_AVAILABLE
|
||||
self.use_rag = (self.provider.lower() == "rag") and CHROMADB_AVAILABLE and self.cfg.get("use_embedding", False)
|
||||
self.graph_enabled = False # Initialize graph support flag
|
||||
|
||||
# Extract embedding model from config
|
||||
self.embedder_config = self.cfg.get("embedder", {})
|
||||
if isinstance(self.embedder_config, dict):
|
||||
embedder_model_config = self.embedder_config.get("config", {})
|
||||
self.embedding_model = embedder_model_config.get("model", "text-embedding-3-small")
|
||||
else:
|
||||
self.embedding_model = "text-embedding-3-small"
|
||||
|
||||
self._log_verbose(f"Using embedding model: {self.embedding_model}")
|
||||
|
||||
# Create .praison directory if it doesn't exist
|
||||
os.makedirs(".praison", exist_ok=True)
|
||||
@@ -355,7 +365,7 @@ class Memory:
|
||||
import litellm
|
||||
|
||||
response = litellm.embedding(
|
||||
model="text-embedding-3-small",
|
||||
model=self.embedding_model,
|
||||
input=query
|
||||
)
|
||||
query_embedding = response.data[0]["embedding"]
|
||||
@@ -366,7 +376,7 @@ class Memory:
|
||||
|
||||
response = client.embeddings.create(
|
||||
input=query,
|
||||
model="text-embedding-3-small"
|
||||
model=self.embedding_model
|
||||
)
|
||||
query_embedding = response.data[0].embedding
|
||||
else:
|
||||
@@ -496,7 +506,7 @@ class Memory:
|
||||
logger.debug(f"Embedding input text: {text}")
|
||||
|
||||
response = litellm.embedding(
|
||||
model="text-embedding-3-small",
|
||||
model=self.embedding_model,
|
||||
input=text
|
||||
)
|
||||
embedding = response.data[0]["embedding"]
|
||||
@@ -513,7 +523,7 @@ class Memory:
|
||||
|
||||
response = client.embeddings.create(
|
||||
input=text,
|
||||
model="text-embedding-3-small"
|
||||
model=self.embedding_model
|
||||
)
|
||||
embedding = response.data[0].embedding
|
||||
logger.info("Successfully got embeddings from OpenAI")
|
||||
@@ -576,7 +586,7 @@ class Memory:
|
||||
import litellm
|
||||
|
||||
response = litellm.embedding(
|
||||
model="text-embedding-3-small",
|
||||
model=self.embedding_model,
|
||||
input=query
|
||||
)
|
||||
query_embedding = response.data[0]["embedding"]
|
||||
@@ -587,7 +597,7 @@ class Memory:
|
||||
|
||||
response = client.embeddings.create(
|
||||
input=query,
|
||||
model="text-embedding-3-small"
|
||||
model=self.embedding_model
|
||||
)
|
||||
query_embedding = response.data[0].embedding
|
||||
else:
|
||||
|
||||
@@ -6,9 +6,9 @@ from pydantic import BaseModel, ConfigDict
|
||||
from ..agent.agent import Agent
|
||||
from ..task.task import Task
|
||||
from ..main import display_error
|
||||
from ..llm import LLM
|
||||
import csv
|
||||
import os
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
|
||||
class LoopItems(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
@@ -16,6 +16,7 @@ class LoopItems(BaseModel):
|
||||
|
||||
class Process:
|
||||
DEFAULT_RETRY_LIMIT = 3 # Predefined retry limit in a common place
|
||||
VALIDATION_FAILURE_DECISIONS = ["invalid", "retry", "failed", "error", "unsuccessful", "fail", "errors", "reject", "rejected", "incomplete"] # Decision strings that trigger validation feedback
|
||||
|
||||
def __init__(self, tasks: Dict[str, Task], agents: List[Agent], manager_llm: Optional[str] = None, verbose: bool = False, max_iter: int = 10):
|
||||
logging.debug(f"=== Initializing Process ===")
|
||||
@@ -33,12 +34,141 @@ class Process:
|
||||
self.task_retry_counter: Dict[str, int] = {} # Initialize retry counter
|
||||
self.workflow_finished = False # ADDED: Workflow finished flag
|
||||
|
||||
def _create_llm_instance(self):
|
||||
"""Create and return a configured LLM instance for manager tasks."""
|
||||
return LLM(model=self.manager_llm, temperature=0.7)
|
||||
|
||||
def _parse_manager_instructions(self, response, ManagerInstructions):
|
||||
"""Parse LLM response and return ManagerInstructions instance.
|
||||
|
||||
Args:
|
||||
response: String response from LLM
|
||||
ManagerInstructions: Pydantic model class for validation
|
||||
|
||||
Returns:
|
||||
ManagerInstructions instance
|
||||
|
||||
Raises:
|
||||
Exception: If parsing fails
|
||||
"""
|
||||
try:
|
||||
parsed_json = json.loads(response)
|
||||
return ManagerInstructions(**parsed_json)
|
||||
except (json.JSONDecodeError, ValueError, TypeError) as e:
|
||||
raise Exception(f"Failed to parse response: {response}") from e
|
||||
|
||||
def _create_loop_subtasks(self, loop_task: Task):
|
||||
"""Create subtasks for a loop task from input file."""
|
||||
if not loop_task.input_file:
|
||||
logging.warning(f"_create_loop_subtasks called for {loop_task.name} but no input_file specified")
|
||||
return
|
||||
|
||||
try:
|
||||
file_ext = os.path.splitext(loop_task.input_file)[1].lower()
|
||||
new_tasks = []
|
||||
|
||||
if file_ext == ".csv":
|
||||
with open(loop_task.input_file, "r", encoding="utf-8") as f:
|
||||
reader = csv.reader(f, quotechar='"', escapechar='\\')
|
||||
previous_task = None
|
||||
task_count = 0
|
||||
|
||||
for i, row in enumerate(reader):
|
||||
if not row: # Skip empty rows
|
||||
continue
|
||||
|
||||
# Handle Q&A pairs with potential commas
|
||||
task_desc = row[0].strip() if row else ""
|
||||
if len(row) > 1:
|
||||
question = row[0].strip()
|
||||
answer = ",".join(field.strip() for field in row[1:])
|
||||
task_desc = f"Question: {question}\nAnswer: {answer}"
|
||||
|
||||
if not task_desc: # Skip rows with empty content
|
||||
continue
|
||||
|
||||
task_count += 1
|
||||
logging.debug(f"Creating subtask from CSV row {i+1}: {task_desc}")
|
||||
|
||||
row_task = Task(
|
||||
description=f"{loop_task.description}\n{task_desc}" if loop_task.description else task_desc,
|
||||
agent=loop_task.agent,
|
||||
name=f"{loop_task.name}_{task_count}" if loop_task.name else task_desc,
|
||||
expected_output=getattr(loop_task, 'expected_output', None),
|
||||
callback=loop_task.callback, # Inherit callback from parent loop task
|
||||
is_start=(task_count == 1),
|
||||
task_type="task"
|
||||
)
|
||||
self.tasks[row_task.id] = row_task
|
||||
new_tasks.append(row_task)
|
||||
|
||||
if previous_task:
|
||||
previous_task.next_tasks = [row_task.name]
|
||||
previous_task = row_task
|
||||
|
||||
logging.info(f"Created {task_count} subtasks from CSV file for {loop_task.name}")
|
||||
else:
|
||||
# Handle text files
|
||||
with open(loop_task.input_file, "r", encoding="utf-8") as f:
|
||||
lines = f.read().splitlines()
|
||||
previous_task = None
|
||||
for i, line in enumerate(lines):
|
||||
if not line.strip(): # Skip empty lines
|
||||
continue
|
||||
|
||||
row_task = Task(
|
||||
description=f"{loop_task.description}\n{line.strip()}" if loop_task.description else line.strip(),
|
||||
agent=loop_task.agent,
|
||||
name=f"{loop_task.name}_{i+1}" if loop_task.name else line.strip(),
|
||||
expected_output=getattr(loop_task, 'expected_output', None),
|
||||
callback=loop_task.callback, # Inherit callback from parent loop task
|
||||
is_start=(i == 0),
|
||||
task_type="task"
|
||||
)
|
||||
self.tasks[row_task.id] = row_task
|
||||
new_tasks.append(row_task)
|
||||
|
||||
if previous_task:
|
||||
previous_task.next_tasks = [row_task.name]
|
||||
previous_task = row_task
|
||||
|
||||
logging.info(f"Created {len(new_tasks)} subtasks from text file for {loop_task.name}")
|
||||
|
||||
if new_tasks and loop_task.next_tasks:
|
||||
# Connect last subtask to loop task's next tasks
|
||||
last_task = new_tasks[-1]
|
||||
last_task.next_tasks = loop_task.next_tasks
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to create subtasks for loop task {loop_task.name}: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
def _build_task_context(self, current_task: Task) -> str:
|
||||
"""Build context for a task based on its retain_full_context setting"""
|
||||
if not (current_task.previous_tasks or current_task.context):
|
||||
return ""
|
||||
# Check if we have validation feedback to include
|
||||
if current_task.validation_feedback:
|
||||
feedback = current_task.validation_feedback
|
||||
context = f"\nPrevious attempt failed validation with reason: {feedback['validation_response']}"
|
||||
if feedback.get('validated_task'):
|
||||
context += f"\nValidated task: {feedback['validated_task']}"
|
||||
if feedback.get('validation_details'):
|
||||
context += f"\nValidation feedback: {feedback['validation_details']}"
|
||||
if feedback.get('rejected_output'):
|
||||
context += f"\nRejected output: {feedback['rejected_output']}"
|
||||
context += "\nPlease try again with a different approach based on this feedback.\n"
|
||||
# Clear the feedback after including it to prevent it from persisting
|
||||
current_task.validation_feedback = None
|
||||
|
||||
context = "\nInput data from previous tasks:"
|
||||
# If we have validation feedback but no previous tasks context, return just the feedback
|
||||
if not (current_task.previous_tasks or current_task.context):
|
||||
return context
|
||||
# Otherwise, append the regular context
|
||||
context += "\nInput data from previous tasks:"
|
||||
elif not (current_task.previous_tasks or current_task.context):
|
||||
return ""
|
||||
else:
|
||||
context = "\nInput data from previous tasks:"
|
||||
|
||||
if current_task.retain_full_context:
|
||||
# Original behavior: include all previous tasks
|
||||
@@ -140,24 +270,24 @@ class Process:
|
||||
|
||||
def _get_manager_instructions_with_fallback(self, manager_task, manager_prompt, ManagerInstructions):
|
||||
"""Sync version of getting manager instructions with fallback"""
|
||||
# Create OpenAI client
|
||||
client = OpenAI()
|
||||
# Create LLM instance with the manager_llm
|
||||
llm = self._create_llm_instance()
|
||||
|
||||
try:
|
||||
# First try structured output (OpenAI compatible)
|
||||
# Use LLM with output_pydantic for structured output
|
||||
logging.info("Attempting structured output...")
|
||||
manager_response = client.beta.chat.completions.parse(
|
||||
model=self.manager_llm,
|
||||
messages=[
|
||||
{"role": "system", "content": manager_task.description},
|
||||
{"role": "user", "content": manager_prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format=ManagerInstructions
|
||||
response = llm.get_response(
|
||||
prompt=manager_prompt,
|
||||
system_prompt=manager_task.description,
|
||||
output_pydantic=ManagerInstructions
|
||||
)
|
||||
return manager_response.choices[0].message.parsed
|
||||
|
||||
# Parse the response and validate with Pydantic
|
||||
return self._parse_manager_instructions(response, ManagerInstructions)
|
||||
|
||||
except Exception as e:
|
||||
logging.info(f"Structured output failed: {e}, falling back to JSON mode...")
|
||||
# Fallback to regular JSON mode
|
||||
# Fallback to regular JSON mode with explicit JSON instructions
|
||||
try:
|
||||
# Generate JSON structure description from Pydantic model
|
||||
try:
|
||||
@@ -173,23 +303,14 @@ class Process:
|
||||
# Fallback to hardcoded prompt if schema generation fails
|
||||
enhanced_prompt = manager_prompt + "\n\nIMPORTANT: Respond with valid JSON only, using this exact structure: {\"task_id\": <int>, \"agent_name\": \"<string>\", \"action\": \"<execute or stop>\"}"
|
||||
|
||||
manager_response = client.chat.completions.create(
|
||||
model=self.manager_llm,
|
||||
messages=[
|
||||
{"role": "system", "content": manager_task.description},
|
||||
{"role": "user", "content": enhanced_prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format={"type": "json_object"}
|
||||
response = llm.get_response(
|
||||
prompt=enhanced_prompt,
|
||||
system_prompt=manager_task.description,
|
||||
output_json=True
|
||||
)
|
||||
|
||||
# Parse JSON and validate with Pydantic
|
||||
try:
|
||||
json_content = manager_response.choices[0].message.content
|
||||
parsed_json = json.loads(json_content)
|
||||
return ManagerInstructions(**parsed_json)
|
||||
except (json.JSONDecodeError, ValueError) as e:
|
||||
raise Exception(f"Failed to parse JSON response: {json_content}") from e
|
||||
return self._parse_manager_instructions(response, ManagerInstructions)
|
||||
except Exception as fallback_error:
|
||||
error_msg = f"Both structured output and JSON fallback failed: {fallback_error}"
|
||||
logging.error(error_msg, exc_info=True)
|
||||
@@ -197,41 +318,44 @@ class Process:
|
||||
|
||||
async def _get_structured_response_async(self, manager_task, manager_prompt, ManagerInstructions):
|
||||
"""Async version of structured response"""
|
||||
# Create an async client instance for this async method
|
||||
async_client = AsyncOpenAI()
|
||||
manager_response = await async_client.beta.chat.completions.parse(
|
||||
model=self.manager_llm,
|
||||
messages=[
|
||||
{"role": "system", "content": manager_task.description},
|
||||
{"role": "user", "content": manager_prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format=ManagerInstructions
|
||||
# Create LLM instance with the manager_llm
|
||||
llm = self._create_llm_instance()
|
||||
|
||||
# Use async get_response with output_pydantic
|
||||
response = await llm.get_response_async(
|
||||
prompt=manager_prompt,
|
||||
system_prompt=manager_task.description,
|
||||
output_pydantic=ManagerInstructions
|
||||
)
|
||||
return manager_response.choices[0].message.parsed
|
||||
|
||||
# Parse the response and validate with Pydantic
|
||||
return self._parse_manager_instructions(response, ManagerInstructions)
|
||||
|
||||
async def _get_json_response_async(self, manager_task, enhanced_prompt, ManagerInstructions):
|
||||
"""Async version of JSON fallback response"""
|
||||
# Create an async client instance for this async method
|
||||
async_client = AsyncOpenAI()
|
||||
manager_response = await async_client.chat.completions.create(
|
||||
model=self.manager_llm,
|
||||
messages=[
|
||||
{"role": "system", "content": manager_task.description},
|
||||
{"role": "user", "content": enhanced_prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format={"type": "json_object"}
|
||||
# Create LLM instance with the manager_llm
|
||||
llm = self._create_llm_instance()
|
||||
|
||||
response = await llm.get_response_async(
|
||||
prompt=enhanced_prompt,
|
||||
system_prompt=manager_task.description,
|
||||
output_json=True
|
||||
)
|
||||
|
||||
# Parse JSON and validate with Pydantic
|
||||
try:
|
||||
json_content = manager_response.choices[0].message.content
|
||||
parsed_json = json.loads(json_content)
|
||||
return ManagerInstructions(**parsed_json)
|
||||
except (json.JSONDecodeError, ValueError) as e:
|
||||
raise Exception(f"Failed to parse JSON response: {json_content}") from e
|
||||
return self._parse_manager_instructions(response, ManagerInstructions)
|
||||
|
||||
def _check_all_tasks_completed(self) -> bool:
|
||||
"""Check if all tasks are completed and handle workflow completion.
|
||||
|
||||
Returns:
|
||||
bool: True if all tasks are completed and workflow should exit, False otherwise.
|
||||
"""
|
||||
if all(task.status == "completed" for task in self.tasks.values()):
|
||||
logging.info("All tasks are completed.")
|
||||
self.workflow_finished = True
|
||||
return True
|
||||
return False
|
||||
|
||||
async def aworkflow(self) -> AsyncGenerator[str, None]:
|
||||
"""Async version of workflow method"""
|
||||
@@ -294,10 +418,8 @@ Tasks by type:
|
||||
""")
|
||||
|
||||
# ADDED: Check if all tasks are completed and set workflow_finished flag
|
||||
if all(task.status == "completed" for task in self.tasks.values()):
|
||||
logging.info("All tasks are completed.")
|
||||
self.workflow_finished = True
|
||||
# The next iteration loop check will break the workflow
|
||||
if self._check_all_tasks_completed():
|
||||
break # Exit immediately to prevent task reset
|
||||
|
||||
task_id = current_task.id
|
||||
logging.debug(f"""
|
||||
@@ -442,16 +564,18 @@ Subtask: {st.name}
|
||||
logging.debug(f"Task type: {task_to_check.task_type}")
|
||||
logging.debug(f"Task status before reset check: {task_to_check.status}")
|
||||
logging.debug(f"Task rerun: {getattr(task_to_check, 'rerun', True)}") # default to True if not set
|
||||
logging.debug(f"Task async_execution: {task_to_check.async_execution}")
|
||||
|
||||
if (getattr(task_to_check, 'rerun', True) and # Corrected condition - reset only if rerun is True (or default True)
|
||||
task_to_check.task_type != "loop" and # Removed "decision" from exclusion
|
||||
not any(t.task_type == "loop" and subtask_name.startswith(t.name + "_")
|
||||
for t in self.tasks.values())):
|
||||
logging.debug(f"=== Resetting non-loop, non-decision task {subtask_name} to 'not started' ===")
|
||||
for t in self.tasks.values()) and
|
||||
not task_to_check.async_execution): # Don't reset async parallel tasks
|
||||
logging.debug(f"=== Resetting non-loop, non-decision, non-parallel task {subtask_name} to 'not started' ===")
|
||||
self.tasks[task_id].status = "not started"
|
||||
logging.debug(f"Task status after reset: {self.tasks[task_id].status}")
|
||||
else:
|
||||
logging.debug(f"=== Skipping reset for loop/decision/subtask or rerun=False: {subtask_name} ===")
|
||||
logging.debug(f"=== Skipping reset for loop/decision/subtask/parallel or rerun=False: {subtask_name} ===")
|
||||
logging.debug(f"Keeping status as: {self.tasks[task_id].status}")
|
||||
|
||||
# Handle loop progression
|
||||
@@ -496,6 +620,35 @@ Subtask: {st.name}
|
||||
next_task = next((t for t in self.tasks.values() if t.name == task_value), None)
|
||||
if next_task:
|
||||
next_task.status = "not started" # Reset status to allow execution
|
||||
|
||||
# Capture validation feedback for retry scenarios
|
||||
if decision_str in Process.VALIDATION_FAILURE_DECISIONS:
|
||||
if current_task and current_task.result:
|
||||
# Get the rejected output from the task that was validated
|
||||
validated_task = None
|
||||
# Find the task that produced the output being validated
|
||||
if current_task.previous_tasks:
|
||||
# For validation tasks, typically validate the most recent previous task
|
||||
prev_task_name = current_task.previous_tasks[-1]
|
||||
validated_task = next((t for t in self.tasks.values() if t.name == prev_task_name), None)
|
||||
elif current_task.context:
|
||||
# If no previous_tasks, check context for the validated task
|
||||
# Use the most recent task with a result from context
|
||||
for ctx_task in reversed(current_task.context):
|
||||
if ctx_task.result and ctx_task.name != current_task.name:
|
||||
validated_task = ctx_task
|
||||
break
|
||||
|
||||
feedback = {
|
||||
'validation_response': decision_str,
|
||||
'validation_details': current_task.result.raw,
|
||||
'rejected_output': validated_task.result.raw if validated_task and validated_task.result else None,
|
||||
'validator_task': current_task.name,
|
||||
'validated_task': validated_task.name if validated_task else None
|
||||
}
|
||||
next_task.validation_feedback = feedback
|
||||
logging.debug(f"Added validation feedback to {next_task.name}: {feedback['validation_response']} (validated task: {feedback.get('validated_task', 'None')})")
|
||||
|
||||
logging.debug(f"Routing to {next_task.name} based on decision: {decision_str}")
|
||||
# Don't mark workflow as finished when following condition path
|
||||
self.workflow_finished = False
|
||||
@@ -736,6 +889,7 @@ Provide a JSON with the structure:
|
||||
agent=start_task.agent,
|
||||
name=f"{start_task.name}_{task_count}" if start_task.name else task_desc,
|
||||
expected_output=getattr(start_task, 'expected_output', None),
|
||||
callback=start_task.callback, # Inherit callback from parent loop task
|
||||
is_start=(task_count == 1),
|
||||
task_type="decision", # Change to decision type
|
||||
next_tasks=inherited_next_tasks, # Inherit parent's next tasks
|
||||
@@ -769,6 +923,7 @@ Provide a JSON with the structure:
|
||||
agent=start_task.agent,
|
||||
name=f"{start_task.name}_{i+1}" if start_task.name else line.strip(),
|
||||
expected_output=getattr(start_task, 'expected_output', None),
|
||||
callback=start_task.callback, # Inherit callback from parent loop task
|
||||
is_start=(i == 0),
|
||||
task_type="task",
|
||||
condition={
|
||||
@@ -823,10 +978,8 @@ Tasks by type:
|
||||
""")
|
||||
|
||||
# ADDED: Check if all tasks are completed and set workflow_finished flag
|
||||
if all(task.status == "completed" for task in self.tasks.values()):
|
||||
logging.info("All tasks are completed.")
|
||||
self.workflow_finished = True
|
||||
# The next iteration loop check will break the workflow
|
||||
if self._check_all_tasks_completed():
|
||||
break # Exit immediately to prevent task reset
|
||||
|
||||
|
||||
# Handle loop task file reading at runtime
|
||||
@@ -854,6 +1007,7 @@ Tasks by type:
|
||||
agent=current_task.agent,
|
||||
name=f"{current_task.name}_{i+1}" if current_task.name else task_desc,
|
||||
expected_output=getattr(current_task, 'expected_output', None),
|
||||
callback=current_task.callback, # Inherit callback from parent loop task
|
||||
is_start=(i == 0),
|
||||
task_type="task",
|
||||
condition={
|
||||
@@ -878,6 +1032,7 @@ Tasks by type:
|
||||
agent=current_task.agent,
|
||||
name=f"{current_task.name}_{i+1}" if current_task.name else line.strip(),
|
||||
expected_output=getattr(current_task, 'expected_output', None),
|
||||
callback=current_task.callback, # Inherit callback from parent loop task
|
||||
is_start=(i == 0),
|
||||
task_type="task",
|
||||
condition={
|
||||
@@ -1043,16 +1198,18 @@ Subtask: {st.name}
|
||||
logging.debug(f"Task type: {task_to_check.task_type}")
|
||||
logging.debug(f"Task status before reset check: {task_to_check.status}")
|
||||
logging.debug(f"Task rerun: {getattr(task_to_check, 'rerun', True)}") # default to True if not set
|
||||
logging.debug(f"Task async_execution: {task_to_check.async_execution}")
|
||||
|
||||
if (getattr(task_to_check, 'rerun', True) and # Corrected condition - reset only if rerun is True (or default True)
|
||||
task_to_check.task_type != "loop" and # Removed "decision" from exclusion
|
||||
not any(t.task_type == "loop" and subtask_name.startswith(t.name + "_")
|
||||
for t in self.tasks.values())):
|
||||
logging.debug(f"=== Resetting non-loop, non-decision task {subtask_name} to 'not started' ===")
|
||||
for t in self.tasks.values()) and
|
||||
not task_to_check.async_execution): # Don't reset async parallel tasks
|
||||
logging.debug(f"=== Resetting non-loop, non-decision, non-parallel task {subtask_name} to 'not started' ===")
|
||||
self.tasks[task_id].status = "not started"
|
||||
logging.debug(f"Task status after reset: {self.tasks[task_id].status}")
|
||||
else:
|
||||
logging.debug(f"=== Skipping reset for loop/decision/subtask or rerun=False: {subtask_name} ===")
|
||||
logging.debug(f"=== Skipping reset for loop/decision/subtask/parallel or rerun=False: {subtask_name} ===")
|
||||
logging.debug(f"Keeping status as: {self.tasks[task_id].status}")
|
||||
|
||||
|
||||
@@ -1098,6 +1255,35 @@ Subtask: {st.name}
|
||||
next_task = next((t for t in self.tasks.values() if t.name == task_value), None)
|
||||
if next_task:
|
||||
next_task.status = "not started" # Reset status to allow execution
|
||||
|
||||
# Capture validation feedback for retry scenarios
|
||||
if decision_str in Process.VALIDATION_FAILURE_DECISIONS:
|
||||
if current_task and current_task.result:
|
||||
# Get the rejected output from the task that was validated
|
||||
validated_task = None
|
||||
# Find the task that produced the output being validated
|
||||
if current_task.previous_tasks:
|
||||
# For validation tasks, typically validate the most recent previous task
|
||||
prev_task_name = current_task.previous_tasks[-1]
|
||||
validated_task = next((t for t in self.tasks.values() if t.name == prev_task_name), None)
|
||||
elif current_task.context:
|
||||
# If no previous_tasks, check context for the validated task
|
||||
# Use the most recent task with a result from context
|
||||
for ctx_task in reversed(current_task.context):
|
||||
if ctx_task.result and ctx_task.name != current_task.name:
|
||||
validated_task = ctx_task
|
||||
break
|
||||
|
||||
feedback = {
|
||||
'validation_response': decision_str,
|
||||
'validation_details': current_task.result.raw,
|
||||
'rejected_output': validated_task.result.raw if validated_task and validated_task.result else None,
|
||||
'validator_task': current_task.name,
|
||||
'validated_task': validated_task.name if validated_task else None
|
||||
}
|
||||
next_task.validation_feedback = feedback
|
||||
logging.debug(f"Added validation feedback to {next_task.name}: {feedback['validation_response']} (validated task: {feedback.get('validated_task', 'None')})")
|
||||
|
||||
logging.debug(f"Routing to {next_task.name} based on decision: {decision_str}")
|
||||
# Don't mark workflow as finished when following condition path
|
||||
self.workflow_finished = False
|
||||
|
||||
@@ -88,6 +88,7 @@ class Task:
|
||||
self.max_retries = max_retries
|
||||
self.retry_count = retry_count
|
||||
self._guardrail_fn = None
|
||||
self.validation_feedback = None # Store validation failure feedback for retry attempts
|
||||
|
||||
# Set logger level based on config verbose level
|
||||
verbose = self.config.get("verbose", 0)
|
||||
@@ -171,8 +172,15 @@ class Task:
|
||||
# Check return annotation if present
|
||||
return_annotation = sig.return_annotation
|
||||
if return_annotation != inspect.Signature.empty:
|
||||
# Import GuardrailResult for checking
|
||||
from ..guardrails import GuardrailResult
|
||||
|
||||
# Check if it's a GuardrailResult type
|
||||
is_guardrail_result = return_annotation is GuardrailResult
|
||||
|
||||
# Check for tuple return type
|
||||
return_annotation_args = get_args(return_annotation)
|
||||
if not (
|
||||
is_tuple = (
|
||||
get_origin(return_annotation) is tuple
|
||||
and len(return_annotation_args) == 2
|
||||
and return_annotation_args[0] is bool
|
||||
@@ -182,9 +190,11 @@ class Task:
|
||||
or return_annotation_args[1] is TaskOutput
|
||||
or return_annotation_args[1] == Union[str, TaskOutput]
|
||||
)
|
||||
):
|
||||
)
|
||||
|
||||
if not (is_guardrail_result or is_tuple):
|
||||
raise ValueError(
|
||||
"If return type is annotated, it must be Tuple[bool, Any]"
|
||||
"If return type is annotated, it must be GuardrailResult or Tuple[bool, Any]"
|
||||
)
|
||||
|
||||
self._guardrail_fn = self.guardrail
|
||||
@@ -371,13 +381,10 @@ class Task:
|
||||
logger.exception(e) # Print full stack trace
|
||||
# Continue execution even if memory operations fail
|
||||
|
||||
# Execute original callback
|
||||
# Execute original callback with metadata support
|
||||
if self.callback:
|
||||
try:
|
||||
if asyncio.iscoroutinefunction(self.callback):
|
||||
await self.callback(task_output)
|
||||
else:
|
||||
self.callback(task_output)
|
||||
await self._execute_callback_with_metadata(task_output)
|
||||
except Exception as e:
|
||||
logger.error(f"Task {self.id}: Failed to execute callback: {e}")
|
||||
logger.exception(e)
|
||||
@@ -446,7 +453,11 @@ Context:
|
||||
# Call the guardrail function
|
||||
result = self._guardrail_fn(task_output)
|
||||
|
||||
# Convert the result to a GuardrailResult
|
||||
# Check if result is already a GuardrailResult
|
||||
if isinstance(result, GuardrailResult):
|
||||
return result
|
||||
|
||||
# Otherwise, convert the tuple result to a GuardrailResult
|
||||
return GuardrailResult.from_tuple(result)
|
||||
|
||||
except Exception as e:
|
||||
@@ -456,4 +467,63 @@ Context:
|
||||
success=False,
|
||||
result=None,
|
||||
error=f"Guardrail validation error: {str(e)}"
|
||||
)
|
||||
)
|
||||
|
||||
async def _execute_callback_with_metadata(self, task_output):
|
||||
"""Execute callback with metadata support while maintaining backward compatibility.
|
||||
|
||||
This method automatically detects the callback signature:
|
||||
- Single parameter callbacks receive only TaskOutput (backward compatible)
|
||||
- Two parameter callbacks receive TaskOutput and metadata dict (enhanced)
|
||||
|
||||
Args:
|
||||
task_output: The TaskOutput object to pass to the callback
|
||||
"""
|
||||
if not self.callback:
|
||||
return
|
||||
|
||||
try:
|
||||
# Inspect the callback signature to determine parameter count
|
||||
sig = inspect.signature(self.callback)
|
||||
param_count = len(sig.parameters)
|
||||
|
||||
if param_count == 1:
|
||||
# Backward compatible: single parameter callback
|
||||
if asyncio.iscoroutinefunction(self.callback):
|
||||
await self.callback(task_output)
|
||||
else:
|
||||
self.callback(task_output)
|
||||
elif param_count >= 2:
|
||||
# Enhanced: two parameter callback with metadata
|
||||
metadata = {
|
||||
'task_id': self.id,
|
||||
'task_name': self.name,
|
||||
'agent_name': self.agent.name if self.agent else None,
|
||||
'task_type': self.task_type,
|
||||
'task_status': self.status,
|
||||
'task_description': self.description,
|
||||
'expected_output': self.expected_output,
|
||||
'input_file': getattr(self, 'input_file', None),
|
||||
'loop_state': self.loop_state,
|
||||
'retry_count': getattr(self, 'retry_count', 0),
|
||||
'async_execution': self.async_execution
|
||||
}
|
||||
|
||||
if asyncio.iscoroutinefunction(self.callback):
|
||||
await self.callback(task_output, metadata)
|
||||
else:
|
||||
self.callback(task_output, metadata)
|
||||
else:
|
||||
# No parameter callback - unusual but handle gracefully
|
||||
if asyncio.iscoroutinefunction(self.callback):
|
||||
await self.callback()
|
||||
else:
|
||||
self.callback()
|
||||
|
||||
except TypeError as e:
|
||||
# Fallback for signature inspection issues
|
||||
logger.warning(f"Task {self.id}: Callback signature inspection failed, falling back to single parameter: {e}")
|
||||
if asyncio.iscoroutinefunction(self.callback):
|
||||
await self.callback(task_output)
|
||||
else:
|
||||
self.callback(task_output)
|
||||
@@ -99,6 +99,19 @@ Ask yourself these questions:
|
||||
- Each search is independent and secure
|
||||
- Alternative to traditional search engines for privacy
|
||||
|
||||
|
||||
### Spider Tool (Function-Based)
|
||||
- General-purpose web scraping and crawling
|
||||
- Supports CSS selectors for precise content extraction
|
||||
- Can crawl multiple pages and extract links/images
|
||||
- Flexible for various web scraping needs
|
||||
|
||||
### Newspaper Tool (Function-Based)
|
||||
- Specialized for news article extraction
|
||||
- Extracts article title, text, authors, and publish date
|
||||
- Includes NLP processing for keywords and summaries
|
||||
- Categorizes news sources by topic
|
||||
|
||||
### Stock Market Tool (Class-Based)
|
||||
- Does multiple things: check prices, get company info, get historical data
|
||||
- Remembers stock information to avoid repeated downloads
|
||||
|
||||
@@ -39,6 +39,7 @@ TOOL_MAPPINGS = {
|
||||
'extract_text': ('.spider_tools', None),
|
||||
'spider_tools': ('.spider_tools', None),
|
||||
|
||||
|
||||
# DuckDB Tools
|
||||
'query': ('.duckdb_tools', None),
|
||||
'create_table': ('.duckdb_tools', None),
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.0.122"
|
||||
version = "0.0.132"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
requires-python = ">=3.10"
|
||||
authors = [
|
||||
@@ -15,7 +15,8 @@ dependencies = [
|
||||
"rich",
|
||||
"openai",
|
||||
"mcp>=1.6.0",
|
||||
"posthog>=3.0.0"
|
||||
"posthog>=3.0.0",
|
||||
"aiohttp>=3.8.0"
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
@@ -60,6 +61,7 @@ telemetry = [
|
||||
"posthog>=3.0.0"
|
||||
]
|
||||
|
||||
|
||||
# Combined features
|
||||
all = [
|
||||
"praisonaiagents[memory]",
|
||||
@@ -73,4 +75,4 @@ all = [
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["praisonaiagents*"]
|
||||
include = ["praisonaiagents*"]
|
||||
@@ -0,0 +1,341 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Simple test script to isolate the duplicate callback issue #878.
|
||||
This script directly tests the LLM class to determine where the duplicate callbacks originate.
|
||||
|
||||
ROOT CAUSE IDENTIFIED:
|
||||
In llm.py, there are two separate callback execution paths:
|
||||
1. execute_sync_callback() called directly (lines 851-857, 885-891, etc.)
|
||||
2. display_interaction() called when verbose=True (line 896, etc.)
|
||||
|
||||
Since display_interaction() ALSO executes sync callbacks internally (main.py lines 164-190),
|
||||
this results in duplicate callback execution for the same LLM response.
|
||||
|
||||
The issue occurs specifically when verbose=True because both callback paths are active.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add the path to the praisonaiagents module
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src', 'praisonai-agents'))
|
||||
|
||||
from praisonaiagents import register_display_callback
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
# Global counter and log
|
||||
callback_invocations = []
|
||||
|
||||
def track_callback(message=None, response=None, **kwargs):
|
||||
"""
|
||||
Callback function that tracks each invocation with detailed information.
|
||||
"""
|
||||
invocation = {
|
||||
'count': len(callback_invocations) + 1,
|
||||
'message': str(message)[:100] if message else None,
|
||||
'response': str(response)[:100] if response else None,
|
||||
'kwargs_keys': list(kwargs.keys()) if kwargs else [],
|
||||
'stack_trace': get_simplified_stack_trace()
|
||||
}
|
||||
callback_invocations.append(invocation)
|
||||
|
||||
print(f"🔔 CALLBACK #{invocation['count']}")
|
||||
print(f" Message: {invocation['message']}")
|
||||
print(f" Response: {invocation['response']}")
|
||||
print(f" Called from: {invocation['stack_trace']}")
|
||||
print()
|
||||
|
||||
def get_simplified_stack_trace():
|
||||
"""Get a simplified stack trace to see where the callback is called from."""
|
||||
import traceback
|
||||
import inspect
|
||||
|
||||
# Get the current stack
|
||||
stack = inspect.stack()
|
||||
|
||||
# Find relevant frames (skip our own callback function)
|
||||
relevant_frames = []
|
||||
for frame in stack[2:8]: # Skip track_callback and get_simplified_stack_trace
|
||||
filename = os.path.basename(frame.filename)
|
||||
if 'praisonaiagents' in frame.filename:
|
||||
relevant_frames.append(f"{filename}:{frame.lineno} in {frame.function}")
|
||||
|
||||
return " -> ".join(relevant_frames) if relevant_frames else "Unknown"
|
||||
|
||||
def test_llm_streaming():
|
||||
"""Test LLM with streaming enabled."""
|
||||
print("🧪 TEST 1: LLM with streaming")
|
||||
print("-" * 40)
|
||||
|
||||
global callback_invocations
|
||||
callback_invocations = []
|
||||
|
||||
# Register callback
|
||||
register_display_callback('interaction', track_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
|
||||
response = llm.get_response(
|
||||
prompt="Say exactly: Hello World",
|
||||
verbose=True, # This should trigger the callback
|
||||
stream=True # Test with streaming
|
||||
)
|
||||
|
||||
print(f"✅ Response: {response}")
|
||||
print(f"📊 Total callback invocations: {len(callback_invocations)}")
|
||||
|
||||
return len(callback_invocations)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
return -1
|
||||
|
||||
def test_llm_non_streaming():
|
||||
"""Test LLM without streaming."""
|
||||
print("🧪 TEST 2: LLM without streaming")
|
||||
print("-" * 40)
|
||||
|
||||
global callback_invocations
|
||||
callback_invocations = []
|
||||
|
||||
# Register callback
|
||||
register_display_callback('interaction', track_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
|
||||
response = llm.get_response(
|
||||
prompt="Say exactly: Hello World",
|
||||
verbose=True, # This should trigger the callback
|
||||
stream=False # Test without streaming
|
||||
)
|
||||
|
||||
print(f"✅ Response: {response}")
|
||||
print(f"📊 Total callback invocations: {len(callback_invocations)}")
|
||||
|
||||
return len(callback_invocations)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
return -1
|
||||
|
||||
def test_llm_with_self_reflection():
|
||||
"""Test LLM with self-reflection enabled."""
|
||||
print("🧪 TEST 3: LLM with self-reflection")
|
||||
print("-" * 40)
|
||||
|
||||
global callback_invocations
|
||||
callback_invocations = []
|
||||
|
||||
# Register callback
|
||||
register_display_callback('interaction', track_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False, self_reflect=True)
|
||||
|
||||
response = llm.get_response(
|
||||
prompt="Say exactly: Hello World",
|
||||
verbose=True, # This should trigger the callback
|
||||
self_reflect=True,
|
||||
min_reflect=1,
|
||||
max_reflect=1
|
||||
)
|
||||
|
||||
print(f"✅ Response: {response}")
|
||||
print(f"📊 Total callback invocations: {len(callback_invocations)}")
|
||||
|
||||
return len(callback_invocations)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
return -1
|
||||
|
||||
def test_verbose_vs_non_verbose():
|
||||
"""Test the key difference: verbose=True vs verbose=False to isolate the duplicate issue."""
|
||||
print("🧪 TEST 4: Verbose vs Non-Verbose (Key Test for Duplicate Issue)")
|
||||
print("-" * 60)
|
||||
|
||||
global callback_invocations
|
||||
|
||||
# Test with verbose=False (should trigger callback once via execute_sync_callback)
|
||||
print(" 🔹 Testing with verbose=False:")
|
||||
callback_invocations = []
|
||||
register_display_callback('interaction', track_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
response = llm.get_response(
|
||||
prompt="Say exactly: Test Non-Verbose",
|
||||
verbose=False # This should only trigger execute_sync_callback
|
||||
)
|
||||
print(f" Response: {response}")
|
||||
print(f" Callbacks with verbose=False: {len(callback_invocations)}")
|
||||
non_verbose_count = len(callback_invocations)
|
||||
except Exception as e:
|
||||
print(f" ❌ Error: {e}")
|
||||
non_verbose_count = -1
|
||||
|
||||
print()
|
||||
|
||||
# Test with verbose=True (should trigger callback twice: execute_sync_callback + display_interaction)
|
||||
print(" 🔹 Testing with verbose=True:")
|
||||
callback_invocations = []
|
||||
register_display_callback('interaction', track_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
response = llm.get_response(
|
||||
prompt="Say exactly: Test Verbose",
|
||||
verbose=True # This should trigger BOTH execute_sync_callback AND display_interaction
|
||||
)
|
||||
print(f" Response: {response}")
|
||||
print(f" Callbacks with verbose=True: {len(callback_invocations)}")
|
||||
verbose_count = len(callback_invocations)
|
||||
|
||||
# Show the call stack differences
|
||||
if len(callback_invocations) > 1:
|
||||
print(f"\n 📋 Call stack analysis:")
|
||||
for i, inv in enumerate(callback_invocations):
|
||||
print(f" Callback #{i+1}: {inv['stack_trace']}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" ❌ Error: {e}")
|
||||
verbose_count = -1
|
||||
|
||||
print(f"\n 📊 RESULT:")
|
||||
print(f" verbose=False: {non_verbose_count} callback(s)")
|
||||
print(f" verbose=True: {verbose_count} callback(s)")
|
||||
|
||||
if verbose_count > non_verbose_count and verbose_count > 1:
|
||||
print(f" 🚨 DUPLICATE ISSUE CONFIRMED!")
|
||||
print(f" verbose=True triggers {verbose_count - non_verbose_count} extra callback(s)")
|
||||
|
||||
return verbose_count
|
||||
|
||||
def test_multiple_calls():
|
||||
"""Test multiple consecutive calls to see pattern."""
|
||||
print("🧪 TEST 5: Multiple consecutive calls")
|
||||
print("-" * 40)
|
||||
|
||||
global callback_invocations
|
||||
callback_invocations = []
|
||||
|
||||
# Register callback
|
||||
register_display_callback('interaction', track_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
|
||||
for i in range(3):
|
||||
print(f" Call {i+1}:")
|
||||
response = llm.get_response(
|
||||
prompt=f"Say exactly: Call {i+1}",
|
||||
verbose=True
|
||||
)
|
||||
print(f" Response: {response}")
|
||||
|
||||
print(f"📊 Total callback invocations across 3 calls: {len(callback_invocations)}")
|
||||
|
||||
return len(callback_invocations)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
return -1
|
||||
|
||||
def analyze_callback_patterns():
|
||||
"""Analyze the patterns in callback invocations."""
|
||||
print("\n📊 CALLBACK PATTERN ANALYSIS")
|
||||
print("=" * 50)
|
||||
|
||||
if not callback_invocations:
|
||||
print("❌ No callback invocations recorded")
|
||||
return
|
||||
|
||||
print(f"Total invocations: {len(callback_invocations)}")
|
||||
print("\nDetailed breakdown:")
|
||||
|
||||
for inv in callback_invocations:
|
||||
print(f" #{inv['count']}: {inv['stack_trace']}")
|
||||
if inv['message']:
|
||||
print(f" Message: {inv['message']}")
|
||||
if inv['response']:
|
||||
print(f" Response: {inv['response']}")
|
||||
print()
|
||||
|
||||
# Look for duplicate patterns
|
||||
unique_traces = set(inv['stack_trace'] for inv in callback_invocations)
|
||||
if len(unique_traces) < len(callback_invocations):
|
||||
print("⚠️ POTENTIAL DUPLICATE PATTERN DETECTED!")
|
||||
print("Some callbacks are coming from the same call stack.")
|
||||
else:
|
||||
print("✅ All callbacks appear to come from different call paths.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("🔬 ISOLATING DUPLICATE CALLBACK ISSUE #878")
|
||||
print("=" * 60)
|
||||
print("This script tests the LLM class directly to isolate where duplicate callbacks occur.")
|
||||
print()
|
||||
|
||||
# Run different test scenarios
|
||||
results = {}
|
||||
|
||||
# Test 1: Streaming
|
||||
results['streaming'] = test_llm_streaming()
|
||||
print()
|
||||
|
||||
# Test 2: Non-streaming
|
||||
results['non_streaming'] = test_llm_non_streaming()
|
||||
print()
|
||||
|
||||
# Test 3: Self-reflection
|
||||
results['self_reflection'] = test_llm_with_self_reflection()
|
||||
print()
|
||||
|
||||
# Test 4: Verbose vs Non-Verbose (KEY TEST)
|
||||
results['verbose_vs_nonverbose'] = test_verbose_vs_non_verbose()
|
||||
print()
|
||||
|
||||
# Test 5: Multiple calls
|
||||
results['multiple_calls'] = test_multiple_calls()
|
||||
print()
|
||||
|
||||
# Analyze patterns
|
||||
analyze_callback_patterns()
|
||||
|
||||
# Summary
|
||||
print("\n🏁 SUMMARY OF RESULTS")
|
||||
print("=" * 40)
|
||||
|
||||
expected_single_call_tests = ['streaming', 'non_streaming', 'self_reflection']
|
||||
issues_detected = []
|
||||
|
||||
for test_name, count in results.items():
|
||||
if count == -1:
|
||||
print(f"{test_name}: ERROR")
|
||||
elif test_name == 'multiple_calls':
|
||||
expected = 3
|
||||
print(f"{test_name}: {count} invocations (expected ~{expected})")
|
||||
if count > expected * 1.5: # Allow some tolerance
|
||||
issues_detected.append(f"{test_name}: {count} > expected {expected}")
|
||||
elif test_name == 'verbose_vs_nonverbose':
|
||||
# This is the key test - verbose=True should show the duplicate issue
|
||||
print(f"{test_name}: {count} invocations (this test shows verbose=True duplicates)")
|
||||
if count > 1:
|
||||
issues_detected.append(f"{test_name}: verbose=True triggered {count} callbacks (duplicate issue)")
|
||||
else:
|
||||
expected = 1
|
||||
print(f"{test_name}: {count} invocations (expected {expected})")
|
||||
if count > expected:
|
||||
issues_detected.append(f"{test_name}: {count} > expected {expected}")
|
||||
|
||||
if issues_detected:
|
||||
print(f"\n❌ DUPLICATE CALLBACK ISSUES DETECTED:")
|
||||
for issue in issues_detected:
|
||||
print(f" - {issue}")
|
||||
print("\nThe callback system is triggering multiple times when it should trigger once.")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print(f"\n✅ NO DUPLICATE CALLBACK ISSUES DETECTED")
|
||||
print("All callback invocations appear to be working as expected.")
|
||||
sys.exit(0)
|
||||
@@ -0,0 +1,222 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test script to verify async task execution in sequential mode"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import os
|
||||
from datetime import datetime
|
||||
|
||||
# Set up the environment
|
||||
import sys
|
||||
sys.path.insert(0, '/home/runner/work/PraisonAI/PraisonAI/src/praisonai-agents')
|
||||
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
# Create test agents
|
||||
async_agent = Agent(
|
||||
name="AsyncAgent",
|
||||
role="Async task executor",
|
||||
goal="Execute async tasks",
|
||||
backstory="An agent that executes async tasks",
|
||||
llm="openai/gpt-3.5-turbo" # Using a mock/test LLM
|
||||
)
|
||||
|
||||
sync_agent = Agent(
|
||||
name="SyncAgent",
|
||||
role="Sync task executor",
|
||||
goal="Execute sync tasks",
|
||||
backstory="An agent that executes sync tasks",
|
||||
llm="openai/gpt-3.5-turbo"
|
||||
)
|
||||
|
||||
# Track execution times
|
||||
execution_log = []
|
||||
|
||||
def log_execution(task_name, start_time, end_time, is_async):
|
||||
execution_log.append({
|
||||
"task": task_name,
|
||||
"start": start_time,
|
||||
"end": end_time,
|
||||
"duration": end_time - start_time,
|
||||
"type": "async" if is_async else "sync"
|
||||
})
|
||||
|
||||
# Create test tasks
|
||||
async def async_task_function(task_name, delay):
|
||||
start_time = time.time()
|
||||
print(f"{datetime.now().strftime('%H:%M:%S.%f')[:-3]} - {task_name} started (async)")
|
||||
await asyncio.sleep(delay)
|
||||
end_time = time.time()
|
||||
print(f"{datetime.now().strftime('%H:%M:%S.%f')[:-3]} - {task_name} completed (async)")
|
||||
log_execution(task_name, start_time, end_time, True)
|
||||
return f"{task_name} result after {delay}s"
|
||||
|
||||
def sync_task_function(task_name, delay):
|
||||
start_time = time.time()
|
||||
print(f"{datetime.now().strftime('%H:%M:%S.%f')[:-3]} - {task_name} started (sync)")
|
||||
time.sleep(delay)
|
||||
end_time = time.time()
|
||||
print(f"{datetime.now().strftime('%H:%M:%S.%f')[:-3]} - {task_name} completed (sync)")
|
||||
log_execution(task_name, start_time, end_time, False)
|
||||
return f"{task_name} result after {delay}s"
|
||||
|
||||
# Test scenario: Multiple async tasks followed by sync task followed by more async tasks
|
||||
tasks = []
|
||||
|
||||
# First batch of async tasks (should run in parallel)
|
||||
task1 = Task(
|
||||
name="async_task_1",
|
||||
description="First async task - sleep 2s",
|
||||
expected_output="Task result",
|
||||
agent=async_agent,
|
||||
async_execution=True
|
||||
)
|
||||
|
||||
task2 = Task(
|
||||
name="async_task_2",
|
||||
description="Second async task - sleep 2s",
|
||||
expected_output="Task result",
|
||||
agent=async_agent,
|
||||
async_execution=True
|
||||
)
|
||||
|
||||
task3 = Task(
|
||||
name="async_task_3",
|
||||
description="Third async task - sleep 2s",
|
||||
expected_output="Task result",
|
||||
agent=async_agent,
|
||||
async_execution=True
|
||||
)
|
||||
|
||||
# Sync task (should wait for async tasks to complete)
|
||||
task4 = Task(
|
||||
name="sync_task_1",
|
||||
description="First sync task - sleep 1s",
|
||||
expected_output="Task result",
|
||||
agent=sync_agent,
|
||||
async_execution=False
|
||||
)
|
||||
|
||||
# Second batch of async tasks (should run in parallel after sync task)
|
||||
task5 = Task(
|
||||
name="async_task_4",
|
||||
description="Fourth async task - sleep 1s",
|
||||
expected_output="Task result",
|
||||
agent=async_agent,
|
||||
async_execution=True
|
||||
)
|
||||
|
||||
task6 = Task(
|
||||
name="async_task_5",
|
||||
description="Fifth async task - sleep 1s",
|
||||
expected_output="Task result",
|
||||
agent=async_agent,
|
||||
async_execution=True
|
||||
)
|
||||
|
||||
tasks = [task1, task2, task3, task4, task5, task6]
|
||||
|
||||
async def test_sequential_async_execution():
|
||||
"""Test that async tasks run in parallel in sequential mode"""
|
||||
|
||||
print("\n=== Testing Sequential Mode with Async Tasks ===\n")
|
||||
|
||||
# Mock the agent execute methods to use our test functions
|
||||
original_execute = Agent.execute
|
||||
original_aexecute = Agent.aexecute
|
||||
|
||||
def mock_execute(self, task, context=None):
|
||||
if "async" in task.name:
|
||||
# This shouldn't be called for async tasks
|
||||
print(f"WARNING: Sync execute called for async task {task.name}")
|
||||
else:
|
||||
delay = int(task.description.split("sleep ")[1].split("s")[0])
|
||||
return sync_task_function(task.name, delay)
|
||||
|
||||
async def mock_aexecute(self, task, context=None):
|
||||
if "async" in task.name:
|
||||
delay = int(task.description.split("sleep ")[1].split("s")[0])
|
||||
return await async_task_function(task.name, delay)
|
||||
else:
|
||||
# This shouldn't be called for sync tasks
|
||||
print(f"WARNING: Async execute called for sync task {task.name}")
|
||||
|
||||
# Patch the methods
|
||||
Agent.execute = mock_execute
|
||||
Agent.aexecute = mock_aexecute
|
||||
|
||||
try:
|
||||
# Create workflow
|
||||
workflow = PraisonAIAgents(
|
||||
agents=[async_agent, sync_agent],
|
||||
tasks=tasks,
|
||||
process="sequential",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Run the workflow
|
||||
start_time = time.time()
|
||||
await workflow.astart()
|
||||
end_time = time.time()
|
||||
|
||||
print(f"\nTotal execution time: {end_time - start_time:.2f}s")
|
||||
|
||||
# Analyze results
|
||||
print("\n=== Execution Analysis ===")
|
||||
|
||||
# Check if first batch of async tasks ran in parallel
|
||||
async_batch1 = [log for log in execution_log if log["task"] in ["async_task_1", "async_task_2", "async_task_3"]]
|
||||
if async_batch1:
|
||||
earliest_start = min(log["start"] for log in async_batch1)
|
||||
latest_end = max(log["end"] for log in async_batch1)
|
||||
batch1_duration = latest_end - earliest_start
|
||||
print(f"\nFirst async batch (3 tasks, 2s each):")
|
||||
print(f" - Started at: {earliest_start:.2f}")
|
||||
print(f" - Ended at: {latest_end:.2f}")
|
||||
print(f" - Total duration: {batch1_duration:.2f}s")
|
||||
print(f" - Parallel execution: {'YES' if batch1_duration < 4 else 'NO'} (expected ~2s for parallel, 6s for sequential)")
|
||||
|
||||
# Check sync task timing
|
||||
sync_logs = [log for log in execution_log if log["type"] == "sync"]
|
||||
if sync_logs:
|
||||
sync_log = sync_logs[0]
|
||||
print(f"\nSync task:")
|
||||
print(f" - Started at: {sync_log['start']:.2f}")
|
||||
print(f" - Ended at: {sync_log['end']:.2f}")
|
||||
print(f" - Started after async batch: {'YES' if sync_log['start'] >= latest_end - 0.1 else 'NO'}")
|
||||
|
||||
# Check if second batch of async tasks ran in parallel
|
||||
async_batch2 = [log for log in execution_log if log["task"] in ["async_task_4", "async_task_5"]]
|
||||
if async_batch2:
|
||||
earliest_start = min(log["start"] for log in async_batch2)
|
||||
latest_end = max(log["end"] for log in async_batch2)
|
||||
batch2_duration = latest_end - earliest_start
|
||||
print(f"\nSecond async batch (2 tasks, 1s each):")
|
||||
print(f" - Started at: {earliest_start:.2f}")
|
||||
print(f" - Ended at: {latest_end:.2f}")
|
||||
print(f" - Total duration: {batch2_duration:.2f}s")
|
||||
print(f" - Parallel execution: {'YES' if batch2_duration < 1.5 else 'NO'} (expected ~1s for parallel, 2s for sequential)")
|
||||
|
||||
# Overall assessment
|
||||
print(f"\n=== Test Result ===")
|
||||
expected_time = 2 + 1 + 1 # First batch (2s parallel) + sync (1s) + second batch (1s parallel)
|
||||
actual_time = end_time - start_time
|
||||
print(f"Expected total time (with parallel async): ~{expected_time}s")
|
||||
print(f"Actual total time: {actual_time:.2f}s")
|
||||
|
||||
if actual_time < expected_time + 1: # Allow 1s margin for overhead
|
||||
print("✅ PASS: Async tasks executed in parallel!")
|
||||
else:
|
||||
print("❌ FAIL: Async tasks did not execute in parallel as expected")
|
||||
|
||||
finally:
|
||||
# Restore original methods
|
||||
Agent.execute = original_execute
|
||||
Agent.aexecute = original_aexecute
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Set a dummy API key to avoid errors
|
||||
os.environ["OPENAI_API_KEY"] = "test-key"
|
||||
|
||||
# Run the test
|
||||
asyncio.run(test_sequential_async_execution())
|
||||
@@ -0,0 +1,143 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify that AutoAgents hierarchical workflow works with non-OpenAI LLMs
|
||||
without requiring OpenAI API key.
|
||||
|
||||
This tests the fix for issue #873.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Ensure OpenAI API key is not set to verify the fix works
|
||||
if 'OPENAI_API_KEY' in os.environ:
|
||||
del os.environ['OPENAI_API_KEY']
|
||||
|
||||
# Add the parent directory to the Python path
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from praisonaiagents import AutoAgents
|
||||
|
||||
def get_stock_price(company_name: str) -> str:
|
||||
"""
|
||||
Get the stock price of a company
|
||||
|
||||
Args:
|
||||
company_name (str): The name of the company
|
||||
|
||||
Returns:
|
||||
str: The stock price of the company
|
||||
"""
|
||||
if company_name.lower() == "apple" or company_name.lower() == "aapl":
|
||||
return f"The stock price of {company_name} is 100"
|
||||
elif company_name.lower() == "google" or company_name.lower() == "googl":
|
||||
return f"The stock price of {company_name} is 200"
|
||||
else:
|
||||
return f"The stock price of {company_name} is 50"
|
||||
|
||||
def test_hierarchical_without_openai():
|
||||
"""Test hierarchical process without OpenAI API key"""
|
||||
print("Testing AutoAgents with hierarchical process and Gemini LLM...")
|
||||
print("OPENAI_API_KEY is set:", 'OPENAI_API_KEY' in os.environ)
|
||||
|
||||
try:
|
||||
# Create AutoAgents instance with Gemini
|
||||
agents = AutoAgents(
|
||||
instructions="Write a poem on the stock price of apple",
|
||||
tools=[get_stock_price],
|
||||
process="hierarchical",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
self_reflect=True,
|
||||
verbose=True,
|
||||
max_agents=3 # Maximum number of agents to create
|
||||
)
|
||||
|
||||
# Start the agents
|
||||
print("\nStarting agents...")
|
||||
result = agents.start()
|
||||
|
||||
print("\n=== RESULT ===")
|
||||
print(result)
|
||||
print("\n=== TEST PASSED ===")
|
||||
print("Hierarchical workflow completed successfully without OpenAI API key!")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
if "api_key" in str(e) and "OPENAI_API_KEY" in str(e):
|
||||
print("\n=== TEST FAILED ===")
|
||||
print(f"Error: Still requires OpenAI API key: {e}")
|
||||
return False
|
||||
else:
|
||||
print(f"\nUnexpected error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
def test_hierarchical_with_mock_llm():
|
||||
"""Test with a mock LLM to ensure no API calls are made to OpenAI"""
|
||||
print("\n\nTesting with mock LLM configuration...")
|
||||
|
||||
try:
|
||||
# Create AutoAgents instance with a mock configuration
|
||||
agents = AutoAgents(
|
||||
instructions="Simple test task",
|
||||
tools=[get_stock_price],
|
||||
process="hierarchical",
|
||||
llm={
|
||||
"model": "test/mock-model",
|
||||
"api_key": "mock-key",
|
||||
"base_url": "http://localhost:9999" # Non-existent URL
|
||||
},
|
||||
verbose=False,
|
||||
max_agents=2
|
||||
)
|
||||
|
||||
# This should fail with connection error, not OpenAI API key error
|
||||
try:
|
||||
result = agents.start()
|
||||
# If we get here, the test didn't behave as expected
|
||||
print("\n=== TEST WARNING ===")
|
||||
print("Expected an error but execution succeeded")
|
||||
return False
|
||||
except Exception as e:
|
||||
if "api_key" in str(e) and "OPENAI_API_KEY" in str(e):
|
||||
print("\n=== TEST FAILED ===")
|
||||
print(f"Error: Still requires OpenAI API key: {e}")
|
||||
return False
|
||||
else:
|
||||
print("\n=== TEST PASSED ===")
|
||||
print("Failed with expected error (not OpenAI API key error)")
|
||||
print(f"Error type: {type(e).__name__}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"\nSetup error: {e}")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Track test results
|
||||
test_results = []
|
||||
|
||||
# Run both tests
|
||||
mock_test_passed = test_hierarchical_with_mock_llm()
|
||||
test_results.append(("Mock LLM Test", mock_test_passed))
|
||||
|
||||
# Only run the actual Gemini test if we have the API key
|
||||
if 'GOOGLE_API_KEY' in os.environ or 'GEMINI_API_KEY' in os.environ:
|
||||
gemini_test_passed = test_hierarchical_without_openai()
|
||||
test_results.append(("Gemini Test", gemini_test_passed))
|
||||
else:
|
||||
print("\n\nSkipping Gemini test - no GOOGLE_API_KEY or GEMINI_API_KEY found")
|
||||
print("The mock test passed, confirming the fix works!")
|
||||
|
||||
# Summary
|
||||
print("\n=== TEST SUMMARY ===")
|
||||
all_passed = True
|
||||
for test_name, passed in test_results:
|
||||
status = "PASSED" if passed else "FAILED"
|
||||
print(f"{test_name}: {status}")
|
||||
if not passed:
|
||||
all_passed = False
|
||||
|
||||
# Exit with appropriate code
|
||||
sys.exit(0 if all_passed else 1)
|
||||
@@ -0,0 +1,82 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test script for callback enhancement - Issue #896"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src', 'praisonai-agents'))
|
||||
|
||||
import asyncio
|
||||
from praisonaiagents import (
|
||||
register_display_callback,
|
||||
Agent,
|
||||
Task,
|
||||
PraisonAIAgents
|
||||
)
|
||||
|
||||
def enhanced_callback(message=None, response=None, agent_name=None, agent_role=None, agent_tools=None, task_name=None, task_description=None, task_id=None, **kwargs):
|
||||
"""Enhanced callback that demonstrates the new context information"""
|
||||
print("="*80)
|
||||
print("ENHANCED CALLBACK TRIGGERED")
|
||||
print("="*80)
|
||||
print(f"Message: {message}")
|
||||
print(f"Response: {response}")
|
||||
print(f"Agent Name: {agent_name}")
|
||||
print(f"Agent Role: {agent_role}")
|
||||
print(f"Agent Tools: {agent_tools}")
|
||||
print(f"Task Name: {task_name}")
|
||||
print(f"Task Description: {task_description}")
|
||||
print(f"Task ID: {task_id}")
|
||||
print(f"Other kwargs: {kwargs}")
|
||||
print("="*80)
|
||||
print()
|
||||
|
||||
# Register the enhanced callback
|
||||
register_display_callback('interaction', enhanced_callback, is_async=False)
|
||||
|
||||
def test_callback_enhancement():
|
||||
"""Test function to verify the callback enhancement works"""
|
||||
|
||||
# Create an agent with detailed information
|
||||
agent = Agent(
|
||||
name="TestAgent",
|
||||
role="Assistant",
|
||||
goal="Help with testing callbacks",
|
||||
backstory="I am a helpful assistant for testing purposes",
|
||||
llm="openai/gpt-4o-mini",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Create a task with detailed information
|
||||
task = Task(
|
||||
name="callback_test_task",
|
||||
description="Test the enhanced callback functionality",
|
||||
agent=agent,
|
||||
expected_output="A simple response to test callbacks"
|
||||
)
|
||||
|
||||
# Run the agents
|
||||
try:
|
||||
agents = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
verbose=True
|
||||
)
|
||||
result = agents.start()
|
||||
print(f"Test completed successfully! Result: {result}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Test failed with error: {e}")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing callback enhancement (Issue #896)...")
|
||||
print("This test will demonstrate the new agent_name and task context in callbacks")
|
||||
print()
|
||||
|
||||
# Run the test
|
||||
success = test_callback_enhancement()
|
||||
|
||||
if success:
|
||||
print("✅ Test passed! Callback enhancement is working correctly.")
|
||||
else:
|
||||
print("❌ Test failed! Check the error output above.")
|
||||
@@ -0,0 +1,71 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test script to verify callbacks work with verbose=False"""
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, 'src/praisonai-agents')
|
||||
|
||||
from praisonaiagents import (
|
||||
register_display_callback,
|
||||
Agent,
|
||||
Task,
|
||||
PraisonAIAgents
|
||||
)
|
||||
|
||||
# Track if callback was executed
|
||||
callback_executed = False
|
||||
|
||||
def simple_callback(message=None, response=None, **kwargs):
|
||||
global callback_executed
|
||||
callback_executed = True
|
||||
with open('callback_test_log.txt', 'a') as f:
|
||||
f.write(f"CALLBACK EXECUTED!\n")
|
||||
f.write(f"Message: {message}\n")
|
||||
f.write(f"Response: {response}\n")
|
||||
f.write(f"Other kwargs: {kwargs}\n")
|
||||
f.write("-" * 50 + "\n")
|
||||
|
||||
# Register callback
|
||||
register_display_callback('interaction', simple_callback, is_async=False)
|
||||
|
||||
# Test with verbose=False
|
||||
print("Testing with verbose=False...")
|
||||
agent = Agent(
|
||||
name="TestAgent",
|
||||
role="Assistant",
|
||||
goal="Help with tasks",
|
||||
backstory="I am a helpful assistant",
|
||||
llm="gpt-4o-mini", # Using a simple model for testing
|
||||
verbose=False
|
||||
)
|
||||
|
||||
task = Task(
|
||||
name="simple_task",
|
||||
description="Say hello world",
|
||||
agent=agent,
|
||||
expected_output="A greeting"
|
||||
)
|
||||
|
||||
# Clear the log file
|
||||
with open('callback_test_log.txt', 'w') as f:
|
||||
f.write("Starting test with verbose=False\n")
|
||||
|
||||
# Run the agent
|
||||
try:
|
||||
agents = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task]
|
||||
)
|
||||
result = agents.start()
|
||||
print(f"Task completed. Result: {result}")
|
||||
print(f"Callback executed: {callback_executed}")
|
||||
|
||||
# Check if callback was executed
|
||||
if callback_executed:
|
||||
print("✅ SUCCESS: Callback was executed with verbose=False!")
|
||||
else:
|
||||
print("❌ FAILED: Callback was NOT executed with verbose=False")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error during test: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
@@ -0,0 +1,112 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Simple test to verify that the callback fix is working
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__)))
|
||||
|
||||
from praisonaiagents import register_display_callback
|
||||
from praisonaiagents.main import display_interaction, display_self_reflection
|
||||
|
||||
# Track callback calls
|
||||
callback_calls = []
|
||||
|
||||
def test_callback(message, response, **kwargs):
|
||||
"""Test callback function"""
|
||||
callback_calls.append({
|
||||
'type': 'interaction',
|
||||
'message': message,
|
||||
'response': response
|
||||
})
|
||||
print(f"[CALLBACK] Interaction: {message[:30]}... -> {response[:30]}...")
|
||||
|
||||
def test_self_reflection_callback(message, **kwargs):
|
||||
"""Test self-reflection callback function"""
|
||||
callback_calls.append({
|
||||
'type': 'self_reflection',
|
||||
'message': message
|
||||
})
|
||||
print(f"[CALLBACK] Self-reflection: {message[:50]}...")
|
||||
|
||||
def test_display_interaction_callback():
|
||||
"""Test that display_interaction calls the callback"""
|
||||
global callback_calls
|
||||
callback_calls = []
|
||||
|
||||
# Register callback
|
||||
register_display_callback('interaction', test_callback)
|
||||
|
||||
try:
|
||||
# Call display_interaction
|
||||
display_interaction(
|
||||
message="Test message",
|
||||
response="Test response",
|
||||
markdown=True,
|
||||
generation_time=1.5
|
||||
)
|
||||
|
||||
print(f"Callback calls: {len(callback_calls)}")
|
||||
|
||||
# Verify callback was called
|
||||
assert len(callback_calls) == 1, f"Expected 1 callback call, got {len(callback_calls)}"
|
||||
assert callback_calls[0]['type'] == 'interaction'
|
||||
assert callback_calls[0]['message'] == "Test message"
|
||||
assert callback_calls[0]['response'] == "Test response"
|
||||
|
||||
print("✓ display_interaction callback test PASSED")
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
from praisonaiagents.main import sync_display_callbacks
|
||||
sync_display_callbacks.pop('interaction', None)
|
||||
|
||||
def test_display_self_reflection_callback():
|
||||
"""Test that display_self_reflection calls the callback"""
|
||||
global callback_calls
|
||||
callback_calls = []
|
||||
|
||||
# Register callback
|
||||
register_display_callback('self_reflection', test_self_reflection_callback)
|
||||
|
||||
try:
|
||||
# Call display_self_reflection
|
||||
display_self_reflection("This is a self-reflection message")
|
||||
|
||||
print(f"Self-reflection callback calls: {len(callback_calls)}")
|
||||
|
||||
# Verify callback was called
|
||||
assert len(callback_calls) == 1, f"Expected 1 callback call, got {len(callback_calls)}"
|
||||
assert callback_calls[0]['type'] == 'self_reflection'
|
||||
assert callback_calls[0]['message'] == "This is a self-reflection message"
|
||||
|
||||
print("✓ display_self_reflection callback test PASSED")
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
from praisonaiagents.main import sync_display_callbacks
|
||||
sync_display_callbacks.pop('self_reflection', None)
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing callback fix for display functions...\n")
|
||||
|
||||
try:
|
||||
print("1. Testing display_interaction callback...")
|
||||
test_display_interaction_callback()
|
||||
print()
|
||||
|
||||
print("2. Testing display_self_reflection callback...")
|
||||
test_display_self_reflection_callback()
|
||||
print()
|
||||
|
||||
print("All callback tests passed! The fix is working correctly.")
|
||||
|
||||
except AssertionError as e:
|
||||
print(f"✗ FAILED: {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"✗ ERROR: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,44 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Simple test for callback enhancement - Issue #896"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src', 'praisonai-agents'))
|
||||
|
||||
# Test just the callback registration and display functions
|
||||
from praisonaiagents import register_display_callback, display_interaction
|
||||
|
||||
def test_callback(message=None, response=None, agent_name=None, agent_role=None, agent_tools=None, task_name=None, task_description=None, task_id=None, **kwargs):
|
||||
"""Test callback function"""
|
||||
print("Callback triggered with:")
|
||||
print(f" Message: {message}")
|
||||
print(f" Response: {response}")
|
||||
print(f" Agent Name: {agent_name}")
|
||||
print(f" Agent Role: {agent_role}")
|
||||
print(f" Agent Tools: {agent_tools}")
|
||||
print(f" Task Name: {task_name}")
|
||||
print(f" Task Description: {task_description}")
|
||||
print(f" Task ID: {task_id}")
|
||||
print(f" Other kwargs: {kwargs}")
|
||||
print()
|
||||
|
||||
# Register the callback
|
||||
register_display_callback('interaction', test_callback, is_async=False)
|
||||
|
||||
# Test the display_interaction function with new parameters
|
||||
print("Testing display_interaction with enhanced parameters...")
|
||||
|
||||
display_interaction(
|
||||
message="Test message",
|
||||
response="Test response",
|
||||
agent_name="TestAgent",
|
||||
agent_role="Assistant",
|
||||
agent_tools=["tool1", "tool2"],
|
||||
task_name="test_task",
|
||||
task_description="Test task description",
|
||||
task_id="task_123",
|
||||
markdown=True,
|
||||
generation_time=1.5
|
||||
)
|
||||
|
||||
print("✅ Simple callback test completed!")
|
||||
@@ -0,0 +1,230 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify the duplicate callback fix for issue #878
|
||||
This tests that display_interaction is only called once per LLM response
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src/praisonai-agents'))
|
||||
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
from praisonaiagents import register_display_callback
|
||||
from unittest.mock import patch, MagicMock
|
||||
import json
|
||||
|
||||
# Track display_interaction calls
|
||||
display_calls = []
|
||||
|
||||
def mock_display_interaction(message, response, markdown=True, generation_time=0, **kwargs):
|
||||
"""Mock display_interaction to track calls"""
|
||||
display_calls.append({
|
||||
'message': message,
|
||||
'response': response,
|
||||
'markdown': markdown,
|
||||
'generation_time': generation_time
|
||||
})
|
||||
print(f"[DISPLAY] {message[:50]}... -> {response[:50]}...")
|
||||
|
||||
def test_single_display_no_tools():
|
||||
"""Test that display_interaction is called only once without tools"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
# Register callback instead of patching
|
||||
register_display_callback('interaction', mock_display_interaction)
|
||||
|
||||
try:
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
# Mock streaming response
|
||||
mock_completion.return_value = [
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]),
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content=" world"))]),
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))])
|
||||
]
|
||||
|
||||
llm = LLM(model="gpt-4o-mini", verbose=False)
|
||||
response = llm.get_response(
|
||||
prompt="Test prompt",
|
||||
verbose=True,
|
||||
stream=True
|
||||
)
|
||||
|
||||
print(f"\nResponse: {response}")
|
||||
print(f"Display calls: {len(display_calls)}")
|
||||
|
||||
assert len(display_calls) == 1, f"Expected 1 display call, got {len(display_calls)}"
|
||||
assert response == "Hello world!"
|
||||
finally:
|
||||
# Clean up callback to prevent interference with other tests
|
||||
from praisonaiagents.main import sync_display_callbacks
|
||||
sync_display_callbacks.pop('interaction', None)
|
||||
|
||||
def test_single_display_with_reasoning():
|
||||
"""Test that display_interaction is called only once with reasoning steps"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
# Register callback instead of patching
|
||||
register_display_callback('interaction', mock_display_interaction)
|
||||
|
||||
try:
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
# Mock non-streaming response with reasoning
|
||||
mock_completion.return_value = {
|
||||
"choices": [{
|
||||
"message": {
|
||||
"content": "The answer is 42",
|
||||
"provider_specific_fields": {
|
||||
"reasoning_content": "Let me think about this..."
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
|
||||
llm = LLM(model="o1-preview", verbose=False, reasoning_steps=True)
|
||||
response = llm.get_response(
|
||||
prompt="What is the meaning of life?",
|
||||
verbose=True,
|
||||
stream=False,
|
||||
reasoning_steps=True
|
||||
)
|
||||
|
||||
print(f"\nResponse: {response}")
|
||||
print(f"Display calls: {len(display_calls)}")
|
||||
|
||||
assert len(display_calls) == 1, f"Expected 1 display call, got {len(display_calls)}"
|
||||
finally:
|
||||
# Clean up callback to prevent interference with other tests
|
||||
from praisonaiagents.main import sync_display_callbacks
|
||||
sync_display_callbacks.pop('interaction', None)
|
||||
|
||||
def test_single_display_with_self_reflection():
|
||||
"""Test that display_interaction is called appropriately with self-reflection"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
# Register callback instead of patching
|
||||
register_display_callback('interaction', mock_display_interaction)
|
||||
|
||||
try:
|
||||
with patch('praisonaiagents.main.display_self_reflection'):
|
||||
with patch('litellm.completion') as mock_completion:
|
||||
# First call - initial response
|
||||
# Second call - reflection
|
||||
# Third call - regenerated response
|
||||
call_count = 0
|
||||
def mock_streaming(*args, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
|
||||
if call_count == 1:
|
||||
# Initial response
|
||||
return [
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content="Initial"))]),
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content=" response"))])
|
||||
]
|
||||
elif call_count == 2:
|
||||
# Reflection
|
||||
reflection = {"reflection": "Could be better", "satisfactory": "no"}
|
||||
return [
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content=json.dumps(reflection)))])
|
||||
]
|
||||
else:
|
||||
# Final response
|
||||
return [
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content="Better"))]),
|
||||
MagicMock(choices=[MagicMock(delta=MagicMock(content=" response"))])
|
||||
]
|
||||
|
||||
mock_completion.side_effect = mock_streaming
|
||||
|
||||
llm = LLM(model="gpt-4o-mini", verbose=False, self_reflect=True, min_reflect=1, max_reflect=2)
|
||||
response = llm.get_response(
|
||||
prompt="Test prompt",
|
||||
verbose=True,
|
||||
stream=True,
|
||||
self_reflect=True,
|
||||
min_reflect=1,
|
||||
max_reflect=2
|
||||
)
|
||||
|
||||
print(f"\nResponse: {response}")
|
||||
print(f"Display calls: {len(display_calls)}")
|
||||
|
||||
# Should display only the final response
|
||||
assert len(display_calls) == 1, f"Expected 1 display call, got {len(display_calls)}"
|
||||
assert response == "Better response"
|
||||
finally:
|
||||
# Clean up callback to prevent interference with other tests
|
||||
from praisonaiagents.main import sync_display_callbacks
|
||||
sync_display_callbacks.pop('interaction', None)
|
||||
|
||||
def test_async_single_display():
|
||||
"""Test async version also prevents duplicate displays"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
import asyncio
|
||||
|
||||
# Register callback instead of patching
|
||||
register_display_callback('interaction', mock_display_interaction)
|
||||
|
||||
async def run_test():
|
||||
with patch('litellm.acompletion') as mock_acompletion:
|
||||
# Mock async streaming response
|
||||
async def async_generator():
|
||||
yield MagicMock(choices=[MagicMock(delta=MagicMock(content="Async"))])
|
||||
yield MagicMock(choices=[MagicMock(delta=MagicMock(content=" response"))])
|
||||
|
||||
mock_acompletion.return_value = async_generator()
|
||||
|
||||
llm = LLM(model="gpt-4o-mini", verbose=False)
|
||||
response = await llm.get_response_async(
|
||||
prompt="Test async",
|
||||
verbose=True,
|
||||
stream=True
|
||||
)
|
||||
|
||||
print(f"\nAsync Response: {response}")
|
||||
print(f"Display calls: {len(display_calls)}")
|
||||
|
||||
assert len(display_calls) == 1, f"Expected 1 display call, got {len(display_calls)}"
|
||||
assert response == "Async response"
|
||||
|
||||
try:
|
||||
asyncio.run(run_test())
|
||||
finally:
|
||||
# Clean up callback to prevent interference with other tests
|
||||
from praisonaiagents.main import sync_display_callbacks
|
||||
sync_display_callbacks.pop('interaction', None)
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing duplicate callback fix for issue #878...\n")
|
||||
|
||||
try:
|
||||
print("1. Testing single display without tools...")
|
||||
test_single_display_no_tools()
|
||||
print("✓ PASSED\n")
|
||||
|
||||
print("2. Testing single display with reasoning...")
|
||||
test_single_display_with_reasoning()
|
||||
print("✓ PASSED\n")
|
||||
|
||||
print("3. Testing single display with self-reflection...")
|
||||
test_single_display_with_self_reflection()
|
||||
print("✓ PASSED\n")
|
||||
|
||||
print("4. Testing async single display...")
|
||||
test_async_single_display()
|
||||
print("✓ PASSED\n")
|
||||
|
||||
print("All tests passed! The duplicate callback issue is fixed.")
|
||||
except AssertionError as e:
|
||||
print(f"✗ FAILED: {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"✗ ERROR: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,36 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Simple test to verify the self-reflection fix works"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.tools import evaluate
|
||||
|
||||
def test_self_reflection_fix():
|
||||
"""Test that self-reflection works with tools after the fix"""
|
||||
print("=== Testing Self-Reflection Fix ===")
|
||||
|
||||
# Create an agent with self-reflection and a simple tool
|
||||
agent = Agent(
|
||||
role="Math Assistant",
|
||||
goal="Solve math problems accurately",
|
||||
backstory="You are a helpful math assistant",
|
||||
self_reflect=True,
|
||||
llm="gpt-4o-mini", # Use a more widely available model
|
||||
verbose=True,
|
||||
tools=[evaluate],
|
||||
min_reflect=1,
|
||||
max_reflect=2
|
||||
)
|
||||
|
||||
# Test with a simple calculation that might trigger self-reflection
|
||||
response = agent.start("What is 25 * 17? Show your work and double-check the answer.")
|
||||
print(f"\nResponse: {response}")
|
||||
|
||||
if response:
|
||||
print("\n✅ SUCCESS: Self-reflection with tools is working!")
|
||||
return True
|
||||
else:
|
||||
print("\n❌ FAILED: Self-reflection with tools is not working")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_self_reflection_fix()
|
||||
@@ -0,0 +1,153 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify the duplicate callback fix for issue #878.
|
||||
This tests that callbacks are triggered exactly once per LLM response.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add the path to the praisonaiagents module
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src', 'praisonai-agents'))
|
||||
|
||||
from praisonaiagents import register_display_callback
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
# Track callback invocations
|
||||
callback_count = 0
|
||||
callback_details = []
|
||||
|
||||
def test_callback(message=None, response=None, **kwargs):
|
||||
"""Callback function to track invocations."""
|
||||
global callback_count, callback_details
|
||||
callback_count += 1
|
||||
callback_details.append({
|
||||
'count': callback_count,
|
||||
'message': str(message)[:50] if message else None,
|
||||
'response': str(response)[:50] if response else None,
|
||||
'kwargs': list(kwargs.keys()) if kwargs else []
|
||||
})
|
||||
print(f"🔔 CALLBACK #{callback_count}: {response[:50] if response else 'No response'}")
|
||||
|
||||
def test_verbose_true():
|
||||
"""Test with verbose=True (this was causing duplicates before the fix)."""
|
||||
global callback_count, callback_details
|
||||
callback_count = 0
|
||||
callback_details = []
|
||||
|
||||
print("🧪 Testing with verbose=True...")
|
||||
register_display_callback('interaction', test_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
response = llm.get_response(
|
||||
prompt="Say exactly: Hello World",
|
||||
verbose=True # This should NOT cause duplicate callbacks after the fix
|
||||
)
|
||||
|
||||
print(f"✅ Response: {response}")
|
||||
print(f"📊 Callback count: {callback_count}")
|
||||
|
||||
if callback_count == 1:
|
||||
print("✅ SUCCESS: Exactly 1 callback triggered (no duplicates)")
|
||||
return True
|
||||
else:
|
||||
print(f"❌ FAIL: Expected 1 callback, got {callback_count}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ ERROR: {e}")
|
||||
return False
|
||||
|
||||
def test_verbose_false():
|
||||
"""Test with verbose=False (should still work)."""
|
||||
global callback_count, callback_details
|
||||
callback_count = 0
|
||||
callback_details = []
|
||||
|
||||
print("\n🧪 Testing with verbose=False...")
|
||||
register_display_callback('interaction', test_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
response = llm.get_response(
|
||||
prompt="Say exactly: Hello Non-Verbose",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
print(f"✅ Response: {response}")
|
||||
print(f"📊 Callback count: {callback_count}")
|
||||
|
||||
if callback_count == 1:
|
||||
print("✅ SUCCESS: Exactly 1 callback triggered")
|
||||
return True
|
||||
else:
|
||||
print(f"❌ FAIL: Expected 1 callback, got {callback_count}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ ERROR: {e}")
|
||||
return False
|
||||
|
||||
def test_multiple_calls():
|
||||
"""Test multiple consecutive calls."""
|
||||
global callback_count, callback_details
|
||||
callback_count = 0
|
||||
callback_details = []
|
||||
|
||||
print("\n🧪 Testing multiple consecutive calls...")
|
||||
register_display_callback('interaction', test_callback, is_async=False)
|
||||
|
||||
try:
|
||||
llm = LLM(model="gemini/gemini-2.5-flash-lite-preview-06-17", verbose=False)
|
||||
|
||||
for i in range(3):
|
||||
response = llm.get_response(
|
||||
prompt=f"Say exactly: Call {i+1}",
|
||||
verbose=True
|
||||
)
|
||||
print(f" Call {i+1}: {response}")
|
||||
|
||||
print(f"📊 Total callback count: {callback_count}")
|
||||
|
||||
if callback_count == 3:
|
||||
print("✅ SUCCESS: Exactly 3 callbacks for 3 calls (1 each)")
|
||||
return True
|
||||
else:
|
||||
print(f"❌ FAIL: Expected 3 callbacks, got {callback_count}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ ERROR: {e}")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("🔧 TESTING DUPLICATE CALLBACK FIX FOR ISSUE #878")
|
||||
print("=" * 60)
|
||||
|
||||
results = []
|
||||
|
||||
# Test 1: verbose=True (the main issue)
|
||||
results.append(test_verbose_true())
|
||||
|
||||
# Test 2: verbose=False (should still work)
|
||||
results.append(test_verbose_false())
|
||||
|
||||
# Test 3: multiple calls
|
||||
results.append(test_multiple_calls())
|
||||
|
||||
# Summary
|
||||
print(f"\n🏁 SUMMARY")
|
||||
print("=" * 30)
|
||||
|
||||
success_count = sum(results)
|
||||
total_tests = len(results)
|
||||
|
||||
print(f"Tests passed: {success_count}/{total_tests}")
|
||||
|
||||
if success_count == total_tests:
|
||||
print("✅ ALL TESTS PASSED! The duplicate callback issue is FIXED.")
|
||||
sys.exit(0)
|
||||
else:
|
||||
print("❌ SOME TESTS FAILED! The duplicate callback issue still exists.")
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,154 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test script for Gemini embedding support in PraisonAI"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import logging
|
||||
|
||||
# Add the source directory to Python path
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src/praisonai-agents'))
|
||||
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def test_gemini_embedding():
|
||||
"""Test Gemini embedding model integration"""
|
||||
|
||||
# Test with OpenAI embedding (backward compatibility)
|
||||
print("\n" + "="*50)
|
||||
print("Testing OpenAI Embedding (Backward Compatibility)")
|
||||
print("="*50)
|
||||
|
||||
# Create a simple agent
|
||||
agent = Agent(
|
||||
name="Test Agent",
|
||||
role="Memory Tester",
|
||||
goal="Test memory storage with embeddings",
|
||||
backstory="An agent that tests memory functionality",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create a task
|
||||
task = Task(
|
||||
description="Remember this fact: The capital of France is Paris. This is a test fact for memory storage.",
|
||||
expected_output="Confirmation that the fact has been stored",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
# Test with default OpenAI embeddings
|
||||
agents_openai = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
verbose=5,
|
||||
memory=True
|
||||
)
|
||||
|
||||
print("\nRunning with default OpenAI embeddings...")
|
||||
agents_openai.start()
|
||||
|
||||
# Test with explicit OpenAI embedding config
|
||||
print("\n" + "="*50)
|
||||
print("Testing Explicit OpenAI Embedding Configuration")
|
||||
print("="*50)
|
||||
|
||||
agent2 = Agent(
|
||||
name="Test Agent 2",
|
||||
role="Memory Tester",
|
||||
goal="Test memory storage with explicit OpenAI embeddings",
|
||||
backstory="An agent that tests memory functionality",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
task2 = Task(
|
||||
description="Remember this fact: The Eiffel Tower is 330 meters tall. This is another test fact.",
|
||||
expected_output="Confirmation that the fact has been stored",
|
||||
agent=agent2
|
||||
)
|
||||
|
||||
agents_openai_explicit = PraisonAIAgents(
|
||||
agents=[agent2],
|
||||
tasks=[task2],
|
||||
verbose=5,
|
||||
memory=True,
|
||||
embedder={
|
||||
"provider": "openai",
|
||||
"config": {
|
||||
"model": "text-embedding-3-small"
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
print("\nRunning with explicit OpenAI embedding config...")
|
||||
agents_openai_explicit.start()
|
||||
|
||||
# Test with Gemini embeddings
|
||||
print("\n" + "="*50)
|
||||
print("Testing Gemini Embedding Model")
|
||||
print("="*50)
|
||||
|
||||
agent3 = Agent(
|
||||
name="Gemini Test Agent",
|
||||
role="Memory Tester with Gemini",
|
||||
goal="Test memory storage with Gemini embeddings",
|
||||
backstory="An agent that tests Gemini embedding functionality",
|
||||
llm="gpt-4o-mini" # Still use GPT for chat, but Gemini for embeddings
|
||||
)
|
||||
|
||||
task3 = Task(
|
||||
description="Remember this fact: Google released Gemini in 2023. This tests Gemini embedding storage.",
|
||||
expected_output="Confirmation that the fact has been stored using Gemini embeddings",
|
||||
agent=agent3
|
||||
)
|
||||
|
||||
# Configure with Gemini embeddings
|
||||
agents_gemini = PraisonAIAgents(
|
||||
agents=[agent3],
|
||||
tasks=[task3],
|
||||
verbose=5,
|
||||
memory=True,
|
||||
embedder={
|
||||
"provider": "gemini",
|
||||
"config": {
|
||||
"model": "text-embedding-004" # Gemini embedding model
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
print("\nRunning with Gemini embeddings...")
|
||||
try:
|
||||
agents_gemini.start()
|
||||
print("\n✅ Gemini embedding test completed successfully!")
|
||||
except Exception as e:
|
||||
print(f"\n❌ Gemini embedding test failed: {e}")
|
||||
logger.error(f"Error details: {e}", exc_info=True)
|
||||
|
||||
# Test memory search across different embedding models
|
||||
print("\n" + "="*50)
|
||||
print("Testing Memory Search")
|
||||
print("="*50)
|
||||
|
||||
if agents_gemini.shared_memory:
|
||||
print("\nSearching for 'Gemini' in memory...")
|
||||
results = agents_gemini.shared_memory.search_long_term("Gemini", limit=5)
|
||||
print(f"Found {len(results)} results")
|
||||
for i, result in enumerate(results):
|
||||
print(f"\nResult {i+1}:")
|
||||
print(f" Text: {result.get('text', '')[:100]}...")
|
||||
print(f" Score: {result.get('score', 'N/A')}")
|
||||
|
||||
print("\n" + "="*50)
|
||||
print("All tests completed!")
|
||||
print("="*50)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Check if Google API key is set
|
||||
if not os.environ.get("GOOGLE_API_KEY"):
|
||||
print("\n⚠️ Warning: GOOGLE_API_KEY not found in environment variables.")
|
||||
print(" Gemini embedding test will likely fail without it.")
|
||||
print(" Set it with: export GOOGLE_API_KEY='your-api-key'")
|
||||
|
||||
test_gemini_embedding()
|
||||
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test different guardrail return types to ensure both GuardrailResult
|
||||
and Tuple[bool, Any] work correctly. This addresses issue #875.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
from typing import Tuple, Any
|
||||
|
||||
from praisonaiagents import Agent, Task, TaskOutput, GuardrailResult
|
||||
import inspect
|
||||
|
||||
|
||||
def test_guardrail_return_types():
|
||||
"""Test that both GuardrailResult and Tuple[bool, Any] return types are accepted."""
|
||||
print("Testing guardrail return type validation...")
|
||||
|
||||
# Test 1: GuardrailResult return type
|
||||
def validate_with_guardrailresult(task_output: TaskOutput) -> GuardrailResult:
|
||||
"""Validation function returning GuardrailResult."""
|
||||
if len(task_output.raw) < 10:
|
||||
return GuardrailResult(
|
||||
success=False,
|
||||
result=None,
|
||||
error="Output is too short"
|
||||
)
|
||||
return GuardrailResult(
|
||||
success=True,
|
||||
result=task_output,
|
||||
error=""
|
||||
)
|
||||
|
||||
# Test 2: Tuple[bool, Any] return type
|
||||
def validate_with_tuple(task_output: TaskOutput) -> Tuple[bool, Any]:
|
||||
"""Validation function returning tuple."""
|
||||
if len(task_output.raw) < 10:
|
||||
return False, "Output is too short"
|
||||
return True, task_output
|
||||
|
||||
# Test 3: No return type annotation
|
||||
def validate_no_annotation(task_output: TaskOutput):
|
||||
"""Validation function without return type annotation."""
|
||||
if len(task_output.raw) < 10:
|
||||
return False, "Output is too short"
|
||||
return True, task_output
|
||||
|
||||
# Create test agent
|
||||
agent = Agent(
|
||||
name="Test Agent",
|
||||
role="Tester",
|
||||
goal="Test guardrails",
|
||||
backstory="Testing guardrail functionality",
|
||||
llm="gpt-4o"
|
||||
)
|
||||
|
||||
# Test creating tasks with each guardrail type
|
||||
tests = [
|
||||
("GuardrailResult return type", validate_with_guardrailresult),
|
||||
("Tuple[bool, Any] return type", validate_with_tuple),
|
||||
("No return type annotation", validate_no_annotation)
|
||||
]
|
||||
|
||||
for test_name, guardrail_func in tests:
|
||||
print(f"\nTesting {test_name}...")
|
||||
try:
|
||||
task = Task(
|
||||
description="Test task",
|
||||
expected_output="Test output",
|
||||
agent=agent,
|
||||
guardrail=guardrail_func
|
||||
)
|
||||
print(f"✓ Task created successfully with {test_name}")
|
||||
|
||||
# Verify guardrail was set up
|
||||
assert task._guardrail_fn is not None, "Guardrail function not set"
|
||||
|
||||
# Test guardrail processing
|
||||
test_output = TaskOutput(
|
||||
description="Test",
|
||||
raw="This is a test output that is long enough",
|
||||
agent="Test Agent"
|
||||
)
|
||||
|
||||
result = task._process_guardrail(test_output)
|
||||
assert isinstance(result, GuardrailResult), f"Expected GuardrailResult, got {type(result)}"
|
||||
assert result.success, "Guardrail validation should pass"
|
||||
print(f"✓ Guardrail processing works with {test_name}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Failed with {test_name}: {e}")
|
||||
return False
|
||||
|
||||
print("\n✓ All guardrail return type tests passed!")
|
||||
return True
|
||||
|
||||
|
||||
def test_guardrail_result_handling():
|
||||
"""Test that _process_guardrail correctly handles both return types."""
|
||||
print("\nTesting guardrail result handling...")
|
||||
|
||||
# Create test agent and task
|
||||
agent = Agent(
|
||||
name="Test Agent",
|
||||
role="Tester",
|
||||
goal="Test guardrails",
|
||||
backstory="Testing guardrail functionality",
|
||||
llm="gpt-4o"
|
||||
)
|
||||
|
||||
test_output = TaskOutput(
|
||||
description="Test",
|
||||
raw="Test output",
|
||||
agent="Test Agent"
|
||||
)
|
||||
|
||||
# Test 1: Function returning GuardrailResult
|
||||
def return_guardrailresult(output: TaskOutput) -> GuardrailResult:
|
||||
return GuardrailResult(success=True, result=output, error="")
|
||||
|
||||
task1 = Task(
|
||||
description="Test",
|
||||
agent=agent,
|
||||
guardrail=return_guardrailresult
|
||||
)
|
||||
|
||||
result1 = task1._process_guardrail(test_output)
|
||||
assert isinstance(result1, GuardrailResult), "Should return GuardrailResult"
|
||||
assert result1.success, "Should be successful"
|
||||
print("✓ GuardrailResult return handled correctly")
|
||||
|
||||
# Test 2: Function returning tuple
|
||||
def return_tuple(output: TaskOutput) -> Tuple[bool, Any]:
|
||||
return True, output
|
||||
|
||||
task2 = Task(
|
||||
description="Test",
|
||||
agent=agent,
|
||||
guardrail=return_tuple
|
||||
)
|
||||
|
||||
result2 = task2._process_guardrail(test_output)
|
||||
assert isinstance(result2, GuardrailResult), "Should convert tuple to GuardrailResult"
|
||||
assert result2.success, "Should be successful"
|
||||
print("✓ Tuple return converted correctly")
|
||||
|
||||
print("✓ Guardrail result handling tests passed!")
|
||||
return True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Running guardrail return type tests for issue #875...")
|
||||
|
||||
success = True
|
||||
success &= test_guardrail_return_types()
|
||||
success &= test_guardrail_result_handling()
|
||||
|
||||
if success:
|
||||
print("\n✅ All tests passed! The fix for issue #875 is working correctly.")
|
||||
sys.exit(0)
|
||||
else:
|
||||
print("\n❌ Some tests failed!")
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,122 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to reproduce and verify the fix for issue #877
|
||||
where callbacks only work when verbose=True.
|
||||
"""
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, 'src/praisonai-agents')
|
||||
|
||||
from praisonaiagents import (
|
||||
register_display_callback,
|
||||
Agent,
|
||||
Task,
|
||||
PraisonAIAgents
|
||||
)
|
||||
|
||||
# Track callback execution
|
||||
callback_log = []
|
||||
|
||||
def simple_callback(message=None, response=None, **kwargs):
|
||||
callback_log.append({
|
||||
'message': message,
|
||||
'response': response,
|
||||
'kwargs': kwargs
|
||||
})
|
||||
with open('callback_log.txt', 'a') as f:
|
||||
f.write(f"Received message: {message}\n")
|
||||
f.write(f"Got response: {response}\n")
|
||||
f.write(f"Other stuff: {kwargs}\n\n")
|
||||
|
||||
# Clear the log file
|
||||
with open('callback_log.txt', 'w') as f:
|
||||
f.write("=== Testing Callback Issue #877 ===\n\n")
|
||||
|
||||
# Register as synchronous callback
|
||||
register_display_callback('interaction', simple_callback, is_async=False)
|
||||
|
||||
print("=" * 60)
|
||||
print("Testing Issue #877: Callbacks Only Work When verbose=True")
|
||||
print("=" * 60)
|
||||
|
||||
# Test Case 1: verbose=False
|
||||
print("\nTest Case 1: verbose=False")
|
||||
print("-" * 30)
|
||||
callback_log.clear()
|
||||
|
||||
agent1 = Agent(
|
||||
name="MyAgent",
|
||||
role="Assistant",
|
||||
goal="Help with tasks",
|
||||
backstory="I am a helpful assistant",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17", # Using the same model as in the issue
|
||||
verbose=False
|
||||
)
|
||||
|
||||
task1 = Task(
|
||||
name="simple_task",
|
||||
description="Say the number 1",
|
||||
agent=agent1,
|
||||
expected_output="1"
|
||||
)
|
||||
|
||||
try:
|
||||
agents1 = PraisonAIAgents(
|
||||
agents=[agent1],
|
||||
tasks=[task1]
|
||||
)
|
||||
result1 = agents1.start()
|
||||
print(f"Task completed. Callbacks executed: {len(callback_log)}")
|
||||
if len(callback_log) > 0:
|
||||
print("✅ SUCCESS: Callbacks were executed with verbose=False!")
|
||||
print(f" First callback: message='{callback_log[0]['message'][:50]}...', response='{callback_log[0]['response'][:50]}...'")
|
||||
else:
|
||||
print("❌ FAILED: No callbacks were executed with verbose=False")
|
||||
except Exception as e:
|
||||
print(f"❌ ERROR during test: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
# Test Case 2: verbose=True (for comparison)
|
||||
print("\n\nTest Case 2: verbose=True (for comparison)")
|
||||
print("-" * 30)
|
||||
callback_log.clear()
|
||||
|
||||
agent2 = Agent(
|
||||
name="MyAgent",
|
||||
role="Assistant",
|
||||
goal="Help with tasks",
|
||||
backstory="I am a helpful assistant",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
task2 = Task(
|
||||
name="simple_task",
|
||||
description="Say the number 2",
|
||||
agent=agent2,
|
||||
expected_output="2"
|
||||
)
|
||||
|
||||
try:
|
||||
agents2 = PraisonAIAgents(
|
||||
agents=[agent2],
|
||||
tasks=[task2]
|
||||
)
|
||||
result2 = agents2.start()
|
||||
print(f"\nTask completed. Callbacks executed: {len(callback_log)}")
|
||||
if len(callback_log) > 0:
|
||||
print("✅ Callbacks were executed with verbose=True as expected")
|
||||
print(f" First callback: message='{callback_log[0]['message'][:50]}...', response='{callback_log[0]['response'][:50]}...'")
|
||||
else:
|
||||
print("⚠️ WARNING: No callbacks were executed even with verbose=True")
|
||||
except Exception as e:
|
||||
print(f"❌ ERROR during test: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Test Results Summary:")
|
||||
print("-" * 30)
|
||||
print("Check callback_log.txt for detailed callback logs")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,121 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Demonstration of Issue #896 Enhancement: Callback Needs More Contextual Information
|
||||
|
||||
This script demonstrates how the enhanced callback system now provides agent_name,
|
||||
agent_role, agent_tools, task_name, task_description, and task_id to callbacks.
|
||||
|
||||
Usage:
|
||||
python test_issue_896_demo.py
|
||||
|
||||
Author: Generated by Claude Code
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src', 'praisonai-agents'))
|
||||
|
||||
from praisonaiagents import register_display_callback, display_interaction
|
||||
|
||||
def enhanced_callback(message=None, response=None, agent_name=None, agent_role=None,
|
||||
agent_tools=None, task_name=None, task_description=None,
|
||||
task_id=None, **kwargs):
|
||||
"""
|
||||
Enhanced callback function that demonstrates the new contextual information.
|
||||
|
||||
This callback now receives:
|
||||
- agent_name: The name of the agent
|
||||
- agent_role: The role of the agent
|
||||
- agent_tools: List of tools available to the agent
|
||||
- task_name: The name of the task being executed
|
||||
- task_description: The description of the task
|
||||
- task_id: The unique identifier of the task
|
||||
"""
|
||||
print("="*80)
|
||||
print("🔔 ENHANCED CALLBACK TRIGGERED")
|
||||
print("="*80)
|
||||
print(f"📝 Message: {message}")
|
||||
print(f"💬 Response: {response}")
|
||||
print(f"🤖 Agent Name: {agent_name}")
|
||||
print(f"👤 Agent Role: {agent_role}")
|
||||
print(f"🛠️ Agent Tools: {agent_tools}")
|
||||
print(f"📋 Task Name: {task_name}")
|
||||
print(f"📄 Task Description: {task_description}")
|
||||
print(f"🔢 Task ID: {task_id}")
|
||||
print(f"⚙️ Other kwargs: {kwargs}")
|
||||
print("="*80)
|
||||
print()
|
||||
|
||||
def legacy_callback(message=None, response=None, **kwargs):
|
||||
"""
|
||||
Legacy callback function that demonstrates backward compatibility.
|
||||
|
||||
This callback works exactly as before, proving that existing callbacks
|
||||
continue to function without modification.
|
||||
"""
|
||||
print("🔄 Legacy callback still works!")
|
||||
print(f"Message: {message}")
|
||||
print(f"Response: {response}")
|
||||
print(f"Other kwargs: {kwargs}")
|
||||
print()
|
||||
|
||||
def demo_enhanced_callbacks():
|
||||
"""Demonstrate the enhanced callback functionality"""
|
||||
|
||||
print("🚀 Demonstrating Enhanced Callbacks for Issue #896")
|
||||
print("="*60)
|
||||
print()
|
||||
|
||||
# Register both enhanced and legacy callbacks
|
||||
register_display_callback('interaction', enhanced_callback, is_async=False)
|
||||
|
||||
print("1. Testing enhanced callback with full context:")
|
||||
print("-" * 50)
|
||||
|
||||
# Test with full context
|
||||
display_interaction(
|
||||
message="You need to do the following task: Say the number 1.\nExpected Output: 1.\nPlease provide only the final result of your work.",
|
||||
response="1",
|
||||
agent_name="MyAgent",
|
||||
agent_role="Assistant",
|
||||
agent_tools=["calculator", "text_processor"],
|
||||
task_name="simple_task",
|
||||
task_description="Say the number 1",
|
||||
task_id="task_001",
|
||||
markdown=True,
|
||||
generation_time=1.3638858795166016
|
||||
)
|
||||
|
||||
print("2. Testing backward compatibility (missing context):")
|
||||
print("-" * 50)
|
||||
|
||||
# Test with minimal context (backward compatibility)
|
||||
display_interaction(
|
||||
message="Test message without context",
|
||||
response="Test response",
|
||||
markdown=True,
|
||||
generation_time=0.5
|
||||
)
|
||||
|
||||
print("3. Testing with partial context:")
|
||||
print("-" * 50)
|
||||
|
||||
# Test with partial context
|
||||
display_interaction(
|
||||
message="Partial context test",
|
||||
response="Response with partial context",
|
||||
agent_name="PartialAgent",
|
||||
task_name="partial_task",
|
||||
markdown=True,
|
||||
generation_time=0.8
|
||||
)
|
||||
|
||||
print("✅ All tests completed successfully!")
|
||||
print()
|
||||
print("🎉 The enhanced callback system is working correctly!")
|
||||
print(" - Callbacks now receive agent and task context")
|
||||
print(" - Backward compatibility is maintained")
|
||||
print(" - All parameters are optional with sensible defaults")
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo_enhanced_callbacks()
|
||||
@@ -0,0 +1,35 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test LLM class directly to verify self-reflection fix"""
|
||||
|
||||
from praisonaiagents.llm import LLM
|
||||
from praisonaiagents.tools import evaluate
|
||||
|
||||
def test_llm_direct():
|
||||
"""Test LLM class directly with self-reflection and tools"""
|
||||
print("=== Testing LLM Direct with Self-Reflection and Tools ===")
|
||||
|
||||
# Create LLM instance
|
||||
llm = LLM(model="gpt-4o-mini")
|
||||
|
||||
# Test with self-reflection and tools
|
||||
response = llm.get_response(
|
||||
prompt="Calculate 15 * 23 and verify your answer",
|
||||
system_prompt="You are a helpful math assistant. Use tools when needed.",
|
||||
tools=[evaluate],
|
||||
self_reflect=True,
|
||||
min_reflect=1,
|
||||
max_reflect=2,
|
||||
verbose=True
|
||||
)
|
||||
|
||||
print(f"\nResponse: {response}")
|
||||
|
||||
if response:
|
||||
print("\n✅ SUCCESS: LLM self-reflection with tools is working!")
|
||||
return True
|
||||
else:
|
||||
print("\n❌ FAILED: LLM self-reflection with tools is not working")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_llm_direct()
|
||||
@@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify that the manager LLM methods correctly handle Pydantic models.
|
||||
This tests the fix for the LLM interface misuse bug.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
from pydantic import BaseModel
|
||||
|
||||
# Add the parent directory to the Python path
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from praisonaiagents.process.process import Process
|
||||
from praisonaiagents.task.task import Task
|
||||
from praisonaiagents.agent.agent import Agent
|
||||
|
||||
class ManagerInstructions(BaseModel):
|
||||
task_id: int
|
||||
agent_name: str
|
||||
action: str
|
||||
|
||||
def test_parse_manager_instructions():
|
||||
"""Test the _parse_manager_instructions helper method"""
|
||||
print("Testing _parse_manager_instructions method...")
|
||||
|
||||
# Create a dummy process instance
|
||||
tasks = {"1": Task(name="test_task", description="Test task")}
|
||||
agents = [Agent(name="TestAgent", role="Test role", goal="Test goal")]
|
||||
process = Process(tasks=tasks, agents=agents, manager_llm="test/model")
|
||||
|
||||
# Test valid JSON response
|
||||
valid_json = '{"task_id": 1, "agent_name": "TestAgent", "action": "execute"}'
|
||||
try:
|
||||
result = process._parse_manager_instructions(valid_json, ManagerInstructions)
|
||||
assert isinstance(result, ManagerInstructions)
|
||||
assert result.task_id == 1
|
||||
assert result.agent_name == "TestAgent"
|
||||
assert result.action == "execute"
|
||||
print("✓ Valid JSON parsing works correctly")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to parse valid JSON: {e}")
|
||||
return False
|
||||
|
||||
# Test invalid JSON
|
||||
invalid_json = "not a json string"
|
||||
try:
|
||||
result = process._parse_manager_instructions(invalid_json, ManagerInstructions)
|
||||
print("✗ Should have failed on invalid JSON")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✓ Correctly raised exception for invalid JSON: {type(e).__name__}")
|
||||
|
||||
# Test JSON with missing fields
|
||||
incomplete_json = '{"task_id": 1}'
|
||||
try:
|
||||
result = process._parse_manager_instructions(incomplete_json, ManagerInstructions)
|
||||
print("✗ Should have failed on incomplete JSON")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✓ Correctly raised exception for incomplete JSON: {type(e).__name__}")
|
||||
|
||||
return True
|
||||
|
||||
def test_create_llm_instance():
|
||||
"""Test the _create_llm_instance helper method"""
|
||||
print("\nTesting _create_llm_instance method...")
|
||||
|
||||
# Create a dummy process instance
|
||||
tasks = {"1": Task(name="test_task", description="Test task")}
|
||||
agents = [Agent(name="TestAgent", role="Test role", goal="Test goal")]
|
||||
manager_llm = "gemini/gemini-2.5-flash-lite-preview-06-17"
|
||||
process = Process(tasks=tasks, agents=agents, manager_llm=manager_llm)
|
||||
|
||||
# Test LLM instance creation
|
||||
try:
|
||||
llm = process._create_llm_instance()
|
||||
assert llm.model == manager_llm
|
||||
assert llm.temperature == 0.7
|
||||
print("✓ LLM instance created successfully")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to create LLM instance: {e}")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=== Testing Manager LLM Fix ===\n")
|
||||
|
||||
# Run tests
|
||||
test_results = []
|
||||
test_results.append(("_parse_manager_instructions", test_parse_manager_instructions()))
|
||||
test_results.append(("_create_llm_instance", test_create_llm_instance()))
|
||||
|
||||
# Summary
|
||||
print("\n=== TEST SUMMARY ===")
|
||||
all_passed = True
|
||||
for test_name, passed in test_results:
|
||||
status = "PASSED" if passed else "FAILED"
|
||||
print(f"{test_name}: {status}")
|
||||
if not passed:
|
||||
all_passed = False
|
||||
|
||||
sys.exit(0 if all_passed else 1)
|
||||
@@ -0,0 +1,141 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script for multi-provider/multi-model support
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
from praisonaiagents.agent import RouterAgent
|
||||
from praisonaiagents.llm import ModelRouter, TaskComplexity
|
||||
|
||||
def test_model_router():
|
||||
"""Test the ModelRouter functionality"""
|
||||
print("=== Testing ModelRouter ===\n")
|
||||
|
||||
router = ModelRouter()
|
||||
|
||||
# Test complexity analysis
|
||||
test_tasks = [
|
||||
("What is 2 + 2?", TaskComplexity.SIMPLE),
|
||||
("Summarize the benefits of renewable energy", TaskComplexity.MODERATE),
|
||||
("Implement a binary search tree in Python", TaskComplexity.COMPLEX),
|
||||
("Design a comprehensive microservices architecture with detailed analysis", TaskComplexity.VERY_COMPLEX)
|
||||
]
|
||||
|
||||
for task, expected in test_tasks:
|
||||
complexity = router.analyze_task_complexity(task)
|
||||
print(f"Task: {task[:50]}...")
|
||||
print(f"Expected: {expected.value}, Got: {complexity.value}")
|
||||
print(f"Match: {'✓' if complexity == expected else '✗'}\n")
|
||||
|
||||
# Test model selection
|
||||
print("\n=== Testing Model Selection ===\n")
|
||||
|
||||
selected = router.select_model(
|
||||
task_description="Calculate the sum of numbers",
|
||||
budget_conscious=True
|
||||
)
|
||||
print(f"Simple task (budget mode): {selected}")
|
||||
|
||||
selected = router.select_model(
|
||||
task_description="Write a complex algorithm for graph traversal",
|
||||
budget_conscious=False
|
||||
)
|
||||
print(f"Complex task (performance mode): {selected}")
|
||||
|
||||
print("\n✓ ModelRouter tests completed")
|
||||
|
||||
|
||||
def test_router_agent():
|
||||
"""Test the RouterAgent functionality"""
|
||||
print("\n=== Testing RouterAgent ===\n")
|
||||
|
||||
# Create a router agent
|
||||
agent = RouterAgent(
|
||||
name="Test Agent",
|
||||
role="Test Assistant",
|
||||
goal="Test multi-model functionality",
|
||||
models=["gpt-4o-mini"], # Using single model for testing
|
||||
routing_strategy="auto",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
print(f"Agent created: {agent.name}")
|
||||
print(f"Available models: {list(agent.available_models.keys())}")
|
||||
print(f"Routing strategy: {agent.routing_strategy}")
|
||||
|
||||
# Test model selection
|
||||
selected = agent._select_model_for_task(
|
||||
task_description="Simple math problem",
|
||||
tools=None
|
||||
)
|
||||
print(f"Selected model for simple task: {selected}")
|
||||
|
||||
print("\n✓ RouterAgent tests completed")
|
||||
|
||||
|
||||
def test_integration():
|
||||
"""Test integration with PraisonAIAgents"""
|
||||
print("\n=== Testing Integration ===\n")
|
||||
|
||||
# Create a simple router agent
|
||||
agent = RouterAgent(
|
||||
name="Integration Test Agent",
|
||||
role="Tester",
|
||||
goal="Test the integration",
|
||||
models=["gpt-4o-mini"], # Single model for testing
|
||||
routing_strategy="manual", # Use manual to avoid complexity
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Create a simple task
|
||||
task = Task(
|
||||
name="test_task",
|
||||
description="Return the word 'success'",
|
||||
expected_output="The word success",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
# Create agents system
|
||||
agents_system = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
process="sequential",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
print("Created PraisonAIAgents with RouterAgent")
|
||||
print("✓ Integration test setup completed")
|
||||
|
||||
# Note: Actual execution would require API keys
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all tests"""
|
||||
print("\n🚀 Multi-Provider/Multi-Model Support Test Suite\n")
|
||||
|
||||
try:
|
||||
test_model_router()
|
||||
test_router_agent()
|
||||
test_integration()
|
||||
|
||||
print("\n✅ All tests completed successfully!")
|
||||
print("\n📝 Summary:")
|
||||
print("- ModelRouter can analyze task complexity")
|
||||
print("- ModelRouter can select appropriate models")
|
||||
print("- RouterAgent can be created and configured")
|
||||
print("- Integration with PraisonAIAgents works")
|
||||
print("\n🎉 Multi-provider support is ready to use!")
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n❌ Test failed with error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,177 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify self-reflection fix for issue #901
|
||||
Tests both with and without tools to ensure fix works correctly
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Add the src directory to the path
|
||||
sys.path.insert(0, str(Path(__file__).parent / "src" / "praisonai-agents"))
|
||||
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
def create_mock_tool():
|
||||
"""Create a mock tool that mimics the user's custom search tool"""
|
||||
def mock_search_tool(query: str) -> str:
|
||||
"""Mock tool that returns a fake search result"""
|
||||
return f"Mock search result for: {query}"
|
||||
|
||||
# Set the function name as expected by the tool system
|
||||
mock_search_tool.__name__ = "mock_search_tool"
|
||||
|
||||
return mock_search_tool
|
||||
|
||||
def test_self_reflection_with_tools():
|
||||
"""Test that self-reflection works when tools are present"""
|
||||
print("Testing self-reflection WITH tools...")
|
||||
|
||||
# Create mock tool
|
||||
mock_tool = create_mock_tool()
|
||||
|
||||
# Create an agent with self-reflection and tools
|
||||
agent = Agent(
|
||||
role="Senior Research Analyst",
|
||||
goal="Analyze and provide insights on given topics",
|
||||
backstory="You are an expert analyst with strong critical thinking skills",
|
||||
self_reflect=True,
|
||||
llm="gpt-4o-mini", # Using OpenAI model for testing
|
||||
verbose=True,
|
||||
tools=[mock_tool],
|
||||
max_reflect=2, # Keep it small for testing
|
||||
min_reflect=1
|
||||
)
|
||||
|
||||
# Create a task
|
||||
task = Task(
|
||||
description="Analyze the importance of AI in modern business",
|
||||
expected_output="A brief analysis report",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
# Create and start the agents
|
||||
try:
|
||||
agents = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
process="sequential"
|
||||
)
|
||||
|
||||
result = agents.start()
|
||||
print("✅ Self-reflection with tools completed successfully!")
|
||||
print(f"Result: {result}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"❌ Self-reflection with tools failed: {e}")
|
||||
return False
|
||||
|
||||
def test_self_reflection_without_tools():
|
||||
"""Test that self-reflection still works without tools (baseline)"""
|
||||
print("\nTesting self-reflection WITHOUT tools...")
|
||||
|
||||
# Create an agent with self-reflection but no tools
|
||||
agent = Agent(
|
||||
role="Senior Research Analyst",
|
||||
goal="Analyze and provide insights on given topics",
|
||||
backstory="You are an expert analyst with strong critical thinking skills",
|
||||
self_reflect=True,
|
||||
llm="gpt-4o-mini", # Using OpenAI model for testing
|
||||
verbose=True,
|
||||
max_reflect=2, # Keep it small for testing
|
||||
min_reflect=1
|
||||
)
|
||||
|
||||
# Create a task
|
||||
task = Task(
|
||||
description="Analyze the importance of AI in modern business",
|
||||
expected_output="A brief analysis report",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
# Create and start the agents
|
||||
try:
|
||||
agents = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
process="sequential"
|
||||
)
|
||||
|
||||
result = agents.start()
|
||||
print("✅ Self-reflection without tools completed successfully!")
|
||||
print(f"Result: {result}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"❌ Self-reflection without tools failed: {e}")
|
||||
return False
|
||||
|
||||
def test_llm_directly():
|
||||
"""Test LLM class directly to isolate the issue"""
|
||||
print("\nTesting LLM class directly...")
|
||||
|
||||
from praisonaiagents.llm import LLM
|
||||
|
||||
# Create mock tool
|
||||
mock_tool = create_mock_tool()
|
||||
|
||||
# Create LLM instance
|
||||
llm = LLM(model="gpt-4o-mini")
|
||||
|
||||
def mock_tool_executor(function_name, arguments):
|
||||
"""Mock tool executor that simulates tool execution"""
|
||||
if function_name == "mock_search_tool":
|
||||
return mock_tool(arguments.get("query", "test"))
|
||||
return None
|
||||
|
||||
try:
|
||||
# Test with tools and self-reflection
|
||||
response = llm.get_response(
|
||||
prompt="Analyze the importance of AI in modern business",
|
||||
tools=[mock_tool],
|
||||
self_reflect=True,
|
||||
max_reflect=2,
|
||||
min_reflect=1,
|
||||
verbose=True,
|
||||
execute_tool_fn=mock_tool_executor
|
||||
)
|
||||
|
||||
print("✅ LLM direct test with tools completed successfully!")
|
||||
print(f"Response: {response}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"❌ LLM direct test with tools failed: {e}")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Check if we have required environment variables
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
print("❌ OPENAI_API_KEY environment variable is required for testing")
|
||||
sys.exit(1)
|
||||
|
||||
print("Testing self-reflection fix for issue #901")
|
||||
print("=" * 50)
|
||||
|
||||
# Run tests
|
||||
test_results = []
|
||||
|
||||
# Test 1: Without tools (baseline)
|
||||
test_results.append(test_self_reflection_without_tools())
|
||||
|
||||
# Test 2: With tools (the fixed case)
|
||||
test_results.append(test_self_reflection_with_tools())
|
||||
|
||||
# Test 3: Direct LLM test
|
||||
test_results.append(test_llm_directly())
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 50)
|
||||
print("Test Summary:")
|
||||
print(f"✅ Tests passed: {sum(test_results)}")
|
||||
print(f"❌ Tests failed: {len(test_results) - sum(test_results)}")
|
||||
|
||||
if all(test_results):
|
||||
print("🎉 All tests passed! Self-reflection fix is working correctly.")
|
||||
sys.exit(0)
|
||||
else:
|
||||
print("💥 Some tests failed. Please check the output above.")
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,90 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test script to verify self-reflection works with tools after the fix"""
|
||||
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
from praisonaiagents.tools import duckduckgo_search
|
||||
|
||||
def test_self_reflection_with_tools():
|
||||
"""Test self-reflection with tools - should work after the fix"""
|
||||
print("=== Testing Self-Reflection WITH Tools ===")
|
||||
|
||||
# Create an agent with self-reflection and tools
|
||||
agent = Agent(
|
||||
role="Senior Research Analyst",
|
||||
goal="Analyze and provide insights on given topics",
|
||||
backstory="You are an expert analyst with strong critical thinking skills",
|
||||
self_reflect=True,
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
verbose=True,
|
||||
tools=[duckduckgo_search]
|
||||
)
|
||||
|
||||
# Create a task
|
||||
task = Task(
|
||||
description="Search for recent developments in AI and provide a brief analysis",
|
||||
expected_output="A detailed analysis report",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
# Create and start the agents
|
||||
agents = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
process="sequential"
|
||||
)
|
||||
|
||||
# Start execution
|
||||
result = agents.start()
|
||||
print(f"Result with tools: {result}")
|
||||
return result
|
||||
|
||||
def test_self_reflection_without_tools():
|
||||
"""Test self-reflection without tools - should work (baseline)"""
|
||||
print("\n=== Testing Self-Reflection WITHOUT Tools ===")
|
||||
|
||||
# Create an agent with self-reflection but no tools
|
||||
agent = Agent(
|
||||
role="Senior Research Analyst",
|
||||
goal="Analyze and provide insights on given topics",
|
||||
backstory="You are an expert analyst with strong critical thinking skills",
|
||||
self_reflect=True,
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Create a task
|
||||
task = Task(
|
||||
description="Analyze recent developments in AI",
|
||||
expected_output="A detailed analysis report",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
# Create and start the agents
|
||||
agents = PraisonAIAgents(
|
||||
agents=[agent],
|
||||
tasks=[task],
|
||||
process="sequential"
|
||||
)
|
||||
|
||||
# Start execution
|
||||
result = agents.start()
|
||||
print(f"Result without tools: {result}")
|
||||
return result
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing self-reflection fix...")
|
||||
|
||||
# Test without tools (should work)
|
||||
result_without_tools = test_self_reflection_without_tools()
|
||||
|
||||
# Test with tools (should work after fix)
|
||||
result_with_tools = test_self_reflection_with_tools()
|
||||
|
||||
print("\n=== Test Summary ===")
|
||||
print(f"Without tools: {'SUCCESS' if result_without_tools else 'FAILED'}")
|
||||
print(f"With tools: {'SUCCESS' if result_with_tools else 'FAILED'}")
|
||||
|
||||
if result_with_tools:
|
||||
print("\n✅ Fix verified: Self-reflection now works with tools!")
|
||||
else:
|
||||
print("\n❌ Fix failed: Self-reflection still not working with tools")
|
||||
@@ -0,0 +1,131 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test self-reflection functionality with tools to verify the fix for issue #901"""
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Add the src directory to the path
|
||||
sys.path.insert(0, str(Path(__file__).parent / "src" / "praisonai-agents"))
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.llm import LLM
|
||||
|
||||
def simple_calculator(operation: str, a: int, b: int) -> int:
|
||||
"""A simple calculator tool"""
|
||||
if operation == "add":
|
||||
return a + b
|
||||
elif operation == "subtract":
|
||||
return a - b
|
||||
elif operation == "multiply":
|
||||
return a * b
|
||||
elif operation == "divide":
|
||||
return a / b if b != 0 else "Cannot divide by zero"
|
||||
else:
|
||||
return "Unknown operation"
|
||||
|
||||
def test_llm_self_reflection_with_tools():
|
||||
"""Test self-reflection in LLM class directly with tools"""
|
||||
print("=== Testing LLM Self-Reflection WITH Tools ===")
|
||||
llm = LLM(model="gpt-4o-mini")
|
||||
|
||||
def mock_tool_executor(function_name, arguments):
|
||||
"""Mock tool executor for testing"""
|
||||
if function_name == "simple_calculator":
|
||||
return simple_calculator(
|
||||
arguments.get("operation", "add"),
|
||||
arguments.get("a", 0),
|
||||
arguments.get("b", 0)
|
||||
)
|
||||
return None
|
||||
|
||||
response = llm.get_response(
|
||||
prompt="Calculate 5 + 3 and then reflect on your answer",
|
||||
system_prompt="You are a helpful assistant with access to a calculator tool.",
|
||||
tools=[simple_calculator],
|
||||
self_reflect=True,
|
||||
min_reflect=1,
|
||||
max_reflect=2,
|
||||
verbose=True,
|
||||
execute_tool_fn=mock_tool_executor
|
||||
)
|
||||
print(f"LLM Response with tools: {response}")
|
||||
print()
|
||||
|
||||
def test_llm_self_reflection_without_tools():
|
||||
"""Test self-reflection in LLM class directly without tools"""
|
||||
print("=== Testing LLM Self-Reflection WITHOUT Tools ===")
|
||||
llm = LLM(model="gpt-4o-mini")
|
||||
|
||||
response = llm.get_response(
|
||||
prompt="Calculate 5 + 3 and then reflect on your answer",
|
||||
system_prompt="You are a helpful assistant.",
|
||||
self_reflect=True,
|
||||
min_reflect=1,
|
||||
max_reflect=2,
|
||||
verbose=True
|
||||
)
|
||||
print(f"LLM Response without tools: {response}")
|
||||
print()
|
||||
|
||||
def test_agent_self_reflection_with_tools():
|
||||
"""Test self-reflection in Agent class with tools"""
|
||||
print("=== Testing Agent Self-Reflection WITH Tools ===")
|
||||
agent = Agent(
|
||||
name="CalculatorAgent",
|
||||
instructions="You are a helpful assistant with access to a calculator tool.",
|
||||
llm="gpt-4o-mini",
|
||||
self_reflect=True,
|
||||
min_reflect=1,
|
||||
max_reflect=2,
|
||||
tools=[simple_calculator]
|
||||
)
|
||||
|
||||
response = agent.start("Calculate 5 + 3 and then reflect on your answer")
|
||||
print(f"Agent Response with tools: {response}")
|
||||
print()
|
||||
|
||||
def test_agent_self_reflection_without_tools():
|
||||
"""Test self-reflection in Agent class without tools"""
|
||||
print("=== Testing Agent Self-Reflection WITHOUT Tools ===")
|
||||
agent = Agent(
|
||||
name="BasicAgent",
|
||||
instructions="You are a helpful assistant.",
|
||||
llm="gpt-4o-mini",
|
||||
self_reflect=True,
|
||||
min_reflect=1,
|
||||
max_reflect=2
|
||||
)
|
||||
|
||||
response = agent.start("Calculate 5 + 3 and then reflect on your answer")
|
||||
print(f"Agent Response without tools: {response}")
|
||||
print()
|
||||
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
|
||||
# Check if we have the required API key
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
print("❌ OPENAI_API_KEY environment variable is required for testing")
|
||||
print("Set it with: export OPENAI_API_KEY='your-api-key-here'")
|
||||
exit(1)
|
||||
|
||||
print("Testing self-reflection with tools fix for issue #901")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
# Test baseline (without tools)
|
||||
test_llm_self_reflection_without_tools()
|
||||
test_agent_self_reflection_without_tools()
|
||||
|
||||
# Test the fix (with tools)
|
||||
test_llm_self_reflection_with_tools()
|
||||
test_agent_self_reflection_with_tools()
|
||||
|
||||
print("✅ All tests completed successfully!")
|
||||
print("🎉 Self-reflection with tools is now working correctly!")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Test failed: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
exit(1)
|
||||
Loaded 100 of 126 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user