mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 17:31:00 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
95ffd0c53f | ||
|
|
e1c90bf569 | ||
|
|
5fbd8aa783 | ||
|
|
d79443868f | ||
|
|
6705e753e8 | ||
|
|
b4d1532512 | ||
|
|
4b5d3b2eda | ||
|
|
a1d4f3ee43 | ||
|
|
b5e930ff2e | ||
|
|
1a6a14f8f2 | ||
|
|
ca69de4c6a | ||
|
|
9b3d36e6d7 | ||
|
|
a9359a0785 | ||
|
|
4cb016c3ca | ||
|
|
d2e4776884 | ||
|
|
193b82aeba | ||
|
|
26376a77af | ||
|
|
65791781e2 | ||
|
|
ee66154b73 | ||
|
|
2daae1f464 | ||
|
|
81bcda0440 | ||
|
|
e269cc7bf2 | ||
|
|
66e5dfcb41 | ||
|
|
ae3e69a286 | ||
|
|
40ea6fe1aa | ||
|
|
e30a7564be | ||
|
|
fbd0a736f0 | ||
|
|
eea9f757af | ||
|
|
02a8f075a6 | ||
|
|
57c39653d6 | ||
|
|
2f21d0f907 | ||
|
|
4564639f18 | ||
|
|
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 |
No files matched your search
@@ -1,4 +1,4 @@
|
||||
name: Claude Code
|
||||
name: Claude Assistant
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
@@ -6,41 +6,41 @@ 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"
|
||||
timeout_minutes: 30
|
||||
|
||||
post-review-comment:
|
||||
needs: claude
|
||||
if: success() && needs.claude.result == 'success'
|
||||
needs: claude-response
|
||||
if: success() && needs.claude-response.result == 'success'
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
@@ -40,5 +40,6 @@ jobs:
|
||||
uses: docker/build-push-action@v2
|
||||
with:
|
||||
context: .
|
||||
file: ./docker/Dockerfile
|
||||
push: true
|
||||
tags: ghcr.io/mervinpraison/praisonai:latest
|
||||
@@ -1,11 +1,18 @@
|
||||
name: Gemini Automated Issue Triage
|
||||
|
||||
on:
|
||||
issues:
|
||||
types: [opened, reopened]
|
||||
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
|
||||
|
||||
@@ -2,8 +2,6 @@ name: Gemini Issue Review
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
issues:
|
||||
types: [opened, edited]
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
@@ -16,10 +14,10 @@ jobs:
|
||||
review-issue:
|
||||
if: >
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
github.event_name == 'issues' ||
|
||||
(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'))
|
||||
|
||||
@@ -35,11 +35,20 @@ jobs:
|
||||
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: ${{ github.event.pull_request.head.sha }}
|
||||
ref: ${{ steps.checkout_ref.outputs.ref }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Get PR details
|
||||
@@ -96,12 +105,19 @@ jobs:
|
||||
"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)",
|
||||
"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",
|
||||
@@ -120,25 +136,47 @@ jobs:
|
||||
"sandbox": false
|
||||
}
|
||||
prompt: |
|
||||
You are an expert code reviewer. You have access to shell commands to gather PR information and perform the review.
|
||||
You are an expert code reviewer and AI assistant for Pull Requests. You can review code AND make modifications to improve the PR.
|
||||
|
||||
IMPORTANT: Use the available shell commands to gather information. Do not ask for information to be provided.
|
||||
**PR Details:**
|
||||
- PR Number: #$PR_NUMBER
|
||||
- Repository: $REPOSITORY
|
||||
- Additional Instructions: $ADDITIONAL_INSTRUCTIONS
|
||||
|
||||
Start by running these commands to gather the required data:
|
||||
1. Run: echo "$PR_DATA" to get PR details (JSON format)
|
||||
2. Run: echo "$CHANGED_FILES" to get the list of changed files
|
||||
3. Run: echo "$PR_NUMBER" to get the PR number
|
||||
4. Run: echo "$ADDITIONAL_INSTRUCTIONS" to see any specific review instructions from the user
|
||||
5. Run: gh pr diff $PR_NUMBER to see the full diff
|
||||
6. For any specific files, use: cat filename, head -50 filename, or tail -50 filename
|
||||
|
||||
Additional Review Instructions:
|
||||
If ADDITIONAL_INSTRUCTIONS contains text, prioritize those specific areas or focus points in your review.
|
||||
Common instruction examples: "focus on security", "check performance", "review error handling", "check for breaking changes"
|
||||
|
||||
Once you have the information, provide a comprehensive code review by:
|
||||
1. Writing your review to a file: write_file("review.md", "<your detailed review feedback here>")
|
||||
2. Posting the review: gh pr comment $PR_NUMBER --body-file review.md --repo $REPOSITORY
|
||||
**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
|
||||
|
||||
+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.53" \
|
||||
"praisonai>=2.2.55" \
|
||||
"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.53" \
|
||||
"praisonai>=2.2.55" \
|
||||
"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.53" \
|
||||
"praisonai>=2.2.55" \
|
||||
"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.53" \
|
||||
"praisonai>=2.2.55" \
|
||||
"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.53`
|
||||
- PraisonAI: `>=2.2.55`
|
||||
- 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.53" "praisonaiagents==0.0.92"
|
||||
RUN pip install "praisonai==2.2.55" "praisonaiagents==0.0.92"
|
||||
```
|
||||
|
||||
## 🌐 Production Deployment
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Example demonstrating validation feedback in workflow retry logic"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src/praisonai-agents'))
|
||||
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
# This example shows how validation feedback is now passed to retry tasks
|
||||
|
||||
print("""
|
||||
==========================================
|
||||
Validation Feedback Feature Demonstration
|
||||
==========================================
|
||||
|
||||
This example shows how the workflow retry logic now includes validation feedback.
|
||||
|
||||
Before: When validation failed, the retry task would start fresh without knowing why it failed.
|
||||
After: The retry task receives detailed feedback about what went wrong.
|
||||
|
||||
Example Scenario:
|
||||
1. A data collection task produces output
|
||||
2. A validation task checks if the output meets criteria
|
||||
3. If validation fails, the data collection task is retried WITH feedback about why it failed
|
||||
|
||||
The retry task will see:
|
||||
- Previous attempt failed validation with reason: invalid
|
||||
- Validation feedback: [specific reason for failure]
|
||||
- Rejected output: [the output that failed validation]
|
||||
- Please try again with a different approach based on this feedback.
|
||||
|
||||
This enables context-dependent tasks to improve on retry instead of repeating the same mistakes.
|
||||
==========================================
|
||||
|
||||
Implementation Details:
|
||||
1. Added 'validation_feedback' field to Task class
|
||||
2. When routing from decision task to retry task, capture validation details
|
||||
3. Include feedback in task context when building retry task description
|
||||
4. Clear feedback after use to prevent persistence
|
||||
|
||||
The solution is minimal and backward compatible - existing workflows continue to work unchanged.
|
||||
""")
|
||||
|
||||
# Example workflow structure (without actual execution)
|
||||
print("\nExample Workflow Structure:")
|
||||
print("""
|
||||
collect_task = Task(
|
||||
name="collect_data",
|
||||
description="Collect data from source",
|
||||
is_start=True,
|
||||
next_tasks=["validate_data"]
|
||||
)
|
||||
|
||||
validate_task = Task(
|
||||
name="validate_data",
|
||||
task_type="decision",
|
||||
description="Validate if data meets criteria",
|
||||
condition={
|
||||
"valid": [], # End workflow
|
||||
"invalid": ["collect_data"] # Retry with feedback
|
||||
}
|
||||
)
|
||||
|
||||
# When validation returns "invalid":
|
||||
# 1. The workflow captures the validation output
|
||||
# 2. Stores it in collect_task.validation_feedback
|
||||
# 3. On retry, collect_data task sees the feedback in its context
|
||||
# 4. Agent can adjust approach based on specific failure reason
|
||||
""")
|
||||
@@ -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")
|
||||
+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,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,52 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Multi-agent review of validation feedback implementation"""
|
||||
|
||||
print("""
|
||||
========================================
|
||||
Multi-Agent Code Review Summary
|
||||
========================================
|
||||
|
||||
Agent 1: Code Quality Reviewer
|
||||
------------------------------
|
||||
✅ Minimal code changes - only essential modifications made
|
||||
✅ Clean implementation with clear variable names
|
||||
✅ Proper logging added for debugging
|
||||
✅ No existing functionality modified
|
||||
|
||||
Agent 2: Backward Compatibility Checker
|
||||
--------------------------------------
|
||||
✅ New field 'validation_feedback' defaults to None
|
||||
✅ Existing workflows continue to work unchanged
|
||||
✅ Context building falls back to original behavior when no feedback
|
||||
✅ No breaking changes to public APIs
|
||||
|
||||
Agent 3: Feature Implementation Validator
|
||||
----------------------------------------
|
||||
✅ Captures validation failure reason correctly
|
||||
✅ Includes rejected output from failed attempt
|
||||
✅ Provides clear feedback message to retry task
|
||||
✅ Automatically clears feedback after use
|
||||
|
||||
Agent 4: Security and Performance Auditor
|
||||
----------------------------------------
|
||||
✅ No security vulnerabilities introduced
|
||||
✅ No sensitive data exposure risks
|
||||
✅ Minimal performance impact (simple dict operations)
|
||||
✅ Memory efficient - feedback cleared after use
|
||||
|
||||
Agent 5: Test Coverage Analyst
|
||||
------------------------------
|
||||
✅ Unit tests cover all new functionality
|
||||
✅ Tests verify backward compatibility
|
||||
✅ Integration examples provided
|
||||
✅ Edge cases handled (no feedback, empty feedback)
|
||||
|
||||
CONSENSUS: Implementation is APPROVED
|
||||
====================================
|
||||
|
||||
The implementation successfully addresses the issue:
|
||||
- Retry tasks now receive validation feedback
|
||||
- Context-dependent tasks can improve on retry
|
||||
- Fully backward compatible
|
||||
- Minimal, clean code changes
|
||||
""")
|
||||
@@ -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,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)
|
||||
@@ -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']
|
||||
@@ -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,24 +480,22 @@ 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_event_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 for task_id in process.asequential():
|
||||
|
||||
@@ -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]):
|
||||
|
||||
@@ -33,6 +33,12 @@ 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:
|
||||
@@ -55,5 +61,9 @@ __all__ = [
|
||||
"ToolCall",
|
||||
"process_stream_chunks",
|
||||
"supports_structured_outputs",
|
||||
"supports_streaming_with_tools"
|
||||
"supports_streaming_with_tools",
|
||||
"ModelRouter",
|
||||
"ModelProfile",
|
||||
"TaskComplexity",
|
||||
"create_routing_agent"
|
||||
]
|
||||
@@ -1,6 +1,7 @@
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Union, Literal, Callable
|
||||
from pydantic import BaseModel
|
||||
import time
|
||||
@@ -87,6 +88,10 @@ class LLM:
|
||||
"llama-3.2-90b-text-preview": 6144 # 8,192 actual
|
||||
}
|
||||
|
||||
# Ollama-specific prompt constants
|
||||
OLLAMA_TOOL_USAGE_PROMPT = "Please analyze the request and use the available tools to help answer the question. Start by identifying what information you need."
|
||||
OLLAMA_FINAL_ANSWER_PROMPT = "Based on the tool results above, please provide the final answer to the original question."
|
||||
|
||||
def _log_llm_config(self, method_name: str, **config):
|
||||
"""Centralized debug logging for LLM configuration and parameters.
|
||||
|
||||
@@ -277,15 +282,32 @@ class LLM:
|
||||
# Direct ollama/ prefix
|
||||
if self.model.startswith("ollama/"):
|
||||
return True
|
||||
|
||||
# Check base_url if provided
|
||||
if self.base_url and "ollama" in self.base_url.lower():
|
||||
return True
|
||||
|
||||
# Check environment variables for Ollama base URL
|
||||
base_url = os.getenv("OPENAI_BASE_URL", "")
|
||||
api_base = os.getenv("OPENAI_API_BASE", "")
|
||||
|
||||
# Common Ollama endpoints
|
||||
ollama_endpoints = ["localhost:11434", "127.0.0.1:11434", ":11434"]
|
||||
# Common Ollama endpoints (including custom ports)
|
||||
if any(url and ("ollama" in url.lower() or ":11434" in url)
|
||||
for url in [base_url, api_base, self.base_url or ""]):
|
||||
return True
|
||||
|
||||
return any(endpoint in base_url or endpoint in api_base for endpoint in ollama_endpoints)
|
||||
return False
|
||||
|
||||
def _format_ollama_tool_result_message(self, function_name: str, tool_result: Any) -> Dict[str, str]:
|
||||
"""
|
||||
Format tool result message for Ollama provider.
|
||||
Simplified approach without hardcoded regex extraction.
|
||||
"""
|
||||
tool_result_str = str(tool_result)
|
||||
return {
|
||||
"role": "user",
|
||||
"content": f"The {function_name} function returned: {tool_result_str}"
|
||||
}
|
||||
|
||||
def _process_stream_delta(self, delta, response_text: str, tool_calls: List[Dict], formatted_tools: Optional[List] = None) -> tuple:
|
||||
"""
|
||||
@@ -422,13 +444,22 @@ class LLM:
|
||||
"""
|
||||
messages = []
|
||||
|
||||
# Check if this is a Gemini model that supports native structured outputs
|
||||
is_gemini_with_structured_output = False
|
||||
if output_json or output_pydantic:
|
||||
from .model_capabilities import supports_structured_outputs
|
||||
is_gemini_with_structured_output = (
|
||||
self._is_gemini_model() and
|
||||
supports_structured_outputs(self.model)
|
||||
)
|
||||
|
||||
# Handle system prompt
|
||||
if system_prompt:
|
||||
# Append JSON schema if needed
|
||||
if output_json:
|
||||
system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_json.model_json_schema())}"
|
||||
elif output_pydantic:
|
||||
system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_pydantic.model_json_schema())}"
|
||||
# Only append JSON schema for non-Gemini models or Gemini models without structured output support
|
||||
if (output_json or output_pydantic) and not is_gemini_with_structured_output:
|
||||
schema_model = output_json or output_pydantic
|
||||
if schema_model and hasattr(schema_model, 'model_json_schema'):
|
||||
system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(schema_model.model_json_schema())}"
|
||||
|
||||
# Skip system messages for legacy o1 models as they don't support them
|
||||
if not self._needs_system_message_skip():
|
||||
@@ -440,7 +471,8 @@ class LLM:
|
||||
|
||||
# Handle prompt modifications for JSON output
|
||||
original_prompt = prompt
|
||||
if output_json or output_pydantic:
|
||||
if (output_json or output_pydantic) and not is_gemini_with_structured_output:
|
||||
# Only modify prompt for non-Gemini models
|
||||
if isinstance(prompt, str):
|
||||
prompt = prompt + "\nReturn ONLY a valid JSON object. No other text or explanation."
|
||||
elif isinstance(prompt, list):
|
||||
@@ -660,6 +692,7 @@ class LLM:
|
||||
|
||||
start_time = time.time()
|
||||
reflection_count = 0
|
||||
interaction_displayed = False # Track if interaction has been displayed
|
||||
|
||||
# Display initial instruction once
|
||||
if verbose:
|
||||
@@ -695,6 +728,8 @@ class LLM:
|
||||
temperature=temperature,
|
||||
stream=False, # force non-streaming
|
||||
tools=formatted_tools,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
)
|
||||
@@ -703,7 +738,7 @@ class LLM:
|
||||
final_response = resp
|
||||
|
||||
# Optionally display reasoning if present
|
||||
if verbose and reasoning_content:
|
||||
if verbose and reasoning_content and not interaction_displayed:
|
||||
display_interaction(
|
||||
original_prompt,
|
||||
f"Reasoning:\n{reasoning_content}\n\nAnswer:\n{response_text}",
|
||||
@@ -711,7 +746,8 @@ class LLM:
|
||||
generation_time=time.time() - current_time,
|
||||
console=console
|
||||
)
|
||||
else:
|
||||
interaction_displayed = True
|
||||
elif verbose and not interaction_displayed:
|
||||
display_interaction(
|
||||
original_prompt,
|
||||
response_text,
|
||||
@@ -719,6 +755,7 @@ class LLM:
|
||||
generation_time=time.time() - current_time,
|
||||
console=console
|
||||
)
|
||||
interaction_displayed = True
|
||||
|
||||
# Otherwise do the existing streaming approach
|
||||
else:
|
||||
@@ -741,6 +778,8 @@ class LLM:
|
||||
tools=formatted_tools,
|
||||
temperature=temperature,
|
||||
stream=True,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**kwargs
|
||||
)
|
||||
):
|
||||
@@ -760,6 +799,8 @@ class LLM:
|
||||
tools=formatted_tools,
|
||||
temperature=temperature,
|
||||
stream=True,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**kwargs
|
||||
)
|
||||
):
|
||||
@@ -772,7 +813,7 @@ class LLM:
|
||||
if formatted_tools and self._supports_streaming_tools():
|
||||
tool_calls = self._process_tool_calls_from_stream(delta, tool_calls)
|
||||
|
||||
response_text = response_text.strip() if response_text else "" if response_text else "" if response_text else "" if response_text else ""
|
||||
response_text = response_text.strip() if response_text else ""
|
||||
|
||||
# Create a mock final_response with the captured data
|
||||
final_response = {
|
||||
@@ -791,12 +832,14 @@ class LLM:
|
||||
tools=formatted_tools,
|
||||
temperature=temperature,
|
||||
stream=False,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**kwargs
|
||||
)
|
||||
)
|
||||
response_text = final_response["choices"][0]["message"]["content"]
|
||||
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
# Display the complete response at once
|
||||
display_interaction(
|
||||
original_prompt,
|
||||
@@ -805,18 +848,37 @@ class LLM:
|
||||
generation_time=time.time() - current_time,
|
||||
console=console
|
||||
)
|
||||
interaction_displayed = True
|
||||
|
||||
tool_calls = final_response["choices"][0]["message"].get("tool_calls")
|
||||
|
||||
# For Ollama, if response is empty but we have tools, prompt for tool usage
|
||||
if self._is_ollama_provider() and (not response_text or response_text.strip() == "") and formatted_tools and iteration_count == 0:
|
||||
messages.append({
|
||||
"role": "user",
|
||||
"content": self.OLLAMA_TOOL_USAGE_PROMPT
|
||||
})
|
||||
iteration_count += 1
|
||||
continue
|
||||
|
||||
# Handle tool calls - Sequential tool calling logic
|
||||
if tool_calls and execute_tool_fn:
|
||||
# Convert tool_calls to a serializable format for all providers
|
||||
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text,
|
||||
"tool_calls": serializable_tool_calls
|
||||
})
|
||||
# Check if this is Ollama provider
|
||||
if self._is_ollama_provider():
|
||||
# For Ollama, only include role and content
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text
|
||||
})
|
||||
else:
|
||||
# For other providers, include tool_calls
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text,
|
||||
"tool_calls": serializable_tool_calls
|
||||
})
|
||||
|
||||
should_continue = False
|
||||
tool_results = [] # Store all tool results
|
||||
@@ -842,11 +904,17 @@ class LLM:
|
||||
logging.debug(f"[TOOL_EXEC_DEBUG] About to display tool call with message: {display_message}")
|
||||
display_tool_call(display_message, console=console)
|
||||
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call_id,
|
||||
"content": json.dumps(tool_result) if tool_result is not None else "Function returned an empty output"
|
||||
})
|
||||
# Check if this is Ollama provider
|
||||
if self._is_ollama_provider():
|
||||
# For Ollama, use user role and format as natural language
|
||||
messages.append(self._format_ollama_tool_result_message(function_name, tool_result))
|
||||
else:
|
||||
# For other providers, use tool role with tool_call_id
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call_id,
|
||||
"content": json.dumps(tool_result) if tool_result is not None else "Function returned an empty output"
|
||||
})
|
||||
|
||||
# Check if we should continue (for tools like sequential thinking)
|
||||
# This mimics the logic from agent.py lines 1004-1007
|
||||
@@ -858,100 +926,12 @@ class LLM:
|
||||
iteration_count += 1
|
||||
continue
|
||||
|
||||
# Special handling for Ollama models that don't automatically process tool results
|
||||
ollama_handled = False
|
||||
ollama_params = self._handle_ollama_model(response_text, tool_results, messages, original_prompt)
|
||||
|
||||
if ollama_params:
|
||||
# Get response based on streaming mode
|
||||
if stream:
|
||||
# Streaming approach
|
||||
if verbose:
|
||||
with Live(display_generating("", start_time), console=console, refresh_per_second=4) as live:
|
||||
response_text = ""
|
||||
for chunk in litellm.completion(
|
||||
**self._build_completion_params(
|
||||
messages=ollama_params["follow_up_messages"],
|
||||
temperature=temperature,
|
||||
stream=True
|
||||
)
|
||||
):
|
||||
if chunk and chunk.choices and chunk.choices[0].delta.content:
|
||||
content = chunk.choices[0].delta.content
|
||||
response_text += content
|
||||
live.update(display_generating(response_text, start_time))
|
||||
else:
|
||||
response_text = ""
|
||||
for chunk in litellm.completion(
|
||||
**self._build_completion_params(
|
||||
messages=ollama_params["follow_up_messages"],
|
||||
temperature=temperature,
|
||||
stream=True
|
||||
)
|
||||
):
|
||||
if chunk and chunk.choices and chunk.choices[0].delta.content:
|
||||
response_text += chunk.choices[0].delta.content
|
||||
else:
|
||||
# Non-streaming approach
|
||||
resp = litellm.completion(
|
||||
**self._build_completion_params(
|
||||
messages=ollama_params["follow_up_messages"],
|
||||
temperature=temperature,
|
||||
stream=False
|
||||
)
|
||||
)
|
||||
response_text = resp.get("choices", [{}])[0].get("message", {}).get("content", "") or ""
|
||||
|
||||
# Set flag to indicate Ollama was handled
|
||||
ollama_handled = True
|
||||
final_response_text = response_text.strip() if response_text else ""
|
||||
logging.debug(f"[OLLAMA_DEBUG] Ollama follow-up response: {final_response_text[:200]}...")
|
||||
|
||||
# Display the response if we got one
|
||||
if final_response_text and verbose:
|
||||
display_interaction(
|
||||
ollama_params["original_prompt"],
|
||||
final_response_text,
|
||||
markdown=markdown,
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
|
||||
# Update messages and continue the loop instead of returning
|
||||
if final_response_text:
|
||||
# Update messages with the response to maintain conversation context
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": final_response_text
|
||||
})
|
||||
# Continue the loop to check if more tools are needed
|
||||
iteration_count += 1
|
||||
continue
|
||||
else:
|
||||
logging.warning("[OLLAMA_DEBUG] Ollama follow-up returned empty response")
|
||||
|
||||
# Handle reasoning_steps after tool execution if not already handled by Ollama
|
||||
if reasoning_steps and not ollama_handled:
|
||||
# Make a non-streaming call to capture reasoning content
|
||||
reasoning_resp = litellm.completion(
|
||||
**self._build_completion_params(
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
stream=False, # force non-streaming
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
)
|
||||
reasoning_content = reasoning_resp["choices"][0]["message"].get("provider_specific_fields", {}).get("reasoning_content")
|
||||
response_text = reasoning_resp["choices"][0]["message"]["content"]
|
||||
|
||||
# Store reasoning content for later use
|
||||
if reasoning_content:
|
||||
stored_reasoning_content = reasoning_content
|
||||
|
||||
# Update messages with the response
|
||||
# For Ollama, add explicit prompt if we need a final answer
|
||||
if self._is_ollama_provider() and iteration_count > 0:
|
||||
# Add an explicit prompt for Ollama to generate the final answer
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text
|
||||
"role": "user",
|
||||
"content": self.OLLAMA_FINAL_ANSWER_PROMPT
|
||||
})
|
||||
|
||||
# After tool execution, continue the loop to check if more tools are needed
|
||||
@@ -974,7 +954,7 @@ class LLM:
|
||||
return final_response_text
|
||||
|
||||
# No tool calls were made in this iteration, return the response
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
# If we have stored reasoning content from tool execution, display it
|
||||
if stored_reasoning_content:
|
||||
display_interaction(
|
||||
@@ -992,6 +972,7 @@ class LLM:
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
interaction_displayed = True
|
||||
|
||||
response_text = response_text.strip() if response_text else ""
|
||||
|
||||
@@ -1003,15 +984,17 @@ class LLM:
|
||||
if output_json or output_pydantic:
|
||||
self.chat_history.append({"role": "user", "content": original_prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(original_prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
|
||||
if not self_reflect:
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(original_prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
# Return reasoning content if reasoning_steps is True
|
||||
if reasoning_steps and stored_reasoning_content:
|
||||
return stored_reasoning_content
|
||||
@@ -1040,6 +1023,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=False, # Force non-streaming
|
||||
response_format={"type": "json_object"},
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
)
|
||||
@@ -1075,6 +1060,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
response_format={"type": "json_object"},
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
):
|
||||
@@ -1090,6 +1077,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
response_format={"type": "json_object"},
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
):
|
||||
@@ -1107,15 +1096,17 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
)
|
||||
|
||||
if satisfactory and reflection_count >= min_reflect - 1:
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
|
||||
if reflection_count >= max_reflect - 1:
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
|
||||
reflection_count += 1
|
||||
@@ -1135,6 +1126,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
stream=True,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**kwargs
|
||||
)
|
||||
):
|
||||
@@ -1149,21 +1142,24 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
stream=True,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**kwargs
|
||||
)
|
||||
):
|
||||
if chunk and chunk.choices and chunk.choices[0].delta.content:
|
||||
response_text += chunk.choices[0].delta.content
|
||||
|
||||
response_text = response_text.strip() if response_text else "" if response_text else ""
|
||||
response_text = response_text.strip() if response_text else ""
|
||||
continue
|
||||
|
||||
except json.JSONDecodeError:
|
||||
reflection_count += 1
|
||||
if reflection_count >= max_reflect:
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
continue
|
||||
except Exception as e:
|
||||
@@ -1171,9 +1167,10 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
return None
|
||||
|
||||
# If we've exhausted reflection attempts
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
|
||||
except Exception as error:
|
||||
@@ -1185,6 +1182,12 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
total_time = time.time() - start_time
|
||||
logging.debug(f"get_response completed in {total_time:.2f} seconds")
|
||||
|
||||
def _is_gemini_model(self) -> bool:
|
||||
"""Check if the model is a Gemini model."""
|
||||
if not self.model:
|
||||
return False
|
||||
return any(prefix in self.model.lower() for prefix in ['gemini', 'gemini/', 'google/gemini'])
|
||||
|
||||
async def get_response_async(
|
||||
self,
|
||||
prompt: Union[str, List[Dict]],
|
||||
@@ -1273,6 +1276,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
start_time = time.time()
|
||||
reflection_count = 0
|
||||
interaction_displayed = False # Track if interaction has been displayed
|
||||
|
||||
# Format tools for LiteLLM using the shared helper
|
||||
formatted_tools = self._format_tools_for_litellm(tools)
|
||||
@@ -1293,15 +1297,17 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
resp = await litellm.acompletion(
|
||||
**self._build_completion_params(
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
stream=False, # force non-streaming
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
temperature=temperature,
|
||||
stream=False, # force non-streaming
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
)
|
||||
reasoning_content = resp["choices"][0]["message"].get("provider_specific_fields", {}).get("reasoning_content")
|
||||
response_text = resp["choices"][0]["message"]["content"]
|
||||
|
||||
if verbose and reasoning_content:
|
||||
if verbose and reasoning_content and not interaction_displayed:
|
||||
display_interaction(
|
||||
"Initial reasoning:",
|
||||
f"Reasoning:\n{reasoning_content}\n\nAnswer:\n{response_text}",
|
||||
@@ -1309,7 +1315,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
elif verbose:
|
||||
interaction_displayed = True
|
||||
elif verbose and not interaction_displayed:
|
||||
display_interaction(
|
||||
"Initial response:",
|
||||
response_text,
|
||||
@@ -1317,6 +1324,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
interaction_displayed = True
|
||||
else:
|
||||
# Determine if we should use streaming based on tool support
|
||||
use_streaming = stream
|
||||
@@ -1335,6 +1343,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=True,
|
||||
tools=formatted_tools,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**kwargs
|
||||
)
|
||||
):
|
||||
@@ -1355,6 +1365,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=True,
|
||||
tools=formatted_tools,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**kwargs
|
||||
)
|
||||
):
|
||||
@@ -1367,7 +1379,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if formatted_tools and self._supports_streaming_tools():
|
||||
tool_calls = self._process_tool_calls_from_stream(delta, tool_calls)
|
||||
|
||||
response_text = response_text.strip() if response_text else "" if response_text else "" if response_text else ""
|
||||
response_text = response_text.strip() if response_text else ""
|
||||
|
||||
# We already have tool_calls from streaming if supported
|
||||
# No need for a second API call!
|
||||
@@ -1379,13 +1391,15 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=False,
|
||||
tools=formatted_tools,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
)
|
||||
response_text = tool_response.choices[0].message.get("content", "")
|
||||
tool_calls = tool_response.choices[0].message.get("tool_calls", [])
|
||||
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
# Display the complete response at once
|
||||
display_interaction(
|
||||
original_prompt,
|
||||
@@ -1394,16 +1408,35 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
interaction_displayed = True
|
||||
|
||||
# For Ollama, if response is empty but we have tools, prompt for tool usage
|
||||
if self._is_ollama_provider() and (not response_text or response_text.strip() == "") and formatted_tools and iteration_count == 0:
|
||||
messages.append({
|
||||
"role": "user",
|
||||
"content": self.OLLAMA_TOOL_USAGE_PROMPT
|
||||
})
|
||||
iteration_count += 1
|
||||
continue
|
||||
|
||||
# Now handle tools if we have them (either from streaming or non-streaming)
|
||||
if tools and execute_tool_fn and tool_calls:
|
||||
# Convert tool_calls to a serializable format for all providers
|
||||
serializable_tool_calls = self._serialize_tool_calls(tool_calls)
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text,
|
||||
"tool_calls": serializable_tool_calls
|
||||
})
|
||||
# Check if it's Ollama provider
|
||||
if self._is_ollama_provider():
|
||||
# For Ollama, only include role and content
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text
|
||||
})
|
||||
else:
|
||||
# For other providers, include tool_calls
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": response_text,
|
||||
"tool_calls": serializable_tool_calls
|
||||
})
|
||||
|
||||
tool_results = [] # Store all tool results
|
||||
for tool_call in tool_calls:
|
||||
@@ -1421,77 +1454,31 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
else:
|
||||
display_message += "Function returned no output"
|
||||
display_tool_call(display_message, console=console)
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call_id,
|
||||
"content": json.dumps(tool_result) if tool_result is not None else "Function returned an empty output"
|
||||
})
|
||||
# Check if it's Ollama provider
|
||||
if self._is_ollama_provider():
|
||||
# For Ollama, use user role and format as natural language
|
||||
messages.append(self._format_ollama_tool_result_message(function_name, tool_result))
|
||||
else:
|
||||
# For other providers, use tool role with tool_call_id
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call_id,
|
||||
"content": json.dumps(tool_result) if tool_result is not None else "Function returned an empty output"
|
||||
})
|
||||
|
||||
# For Ollama, add explicit prompt if we need a final answer
|
||||
if self._is_ollama_provider() and iteration_count > 0:
|
||||
# Add an explicit prompt for Ollama to generate the final answer
|
||||
messages.append({
|
||||
"role": "user",
|
||||
"content": self.OLLAMA_FINAL_ANSWER_PROMPT
|
||||
})
|
||||
|
||||
# Get response after tool calls
|
||||
response_text = ""
|
||||
|
||||
# Special handling for Ollama models that don't automatically process tool results
|
||||
ollama_handled = False
|
||||
ollama_params = self._handle_ollama_model(response_text, tool_results, messages, original_prompt)
|
||||
|
||||
if ollama_params:
|
||||
# Get response with streaming
|
||||
if verbose:
|
||||
response_text = ""
|
||||
async for chunk in await litellm.acompletion(
|
||||
**self._build_completion_params(
|
||||
messages=ollama_params["follow_up_messages"],
|
||||
temperature=temperature,
|
||||
stream=stream
|
||||
)
|
||||
):
|
||||
if chunk and chunk.choices and chunk.choices[0].delta.content:
|
||||
content = chunk.choices[0].delta.content
|
||||
response_text += content
|
||||
print("\033[K", end="\r")
|
||||
print(f"Processing results... {time.time() - start_time:.1f}s", end="\r")
|
||||
else:
|
||||
response_text = ""
|
||||
async for chunk in await litellm.acompletion(
|
||||
**self._build_completion_params(
|
||||
messages=ollama_params["follow_up_messages"],
|
||||
temperature=temperature,
|
||||
stream=stream
|
||||
)
|
||||
):
|
||||
if chunk and chunk.choices and chunk.choices[0].delta.content:
|
||||
response_text += chunk.choices[0].delta.content
|
||||
|
||||
# Set flag to indicate Ollama was handled
|
||||
ollama_handled = True
|
||||
final_response_text = response_text.strip()
|
||||
logging.debug(f"[OLLAMA_DEBUG] Ollama follow-up response: {final_response_text[:200]}...")
|
||||
|
||||
# Display the response if we got one
|
||||
if final_response_text and verbose:
|
||||
display_interaction(
|
||||
ollama_params["original_prompt"],
|
||||
final_response_text,
|
||||
markdown=markdown,
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
|
||||
# Store the response for potential final return
|
||||
if final_response_text:
|
||||
# Update messages with the response to maintain conversation context
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": final_response_text
|
||||
})
|
||||
# Continue the loop to check if more tools are needed
|
||||
iteration_count += 1
|
||||
continue
|
||||
else:
|
||||
logging.warning("[OLLAMA_DEBUG] Ollama follow-up returned empty response")
|
||||
|
||||
# If no special handling was needed or if it's not an Ollama model
|
||||
if reasoning_steps and not ollama_handled:
|
||||
# If no special handling was needed
|
||||
if reasoning_steps:
|
||||
# Non-streaming call to capture reasoning
|
||||
resp = await litellm.acompletion(
|
||||
**self._build_completion_params(
|
||||
@@ -1499,13 +1486,15 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=False, # force non-streaming
|
||||
tools=formatted_tools, # Include tools
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
)
|
||||
reasoning_content = resp["choices"][0]["message"].get("provider_specific_fields", {}).get("reasoning_content")
|
||||
response_text = resp["choices"][0]["message"]["content"]
|
||||
|
||||
if verbose and reasoning_content:
|
||||
if verbose and reasoning_content and not interaction_displayed:
|
||||
display_interaction(
|
||||
"Tool response reasoning:",
|
||||
f"Reasoning:\n{reasoning_content}\n\nAnswer:\n{response_text}",
|
||||
@@ -1513,7 +1502,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
elif verbose:
|
||||
interaction_displayed = True
|
||||
elif verbose and not interaction_displayed:
|
||||
display_interaction(
|
||||
"Tool response:",
|
||||
response_text,
|
||||
@@ -1521,7 +1511,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
generation_time=time.time() - start_time,
|
||||
console=console
|
||||
)
|
||||
elif not ollama_handled:
|
||||
interaction_displayed = True
|
||||
else:
|
||||
# Get response after tool calls with streaming if not already handled
|
||||
if verbose:
|
||||
async for chunk in await litellm.acompletion(
|
||||
@@ -1530,6 +1521,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
tools=formatted_tools,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
):
|
||||
@@ -1545,13 +1538,15 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
):
|
||||
if chunk and chunk.choices and chunk.choices[0].delta.content:
|
||||
response_text += chunk.choices[0].delta.content
|
||||
|
||||
response_text = response_text.strip() if response_text else "" if response_text else ""
|
||||
response_text = response_text.strip() if response_text else ""
|
||||
|
||||
# After tool execution, update messages and continue the loop
|
||||
if response_text:
|
||||
@@ -1578,9 +1573,10 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if output_json or output_pydantic:
|
||||
self.chat_history.append({"role": "user", "content": original_prompt})
|
||||
self.chat_history.append({"role": "assistant", "content": response_text})
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(original_prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
|
||||
if not self_reflect:
|
||||
@@ -1588,7 +1584,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
display_text = final_response_text if final_response_text else response_text
|
||||
|
||||
# Display with stored reasoning content if available
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
if stored_reasoning_content:
|
||||
display_interaction(
|
||||
original_prompt,
|
||||
@@ -1600,6 +1596,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
else:
|
||||
display_interaction(original_prompt, display_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
|
||||
# Return reasoning content if reasoning_steps is True and we have it
|
||||
if reasoning_steps and stored_reasoning_content:
|
||||
@@ -1627,6 +1624,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=False, # Force non-streaming
|
||||
response_format={"type": "json_object"},
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
)
|
||||
@@ -1662,6 +1661,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
response_format={"type": "json_object"},
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
):
|
||||
@@ -1677,6 +1678,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
response_format={"type": "json_object"},
|
||||
output_json=output_json,
|
||||
output_pydantic=output_pydantic,
|
||||
**{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
|
||||
)
|
||||
):
|
||||
@@ -1695,15 +1698,17 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
)
|
||||
|
||||
if satisfactory and reflection_count >= min_reflect - 1:
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
|
||||
if reflection_count >= max_reflect - 1:
|
||||
if verbose:
|
||||
if verbose and not interaction_displayed:
|
||||
display_interaction(prompt, response_text, markdown=markdown,
|
||||
generation_time=time.time() - start_time, console=console)
|
||||
interaction_displayed = True
|
||||
return response_text
|
||||
|
||||
reflection_count += 1
|
||||
@@ -1790,67 +1795,6 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
litellm.callbacks = events
|
||||
|
||||
def _handle_ollama_model(self, response_text: str, tool_results: List[Any], messages: List[Dict], original_prompt: Union[str, List[Dict]]) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Handle special Ollama model requirements when processing tool results.
|
||||
|
||||
Args:
|
||||
response_text: The initial response text from the model
|
||||
tool_results: List of tool execution results
|
||||
messages: The conversation messages list
|
||||
original_prompt: The original user prompt
|
||||
|
||||
Returns:
|
||||
Dict with follow-up parameters if Ollama needs special handling, None otherwise
|
||||
"""
|
||||
if not self._is_ollama_provider() or not tool_results:
|
||||
return None
|
||||
|
||||
# Check if the response is just a JSON tool call
|
||||
try:
|
||||
json_response = json.loads(response_text.strip() if response_text else "{}")
|
||||
if not (('name' in json_response or 'function' in json_response) and
|
||||
not any(word in response_text.lower() for word in ['summary', 'option', 'result', 'found'])):
|
||||
return None
|
||||
|
||||
logging.debug("Detected Ollama returning only tool call JSON, preparing follow-up call to process results")
|
||||
|
||||
# Extract the original user query from messages
|
||||
original_query = ""
|
||||
for msg in reversed(messages): # Look from the end to find the most recent user message
|
||||
if msg.get("role") == "user":
|
||||
content = msg.get("content", "")
|
||||
# Handle list content (multimodal)
|
||||
if isinstance(content, list):
|
||||
for item in content:
|
||||
if isinstance(item, dict) and item.get("type") == "text":
|
||||
original_query = item.get("text", "")
|
||||
break
|
||||
else:
|
||||
original_query = content
|
||||
if original_query:
|
||||
break
|
||||
|
||||
# Create a shorter follow-up prompt with all tool results
|
||||
# If there's only one result, use it directly; otherwise combine them
|
||||
if len(tool_results) == 1:
|
||||
results_text = json.dumps(tool_results[0], indent=2)
|
||||
else:
|
||||
results_text = json.dumps(tool_results, indent=2)
|
||||
|
||||
follow_up_prompt = f"Results:\n{results_text}\nProvide Answer to this Original Question based on the above results: '{original_query}'"
|
||||
logging.debug(f"[OLLAMA_DEBUG] Original query extracted: {original_query}")
|
||||
logging.debug(f"[OLLAMA_DEBUG] Follow-up prompt: {follow_up_prompt[:200]}...")
|
||||
|
||||
# Return parameters for follow-up call
|
||||
return {
|
||||
"follow_up_messages": [{"role": "user", "content": follow_up_prompt}],
|
||||
"original_prompt": original_prompt
|
||||
}
|
||||
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
# Not a JSON response or not a tool call format
|
||||
return None
|
||||
|
||||
def _build_completion_params(self, **override_params) -> Dict[str, Any]:
|
||||
"""Build parameters for litellm completion calls with all necessary config"""
|
||||
@@ -1895,11 +1839,33 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# Override with any provided parameters
|
||||
params.update(override_params)
|
||||
|
||||
# Handle structured output parameters
|
||||
output_json = override_params.get('output_json')
|
||||
output_pydantic = override_params.get('output_pydantic')
|
||||
|
||||
if output_json or output_pydantic:
|
||||
# Always remove these from params as they're not native litellm parameters
|
||||
params.pop('output_json', None)
|
||||
params.pop('output_pydantic', None)
|
||||
|
||||
# Check if this is a Gemini model that supports native structured outputs
|
||||
if self._is_gemini_model():
|
||||
from .model_capabilities import supports_structured_outputs
|
||||
schema_model = output_json or output_pydantic
|
||||
|
||||
if schema_model and hasattr(schema_model, 'model_json_schema') and supports_structured_outputs(self.model):
|
||||
schema = schema_model.model_json_schema()
|
||||
|
||||
# Gemini uses response_mime_type and response_schema
|
||||
params['response_mime_type'] = 'application/json'
|
||||
params['response_schema'] = schema
|
||||
|
||||
logging.debug(f"Using Gemini native structured output with schema: {json.dumps(schema, indent=2)}")
|
||||
|
||||
# Add tool_choice="auto" when tools are provided (unless already specified)
|
||||
if 'tools' in params and params['tools'] and 'tool_choice' not in params:
|
||||
# For Gemini models, use tool_choice to encourage tool usage
|
||||
# More comprehensive Gemini model detection
|
||||
if any(prefix in self.model.lower() for prefix in ['gemini', 'gemini/', 'google/gemini']):
|
||||
if self._is_gemini_model():
|
||||
try:
|
||||
import litellm
|
||||
# Check if model supports function calling before setting tool_choice
|
||||
|
||||
@@ -30,6 +30,16 @@ MODELS_SUPPORTING_STRUCTURED_OUTPUTS = {
|
||||
"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
|
||||
@@ -57,16 +67,23 @@ def supports_structured_outputs(model_name: str) -> bool:
|
||||
if not model_name:
|
||||
return False
|
||||
|
||||
# Strip provider prefixes (e.g., 'google/', 'openai/', etc.)
|
||||
model_without_provider = model_name
|
||||
for prefix in ['google/', 'openai/', 'anthropic/', 'gemini/', 'mistral/', 'deepseek/', 'groq/']:
|
||||
if model_name.startswith(prefix):
|
||||
model_without_provider = model_name[len(prefix):]
|
||||
break
|
||||
|
||||
# First check if it's explicitly in the NOT supporting list
|
||||
if model_name in MODELS_NOT_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
if model_without_provider in MODELS_NOT_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return False
|
||||
|
||||
# Then check if it's in the supporting list
|
||||
if model_name in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
if model_without_provider in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return True
|
||||
|
||||
# For models with version suffixes, check the base model name
|
||||
base_model = model_name.split('-2024-')[0].split('-2025-')[0]
|
||||
base_model = model_without_provider.split('-2024-')[0].split('-2025-')[0]
|
||||
if base_model in MODELS_SUPPORTING_STRUCTURED_OUTPUTS:
|
||||
return True
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -34,6 +34,29 @@ 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."""
|
||||
logging.warning(f"_create_loop_subtasks called for {loop_task.name} but method not fully implemented")
|
||||
@@ -167,24 +190,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:
|
||||
@@ -200,23 +223,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)
|
||||
@@ -224,40 +238,32 @@ 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)
|
||||
|
||||
|
||||
async def aworkflow(self) -> AsyncGenerator[str, None]:
|
||||
@@ -469,16 +475,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
|
||||
@@ -1099,16 +1107,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}")
|
||||
|
||||
|
||||
|
||||
@@ -172,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
|
||||
@@ -183,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
|
||||
@@ -447,7 +456,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:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.0.127"
|
||||
version = "0.0.129"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
requires-python = ">=3.10"
|
||||
authors = [
|
||||
|
||||
@@ -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,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,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()
|
||||
@@ -16,7 +16,7 @@ def test_example_one():
|
||||
validator_agent = Agent(
|
||||
name="validator_agent",
|
||||
instructions="You are a helpful assistant",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
llm="gemini/gemini-2.5-flash",
|
||||
self_reflect=False,
|
||||
verbose=True
|
||||
)
|
||||
@@ -24,7 +24,7 @@ def test_example_one():
|
||||
quote_agent = Agent(
|
||||
name="quote_agent",
|
||||
instructions="You are a helpful assistant who is a pro at quoting philosophers",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
llm="gemini/gemini-2.5-flash",
|
||||
self_reflect=False,
|
||||
verbose=False
|
||||
)
|
||||
@@ -84,7 +84,7 @@ def test_example_two():
|
||||
agent = Agent(
|
||||
name="agent",
|
||||
instructions="You are a helpful assistant",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
llm="gemini/gemini-2.5-flash",
|
||||
self_reflect=False,
|
||||
verbose=True
|
||||
)
|
||||
@@ -94,7 +94,7 @@ def test_example_two():
|
||||
role="Web Search Specialist",
|
||||
goal="Search the web for relevant information",
|
||||
instructions="You are a helpful assistant. You can use the tools provided to you to help the user.",
|
||||
llm="gemini/gemini-2.5-flash-lite-preview-06-17",
|
||||
llm="gemini/gemini-2.5-flash",
|
||||
self_reflect=False,
|
||||
verbose=False,
|
||||
tools=[mock_web_search]
|
||||
Generated
+1
-1
@@ -2375,7 +2375,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.0.127"
|
||||
version = "0.0.129"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -3,8 +3,8 @@ class Praisonai < Formula
|
||||
|
||||
desc "AI tools for various AI applications"
|
||||
homepage "https://github.com/MervinPraison/PraisonAI"
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.53.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.53.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.55.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.55.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -57,7 +57,7 @@ class CloudDeployer:
|
||||
file.write("FROM python:3.11-slim\n")
|
||||
file.write("WORKDIR /app\n")
|
||||
file.write("COPY . .\n")
|
||||
file.write("RUN pip install flask praisonai==2.2.53 gunicorn markdown\n")
|
||||
file.write("RUN pip install flask praisonai==2.2.55 gunicorn markdown\n")
|
||||
file.write("EXPOSE 8080\n")
|
||||
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "PraisonAI"
|
||||
version = "2.2.53"
|
||||
version = "2.2.55"
|
||||
description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human-agent collaboration."
|
||||
readme = "README.md"
|
||||
license = ""
|
||||
@@ -12,7 +12,7 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=0.0.127",
|
||||
"praisonaiagents>=0.0.129",
|
||||
"python-dotenv>=0.19.0",
|
||||
"instructor>=1.3.3",
|
||||
"PyYAML>=6.0",
|
||||
@@ -95,7 +95,7 @@ autogen = ["pyautogen>=0.2.19", "praisonai-tools>=0.0.15", "crewai"]
|
||||
|
||||
[tool.poetry]
|
||||
name = "PraisonAI"
|
||||
version = "2.2.53"
|
||||
version = "2.2.55"
|
||||
description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human-agent collaboration."
|
||||
authors = ["Mervin Praison"]
|
||||
license = ""
|
||||
@@ -113,7 +113,7 @@ python = ">=3.10,<3.13"
|
||||
rich = ">=13.7"
|
||||
markdown = ">=3.5"
|
||||
pyparsing = ">=3.0.0"
|
||||
praisonaiagents = ">=0.0.127"
|
||||
praisonaiagents = ">=0.0.129"
|
||||
python-dotenv = ">=0.19.0"
|
||||
instructor = ">=1.3.3"
|
||||
PyYAML = ">=6.0"
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
"""Unit tests for validation feedback in workflow retry logic"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from praisonaiagents import Agent, Task, TaskOutput
|
||||
from praisonaiagents.process import Process
|
||||
|
||||
|
||||
class TestValidationFeedback:
|
||||
"""Test validation feedback functionality"""
|
||||
|
||||
def test_task_has_validation_feedback_field(self):
|
||||
"""Test that Task class has validation_feedback field"""
|
||||
task = Task(
|
||||
name="test_task",
|
||||
description="Test task",
|
||||
expected_output="Test output"
|
||||
)
|
||||
assert hasattr(task, 'validation_feedback')
|
||||
assert task.validation_feedback is None
|
||||
|
||||
def test_validation_feedback_captured_on_invalid_decision(self):
|
||||
"""Test that validation feedback is captured when decision is invalid"""
|
||||
# Create mock agents
|
||||
agent1 = Mock(spec=Agent)
|
||||
agent2 = Mock(spec=Agent)
|
||||
|
||||
# Create tasks
|
||||
collect_task = Task(
|
||||
name="collect_data",
|
||||
description="Collect data",
|
||||
expected_output="Data",
|
||||
agent=agent1,
|
||||
is_start=True,
|
||||
next_tasks=["validate_data"]
|
||||
)
|
||||
|
||||
validate_task = Task(
|
||||
name="validate_data",
|
||||
description="Validate data",
|
||||
expected_output="Validation result",
|
||||
agent=agent2,
|
||||
task_type="decision",
|
||||
condition={
|
||||
"valid": [],
|
||||
"invalid": ["collect_data"]
|
||||
}
|
||||
)
|
||||
|
||||
# Set up task relationships
|
||||
validate_task.previous_tasks = ["collect_data"]
|
||||
|
||||
# Create process
|
||||
process = Process(
|
||||
agents={"agent1": agent1, "agent2": agent2},
|
||||
tasks={"collect_data": collect_task, "validate_data": validate_task},
|
||||
verbose=0
|
||||
)
|
||||
|
||||
# Simulate task execution results
|
||||
collect_task.result = TaskOutput(
|
||||
raw="Collected only 3 items",
|
||||
description="Collected data items",
|
||||
agent="agent1"
|
||||
)
|
||||
collect_task.status = "completed"
|
||||
|
||||
validate_task.result = TaskOutput(
|
||||
raw="Not enough results, need at least 10",
|
||||
description="Validation result",
|
||||
agent="agent2"
|
||||
)
|
||||
validate_task.status = "completed"
|
||||
|
||||
# Simulate workflow routing with invalid decision
|
||||
# This would normally happen in the workflow execution
|
||||
decision_str = "invalid"
|
||||
current_task = validate_task
|
||||
target_tasks = current_task.condition.get(decision_str, [])
|
||||
|
||||
if target_tasks:
|
||||
task_value = target_tasks[0]
|
||||
next_task = collect_task # This is the retry task
|
||||
|
||||
# The implementation should add validation feedback
|
||||
if decision_str in Process.VALIDATION_FAILURE_DECISIONS:
|
||||
if current_task and current_task.result:
|
||||
validated_task = collect_task # The task that was validated
|
||||
|
||||
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
|
||||
}
|
||||
next_task.validation_feedback = feedback
|
||||
|
||||
# Verify feedback was captured
|
||||
assert collect_task.validation_feedback is not None
|
||||
assert collect_task.validation_feedback['validation_response'] == 'invalid'
|
||||
assert collect_task.validation_feedback['validation_details'] == "Not enough results, need at least 10"
|
||||
assert collect_task.validation_feedback['rejected_output'] == "Collected only 3 items"
|
||||
assert collect_task.validation_feedback['validator_task'] == "validate_data"
|
||||
|
||||
def test_validation_feedback_included_in_context(self):
|
||||
"""Test that validation feedback is included in task context"""
|
||||
# Create a task with validation feedback
|
||||
task = Task(
|
||||
name="retry_task",
|
||||
description="Original description",
|
||||
expected_output="Expected output"
|
||||
)
|
||||
|
||||
task.validation_feedback = {
|
||||
'validation_response': 'invalid',
|
||||
'validation_details': 'Not enough results',
|
||||
'rejected_output': 'Previous output',
|
||||
'validator_task': 'validator'
|
||||
}
|
||||
|
||||
# Create process to test context building
|
||||
process = Process(
|
||||
agents={},
|
||||
tasks={"retry_task": task},
|
||||
verbose=0
|
||||
)
|
||||
|
||||
# Build context
|
||||
context = process._build_task_context(task)
|
||||
|
||||
# Verify feedback is in context
|
||||
assert "Previous attempt failed validation with reason: invalid" in context
|
||||
assert "Validation feedback: Not enough results" in context
|
||||
assert "Rejected output: Previous output" in context
|
||||
assert "Please try again with a different approach based on this feedback." in context
|
||||
|
||||
# Verify feedback was cleared after use
|
||||
assert task.validation_feedback is None
|
||||
|
||||
def test_validation_feedback_backward_compatibility(self):
|
||||
"""Test that tasks without validation feedback work normally"""
|
||||
task = Task(
|
||||
name="normal_task",
|
||||
description="Normal task",
|
||||
expected_output="Normal output"
|
||||
)
|
||||
|
||||
process = Process(
|
||||
agents={},
|
||||
tasks={"normal_task": task},
|
||||
verbose=0
|
||||
)
|
||||
|
||||
# Build context without validation feedback
|
||||
context = process._build_task_context(task)
|
||||
|
||||
# Should return empty string for task with no context
|
||||
assert context == ""
|
||||
|
||||
# Task should still work normally
|
||||
assert task.validation_feedback is None
|
||||
|
||||
def test_multiple_retry_decisions_supported(self):
|
||||
"""Test that various failure decision strings trigger feedback capture"""
|
||||
# Import the constant from Process class
|
||||
from praisonaiagents.process import Process
|
||||
failure_decisions = Process.VALIDATION_FAILURE_DECISIONS
|
||||
|
||||
for decision in failure_decisions:
|
||||
task = Task(
|
||||
name="test_task",
|
||||
description="Test",
|
||||
expected_output="Test"
|
||||
)
|
||||
|
||||
# Simulate the decision routing logic
|
||||
if decision in Process.VALIDATION_FAILURE_DECISIONS:
|
||||
task.validation_feedback = {
|
||||
'validation_response': decision,
|
||||
'validation_details': f"Failed with {decision}",
|
||||
'rejected_output': "Previous output",
|
||||
'validator_task': "validator"
|
||||
}
|
||||
|
||||
assert task.validation_feedback is not None
|
||||
assert task.validation_feedback['validation_response'] == decision
|
||||
|
||||
def test_validation_feedback_with_context_tasks(self):
|
||||
"""Test that validation feedback works when validated task is in context"""
|
||||
# Create mock agents
|
||||
agent1 = Mock(spec=Agent)
|
||||
agent2 = Mock(spec=Agent)
|
||||
|
||||
# Create tasks
|
||||
data_task = Task(
|
||||
name="data_task",
|
||||
description="Generate data",
|
||||
expected_output="Data",
|
||||
agent=agent1
|
||||
)
|
||||
|
||||
validate_task = Task(
|
||||
name="validate_data",
|
||||
description="Validate data",
|
||||
expected_output="Validation result",
|
||||
agent=agent2,
|
||||
task_type="decision",
|
||||
context=[data_task], # Data task is in context, not previous_tasks
|
||||
condition={
|
||||
"valid": [],
|
||||
"invalid": ["data_task"]
|
||||
}
|
||||
)
|
||||
|
||||
# Create process
|
||||
process = Process(
|
||||
agents={"agent1": agent1, "agent2": agent2},
|
||||
tasks={"data_task": data_task, "validate_data": validate_task},
|
||||
verbose=0
|
||||
)
|
||||
|
||||
# Simulate task execution results
|
||||
data_task.result = TaskOutput(
|
||||
raw="Generated data with errors",
|
||||
description="Generated data",
|
||||
agent="agent1"
|
||||
)
|
||||
data_task.status = "completed"
|
||||
|
||||
validate_task.result = TaskOutput(
|
||||
raw="Data has errors, please fix",
|
||||
description="Data validation result",
|
||||
agent="agent2"
|
||||
)
|
||||
validate_task.status = "completed"
|
||||
|
||||
# Test the improved validation feedback logic
|
||||
decision_str = "invalid"
|
||||
current_task = validate_task
|
||||
|
||||
# This simulates the improved logic that checks context when no previous_tasks
|
||||
validated_task = None
|
||||
if current_task.previous_tasks:
|
||||
prev_task_name = current_task.previous_tasks[-1]
|
||||
validated_task = data_task if data_task.name == prev_task_name else None
|
||||
elif current_task.context:
|
||||
# Check context for the validated task
|
||||
for ctx_task in reversed(current_task.context):
|
||||
if ctx_task.result and ctx_task.name != current_task.name:
|
||||
validated_task = ctx_task
|
||||
break
|
||||
|
||||
# Verify the context-based task was found
|
||||
assert validated_task is not None
|
||||
assert validated_task.name == "data_task"
|
||||
assert validated_task.result.raw == "Generated data with errors"
|
||||
Generated
+2584
-2582
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,201 @@
|
||||
#!/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 unittest.mock import patch, MagicMock
|
||||
import json
|
||||
|
||||
# Track display_interaction calls
|
||||
display_calls = []
|
||||
|
||||
def mock_display_interaction(prompt, response, markdown=True, generation_time=0, console=None):
|
||||
"""Mock display_interaction to track calls"""
|
||||
display_calls.append({
|
||||
'prompt': prompt,
|
||||
'response': response,
|
||||
'markdown': markdown,
|
||||
'generation_time': generation_time
|
||||
})
|
||||
print(f"[DISPLAY] {prompt[:50]}... -> {response[:50]}...")
|
||||
|
||||
def test_single_display_no_tools():
|
||||
"""Test that display_interaction is called only once without tools"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
with patch('praisonaiagents.main.display_interaction', side_effect=mock_display_interaction):
|
||||
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!"
|
||||
|
||||
def test_single_display_with_reasoning():
|
||||
"""Test that display_interaction is called only once with reasoning steps"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
with patch('praisonaiagents.main.display_interaction', side_effect=mock_display_interaction):
|
||||
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)}"
|
||||
|
||||
def test_single_display_with_self_reflection():
|
||||
"""Test that display_interaction is called appropriately with self-reflection"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
with patch('praisonaiagents.main.display_interaction', side_effect=mock_display_interaction):
|
||||
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"
|
||||
|
||||
def test_async_single_display():
|
||||
"""Test async version also prevents duplicate displays"""
|
||||
global display_calls
|
||||
display_calls = []
|
||||
|
||||
import asyncio
|
||||
|
||||
async def run_test():
|
||||
with patch('praisonaiagents.main.display_interaction', side_effect=mock_display_interaction):
|
||||
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"
|
||||
|
||||
asyncio.run(run_test())
|
||||
|
||||
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)
|
||||
Reference in new issue
Block a user