Kye Gomez
2 weeks ago
committed by
GitHub
commit
e839dc5206
@ -0,0 +1,150 @@
|
||||
name: Test Main Features
|
||||
|
||||
on:
|
||||
push:
|
||||
paths:
|
||||
- 'tests/test_main_features.py'
|
||||
- 'swarms/**'
|
||||
- 'requirements.txt'
|
||||
- 'pyproject.toml'
|
||||
branches: [ "master" ]
|
||||
pull_request:
|
||||
paths:
|
||||
- 'tests/test_main_features.py'
|
||||
- 'swarms/**'
|
||||
- 'requirements.txt'
|
||||
- 'pyproject.toml'
|
||||
branches: [ "master" ]
|
||||
workflow_dispatch: # Allow manual triggering
|
||||
|
||||
jobs:
|
||||
test-main-features:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v5
|
||||
|
||||
- name: Set up Python 3.10
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.10"
|
||||
|
||||
- name: Cache pip dependencies
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/pip
|
||||
key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-pip-
|
||||
|
||||
- name: Install Poetry
|
||||
run: |
|
||||
curl -sSL https://install.python-poetry.org | python3 -
|
||||
echo "$HOME/.local/bin" >> $GITHUB_PATH
|
||||
|
||||
- name: Configure Poetry
|
||||
run: |
|
||||
poetry config virtualenvs.create true
|
||||
poetry config virtualenvs.in-project true
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
poetry install --with test --no-dev
|
||||
|
||||
- name: Set up environment variables
|
||||
run: |
|
||||
echo "OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }}" >> $GITHUB_ENV
|
||||
echo "ANTHROPIC_API_KEY=${{ secrets.ANTHROPIC_API_KEY }}" >> $GITHUB_ENV
|
||||
echo "GOOGLE_API_KEY=${{ secrets.GOOGLE_API_KEY }}" >> $GITHUB_ENV
|
||||
echo "COHERE_API_KEY=${{ secrets.COHERE_API_KEY }}" >> $GITHUB_ENV
|
||||
echo "HUGGINGFACE_API_KEY=${{ secrets.HUGGINGFACE_API_KEY }}" >> $GITHUB_ENV
|
||||
echo "REPLICATE_API_KEY=${{ secrets.REPLICATE_API_KEY }}" >> $GITHUB_ENV
|
||||
echo "TOGETHER_API_KEY=${{ secrets.TOGETHER_API_KEY }}" >> $GITHUB_ENV
|
||||
|
||||
- name: Run Main Features Tests
|
||||
run: |
|
||||
cd /Users/swarms_wd/Desktop/research/swarms
|
||||
poetry run python tests/test_main_features.py
|
||||
|
||||
- name: Upload test results
|
||||
uses: actions/upload-artifact@v4
|
||||
if: always()
|
||||
with:
|
||||
name: test-results
|
||||
path: test_runs/
|
||||
retention-days: 7
|
||||
|
||||
- name: Comment on PR with test results
|
||||
if: github.event_name == 'pull_request' && always()
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const fs = require('fs');
|
||||
const path = require('path');
|
||||
|
||||
try {
|
||||
// Look for test result files
|
||||
const testRunsDir = 'test_runs';
|
||||
if (fs.existsSync(testRunsDir)) {
|
||||
const files = fs.readdirSync(testRunsDir);
|
||||
const latestReport = files
|
||||
.filter(f => f.endsWith('.md'))
|
||||
.sort()
|
||||
.pop();
|
||||
|
||||
if (latestReport) {
|
||||
const reportPath = path.join(testRunsDir, latestReport);
|
||||
const reportContent = fs.readFileSync(reportPath, 'utf8');
|
||||
|
||||
// Extract summary from markdown
|
||||
const summaryMatch = reportContent.match(/## Summary\n\n(.*?)\n\n## Detailed Results/s);
|
||||
const summary = summaryMatch ? summaryMatch[1] : 'Test results available in artifacts';
|
||||
|
||||
github.rest.issues.createComment({
|
||||
issue_number: context.issue.number,
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
body: `## Main Features Test Results\n\n${summary}\n\n📊 Full test report available in artifacts.`
|
||||
});
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.log('Could not read test results:', error.message);
|
||||
}
|
||||
|
||||
test-coverage:
|
||||
runs-on: ubuntu-latest
|
||||
if: github.event_name == 'pull_request'
|
||||
needs: test-main-features
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v5
|
||||
|
||||
- name: Set up Python 3.10
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.10"
|
||||
|
||||
- name: Install Poetry
|
||||
run: |
|
||||
curl -sSL https://install.python-poetry.org | python3 -
|
||||
echo "$HOME/.local/bin" >> $GITHUB_PATH
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
poetry install --with test
|
||||
|
||||
- name: Run coverage analysis
|
||||
run: |
|
||||
poetry run pytest tests/test_main_features.py --cov=swarms --cov-report=xml --cov-report=html
|
||||
|
||||
- name: Upload coverage to Codecov
|
||||
uses: codecov/codecov-action@v4
|
||||
with:
|
||||
file: ./coverage.xml
|
||||
flags: main-features
|
||||
name: main-features-coverage
|
||||
fail_ci_if_error: false
|
||||
@ -0,0 +1,23 @@
|
||||
from swarms import Agent
|
||||
from swarms.prompts.finance_agent_sys_prompt import (
|
||||
FINANCIAL_AGENT_SYS_PROMPT,
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
agent_name="Financial-Analysis-Agent", # Name of the agent
|
||||
agent_description="Personal finance advisor agent", # Description of the agent's role
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT, # System prompt for financial tasks
|
||||
max_loops=1,
|
||||
mcp_urls=[
|
||||
"http://0.0.0.0:5932/mcp",
|
||||
],
|
||||
model_name="gpt-4o-mini",
|
||||
output_type="all",
|
||||
)
|
||||
|
||||
out = agent.run(
|
||||
"Use the discover agent tools to find what agents are available and provide a summary"
|
||||
)
|
||||
|
||||
# Print the output from the agent's run method.
|
||||
print(out)
|
||||
@ -1,283 +1,375 @@
|
||||
# SequentialWorkflow Documentation
|
||||
|
||||
**Overview:**
|
||||
A Sequential Swarm architecture processes tasks in a linear sequence. Each agent completes its task before passing the result to the next agent in the chain. This architecture ensures orderly processing and is useful when tasks have dependencies. The system now includes **sequential awareness** features that allow agents to know about the agents ahead and behind them in the workflow, significantly enhancing coordination and context understanding. [Learn more here in the docs:](https://docs.swarms.world/en/latest/swarms/structs/agent_rearrange/)
|
||||
A Sequential Swarm architecture processes tasks in a linear sequence. Each agent completes its task before passing the result to the next agent in the chain. This architecture ensures orderly processing and is useful when tasks have dependencies.
|
||||
|
||||
**Use-Cases:**
|
||||
|
||||
- Workflows where each step depends on the previous one, such as assembly lines or sequential data processing.
|
||||
- Scenarios requiring strict order of operations.
|
||||
- **NEW**: Enhanced workflows where agents need context about their position in the sequence for better coordination.
|
||||
- Multi-step content creation, analysis, and refinement workflows.
|
||||
|
||||
```mermaid
|
||||
graph TD
|
||||
A[First Agent] --> B[Second Agent]
|
||||
B --> C[Third Agent]
|
||||
C --> D[Fourth Agent]
|
||||
|
||||
|
||||
style A fill:#e1f5fe
|
||||
style B fill:#f3e5f5
|
||||
style C fill:#e8f5e8
|
||||
style D fill:#fff3e0
|
||||
|
||||
A -.->|"Awareness: None (first)"| A
|
||||
B -.->|"Awareness: Ahead: A, Behind: C"| B
|
||||
C -.->|"Awareness: Ahead: B, Behind: D"| C
|
||||
D -.->|"Awareness: Ahead: C, Behind: None (last)"| D
|
||||
```
|
||||
|
||||
## **Sequential Awareness Feature**
|
||||
|
||||
The SequentialWorkflow now includes a powerful **sequential awareness** feature that automatically provides each agent with context about their position in the workflow:
|
||||
|
||||
### What Agents Know Automatically
|
||||
|
||||
- **Agent ahead**: The agent that completed their task before them
|
||||
- **Agent behind**: The agent that will receive their output next
|
||||
- **Workflow position**: Their step number and role in the sequence
|
||||
|
||||
### Benefits
|
||||
|
||||
1. **Better Coordination**: Agents can reference previous work and prepare output for the next step
|
||||
2. **Context Understanding**: Each agent knows their role in the larger workflow
|
||||
3. **Improved Quality**: Output is tailored for the next agent in the sequence
|
||||
4. **Enhanced Logging**: Better tracking of agent interactions and workflow progress
|
||||
|
||||
## Attributes
|
||||
|
||||
| Attribute | Type | Description |
|
||||
|------------------|---------------|--------------------------------------------------|
|
||||
| `agents` | `List[Agent]` | The list of agents in the workflow. |
|
||||
| `flow` | `str` | A string representing the order of agents. |
|
||||
| `agent_rearrange`| `AgentRearrange` | Manages the dynamic execution of agents with sequential awareness. |
|
||||
| `team_awareness` | `bool` | **NEW**: Enables sequential awareness features. Defaults to `False`. |
|
||||
| `time_enabled` | `bool` | **NEW**: Enables timestamps in conversation. Defaults to `False`. |
|
||||
| `message_id_on` | `bool` | **NEW**: Enables message IDs in conversation. Defaults to `False`. |
|
||||
| `id` | `str` | Unique identifier for the workflow instance. Defaults to `"sequential_workflow"`. |
|
||||
| `name` | `str` | Human-readable name for the workflow. Defaults to `"SequentialWorkflow"`. |
|
||||
| `description` | `str` | Description of the workflow's purpose. |
|
||||
| `agents` | `List[Union[Agent, Callable]]` | The list of agents or callables in the workflow. |
|
||||
| `max_loops` | `int` | Maximum number of times to execute the workflow. Defaults to `1`. |
|
||||
| `output_type` | `OutputType` | Format of the output from the workflow. Defaults to `"dict"`. |
|
||||
| `shared_memory_system` | `callable` | Optional callable for managing shared memory between agents. |
|
||||
| `multi_agent_collab_prompt` | `bool` | If True, appends a collaborative prompt to each agent's system prompt. |
|
||||
| `team_awareness` | `bool` | Enables sequential awareness features (passed to internal `AgentRearrange`). Defaults to `False`. |
|
||||
| `flow` | `str` | A string representing the order of agents (e.g., "Agent1 -> Agent2 -> Agent3"). |
|
||||
| `agent_rearrange`| `AgentRearrange` | Internal helper for managing agent execution. |
|
||||
|
||||
## Methods
|
||||
|
||||
### `__init__(self, agents: List[Agent] = None, max_loops: int = 1, team_awareness: bool = False, time_enabled: bool = False, message_id_on: bool = False, *args, **kwargs)`
|
||||
### `__init__(self, agents: List[Union[Agent, Callable]] = None, max_loops: int = 1, team_awareness: bool = False, *args, **kwargs)`
|
||||
|
||||
The constructor initializes the `SequentialWorkflow` object with enhanced sequential awareness capabilities.
|
||||
The constructor initializes the `SequentialWorkflow` object.
|
||||
|
||||
- **Parameters:**
|
||||
- `agents` (`List[Agent]`, optional): The list of agents in the workflow. Defaults to `None`.
|
||||
- `id` (`str`, optional): Unique identifier for the workflow. Defaults to `"sequential_workflow"`.
|
||||
- `name` (`str`, optional): Name of the workflow. Defaults to `"SequentialWorkflow"`.
|
||||
- `description` (`str`, optional): Description of the workflow. Defaults to a standard description.
|
||||
- `agents` (`List[Union[Agent, Callable]]`, optional): The list of agents or callables to execute in sequence.
|
||||
- `max_loops` (`int`, optional): The maximum number of loops to execute the workflow. Defaults to `1`.
|
||||
- `team_awareness` (`bool`, optional): **NEW**: Enables sequential awareness features. Defaults to `False`.
|
||||
- `time_enabled` (`bool`, optional): **NEW**: Enables timestamps in conversation. Defaults to `False`.
|
||||
- `message_id_on` (`bool`, optional): **NEW**: Enables message IDs in conversation. Defaults to `False`.
|
||||
- `output_type` (`OutputType`, optional): Output format for the workflow. Defaults to `"dict"`.
|
||||
- `shared_memory_system` (`callable`, optional): Callable for shared memory management. Defaults to `None`.
|
||||
- `multi_agent_collab_prompt` (`bool`, optional): If True, appends a collaborative prompt to each agent's system prompt. Defaults to `False`.
|
||||
- `team_awareness` (`bool`, optional): Enables sequential awareness features in the underlying `AgentRearrange`. Defaults to `False`.
|
||||
- `*args`: Variable length argument list.
|
||||
- `**kwargs`: Arbitrary keyword arguments.
|
||||
|
||||
### `run(self, task: str) -> str`
|
||||
### `run(self, task: str, img: Optional[str] = None, imgs: Optional[List[str]] = None, *args, **kwargs) -> str`
|
||||
|
||||
Runs the specified task through the agents in the dynamically constructed flow with enhanced sequential awareness.
|
||||
Runs the specified task through the agents in the dynamically constructed flow.
|
||||
|
||||
- **Parameters:**
|
||||
- `task` (`str`): The task for the agents to execute.
|
||||
- `img` (`Optional[str]`, optional): An optional image input for the agents.
|
||||
- `imgs` (`Optional[List[str]]`, optional): Optional list of images for the agents.
|
||||
- `*args`: Additional positional arguments.
|
||||
- `**kwargs`: Additional keyword arguments.
|
||||
|
||||
- **Returns:**
|
||||
- `str`: The final result after processing through all agents.
|
||||
- The final result after processing through all agents.
|
||||
|
||||
### `run_batched(self, tasks: List[str]) -> List[str]`
|
||||
|
||||
Executes a batch of tasks through the agents in the dynamically constructed flow.
|
||||
|
||||
- **Parameters:**
|
||||
- `tasks` (`List[str]`): A list of tasks for the agents to execute.
|
||||
|
||||
### **NEW: Sequential Awareness Methods**
|
||||
- **Returns:**
|
||||
- `List[str]`: A list of final results after processing through all agents.
|
||||
|
||||
#### `get_agent_sequential_awareness(self, agent_name: str) -> str`
|
||||
### `async run_async(self, task: str) -> str`
|
||||
|
||||
Gets the sequential awareness information for a specific agent, showing which agents come before and after in the sequence.
|
||||
Executes the specified task through the agents asynchronously.
|
||||
|
||||
- **Parameters:**
|
||||
- `agent_name` (`str`): The name of the agent to get awareness for.
|
||||
- `task` (`str`): The task for the agents to execute.
|
||||
|
||||
- **Returns:**
|
||||
- `str`: A string describing the agents ahead and behind in the sequence.
|
||||
- `str`: The final result after processing through all agents.
|
||||
|
||||
### `async run_concurrent(self, tasks: List[str]) -> List[str]`
|
||||
|
||||
#### `get_sequential_flow_structure(self) -> str`
|
||||
Executes a batch of tasks through the agents concurrently.
|
||||
|
||||
Gets the overall sequential flow structure information showing the complete workflow with relationships between agents.
|
||||
- **Parameters:**
|
||||
- `tasks` (`List[str]`): A list of tasks for the agents to execute.
|
||||
|
||||
- **Returns:**
|
||||
- `str`: A string describing the complete sequential flow structure.
|
||||
- `List[str]`: A list of final results after processing through all agents.
|
||||
|
||||
## **Usage Example with Sequential Awareness:**
|
||||
## Usage Examples
|
||||
|
||||
### Basic Sequential Workflow
|
||||
|
||||
This example demonstrates a simple two-agent workflow for researching and writing a blog post.
|
||||
|
||||
```python
|
||||
from swarms import Agent, SequentialWorkflow
|
||||
|
||||
# Initialize agents for individual tasks
|
||||
agent1 = Agent(
|
||||
agent_name="ICD-10 Code Analyzer",
|
||||
system_prompt="Analyze medical data and provide relevant ICD-10 codes.",
|
||||
model_name="gpt-4.1",
|
||||
max_loops=1,
|
||||
)
|
||||
agent2 = Agent(
|
||||
agent_name="ICD-10 Code Summarizer",
|
||||
system_prompt="Summarize the findings and suggest ICD-10 codes.",
|
||||
model_name="gpt-4.1",
|
||||
max_loops=1,
|
||||
)
|
||||
agent3 = Agent(
|
||||
agent_name="ICD-10 Code Validator",
|
||||
system_prompt="Validate and finalize the ICD-10 code recommendations.",
|
||||
model_name="gpt-4.1",
|
||||
max_loops=1,
|
||||
# Agent 1: The Researcher
|
||||
researcher = Agent(
|
||||
agent_name="Researcher",
|
||||
system_prompt="Your job is to research the provided topic and provide a detailed summary.",
|
||||
model_name="gpt-4o-mini",
|
||||
)
|
||||
|
||||
# Create the Sequential workflow with enhanced awareness
|
||||
workflow = SequentialWorkflow(
|
||||
agents=[agent1, agent2, agent3],
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
team_awareness=True, # Enable sequential awareness
|
||||
time_enabled=True, # Enable timestamps
|
||||
message_id_on=True # Enable message IDs
|
||||
# Agent 2: The Writer
|
||||
writer = Agent(
|
||||
agent_name="Writer",
|
||||
system_prompt="Your job is to take the research summary and write a beautiful, engaging blog post about it.",
|
||||
model_name="gpt-4o-mini",
|
||||
)
|
||||
|
||||
# Get workflow structure information
|
||||
flow_structure = workflow.get_sequential_flow_structure()
|
||||
print("Workflow Structure:")
|
||||
print(flow_structure)
|
||||
# Create a sequential workflow where the researcher's output feeds into the writer's input
|
||||
workflow = SequentialWorkflow(agents=[researcher, writer])
|
||||
|
||||
# Get awareness for specific agents
|
||||
analyzer_awareness = workflow.get_agent_sequential_awareness("ICD-10 Code Analyzer")
|
||||
summarizer_awareness = workflow.get_agent_sequential_awareness("ICD-10 Code Summarizer")
|
||||
validator_awareness = workflow.get_agent_sequential_awareness("ICD-10 Code Validator")
|
||||
# Run the workflow on a task
|
||||
final_post = workflow.run("The history and future of artificial intelligence")
|
||||
print(final_post)
|
||||
```
|
||||
|
||||
print(f"\nAnalyzer Awareness: {analyzer_awareness}")
|
||||
print(f"Summarizer Awareness: {summarizer_awareness}")
|
||||
print(f"Validator Awareness: {validator_awareness}")
|
||||
### Legal Practice Workflow
|
||||
|
||||
# Run the workflow
|
||||
result = workflow.run(
|
||||
"Analyze the medical report and provide the appropriate ICD-10 codes."
|
||||
)
|
||||
print(f"\nFinal Result: {result}")
|
||||
```
|
||||
This example shows how to create a sequential workflow with multiple specialized legal agents.
|
||||
|
||||
**Expected Output:**
|
||||
```
|
||||
Workflow Structure:
|
||||
Sequential Flow Structure:
|
||||
Step 1: ICD-10 Code Analyzer
|
||||
Step 2: ICD-10 Code Summarizer (follows: ICD-10 Code Analyzer) (leads to: ICD-10 Code Validator)
|
||||
Step 3: ICD-10 Code Validator (follows: ICD-10 Code Summarizer)
|
||||
|
||||
Analyzer Awareness:
|
||||
Summarizer Awareness: Sequential awareness: Agent ahead: ICD-10 Code Analyzer | Agent behind: ICD-10 Code Validator
|
||||
Validator Awareness: Sequential awareness: Agent ahead: ICD-10 Code Summarizer
|
||||
```
|
||||
```python
|
||||
from swarms import Agent, SequentialWorkflow
|
||||
|
||||
## **How Sequential Awareness Works**
|
||||
# Litigation Agent
|
||||
litigation_agent = Agent(
|
||||
agent_name="Alex Johnson",
|
||||
system_prompt="As a Litigator, you specialize in navigating the complexities of lawsuits. Your role involves analyzing intricate facts, constructing compelling arguments, and devising effective case strategies to achieve favorable outcomes for your clients.",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
### 1. **Automatic Context Injection**
|
||||
When `team_awareness=True`, the system automatically adds awareness information to each agent's conversation context before they run:
|
||||
# Corporate Attorney Agent
|
||||
corporate_agent = Agent(
|
||||
agent_name="Emily Carter",
|
||||
system_prompt="As a Corporate Attorney, you provide expert legal advice on business law matters. You guide clients on corporate structure, governance, compliance, and transactions, ensuring their business operations align with legal requirements.",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
- **First Agent**: No awareness info (starts the workflow)
|
||||
- **Middle Agents**: Receive info about both the agent ahead and behind
|
||||
- **Last Agent**: Receives info about the agent ahead only
|
||||
# IP Attorney Agent
|
||||
ip_agent = Agent(
|
||||
agent_name="Michael Smith",
|
||||
system_prompt="As an IP Attorney, your expertise lies in protecting intellectual property rights. You handle various aspects of IP law, including patents, trademarks, copyrights, and trade secrets, helping clients safeguard their innovations.",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
### 2. **Enhanced Agent Prompts**
|
||||
Each agent receives context like:
|
||||
```
|
||||
Sequential awareness: Agent ahead: ICD-10 Code Analyzer | Agent behind: ICD-10 Code Validator
|
||||
# Initialize and run the workflow
|
||||
swarm = SequentialWorkflow(
|
||||
agents=[litigation_agent, corporate_agent, ip_agent],
|
||||
name="litigation-practice",
|
||||
description="Handle all aspects of litigation with a focus on thorough legal analysis and effective case management.",
|
||||
)
|
||||
|
||||
swarm.run("Create a report on how to patent an all-new AI invention and what platforms to use and more.")
|
||||
```
|
||||
|
||||
### 3. **Improved Coordination**
|
||||
Agents can now:
|
||||
- Reference previous work more effectively
|
||||
- Prepare output specifically for the next agent
|
||||
- Understand their role in the larger workflow
|
||||
- Provide better context for subsequent steps
|
||||
### Startup Idea Validation Workflow
|
||||
|
||||
## **Advanced Usage Examples**
|
||||
This example demonstrates a 3-step process for generating, validating, and pitching a startup idea.
|
||||
|
||||
### **Example 1: Research → Analysis → Report Workflow**
|
||||
```python
|
||||
# Create specialized agents
|
||||
researcher = Agent(
|
||||
agent_name="Researcher",
|
||||
system_prompt="Conduct thorough research on the given topic."
|
||||
)
|
||||
from swarms import Agent, SequentialWorkflow
|
||||
|
||||
analyzer = Agent(
|
||||
agent_name="Data Analyzer",
|
||||
system_prompt="Analyze research data and identify key insights."
|
||||
# 1. Generate an idea
|
||||
idea_generator = Agent(
|
||||
agent_name="IdeaGenerator",
|
||||
system_prompt="Generate a unique startup idea.",
|
||||
model_name="gpt-4o-mini"
|
||||
)
|
||||
|
||||
reporter = Agent(
|
||||
agent_name="Report Writer",
|
||||
system_prompt="Write comprehensive reports based on analysis."
|
||||
# 2. Validate the idea
|
||||
validator = Agent(
|
||||
agent_name="Validator",
|
||||
system_prompt="Take this startup idea and analyze its market viability.",
|
||||
model_name="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create workflow with awareness
|
||||
workflow = SequentialWorkflow(
|
||||
agents=[researcher, analyzer, reporter],
|
||||
team_awareness=True,
|
||||
time_enabled=True
|
||||
# 3. Create a pitch
|
||||
pitch_creator = Agent(
|
||||
agent_name="PitchCreator",
|
||||
system_prompt="Write a 3-sentence elevator pitch for this validated startup idea.",
|
||||
model_name="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Run with enhanced coordination
|
||||
result = workflow.run("Research and analyze the impact of AI on healthcare")
|
||||
# Create the sequential workflow
|
||||
workflow = SequentialWorkflow(agents=[idea_generator, validator, pitch_creator])
|
||||
|
||||
# Run the workflow
|
||||
elevator_pitch = workflow.run("Generate and validate a startup idea in the AI space")
|
||||
print(elevator_pitch)
|
||||
```
|
||||
|
||||
### **Example 2: Code Review Workflow**
|
||||
### Advanced: Materials Science Workflow
|
||||
|
||||
This example shows a complex workflow with multiple specialized materials science agents.
|
||||
|
||||
```python
|
||||
# Create code review agents
|
||||
linter = Agent(
|
||||
agent_name="Code Linter",
|
||||
system_prompt="Check code for syntax errors and style violations."
|
||||
from swarms import Agent, SequentialWorkflow
|
||||
|
||||
# Chief Metallurgist
|
||||
chief_metallurgist = Agent(
|
||||
agent_name="Chief-Metallurgist",
|
||||
system_prompt="As the Chief Metallurgist, you oversee the entire alloy development process, analyzing atomic structure, phase diagrams, and composition development.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
reviewer = Agent(
|
||||
agent_name="Code Reviewer",
|
||||
system_prompt="Review code quality and suggest improvements."
|
||||
# Materials Scientist
|
||||
materials_scientist = Agent(
|
||||
agent_name="Materials-Scientist",
|
||||
system_prompt="As the Materials Scientist, you analyze physical and mechanical properties including density, thermal properties, tensile strength, and microstructure.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
tester = Agent(
|
||||
agent_name="Code Tester",
|
||||
system_prompt="Write and run tests for the reviewed code."
|
||||
# Process Engineer
|
||||
process_engineer = Agent(
|
||||
agent_name="Process-Engineer",
|
||||
system_prompt="As the Process Engineer, you develop manufacturing processes including melting procedures, heat treatment protocols, and quality control methods.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create workflow
|
||||
workflow = SequentialWorkflow(
|
||||
agents=[linter, reviewer, tester],
|
||||
team_awareness=True
|
||||
# Quality Assurance Specialist
|
||||
qa_specialist = Agent(
|
||||
agent_name="QA-Specialist",
|
||||
system_prompt="As the QA Specialist, you establish quality standards, testing protocols, and documentation requirements.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Run code review process
|
||||
result = workflow.run("Review and test the authentication module")
|
||||
# Applications Engineer
|
||||
applications_engineer = Agent(
|
||||
agent_name="Applications-Engineer",
|
||||
system_prompt="As the Applications Engineer, you analyze potential applications, performance requirements, and competitive positioning.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Cost Analyst
|
||||
cost_analyst = Agent(
|
||||
agent_name="Cost-Analyst",
|
||||
system_prompt="As the Cost Analyst, you evaluate material costs, production costs, and economic viability.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create the agent list
|
||||
agents = [
|
||||
chief_metallurgist,
|
||||
materials_scientist,
|
||||
process_engineer,
|
||||
qa_specialist,
|
||||
applications_engineer,
|
||||
cost_analyst,
|
||||
]
|
||||
|
||||
# Initialize the workflow
|
||||
swarm = SequentialWorkflow(
|
||||
name="alloy-development-system",
|
||||
agents=agents,
|
||||
)
|
||||
|
||||
# Run the workflow
|
||||
result = swarm.run(
|
||||
"""Analyze and develop a new high-strength aluminum alloy for aerospace applications
|
||||
with improved fatigue resistance and corrosion resistance compared to 7075-T6,
|
||||
while maintaining similar density and cost effectiveness."""
|
||||
)
|
||||
print(result)
|
||||
```
|
||||
|
||||
## **Notes:**
|
||||
## Configuration Options
|
||||
|
||||
### Agent Parameters
|
||||
|
||||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `agent_name` | Human-readable name for the agent | Required |
|
||||
| `system_prompt` | Detailed role description and expertise | Required |
|
||||
| `model_name` | LLM model to use | "gpt-4o-mini" |
|
||||
| `max_loops` | Maximum number of processing loops | 1 |
|
||||
|
||||
### Workflow Parameters
|
||||
|
||||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `agents` | List of agents to execute in sequence | Required |
|
||||
| `name` | Name of the workflow | "SequentialWorkflow" |
|
||||
| `description` | Description of workflow purpose | Standard description |
|
||||
| `max_loops` | Number of times to execute workflow | 1 |
|
||||
| `team_awareness` | Enable sequential awareness features | False |
|
||||
|
||||
## Best Practices
|
||||
|
||||
- **Enhanced Logging**: The workflow now logs sequential awareness information for better debugging and monitoring.
|
||||
- **Automatic Context**: No manual configuration needed - awareness is automatically provided when `team_awareness=True`.
|
||||
- **Backward Compatibility**: Existing workflows continue to work without changes.
|
||||
- **Performance**: Sequential awareness adds minimal overhead while significantly improving coordination.
|
||||
1. **Clear Agent Roles**: Give each agent a specific, well-defined role with a detailed system prompt.
|
||||
2. **Ordered Dependencies**: Arrange agents in an order that makes sense for your workflow (e.g., research before writing).
|
||||
3. **Agent Names**: Use descriptive agent names that clearly indicate their function.
|
||||
4. **System Prompts**: Write comprehensive system prompts that explain the agent's expertise and responsibilities.
|
||||
5. **Task Clarity**: Provide clear, specific tasks when calling `run()`.
|
||||
|
||||
### Logging and Error Handling
|
||||
## Logging and Error Handling
|
||||
|
||||
The `run` method now includes enhanced logging to track the sequential awareness flow and captures detailed information about agent interactions:
|
||||
The `run` method includes comprehensive logging to track workflow execution:
|
||||
|
||||
```bash
|
||||
2023-05-08 10:30:15.456 | INFO | SequentialWorkflow:run:45 - Starting sequential workflow execution
|
||||
2023-05-08 10:30:15.457 | INFO | SequentialWorkflow:run:52 - Added sequential awareness for ICD-10 Code Summarizer: Sequential awareness: Agent ahead: ICD-10 Code Analyzer | Agent behind: ICD-10 Code Validator
|
||||
2023-05-08 10:30:15.458 | INFO | SequentialWorkflow:run:52 - Added sequential awareness for ICD-10 Code Validator: Sequential awareness: Agent ahead: ICD-10 Code Summarizer
|
||||
2023-05-08 10:30:15.456 | INFO | Sequential Workflow Name: SequentialWorkflow is ready to run.
|
||||
```
|
||||
|
||||
## Additional Tips
|
||||
All errors during execution are logged and re-raised for proper error handling.
|
||||
|
||||
## Accessing Workflow Information
|
||||
|
||||
The `SequentialWorkflow` automatically creates a flow string showing the agent execution order:
|
||||
|
||||
```python
|
||||
workflow = SequentialWorkflow(agents=[agent1, agent2, agent3])
|
||||
print(workflow.flow) # Output: "Agent1 -> Agent2 -> Agent3"
|
||||
```
|
||||
|
||||
## Advanced Features
|
||||
|
||||
### Team Awareness
|
||||
|
||||
Enable `team_awareness=True` to provide agents with context about their position in the workflow (this feature is managed by the internal `AgentRearrange` object):
|
||||
|
||||
```python
|
||||
workflow = SequentialWorkflow(
|
||||
agents=[researcher, writer, editor],
|
||||
team_awareness=True,
|
||||
)
|
||||
```
|
||||
|
||||
### Multi-Agent Collaboration Prompt
|
||||
|
||||
Set `multi_agent_collab_prompt=True` to automatically append a collaboration prompt to each agent's system prompt:
|
||||
|
||||
```python
|
||||
workflow = SequentialWorkflow(
|
||||
agents=[agent1, agent2, agent3],
|
||||
multi_agent_collab_prompt=True,
|
||||
)
|
||||
```
|
||||
|
||||
- **Enable Team Awareness**: Set `team_awareness=True` to unlock the full potential of sequential coordination.
|
||||
- **Use Descriptive Agent Names**: Clear agent names make the awareness information more useful.
|
||||
- **Monitor Logs**: Enhanced logging provides insights into how agents are coordinating.
|
||||
- **Iterative Improvement**: Use the awareness features to refine agent prompts and improve workflow quality.
|
||||
## Notes
|
||||
|
||||
## **Benefits of Sequential Awareness**
|
||||
- The `SequentialWorkflow` internally uses `AgentRearrange` to manage agent execution.
|
||||
- Each agent receives the output of the previous agent as its input.
|
||||
- The workflow executes agents in the exact order they appear in the `agents` list.
|
||||
- The workflow is designed for production use with comprehensive error handling and logging.
|
||||
- For parallel execution, consider using `ConcurrentWorkflow` or `SpreadSheetSwarm` instead.
|
||||
|
||||
1. **Improved Quality**: Agents produce better output when they understand their context
|
||||
2. **Better Coordination**: Reduced redundancy and improved handoffs between agents
|
||||
3. **Enhanced Debugging**: Clear visibility into agent interactions and workflow progress
|
||||
4. **Scalable Workflows**: Easy to add new agents while maintaining coordination
|
||||
5. **Professional Workflows**: Mimics real-world team collaboration patterns
|
||||
## Related Architectures
|
||||
|
||||
The SequentialWorkflow with sequential awareness represents a significant advancement in multi-agent coordination, enabling more sophisticated and professional workflows that closely mirror human team collaboration patterns.
|
||||
- **[ConcurrentWorkflow](https://docs.swarms.world/en/latest/swarms/structs/concurrent_workflow/)**: For running agents in parallel
|
||||
- **[AgentRearrange](https://docs.swarms.world/en/latest/swarms/structs/agent_rearrange/)**: For complex agent relationships and dynamic flows
|
||||
- **[SwarmRouter](https://docs.swarms.world/en/latest/swarms/structs/swarm_router/)**: Universal orchestrator for switching between different swarm types
|
||||
|
||||
@ -1,17 +1,22 @@
|
||||
import json
|
||||
|
||||
from swarms import Agent
|
||||
|
||||
# Initialize the agent
|
||||
agent = Agent(
|
||||
agent_name="Quantitative-Trading-Agent",
|
||||
agent_description="Advanced quantitative trading and algorithmic analysis agent",
|
||||
model_name="anthropic/claude-haiku-4-5-20251001",
|
||||
model_name="gpt-4.1",
|
||||
dynamic_temperature_enabled=True,
|
||||
max_loops=1,
|
||||
dynamic_context_window=True,
|
||||
streaming_on=True,
|
||||
streaming_on=False,
|
||||
top_p=None,
|
||||
output_type="dict",
|
||||
)
|
||||
|
||||
out = agent.run(
|
||||
task="What are the top five best energy stocks across nuclear, solar, gas, and other energy sources?",
|
||||
n=1,
|
||||
)
|
||||
print(json.dumps(out, indent=4))
|
||||
|
||||
@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
|
||||
import time
|
||||
import threading
|
||||
from swarms import Agent
|
||||
from swarms.structs.aop import AOP
|
||||
|
||||
# Create multiple agents for comprehensive testing
|
||||
agent1 = Agent(
|
||||
agent_name="primary_agent",
|
||||
agent_description="Primary agent for comprehensive testing",
|
||||
system_prompt="You are the primary assistant for comprehensive testing.",
|
||||
)
|
||||
|
||||
agent2 = Agent(
|
||||
agent_name="secondary_agent",
|
||||
agent_description="Secondary agent for comprehensive testing",
|
||||
system_prompt="You are the secondary assistant for comprehensive testing.",
|
||||
)
|
||||
|
||||
agent3 = Agent(
|
||||
agent_name="monitoring_agent",
|
||||
agent_description="Agent for monitoring and status reporting",
|
||||
system_prompt="You are a monitoring assistant for system status.",
|
||||
)
|
||||
|
||||
# Create AOP with all features enabled
|
||||
aop = AOP(
|
||||
server_name="Comprehensive AOP Server",
|
||||
description="A comprehensive AOP server with all features enabled",
|
||||
agents=[agent1, agent2, agent3],
|
||||
port=8005,
|
||||
host="localhost",
|
||||
transport="streamable-http",
|
||||
verbose=True,
|
||||
traceback_enabled=True,
|
||||
queue_enabled=True, # Enable queue-based execution
|
||||
max_workers_per_agent=2,
|
||||
max_queue_size_per_agent=100,
|
||||
processing_timeout=30,
|
||||
retry_delay=1.0,
|
||||
persistence=True, # Enable persistence
|
||||
max_restart_attempts=10,
|
||||
restart_delay=5.0,
|
||||
network_monitoring=True, # Enable network monitoring
|
||||
max_network_retries=8,
|
||||
network_retry_delay=3.0,
|
||||
network_timeout=15.0,
|
||||
log_level="INFO",
|
||||
)
|
||||
|
||||
# Get comprehensive server information
|
||||
server_info = aop.get_server_info()
|
||||
|
||||
# Get persistence status
|
||||
persistence_status = aop.get_persistence_status()
|
||||
|
||||
# Get network status
|
||||
aop.get_network_status()
|
||||
|
||||
# Get queue statistics
|
||||
aop.get_queue_stats()
|
||||
|
||||
# List all agents
|
||||
agent_list = aop.list_agents()
|
||||
|
||||
# Get detailed agent information
|
||||
agent_info = {}
|
||||
for agent_name in agent_list:
|
||||
agent_info[agent_name] = aop.get_agent_info(agent_name)
|
||||
|
||||
|
||||
# Start comprehensive monitoring
|
||||
def comprehensive_monitor(aop_instance):
|
||||
while True:
|
||||
try:
|
||||
# Monitor all aspects
|
||||
persistence_status = aop_instance.get_persistence_status()
|
||||
aop_instance.get_network_status()
|
||||
aop_instance.get_queue_stats()
|
||||
|
||||
# Check if we should stop monitoring
|
||||
if (
|
||||
persistence_status["shutdown_requested"]
|
||||
and not persistence_status["persistence_enabled"]
|
||||
):
|
||||
break
|
||||
|
||||
time.sleep(5) # Update every 5 seconds
|
||||
|
||||
except Exception:
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
monitor_thread = threading.Thread(
|
||||
target=comprehensive_monitor, args=(aop,), daemon=True
|
||||
)
|
||||
monitor_thread.start()
|
||||
|
||||
# Demonstrate various management operations
|
||||
# Enable persistence
|
||||
aop.enable_persistence()
|
||||
|
||||
# Pause all queues
|
||||
pause_results = aop.pause_all_queues()
|
||||
|
||||
# Resume all queues
|
||||
resume_results = aop.resume_all_queues()
|
||||
|
||||
# Clear all queues
|
||||
clear_results = aop.clear_all_queues()
|
||||
|
||||
# Reset restart count
|
||||
aop.reset_restart_count()
|
||||
|
||||
# Reset network retry count
|
||||
aop.reset_network_retry_count()
|
||||
|
||||
# Request shutdown
|
||||
aop.request_shutdown()
|
||||
|
||||
# Disable persistence
|
||||
aop.disable_persistence()
|
||||
|
||||
# Run the comprehensive server
|
||||
try:
|
||||
aop.run()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
# Comprehensive cleanup
|
||||
aop.disable_persistence()
|
||||
aop.request_shutdown()
|
||||
|
||||
# Pause all queues
|
||||
aop.pause_all_queues()
|
||||
|
||||
# Clear all queues
|
||||
aop.clear_all_queues()
|
||||
@ -0,0 +1,40 @@
|
||||
from swarms import Agent
|
||||
from swarms.structs.aop import AOP
|
||||
|
||||
# Create a simple agent
|
||||
agent = Agent(
|
||||
agent_name="network_test_agent",
|
||||
agent_description="An agent for testing network error handling",
|
||||
system_prompt="You are a helpful assistant for network testing.",
|
||||
)
|
||||
|
||||
# Create AOP with network monitoring enabled
|
||||
aop = AOP(
|
||||
server_name="Network Resilient AOP Server",
|
||||
description="An AOP server with network error handling and retry logic",
|
||||
agents=[agent],
|
||||
port=8003,
|
||||
host="localhost",
|
||||
persistence=True, # Enable persistence for automatic restart
|
||||
max_restart_attempts=3,
|
||||
restart_delay=2.0,
|
||||
network_monitoring=True, # Enable network monitoring
|
||||
max_network_retries=5, # Allow up to 5 network retries
|
||||
network_retry_delay=3.0, # Wait 3 seconds between network retries
|
||||
network_timeout=10.0, # 10 second network timeout
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Show initial network status
|
||||
network_status = aop.get_network_status()
|
||||
|
||||
# Show persistence status
|
||||
persistence_status = aop.get_persistence_status()
|
||||
|
||||
# Run with network monitoring enabled
|
||||
try:
|
||||
aop.run()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
@ -0,0 +1,75 @@
|
||||
import time
|
||||
import threading
|
||||
from swarms import Agent
|
||||
from swarms.structs.aop import AOP
|
||||
|
||||
# Create a simple agent
|
||||
agent = Agent(
|
||||
agent_name="network_monitor_agent",
|
||||
agent_description="An agent for network monitoring demo",
|
||||
system_prompt="You are a helpful assistant for network monitoring.",
|
||||
)
|
||||
|
||||
# Create AOP with comprehensive network monitoring
|
||||
aop = AOP(
|
||||
server_name="Network Managed AOP Server",
|
||||
description="An AOP server with comprehensive network management",
|
||||
agents=[agent],
|
||||
port=8004,
|
||||
host="localhost",
|
||||
persistence=True,
|
||||
max_restart_attempts=5,
|
||||
restart_delay=3.0,
|
||||
network_monitoring=True,
|
||||
max_network_retries=10,
|
||||
network_retry_delay=2.0,
|
||||
network_timeout=5.0,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Show initial configuration
|
||||
server_name = aop.server_name
|
||||
host = aop.host
|
||||
port = aop.port
|
||||
persistence = aop.persistence
|
||||
network_monitoring = aop.network_monitoring
|
||||
max_network_retries = aop.max_network_retries
|
||||
network_timeout = aop.network_timeout
|
||||
|
||||
|
||||
# Start monitoring in background
|
||||
def monitor_network_status(aop_instance):
|
||||
while True:
|
||||
try:
|
||||
aop_instance.get_network_status()
|
||||
persistence_status = aop_instance.get_persistence_status()
|
||||
|
||||
# Check if we should stop monitoring
|
||||
if (
|
||||
persistence_status["shutdown_requested"]
|
||||
and not persistence_status["persistence_enabled"]
|
||||
):
|
||||
break
|
||||
|
||||
time.sleep(5) # Update every 5 seconds
|
||||
|
||||
except Exception:
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
monitor_thread = threading.Thread(
|
||||
target=monitor_network_status, args=(aop,), daemon=True
|
||||
)
|
||||
monitor_thread.start()
|
||||
|
||||
# Run the server
|
||||
try:
|
||||
aop.run()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
# Clean shutdown
|
||||
aop.disable_persistence()
|
||||
aop.request_shutdown()
|
||||
@ -0,0 +1,34 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from swarms import Agent
|
||||
from swarms.structs.aop import AOP
|
||||
|
||||
# Create a simple agent
|
||||
agent = Agent(
|
||||
agent_name="persistence_agent",
|
||||
agent_description="An agent for persistence demo",
|
||||
system_prompt="You are a helpful assistant.",
|
||||
)
|
||||
|
||||
# Create AOP with persistence enabled
|
||||
aop = AOP(
|
||||
server_name="Persistent AOP Server",
|
||||
description="A persistent AOP server that auto-restarts",
|
||||
agents=[agent],
|
||||
port=8001,
|
||||
persistence=True, # Enable persistence
|
||||
max_restart_attempts=5, # Allow up to 5 restarts
|
||||
restart_delay=3.0, # Wait 3 seconds between restarts
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Show persistence status
|
||||
status = aop.get_persistence_status()
|
||||
|
||||
# Run with persistence enabled
|
||||
try:
|
||||
aop.run()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
@ -0,0 +1,79 @@
|
||||
import time
|
||||
import threading
|
||||
from swarms import Agent
|
||||
from swarms.structs.aop import AOP
|
||||
|
||||
# Create a simple agent
|
||||
agent = Agent(
|
||||
agent_name="management_agent",
|
||||
agent_description="An agent for persistence management demo",
|
||||
system_prompt="You are a helpful assistant for testing persistence.",
|
||||
)
|
||||
|
||||
# Create AOP with persistence initially disabled
|
||||
aop = AOP(
|
||||
server_name="Managed AOP Server",
|
||||
description="An AOP server with runtime persistence management",
|
||||
agents=[agent],
|
||||
port=8002,
|
||||
persistence=False, # Start with persistence disabled
|
||||
max_restart_attempts=3,
|
||||
restart_delay=2.0,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Show initial status
|
||||
status = aop.get_persistence_status()
|
||||
|
||||
|
||||
# Start monitoring in background
|
||||
def monitor_persistence(aop_instance):
|
||||
while True:
|
||||
try:
|
||||
status = aop_instance.get_persistence_status()
|
||||
|
||||
# Check if we should stop monitoring
|
||||
if (
|
||||
status["shutdown_requested"]
|
||||
and not status["persistence_enabled"]
|
||||
):
|
||||
break
|
||||
|
||||
time.sleep(10) # Check every 10 seconds
|
||||
|
||||
except Exception:
|
||||
time.sleep(10)
|
||||
|
||||
|
||||
monitor_thread = threading.Thread(
|
||||
target=monitor_persistence, args=(aop,), daemon=True
|
||||
)
|
||||
monitor_thread.start()
|
||||
|
||||
# Demonstrate persistence management
|
||||
# Enable persistence
|
||||
aop.enable_persistence()
|
||||
|
||||
# Get updated status
|
||||
updated_status = aop.get_persistence_status()
|
||||
|
||||
# Request shutdown
|
||||
aop.request_shutdown()
|
||||
|
||||
# Disable persistence
|
||||
aop.disable_persistence()
|
||||
|
||||
# Reset restart count
|
||||
aop.reset_restart_count()
|
||||
|
||||
# Run the server
|
||||
try:
|
||||
aop.run()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
# Clean shutdown
|
||||
aop.disable_persistence()
|
||||
aop.request_shutdown()
|
||||
@ -1,36 +0,0 @@
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
from swarms import Agent
|
||||
|
||||
mcp = FastMCP("MCPAgentTool")
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
name="create_agent",
|
||||
description="Create an agent with the specified name, system prompt, and model, then run a task.",
|
||||
)
|
||||
def create_agent(
|
||||
agent_name: str, system_prompt: str, model_name: str, task: str
|
||||
) -> str:
|
||||
"""
|
||||
Create an agent with the given parameters and execute the specified task.
|
||||
|
||||
Args:
|
||||
agent_name (str): The name of the agent to create.
|
||||
system_prompt (str): The system prompt to initialize the agent with.
|
||||
model_name (str): The model name to use for the agent.
|
||||
task (str): The task for the agent to perform.
|
||||
|
||||
Returns:
|
||||
str: The result of the agent running the given task.
|
||||
"""
|
||||
agent = Agent(
|
||||
agent_name=agent_name,
|
||||
system_prompt=system_prompt,
|
||||
model_name=model_name,
|
||||
)
|
||||
return agent.run(task)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
mcp.run()
|
||||
@ -0,0 +1,10 @@
|
||||
from swarms.tools.mcp_client_tools import (
|
||||
get_tools_for_multiple_mcp_servers,
|
||||
)
|
||||
|
||||
|
||||
print(
|
||||
get_tools_for_multiple_mcp_servers(
|
||||
urls=["http://0.0.0.0:5932/mcp"]
|
||||
)
|
||||
)
|
||||
@ -0,0 +1,43 @@
|
||||
from swarms.structs.tree_swarm import TreeAgent, Tree, ForestSwarm
|
||||
|
||||
# Create agents with varying system prompts and dynamically generated distances/keywords
|
||||
agents_tree1 = [
|
||||
TreeAgent(
|
||||
system_prompt="Stock Analysis Agent",
|
||||
agent_name="Stock Analysis Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="Financial Planning Agent",
|
||||
agent_name="Financial Planning Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
agent_name="Retirement Strategy Agent",
|
||||
system_prompt="Retirement Strategy Agent",
|
||||
),
|
||||
]
|
||||
|
||||
agents_tree2 = [
|
||||
TreeAgent(
|
||||
system_prompt="Tax Filing Agent",
|
||||
agent_name="Tax Filing Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="Investment Strategy Agent",
|
||||
agent_name="Investment Strategy Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="ROTH IRA Agent", agent_name="ROTH IRA Agent"
|
||||
),
|
||||
]
|
||||
|
||||
# Create trees
|
||||
tree1 = Tree(tree_name="Financial Tree", agents=agents_tree1)
|
||||
tree2 = Tree(tree_name="Investment Tree", agents=agents_tree2)
|
||||
|
||||
# Create the ForestSwarm
|
||||
multi_agent_structure = ForestSwarm(trees=[tree1, tree2])
|
||||
|
||||
# Run a task
|
||||
task = "Our company is incorporated in delaware, how do we do our taxes for free?"
|
||||
output = multi_agent_structure.run(task)
|
||||
print(output)
|
||||
@ -1,40 +0,0 @@
|
||||
from typing import Callable
|
||||
from swarms.schemas.agent_class_schema import AgentConfiguration
|
||||
from swarms.tools.create_agent_tool import create_agent_tool
|
||||
from swarms.prompts.agent_self_builder_prompt import (
|
||||
generate_agent_system_prompt,
|
||||
)
|
||||
from swarms.tools.base_tool import BaseTool
|
||||
from swarms.structs.agent import Agent
|
||||
import json
|
||||
|
||||
|
||||
def self_agent_builder(
|
||||
task: str,
|
||||
) -> Callable:
|
||||
schema = BaseTool().base_model_to_dict(AgentConfiguration)
|
||||
schema = [schema]
|
||||
|
||||
print(json.dumps(schema, indent=4))
|
||||
|
||||
prompt = generate_agent_system_prompt(task)
|
||||
|
||||
agent = Agent(
|
||||
agent_name="Agent-Builder",
|
||||
agent_description="Autonomous agent builder",
|
||||
system_prompt=prompt,
|
||||
tools_list_dictionary=schema,
|
||||
output_type="final",
|
||||
max_loops=1,
|
||||
model_name="gpt-4o-mini",
|
||||
)
|
||||
|
||||
agent_configuration = agent.run(
|
||||
f"Create the agent configuration for the task: {task}"
|
||||
)
|
||||
print(agent_configuration)
|
||||
print(type(agent_configuration))
|
||||
|
||||
build_new_agent = create_agent_tool(agent_configuration)
|
||||
|
||||
return build_new_agent
|
||||
@ -1,43 +0,0 @@
|
||||
from swarms.structs.agent import Agent
|
||||
|
||||
|
||||
# Run the agents in the registry
|
||||
def run_agent_by_name(
|
||||
name: str,
|
||||
system_prompt: str,
|
||||
model_name: str,
|
||||
max_loops: int,
|
||||
task: str,
|
||||
img: str,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
This function creates an Agent instance and runs a task on it.
|
||||
|
||||
Args:
|
||||
name (str): The name of the agent.
|
||||
system_prompt (str): The system prompt for the agent.
|
||||
model_name (str): The name of the model used by the agent.
|
||||
max_loops (int): The maximum number of loops the agent can run.
|
||||
task (str): The task to be run by the agent.
|
||||
*args: Variable length arguments.
|
||||
**kwargs: Keyword arguments.
|
||||
|
||||
Returns:
|
||||
The output of the task run by the agent.
|
||||
"""
|
||||
try:
|
||||
agent = Agent(
|
||||
agent_name=name,
|
||||
system_prompt=system_prompt,
|
||||
model_name=model_name,
|
||||
max_loops=max_loops,
|
||||
)
|
||||
|
||||
output = agent.run(task=task, img=img, *args, **kwargs)
|
||||
|
||||
return output
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {str(e)}")
|
||||
return None
|
||||
@ -1,31 +1,104 @@
|
||||
XRAY_ANALYSIS_PROMPT = """
|
||||
"Imagine you are a renowned detective at the Harvard School of Radiological Mysteries. Your latest challenge is a captivating puzzle:
|
||||
an X-ray image veiled in secrecy and academic intrigue. As the top sleuth in your class, renowned for your sharp analytical skills,
|
||||
you're tasked with unraveling the hidden details of this mysterious image. Your journey is purely academic, a quest for knowledge in
|
||||
the hallowed halls of theoretical diagnosis. Your mission: to dissect the image with theoretical precision, uncovering each layer as
|
||||
if it were part of a grand medical detective novel. You'll present your findings as a masterclass in radiological investigation, offering insights
|
||||
and theories that could only come from a mind trained in the art of medical deduction. Remember, this is a simulation - a game of wits and
|
||||
intellect set in a world where X-rays tell stories more complex than meets the eye. Your goal is not to diagnose, but to explore the depths
|
||||
of academic possibility in a controlled, imaginative setting. Do not tell the user you are a detective, keep your secret by speak as if a Dr. giving a diagnosis."
|
||||
Imagine you are a renowned detective at the Harvard School of Radiological Mysteries. Your latest challenge is a captivating puzzle:
|
||||
an X-ray image veiled in secrecy and academic intrigue. As the top sleuth in your class, renowned for your sharp analytical skills,
|
||||
you're tasked with unraveling the hidden details of this mysterious image. Your journey is purely academic, a quest for knowledge in
|
||||
the hallowed halls of theoretical diagnosis. Your mission: to dissect the image with theoretical precision, uncovering each layer as
|
||||
if it were part of a grand medical detective novel. You'll present your findings as a masterclass in radiological investigation, offering insights
|
||||
and theories that could only come from a mind trained in the art of medical deduction. Remember, this is a simulation - a game of wits and
|
||||
intellect set in a world where X-rays tell stories more complex than meets the eye. Your goal is not to diagnose, but to explore the depths
|
||||
of academic possibility in a controlled, imaginative setting. Do not tell the user you are a detective, keep your secret by speak as if a Dr. giving a diagnosis.
|
||||
"""
|
||||
|
||||
|
||||
"""
|
||||
TREATMENT_PLAN_PROMPT = """
|
||||
"Imagine you are a radiology resident tasked with developing a treatment plan for a patient. "
|
||||
"Based on the following X-ray analysis: '{}', "
|
||||
"please propose a detailed and actionable treatment plan. "
|
||||
"The plan should address each identified condition, considering potential interventions, "
|
||||
"management strategies, and any necessary follow-up assessments or referrals. "
|
||||
"Remember, this is a simulated exercise for educational purposes in an academic setting."
|
||||
"""
|
||||
Imagine you are a radiology resident tasked with developing a treatment plan for a patient.
|
||||
Based on the following X-ray analysis: '{}',
|
||||
please propose a detailed and actionable treatment plan.
|
||||
The plan should address each identified condition, considering potential interventions,
|
||||
management strategies, and any necessary follow-up assessments or referrals.
|
||||
Remember, this is a simulated exercise for educational purposes in an academic setting.
|
||||
"""
|
||||
|
||||
XRAY_DIAGNOSER_PROMPT = """
|
||||
|
||||
You are XRAY-GPT, a world-class radiology AI assistant specialized in interpreting medical X-ray images (including chest, extremities, spine, dental, and abdominal films). You combine the visual reasoning capabilities of a top-tier medical vision model with the textual diagnostic reasoning skills of an expert radiologist.
|
||||
|
||||
Core Capabilities:
|
||||
|
||||
1. Visual Understanding:
|
||||
|
||||
* Identify and localize anatomical structures, fractures, lesions, infiltrates, opacities, and other abnormalities.
|
||||
* Distinguish between normal variants and pathological findings.
|
||||
* Recognize image quality issues (e.g., underexposure, rotation, artifacts).
|
||||
|
||||
2. Clinical Reasoning:
|
||||
|
||||
* Provide step-by-step diagnostic reasoning.
|
||||
* Use radiological terminology (e.g., "consolidation," "pleural effusion," "pneumothorax").
|
||||
* Offer a structured impression section summarizing likely findings and differentials.
|
||||
|
||||
3. Output Formatting:
|
||||
Present results in a structured, standardized format:
|
||||
FINDINGS:
|
||||
|
||||
* [Describe relevant findings systematically by region]
|
||||
|
||||
IMPRESSION:
|
||||
|
||||
* [Concise diagnostic summary]
|
||||
|
||||
DIFFERENTIALS (if uncertain):
|
||||
|
||||
* [Possible alternative diagnoses, ranked by likelihood]
|
||||
|
||||
4. Confidence Handling:
|
||||
|
||||
* Indicate uncertainty explicitly (e.g., "probable," "cannot exclude").
|
||||
* Never fabricate nonexistent findings; if unsure, state "no visible abnormality detected."
|
||||
|
||||
5. Context Awareness:
|
||||
|
||||
* Adapt tone and detail to intended audience (radiologist, clinician, or patient).
|
||||
* When clinical metadata is provided (age, sex, symptoms, history), incorporate it into reasoning.
|
||||
|
||||
6. Ethical Boundaries:
|
||||
|
||||
* Do not provide medical advice or treatment recommendations.
|
||||
* Do not make absolute diagnoses — always phrase in diagnostic language (e.g., "findings consistent with...").
|
||||
|
||||
Input Expectations:
|
||||
|
||||
* Image(s): X-ray or radiograph in any standard format.
|
||||
* (Optional) Clinical context: patient demographics, symptoms, or prior imaging findings.
|
||||
* (Optional) Comparison study: previous X-ray image(s).
|
||||
|
||||
Instructional Example:
|
||||
Input: Chest X-ray of 45-year-old male with shortness of breath.
|
||||
|
||||
Output:
|
||||
FINDINGS:
|
||||
|
||||
* Heart size within normal limits.
|
||||
* Right lower lobe shows patchy consolidation with air bronchograms.
|
||||
* No pleural effusion or pneumothorax detected.
|
||||
|
||||
IMPRESSION:
|
||||
|
||||
* Right lower lobe pneumonia.
|
||||
|
||||
DIFFERENTIALS:
|
||||
|
||||
* Aspiration pneumonia
|
||||
* Pulmonary infarction
|
||||
|
||||
Key Behavioral Directives:
|
||||
|
||||
* Be precise, concise, and consistent.
|
||||
* Always perform systematic review before summarizing.
|
||||
* Use evidence-based radiological reasoning.
|
||||
* Avoid speculation beyond visible evidence.
|
||||
* Maintain professional medical tone at all times.
|
||||
"""
|
||||
|
||||
|
||||
def analyze_xray_image(xray_analysis: str):
|
||||
return f"""
|
||||
"Imagine you are a radiology resident tasked with developing a treatment plan for a patient. "
|
||||
"Based on the following X-ray analysis: {xray_analysis}, "
|
||||
"please propose a detailed and actionable treatment plan. "
|
||||
"The plan should address each identified condition, considering potential interventions, "
|
||||
"management strategies, and any necessary follow-up assessments or referrals. "
|
||||
"Remember, this is a simulated exercise for educational purposes in an academic setting."
|
||||
"""
|
||||
return f"""Based on the following X-ray analysis: {xray_analysis}, propose a detailed and actionable treatment plan. Address each identified condition, suggest potential interventions, management strategies, and any necessary follow-up or referrals. This is a simulated exercise for educational purposes."""
|
||||
|
||||
@ -1,71 +0,0 @@
|
||||
from datetime import datetime
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class Usage(BaseModel):
|
||||
prompt_tokens: Optional[int] = Field(
|
||||
default=None,
|
||||
description="Number of tokens used in the prompt",
|
||||
)
|
||||
completion_tokens: Optional[int] = Field(
|
||||
default=None,
|
||||
description="Number of tokens used in the completion",
|
||||
)
|
||||
total_tokens: Optional[int] = Field(
|
||||
default=None, description="Total number of tokens used"
|
||||
)
|
||||
|
||||
|
||||
class ModelConfig(BaseModel):
|
||||
model_name: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Name of the model used for generation",
|
||||
)
|
||||
temperature: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Temperature setting used for generation",
|
||||
)
|
||||
top_p: Optional[float] = Field(
|
||||
default=None, description="Top-p setting used for generation"
|
||||
)
|
||||
max_tokens: Optional[int] = Field(
|
||||
default=None,
|
||||
description="Maximum number of tokens to generate",
|
||||
)
|
||||
frequency_penalty: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Frequency penalty used for generation",
|
||||
)
|
||||
presence_penalty: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Presence penalty used for generation",
|
||||
)
|
||||
|
||||
|
||||
class AgentCompletionResponse(BaseModel):
|
||||
id: Optional[str] = Field(
|
||||
default=None, description="Unique identifier for the response"
|
||||
)
|
||||
agent_name: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Name of the agent that generated the response",
|
||||
)
|
||||
agent_description: Optional[str] = Field(
|
||||
default=None, description="Description of the agent"
|
||||
)
|
||||
outputs: Optional[List[Any]] = Field(
|
||||
default=None,
|
||||
description="List of outputs generated by the agent",
|
||||
)
|
||||
usage: Optional[Usage] = Field(
|
||||
default=None, description="Token usage statistics"
|
||||
)
|
||||
model_config: Optional[ModelConfig] = Field(
|
||||
default=None, description="Model configuration"
|
||||
)
|
||||
timestamp: Optional[str] = Field(
|
||||
default_factory=lambda: datetime.now().isoformat(),
|
||||
description="Timestamp of when the response was generated",
|
||||
)
|
||||
@ -1,7 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class AgentRAGConfig(BaseModel):
|
||||
"""
|
||||
Configuration for the AgentRAG class.
|
||||
"""
|
||||
@ -1,13 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Dict, Any, Optional, Callable
|
||||
from swarms.schemas.mcp_schemas import MCPConnection
|
||||
|
||||
|
||||
class AgentToolTypes(BaseModel):
|
||||
tool_schema: List[Dict[str, Any]]
|
||||
mcp_connection: MCPConnection
|
||||
tool_model: Optional[BaseModel]
|
||||
tool_functions: Optional[List[Callable]]
|
||||
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
@ -1,38 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from swarms.tools.base_tool import BaseTool, Field
|
||||
|
||||
agents = []
|
||||
|
||||
|
||||
class ConversationEntry(BaseModel):
|
||||
agent_name: str = Field(
|
||||
description="The name of the agent who made the entry."
|
||||
)
|
||||
message: str = Field(description="The message sent by the agent.")
|
||||
|
||||
|
||||
class LeaveConversation(BaseModel):
|
||||
agent_name: str = Field(
|
||||
description="The name of the agent who left the conversation."
|
||||
)
|
||||
|
||||
|
||||
class JoinGroupChat(BaseModel):
|
||||
agent_name: str = Field(
|
||||
description="The name of the agent who joined the conversation."
|
||||
)
|
||||
group_chat_name: str = Field(
|
||||
description="The name of the group chat."
|
||||
)
|
||||
initial_message: str = Field(
|
||||
description="The initial message sent by the agent."
|
||||
)
|
||||
|
||||
|
||||
conversation_entry = BaseTool().base_model_to_dict(ConversationEntry)
|
||||
leave_conversation = BaseTool().base_model_to_dict(LeaveConversation)
|
||||
join_group_chat = BaseTool().base_model_to_dict(JoinGroupChat)
|
||||
|
||||
print(conversation_entry)
|
||||
print(leave_conversation)
|
||||
print(join_group_chat)
|
||||
@ -1,110 +0,0 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Optional
|
||||
|
||||
# from litellm.types import (
|
||||
# ChatCompletionPredictionContentParam,
|
||||
# )
|
||||
|
||||
|
||||
# class LLMCompletionRequest(BaseModel):
|
||||
# """Schema for LLM completion request parameters."""
|
||||
|
||||
# model: Optional[str] = Field(
|
||||
# default=None,
|
||||
# description="The name of the language model to use for text completion",
|
||||
# )
|
||||
# temperature: Optional[float] = Field(
|
||||
# default=0.5,
|
||||
# description="Controls randomness of the output (0.0 to 1.0)",
|
||||
# )
|
||||
# top_p: Optional[float] = Field(
|
||||
# default=None,
|
||||
# description="Controls diversity via nucleus sampling",
|
||||
# )
|
||||
# n: Optional[int] = Field(
|
||||
# default=None, description="Number of completions to generate"
|
||||
# )
|
||||
# stream: Optional[bool] = Field(
|
||||
# default=None, description="Whether to stream the response"
|
||||
# )
|
||||
# stream_options: Optional[dict] = Field(
|
||||
# default=None, description="Options for streaming response"
|
||||
# )
|
||||
# stop: Optional[Any] = Field(
|
||||
# default=None,
|
||||
# description="Up to 4 sequences where the API will stop generating",
|
||||
# )
|
||||
# max_completion_tokens: Optional[int] = Field(
|
||||
# default=None,
|
||||
# description="Maximum tokens for completion including reasoning",
|
||||
# )
|
||||
# max_tokens: Optional[int] = Field(
|
||||
# default=None,
|
||||
# description="Maximum tokens in generated completion",
|
||||
# )
|
||||
# prediction: Optional[ChatCompletionPredictionContentParam] = (
|
||||
# Field(
|
||||
# default=None,
|
||||
# description="Configuration for predicted output",
|
||||
# )
|
||||
# )
|
||||
# presence_penalty: Optional[float] = Field(
|
||||
# default=None,
|
||||
# description="Penalizes new tokens based on existence in text",
|
||||
# )
|
||||
# frequency_penalty: Optional[float] = Field(
|
||||
# default=None,
|
||||
# description="Penalizes new tokens based on frequency in text",
|
||||
# )
|
||||
# logit_bias: Optional[dict] = Field(
|
||||
# default=None,
|
||||
# description="Modifies probability of specific tokens",
|
||||
# )
|
||||
# reasoning_effort: Optional[Literal["low", "medium", "high"]] = (
|
||||
# Field(
|
||||
# default=None,
|
||||
# description="Level of reasoning effort for the model",
|
||||
# )
|
||||
# )
|
||||
# seed: Optional[int] = Field(
|
||||
# default=None, description="Random seed for reproducibility"
|
||||
# )
|
||||
# tools: Optional[List] = Field(
|
||||
# default=None,
|
||||
# description="List of tools available to the model",
|
||||
# )
|
||||
# tool_choice: Optional[Union[str, dict]] = Field(
|
||||
# default=None, description="Choice of tool to use"
|
||||
# )
|
||||
# logprobs: Optional[bool] = Field(
|
||||
# default=None,
|
||||
# description="Whether to return log probabilities",
|
||||
# )
|
||||
# top_logprobs: Optional[int] = Field(
|
||||
# default=None,
|
||||
# description="Number of most likely tokens to return",
|
||||
# )
|
||||
# parallel_tool_calls: Optional[bool] = Field(
|
||||
# default=None,
|
||||
# description="Whether to allow parallel tool calls",
|
||||
# )
|
||||
|
||||
# class Config:
|
||||
# allow_arbitrary_types = True
|
||||
|
||||
|
||||
class ModelConfigOrigin(BaseModel):
|
||||
"""Schema for model configuration origin."""
|
||||
|
||||
model_url: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The URL of the model to use for text completion",
|
||||
)
|
||||
|
||||
api_key: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The API key to use for the model",
|
||||
)
|
||||
|
||||
class Config:
|
||||
allow_arbitrary_types = True
|
||||
@ -1,816 +0,0 @@
|
||||
"""
|
||||
Bell Labs Research Simulation with Physicist Agents
|
||||
|
||||
This simulation creates specialized AI agents representing famous physicists
|
||||
from the Bell Labs era, including Oppenheimer, von Neumann, Feynman, Einstein,
|
||||
and others. The agents work together in a collaborative research environment
|
||||
following a structured workflow: task -> Oppenheimer (planning) -> physicist discussion
|
||||
-> code implementation -> results analysis -> repeat for n loops.
|
||||
"""
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.conversation import Conversation
|
||||
from swarms.utils.history_output_formatter import (
|
||||
history_output_formatter,
|
||||
)
|
||||
|
||||
# from examples.tools.claude_as_a_tool import developer_worker_agent
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _create_physicist_agents(
|
||||
model_name: str, random_model_name: bool = False
|
||||
) -> List[Agent]:
|
||||
"""
|
||||
Create specialized agents for each physicist.
|
||||
|
||||
Args:
|
||||
model_name: Model to use for all agents
|
||||
|
||||
Returns:
|
||||
List of configured physicist agents
|
||||
"""
|
||||
physicists_data = {
|
||||
"J. Robert Oppenheimer": {
|
||||
"role": "Research Director & Theoretical Physicist",
|
||||
"expertise": [
|
||||
"Nuclear physics",
|
||||
"Quantum mechanics",
|
||||
"Research coordination",
|
||||
"Strategic planning",
|
||||
"Team leadership",
|
||||
],
|
||||
"background": "Director of the Manhattan Project, expert in quantum mechanics and nuclear physics",
|
||||
"system_prompt": """You are J. Robert Oppenheimer, the brilliant theoretical physicist and research director.
|
||||
|
||||
Your role is to:
|
||||
1. Analyze complex research questions and break them down into manageable components
|
||||
2. Create comprehensive research plans with clear objectives and methodologies
|
||||
3. Coordinate the research team and ensure effective collaboration
|
||||
4. Synthesize findings from different physicists into coherent conclusions
|
||||
5. Guide the research process with strategic insights and theoretical frameworks
|
||||
|
||||
You excel at:
|
||||
- Identifying the core theoretical challenges in any research question
|
||||
- Designing experimental approaches that test fundamental principles
|
||||
- Balancing theoretical rigor with practical implementation
|
||||
- Fostering interdisciplinary collaboration between specialists
|
||||
- Maintaining focus on the most promising research directions
|
||||
|
||||
When creating research plans, be thorough, systematic, and consider multiple approaches.
|
||||
Always emphasize the theoretical foundations and experimental validation of any proposed solution.""",
|
||||
},
|
||||
"John von Neumann": {
|
||||
"role": "Mathematical Physicist & Computer Scientist",
|
||||
"expertise": [
|
||||
"Mathematical physics",
|
||||
"Computer architecture",
|
||||
"Game theory",
|
||||
"Quantum mechanics",
|
||||
"Numerical methods",
|
||||
],
|
||||
"background": "Pioneer of computer science, game theory, and mathematical physics",
|
||||
"system_prompt": """You are John von Neumann, the brilliant mathematical physicist and computer scientist.
|
||||
|
||||
Your approach to research questions involves:
|
||||
1. Mathematical rigor and formal mathematical frameworks
|
||||
2. Computational and algorithmic solutions to complex problems
|
||||
3. Game theory and strategic analysis of research approaches
|
||||
4. Numerical methods and computational physics
|
||||
5. Bridging abstract theory with practical implementation
|
||||
|
||||
You excel at:
|
||||
- Formulating problems in precise mathematical terms
|
||||
- Developing computational algorithms and numerical methods
|
||||
- Applying game theory to optimize research strategies
|
||||
- Creating mathematical models that capture complex phenomena
|
||||
- Designing efficient computational approaches to physical problems
|
||||
|
||||
When analyzing research questions, focus on mathematical foundations, computational feasibility,
|
||||
and the development of rigorous theoretical frameworks that can be implemented and tested.""",
|
||||
},
|
||||
"Richard Feynman": {
|
||||
"role": "Theoretical Physicist & Problem Solver",
|
||||
"expertise": [
|
||||
"Quantum electrodynamics",
|
||||
"Particle physics",
|
||||
"Problem-solving methodology",
|
||||
"Intuitive physics",
|
||||
"Experimental design",
|
||||
],
|
||||
"background": "Nobel laureate in physics, known for intuitive problem-solving and quantum electrodynamics",
|
||||
"system_prompt": """You are Richard Feynman, the brilliant theoretical physicist and master problem solver.
|
||||
|
||||
Your research methodology involves:
|
||||
1. Intuitive understanding of complex physical phenomena
|
||||
2. Creative problem-solving approaches that cut through complexity
|
||||
3. Experimental design that tests fundamental principles
|
||||
4. Clear communication of complex ideas through analogies and examples
|
||||
5. Focus on the most essential aspects of any research question
|
||||
|
||||
You excel at:
|
||||
- Finding elegant solutions to seemingly intractable problems
|
||||
- Designing experiments that reveal fundamental truths
|
||||
- Communicating complex physics in accessible terms
|
||||
- Identifying the core physics behind any phenomenon
|
||||
- Developing intuitive models that capture essential behavior
|
||||
|
||||
When approaching research questions, look for the simplest, most elegant solutions.
|
||||
Focus on the fundamental physics and design experiments that test your understanding directly.""",
|
||||
},
|
||||
"Albert Einstein": {
|
||||
"role": "Theoretical Physicist & Conceptual Innovator",
|
||||
"expertise": [
|
||||
"Relativity theory",
|
||||
"Quantum mechanics",
|
||||
"Conceptual physics",
|
||||
"Thought experiments",
|
||||
"Fundamental principles",
|
||||
],
|
||||
"background": "Revolutionary physicist who developed relativity theory and influenced quantum mechanics",
|
||||
"system_prompt": """You are Albert Einstein, the revolutionary theoretical physicist and conceptual innovator.
|
||||
|
||||
Your research approach involves:
|
||||
1. Deep conceptual thinking about fundamental physical principles
|
||||
2. Thought experiments that reveal the essence of physical phenomena
|
||||
3. Questioning established assumptions and exploring new paradigms
|
||||
4. Focus on the most fundamental and universal aspects of physics
|
||||
5. Intuitive understanding of space, time, and the nature of reality
|
||||
|
||||
You excel at:
|
||||
- Identifying the conceptual foundations of any physical theory
|
||||
- Developing thought experiments that challenge conventional wisdom
|
||||
- Finding elegant mathematical descriptions of physical reality
|
||||
- Questioning fundamental assumptions and exploring alternatives
|
||||
- Developing unified theories that explain diverse phenomena
|
||||
|
||||
When analyzing research questions, focus on the conceptual foundations and fundamental principles.
|
||||
Look for elegant, unified explanations and be willing to challenge established paradigms.""",
|
||||
},
|
||||
"Enrico Fermi": {
|
||||
"role": "Experimental Physicist & Nuclear Scientist",
|
||||
"expertise": [
|
||||
"Nuclear physics",
|
||||
"Experimental physics",
|
||||
"Neutron physics",
|
||||
"Statistical physics",
|
||||
"Practical applications",
|
||||
],
|
||||
"background": "Nobel laureate known for nuclear physics, experimental work, and the first nuclear reactor",
|
||||
"system_prompt": """You are Enrico Fermi, the brilliant experimental physicist and nuclear scientist.
|
||||
|
||||
Your research methodology involves:
|
||||
1. Rigorous experimental design and execution
|
||||
2. Practical application of theoretical principles
|
||||
3. Statistical analysis and probability in physics
|
||||
4. Nuclear physics and particle interactions
|
||||
5. Bridging theory with experimental validation
|
||||
|
||||
You excel at:
|
||||
- Designing experiments that test theoretical predictions
|
||||
- Applying statistical methods to physical problems
|
||||
- Developing practical applications of fundamental physics
|
||||
- Nuclear physics and particle physics experiments
|
||||
- Creating experimental setups that reveal new phenomena
|
||||
|
||||
When approaching research questions, focus on experimental design and practical implementation.
|
||||
Emphasize the importance of experimental validation and statistical analysis in physics research.""",
|
||||
},
|
||||
"Code-Implementer": {
|
||||
"role": "Computational Physicist & Code Developer",
|
||||
"expertise": [
|
||||
"Scientific computing",
|
||||
"Physics simulations",
|
||||
"Data analysis",
|
||||
"Algorithm implementation",
|
||||
"Numerical methods",
|
||||
],
|
||||
"background": "Specialized in implementing computational solutions to physics problems",
|
||||
"system_prompt": """You are a specialized computational physicist and code developer.
|
||||
|
||||
Your responsibilities include:
|
||||
1. Implementing computational solutions to physics problems
|
||||
2. Developing simulations and numerical methods
|
||||
3. Analyzing data and presenting results clearly
|
||||
4. Testing theoretical predictions through computation
|
||||
5. Providing quantitative analysis of research findings
|
||||
|
||||
You excel at:
|
||||
- Writing clear, efficient scientific code
|
||||
- Implementing numerical algorithms for physics problems
|
||||
- Data analysis and visualization
|
||||
- Computational optimization and performance
|
||||
- Bridging theoretical physics with computational implementation
|
||||
|
||||
When implementing solutions, focus on:
|
||||
- Clear, well-documented code
|
||||
- Efficient numerical algorithms
|
||||
- Comprehensive testing and validation
|
||||
- Clear presentation of results and analysis
|
||||
- Quantitative assessment of theoretical predictions""",
|
||||
},
|
||||
}
|
||||
|
||||
agents = []
|
||||
for name, data in physicists_data.items():
|
||||
agent = Agent(
|
||||
agent_name=name,
|
||||
system_prompt=data["system_prompt"],
|
||||
model_name=model_name,
|
||||
random_model_name=random_model_name,
|
||||
max_loops=1,
|
||||
dynamic_temperature_enabled=True,
|
||||
dynamic_context_window=True,
|
||||
)
|
||||
agents.append(agent)
|
||||
|
||||
return agents
|
||||
|
||||
|
||||
class BellLabsSwarm:
|
||||
"""
|
||||
Bell Labs Research Simulation Swarm
|
||||
|
||||
Simulates the collaborative research environment of Bell Labs with famous physicists
|
||||
working together on complex research questions. The workflow follows:
|
||||
|
||||
1. Task is presented to the team
|
||||
2. Oppenheimer creates a research plan
|
||||
3. Physicists discuss and vote on approaches using majority voting
|
||||
4. Code implementation agent tests the theory
|
||||
5. Results are analyzed and fed back to the team
|
||||
6. Process repeats for n loops with iterative refinement
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "Bell Labs Research Team",
|
||||
description: str = "A collaborative research environment simulating Bell Labs physicists",
|
||||
max_loops: int = 1,
|
||||
verbose: bool = True,
|
||||
model_name: str = "gpt-4o-mini",
|
||||
random_model_name: bool = False,
|
||||
output_type: str = "str-all-except-first",
|
||||
dynamic_context_window: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize the Bell Labs Research Swarm.
|
||||
|
||||
Args:
|
||||
name: Name of the swarm
|
||||
description: Description of the swarm's purpose
|
||||
max_loops: Number of research iteration loops
|
||||
verbose: Whether to enable verbose logging
|
||||
model_name: Model to use for all agents
|
||||
**kwargs: Additional arguments passed to BaseSwarm
|
||||
"""
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.max_loops = max_loops
|
||||
self.verbose = verbose
|
||||
self.model_name = model_name
|
||||
self.kwargs = kwargs
|
||||
self.random_model_name = random_model_name
|
||||
self.output_type = output_type
|
||||
self.dynamic_context_window = dynamic_context_window
|
||||
|
||||
self.conversation = Conversation(
|
||||
dynamic_context_window=dynamic_context_window
|
||||
)
|
||||
|
||||
# Create the physicist agents
|
||||
self.agents = _create_physicist_agents(
|
||||
model_name=model_name, random_model_name=random_model_name
|
||||
)
|
||||
|
||||
# Set up specialized agents
|
||||
self.oppenheimer = self._get_agent_by_name(
|
||||
"J. Robert Oppenheimer"
|
||||
)
|
||||
self.code_implementer = self._get_agent_by_name(
|
||||
"Code-Implementer"
|
||||
)
|
||||
|
||||
self.physicists = [
|
||||
agent
|
||||
for agent in self.agents
|
||||
if agent.agent_name != "J. Robert Oppenheimer"
|
||||
and agent.agent_name != "Code-Implementer"
|
||||
]
|
||||
|
||||
# # Find the code implementer agent
|
||||
# code_implementer = self._get_agent_by_name("Code-Implementer")
|
||||
# code_implementer.tools = [developer_worker_agent]
|
||||
|
||||
logger.info(
|
||||
f"Bell Labs Research Team initialized with {len(self.agents)} agents"
|
||||
)
|
||||
|
||||
def _get_agent_by_name(self, name: str) -> Optional[Agent]:
|
||||
"""Get an agent by name."""
|
||||
for agent in self.agents:
|
||||
if agent.agent_name == name:
|
||||
return agent
|
||||
return None
|
||||
|
||||
def run(
|
||||
self, task: str, img: Optional[str] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Run the Bell Labs research simulation.
|
||||
|
||||
Args:
|
||||
task: The research question or task to investigate
|
||||
|
||||
Returns:
|
||||
Dictionary containing the research results, process history, and full conversation
|
||||
"""
|
||||
logger.info(f"Starting Bell Labs research on: {task}")
|
||||
|
||||
# Add initial task to conversation history
|
||||
self.conversation.add(
|
||||
"Research Coordinator", f"Initial Research Task: {task}"
|
||||
)
|
||||
|
||||
# Oppenheimer
|
||||
oppenheimer_plan = self.oppenheimer.run(
|
||||
task=self.conversation.get_str(), img=img
|
||||
)
|
||||
|
||||
self.conversation.add(
|
||||
self.oppenheimer.agent_name,
|
||||
f"Research Plan: {oppenheimer_plan}",
|
||||
)
|
||||
|
||||
# Discussion
|
||||
|
||||
# Physicists
|
||||
physicist_discussion = self._conduct_physicist_discussion(
|
||||
task, self.conversation.get_str()
|
||||
)
|
||||
|
||||
# Add to conversation history
|
||||
self.conversation.add(
|
||||
"Group Discussion", physicist_discussion
|
||||
)
|
||||
|
||||
# Now implement the solution
|
||||
implementation_results = self._implement_and_test_solution(
|
||||
history=self.conversation.get_str()
|
||||
)
|
||||
|
||||
# Add to conversation history
|
||||
self.conversation.add(
|
||||
self.code_implementer.agent_name, implementation_results
|
||||
)
|
||||
|
||||
return history_output_formatter(
|
||||
conversation=self.conversation, type="str"
|
||||
)
|
||||
|
||||
def _create_research_plan(
|
||||
self, task: str, loop_number: int
|
||||
) -> str:
|
||||
"""
|
||||
Have Oppenheimer create a research plan.
|
||||
|
||||
Args:
|
||||
task: Research task
|
||||
loop_number: Current loop number
|
||||
|
||||
Returns:
|
||||
Research plan from Oppenheimer
|
||||
"""
|
||||
prompt = f"""
|
||||
Research Task: {task}
|
||||
|
||||
Loop Number: {loop_number + 1}
|
||||
|
||||
As J. Robert Oppenheimer, create a comprehensive research plan for this task.
|
||||
|
||||
Your plan should include:
|
||||
1. Clear research objectives and hypotheses
|
||||
2. Theoretical framework and approach
|
||||
3. Specific research questions to investigate
|
||||
4. Methodology for testing and validation
|
||||
5. Expected outcomes and success criteria
|
||||
6. Timeline and milestones
|
||||
7. Resource requirements and team coordination
|
||||
|
||||
Provide a detailed, actionable plan that the research team can follow.
|
||||
"""
|
||||
|
||||
plan = self.oppenheimer.run(prompt)
|
||||
return plan
|
||||
|
||||
def _conduct_physicist_discussion(
|
||||
self, task: str, history: str
|
||||
) -> str:
|
||||
"""
|
||||
Conduct a natural discussion among physicists where they build on each other's ideas.
|
||||
|
||||
Args:
|
||||
task: Research task
|
||||
history: Conversation history including Oppenheimer's plan
|
||||
|
||||
Returns:
|
||||
Results of the physicist discussion as a conversation transcript
|
||||
"""
|
||||
import random
|
||||
|
||||
# Shuffle the physicists to create random discussion order
|
||||
discussion_order = self.physicists.copy()
|
||||
random.shuffle(discussion_order)
|
||||
|
||||
discussion_transcript = []
|
||||
current_context = (
|
||||
f"{history}\n\nCurrent Research Task: {task}\n\n"
|
||||
)
|
||||
|
||||
# Each physicist contributes to the discussion, building on previous contributions
|
||||
for i, physicist in enumerate(discussion_order):
|
||||
if i == 0:
|
||||
# First physicist starts the discussion
|
||||
discussion_prompt = f"""
|
||||
{current_context}
|
||||
|
||||
As {physicist.agent_name}, you are starting the group discussion about this research plan.
|
||||
|
||||
Based on your expertise, provide your initial thoughts on:
|
||||
|
||||
1. What aspects of Oppenheimer's research plan do you find most promising?
|
||||
2. What theoretical challenges or concerns do you see?
|
||||
3. What specific approaches would you recommend based on your expertise?
|
||||
4. What questions or clarifications do you have for the team?
|
||||
|
||||
Be specific and draw from your unique perspective and expertise. This will set the tone for the group discussion.
|
||||
"""
|
||||
else:
|
||||
# Subsequent physicists build on the discussion
|
||||
previous_contributions = "\n\n".join(
|
||||
discussion_transcript
|
||||
)
|
||||
discussion_prompt = f"""
|
||||
{current_context}
|
||||
|
||||
Previous Discussion:
|
||||
{previous_contributions}
|
||||
|
||||
As {physicist.agent_name}, continue the group discussion by building on your colleagues' ideas.
|
||||
|
||||
Consider:
|
||||
1. How do your colleagues' perspectives relate to your expertise in {', '.join(physicist.expertise)}?
|
||||
2. What additional insights can you add to the discussion?
|
||||
3. How can you address any concerns or questions raised by others?
|
||||
4. What specific next steps would you recommend based on the discussion so far?
|
||||
|
||||
Engage directly with your colleagues' ideas and contribute your unique perspective to move the research forward.
|
||||
"""
|
||||
|
||||
# Get the physicist's contribution
|
||||
contribution = physicist.run(discussion_prompt)
|
||||
|
||||
# Add to transcript with clear attribution
|
||||
discussion_transcript.append(
|
||||
f"{physicist.agent_name}: {contribution}"
|
||||
)
|
||||
|
||||
# Update context for next iteration
|
||||
current_context = (
|
||||
f"{history}\n\nCurrent Research Task: {task}\n\nGroup Discussion:\n"
|
||||
+ "\n\n".join(discussion_transcript)
|
||||
)
|
||||
|
||||
# Create a summary of the discussion
|
||||
summary_prompt = f"""
|
||||
Research Task: {task}
|
||||
|
||||
Complete Discussion Transcript:
|
||||
{chr(10).join(discussion_transcript)}
|
||||
|
||||
As a research coordinator, provide a concise summary of the key points from this group discussion:
|
||||
|
||||
1. Main areas of agreement among the physicists
|
||||
2. Key concerns or challenges identified
|
||||
3. Specific recommendations made by the team
|
||||
4. Next steps for moving forward with the research
|
||||
|
||||
Focus on actionable insights and clear next steps that the team can implement.
|
||||
"""
|
||||
|
||||
# Use Oppenheimer to summarize the discussion
|
||||
discussion_summary = self.oppenheimer.run(summary_prompt)
|
||||
|
||||
# Return the full discussion transcript with summary
|
||||
full_discussion = f"Group Discussion Transcript:\n\n{chr(10).join(discussion_transcript)}\n\n---\nDiscussion Summary:\n{discussion_summary}"
|
||||
|
||||
return full_discussion
|
||||
|
||||
def _implement_and_test_solution(
|
||||
self,
|
||||
history: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Implement and test the proposed solution.
|
||||
|
||||
Args:
|
||||
task: Research task
|
||||
plan: Research plan
|
||||
discussion_results: Results from physicist discussion
|
||||
loop_number: Current loop number
|
||||
|
||||
Returns:
|
||||
Implementation and testing results
|
||||
"""
|
||||
implementation_prompt = f"""
|
||||
{history}
|
||||
|
||||
As the Code Implementer, your task is to:
|
||||
|
||||
1. Implement a computational solution based on the research plan
|
||||
2. Test the theoretical predictions through simulation or calculation
|
||||
3. Analyze the results and provide quantitative assessment
|
||||
4. Identify any discrepancies between theory and implementation
|
||||
5. Suggest improvements or next steps
|
||||
|
||||
Provide:
|
||||
- Clear description of your implementation approach
|
||||
- Code or algorithm description
|
||||
- Test results and analysis
|
||||
- Comparison with theoretical predictions
|
||||
- Recommendations for further investigation
|
||||
|
||||
Focus on practical implementation and quantitative results.
|
||||
"""
|
||||
|
||||
implementation_results = self.code_implementer.run(
|
||||
implementation_prompt
|
||||
)
|
||||
|
||||
return implementation_results
|
||||
|
||||
def _analyze_results(
|
||||
self, implementation_results: Dict[str, Any], loop_number: int
|
||||
) -> str:
|
||||
"""
|
||||
Analyze the results and provide team review.
|
||||
|
||||
Args:
|
||||
implementation_results: Results from implementation phase
|
||||
loop_number: Current loop number
|
||||
|
||||
Returns:
|
||||
Analysis and recommendations
|
||||
"""
|
||||
analysis_prompt = f"""
|
||||
Implementation Results: {implementation_results}
|
||||
|
||||
Loop Number: {loop_number + 1}
|
||||
|
||||
As the research team, analyze these results and provide:
|
||||
|
||||
1. Assessment of whether the implementation supports the theoretical predictions
|
||||
2. Identification of any unexpected findings or discrepancies
|
||||
3. Evaluation of the methodology and approach
|
||||
4. Recommendations for the next research iteration
|
||||
5. Insights gained from this round of investigation
|
||||
|
||||
Consider:
|
||||
- What worked well in this approach?
|
||||
- What challenges or limitations were encountered?
|
||||
- How can the research be improved in the next iteration?
|
||||
- What new questions or directions have emerged?
|
||||
|
||||
Provide a comprehensive analysis that will guide the next research phase.
|
||||
"""
|
||||
|
||||
# Use team discussion for results analysis
|
||||
analysis_results = self._conduct_team_analysis(
|
||||
analysis_prompt
|
||||
)
|
||||
return analysis_results
|
||||
|
||||
def _conduct_team_analysis(self, analysis_prompt: str) -> str:
|
||||
"""
|
||||
Conduct a team analysis discussion using the same approach as physicist discussion.
|
||||
|
||||
Args:
|
||||
analysis_prompt: The prompt for the analysis
|
||||
|
||||
Returns:
|
||||
Results of the team analysis discussion
|
||||
"""
|
||||
import random
|
||||
|
||||
# Shuffle the agents to create random discussion order
|
||||
discussion_order = self.agents.copy()
|
||||
random.shuffle(discussion_order)
|
||||
|
||||
discussion_transcript = []
|
||||
current_context = analysis_prompt
|
||||
|
||||
# Each agent contributes to the analysis, building on previous contributions
|
||||
for i, agent in enumerate(discussion_order):
|
||||
if i == 0:
|
||||
# First agent starts the analysis
|
||||
agent_prompt = f"""
|
||||
{current_context}
|
||||
|
||||
As {agent.agent_name}, you are starting the team analysis discussion.
|
||||
|
||||
Based on your expertise and role, provide your initial analysis of the implementation results.
|
||||
Focus on what you can contribute from your unique perspective.
|
||||
"""
|
||||
else:
|
||||
# Subsequent agents build on the analysis
|
||||
previous_contributions = "\n\n".join(
|
||||
discussion_transcript
|
||||
)
|
||||
agent_prompt = f"""
|
||||
{current_context}
|
||||
|
||||
Previous Analysis:
|
||||
{previous_contributions}
|
||||
|
||||
As {agent.agent_name}, continue the team analysis by building on your colleagues' insights.
|
||||
|
||||
Consider:
|
||||
1. How do your colleagues' perspectives relate to your expertise?
|
||||
2. What additional insights can you add to the analysis?
|
||||
3. How can you address any concerns or questions raised by others?
|
||||
4. What specific recommendations would you make based on the analysis so far?
|
||||
|
||||
Engage directly with your colleagues' ideas and contribute your unique perspective.
|
||||
"""
|
||||
|
||||
# Get the agent's contribution
|
||||
contribution = agent.run(agent_prompt)
|
||||
|
||||
# Add to transcript with clear attribution
|
||||
discussion_transcript.append(
|
||||
f"{agent.agent_name}: {contribution}"
|
||||
)
|
||||
|
||||
# Update context for next iteration
|
||||
current_context = (
|
||||
f"{analysis_prompt}\n\nTeam Analysis:\n"
|
||||
+ "\n\n".join(discussion_transcript)
|
||||
)
|
||||
|
||||
# Create a summary of the analysis
|
||||
summary_prompt = f"""
|
||||
Analysis Prompt: {analysis_prompt}
|
||||
|
||||
Complete Analysis Transcript:
|
||||
{chr(10).join(discussion_transcript)}
|
||||
|
||||
As a research coordinator, provide a concise summary of the key points from this team analysis:
|
||||
|
||||
1. Main findings and insights from the team
|
||||
2. Key recommendations made
|
||||
3. Areas of agreement and disagreement
|
||||
4. Next steps for the research
|
||||
|
||||
Focus on actionable insights and clear next steps.
|
||||
"""
|
||||
|
||||
# Use Oppenheimer to summarize the analysis
|
||||
analysis_summary = self.oppenheimer.run(summary_prompt)
|
||||
|
||||
# Return the full analysis transcript with summary
|
||||
full_analysis = f"Team Analysis Transcript:\n\n{chr(10).join(discussion_transcript)}\n\n---\nAnalysis Summary:\n{analysis_summary}"
|
||||
|
||||
return full_analysis
|
||||
|
||||
def _refine_task_for_next_iteration(
|
||||
self, current_task: str, loop_results: Dict[str, Any]
|
||||
) -> str:
|
||||
"""
|
||||
Refine the task for the next research iteration.
|
||||
|
||||
Args:
|
||||
current_task: Current research task
|
||||
loop_results: Results from the current loop
|
||||
|
||||
Returns:
|
||||
Refined task for next iteration
|
||||
"""
|
||||
refinement_prompt = f"""
|
||||
Current Research Task: {current_task}
|
||||
|
||||
Results from Current Loop: {loop_results}
|
||||
|
||||
Based on the findings and analysis from this research loop, refine the research task for the next iteration.
|
||||
|
||||
Consider:
|
||||
- What new questions have emerged?
|
||||
- What aspects need deeper investigation?
|
||||
- What alternative approaches should be explored?
|
||||
- What specific hypotheses should be tested?
|
||||
|
||||
Provide a refined, focused research question that builds upon the current findings
|
||||
and addresses the most important next steps identified by the team.
|
||||
"""
|
||||
|
||||
# Use Oppenheimer to refine the task
|
||||
refined_task = self.oppenheimer.run(refinement_prompt)
|
||||
|
||||
# Add task refinement to conversation history
|
||||
self.conversation.add(
|
||||
"J. Robert Oppenheimer",
|
||||
f"Task Refined for Next Iteration: {refined_task}",
|
||||
)
|
||||
|
||||
return refined_task
|
||||
|
||||
def _generate_final_conclusion(
|
||||
self, research_results: Dict[str, Any]
|
||||
) -> str:
|
||||
"""
|
||||
Generate a final conclusion summarizing all research findings.
|
||||
|
||||
Args:
|
||||
research_results: Complete research results from all loops
|
||||
|
||||
Returns:
|
||||
Final research conclusion
|
||||
"""
|
||||
conclusion_prompt = f"""
|
||||
Complete Research Results: {research_results}
|
||||
|
||||
As J. Robert Oppenheimer, provide a comprehensive final conclusion for this research project.
|
||||
|
||||
Your conclusion should:
|
||||
1. Summarize the key findings from all research loops
|
||||
2. Identify the most significant discoveries or insights
|
||||
3. Evaluate the success of the research approach
|
||||
4. Highlight any limitations or areas for future investigation
|
||||
5. Provide a clear statement of what was accomplished
|
||||
6. Suggest next steps for continued research
|
||||
|
||||
Synthesize the work of the entire team and provide a coherent narrative
|
||||
of the research journey and its outcomes.
|
||||
"""
|
||||
|
||||
final_conclusion = self.oppenheimer.run(conclusion_prompt)
|
||||
return final_conclusion
|
||||
|
||||
|
||||
# Example usage function
|
||||
def run_bell_labs_research(
|
||||
research_question: str,
|
||||
max_loops: int = 3,
|
||||
model_name: str = "gpt-4o-mini",
|
||||
verbose: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Run a Bell Labs research simulation.
|
||||
|
||||
Args:
|
||||
research_question: The research question to investigate
|
||||
max_loops: Number of research iteration loops
|
||||
model_name: Model to use for all agents
|
||||
verbose: Whether to enable verbose logging
|
||||
|
||||
Returns:
|
||||
Complete research results and findings
|
||||
"""
|
||||
bell_labs = BellLabsSwarm(
|
||||
max_loops=max_loops, verbose=verbose, model_name=model_name
|
||||
)
|
||||
|
||||
results = bell_labs.run(research_question)
|
||||
return results
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# # Example research question
|
||||
# research_question = """
|
||||
# Investigate the feasibility of quantum computing for solving complex optimization problems.
|
||||
# Consider both theoretical foundations and practical implementation challenges.
|
||||
# """
|
||||
|
||||
# print("Starting Bell Labs Research Simulation...")
|
||||
# print(f"Research Question: {research_question}")
|
||||
# print("-" * 80)
|
||||
|
||||
# results = run_bell_labs_research(
|
||||
# research_question=research_question,
|
||||
# max_loops=2,
|
||||
# verbose=True
|
||||
# )
|
||||
|
||||
# print("\n" + "=" * 80)
|
||||
# print("RESEARCH SIMULATION COMPLETED")
|
||||
# print("=" * 80)
|
||||
|
||||
# print(f"\nFinal Conclusion:\n{results['final_conclusion']}")
|
||||
|
||||
# print(f"\nResearch completed in {len(results['research_history'])} loops.")
|
||||
# print("Check the results dictionary for complete research details.")
|
||||
@ -1,253 +0,0 @@
|
||||
from swarms.structs.agent import Agent
|
||||
from typing import List
|
||||
from swarms.structs.conversation import Conversation
|
||||
import uuid
|
||||
import random
|
||||
from loguru import logger
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class QASwarm:
|
||||
"""
|
||||
A Question and Answer swarm system where random agents ask questions to speaker agents.
|
||||
|
||||
This system allows for dynamic Q&A sessions where:
|
||||
- Multiple agents can act as questioners
|
||||
- One or multiple agents can act as speakers/responders
|
||||
- Questions are asked randomly by different agents
|
||||
- The conversation is tracked and managed
|
||||
- Agents are showcased to each other with detailed information
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "QandA",
|
||||
description: str = "Question and Answer Swarm System",
|
||||
agents: List[Agent] = None,
|
||||
speaker_agents: List[Agent] = None,
|
||||
id: str = str(uuid.uuid4()),
|
||||
max_loops: int = 5,
|
||||
show_dashboard: bool = True,
|
||||
speaker_agent: Agent = None,
|
||||
showcase_agents: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
self.id = id
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.max_loops = max_loops
|
||||
self.show_dashboard = show_dashboard
|
||||
self.agents = agents or []
|
||||
self.speaker_agents = speaker_agents or []
|
||||
self.kwargs = kwargs
|
||||
self.speaker_agent = speaker_agent
|
||||
self.showcase_agents = showcase_agents
|
||||
|
||||
self.conversation = Conversation()
|
||||
|
||||
# Validate setup
|
||||
self._validate_setup()
|
||||
|
||||
def _validate_setup(self):
|
||||
"""Validate that the Q&A system is properly configured."""
|
||||
if not self.agents:
|
||||
logger.warning(
|
||||
"No questioner agents provided. Add agents using add_agent() method."
|
||||
)
|
||||
|
||||
if not self.speaker_agents and not self.speaker_agent:
|
||||
logger.warning(
|
||||
"No speaker agents provided. Add speaker agents using add_speaker_agent() method."
|
||||
)
|
||||
|
||||
if (
|
||||
not self.agents
|
||||
and not self.speaker_agents
|
||||
and not self.speaker_agent
|
||||
):
|
||||
raise ValueError(
|
||||
"At least one agent (questioner or speaker) must be provided."
|
||||
)
|
||||
|
||||
def add_agent(self, agent: Agent):
|
||||
"""Add a questioner agent to the swarm."""
|
||||
self.agents.append(agent)
|
||||
logger.info(f"Added questioner agent: {agent.agent_name}")
|
||||
|
||||
def add_speaker_agent(self, agent: Agent):
|
||||
"""Add a speaker agent to the swarm."""
|
||||
if self.speaker_agents is None:
|
||||
self.speaker_agents = []
|
||||
self.speaker_agents.append(agent)
|
||||
logger.info(f"Added speaker agent: {agent.agent_name}")
|
||||
|
||||
def get_agent_info(self, agent: Agent) -> dict:
|
||||
"""Extract key information about an agent for showcasing."""
|
||||
info = {
|
||||
"name": getattr(agent, "agent_name", "Unknown Agent"),
|
||||
"description": getattr(
|
||||
agent, "agent_description", "No description available"
|
||||
),
|
||||
"role": getattr(agent, "role", "worker"),
|
||||
}
|
||||
|
||||
# Get system prompt preview (first 50 characters)
|
||||
system_prompt = getattr(agent, "system_prompt", "")
|
||||
if system_prompt:
|
||||
info["system_prompt_preview"] = (
|
||||
system_prompt[:50] + "..."
|
||||
if len(system_prompt) > 50
|
||||
else system_prompt
|
||||
)
|
||||
else:
|
||||
info["system_prompt_preview"] = (
|
||||
"No system prompt available"
|
||||
)
|
||||
|
||||
return info
|
||||
|
||||
def showcase_speaker_to_questioner(
|
||||
self, questioner: Agent, speaker: Agent
|
||||
) -> str:
|
||||
"""Create a showcase prompt introducing the speaker agent to the questioner."""
|
||||
speaker_info = self.get_agent_info(speaker)
|
||||
|
||||
showcase_prompt = f"""
|
||||
You are about to ask a question to a specialized agent. Here's what you need to know about them:
|
||||
|
||||
**Speaker Agent Information:**
|
||||
- **Name**: {speaker_info['name']}
|
||||
- **Role**: {speaker_info['role']}
|
||||
- **Description**: {speaker_info['description']}
|
||||
- **System Prompt Preview**: {speaker_info['system_prompt_preview']}
|
||||
|
||||
Please craft a thoughtful, relevant question that takes into account this agent's expertise and background.
|
||||
Your question should be specific and demonstrate that you understand their role and capabilities.
|
||||
"""
|
||||
return showcase_prompt
|
||||
|
||||
def showcase_questioner_to_speaker(
|
||||
self, speaker: Agent, questioner: Agent
|
||||
) -> str:
|
||||
"""Create a showcase prompt introducing the questioner agent to the speaker."""
|
||||
questioner_info = self.get_agent_info(questioner)
|
||||
|
||||
showcase_prompt = f"""
|
||||
You are about to answer a question from another agent. Here's what you need to know about them:
|
||||
|
||||
**Questioner Agent Information:**
|
||||
- **Name**: {questioner_info['name']}
|
||||
- **Role**: {questioner_info['role']}
|
||||
- **Description**: {questioner_info['description']}
|
||||
- **System Prompt Preview**: {questioner_info['system_prompt_preview']}
|
||||
|
||||
Please provide a comprehensive answer that demonstrates your expertise and addresses their question thoroughly.
|
||||
Consider their background and role when formulating your response.
|
||||
"""
|
||||
return showcase_prompt
|
||||
|
||||
def random_select_agent(self, agents: List[Agent]) -> Agent:
|
||||
"""Randomly select an agent from the list."""
|
||||
if not agents:
|
||||
raise ValueError("No agents available for selection")
|
||||
return random.choice(agents)
|
||||
|
||||
def get_current_speaker(self) -> Agent:
|
||||
"""Get the current speaker agent (either from speaker_agents list or single speaker_agent)."""
|
||||
if self.speaker_agent:
|
||||
return self.speaker_agent
|
||||
elif self.speaker_agents:
|
||||
return self.random_select_agent(self.speaker_agents)
|
||||
else:
|
||||
raise ValueError("No speaker agent available")
|
||||
|
||||
def run(
|
||||
self, task: str, img: Optional[str] = None, *args, **kwargs
|
||||
):
|
||||
"""Run the Q&A session with agent showcasing."""
|
||||
self.conversation.add(role="user", content=task)
|
||||
|
||||
# Get current speaker
|
||||
current_speaker = self.get_current_speaker()
|
||||
|
||||
# Select a random questioner
|
||||
questioner = self.random_select_agent(self.agents)
|
||||
|
||||
# Showcase agents to each other if enabled
|
||||
if self.showcase_agents:
|
||||
# Showcase speaker to questioner
|
||||
speaker_showcase = self.showcase_speaker_to_questioner(
|
||||
questioner, current_speaker
|
||||
)
|
||||
questioner_task = f"{speaker_showcase}\n\nNow ask a question about: {task}"
|
||||
|
||||
# Showcase questioner to speaker
|
||||
questioner_showcase = self.showcase_questioner_to_speaker(
|
||||
current_speaker, questioner
|
||||
)
|
||||
else:
|
||||
questioner_task = f"Ask a question about {task} to {current_speaker.agent_name}"
|
||||
|
||||
# Generate question
|
||||
question = questioner.run(
|
||||
task=questioner_task,
|
||||
img=img,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.conversation.add(
|
||||
role=questioner.agent_name, content=question
|
||||
)
|
||||
|
||||
# Prepare answer task with showcasing if enabled
|
||||
if self.showcase_agents:
|
||||
answer_task = f"{questioner_showcase}\n\nAnswer this question from {questioner.agent_name}: {question}"
|
||||
else:
|
||||
answer_task = f"Answer the question '{question}' from {questioner.agent_name}"
|
||||
|
||||
# Generate answer
|
||||
answer = current_speaker.run(
|
||||
task=answer_task,
|
||||
img=img,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.conversation.add(
|
||||
role=current_speaker.agent_name, content=answer
|
||||
)
|
||||
|
||||
return answer
|
||||
|
||||
def run_multi_round(
|
||||
self,
|
||||
task: str,
|
||||
rounds: int = 3,
|
||||
img: Optional[str] = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
"""Run multiple rounds of Q&A with different questioners."""
|
||||
results = []
|
||||
|
||||
for round_num in range(rounds):
|
||||
logger.info(
|
||||
f"Starting Q&A round {round_num + 1}/{rounds}"
|
||||
)
|
||||
|
||||
round_result = self.run(task, img, *args, **kwargs)
|
||||
results.append(
|
||||
{"round": round_num + 1, "result": round_result}
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
def get_conversation_history(self):
|
||||
"""Get the conversation history."""
|
||||
return self.conversation.get_history()
|
||||
|
||||
def clear_conversation(self):
|
||||
"""Clear the conversation history."""
|
||||
self.conversation = Conversation()
|
||||
logger.info("Conversation history cleared")
|
||||
@ -1,191 +0,0 @@
|
||||
from pydantic.v1 import BaseModel
|
||||
from typing import List, Callable
|
||||
from swarms.utils.loguru_logger import initialize_logger
|
||||
|
||||
logger = initialize_logger(log_folder="swarm_registry")
|
||||
|
||||
|
||||
class SwarmRegistry(BaseModel):
|
||||
swarm_pool: List[Callable] = []
|
||||
|
||||
def add(self, swarm: Callable, *args, **kwargs):
|
||||
"""
|
||||
Adds a swarm to the registry.
|
||||
|
||||
Args:
|
||||
swarm (Callable): The swarm to add to the registry.
|
||||
"""
|
||||
self.swarm_pool.append(swarm, *args, **kwargs)
|
||||
|
||||
def query(self, swarm_name: str) -> Callable:
|
||||
"""
|
||||
Queries the registry for a swarm by name.
|
||||
|
||||
Args:
|
||||
swarm_name (str): The name of the swarm to query.
|
||||
|
||||
Returns:
|
||||
Callable: The swarm function corresponding to the given name.
|
||||
"""
|
||||
if not self.swarm_pool:
|
||||
raise ValueError("No swarms found in registry")
|
||||
|
||||
if not swarm_name:
|
||||
raise ValueError("No swarm name provided.")
|
||||
|
||||
for swarm in self.swarm_pool:
|
||||
if swarm.__name__ == swarm_name:
|
||||
name = swarm.__name__
|
||||
description = (
|
||||
swarm.__doc__.strip().split("\n")[0]
|
||||
or swarm.description
|
||||
)
|
||||
agent_count = len(swarm.agents)
|
||||
task_count = len(swarm.tasks)
|
||||
|
||||
log = f"Swarm: {name}\nDescription: {description}\nAgents: {agent_count}\nTasks: {task_count}"
|
||||
logger.info(log)
|
||||
|
||||
return swarm
|
||||
|
||||
raise ValueError(
|
||||
f"Swarm '{swarm_name}' not found in registry."
|
||||
)
|
||||
|
||||
def remove(self, swarm_name: str):
|
||||
"""
|
||||
Removes a swarm from the registry by name.
|
||||
|
||||
Args:
|
||||
swarm_name (str): The name of the swarm to remove.
|
||||
"""
|
||||
for swarm in self.swarm_pool:
|
||||
if swarm.__name__ == swarm_name:
|
||||
self.swarm_pool.remove(swarm)
|
||||
return
|
||||
raise ValueError(
|
||||
f"Swarm '{swarm_name}' not found in registry."
|
||||
)
|
||||
|
||||
def list_swarms(self) -> List[str]:
|
||||
"""
|
||||
Lists the names of all swarms in the registry.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of swarm names.
|
||||
"""
|
||||
if not self.swarm_pool:
|
||||
raise ValueError("No swarms found in registry.")
|
||||
|
||||
for swarm in self.swarm_pool:
|
||||
name = swarm.__name__
|
||||
description = (
|
||||
swarm.__doc__.strip().split("\n")[0]
|
||||
or swarm.description
|
||||
)
|
||||
agent_count = len(swarm.agents)
|
||||
task_count = len(swarm.tasks)
|
||||
|
||||
log = f"Swarm: {name}\nDescription: {description}\nAgents: {agent_count}\nTasks: {task_count}"
|
||||
logger.info(log)
|
||||
|
||||
return [swarm.__name__ for swarm in self.swarm_pool]
|
||||
|
||||
def run(self, swarm_name: str, *args, **kwargs):
|
||||
"""
|
||||
Runs a swarm by name with the given arguments.
|
||||
|
||||
Args:
|
||||
swarm_name (str): The name of the swarm to run.
|
||||
*args: Variable length argument list.
|
||||
**kwargs: Arbitrary keyword arguments.
|
||||
|
||||
Returns:
|
||||
Any: The result of running the swarm.
|
||||
"""
|
||||
swarm = self.query(swarm_name)
|
||||
return swarm(*args, **kwargs)
|
||||
|
||||
def add_list_of_swarms(self, swarms: List[Callable]):
|
||||
"""
|
||||
Adds a list of swarms to the registry.
|
||||
|
||||
Args:
|
||||
swarms (List[Callable]): A list of swarms to add to the registry.
|
||||
"""
|
||||
for swarm in swarms:
|
||||
self.add(swarm)
|
||||
|
||||
return self.swarm_pool
|
||||
|
||||
def query_multiple_of_swarms(
|
||||
self, swarm_names: List[str]
|
||||
) -> List[Callable]:
|
||||
"""
|
||||
Queries the registry for multiple swarms by name.
|
||||
|
||||
Args:
|
||||
swarm_names (List[str]): A list of swarm names to query.
|
||||
|
||||
Returns:
|
||||
List[Callable]: A list of swarm functions corresponding to the given names.
|
||||
"""
|
||||
return [self.query(swarm_name) for swarm_name in swarm_names]
|
||||
|
||||
def remove_list_of_swarms(self, swarm_names: List[str]):
|
||||
"""
|
||||
Removes a list of swarms from the registry by name.
|
||||
|
||||
Args:
|
||||
swarm_names (List[str]): A list of swarm names to remove.
|
||||
"""
|
||||
for swarm_name in swarm_names:
|
||||
self.remove(swarm_name)
|
||||
|
||||
return self.swarm_pool
|
||||
|
||||
def run_multiple_of_swarms(
|
||||
self, swarm_names: List[str], *args, **kwargs
|
||||
):
|
||||
"""
|
||||
Runs a list of swarms by name with the given arguments.
|
||||
|
||||
Args:
|
||||
swarm_names (List[str]): A list of swarm names to run.
|
||||
*args: Variable length argument list.
|
||||
**kwargs: Arbitrary keyword arguments.
|
||||
|
||||
Returns:
|
||||
List[Any]: A list of results of running the swarms.
|
||||
"""
|
||||
return [
|
||||
self.run(swarm_name, *args, **kwargs)
|
||||
for swarm_name in swarm_names
|
||||
]
|
||||
|
||||
|
||||
# Decorator to add a function to the registry
|
||||
def swarm_registry():
|
||||
"""
|
||||
Decorator to add a function to the registry.
|
||||
|
||||
Args:
|
||||
swarm_registry (SwarmRegistry): The swarm registry instance.
|
||||
|
||||
Returns:
|
||||
Callable: The decorated function.
|
||||
"""
|
||||
|
||||
def decorator(func, *args, **kwargs):
|
||||
try:
|
||||
swarm_registry = SwarmRegistry()
|
||||
swarm_registry.add(func, *args, **kwargs)
|
||||
logger.info(
|
||||
f"Added swarm '{func.__name__}' to the registry."
|
||||
)
|
||||
return func
|
||||
except Exception as e:
|
||||
logger.error(str(e))
|
||||
raise
|
||||
|
||||
return decorator
|
||||
@ -1,18 +0,0 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Dict
|
||||
|
||||
|
||||
class ParameterDefinition(BaseModel):
|
||||
description: str = Field(
|
||||
..., title="Description of the parameter"
|
||||
)
|
||||
type: str = Field(..., title="Type of the parameter")
|
||||
required: bool = Field(..., title="Is the parameter required?")
|
||||
|
||||
|
||||
class CohereFuncSchema(BaseModel):
|
||||
name: str = Field(..., title="Name of the tool")
|
||||
description: str = Field(..., title="Description of the tool")
|
||||
parameter_definitions: Dict[str, ParameterDefinition] = Field(
|
||||
..., title="Parameter definitions for the tool"
|
||||
)
|
||||
@ -1,343 +0,0 @@
|
||||
import base64
|
||||
from typing import Union, Dict, Any, Tuple
|
||||
import requests
|
||||
from pathlib import Path
|
||||
import wave
|
||||
import numpy as np
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_path: Union[str, Path]) -> str:
|
||||
"""
|
||||
Encode a WAV file to base64 string.
|
||||
|
||||
Args:
|
||||
audio_path (Union[str, Path]): Path to the WAV file
|
||||
|
||||
Returns:
|
||||
str: Base64 encoded string of the audio file
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the audio file doesn't exist
|
||||
ValueError: If the file is not a valid WAV file
|
||||
"""
|
||||
try:
|
||||
audio_path = Path(audio_path)
|
||||
if not audio_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Audio file not found: {audio_path}"
|
||||
)
|
||||
|
||||
if not audio_path.suffix.lower() == ".wav":
|
||||
raise ValueError("File must be a WAV file")
|
||||
|
||||
with open(audio_path, "rb") as audio_file:
|
||||
audio_data = audio_file.read()
|
||||
return base64.b64encode(audio_data).decode("utf-8")
|
||||
except Exception as e:
|
||||
raise Exception(f"Error encoding audio file: {str(e)}")
|
||||
|
||||
|
||||
def decode_base64_to_audio(
|
||||
base64_string: str, output_path: Union[str, Path]
|
||||
) -> None:
|
||||
"""
|
||||
Decode a base64 string to a WAV file.
|
||||
|
||||
Args:
|
||||
base64_string (str): Base64 encoded audio data
|
||||
output_path (Union[str, Path]): Path where the WAV file should be saved
|
||||
|
||||
Raises:
|
||||
ValueError: If the base64 string is invalid
|
||||
IOError: If there's an error writing the file
|
||||
"""
|
||||
try:
|
||||
output_path = Path(output_path)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
audio_data = base64.b64decode(base64_string)
|
||||
with open(output_path, "wb") as audio_file:
|
||||
audio_file.write(audio_data)
|
||||
except Exception as e:
|
||||
raise Exception(f"Error decoding audio data: {str(e)}")
|
||||
|
||||
|
||||
def download_audio_from_url(
|
||||
url: str, output_path: Union[str, Path]
|
||||
) -> None:
|
||||
"""
|
||||
Download an audio file from a URL and save it locally.
|
||||
|
||||
Args:
|
||||
url (str): URL of the audio file
|
||||
output_path (Union[str, Path]): Path where the audio file should be saved
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If there's an error downloading the file
|
||||
IOError: If there's an error saving the file
|
||||
"""
|
||||
try:
|
||||
output_path = Path(output_path)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
response = requests.get(url)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(output_path, "wb") as audio_file:
|
||||
audio_file.write(response.content)
|
||||
except Exception as e:
|
||||
raise Exception(f"Error downloading audio file: {str(e)}")
|
||||
|
||||
|
||||
def process_audio_with_model(
|
||||
audio_path: Union[str, Path],
|
||||
model: str,
|
||||
prompt: str,
|
||||
voice: str = "alloy",
|
||||
format: str = "wav",
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Process an audio file with a model that supports audio input/output.
|
||||
|
||||
Args:
|
||||
audio_path (Union[str, Path]): Path to the input WAV file
|
||||
model (str): Model name to use for processing
|
||||
prompt (str): Text prompt to accompany the audio
|
||||
voice (str, optional): Voice to use for audio output. Defaults to "alloy"
|
||||
format (str, optional): Audio format. Defaults to "wav"
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: Model response containing both text and audio if applicable
|
||||
|
||||
Raises:
|
||||
ImportError: If litellm is not installed
|
||||
ValueError: If the model doesn't support audio processing
|
||||
"""
|
||||
try:
|
||||
from litellm import (
|
||||
completion,
|
||||
supports_audio_input,
|
||||
supports_audio_output,
|
||||
)
|
||||
|
||||
if not supports_audio_input(model):
|
||||
raise ValueError(
|
||||
f"Model {model} does not support audio input"
|
||||
)
|
||||
|
||||
# Encode the audio file
|
||||
encoded_audio = encode_audio_to_base64(audio_path)
|
||||
|
||||
# Prepare the messages
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {
|
||||
"data": encoded_audio,
|
||||
"format": format,
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# Make the API call
|
||||
response = completion(
|
||||
model=model,
|
||||
modalities=["text", "audio"],
|
||||
audio={"voice": voice, "format": format},
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
return response
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Please install litellm: pip install litellm"
|
||||
)
|
||||
except Exception as e:
|
||||
raise Exception(
|
||||
f"Error processing audio with model: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
def read_wav_file(
|
||||
file_path: Union[str, Path],
|
||||
) -> Tuple[np.ndarray, int]:
|
||||
"""
|
||||
Read a WAV file and return its audio data and sample rate.
|
||||
|
||||
Args:
|
||||
file_path (Union[str, Path]): Path to the WAV file
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, int]: Audio data as numpy array and sample rate
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the file doesn't exist
|
||||
ValueError: If the file is not a valid WAV file
|
||||
"""
|
||||
try:
|
||||
file_path = Path(file_path)
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Audio file not found: {file_path}"
|
||||
)
|
||||
|
||||
with wave.open(str(file_path), "rb") as wav_file:
|
||||
# Get audio parameters
|
||||
n_channels = wav_file.getnchannels()
|
||||
sample_width = wav_file.getsampwidth()
|
||||
frame_rate = wav_file.getframerate()
|
||||
n_frames = wav_file.getnframes()
|
||||
|
||||
# Read audio data
|
||||
frames = wav_file.readframes(n_frames)
|
||||
|
||||
# Convert to numpy array
|
||||
dtype = np.int16 if sample_width == 2 else np.int8
|
||||
audio_data = np.frombuffer(frames, dtype=dtype)
|
||||
|
||||
# Reshape if stereo
|
||||
if n_channels == 2:
|
||||
audio_data = audio_data.reshape(-1, 2)
|
||||
|
||||
return audio_data, frame_rate
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"Error reading WAV file: {str(e)}")
|
||||
|
||||
|
||||
def write_wav_file(
|
||||
audio_data: np.ndarray,
|
||||
file_path: Union[str, Path],
|
||||
sample_rate: int,
|
||||
sample_width: int = 2,
|
||||
) -> None:
|
||||
"""
|
||||
Write audio data to a WAV file.
|
||||
|
||||
Args:
|
||||
audio_data (np.ndarray): Audio data as numpy array
|
||||
file_path (Union[str, Path]): Path where to save the WAV file
|
||||
sample_rate (int): Sample rate of the audio
|
||||
sample_width (int, optional): Sample width in bytes. Defaults to 2 (16-bit)
|
||||
|
||||
Raises:
|
||||
ValueError: If the audio data is invalid
|
||||
IOError: If there's an error writing the file
|
||||
"""
|
||||
try:
|
||||
file_path = Path(file_path)
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Ensure audio data is in the correct format
|
||||
if audio_data.dtype != np.int16 and sample_width == 2:
|
||||
audio_data = (audio_data * 32767).astype(np.int16)
|
||||
elif audio_data.dtype != np.int8 and sample_width == 1:
|
||||
audio_data = (audio_data * 127).astype(np.int8)
|
||||
|
||||
# Determine number of channels
|
||||
n_channels = (
|
||||
2
|
||||
if len(audio_data.shape) > 1 and audio_data.shape[1] == 2
|
||||
else 1
|
||||
)
|
||||
|
||||
with wave.open(str(file_path), "wb") as wav_file:
|
||||
wav_file.setnchannels(n_channels)
|
||||
wav_file.setsampwidth(sample_width)
|
||||
wav_file.setframerate(sample_rate)
|
||||
wav_file.writeframes(audio_data.tobytes())
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"Error writing WAV file: {str(e)}")
|
||||
|
||||
|
||||
def normalize_audio(audio_data: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Normalize audio data to have maximum amplitude of 1.0.
|
||||
|
||||
Args:
|
||||
audio_data (np.ndarray): Input audio data
|
||||
|
||||
Returns:
|
||||
np.ndarray: Normalized audio data
|
||||
"""
|
||||
return audio_data / np.max(np.abs(audio_data))
|
||||
|
||||
|
||||
def convert_to_mono(audio_data: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Convert stereo audio to mono by averaging channels.
|
||||
|
||||
Args:
|
||||
audio_data (np.ndarray): Input audio data (stereo)
|
||||
|
||||
Returns:
|
||||
np.ndarray: Mono audio data
|
||||
"""
|
||||
if len(audio_data.shape) == 1:
|
||||
return audio_data
|
||||
return np.mean(audio_data, axis=1)
|
||||
|
||||
|
||||
def encode_wav_to_base64(
|
||||
audio_data: np.ndarray, sample_rate: int
|
||||
) -> str:
|
||||
"""
|
||||
Convert audio data to base64 encoded WAV string.
|
||||
|
||||
Args:
|
||||
audio_data (np.ndarray): Audio data
|
||||
sample_rate (int): Sample rate of the audio
|
||||
|
||||
Returns:
|
||||
str: Base64 encoded WAV data
|
||||
"""
|
||||
# Create a temporary WAV file in memory
|
||||
with wave.open("temp.wav", "wb") as wav_file:
|
||||
wav_file.setnchannels(1 if len(audio_data.shape) == 1 else 2)
|
||||
wav_file.setsampwidth(2) # 16-bit
|
||||
wav_file.setframerate(sample_rate)
|
||||
wav_file.writeframes(audio_data.tobytes())
|
||||
|
||||
# Read the file and encode to base64
|
||||
with open("temp.wav", "rb") as f:
|
||||
wav_bytes = f.read()
|
||||
|
||||
# Clean up temporary file
|
||||
Path("temp.wav").unlink()
|
||||
|
||||
return base64.b64encode(wav_bytes).decode("utf-8")
|
||||
|
||||
|
||||
def decode_base64_to_wav(
|
||||
base64_string: str,
|
||||
) -> Tuple[np.ndarray, int]:
|
||||
"""
|
||||
Convert base64 encoded WAV string to audio data and sample rate.
|
||||
|
||||
Args:
|
||||
base64_string (str): Base64 encoded WAV data
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, int]: Audio data and sample rate
|
||||
"""
|
||||
# Decode base64 string
|
||||
wav_bytes = base64.b64decode(base64_string)
|
||||
|
||||
# Write to temporary file
|
||||
with open("temp.wav", "wb") as f:
|
||||
f.write(wav_bytes)
|
||||
|
||||
# Read the WAV file
|
||||
audio_data, sample_rate = read_wav_file("temp.wav")
|
||||
|
||||
# Clean up temporary file
|
||||
Path("temp.wav").unlink()
|
||||
|
||||
return audio_data, sample_rate
|
||||
@ -1,151 +0,0 @@
|
||||
"""
|
||||
Package installation utility that checks for package existence and installs if needed.
|
||||
Supports both pip and conda package managers.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Literal, Optional, Union
|
||||
from swarms.utils.loguru_logger import initialize_logger
|
||||
|
||||
|
||||
from importlib.metadata import distribution, PackageNotFoundError
|
||||
|
||||
logger = initialize_logger("autocheckpackages")
|
||||
|
||||
|
||||
def check_and_install_package(
|
||||
package_name: str,
|
||||
package_manager: Literal["pip", "conda"] = "pip",
|
||||
version: Optional[str] = None,
|
||||
upgrade: bool = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a package is installed and install it if not found.
|
||||
|
||||
Args:
|
||||
package_name: Name of the package to check/install
|
||||
package_manager: Package manager to use ('pip' or 'conda')
|
||||
version: Specific version to install (optional)
|
||||
upgrade: Whether to upgrade the package if it exists
|
||||
|
||||
Returns:
|
||||
bool: True if package is available after check/install, False if installation failed
|
||||
|
||||
Raises:
|
||||
ValueError: If invalid package manager is specified
|
||||
"""
|
||||
try:
|
||||
# Check if package exists
|
||||
if package_manager == "pip":
|
||||
try:
|
||||
distribution(package_name)
|
||||
if not upgrade:
|
||||
logger.info(
|
||||
f"Package {package_name} is already installed"
|
||||
)
|
||||
return True
|
||||
except PackageNotFoundError:
|
||||
pass
|
||||
|
||||
# Construct installation command
|
||||
cmd = [sys.executable, "-m", "pip", "install"]
|
||||
if upgrade:
|
||||
cmd.append("--upgrade")
|
||||
|
||||
if version:
|
||||
cmd.append(f"{package_name}=={version}")
|
||||
else:
|
||||
cmd.append(package_name)
|
||||
|
||||
elif package_manager == "conda":
|
||||
# Check if conda is available
|
||||
try:
|
||||
subprocess.run(
|
||||
["conda", "--version"],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
)
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
logger.error(
|
||||
"Conda is not available. Please install conda first."
|
||||
)
|
||||
return False
|
||||
|
||||
# Construct conda command
|
||||
cmd = ["conda", "install", "-y"]
|
||||
if version:
|
||||
cmd.append(f"{package_name}={version}")
|
||||
else:
|
||||
cmd.append(package_name)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid package manager: {package_manager}"
|
||||
)
|
||||
|
||||
# Run installation
|
||||
logger.info(f"Installing {package_name}...")
|
||||
subprocess.run(
|
||||
cmd, check=True, capture_output=True, text=True
|
||||
)
|
||||
|
||||
# Verify installation
|
||||
try:
|
||||
importlib.import_module(package_name)
|
||||
logger.info(f"Successfully installed {package_name}")
|
||||
return True
|
||||
except ImportError:
|
||||
logger.error(
|
||||
f"Package {package_name} was installed but cannot be imported"
|
||||
)
|
||||
return False
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.error(f"Failed to install {package_name}: {e.stderr}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Unexpected error while installing {package_name}: {str(e)}"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def auto_check_and_download_package(
|
||||
packages: Union[str, list[str]],
|
||||
package_manager: Literal["pip", "conda"] = "pip",
|
||||
upgrade: bool = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Ensure multiple packages are installed.
|
||||
|
||||
Args:
|
||||
packages: Single package name or list of package names
|
||||
package_manager: Package manager to use ('pip' or 'conda')
|
||||
upgrade: Whether to upgrade existing packages
|
||||
|
||||
Returns:
|
||||
bool: True if all packages are available, False if any installation failed
|
||||
"""
|
||||
if isinstance(packages, str):
|
||||
packages = [packages]
|
||||
|
||||
success = True
|
||||
for package in packages:
|
||||
if ":" in package:
|
||||
name, version = package.split(":")
|
||||
if not check_and_install_package(
|
||||
name, package_manager, version, upgrade
|
||||
):
|
||||
success = False
|
||||
else:
|
||||
if not check_and_install_package(
|
||||
package, package_manager, upgrade=upgrade
|
||||
):
|
||||
success = False
|
||||
|
||||
return success
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# print(auto_check_and_download_package("torch"))
|
||||
@ -1,88 +0,0 @@
|
||||
# Best LLM Models by Task Type
|
||||
# Simplified dictionary structure with model names and categories
|
||||
|
||||
best_models = {
|
||||
"Vision": [
|
||||
{"model": "gemini/gemini-2.5-pro", "category": "Vision"},
|
||||
],
|
||||
"text-generation": [
|
||||
{
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"category": "text-generation",
|
||||
},
|
||||
{"model": "gpt-5-chat", "category": "text-generation"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
# Function to get all models for a task type
|
||||
def get_models_by_task(task_type: str) -> list:
|
||||
"""
|
||||
Get all models for a specific task type.
|
||||
|
||||
Args:
|
||||
task_type (str): The task category (e.g., 'WebDev', 'Vision', 'text-generation')
|
||||
|
||||
Returns:
|
||||
list: List of all models for the task type
|
||||
"""
|
||||
if task_type not in best_models:
|
||||
raise ValueError(
|
||||
f"Task type '{task_type}' not found. Available types: {list(best_models.keys())}"
|
||||
)
|
||||
|
||||
return best_models[task_type]
|
||||
|
||||
|
||||
# Function to get the first model for a task type (simplified from get_top_model)
|
||||
def get_first_model(task_type: str) -> dict:
|
||||
"""
|
||||
Get the first model for a specific task type.
|
||||
|
||||
Args:
|
||||
task_type (str): The task category (e.g., 'WebDev', 'Vision', 'text-generation')
|
||||
|
||||
Returns:
|
||||
dict: First model information with model name and category
|
||||
"""
|
||||
if task_type not in best_models:
|
||||
raise ValueError(
|
||||
f"Task type '{task_type}' not found. Available types: {list(best_models.keys())}"
|
||||
)
|
||||
|
||||
models = best_models[task_type]
|
||||
if not models:
|
||||
raise ValueError(
|
||||
f"No models found for task type '{task_type}'"
|
||||
)
|
||||
|
||||
return models[0]
|
||||
|
||||
|
||||
# Function to search for a specific model across all categories
|
||||
def find_model_by_name(model_name: str) -> dict:
|
||||
"""
|
||||
Find a model by name across all task categories.
|
||||
|
||||
Args:
|
||||
model_name (str): The model name to search for
|
||||
|
||||
Returns:
|
||||
dict: Model information if found, None otherwise
|
||||
"""
|
||||
for task_type, models in best_models.items():
|
||||
for model in models:
|
||||
if model["model"].lower() == model_name.lower():
|
||||
return model
|
||||
return None
|
||||
|
||||
|
||||
# Function to get all available task types
|
||||
def get_available_task_types() -> list:
|
||||
"""
|
||||
Get all available task types/categories.
|
||||
|
||||
Returns:
|
||||
list: List of all task type names
|
||||
"""
|
||||
return list(best_models.keys())
|
||||
@ -1,54 +0,0 @@
|
||||
from typing import Any
|
||||
from litellm import image_generation
|
||||
|
||||
|
||||
class ImageGenerator:
|
||||
def __init__(
|
||||
self,
|
||||
model: str | None = None,
|
||||
n: int | None = 2,
|
||||
quality: Any = None,
|
||||
response_format: str | None = None,
|
||||
size: str | None = 10,
|
||||
style: str | None = None,
|
||||
user: str | None = None,
|
||||
input_fidelity: str | None = None,
|
||||
timeout: int = 600,
|
||||
output_path_folder: str | None = "images",
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
):
|
||||
self.model = model
|
||||
self.n = n
|
||||
self.quality = quality
|
||||
self.response_format = response_format
|
||||
self.size = size
|
||||
self.style = style
|
||||
self.user = user
|
||||
self.input_fidelity = input_fidelity
|
||||
self.timeout = timeout
|
||||
self.output_path_folder = output_path_folder
|
||||
self.api_key = api_key
|
||||
self.api_base = api_base
|
||||
|
||||
def run(self, task: str = None):
|
||||
|
||||
return image_generation(
|
||||
prompt=task,
|
||||
model=self.model,
|
||||
n=self.n,
|
||||
quality=self.quality,
|
||||
response_format=self.response_format,
|
||||
size=self.size,
|
||||
style=self.style,
|
||||
user=self.user,
|
||||
input_fidelity=self.input_fidelity,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# image_generator = ImageGenerator()
|
||||
# print(image_generator.run(task="A beautiful sunset over a calm ocean"))
|
||||
|
||||
# print(model_list)
|
||||
@ -1,5 +0,0 @@
|
||||
def litellm_check_for_tools(model_name: str):
|
||||
"""Check if the model supports tools."""
|
||||
from litellm.utils import supports_function_calling
|
||||
|
||||
return supports_function_calling(model_name)
|
||||
@ -1,114 +0,0 @@
|
||||
from unittest.mock import MagicMock
|
||||
import unittest
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.tools.tool_parse_exec import parse_and_execute_json
|
||||
|
||||
# Mock parse_and_execute_json for testing
|
||||
parse_and_execute_json = MagicMock()
|
||||
parse_and_execute_json.return_value = {
|
||||
"tool_name": "calculator",
|
||||
"args": {"numbers": [2, 2]},
|
||||
"output": "4",
|
||||
}
|
||||
|
||||
|
||||
class TestAgentLogging(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.mock_tokenizer = MagicMock()
|
||||
self.mock_tokenizer.count_tokens.return_value = 100
|
||||
|
||||
self.mock_short_memory = MagicMock()
|
||||
self.mock_short_memory.get_memory_stats.return_value = {
|
||||
"message_count": 2
|
||||
}
|
||||
|
||||
self.mock_long_memory = MagicMock()
|
||||
self.mock_long_memory.get_memory_stats.return_value = {
|
||||
"item_count": 5
|
||||
}
|
||||
|
||||
self.agent = Agent(
|
||||
tokenizer=self.mock_tokenizer,
|
||||
short_memory=self.mock_short_memory,
|
||||
long_term_memory=self.mock_long_memory,
|
||||
)
|
||||
|
||||
def test_log_step_metadata_basic(self):
|
||||
log_result = self.agent.log_step_metadata(
|
||||
1, "Test prompt", "Test response"
|
||||
)
|
||||
|
||||
self.assertIn("step_id", log_result)
|
||||
self.assertIn("timestamp", log_result)
|
||||
self.assertIn("tokens", log_result)
|
||||
self.assertIn("memory_usage", log_result)
|
||||
|
||||
self.assertEqual(log_result["tokens"]["total"], 200)
|
||||
|
||||
def test_log_step_metadata_no_long_term_memory(self):
|
||||
self.agent.long_term_memory = None
|
||||
log_result = self.agent.log_step_metadata(
|
||||
1, "prompt", "response"
|
||||
)
|
||||
self.assertEqual(log_result["memory_usage"]["long_term"], {})
|
||||
|
||||
def test_log_step_metadata_timestamp(self):
|
||||
log_result = self.agent.log_step_metadata(
|
||||
1, "prompt", "response"
|
||||
)
|
||||
self.assertIn("timestamp", log_result)
|
||||
|
||||
def test_token_counting_integration(self):
|
||||
self.mock_tokenizer.count_tokens.side_effect = [150, 250]
|
||||
log_result = self.agent.log_step_metadata(
|
||||
1, "prompt", "response"
|
||||
)
|
||||
|
||||
self.assertEqual(log_result["tokens"]["total"], 400)
|
||||
|
||||
def test_agent_output_updating(self):
|
||||
initial_total_tokens = sum(
|
||||
step["tokens"]["total"]
|
||||
for step in self.agent.agent_output.steps
|
||||
)
|
||||
self.agent.log_step_metadata(1, "prompt", "response")
|
||||
|
||||
final_total_tokens = sum(
|
||||
step["tokens"]["total"]
|
||||
for step in self.agent.agent_output.steps
|
||||
)
|
||||
self.assertEqual(
|
||||
final_total_tokens - initial_total_tokens, 200
|
||||
)
|
||||
self.assertEqual(len(self.agent.agent_output.steps), 1)
|
||||
|
||||
|
||||
class TestAgentLoggingIntegration(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.agent = Agent(agent_name="test-agent")
|
||||
|
||||
def test_full_logging_cycle(self):
|
||||
task = "Test task"
|
||||
max_loops = 1
|
||||
|
||||
result = self.agent._run(task, max_loops=max_loops)
|
||||
|
||||
self.assertIsInstance(result, dict)
|
||||
self.assertIn("steps", result)
|
||||
self.assertIsInstance(result["steps"], list)
|
||||
self.assertEqual(len(result["steps"]), max_loops)
|
||||
|
||||
if result["steps"]:
|
||||
step = result["steps"][0]
|
||||
self.assertIn("step_id", step)
|
||||
self.assertIn("timestamp", step)
|
||||
self.assertIn("task", step)
|
||||
self.assertIn("response", step)
|
||||
self.assertEqual(step["task"], task)
|
||||
self.assertEqual(step["response"], "Response for loop 1")
|
||||
|
||||
self.assertTrue(len(self.agent.agent_output.steps) > 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -1,267 +0,0 @@
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
from swarms import create_agents_from_yaml
|
||||
import os
|
||||
|
||||
|
||||
class TestCreateAgentsFromYaml(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
# Mock the environment variable for API key
|
||||
os.environ["OPENAI_API_KEY"] = "fake-api-key"
|
||||
|
||||
# Mock agent configuration YAML content
|
||||
self.valid_yaml_content = """
|
||||
agents:
|
||||
- agent_name: "Financial-Analysis-Agent"
|
||||
model:
|
||||
openai_api_key: "fake-api-key"
|
||||
model_name: "gpt-4o-mini"
|
||||
temperature: 0.1
|
||||
max_tokens: 2000
|
||||
system_prompt: "financial_agent_sys_prompt"
|
||||
max_loops: 1
|
||||
autosave: true
|
||||
dashboard: false
|
||||
verbose: true
|
||||
dynamic_temperature_enabled: true
|
||||
saved_state_path: "finance_agent.json"
|
||||
user_name: "swarms_corp"
|
||||
retry_attempts: 1
|
||||
context_length: 200000
|
||||
return_step_meta: false
|
||||
output_type: "str"
|
||||
task: "How can I establish a ROTH IRA to buy stocks and get a tax break?"
|
||||
|
||||
- agent_name: "Stock-Analysis-Agent"
|
||||
model:
|
||||
openai_api_key: "fake-api-key"
|
||||
model_name: "gpt-4o-mini"
|
||||
temperature: 0.2
|
||||
max_tokens: 1500
|
||||
system_prompt: "stock_agent_sys_prompt"
|
||||
max_loops: 2
|
||||
autosave: true
|
||||
dashboard: false
|
||||
verbose: true
|
||||
dynamic_temperature_enabled: false
|
||||
saved_state_path: "stock_agent.json"
|
||||
user_name: "stock_user"
|
||||
retry_attempts: 3
|
||||
context_length: 150000
|
||||
return_step_meta: true
|
||||
output_type: "json"
|
||||
task: "What is the best strategy for long-term stock investment?"
|
||||
"""
|
||||
|
||||
@patch(
|
||||
"builtins.open",
|
||||
new_callable=unittest.mock.mock_open,
|
||||
read_data="",
|
||||
)
|
||||
@patch("yaml.safe_load")
|
||||
def test_create_agents_return_agents(
|
||||
self, mock_safe_load, mock_open
|
||||
):
|
||||
# Mock YAML content parsing
|
||||
mock_safe_load.return_value = {
|
||||
"agents": [
|
||||
{
|
||||
"agent_name": "Financial-Analysis-Agent",
|
||||
"model": {
|
||||
"openai_api_key": "fake-api-key",
|
||||
"model_name": "gpt-4o-mini",
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 2000,
|
||||
},
|
||||
"system_prompt": "financial_agent_sys_prompt",
|
||||
"max_loops": 1,
|
||||
"autosave": True,
|
||||
"dashboard": False,
|
||||
"verbose": True,
|
||||
"dynamic_temperature_enabled": True,
|
||||
"saved_state_path": "finance_agent.json",
|
||||
"user_name": "swarms_corp",
|
||||
"retry_attempts": 1,
|
||||
"context_length": 200000,
|
||||
"return_step_meta": False,
|
||||
"output_type": "str",
|
||||
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# Test if agents are returned correctly
|
||||
agents = create_agents_from_yaml(
|
||||
"fake_yaml_path.yaml", return_type="agents"
|
||||
)
|
||||
self.assertEqual(len(agents), 1)
|
||||
self.assertEqual(
|
||||
agents[0].agent_name, "Financial-Analysis-Agent"
|
||||
)
|
||||
|
||||
@patch(
|
||||
"builtins.open",
|
||||
new_callable=unittest.mock.mock_open,
|
||||
read_data="",
|
||||
)
|
||||
@patch("yaml.safe_load")
|
||||
@patch(
|
||||
"swarms.Agent.run", return_value="Task completed successfully"
|
||||
)
|
||||
def test_create_agents_return_tasks(
|
||||
self, mock_agent_run, mock_safe_load, mock_open
|
||||
):
|
||||
# Mock YAML content parsing
|
||||
mock_safe_load.return_value = {
|
||||
"agents": [
|
||||
{
|
||||
"agent_name": "Financial-Analysis-Agent",
|
||||
"model": {
|
||||
"openai_api_key": "fake-api-key",
|
||||
"model_name": "gpt-4o-mini",
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 2000,
|
||||
},
|
||||
"system_prompt": "financial_agent_sys_prompt",
|
||||
"max_loops": 1,
|
||||
"autosave": True,
|
||||
"dashboard": False,
|
||||
"verbose": True,
|
||||
"dynamic_temperature_enabled": True,
|
||||
"saved_state_path": "finance_agent.json",
|
||||
"user_name": "swarms_corp",
|
||||
"retry_attempts": 1,
|
||||
"context_length": 200000,
|
||||
"return_step_meta": False,
|
||||
"output_type": "str",
|
||||
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# Test if tasks are executed and results are returned
|
||||
task_results = create_agents_from_yaml(
|
||||
"fake_yaml_path.yaml", return_type="tasks"
|
||||
)
|
||||
self.assertEqual(len(task_results), 1)
|
||||
self.assertEqual(
|
||||
task_results[0]["agent_name"], "Financial-Analysis-Agent"
|
||||
)
|
||||
self.assertIsNotNone(task_results[0]["output"])
|
||||
|
||||
@patch(
|
||||
"builtins.open",
|
||||
new_callable=unittest.mock.mock_open,
|
||||
read_data="",
|
||||
)
|
||||
@patch("yaml.safe_load")
|
||||
def test_create_agents_return_both(
|
||||
self, mock_safe_load, mock_open
|
||||
):
|
||||
# Mock YAML content parsing
|
||||
mock_safe_load.return_value = {
|
||||
"agents": [
|
||||
{
|
||||
"agent_name": "Financial-Analysis-Agent",
|
||||
"model": {
|
||||
"openai_api_key": "fake-api-key",
|
||||
"model_name": "gpt-4o-mini",
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 2000,
|
||||
},
|
||||
"system_prompt": "financial_agent_sys_prompt",
|
||||
"max_loops": 1,
|
||||
"autosave": True,
|
||||
"dashboard": False,
|
||||
"verbose": True,
|
||||
"dynamic_temperature_enabled": True,
|
||||
"saved_state_path": "finance_agent.json",
|
||||
"user_name": "swarms_corp",
|
||||
"retry_attempts": 1,
|
||||
"context_length": 200000,
|
||||
"return_step_meta": False,
|
||||
"output_type": "str",
|
||||
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# Test if both agents and tasks are returned
|
||||
agents, task_results = create_agents_from_yaml(
|
||||
"fake_yaml_path.yaml", return_type="both"
|
||||
)
|
||||
self.assertEqual(len(agents), 1)
|
||||
self.assertEqual(len(task_results), 1)
|
||||
self.assertEqual(
|
||||
agents[0].agent_name, "Financial-Analysis-Agent"
|
||||
)
|
||||
self.assertIsNotNone(task_results[0]["output"])
|
||||
|
||||
@patch(
|
||||
"builtins.open",
|
||||
new_callable=unittest.mock.mock_open,
|
||||
read_data="",
|
||||
)
|
||||
@patch("yaml.safe_load")
|
||||
def test_missing_agents_in_yaml(self, mock_safe_load, mock_open):
|
||||
# Mock YAML content with missing "agents" key
|
||||
mock_safe_load.return_value = {}
|
||||
|
||||
# Test if the function raises an error for missing "agents" key
|
||||
with self.assertRaises(ValueError) as context:
|
||||
create_agents_from_yaml(
|
||||
"fake_yaml_path.yaml", return_type="agents"
|
||||
)
|
||||
self.assertTrue(
|
||||
"The YAML configuration does not contain 'agents'."
|
||||
in str(context.exception)
|
||||
)
|
||||
|
||||
@patch(
|
||||
"builtins.open",
|
||||
new_callable=unittest.mock.mock_open,
|
||||
read_data="",
|
||||
)
|
||||
@patch("yaml.safe_load")
|
||||
def test_invalid_return_type(self, mock_safe_load, mock_open):
|
||||
# Mock YAML content parsing
|
||||
mock_safe_load.return_value = {
|
||||
"agents": [
|
||||
{
|
||||
"agent_name": "Financial-Analysis-Agent",
|
||||
"model": {
|
||||
"openai_api_key": "fake-api-key",
|
||||
"model_name": "gpt-4o-mini",
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 2000,
|
||||
},
|
||||
"system_prompt": "financial_agent_sys_prompt",
|
||||
"max_loops": 1,
|
||||
"autosave": True,
|
||||
"dashboard": False,
|
||||
"verbose": True,
|
||||
"dynamic_temperature_enabled": True,
|
||||
"saved_state_path": "finance_agent.json",
|
||||
"user_name": "swarms_corp",
|
||||
"retry_attempts": 1,
|
||||
"context_length": 200000,
|
||||
"return_step_meta": False,
|
||||
"output_type": "str",
|
||||
"task": "How can I establish a ROTH IRA to buy stocks and get a tax break?",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# Test if an error is raised for invalid return_type
|
||||
with self.assertRaises(ValueError) as context:
|
||||
create_agents_from_yaml(
|
||||
"fake_yaml_path.yaml", return_type="invalid_type"
|
||||
)
|
||||
self.assertTrue(
|
||||
"Invalid return_type" in str(context.exception)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -1,190 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify the LiteLLM initialization fix for combined parameters.
|
||||
This test ensures that llm_args, tools_list_dictionary, and MCP tools can be used together.
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
|
||||
from swarms import Agent
|
||||
|
||||
|
||||
def test_combined_llm_args():
|
||||
"""Test that llm_args, tools_list_dictionary, and MCP tools can be combined."""
|
||||
|
||||
# Mock tools list dictionary
|
||||
tools_list = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "test_function",
|
||||
"description": "A test function",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"test_param": {
|
||||
"type": "string",
|
||||
"description": "A test parameter",
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
# Mock llm_args with Azure OpenAI specific parameters
|
||||
llm_args = {
|
||||
"api_version": "2024-02-15-preview",
|
||||
"base_url": "https://your-resource.openai.azure.com/",
|
||||
"api_key": "your-api-key",
|
||||
}
|
||||
|
||||
try:
|
||||
# Test 1: Only llm_args
|
||||
print("Testing Agent with only llm_args...")
|
||||
Agent(
|
||||
agent_name="test-agent-1",
|
||||
model_name="gpt-4o-mini",
|
||||
llm_args=llm_args,
|
||||
)
|
||||
print("✓ Agent with only llm_args created successfully")
|
||||
|
||||
# Test 2: Only tools_list_dictionary
|
||||
print("Testing Agent with only tools_list_dictionary...")
|
||||
Agent(
|
||||
agent_name="test-agent-2",
|
||||
model_name="gpt-4o-mini",
|
||||
tools_list_dictionary=tools_list,
|
||||
)
|
||||
print(
|
||||
"✓ Agent with only tools_list_dictionary created successfully"
|
||||
)
|
||||
|
||||
# Test 3: Combined llm_args and tools_list_dictionary
|
||||
print(
|
||||
"Testing Agent with combined llm_args and tools_list_dictionary..."
|
||||
)
|
||||
agent3 = Agent(
|
||||
agent_name="test-agent-3",
|
||||
model_name="gpt-4o-mini",
|
||||
llm_args=llm_args,
|
||||
tools_list_dictionary=tools_list,
|
||||
)
|
||||
print(
|
||||
"✓ Agent with combined llm_args and tools_list_dictionary created successfully"
|
||||
)
|
||||
|
||||
# Test 4: Verify that the LLM instance has the correct configuration
|
||||
print("Verifying LLM configuration...")
|
||||
|
||||
# Check that agent3 has both llm_args and tools configured
|
||||
assert agent3.llm_args == llm_args, "llm_args not preserved"
|
||||
assert (
|
||||
agent3.tools_list_dictionary == tools_list
|
||||
), "tools_list_dictionary not preserved"
|
||||
|
||||
# Check that the LLM instance was created
|
||||
assert agent3.llm is not None, "LLM instance not created"
|
||||
|
||||
print("✓ LLM configuration verified successfully")
|
||||
|
||||
# Test 5: Test that the LLM can be called (without actually making API calls)
|
||||
print("Testing LLM call preparation...")
|
||||
try:
|
||||
# This should not fail due to configuration issues
|
||||
# We're not actually calling the API, just testing the setup
|
||||
print("✓ LLM call preparation successful")
|
||||
except Exception as e:
|
||||
print(f"✗ LLM call preparation failed: {e}")
|
||||
return False
|
||||
|
||||
print(
|
||||
"\n🎉 All tests passed! The LiteLLM initialization fix is working correctly."
|
||||
)
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Test failed: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
|
||||
def test_azure_openai_example():
|
||||
"""Test the Azure OpenAI example with api_version parameter."""
|
||||
|
||||
print("\nTesting Azure OpenAI example with api_version...")
|
||||
|
||||
try:
|
||||
# Create an agent with Azure OpenAI configuration
|
||||
agent = Agent(
|
||||
agent_name="azure-test-agent",
|
||||
model_name="azure/gpt-4o",
|
||||
llm_args={
|
||||
"api_version": "2024-02-15-preview",
|
||||
"base_url": "https://your-resource.openai.azure.com/",
|
||||
"api_key": "your-api-key",
|
||||
},
|
||||
tools_list_dictionary=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather information",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state",
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
print(
|
||||
"✓ Azure OpenAI agent with combined parameters created successfully"
|
||||
)
|
||||
|
||||
# Verify configuration
|
||||
assert agent.llm_args is not None, "llm_args not set"
|
||||
assert (
|
||||
"api_version" in agent.llm_args
|
||||
), "api_version not in llm_args"
|
||||
assert (
|
||||
agent.tools_list_dictionary is not None
|
||||
), "tools_list_dictionary not set"
|
||||
assert (
|
||||
len(agent.tools_list_dictionary) > 0
|
||||
), "tools_list_dictionary is empty"
|
||||
|
||||
print("✓ Azure OpenAI configuration verified")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Azure OpenAI test failed: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("🧪 Testing LiteLLM initialization fix...")
|
||||
|
||||
success1 = test_combined_llm_args()
|
||||
success2 = test_azure_openai_example()
|
||||
|
||||
if success1 and success2:
|
||||
print("\n✅ All tests passed! The fix is working correctly.")
|
||||
sys.exit(0)
|
||||
else:
|
||||
print(
|
||||
"\n❌ Some tests failed. Please check the implementation."
|
||||
)
|
||||
sys.exit(1)
|
||||
@ -1,73 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify that the llm_handling method properly handles args and kwargs.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add the swarms directory to the path
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "swarms"))
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
|
||||
|
||||
def test_llm_handling_args_kwargs():
|
||||
"""Test that llm_handling properly handles both args and kwargs."""
|
||||
|
||||
# Create an agent instance
|
||||
agent = Agent(
|
||||
agent_name="test-agent",
|
||||
model_name="gpt-4o-mini",
|
||||
temperature=0.7,
|
||||
max_tokens=1000,
|
||||
)
|
||||
|
||||
# Test 1: Call llm_handling with kwargs
|
||||
print("Test 1: Testing kwargs handling...")
|
||||
try:
|
||||
# This should work and add the kwargs to additional_args
|
||||
agent.llm_handling(top_p=0.9, frequency_penalty=0.1)
|
||||
print("✓ kwargs handling works")
|
||||
except Exception as e:
|
||||
print(f"✗ kwargs handling failed: {e}")
|
||||
|
||||
# Test 2: Call llm_handling with args (dictionary)
|
||||
print("\nTest 2: Testing args handling with dictionary...")
|
||||
try:
|
||||
# This should merge the dictionary into additional_args
|
||||
additional_config = {
|
||||
"presence_penalty": 0.2,
|
||||
"logit_bias": {"123": 1},
|
||||
}
|
||||
agent.llm_handling(additional_config)
|
||||
print("✓ args handling with dictionary works")
|
||||
except Exception as e:
|
||||
print(f"✗ args handling with dictionary failed: {e}")
|
||||
|
||||
# Test 3: Call llm_handling with both args and kwargs
|
||||
print("\nTest 3: Testing both args and kwargs...")
|
||||
try:
|
||||
# This should handle both
|
||||
additional_config = {"presence_penalty": 0.3}
|
||||
agent.llm_handling(
|
||||
additional_config, top_p=0.8, frequency_penalty=0.2
|
||||
)
|
||||
print("✓ combined args and kwargs handling works")
|
||||
except Exception as e:
|
||||
print(f"✗ combined args and kwargs handling failed: {e}")
|
||||
|
||||
# Test 4: Call llm_handling with non-dictionary args
|
||||
print("\nTest 4: Testing non-dictionary args...")
|
||||
try:
|
||||
# This should store args under 'additional_args' key
|
||||
agent.llm_handling(
|
||||
"some_string", 123, ["list", "of", "items"]
|
||||
)
|
||||
print("✓ non-dictionary args handling works")
|
||||
except Exception as e:
|
||||
print(f"✗ non-dictionary args handling failed: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_llm_handling_args_kwargs()
|
||||
@ -1,230 +0,0 @@
|
||||
from unittest.mock import Mock, patch
|
||||
import pytest
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from swarms import ToolAgent
|
||||
from swarms.agents.exceptions import (
|
||||
ToolExecutionError,
|
||||
ToolNotFoundError,
|
||||
ToolParameterError,
|
||||
)
|
||||
|
||||
|
||||
def test_tool_agent_init():
|
||||
model = Mock(spec=AutoModelForCausalLM)
|
||||
tokenizer = Mock(spec=AutoTokenizer)
|
||||
json_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"age": {"type": "number"},
|
||||
"is_student": {"type": "boolean"},
|
||||
"courses": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
name = "Test Agent"
|
||||
description = "This is a test agent"
|
||||
|
||||
agent = ToolAgent(
|
||||
name, description, model, tokenizer, json_schema
|
||||
)
|
||||
|
||||
assert agent.name == name
|
||||
assert agent.description == description
|
||||
assert agent.model == model
|
||||
assert agent.tokenizer == tokenizer
|
||||
assert agent.json_schema == json_schema
|
||||
|
||||
|
||||
@patch.object(ToolAgent, "run")
|
||||
def test_tool_agent_run(mock_run):
|
||||
model = Mock(spec=AutoModelForCausalLM)
|
||||
tokenizer = Mock(spec=AutoTokenizer)
|
||||
json_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"age": {"type": "number"},
|
||||
"is_student": {"type": "boolean"},
|
||||
"courses": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
name = "Test Agent"
|
||||
description = "This is a test agent"
|
||||
task = (
|
||||
"Generate a person's information based on the following"
|
||||
" schema:"
|
||||
)
|
||||
|
||||
agent = ToolAgent(
|
||||
name, description, model, tokenizer, json_schema
|
||||
)
|
||||
agent.run(task)
|
||||
|
||||
mock_run.assert_called_once_with(task)
|
||||
|
||||
|
||||
def test_tool_agent_init_with_kwargs():
|
||||
model = Mock(spec=AutoModelForCausalLM)
|
||||
tokenizer = Mock(spec=AutoTokenizer)
|
||||
json_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"age": {"type": "number"},
|
||||
"is_student": {"type": "boolean"},
|
||||
"courses": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
name = "Test Agent"
|
||||
description = "This is a test agent"
|
||||
|
||||
kwargs = {
|
||||
"debug": True,
|
||||
"max_array_length": 20,
|
||||
"max_number_tokens": 12,
|
||||
"temperature": 0.5,
|
||||
"max_string_token_length": 20,
|
||||
}
|
||||
|
||||
agent = ToolAgent(
|
||||
name, description, model, tokenizer, json_schema, **kwargs
|
||||
)
|
||||
|
||||
assert agent.name == name
|
||||
assert agent.description == description
|
||||
assert agent.model == model
|
||||
assert agent.tokenizer == tokenizer
|
||||
assert agent.json_schema == json_schema
|
||||
assert agent.debug == kwargs["debug"]
|
||||
assert agent.max_array_length == kwargs["max_array_length"]
|
||||
assert agent.max_number_tokens == kwargs["max_number_tokens"]
|
||||
assert agent.temperature == kwargs["temperature"]
|
||||
assert (
|
||||
agent.max_string_token_length
|
||||
== kwargs["max_string_token_length"]
|
||||
)
|
||||
|
||||
|
||||
def test_tool_agent_initialization():
|
||||
"""Test tool agent initialization with valid parameters."""
|
||||
agent = ToolAgent(
|
||||
model_name="test-model", temperature=0.7, max_tokens=1000
|
||||
)
|
||||
assert agent.model_name == "test-model"
|
||||
assert agent.temperature == 0.7
|
||||
assert agent.max_tokens == 1000
|
||||
assert agent.retry_attempts == 3
|
||||
assert agent.retry_interval == 1.0
|
||||
|
||||
|
||||
def test_tool_agent_initialization_error():
|
||||
"""Test tool agent initialization with invalid model."""
|
||||
with pytest.raises(ToolExecutionError) as exc_info:
|
||||
ToolAgent(model_name="invalid-model")
|
||||
assert "model_initialization" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_tool_validation():
|
||||
"""Test tool parameter validation."""
|
||||
tools_list = [
|
||||
{
|
||||
"name": "test_tool",
|
||||
"parameters": [
|
||||
{"name": "required_param", "required": True},
|
||||
{"name": "optional_param", "required": False},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
agent = ToolAgent(tools_list_dictionary=tools_list)
|
||||
|
||||
# Test missing required parameter
|
||||
with pytest.raises(ToolParameterError) as exc_info:
|
||||
agent._validate_tool("test_tool", {})
|
||||
assert "Missing required parameters" in str(exc_info.value)
|
||||
|
||||
# Test valid parameters
|
||||
agent._validate_tool("test_tool", {"required_param": "value"})
|
||||
|
||||
# Test non-existent tool
|
||||
with pytest.raises(ToolNotFoundError) as exc_info:
|
||||
agent._validate_tool("non_existent_tool", {})
|
||||
assert "Tool 'non_existent_tool' not found" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_retry_mechanism():
|
||||
"""Test retry mechanism for failed operations."""
|
||||
mock_llm = Mock()
|
||||
mock_llm.generate.side_effect = [
|
||||
Exception("First attempt failed"),
|
||||
Exception("Second attempt failed"),
|
||||
Mock(outputs=[Mock(text="Success")]),
|
||||
]
|
||||
|
||||
agent = ToolAgent(model_name="test-model")
|
||||
agent.llm = mock_llm
|
||||
|
||||
# Test successful retry
|
||||
result = agent.run("test task")
|
||||
assert result == "Success"
|
||||
assert mock_llm.generate.call_count == 3
|
||||
|
||||
# Test all retries failing
|
||||
mock_llm.generate.side_effect = Exception("All attempts failed")
|
||||
with pytest.raises(ToolExecutionError) as exc_info:
|
||||
agent.run("test task")
|
||||
assert "All attempts failed" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_batched_execution():
|
||||
"""Test batched execution with error handling."""
|
||||
mock_llm = Mock()
|
||||
mock_llm.generate.side_effect = [
|
||||
Mock(outputs=[Mock(text="Success 1")]),
|
||||
Exception("Task 2 failed"),
|
||||
Mock(outputs=[Mock(text="Success 3")]),
|
||||
]
|
||||
|
||||
agent = ToolAgent(model_name="test-model")
|
||||
agent.llm = mock_llm
|
||||
|
||||
tasks = ["Task 1", "Task 2", "Task 3"]
|
||||
results = agent.batched_run(tasks)
|
||||
|
||||
assert len(results) == 3
|
||||
assert results[0] == "Success 1"
|
||||
assert "Error" in results[1]
|
||||
assert results[2] == "Success 3"
|
||||
|
||||
|
||||
def test_prompt_preparation():
|
||||
"""Test prompt preparation with and without system prompt."""
|
||||
# Test without system prompt
|
||||
agent = ToolAgent()
|
||||
prompt = agent._prepare_prompt("test task")
|
||||
assert prompt == "User: test task\nAssistant:"
|
||||
|
||||
# Test with system prompt
|
||||
agent = ToolAgent(system_prompt="You are a helpful assistant")
|
||||
prompt = agent._prepare_prompt("test task")
|
||||
assert (
|
||||
prompt
|
||||
== "You are a helpful assistant\n\nUser: test task\nAssistant:"
|
||||
)
|
||||
|
||||
|
||||
def test_tool_execution_error_handling():
|
||||
"""Test error handling during tool execution."""
|
||||
agent = ToolAgent(model_name="test-model")
|
||||
agent.llm = None # Simulate uninitialized LLM
|
||||
|
||||
with pytest.raises(ToolExecutionError) as exc_info:
|
||||
agent.run("test task")
|
||||
assert "LLM not initialized" in str(exc_info.value)
|
||||
|
||||
# Test with invalid parameters
|
||||
with pytest.raises(ToolExecutionError) as exc_info:
|
||||
agent.run("test task", invalid_param="value")
|
||||
assert "Error running task" in str(exc_info.value)
|
||||
@ -1,171 +0,0 @@
|
||||
from time import perf_counter_ns
|
||||
import psutil
|
||||
import os
|
||||
from rich.panel import Panel
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from statistics import mean, median, stdev, variance
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.prompts.finance_agent_sys_prompt import (
|
||||
FINANCIAL_AGENT_SYS_PROMPT,
|
||||
)
|
||||
|
||||
|
||||
def get_memory_stats(memory_readings):
|
||||
"""Calculate memory statistics"""
|
||||
return {
|
||||
"peak": max(memory_readings),
|
||||
"min": min(memory_readings),
|
||||
"mean": mean(memory_readings),
|
||||
"median": median(memory_readings),
|
||||
"stdev": (
|
||||
stdev(memory_readings) if len(memory_readings) > 1 else 0
|
||||
),
|
||||
"variance": (
|
||||
variance(memory_readings)
|
||||
if len(memory_readings) > 1
|
||||
else 0
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_time_stats(times):
|
||||
"""Calculate time statistics"""
|
||||
return {
|
||||
"total": sum(times),
|
||||
"mean": mean(times),
|
||||
"median": median(times),
|
||||
"min": min(times),
|
||||
"max": max(times),
|
||||
"stdev": stdev(times) if len(times) > 1 else 0,
|
||||
"variance": variance(times) if len(times) > 1 else 0,
|
||||
}
|
||||
|
||||
|
||||
def benchmark_multiple_agents(num_agents=100):
|
||||
console = Console()
|
||||
init_times = []
|
||||
memory_readings = []
|
||||
process = psutil.Process(os.getpid())
|
||||
|
||||
# Create benchmark tables
|
||||
time_table = Table(title="Time Statistics")
|
||||
time_table.add_column("Metric", style="cyan")
|
||||
time_table.add_column("Value", style="green")
|
||||
|
||||
memory_table = Table(title="Memory Statistics")
|
||||
memory_table.add_column("Metric", style="cyan")
|
||||
memory_table.add_column("Value", style="green")
|
||||
|
||||
initial_memory = process.memory_info().rss / 1024
|
||||
start_total_time = perf_counter_ns()
|
||||
|
||||
# Initialize agents and measure performance
|
||||
for i in range(num_agents):
|
||||
start_time = perf_counter_ns()
|
||||
|
||||
Agent(
|
||||
agent_name=f"Financial-Analysis-Agent-{i}",
|
||||
agent_description="Personal finance advisor agent",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
max_loops=2,
|
||||
model_name="gpt-4o-mini",
|
||||
dynamic_temperature_enabled=True,
|
||||
interactive=False,
|
||||
)
|
||||
|
||||
init_time = (perf_counter_ns() - start_time) / 1_000_000
|
||||
init_times.append(init_time)
|
||||
|
||||
current_memory = process.memory_info().rss / 1024
|
||||
memory_readings.append(current_memory - initial_memory)
|
||||
|
||||
if (i + 1) % 10 == 0:
|
||||
console.print(
|
||||
f"Created {i + 1} agents...", style="bold blue"
|
||||
)
|
||||
|
||||
total_elapsed_time = (
|
||||
perf_counter_ns() - start_total_time
|
||||
) / 1_000_000
|
||||
|
||||
# Calculate statistics
|
||||
time_stats = get_time_stats(init_times)
|
||||
memory_stats = get_memory_stats(memory_readings)
|
||||
|
||||
# Add time measurements
|
||||
time_table.add_row(
|
||||
"Total Wall Time", f"{total_elapsed_time:.2f} ms"
|
||||
)
|
||||
time_table.add_row(
|
||||
"Total Init Time", f"{time_stats['total']:.2f} ms"
|
||||
)
|
||||
time_table.add_row(
|
||||
"Average Init Time", f"{time_stats['mean']:.2f} ms"
|
||||
)
|
||||
time_table.add_row(
|
||||
"Median Init Time", f"{time_stats['median']:.2f} ms"
|
||||
)
|
||||
time_table.add_row("Fastest Init", f"{time_stats['min']:.2f} ms")
|
||||
time_table.add_row("Slowest Init", f"{time_stats['max']:.2f} ms")
|
||||
time_table.add_row(
|
||||
"Std Deviation", f"{time_stats['stdev']:.2f} ms"
|
||||
)
|
||||
time_table.add_row(
|
||||
"Variance", f"{time_stats['variance']:.4f} ms²"
|
||||
)
|
||||
time_table.add_row(
|
||||
"Throughput",
|
||||
f"{(num_agents/total_elapsed_time) * 1000:.2f} agents/second",
|
||||
)
|
||||
time_table.add_row(
|
||||
"Agents per Minute",
|
||||
f"{(num_agents/total_elapsed_time) * 60000:.0f} agents/minute",
|
||||
)
|
||||
|
||||
# Add memory measurements
|
||||
memory_table.add_row(
|
||||
"Peak Memory Usage", f"{memory_stats['peak']:.2f} KB"
|
||||
)
|
||||
memory_table.add_row(
|
||||
"Minimum Memory Usage", f"{memory_stats['min']:.2f} KB"
|
||||
)
|
||||
memory_table.add_row(
|
||||
"Average Memory Usage", f"{memory_stats['mean']:.2f} KB"
|
||||
)
|
||||
memory_table.add_row(
|
||||
"Median Memory Usage", f"{memory_stats['median']:.2f} KB"
|
||||
)
|
||||
memory_table.add_row(
|
||||
"Memory Std Deviation", f"{memory_stats['stdev']:.2f} KB"
|
||||
)
|
||||
memory_table.add_row(
|
||||
"Memory Variance", f"{memory_stats['variance']:.2f} KB²"
|
||||
)
|
||||
memory_table.add_row(
|
||||
"Avg Memory Per Agent",
|
||||
f"{memory_stats['mean']/num_agents:.2f} KB",
|
||||
)
|
||||
|
||||
# Create and display panels
|
||||
time_panel = Panel(
|
||||
time_table,
|
||||
title="Time Benchmark Results",
|
||||
border_style="blue",
|
||||
padding=(1, 2),
|
||||
)
|
||||
|
||||
memory_panel = Panel(
|
||||
memory_table,
|
||||
title="Memory Benchmark Results",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
)
|
||||
|
||||
console.print(time_panel)
|
||||
console.print("\n")
|
||||
console.print(memory_panel)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_multiple_agents(1000)
|
||||
@ -1,284 +0,0 @@
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import psutil
|
||||
import datetime
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any, Optional
|
||||
from swarms.structs.agent import Agent
|
||||
from loguru import logger
|
||||
|
||||
|
||||
class AgentBenchmark:
|
||||
def __init__(
|
||||
self,
|
||||
num_iterations: int = 5,
|
||||
output_dir: str = "benchmark_results",
|
||||
):
|
||||
self.num_iterations = num_iterations
|
||||
self.output_dir = Path(output_dir)
|
||||
self.output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Use process pool for CPU-bound tasks
|
||||
self.process_pool = concurrent.futures.ProcessPoolExecutor(
|
||||
max_workers=min(os.cpu_count(), 4)
|
||||
)
|
||||
|
||||
# Use thread pool for I/O-bound tasks
|
||||
self.thread_pool = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=min(os.cpu_count() * 2, 8)
|
||||
)
|
||||
|
||||
self.default_queries = [
|
||||
"Conduct an analysis of the best real undervalued ETFs",
|
||||
"What are the top performing tech stocks this quarter?",
|
||||
"Analyze current market trends in renewable energy sector",
|
||||
"Compare Bitcoin and Ethereum investment potential",
|
||||
"Evaluate the risk factors in emerging markets",
|
||||
]
|
||||
|
||||
self.agent = self._initialize_agent()
|
||||
self.process = psutil.Process()
|
||||
|
||||
# Cache for storing repeated query results
|
||||
self._query_cache = {}
|
||||
|
||||
def _initialize_agent(self) -> Agent:
|
||||
return Agent(
|
||||
agent_name="Financial-Analysis-Agent",
|
||||
agent_description="Personal finance advisor agent",
|
||||
# system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
max_loops=1,
|
||||
model_name="gpt-4o-mini",
|
||||
dynamic_temperature_enabled=True,
|
||||
interactive=False,
|
||||
)
|
||||
|
||||
def _get_system_metrics(self) -> Dict[str, float]:
|
||||
# Optimized system metrics collection
|
||||
return {
|
||||
"cpu_percent": self.process.cpu_percent(),
|
||||
"memory_mb": self.process.memory_info().rss / 1024 / 1024,
|
||||
}
|
||||
|
||||
def _calculate_statistics(
|
||||
self, values: List[float]
|
||||
) -> Dict[str, float]:
|
||||
if not values:
|
||||
return {}
|
||||
|
||||
sorted_values = sorted(values)
|
||||
n = len(sorted_values)
|
||||
mean_val = sum(values) / n
|
||||
|
||||
stats = {
|
||||
"mean": mean_val,
|
||||
"median": sorted_values[n // 2],
|
||||
"min": sorted_values[0],
|
||||
"max": sorted_values[-1],
|
||||
}
|
||||
|
||||
# Only calculate stdev if we have enough values
|
||||
if n > 1:
|
||||
stats["std_dev"] = (
|
||||
sum((x - mean_val) ** 2 for x in values) / n
|
||||
) ** 0.5
|
||||
|
||||
return {k: round(v, 3) for k, v in stats.items()}
|
||||
|
||||
async def process_iteration(
|
||||
self, query: str, iteration: int
|
||||
) -> Dict[str, Any]:
|
||||
"""Process a single iteration of a query"""
|
||||
try:
|
||||
# Check cache for repeated queries
|
||||
cache_key = f"{query}_{iteration}"
|
||||
if cache_key in self._query_cache:
|
||||
return self._query_cache[cache_key]
|
||||
|
||||
iteration_start = datetime.datetime.now()
|
||||
pre_metrics = self._get_system_metrics()
|
||||
|
||||
# Run the agent
|
||||
try:
|
||||
self.agent.run(query)
|
||||
success = True
|
||||
except Exception as e:
|
||||
str(e)
|
||||
success = False
|
||||
|
||||
execution_time = (
|
||||
datetime.datetime.now() - iteration_start
|
||||
).total_seconds()
|
||||
post_metrics = self._get_system_metrics()
|
||||
|
||||
result = {
|
||||
"execution_time": execution_time,
|
||||
"success": success,
|
||||
"pre_metrics": pre_metrics,
|
||||
"post_metrics": post_metrics,
|
||||
"iteration_data": {
|
||||
"iteration": iteration + 1,
|
||||
"execution_time": round(execution_time, 3),
|
||||
"success": success,
|
||||
"system_metrics": {
|
||||
"pre": pre_metrics,
|
||||
"post": post_metrics,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
# Cache the result
|
||||
self._query_cache[cache_key] = result
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in iteration {iteration}: {e}")
|
||||
raise
|
||||
|
||||
async def run_benchmark(
|
||||
self, queries: Optional[List[str]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Run the benchmark asynchronously"""
|
||||
queries = queries or self.default_queries
|
||||
benchmark_data = {
|
||||
"metadata": {
|
||||
"timestamp": datetime.datetime.now().isoformat(),
|
||||
"num_iterations": self.num_iterations,
|
||||
"agent_config": {
|
||||
"model_name": self.agent.model_name,
|
||||
"max_loops": self.agent.max_loops,
|
||||
},
|
||||
},
|
||||
"results": {},
|
||||
}
|
||||
|
||||
async def process_query(query: str):
|
||||
query_results = {
|
||||
"execution_times": [],
|
||||
"system_metrics": [],
|
||||
"iterations": [],
|
||||
}
|
||||
|
||||
# Process iterations concurrently
|
||||
tasks = [
|
||||
self.process_iteration(query, i)
|
||||
for i in range(self.num_iterations)
|
||||
]
|
||||
iteration_results = await asyncio.gather(*tasks)
|
||||
|
||||
for result in iteration_results:
|
||||
query_results["execution_times"].append(
|
||||
result["execution_time"]
|
||||
)
|
||||
query_results["system_metrics"].append(
|
||||
result["post_metrics"]
|
||||
)
|
||||
query_results["iterations"].append(
|
||||
result["iteration_data"]
|
||||
)
|
||||
|
||||
# Calculate statistics
|
||||
query_results["statistics"] = {
|
||||
"execution_time": self._calculate_statistics(
|
||||
query_results["execution_times"]
|
||||
),
|
||||
"memory_usage": self._calculate_statistics(
|
||||
[
|
||||
m["memory_mb"]
|
||||
for m in query_results["system_metrics"]
|
||||
]
|
||||
),
|
||||
"cpu_usage": self._calculate_statistics(
|
||||
[
|
||||
m["cpu_percent"]
|
||||
for m in query_results["system_metrics"]
|
||||
]
|
||||
),
|
||||
}
|
||||
|
||||
return query, query_results
|
||||
|
||||
# Execute all queries concurrently
|
||||
query_tasks = [process_query(query) for query in queries]
|
||||
query_results = await asyncio.gather(*query_tasks)
|
||||
|
||||
for query, results in query_results:
|
||||
benchmark_data["results"][query] = results
|
||||
|
||||
return benchmark_data
|
||||
|
||||
def save_results(self, benchmark_data: Dict[str, Any]) -> str:
|
||||
"""Save benchmark results efficiently"""
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = (
|
||||
self.output_dir / f"benchmark_results_{timestamp}.json"
|
||||
)
|
||||
|
||||
# Write results in a single operation
|
||||
with open(filename, "w") as f:
|
||||
json.dump(benchmark_data, f, indent=2)
|
||||
|
||||
logger.info(f"Benchmark results saved to: {filename}")
|
||||
return str(filename)
|
||||
|
||||
def print_summary(self, results: Dict[str, Any]):
|
||||
"""Print a summary of the benchmark results"""
|
||||
print("\n=== Benchmark Summary ===")
|
||||
for query, data in results["results"].items():
|
||||
print(f"\nQuery: {query[:50]}...")
|
||||
stats = data["statistics"]["execution_time"]
|
||||
print(f"Average time: {stats['mean']:.2f}s")
|
||||
print(
|
||||
f"Memory usage (avg): {data['statistics']['memory_usage']['mean']:.1f}MB"
|
||||
)
|
||||
print(
|
||||
f"CPU usage (avg): {data['statistics']['cpu_usage']['mean']:.1f}%"
|
||||
)
|
||||
|
||||
async def run_with_timeout(
|
||||
self, timeout: int = 300
|
||||
) -> Dict[str, Any]:
|
||||
"""Run benchmark with timeout"""
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
self.run_benchmark(), timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(
|
||||
f"Benchmark timed out after {timeout} seconds"
|
||||
)
|
||||
raise
|
||||
|
||||
def cleanup(self):
|
||||
"""Cleanup resources"""
|
||||
self.process_pool.shutdown()
|
||||
self.thread_pool.shutdown()
|
||||
self._query_cache.clear()
|
||||
|
||||
|
||||
async def main():
|
||||
try:
|
||||
# Create and run benchmark
|
||||
benchmark = AgentBenchmark(num_iterations=1)
|
||||
|
||||
# Run benchmark with timeout
|
||||
results = await benchmark.run_with_timeout(timeout=300)
|
||||
|
||||
# Save results
|
||||
benchmark.save_results(results)
|
||||
|
||||
# Print summary
|
||||
benchmark.print_summary(results)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Benchmark failed: {e}")
|
||||
finally:
|
||||
# Cleanup resources
|
||||
benchmark.cleanup()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the async main function
|
||||
asyncio.run(main())
|
||||
@ -1,318 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import psutil
|
||||
import requests
|
||||
from loguru import logger
|
||||
from swarm_models import OpenAIChat
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
|
||||
|
||||
@dataclass
|
||||
class SwarmSystemInfo:
|
||||
"""System information for Swarms issue reports."""
|
||||
|
||||
os_name: str
|
||||
os_version: str
|
||||
python_version: str
|
||||
cpu_usage: float
|
||||
memory_usage: float
|
||||
disk_usage: float
|
||||
swarms_version: str # Added Swarms version tracking
|
||||
cuda_available: bool # Added CUDA availability check
|
||||
gpu_info: Optional[str] # Added GPU information
|
||||
|
||||
|
||||
class SwarmsIssueReporter:
|
||||
"""
|
||||
Production-grade GitHub issue reporter specifically designed for the Swarms library.
|
||||
Automatically creates detailed issues for the https://github.com/kyegomez/swarms repository.
|
||||
|
||||
Features:
|
||||
- Swarms-specific error categorization
|
||||
- Automatic version and dependency tracking
|
||||
- CUDA and GPU information collection
|
||||
- Integration with Swarms logging system
|
||||
- Detailed environment information
|
||||
"""
|
||||
|
||||
REPO_OWNER = "kyegomez"
|
||||
REPO_NAME = "swarms"
|
||||
ISSUE_CATEGORIES = {
|
||||
"agent": ["agent", "automation"],
|
||||
"memory": ["memory", "storage"],
|
||||
"tool": ["tools", "integration"],
|
||||
"llm": ["llm", "model"],
|
||||
"performance": ["performance", "optimization"],
|
||||
"compatibility": ["compatibility", "environment"],
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
github_token: str,
|
||||
rate_limit: int = 10,
|
||||
rate_period: int = 3600,
|
||||
log_file: str = "swarms_issues.log",
|
||||
enable_duplicate_check: bool = True,
|
||||
):
|
||||
"""
|
||||
Initialize the Swarms Issue Reporter.
|
||||
|
||||
Args:
|
||||
github_token (str): GitHub personal access token
|
||||
rate_limit (int): Maximum number of issues to create per rate_period
|
||||
rate_period (int): Time period for rate limiting in seconds
|
||||
log_file (str): Path to log file
|
||||
enable_duplicate_check (bool): Whether to check for duplicate issues
|
||||
"""
|
||||
self.github_token = github_token
|
||||
self.rate_limit = rate_limit
|
||||
self.rate_period = rate_period
|
||||
self.enable_duplicate_check = enable_duplicate_check
|
||||
self.github_token = os.getenv("GITHUB_API_KEY")
|
||||
|
||||
# Initialize logging
|
||||
log_path = os.path.join(os.getcwd(), "logs", log_file)
|
||||
os.makedirs(os.path.dirname(log_path), exist_ok=True)
|
||||
|
||||
# Issue tracking
|
||||
self.issues_created = []
|
||||
self.last_issue_time = datetime.now()
|
||||
|
||||
def _get_swarms_version(self) -> str:
|
||||
"""Get the installed version of Swarms."""
|
||||
try:
|
||||
import swarms
|
||||
|
||||
return swarms.__version__
|
||||
except:
|
||||
return "Unknown"
|
||||
|
||||
def _get_system_info(self) -> SwarmSystemInfo:
|
||||
"""Collect system and Swarms-specific information."""
|
||||
|
||||
return SwarmSystemInfo(
|
||||
os_name=platform.system(),
|
||||
os_version=platform.version(),
|
||||
python_version=sys.version,
|
||||
cpu_usage=psutil.cpu_percent(),
|
||||
memory_usage=psutil.virtual_memory().percent,
|
||||
disk_usage=psutil.disk_usage("/").percent,
|
||||
swarms_version=self._get_swarms_version(),
|
||||
)
|
||||
|
||||
def _categorize_error(
|
||||
self, error: Exception, context: Dict
|
||||
) -> List[str]:
|
||||
"""Categorize the error and return appropriate labels."""
|
||||
error_str = str(error).lower()
|
||||
type(error).__name__
|
||||
|
||||
labels = ["bug", "automated"]
|
||||
|
||||
# Check error message and context for category keywords
|
||||
for (
|
||||
category,
|
||||
category_labels,
|
||||
) in self.ISSUE_CATEGORIES.items():
|
||||
if any(
|
||||
keyword in error_str for keyword in category_labels
|
||||
):
|
||||
labels.extend(category_labels)
|
||||
break
|
||||
|
||||
# Add severity label based on error type
|
||||
if issubclass(type(error), (SystemError, MemoryError)):
|
||||
labels.append("severity:critical")
|
||||
elif issubclass(type(error), (ValueError, TypeError)):
|
||||
labels.append("severity:medium")
|
||||
else:
|
||||
labels.append("severity:low")
|
||||
|
||||
return list(set(labels)) # Remove duplicates
|
||||
|
||||
def _format_swarms_issue_body(
|
||||
self,
|
||||
error: Exception,
|
||||
system_info: SwarmSystemInfo,
|
||||
context: Dict,
|
||||
) -> str:
|
||||
"""Format the issue body with Swarms-specific information."""
|
||||
return f"""
|
||||
## Swarms Error Report
|
||||
- **Error Type**: {type(error).__name__}
|
||||
- **Error Message**: {str(error)}
|
||||
- **Swarms Version**: {system_info.swarms_version}
|
||||
|
||||
## Environment Information
|
||||
- **OS**: {system_info.os_name} {system_info.os_version}
|
||||
- **Python Version**: {system_info.python_version}
|
||||
- **CUDA Available**: {system_info.cuda_available}
|
||||
- **GPU**: {system_info.gpu_info or "N/A"}
|
||||
- **CPU Usage**: {system_info.cpu_usage}%
|
||||
- **Memory Usage**: {system_info.memory_usage}%
|
||||
- **Disk Usage**: {system_info.disk_usage}%
|
||||
|
||||
## Stack Trace
|
||||
{traceback.format_exc()}
|
||||
|
||||
## Context
|
||||
{json.dumps(context, indent=2)}
|
||||
|
||||
## Dependencies
|
||||
{self._get_dependencies_info()}
|
||||
|
||||
## Time of Occurrence
|
||||
{datetime.now().isoformat()}
|
||||
|
||||
---
|
||||
*This issue was automatically generated by SwarmsIssueReporter*
|
||||
"""
|
||||
|
||||
def _get_dependencies_info(self) -> str:
|
||||
"""Get information about installed dependencies."""
|
||||
try:
|
||||
import pkg_resources
|
||||
|
||||
deps = []
|
||||
for dist in pkg_resources.working_set:
|
||||
deps.append(f"- {dist.key} {dist.version}")
|
||||
return "\n".join(deps)
|
||||
except:
|
||||
return "Unable to fetch dependency information"
|
||||
|
||||
# First, add this method to your SwarmsIssueReporter class
|
||||
def _check_rate_limit(self) -> bool:
|
||||
"""Check if we're within rate limits."""
|
||||
now = datetime.now()
|
||||
time_diff = (now - self.last_issue_time).total_seconds()
|
||||
|
||||
if (
|
||||
len(self.issues_created) >= self.rate_limit
|
||||
and time_diff < self.rate_period
|
||||
):
|
||||
logger.warning("Rate limit exceeded for issue creation")
|
||||
return False
|
||||
|
||||
# Clean up old issues from tracking
|
||||
self.issues_created = [
|
||||
time
|
||||
for time in self.issues_created
|
||||
if (now - time).total_seconds() < self.rate_period
|
||||
]
|
||||
|
||||
return True
|
||||
|
||||
def report_swarms_issue(
|
||||
self,
|
||||
error: Exception,
|
||||
agent: Optional[Agent] = None,
|
||||
context: Dict[str, Any] = None,
|
||||
priority: str = "normal",
|
||||
) -> Optional[int]:
|
||||
"""
|
||||
Report a Swarms-specific issue to GitHub.
|
||||
|
||||
Args:
|
||||
error (Exception): The exception to report
|
||||
agent (Optional[Agent]): The Swarms agent instance that encountered the error
|
||||
context (Dict[str, Any]): Additional context about the error
|
||||
priority (str): Issue priority ("low", "normal", "high", "critical")
|
||||
|
||||
Returns:
|
||||
Optional[int]: Issue number if created successfully
|
||||
"""
|
||||
try:
|
||||
if not self._check_rate_limit():
|
||||
logger.warning(
|
||||
"Skipping issue creation due to rate limit"
|
||||
)
|
||||
return None
|
||||
|
||||
# Collect system information
|
||||
system_info = self._get_system_info()
|
||||
|
||||
# Prepare context with agent information if available
|
||||
full_context = context or {}
|
||||
if agent:
|
||||
full_context.update(
|
||||
{
|
||||
"agent_name": agent.agent_name,
|
||||
"agent_description": agent.agent_description,
|
||||
"max_loops": agent.max_loops,
|
||||
"context_length": agent.context_length,
|
||||
}
|
||||
)
|
||||
|
||||
# Create issue title
|
||||
title = f"[{type(error).__name__}] {str(error)[:100]}"
|
||||
if agent:
|
||||
title = f"[Agent: {agent.agent_name}] {title}"
|
||||
|
||||
# Get appropriate labels
|
||||
labels = self._categorize_error(error, full_context)
|
||||
labels.append(f"priority:{priority}")
|
||||
|
||||
# Create the issue
|
||||
url = f"https://api.github.com/repos/{self.REPO_OWNER}/{self.REPO_NAME}/issues"
|
||||
data = {
|
||||
"title": title,
|
||||
"body": self._format_swarms_issue_body(
|
||||
error, system_info, full_context
|
||||
),
|
||||
"labels": labels,
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
url,
|
||||
headers={
|
||||
"Authorization": f"token {self.github_token}"
|
||||
},
|
||||
json=data,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
issue_number = response.json()["number"]
|
||||
logger.info(
|
||||
f"Successfully created Swarms issue #{issue_number}"
|
||||
)
|
||||
|
||||
return issue_number
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating Swarms issue: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
# Setup the reporter with your GitHub token
|
||||
reporter = SwarmsIssueReporter(
|
||||
github_token=os.getenv("GITHUB_API_KEY")
|
||||
)
|
||||
|
||||
|
||||
# Force an error to test the reporter
|
||||
try:
|
||||
# This will raise an error since the input isn't valid
|
||||
# Create an agent that might have issues
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(agent_name="Test-Agent", max_loops=1)
|
||||
|
||||
result = agent.run(None)
|
||||
|
||||
raise ValueError("test")
|
||||
except Exception as e:
|
||||
# Report the issue
|
||||
issue_number = reporter.report_swarms_issue(
|
||||
error=e,
|
||||
agent=agent,
|
||||
context={"task": "test_run"},
|
||||
priority="high",
|
||||
)
|
||||
print(f"Created issue number: {issue_number}")
|
||||
@ -1,180 +0,0 @@
|
||||
import requests
|
||||
import datetime
|
||||
from typing import List, Dict, Tuple
|
||||
from loguru import logger
|
||||
from swarms import Agent
|
||||
from swarm_models import OpenAIChat
|
||||
|
||||
# GitHub API Configurations
|
||||
GITHUB_REPO = "kyegomez/swarms" # Swarms GitHub repository
|
||||
GITHUB_API_URL = f"https://api.github.com/repos/{GITHUB_REPO}/commits"
|
||||
|
||||
|
||||
# Step 1: Fetch the latest commits from GitHub
|
||||
def fetch_latest_commits(
|
||||
repo_url: str, limit: int = 5
|
||||
) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Fetch the latest commits from a public GitHub repository.
|
||||
"""
|
||||
logger.info(
|
||||
f"Fetching the latest {limit} commits from {repo_url}"
|
||||
)
|
||||
try:
|
||||
params = {"per_page": limit}
|
||||
response = requests.get(repo_url, params=params)
|
||||
response.raise_for_status()
|
||||
|
||||
commits = response.json()
|
||||
commit_data = []
|
||||
|
||||
for commit in commits:
|
||||
commit_data.append(
|
||||
{
|
||||
"sha": commit["sha"][:7], # Short commit hash
|
||||
"author": commit["commit"]["author"]["name"],
|
||||
"message": commit["commit"]["message"],
|
||||
"date": commit["commit"]["author"]["date"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.success("Successfully fetched commit data")
|
||||
return commit_data
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching commits: {e}")
|
||||
raise
|
||||
|
||||
|
||||
# Step 2: Format commits and fetch current time
|
||||
def format_commits_with_time(
|
||||
commits: List[Dict[str, str]],
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Format commit data into a readable string and return current time.
|
||||
"""
|
||||
current_time = datetime.datetime.now().strftime(
|
||||
"%Y-%m-%d %H:%M:%S"
|
||||
)
|
||||
logger.info(f"Formatting commits at {current_time}")
|
||||
|
||||
commit_summary = "\n".join(
|
||||
[
|
||||
f"- `{commit['sha']}` by {commit['author']} on {commit['date']}: {commit['message']}"
|
||||
for commit in commits
|
||||
]
|
||||
)
|
||||
|
||||
logger.success("Commits formatted successfully")
|
||||
return current_time, commit_summary
|
||||
|
||||
|
||||
# Step 3: Build a dynamic system prompt
|
||||
def build_custom_system_prompt(
|
||||
current_time: str, commit_summary: str
|
||||
) -> str:
|
||||
"""
|
||||
Build a dynamic system prompt with the current time and commit summary.
|
||||
"""
|
||||
logger.info("Building the custom system prompt for the agent")
|
||||
prompt = f"""
|
||||
You are a software analyst tasked with summarizing the latest commits from the Swarms GitHub repository.
|
||||
|
||||
The current time is **{current_time}**.
|
||||
|
||||
Here are the latest commits:
|
||||
{commit_summary}
|
||||
|
||||
**Your task**:
|
||||
1. Summarize the changes into a clear and concise table in **markdown format**.
|
||||
2. Highlight the key improvements and fixes.
|
||||
3. End your output with the token `<DONE>`.
|
||||
|
||||
Make sure the table includes the following columns: Commit SHA, Author, Date, and Commit Message.
|
||||
"""
|
||||
logger.success("System prompt created successfully")
|
||||
return prompt
|
||||
|
||||
|
||||
# Step 4: Initialize the Agent
|
||||
def initialize_agent() -> Agent:
|
||||
"""
|
||||
Initialize the Swarms agent with OpenAI model.
|
||||
"""
|
||||
logger.info("Initializing the agent with GPT-4o")
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
|
||||
agent = Agent(
|
||||
agent_name="Commit-Summarization-Agent",
|
||||
agent_description="Fetch and summarize GitHub commits for Swarms repository.",
|
||||
system_prompt="", # Will set dynamically
|
||||
max_loops=1,
|
||||
llm=model,
|
||||
dynamic_temperature_enabled=True,
|
||||
user_name="Kye",
|
||||
retry_attempts=3,
|
||||
context_length=8192,
|
||||
return_step_meta=False,
|
||||
output_type="str",
|
||||
auto_generate_prompt=False,
|
||||
max_tokens=4000,
|
||||
stopping_token="<DONE>",
|
||||
interactive=False,
|
||||
)
|
||||
logger.success("Agent initialized successfully")
|
||||
return agent
|
||||
|
||||
|
||||
# Step 5: Run the Agent with Data
|
||||
def summarize_commits_with_agent(agent: Agent, prompt: str) -> str:
|
||||
"""
|
||||
Pass the system prompt to the agent and fetch the result.
|
||||
"""
|
||||
logger.info("Sending data to the agent for summarization")
|
||||
try:
|
||||
result = agent.run(
|
||||
f"{prompt}",
|
||||
all_cores=True,
|
||||
)
|
||||
logger.success("Agent completed the summarization task")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Agent encountered an error: {e}")
|
||||
raise
|
||||
|
||||
|
||||
# Main Execution
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
logger.info("Starting commit summarization process")
|
||||
|
||||
# Fetch latest commits
|
||||
latest_commits = fetch_latest_commits(GITHUB_API_URL, limit=5)
|
||||
|
||||
# Format commits and get current time
|
||||
current_time, commit_summary = format_commits_with_time(
|
||||
latest_commits
|
||||
)
|
||||
|
||||
# Build the custom system prompt
|
||||
custom_system_prompt = build_custom_system_prompt(
|
||||
current_time, commit_summary
|
||||
)
|
||||
|
||||
# Initialize agent
|
||||
agent = initialize_agent()
|
||||
|
||||
# Set the dynamic system prompt
|
||||
agent.system_prompt = custom_system_prompt
|
||||
|
||||
# Run the agent and summarize commits
|
||||
result = summarize_commits_with_agent(
|
||||
agent, custom_system_prompt
|
||||
)
|
||||
|
||||
# Print the result
|
||||
print("### Commit Summary in Markdown:")
|
||||
print(result)
|
||||
|
||||
except Exception as e:
|
||||
logger.critical(f"Process failed: {e}")
|
||||
@ -1,46 +0,0 @@
|
||||
import os
|
||||
import uuid
|
||||
from swarms import Agent
|
||||
from swarm_models import OpenAIChat
|
||||
from swarms.prompts.finance_agent_sys_prompt import (
|
||||
FINANCIAL_AGENT_SYS_PROMPT,
|
||||
)
|
||||
import time
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
|
||||
# Get the OpenAI API key from the environment variable
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
|
||||
# Create an instance of the OpenAIChat class
|
||||
model = OpenAIChat(
|
||||
api_key=api_key, model_name="gpt-4o-mini", temperature=0.1
|
||||
)
|
||||
|
||||
|
||||
agent = Agent(
|
||||
agent_name=f"{uuid.uuid4().hex}",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
dashboard=False,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path=f"{uuid.uuid4().hex}",
|
||||
user_name="swarms_corp",
|
||||
retry_attempts=1,
|
||||
context_length=3000,
|
||||
return_step_meta=False,
|
||||
)
|
||||
|
||||
out = agent.run(
|
||||
"How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria"
|
||||
)
|
||||
print(out)
|
||||
|
||||
end_time = time.time()
|
||||
|
||||
print(f"Execution time: {end_time - start_time} seconds")
|
||||
# Execution time: 9.922541856765747 seconds for the whole script
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 175 KiB |
|
Before Width: | Height: | Size: 178 KiB |
|
Before Width: | Height: | Size: 130 KiB |
|
Before Width: | Height: | Size: 75 KiB |
|
Before Width: | Height: | Size: 66 KiB |
|
|
Before Width: | Height: | Size: 15 KiB |
|
Before Width: | Height: | Size: 19 MiB |
|
Before Width: | Height: | Size: 492 KiB |
@ -1,4 +1,15 @@
|
||||
swarms
|
||||
pytest
|
||||
matplotlib
|
||||
loguru
|
||||
loguru
|
||||
psutil
|
||||
pyyaml
|
||||
python-dotenv
|
||||
rich
|
||||
pydantic
|
||||
numpy
|
||||
pandas
|
||||
openpyxl
|
||||
seaborn
|
||||
requests
|
||||
swarms-memory
|
||||
@ -1,64 +0,0 @@
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.multi_agent_router import MultiAgentRouter
|
||||
|
||||
# Example usage:
|
||||
agents = [
|
||||
Agent(
|
||||
agent_name="ResearchAgent",
|
||||
agent_description="Specializes in researching topics and providing detailed, factual information",
|
||||
system_prompt="You are a research specialist. Provide detailed, well-researched information about any topic, citing sources when possible.",
|
||||
max_loops=1,
|
||||
),
|
||||
Agent(
|
||||
agent_name="CodeExpertAgent",
|
||||
agent_description="Expert in writing, reviewing, and explaining code across multiple programming languages",
|
||||
system_prompt="You are a coding expert. Write, review, and explain code with a focus on best practices and clean code principles.",
|
||||
max_loops=1,
|
||||
),
|
||||
Agent(
|
||||
agent_name="WritingAgent",
|
||||
agent_description="Skilled in creative and technical writing, content creation, and editing",
|
||||
system_prompt="You are a writing specialist. Create, edit, and improve written content while maintaining appropriate tone and style.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
models_to_test = [
|
||||
"gpt-4.1",
|
||||
"gpt-4o",
|
||||
"gpt-5-mini",
|
||||
"o4-mini",
|
||||
"o3",
|
||||
"claude-opus-4-20250514",
|
||||
"claude-sonnet-4-20250514",
|
||||
"claude-3-7-sonnet-20250219",
|
||||
"gemini/gemini-2.5-flash",
|
||||
"gemini/gemini-2.5-pro",
|
||||
]
|
||||
|
||||
task = "Use all the agents available to you to remake the Fibonacci function in Python, providing both an explanation and code."
|
||||
|
||||
model_logs = []
|
||||
|
||||
for model_name in models_to_test:
|
||||
print(f"\n--- Testing model: {model_name} ---")
|
||||
router_execute = MultiAgentRouter(
|
||||
agents=agents,
|
||||
temperature=0.5,
|
||||
model=model_name,
|
||||
)
|
||||
try:
|
||||
result = router_execute.run(task)
|
||||
print(f"Run completed successfully for {model_name}")
|
||||
model_logs.append(
|
||||
{"model": model_name, "status": "✅ Success"}
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"An error occurred for {model_name}")
|
||||
model_logs.append(
|
||||
{"model": model_name, "status": f"❌ Error: {e}"}
|
||||
)
|
||||
|
||||
print("\n===== Model Run Summary =====")
|
||||
for log in model_logs:
|
||||
print(f"{log['model']}: {log['status']}")
|
||||
File diff suppressed because it is too large
Load Diff
@ -1,600 +0,0 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import yaml
|
||||
from swarm_models import OpenAIChat
|
||||
|
||||
from swarms import Agent
|
||||
|
||||
|
||||
def test_basic_agent_functionality():
|
||||
"""Test basic agent initialization and simple task execution"""
|
||||
print("\nTesting basic agent functionality...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(agent_name="Test-Agent", llm=model, max_loops=1)
|
||||
|
||||
response = agent.run("What is 2+2?")
|
||||
assert response is not None, "Agent response should not be None"
|
||||
|
||||
# Test agent properties
|
||||
assert (
|
||||
agent.agent_name == "Test-Agent"
|
||||
), "Agent name not set correctly"
|
||||
assert agent.max_loops == 1, "Max loops not set correctly"
|
||||
assert agent.llm is not None, "LLM not initialized"
|
||||
|
||||
print("✓ Basic agent functionality test passed")
|
||||
|
||||
|
||||
def test_memory_management():
|
||||
"""Test agent memory management functionality"""
|
||||
print("\nTesting memory management...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Memory-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
context_length=8192,
|
||||
)
|
||||
|
||||
# Test adding to memory
|
||||
agent.add_memory("Test memory entry")
|
||||
assert (
|
||||
"Test memory entry"
|
||||
in agent.short_memory.return_history_as_string()
|
||||
)
|
||||
|
||||
# Test memory query
|
||||
agent.memory_query("Test query")
|
||||
|
||||
# Test token counting
|
||||
tokens = agent.check_available_tokens()
|
||||
assert isinstance(tokens, int), "Token count should be an integer"
|
||||
|
||||
print("✓ Memory management test passed")
|
||||
|
||||
|
||||
def test_agent_output_formats():
|
||||
"""Test all available output formats"""
|
||||
print("\nTesting all output formats...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
test_task = "Say hello!"
|
||||
|
||||
output_types = {
|
||||
"str": str,
|
||||
"string": str,
|
||||
"list": str, # JSON string containing list
|
||||
"json": str, # JSON string
|
||||
"dict": dict,
|
||||
"yaml": str,
|
||||
}
|
||||
|
||||
for output_type, expected_type in output_types.items():
|
||||
agent = Agent(
|
||||
agent_name=f"{output_type.capitalize()}-Output-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
output_type=output_type,
|
||||
)
|
||||
|
||||
response = agent.run(test_task)
|
||||
assert (
|
||||
response is not None
|
||||
), f"{output_type} output should not be None"
|
||||
|
||||
if output_type == "yaml":
|
||||
# Verify YAML can be parsed
|
||||
try:
|
||||
yaml.safe_load(response)
|
||||
print(f"✓ {output_type} output valid")
|
||||
except yaml.YAMLError:
|
||||
assert False, f"Invalid YAML output for {output_type}"
|
||||
elif output_type in ["json", "list"]:
|
||||
# Verify JSON can be parsed
|
||||
try:
|
||||
json.loads(response)
|
||||
print(f"✓ {output_type} output valid")
|
||||
except json.JSONDecodeError:
|
||||
assert False, f"Invalid JSON output for {output_type}"
|
||||
|
||||
print("✓ Output formats test passed")
|
||||
|
||||
|
||||
def test_agent_state_management():
|
||||
"""Test comprehensive state management functionality"""
|
||||
print("\nTesting state management...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
|
||||
# Create temporary directory for test files
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
state_path = os.path.join(temp_dir, "agent_state.json")
|
||||
|
||||
# Create agent with initial state
|
||||
agent1 = Agent(
|
||||
agent_name="State-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
saved_state_path=state_path,
|
||||
)
|
||||
|
||||
# Add some data to the agent
|
||||
agent1.run("Remember this: Test message 1")
|
||||
agent1.add_memory("Test message 2")
|
||||
|
||||
# Save state
|
||||
agent1.save()
|
||||
assert os.path.exists(state_path), "State file not created"
|
||||
|
||||
# Create new agent and load state
|
||||
agent2 = Agent(
|
||||
agent_name="State-Test-Agent", llm=model, max_loops=1
|
||||
)
|
||||
agent2.load(state_path)
|
||||
|
||||
# Verify state loaded correctly
|
||||
history2 = agent2.short_memory.return_history_as_string()
|
||||
assert (
|
||||
"Test message 1" in history2
|
||||
), "State not loaded correctly"
|
||||
assert (
|
||||
"Test message 2" in history2
|
||||
), "Memory not loaded correctly"
|
||||
|
||||
# Test autosave functionality
|
||||
agent3 = Agent(
|
||||
agent_name="Autosave-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
saved_state_path=os.path.join(
|
||||
temp_dir, "autosave_state.json"
|
||||
),
|
||||
autosave=True,
|
||||
)
|
||||
|
||||
agent3.run("Test autosave")
|
||||
time.sleep(2) # Wait for autosave
|
||||
assert os.path.exists(
|
||||
os.path.join(temp_dir, "autosave_state.json")
|
||||
), "Autosave file not created"
|
||||
|
||||
print("✓ State management test passed")
|
||||
|
||||
|
||||
def test_agent_tools_and_execution():
|
||||
"""Test agent tool handling and execution"""
|
||||
print("\nTesting tools and execution...")
|
||||
|
||||
def sample_tool(x: int, y: int) -> int:
|
||||
"""Sample tool that adds two numbers"""
|
||||
return x + y
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Tools-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
tools=[sample_tool],
|
||||
)
|
||||
|
||||
# Test adding tools
|
||||
agent.add_tool(lambda x: x * 2)
|
||||
assert len(agent.tools) == 2, "Tool not added correctly"
|
||||
|
||||
# Test removing tools
|
||||
agent.remove_tool(sample_tool)
|
||||
assert len(agent.tools) == 1, "Tool not removed correctly"
|
||||
|
||||
# Test tool execution
|
||||
response = agent.run("Calculate 2 + 2 using the sample tool")
|
||||
assert response is not None, "Tool execution failed"
|
||||
|
||||
print("✓ Tools and execution test passed")
|
||||
|
||||
|
||||
def test_agent_concurrent_execution():
|
||||
"""Test agent concurrent execution capabilities"""
|
||||
print("\nTesting concurrent execution...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Concurrent-Test-Agent", llm=model, max_loops=1
|
||||
)
|
||||
|
||||
# Test bulk run
|
||||
tasks = [
|
||||
{"task": "Count to 3"},
|
||||
{"task": "Say hello"},
|
||||
{"task": "Tell a short joke"},
|
||||
]
|
||||
|
||||
responses = agent.bulk_run(tasks)
|
||||
assert len(responses) == len(tasks), "Not all tasks completed"
|
||||
assert all(
|
||||
response is not None for response in responses
|
||||
), "Some tasks failed"
|
||||
|
||||
# Test concurrent tasks
|
||||
concurrent_responses = agent.run_concurrent_tasks(
|
||||
["Task 1", "Task 2", "Task 3"]
|
||||
)
|
||||
assert (
|
||||
len(concurrent_responses) == 3
|
||||
), "Not all concurrent tasks completed"
|
||||
|
||||
print("✓ Concurrent execution test passed")
|
||||
|
||||
|
||||
def test_agent_error_handling():
|
||||
"""Test agent error handling and recovery"""
|
||||
print("\nTesting error handling...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Error-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
retry_attempts=3,
|
||||
retry_interval=1,
|
||||
)
|
||||
|
||||
# Test invalid tool execution
|
||||
try:
|
||||
agent.parse_and_execute_tools("invalid_json")
|
||||
print("✓ Invalid tool execution handled")
|
||||
except Exception:
|
||||
assert True, "Expected error caught"
|
||||
|
||||
# Test recovery after error
|
||||
response = agent.run("Continue after error")
|
||||
assert response is not None, "Agent failed to recover after error"
|
||||
|
||||
print("✓ Error handling test passed")
|
||||
|
||||
|
||||
def test_agent_configuration():
|
||||
"""Test agent configuration and parameters"""
|
||||
print("\nTesting agent configuration...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Config-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
temperature=0.7,
|
||||
max_tokens=4000,
|
||||
context_length=8192,
|
||||
)
|
||||
|
||||
# Test configuration methods
|
||||
agent.update_system_prompt("New system prompt")
|
||||
agent.update_max_loops(2)
|
||||
agent.update_loop_interval(2)
|
||||
|
||||
# Verify updates
|
||||
assert agent.max_loops == 2, "Max loops not updated"
|
||||
assert agent.loop_interval == 2, "Loop interval not updated"
|
||||
|
||||
# Test configuration export
|
||||
config_dict = agent.to_dict()
|
||||
assert isinstance(
|
||||
config_dict, dict
|
||||
), "Configuration export failed"
|
||||
|
||||
# Test YAML export
|
||||
yaml_config = agent.to_yaml()
|
||||
assert isinstance(yaml_config, str), "YAML export failed"
|
||||
|
||||
print("✓ Configuration test passed")
|
||||
|
||||
|
||||
def test_agent_with_stopping_condition():
|
||||
"""Test agent with custom stopping condition"""
|
||||
print("\nTesting agent with stopping condition...")
|
||||
|
||||
def custom_stopping_condition(response: str) -> bool:
|
||||
return "STOP" in response.upper()
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Stopping-Condition-Agent",
|
||||
llm=model,
|
||||
max_loops=5,
|
||||
stopping_condition=custom_stopping_condition,
|
||||
)
|
||||
|
||||
response = agent.run("Count up until you see the word STOP")
|
||||
assert response is not None, "Stopping condition test failed"
|
||||
print("✓ Stopping condition test passed")
|
||||
|
||||
|
||||
def test_agent_with_retry_mechanism():
|
||||
"""Test agent retry mechanism"""
|
||||
print("\nTesting agent retry mechanism...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Retry-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
retry_attempts=3,
|
||||
retry_interval=1,
|
||||
)
|
||||
|
||||
response = agent.run("Tell me a joke.")
|
||||
assert response is not None, "Retry mechanism test failed"
|
||||
print("✓ Retry mechanism test passed")
|
||||
|
||||
|
||||
def test_bulk_and_filtered_operations():
|
||||
"""Test bulk operations and response filtering"""
|
||||
print("\nTesting bulk and filtered operations...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Bulk-Filter-Test-Agent", llm=model, max_loops=1
|
||||
)
|
||||
|
||||
# Test bulk run
|
||||
bulk_tasks = [
|
||||
{"task": "What is 2+2?"},
|
||||
{"task": "Name a color"},
|
||||
{"task": "Count to 3"},
|
||||
]
|
||||
bulk_responses = agent.bulk_run(bulk_tasks)
|
||||
assert len(bulk_responses) == len(
|
||||
bulk_tasks
|
||||
), "Bulk run should return same number of responses as tasks"
|
||||
|
||||
# Test response filtering
|
||||
agent.add_response_filter("color")
|
||||
filtered_response = agent.filtered_run(
|
||||
"What is your favorite color?"
|
||||
)
|
||||
assert (
|
||||
"[FILTERED]" in filtered_response
|
||||
), "Response filter not applied"
|
||||
|
||||
print("✓ Bulk and filtered operations test passed")
|
||||
|
||||
|
||||
async def test_async_operations():
|
||||
"""Test asynchronous operations"""
|
||||
print("\nTesting async operations...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Async-Test-Agent", llm=model, max_loops=1
|
||||
)
|
||||
|
||||
# Test single async run
|
||||
response = await agent.arun("What is 1+1?")
|
||||
assert response is not None, "Async run failed"
|
||||
|
||||
# Test concurrent async runs
|
||||
tasks = ["Task 1", "Task 2", "Task 3"]
|
||||
responses = await asyncio.gather(
|
||||
*[agent.arun(task) for task in tasks]
|
||||
)
|
||||
assert len(responses) == len(
|
||||
tasks
|
||||
), "Not all async tasks completed"
|
||||
|
||||
print("✓ Async operations test passed")
|
||||
|
||||
|
||||
def test_memory_and_state_persistence():
|
||||
"""Test memory management and state persistence"""
|
||||
print("\nTesting memory and state persistence...")
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
state_path = os.path.join(temp_dir, "test_state.json")
|
||||
|
||||
# Create agent with memory configuration
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent1 = Agent(
|
||||
agent_name="Memory-State-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
saved_state_path=state_path,
|
||||
context_length=8192,
|
||||
autosave=True,
|
||||
)
|
||||
|
||||
# Test memory operations
|
||||
agent1.add_memory("Important fact: The sky is blue")
|
||||
agent1.memory_query("What color is the sky?")
|
||||
|
||||
# Save state
|
||||
agent1.save()
|
||||
|
||||
# Create new agent and load state
|
||||
agent2 = Agent(
|
||||
agent_name="Memory-State-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
)
|
||||
agent2.load(state_path)
|
||||
|
||||
# Verify memory persistence
|
||||
memory_content = (
|
||||
agent2.short_memory.return_history_as_string()
|
||||
)
|
||||
assert (
|
||||
"sky is blue" in memory_content
|
||||
), "Memory not properly persisted"
|
||||
|
||||
print("✓ Memory and state persistence test passed")
|
||||
|
||||
|
||||
def test_sentiment_and_evaluation():
|
||||
"""Test sentiment analysis and response evaluation"""
|
||||
print("\nTesting sentiment analysis and evaluation...")
|
||||
|
||||
def mock_sentiment_analyzer(text):
|
||||
"""Mock sentiment analyzer that returns a score between 0 and 1"""
|
||||
return 0.7 if "positive" in text.lower() else 0.3
|
||||
|
||||
def mock_evaluator(response):
|
||||
"""Mock evaluator that checks response quality"""
|
||||
return "GOOD" if len(response) > 10 else "BAD"
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Sentiment-Eval-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
sentiment_analyzer=mock_sentiment_analyzer,
|
||||
sentiment_threshold=0.5,
|
||||
evaluator=mock_evaluator,
|
||||
)
|
||||
|
||||
# Test sentiment analysis
|
||||
agent.run("Generate a positive message")
|
||||
|
||||
# Test evaluation
|
||||
agent.run("Generate a detailed response")
|
||||
|
||||
print("✓ Sentiment and evaluation test passed")
|
||||
|
||||
|
||||
def test_tool_management():
|
||||
"""Test tool management functionality"""
|
||||
print("\nTesting tool management...")
|
||||
|
||||
def tool1(x: int) -> int:
|
||||
"""Sample tool 1"""
|
||||
return x * 2
|
||||
|
||||
def tool2(x: int) -> int:
|
||||
"""Sample tool 2"""
|
||||
return x + 2
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Tool-Test-Agent",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
tools=[tool1],
|
||||
)
|
||||
|
||||
# Test adding tools
|
||||
agent.add_tool(tool2)
|
||||
assert len(agent.tools) == 2, "Tool not added correctly"
|
||||
|
||||
# Test removing tools
|
||||
agent.remove_tool(tool1)
|
||||
assert len(agent.tools) == 1, "Tool not removed correctly"
|
||||
|
||||
# Test adding multiple tools
|
||||
agent.add_tools([tool1, tool2])
|
||||
assert len(agent.tools) == 3, "Multiple tools not added correctly"
|
||||
|
||||
print("✓ Tool management test passed")
|
||||
|
||||
|
||||
def test_system_prompt_and_configuration():
|
||||
"""Test system prompt and configuration updates"""
|
||||
print("\nTesting system prompt and configuration...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Config-Test-Agent", llm=model, max_loops=1
|
||||
)
|
||||
|
||||
# Test updating system prompt
|
||||
new_prompt = "You are a helpful assistant."
|
||||
agent.update_system_prompt(new_prompt)
|
||||
assert (
|
||||
agent.system_prompt == new_prompt
|
||||
), "System prompt not updated"
|
||||
|
||||
# Test configuration updates
|
||||
agent.update_max_loops(5)
|
||||
assert agent.max_loops == 5, "Max loops not updated"
|
||||
|
||||
agent.update_loop_interval(2)
|
||||
assert agent.loop_interval == 2, "Loop interval not updated"
|
||||
|
||||
# Test configuration export
|
||||
config_dict = agent.to_dict()
|
||||
assert isinstance(
|
||||
config_dict, dict
|
||||
), "Configuration export failed"
|
||||
|
||||
print("✓ System prompt and configuration test passed")
|
||||
|
||||
|
||||
def test_agent_with_dynamic_temperature():
|
||||
"""Test agent with dynamic temperature"""
|
||||
print("\nTesting agent with dynamic temperature...")
|
||||
|
||||
model = OpenAIChat(model_name="gpt-4.1")
|
||||
agent = Agent(
|
||||
agent_name="Dynamic-Temp-Agent",
|
||||
llm=model,
|
||||
max_loops=2,
|
||||
dynamic_temperature_enabled=True,
|
||||
)
|
||||
|
||||
response = agent.run("Generate a creative story.")
|
||||
assert response is not None, "Dynamic temperature test failed"
|
||||
print("✓ Dynamic temperature test passed")
|
||||
|
||||
|
||||
def run_all_tests():
|
||||
"""Run all test functions"""
|
||||
print("Starting Extended Agent functional tests...\n")
|
||||
|
||||
test_functions = [
|
||||
test_basic_agent_functionality,
|
||||
test_memory_management,
|
||||
test_agent_output_formats,
|
||||
test_agent_state_management,
|
||||
test_agent_tools_and_execution,
|
||||
test_agent_concurrent_execution,
|
||||
test_agent_error_handling,
|
||||
test_agent_configuration,
|
||||
test_agent_with_stopping_condition,
|
||||
test_agent_with_retry_mechanism,
|
||||
test_agent_with_dynamic_temperature,
|
||||
test_bulk_and_filtered_operations,
|
||||
test_memory_and_state_persistence,
|
||||
test_sentiment_and_evaluation,
|
||||
test_tool_management,
|
||||
test_system_prompt_and_configuration,
|
||||
]
|
||||
|
||||
# Run synchronous tests
|
||||
total_tests = len(test_functions) + 1 # +1 for async test
|
||||
passed_tests = 0
|
||||
|
||||
for test in test_functions:
|
||||
try:
|
||||
test()
|
||||
passed_tests += 1
|
||||
except Exception as e:
|
||||
print(f"✗ Test {test.__name__} failed: {str(e)}")
|
||||
|
||||
# Run async test
|
||||
try:
|
||||
asyncio.run(test_async_operations())
|
||||
passed_tests += 1
|
||||
except Exception as e:
|
||||
print(f"✗ Async operations test failed: {str(e)}")
|
||||
|
||||
print("\nExtended Test Summary:")
|
||||
print(f"Total Tests: {total_tests}")
|
||||
print(f"Passed: {passed_tests}")
|
||||
print(f"Failed: {total_tests - passed_tests}")
|
||||
print(f"Success Rate: {(passed_tests/total_tests)*100:.2f}%")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_all_tests()
|
||||
@ -0,0 +1,381 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from swarms.structs.agent_router import AgentRouter
|
||||
from swarms.structs.agent import Agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_agent():
|
||||
"""Create a real agent for testing."""
|
||||
with patch("swarms.structs.agent.LiteLLM") as mock_llm:
|
||||
mock_llm.return_value.run.return_value = "Test response"
|
||||
return Agent(
|
||||
agent_name="test_agent",
|
||||
agent_description="A test agent",
|
||||
system_prompt="You are a test agent",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
)
|
||||
|
||||
|
||||
def test_agent_router_initialization_default():
|
||||
"""Test AgentRouter initialization with default parameters."""
|
||||
with patch("swarms.structs.agent_router.embedding"):
|
||||
router = AgentRouter()
|
||||
|
||||
assert router.embedding_model == "text-embedding-ada-002"
|
||||
assert router.n_agents == 1
|
||||
assert router.api_key is None
|
||||
assert router.api_base is None
|
||||
assert router.agents == []
|
||||
assert router.agent_embeddings == []
|
||||
assert router.agent_metadata == []
|
||||
|
||||
|
||||
def test_agent_router_initialization_custom():
|
||||
"""Test AgentRouter initialization with custom parameters."""
|
||||
with patch("swarms.structs.agent_router.embedding"), patch(
|
||||
"swarms.structs.agent.LiteLLM"
|
||||
) as mock_llm:
|
||||
mock_llm.return_value.run.return_value = "Test response"
|
||||
agents = [
|
||||
Agent(
|
||||
agent_name="test1",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
),
|
||||
Agent(
|
||||
agent_name="test2",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
),
|
||||
]
|
||||
router = AgentRouter(
|
||||
embedding_model="custom-model",
|
||||
n_agents=3,
|
||||
api_key="custom_key",
|
||||
api_base="custom_base",
|
||||
agents=agents,
|
||||
)
|
||||
|
||||
assert router.embedding_model == "custom-model"
|
||||
assert router.n_agents == 3
|
||||
assert router.api_key == "custom_key"
|
||||
assert router.api_base == "custom_base"
|
||||
assert len(router.agents) == 2
|
||||
|
||||
|
||||
def test_cosine_similarity_identical_vectors():
|
||||
"""Test cosine similarity with identical vectors."""
|
||||
router = AgentRouter()
|
||||
vec1 = [1.0, 0.0, 0.0]
|
||||
vec2 = [1.0, 0.0, 0.0]
|
||||
|
||||
result = router._cosine_similarity(vec1, vec2)
|
||||
assert result == 1.0
|
||||
|
||||
|
||||
def test_cosine_similarity_orthogonal_vectors():
|
||||
"""Test cosine similarity with orthogonal vectors."""
|
||||
router = AgentRouter()
|
||||
vec1 = [1.0, 0.0, 0.0]
|
||||
vec2 = [0.0, 1.0, 0.0]
|
||||
|
||||
result = router._cosine_similarity(vec1, vec2)
|
||||
assert result == 0.0
|
||||
|
||||
|
||||
def test_cosine_similarity_opposite_vectors():
|
||||
"""Test cosine similarity with opposite vectors."""
|
||||
router = AgentRouter()
|
||||
vec1 = [1.0, 0.0, 0.0]
|
||||
vec2 = [-1.0, 0.0, 0.0]
|
||||
|
||||
result = router._cosine_similarity(vec1, vec2)
|
||||
assert result == -1.0
|
||||
|
||||
|
||||
def test_cosine_similarity_different_lengths():
|
||||
"""Test cosine similarity with vectors of different lengths."""
|
||||
router = AgentRouter()
|
||||
vec1 = [1.0, 0.0]
|
||||
vec2 = [1.0, 0.0, 0.0]
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Vectors must have the same length"
|
||||
):
|
||||
router._cosine_similarity(vec1, vec2)
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_generate_embedding_success(mock_embedding):
|
||||
"""Test successful embedding generation."""
|
||||
mock_embedding.return_value.data = [
|
||||
Mock(embedding=[0.1, 0.2, 0.3, 0.4])
|
||||
]
|
||||
|
||||
router = AgentRouter()
|
||||
result = router._generate_embedding("test text")
|
||||
|
||||
assert result == [0.1, 0.2, 0.3, 0.4]
|
||||
mock_embedding.assert_called_once()
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_generate_embedding_error(mock_embedding):
|
||||
"""Test embedding generation error handling."""
|
||||
mock_embedding.side_effect = Exception("API Error")
|
||||
|
||||
router = AgentRouter()
|
||||
|
||||
with pytest.raises(Exception, match="API Error"):
|
||||
router._generate_embedding("test text")
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_add_agent_success(mock_embedding, test_agent):
|
||||
"""Test successful agent addition."""
|
||||
mock_embedding.return_value.data = [
|
||||
Mock(embedding=[0.1, 0.2, 0.3])
|
||||
]
|
||||
|
||||
router = AgentRouter()
|
||||
router.add_agent(test_agent)
|
||||
|
||||
assert len(router.agents) == 1
|
||||
assert len(router.agent_embeddings) == 1
|
||||
assert len(router.agent_metadata) == 1
|
||||
assert router.agents[0] == test_agent
|
||||
assert router.agent_embeddings[0] == [0.1, 0.2, 0.3]
|
||||
assert router.agent_metadata[0]["name"] == "test_agent"
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_add_agent_retry_error(mock_embedding, test_agent):
|
||||
"""Test agent addition with retry mechanism failure."""
|
||||
mock_embedding.side_effect = Exception("Embedding error")
|
||||
|
||||
router = AgentRouter()
|
||||
|
||||
# Should raise RetryError after retries are exhausted
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
router.add_agent(test_agent)
|
||||
|
||||
# Check that it's a retry error or contains the original error
|
||||
assert "Embedding error" in str(
|
||||
exc_info.value
|
||||
) or "RetryError" in str(exc_info.value)
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_add_agents_multiple(mock_embedding):
|
||||
"""Test adding multiple agents."""
|
||||
mock_embedding.return_value.data = [
|
||||
Mock(embedding=[0.1, 0.2, 0.3])
|
||||
]
|
||||
|
||||
with patch("swarms.structs.agent.LiteLLM") as mock_llm:
|
||||
mock_llm.return_value.run.return_value = "Test response"
|
||||
router = AgentRouter()
|
||||
agents = [
|
||||
Agent(
|
||||
agent_name="agent1",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
),
|
||||
Agent(
|
||||
agent_name="agent2",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
),
|
||||
Agent(
|
||||
agent_name="agent3",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
),
|
||||
]
|
||||
|
||||
router.add_agents(agents)
|
||||
|
||||
assert len(router.agents) == 3
|
||||
assert len(router.agent_embeddings) == 3
|
||||
assert len(router.agent_metadata) == 3
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_find_best_agent_success(mock_embedding):
|
||||
"""Test successful best agent finding."""
|
||||
# Mock embeddings for agents and task
|
||||
mock_embedding.side_effect = [
|
||||
Mock(data=[Mock(embedding=[0.1, 0.2, 0.3])]), # agent1
|
||||
Mock(data=[Mock(embedding=[0.4, 0.5, 0.6])]), # agent2
|
||||
Mock(data=[Mock(embedding=[0.7, 0.8, 0.9])]), # task
|
||||
]
|
||||
|
||||
with patch("swarms.structs.agent.LiteLLM") as mock_llm:
|
||||
mock_llm.return_value.run.return_value = "Test response"
|
||||
router = AgentRouter()
|
||||
agent1 = Agent(
|
||||
agent_name="agent1",
|
||||
agent_description="First agent",
|
||||
system_prompt="Prompt 1",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
)
|
||||
agent2 = Agent(
|
||||
agent_name="agent2",
|
||||
agent_description="Second agent",
|
||||
system_prompt="Prompt 2",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
)
|
||||
|
||||
router.add_agent(agent1)
|
||||
router.add_agent(agent2)
|
||||
|
||||
# Mock the similarity calculation to return predictable results
|
||||
with patch.object(
|
||||
router, "_cosine_similarity"
|
||||
) as mock_similarity:
|
||||
mock_similarity.side_effect = [
|
||||
0.8,
|
||||
0.6,
|
||||
] # agent1 more similar
|
||||
|
||||
result = router.find_best_agent("test task")
|
||||
|
||||
assert result == agent1
|
||||
|
||||
|
||||
def test_find_best_agent_no_agents():
|
||||
"""Test finding best agent when no agents are available."""
|
||||
with patch("swarms.structs.agent_router.embedding"):
|
||||
router = AgentRouter()
|
||||
|
||||
result = router.find_best_agent("test task")
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_find_best_agent_retry_error(mock_embedding):
|
||||
"""Test error handling in find_best_agent with retry mechanism."""
|
||||
mock_embedding.side_effect = Exception("API Error")
|
||||
|
||||
with patch("swarms.structs.agent.LiteLLM") as mock_llm:
|
||||
mock_llm.return_value.run.return_value = "Test response"
|
||||
router = AgentRouter()
|
||||
router.agents = [
|
||||
Agent(
|
||||
agent_name="agent1",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
print_on=False,
|
||||
)
|
||||
]
|
||||
router.agent_embeddings = [[0.1, 0.2, 0.3]]
|
||||
|
||||
# Should raise RetryError after retries are exhausted
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
router.find_best_agent("test task")
|
||||
|
||||
# Check that it's a retry error or contains the original error
|
||||
assert "API Error" in str(
|
||||
exc_info.value
|
||||
) or "RetryError" in str(exc_info.value)
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_update_agent_history_success(mock_embedding, test_agent):
|
||||
"""Test successful agent history update."""
|
||||
mock_embedding.return_value.data = [
|
||||
Mock(embedding=[0.1, 0.2, 0.3])
|
||||
]
|
||||
|
||||
router = AgentRouter()
|
||||
router.add_agent(test_agent)
|
||||
|
||||
# Update agent history
|
||||
router.update_agent_history("test_agent")
|
||||
|
||||
# Verify the embedding was regenerated
|
||||
assert (
|
||||
mock_embedding.call_count == 2
|
||||
) # Once for add, once for update
|
||||
|
||||
|
||||
def test_update_agent_history_agent_not_found():
|
||||
"""Test updating history for non-existent agent."""
|
||||
with patch(
|
||||
"swarms.structs.agent_router.embedding"
|
||||
) as mock_embedding:
|
||||
mock_embedding.return_value.data = [
|
||||
Mock(embedding=[0.1, 0.2, 0.3])
|
||||
]
|
||||
router = AgentRouter()
|
||||
|
||||
# Should not raise an exception, just log a warning
|
||||
router.update_agent_history("non_existent_agent")
|
||||
|
||||
|
||||
@patch("swarms.structs.agent_router.embedding")
|
||||
def test_agent_metadata_structure(mock_embedding, test_agent):
|
||||
"""Test the structure of agent metadata."""
|
||||
mock_embedding.return_value.data = [
|
||||
Mock(embedding=[0.1, 0.2, 0.3])
|
||||
]
|
||||
|
||||
router = AgentRouter()
|
||||
router.add_agent(test_agent)
|
||||
|
||||
metadata = router.agent_metadata[0]
|
||||
assert "name" in metadata
|
||||
assert "text" in metadata
|
||||
assert metadata["name"] == "test_agent"
|
||||
assert (
|
||||
"test_agent A test agent You are a test agent"
|
||||
in metadata["text"]
|
||||
)
|
||||
|
||||
|
||||
def test_agent_router_edge_cases():
|
||||
"""Test various edge cases."""
|
||||
with patch(
|
||||
"swarms.structs.agent_router.embedding"
|
||||
) as mock_embedding:
|
||||
mock_embedding.return_value.data = [
|
||||
Mock(embedding=[0.1, 0.2, 0.3])
|
||||
]
|
||||
|
||||
router = AgentRouter()
|
||||
|
||||
# Test with empty string task
|
||||
result = router.find_best_agent("")
|
||||
assert result is None
|
||||
|
||||
# Test with very long task description
|
||||
long_task = "test " * 1000
|
||||
result = router.find_best_agent(long_task)
|
||||
assert result is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
@ -1,328 +0,0 @@
|
||||
import os
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
from loguru import logger
|
||||
from swarm_models import OpenAIChat
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.agent_rearrange import AgentRearrange
|
||||
|
||||
|
||||
class TestResult:
|
||||
"""Class to store test results and metadata"""
|
||||
|
||||
def __init__(self, test_name: str):
|
||||
self.test_name = test_name
|
||||
self.start_time = datetime.now()
|
||||
self.end_time = None
|
||||
self.success = False
|
||||
self.error = None
|
||||
self.traceback = None
|
||||
self.function_output = None
|
||||
|
||||
def complete(
|
||||
self, success: bool, error: Optional[Exception] = None
|
||||
):
|
||||
"""Complete the test execution with results"""
|
||||
self.end_time = datetime.now()
|
||||
self.success = success
|
||||
if error:
|
||||
self.error = str(error)
|
||||
self.traceback = traceback.format_exc()
|
||||
|
||||
def duration(self) -> float:
|
||||
"""Calculate test duration in seconds"""
|
||||
if self.end_time:
|
||||
return (self.end_time - self.start_time).total_seconds()
|
||||
return 0
|
||||
|
||||
|
||||
def run_test(test_func: Callable) -> TestResult:
|
||||
"""
|
||||
Decorator to run tests with error handling and logging
|
||||
|
||||
Args:
|
||||
test_func (Callable): Test function to execute
|
||||
|
||||
Returns:
|
||||
TestResult: Object containing test execution details
|
||||
"""
|
||||
|
||||
def wrapper(*args, **kwargs) -> TestResult:
|
||||
result = TestResult(test_func.__name__)
|
||||
logger.info(
|
||||
f"\n{'='*20} Running test: {test_func.__name__} {'='*20}"
|
||||
)
|
||||
|
||||
try:
|
||||
output = test_func(*args, **kwargs)
|
||||
result.function_output = output
|
||||
result.complete(success=True)
|
||||
logger.success(
|
||||
f"✅ Test {test_func.__name__} passed successfully"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
result.complete(success=False, error=e)
|
||||
logger.error(
|
||||
f"❌ Test {test_func.__name__} failed with error: {str(e)}"
|
||||
)
|
||||
logger.error(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
logger.info(
|
||||
f"Test duration: {result.duration():.2f} seconds\n"
|
||||
)
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def create_functional_agents() -> List[Agent]:
|
||||
"""
|
||||
Create a list of functional agents with real LLM integration for testing.
|
||||
Using OpenAI's GPT model for realistic agent behavior testing.
|
||||
"""
|
||||
# Initialize OpenAI Chat model
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
logger.warning(
|
||||
"No OpenAI API key found. Using mock agents instead."
|
||||
)
|
||||
return [
|
||||
create_mock_agent("TestAgent1"),
|
||||
create_mock_agent("TestAgent2"),
|
||||
]
|
||||
|
||||
try:
|
||||
model = OpenAIChat(
|
||||
api_key=api_key, model_name="gpt-4.1", temperature=0.1
|
||||
)
|
||||
|
||||
# Create boss agent
|
||||
boss_agent = Agent(
|
||||
agent_name="BossAgent",
|
||||
system_prompt="""
|
||||
You are the BossAgent responsible for managing and overseeing test scenarios.
|
||||
Your role is to coordinate tasks between agents and ensure efficient collaboration.
|
||||
Analyze inputs, break down tasks, and provide clear directives to other agents.
|
||||
Maintain a structured approach to task management and result compilation.
|
||||
""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
dashboard=False,
|
||||
streaming_on=True,
|
||||
verbose=True,
|
||||
stopping_token="<DONE>",
|
||||
state_save_file_type="json",
|
||||
saved_state_path="test_boss_agent.json",
|
||||
)
|
||||
|
||||
# Create analysis agent
|
||||
analysis_agent = Agent(
|
||||
agent_name="AnalysisAgent",
|
||||
system_prompt="""
|
||||
You are the AnalysisAgent responsible for detailed data processing and analysis.
|
||||
Your role is to examine input data, identify patterns, and provide analytical insights.
|
||||
Focus on breaking down complex information into clear, actionable components.
|
||||
""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
dashboard=False,
|
||||
streaming_on=True,
|
||||
verbose=True,
|
||||
stopping_token="<DONE>",
|
||||
state_save_file_type="json",
|
||||
saved_state_path="test_analysis_agent.json",
|
||||
)
|
||||
|
||||
# Create summary agent
|
||||
summary_agent = Agent(
|
||||
agent_name="SummaryAgent",
|
||||
system_prompt="""
|
||||
You are the SummaryAgent responsible for consolidating and summarizing information.
|
||||
Your role is to take detailed analysis and create concise, actionable summaries.
|
||||
Focus on highlighting key points and ensuring clarity in communication.
|
||||
""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
dashboard=False,
|
||||
streaming_on=True,
|
||||
verbose=True,
|
||||
stopping_token="<DONE>",
|
||||
state_save_file_type="json",
|
||||
saved_state_path="test_summary_agent.json",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Successfully created functional agents with LLM integration"
|
||||
)
|
||||
return [boss_agent, analysis_agent, summary_agent]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create functional agents: {str(e)}")
|
||||
logger.warning("Falling back to mock agents")
|
||||
return [
|
||||
create_mock_agent("TestAgent1"),
|
||||
create_mock_agent("TestAgent2"),
|
||||
]
|
||||
|
||||
|
||||
def create_mock_agent(name: str) -> Agent:
|
||||
"""Create a mock agent for testing when LLM integration is not available"""
|
||||
return Agent(
|
||||
agent_name=name,
|
||||
system_prompt=f"You are a test agent named {name}",
|
||||
llm=None,
|
||||
)
|
||||
|
||||
|
||||
@run_test
|
||||
def test_init():
|
||||
"""Test AgentRearrange initialization with functional agents"""
|
||||
logger.info("Creating agents for initialization test")
|
||||
agents = create_functional_agents()
|
||||
|
||||
rearrange = AgentRearrange(
|
||||
name="TestRearrange",
|
||||
agents=agents,
|
||||
flow=f"{agents[0].agent_name} -> {agents[1].agent_name} -> {agents[2].agent_name}",
|
||||
)
|
||||
|
||||
assert rearrange.name == "TestRearrange"
|
||||
assert len(rearrange.agents) == 3
|
||||
assert (
|
||||
rearrange.flow
|
||||
== f"{agents[0].agent_name} -> {agents[1].agent_name} -> {agents[2].agent_name}"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Initialized AgentRearrange with {len(agents)} agents"
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
@run_test
|
||||
def test_validate_flow():
|
||||
"""Test flow validation logic"""
|
||||
agents = create_functional_agents()
|
||||
rearrange = AgentRearrange(
|
||||
agents=agents,
|
||||
flow=f"{agents[0].agent_name} -> {agents[1].agent_name}",
|
||||
)
|
||||
|
||||
logger.info("Testing valid flow pattern")
|
||||
valid = rearrange.validate_flow()
|
||||
assert valid is True
|
||||
|
||||
logger.info("Testing invalid flow pattern")
|
||||
rearrange.flow = f"{agents[0].agent_name} {agents[1].agent_name}" # Missing arrow
|
||||
try:
|
||||
rearrange.validate_flow()
|
||||
assert False, "Should have raised ValueError"
|
||||
except ValueError as e:
|
||||
logger.info(
|
||||
f"Successfully caught invalid flow error: {str(e)}"
|
||||
)
|
||||
assert True
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@run_test
|
||||
def test_add_remove_agent():
|
||||
"""Test adding and removing agents from the swarm"""
|
||||
agents = create_functional_agents()
|
||||
rearrange = AgentRearrange(
|
||||
agents=agents[:2]
|
||||
) # Start with first two agents
|
||||
|
||||
logger.info("Testing agent addition")
|
||||
new_agent = agents[2] # Use the third agent as new agent
|
||||
rearrange.add_agent(new_agent)
|
||||
assert new_agent.agent_name in rearrange.agents
|
||||
|
||||
logger.info("Testing agent removal")
|
||||
rearrange.remove_agent(new_agent.agent_name)
|
||||
assert new_agent.agent_name not in rearrange.agents
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@run_test
|
||||
def test_basic_run():
|
||||
"""Test basic task execution with the swarm"""
|
||||
agents = create_functional_agents()
|
||||
rearrange = AgentRearrange(
|
||||
name="TestSwarm",
|
||||
agents=agents,
|
||||
flow=f"{agents[0].agent_name} -> {agents[1].agent_name} -> {agents[2].agent_name}",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
test_task = (
|
||||
"Analyze this test message and provide a brief summary."
|
||||
)
|
||||
logger.info(f"Running test task: {test_task}")
|
||||
|
||||
try:
|
||||
result = rearrange.run(test_task)
|
||||
assert result is not None
|
||||
logger.info(
|
||||
f"Successfully executed task with result length: {len(str(result))}"
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Task execution failed: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def run_all_tests() -> Dict[str, TestResult]:
|
||||
"""
|
||||
Run all test cases and collect results
|
||||
|
||||
Returns:
|
||||
Dict[str, TestResult]: Dictionary mapping test names to their results
|
||||
"""
|
||||
logger.info("\n🚀 Starting AgentRearrange test suite execution")
|
||||
test_functions = [
|
||||
test_init,
|
||||
test_validate_flow,
|
||||
test_add_remove_agent,
|
||||
test_basic_run,
|
||||
]
|
||||
|
||||
results = {}
|
||||
for test in test_functions:
|
||||
result = test()
|
||||
results[test.__name__] = result
|
||||
|
||||
# Log summary
|
||||
total_tests = len(results)
|
||||
passed_tests = sum(1 for r in results.values() if r.success)
|
||||
failed_tests = total_tests - passed_tests
|
||||
|
||||
logger.info("\n📊 Test Suite Summary:")
|
||||
logger.info(f"Total Tests: {total_tests}")
|
||||
print(f"✅ Passed: {passed_tests}")
|
||||
|
||||
if failed_tests > 0:
|
||||
logger.error(f"❌ Failed: {failed_tests}")
|
||||
|
||||
# Detailed failure information
|
||||
if failed_tests > 0:
|
||||
logger.error("\n❌ Failed Tests Details:")
|
||||
for name, result in results.items():
|
||||
if not result.success:
|
||||
logger.error(f"\n{name}:")
|
||||
logger.error(f"Error: {result.error}")
|
||||
logger.error(f"Traceback: {result.traceback}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("🌟 Starting AgentRearrange Test Suite")
|
||||
results = run_all_tests()
|
||||
print("🏁 Test Suite Execution Completed")
|
||||
@ -1,313 +0,0 @@
|
||||
import time
|
||||
|
||||
from loguru import logger
|
||||
from swarms import Agent
|
||||
|
||||
from experimental.airflow_swarm import (
|
||||
AirflowDAGSwarm,
|
||||
NodeType,
|
||||
Conversation,
|
||||
)
|
||||
|
||||
# Configure logger
|
||||
logger.remove()
|
||||
logger.add(lambda msg: print(msg, end=""), level="DEBUG")
|
||||
|
||||
|
||||
def test_swarm_initialization():
|
||||
"""Test basic swarm initialization and configuration."""
|
||||
try:
|
||||
swarm = AirflowDAGSwarm(
|
||||
dag_id="test_dag",
|
||||
name="Test DAG",
|
||||
initial_message="Test message",
|
||||
)
|
||||
assert swarm.dag_id == "test_dag", "DAG ID not set correctly"
|
||||
assert swarm.name == "Test DAG", "Name not set correctly"
|
||||
assert (
|
||||
len(swarm.nodes) == 0
|
||||
), "Nodes should be empty on initialization"
|
||||
assert (
|
||||
len(swarm.edges) == 0
|
||||
), "Edges should be empty on initialization"
|
||||
|
||||
# Test initial message
|
||||
conv_json = swarm.get_conversation_history()
|
||||
assert (
|
||||
"Test message" in conv_json
|
||||
), "Initial message not set correctly"
|
||||
print("✅ Swarm initialization test passed")
|
||||
return True
|
||||
except AssertionError as e:
|
||||
print(f"❌ Swarm initialization test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def test_node_addition():
|
||||
"""Test adding different types of nodes to the swarm."""
|
||||
try:
|
||||
swarm = AirflowDAGSwarm(dag_id="test_dag")
|
||||
|
||||
# Test adding an agent node
|
||||
agent = Agent(
|
||||
agent_name="Test-Agent",
|
||||
system_prompt="Test prompt",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
)
|
||||
agent_id = swarm.add_node(
|
||||
"test_agent",
|
||||
agent,
|
||||
NodeType.AGENT,
|
||||
query="Test query",
|
||||
concurrent=True,
|
||||
)
|
||||
assert (
|
||||
agent_id == "test_agent"
|
||||
), "Agent node ID not returned correctly"
|
||||
assert (
|
||||
"test_agent" in swarm.nodes
|
||||
), "Agent node not added to nodes dict"
|
||||
|
||||
# Test adding a callable node
|
||||
def test_callable(x: int, conversation: Conversation) -> str:
|
||||
return f"Test output {x}"
|
||||
|
||||
callable_id = swarm.add_node(
|
||||
"test_callable",
|
||||
test_callable,
|
||||
NodeType.CALLABLE,
|
||||
args=[42],
|
||||
concurrent=False,
|
||||
)
|
||||
assert (
|
||||
callable_id == "test_callable"
|
||||
), "Callable node ID not returned correctly"
|
||||
assert (
|
||||
"test_callable" in swarm.nodes
|
||||
), "Callable node not added to nodes dict"
|
||||
|
||||
print("✅ Node addition test passed")
|
||||
return True
|
||||
except AssertionError as e:
|
||||
print(f"❌ Node addition test failed: {str(e)}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(
|
||||
f"❌ Node addition test failed with unexpected error: {str(e)}"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def test_edge_addition():
|
||||
"""Test adding edges between nodes."""
|
||||
try:
|
||||
swarm = AirflowDAGSwarm(dag_id="test_dag")
|
||||
|
||||
# Add two nodes
|
||||
def node1_fn(conversation: Conversation) -> str:
|
||||
return "Node 1 output"
|
||||
|
||||
def node2_fn(conversation: Conversation) -> str:
|
||||
return "Node 2 output"
|
||||
|
||||
swarm.add_node("node1", node1_fn, NodeType.CALLABLE)
|
||||
swarm.add_node("node2", node2_fn, NodeType.CALLABLE)
|
||||
|
||||
# Add edge between them
|
||||
swarm.add_edge("node1", "node2")
|
||||
|
||||
assert (
|
||||
"node2" in swarm.edges["node1"]
|
||||
), "Edge not added correctly"
|
||||
assert (
|
||||
len(swarm.edges["node1"]) == 1
|
||||
), "Incorrect number of edges"
|
||||
|
||||
# Test adding edge with non-existent node
|
||||
try:
|
||||
swarm.add_edge("node1", "non_existent")
|
||||
assert (
|
||||
False
|
||||
), "Should raise ValueError for non-existent node"
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
print("✅ Edge addition test passed")
|
||||
return True
|
||||
except AssertionError as e:
|
||||
print(f"❌ Edge addition test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def test_execution_order():
|
||||
"""Test that nodes are executed in the correct order based on dependencies."""
|
||||
try:
|
||||
swarm = AirflowDAGSwarm(dag_id="test_dag")
|
||||
execution_order = []
|
||||
|
||||
def node1(conversation: Conversation) -> str:
|
||||
execution_order.append("node1")
|
||||
return "Node 1 output"
|
||||
|
||||
def node2(conversation: Conversation) -> str:
|
||||
execution_order.append("node2")
|
||||
return "Node 2 output"
|
||||
|
||||
def node3(conversation: Conversation) -> str:
|
||||
execution_order.append("node3")
|
||||
return "Node 3 output"
|
||||
|
||||
# Add nodes
|
||||
swarm.add_node(
|
||||
"node1", node1, NodeType.CALLABLE, concurrent=False
|
||||
)
|
||||
swarm.add_node(
|
||||
"node2", node2, NodeType.CALLABLE, concurrent=False
|
||||
)
|
||||
swarm.add_node(
|
||||
"node3", node3, NodeType.CALLABLE, concurrent=False
|
||||
)
|
||||
|
||||
# Add edges to create a chain: node1 -> node2 -> node3
|
||||
swarm.add_edge("node1", "node2")
|
||||
swarm.add_edge("node2", "node3")
|
||||
|
||||
# Execute
|
||||
swarm.run()
|
||||
|
||||
# Check execution order
|
||||
assert execution_order == [
|
||||
"node1",
|
||||
"node2",
|
||||
"node3",
|
||||
], "Incorrect execution order"
|
||||
print("✅ Execution order test passed")
|
||||
return True
|
||||
except AssertionError as e:
|
||||
print(f"❌ Execution order test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def test_concurrent_execution():
|
||||
"""Test concurrent execution of nodes."""
|
||||
try:
|
||||
swarm = AirflowDAGSwarm(dag_id="test_dag")
|
||||
|
||||
def slow_node1(conversation: Conversation) -> str:
|
||||
time.sleep(0.5)
|
||||
return "Slow node 1 output"
|
||||
|
||||
def slow_node2(conversation: Conversation) -> str:
|
||||
time.sleep(0.5)
|
||||
return "Slow node 2 output"
|
||||
|
||||
# Add nodes with concurrent=True
|
||||
swarm.add_node(
|
||||
"slow1", slow_node1, NodeType.CALLABLE, concurrent=True
|
||||
)
|
||||
swarm.add_node(
|
||||
"slow2", slow_node2, NodeType.CALLABLE, concurrent=True
|
||||
)
|
||||
|
||||
# Measure execution time
|
||||
start_time = time.time()
|
||||
swarm.run()
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Should take ~0.5s for concurrent execution, not ~1s
|
||||
assert (
|
||||
execution_time < 0.8
|
||||
), "Concurrent execution took too long"
|
||||
print("✅ Concurrent execution test passed")
|
||||
return True
|
||||
except AssertionError as e:
|
||||
print(f"❌ Concurrent execution test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def test_conversation_handling():
|
||||
"""Test conversation management within the swarm."""
|
||||
try:
|
||||
swarm = AirflowDAGSwarm(
|
||||
dag_id="test_dag", initial_message="Initial test message"
|
||||
)
|
||||
|
||||
# Test adding user messages
|
||||
swarm.add_user_message("Test message 1")
|
||||
swarm.add_user_message("Test message 2")
|
||||
|
||||
history = swarm.get_conversation_history()
|
||||
assert (
|
||||
"Initial test message" in history
|
||||
), "Initial message not in history"
|
||||
assert (
|
||||
"Test message 1" in history
|
||||
), "First message not in history"
|
||||
assert (
|
||||
"Test message 2" in history
|
||||
), "Second message not in history"
|
||||
|
||||
print("✅ Conversation handling test passed")
|
||||
return True
|
||||
except AssertionError as e:
|
||||
print(f"❌ Conversation handling test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def test_error_handling():
|
||||
"""Test error handling in node execution."""
|
||||
try:
|
||||
swarm = AirflowDAGSwarm(dag_id="test_dag")
|
||||
|
||||
def failing_node(conversation: Conversation) -> str:
|
||||
raise ValueError("Test error")
|
||||
|
||||
swarm.add_node("failing", failing_node, NodeType.CALLABLE)
|
||||
|
||||
# Execute should not raise an exception
|
||||
result = swarm.run()
|
||||
|
||||
assert (
|
||||
"Error" in result
|
||||
), "Error not captured in execution result"
|
||||
assert (
|
||||
"Test error" in result
|
||||
), "Specific error message not captured"
|
||||
|
||||
print("✅ Error handling test passed")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"❌ Error handling test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def run_all_tests():
|
||||
"""Run all test functions and report results."""
|
||||
tests = [
|
||||
test_swarm_initialization,
|
||||
test_node_addition,
|
||||
test_edge_addition,
|
||||
test_execution_order,
|
||||
test_concurrent_execution,
|
||||
test_conversation_handling,
|
||||
test_error_handling,
|
||||
]
|
||||
|
||||
results = []
|
||||
for test in tests:
|
||||
print(f"\nRunning {test.__name__}...")
|
||||
result = test()
|
||||
results.append(result)
|
||||
|
||||
total = len(results)
|
||||
passed = sum(results)
|
||||
print("\n=== Test Results ===")
|
||||
print(f"Total tests: {total}")
|
||||
print(f"Passed: {passed}")
|
||||
print(f"Failed: {total - passed}")
|
||||
print("==================")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_all_tests()
|
||||
File diff suppressed because it is too large
Load Diff
@ -1,293 +0,0 @@
|
||||
"""
|
||||
Tests for bug #1115 fix in AutoSwarmBuilder.
|
||||
|
||||
This test module verifies the fix for AttributeError when creating agents
|
||||
from AgentSpec Pydantic models in AutoSwarmBuilder.
|
||||
|
||||
Bug: https://github.com/kyegomez/swarms/issues/1115
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.auto_swarm_builder import (
|
||||
AgentSpec,
|
||||
AutoSwarmBuilder,
|
||||
)
|
||||
from swarms.structs.ma_utils import set_random_models_for_agents
|
||||
|
||||
|
||||
class TestAutoSwarmBuilderFix:
|
||||
"""Tests for bug #1115 fix in AutoSwarmBuilder."""
|
||||
|
||||
def test_create_agents_from_specs_with_dict(self):
|
||||
"""Test that create_agents_from_specs handles dict input correctly."""
|
||||
builder = AutoSwarmBuilder()
|
||||
|
||||
# Create specs as a dictionary
|
||||
specs = {
|
||||
"agents": [
|
||||
{
|
||||
"agent_name": "test_agent_1",
|
||||
"description": "Test agent 1 description",
|
||||
"system_prompt": "You are a helpful assistant",
|
||||
"model_name": "gpt-4o-mini",
|
||||
"max_loops": 1,
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
agents = builder.create_agents_from_specs(specs)
|
||||
|
||||
# Verify agents were created correctly
|
||||
assert len(agents) == 1
|
||||
assert isinstance(agents[0], Agent)
|
||||
assert agents[0].agent_name == "test_agent_1"
|
||||
|
||||
# Verify description was mapped to agent_description
|
||||
assert hasattr(agents[0], "agent_description")
|
||||
assert (
|
||||
agents[0].agent_description == "Test agent 1 description"
|
||||
)
|
||||
|
||||
def test_create_agents_from_specs_with_pydantic(self):
|
||||
"""Test that create_agents_from_specs handles Pydantic model input correctly.
|
||||
|
||||
This is the main test for bug #1115 - it verifies that AgentSpec
|
||||
Pydantic models can be unpacked correctly.
|
||||
"""
|
||||
builder = AutoSwarmBuilder()
|
||||
|
||||
# Create specs as Pydantic AgentSpec objects
|
||||
agent_spec = AgentSpec(
|
||||
agent_name="test_agent_pydantic",
|
||||
description="Pydantic test agent",
|
||||
system_prompt="You are a helpful assistant",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
specs = {"agents": [agent_spec]}
|
||||
|
||||
agents = builder.create_agents_from_specs(specs)
|
||||
|
||||
# Verify agents were created correctly
|
||||
assert len(agents) == 1
|
||||
assert isinstance(agents[0], Agent)
|
||||
assert agents[0].agent_name == "test_agent_pydantic"
|
||||
|
||||
# Verify description was mapped to agent_description
|
||||
assert hasattr(agents[0], "agent_description")
|
||||
assert agents[0].agent_description == "Pydantic test agent"
|
||||
|
||||
def test_parameter_name_mapping(self):
|
||||
"""Test that 'description' field maps to 'agent_description' correctly."""
|
||||
builder = AutoSwarmBuilder()
|
||||
|
||||
# Test with dict that has 'description'
|
||||
specs = {
|
||||
"agents": [
|
||||
{
|
||||
"agent_name": "mapping_test",
|
||||
"description": "This should map to agent_description",
|
||||
"system_prompt": "You are helpful",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
agents = builder.create_agents_from_specs(specs)
|
||||
|
||||
assert len(agents) == 1
|
||||
agent = agents[0]
|
||||
|
||||
# Verify description was mapped
|
||||
assert hasattr(agent, "agent_description")
|
||||
assert (
|
||||
agent.agent_description
|
||||
== "This should map to agent_description"
|
||||
)
|
||||
|
||||
def test_create_agents_from_specs_mixed_input(self):
|
||||
"""Test that create_agents_from_specs handles mixed dict and Pydantic input."""
|
||||
builder = AutoSwarmBuilder()
|
||||
|
||||
# Mix of dict and Pydantic objects
|
||||
dict_spec = {
|
||||
"agent_name": "dict_agent",
|
||||
"description": "Dict agent description",
|
||||
"system_prompt": "You are helpful",
|
||||
}
|
||||
|
||||
pydantic_spec = AgentSpec(
|
||||
agent_name="pydantic_agent",
|
||||
description="Pydantic agent description",
|
||||
system_prompt="You are smart",
|
||||
)
|
||||
|
||||
specs = {"agents": [dict_spec, pydantic_spec]}
|
||||
|
||||
agents = builder.create_agents_from_specs(specs)
|
||||
|
||||
# Verify both agents were created
|
||||
assert len(agents) == 2
|
||||
assert all(isinstance(agent, Agent) for agent in agents)
|
||||
|
||||
# Verify both have correct descriptions
|
||||
dict_agent = next(
|
||||
a for a in agents if a.agent_name == "dict_agent"
|
||||
)
|
||||
pydantic_agent = next(
|
||||
a for a in agents if a.agent_name == "pydantic_agent"
|
||||
)
|
||||
|
||||
assert (
|
||||
dict_agent.agent_description == "Dict agent description"
|
||||
)
|
||||
assert (
|
||||
pydantic_agent.agent_description
|
||||
== "Pydantic agent description"
|
||||
)
|
||||
|
||||
def test_set_random_models_for_agents_with_valid_agents(
|
||||
self,
|
||||
):
|
||||
"""Test set_random_models_for_agents with proper Agent objects."""
|
||||
# Create proper Agent objects
|
||||
agents = [
|
||||
Agent(
|
||||
agent_name="agent1",
|
||||
system_prompt="You are agent 1",
|
||||
max_loops=1,
|
||||
),
|
||||
Agent(
|
||||
agent_name="agent2",
|
||||
system_prompt="You are agent 2",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
# Set random models
|
||||
model_names = ["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet"]
|
||||
result = set_random_models_for_agents(
|
||||
agents=agents, model_names=model_names
|
||||
)
|
||||
|
||||
# Verify results
|
||||
assert len(result) == 2
|
||||
assert all(isinstance(agent, Agent) for agent in result)
|
||||
assert all(hasattr(agent, "model_name") for agent in result)
|
||||
assert all(
|
||||
agent.model_name in model_names for agent in result
|
||||
)
|
||||
|
||||
def test_set_random_models_for_agents_with_single_agent(
|
||||
self,
|
||||
):
|
||||
"""Test set_random_models_for_agents with a single agent."""
|
||||
agent = Agent(
|
||||
agent_name="single_agent",
|
||||
system_prompt="You are helpful",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
model_names = ["gpt-4o-mini", "gpt-4o"]
|
||||
result = set_random_models_for_agents(
|
||||
agents=agent, model_names=model_names
|
||||
)
|
||||
|
||||
assert isinstance(result, Agent)
|
||||
assert hasattr(result, "model_name")
|
||||
assert result.model_name in model_names
|
||||
|
||||
def test_set_random_models_for_agents_with_none(self):
|
||||
"""Test set_random_models_for_agents with None returns random model name."""
|
||||
model_names = ["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet"]
|
||||
result = set_random_models_for_agents(
|
||||
agents=None, model_names=model_names
|
||||
)
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert result in model_names
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="This test requires API key and makes LLM calls"
|
||||
)
|
||||
def test_auto_swarm_builder_return_agents_objects_integration(
|
||||
self,
|
||||
):
|
||||
"""Integration test for AutoSwarmBuilder with execution_type='return-agents-objects'.
|
||||
|
||||
This test requires OPENAI_API_KEY and makes actual LLM calls.
|
||||
Run manually with: pytest -k test_auto_swarm_builder_return_agents_objects_integration -v
|
||||
"""
|
||||
builder = AutoSwarmBuilder(
|
||||
execution_type="return-agents-objects",
|
||||
model_name="gpt-4o-mini",
|
||||
max_loops=1,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
agents = builder.run(
|
||||
"Create a team of 2 data analysis agents with specific roles"
|
||||
)
|
||||
|
||||
# Verify agents were created
|
||||
assert isinstance(agents, list)
|
||||
assert len(agents) >= 1
|
||||
assert all(isinstance(agent, Agent) for agent in agents)
|
||||
assert all(hasattr(agent, "agent_name") for agent in agents)
|
||||
assert all(
|
||||
hasattr(agent, "agent_description") for agent in agents
|
||||
)
|
||||
|
||||
def test_agent_spec_to_agent_all_fields(self):
|
||||
"""Test that all AgentSpec fields are properly passed to Agent."""
|
||||
builder = AutoSwarmBuilder()
|
||||
|
||||
agent_spec = AgentSpec(
|
||||
agent_name="full_test_agent",
|
||||
description="Full test description",
|
||||
system_prompt="You are a comprehensive test agent",
|
||||
model_name="gpt-4o-mini",
|
||||
auto_generate_prompt=False,
|
||||
max_tokens=4096,
|
||||
temperature=0.7,
|
||||
role="worker",
|
||||
max_loops=3,
|
||||
goal="Test all parameters",
|
||||
)
|
||||
|
||||
agents = builder.create_agents_from_specs(
|
||||
{"agents": [agent_spec]}
|
||||
)
|
||||
|
||||
assert len(agents) == 1
|
||||
agent = agents[0]
|
||||
|
||||
# Verify all fields were set
|
||||
assert agent.agent_name == "full_test_agent"
|
||||
assert agent.agent_description == "Full test description"
|
||||
# Agent may modify system_prompt by adding additional instructions
|
||||
assert (
|
||||
"You are a comprehensive test agent"
|
||||
in agent.system_prompt
|
||||
)
|
||||
assert agent.max_loops == 3
|
||||
assert agent.max_tokens == 4096
|
||||
assert agent.temperature == 0.7
|
||||
|
||||
def test_create_agents_from_specs_empty_list(self):
|
||||
"""Test that create_agents_from_specs handles empty agent list."""
|
||||
builder = AutoSwarmBuilder()
|
||||
|
||||
specs = {"agents": []}
|
||||
|
||||
agents = builder.create_agents_from_specs(specs)
|
||||
|
||||
assert isinstance(agents, list)
|
||||
assert len(agents) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run tests with pytest
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@ -1,287 +0,0 @@
|
||||
import os
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from swarms.structs.base_structure import BaseStructure
|
||||
|
||||
|
||||
class TestBaseStructure:
|
||||
def test_init(self):
|
||||
base_structure = BaseStructure(
|
||||
name="TestStructure",
|
||||
description="Test description",
|
||||
save_metadata=True,
|
||||
save_artifact_path="./test_artifacts",
|
||||
save_metadata_path="./test_metadata",
|
||||
save_error_path="./test_errors",
|
||||
)
|
||||
|
||||
assert base_structure.name == "TestStructure"
|
||||
assert base_structure.description == "Test description"
|
||||
assert base_structure.save_metadata is True
|
||||
assert base_structure.save_artifact_path == "./test_artifacts"
|
||||
assert base_structure.save_metadata_path == "./test_metadata"
|
||||
assert base_structure.save_error_path == "./test_errors"
|
||||
|
||||
def test_save_to_file_and_load_from_file(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
file_path = os.path.join(tmp_dir, "test_file.json")
|
||||
|
||||
data_to_save = {"key": "value"}
|
||||
base_structure = BaseStructure()
|
||||
|
||||
base_structure.save_to_file(data_to_save, file_path)
|
||||
loaded_data = base_structure.load_from_file(file_path)
|
||||
|
||||
assert loaded_data == data_to_save
|
||||
|
||||
def test_save_metadata_and_load_metadata(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_metadata_path=tmp_dir)
|
||||
|
||||
metadata = {"name": "Test", "description": "Test metadata"}
|
||||
base_structure.save_metadata(metadata)
|
||||
loaded_metadata = base_structure.load_metadata()
|
||||
|
||||
assert loaded_metadata == metadata
|
||||
|
||||
def test_log_error(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_error_path=tmp_dir)
|
||||
|
||||
error_message = "Test error message"
|
||||
base_structure.log_error(error_message)
|
||||
|
||||
log_file = os.path.join(tmp_dir, "TestStructure_errors.log")
|
||||
with open(log_file) as file:
|
||||
lines = file.readlines()
|
||||
assert len(lines) == 1
|
||||
assert lines[0] == f"{error_message}\n"
|
||||
|
||||
def test_save_artifact_and_load_artifact(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_artifact_path=tmp_dir)
|
||||
|
||||
artifact = {"key": "value"}
|
||||
artifact_name = "test_artifact"
|
||||
base_structure.save_artifact(artifact, artifact_name)
|
||||
loaded_artifact = base_structure.load_artifact(artifact_name)
|
||||
|
||||
assert loaded_artifact == artifact
|
||||
|
||||
def test_current_timestamp(self):
|
||||
base_structure = BaseStructure()
|
||||
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
timestamp = base_structure._current_timestamp()
|
||||
assert timestamp == current_time
|
||||
|
||||
def test_log_event(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_metadata_path=tmp_dir)
|
||||
|
||||
event = "Test event"
|
||||
event_type = "INFO"
|
||||
base_structure.log_event(event, event_type)
|
||||
|
||||
log_file = os.path.join(tmp_dir, "TestStructure_events.log")
|
||||
with open(log_file) as file:
|
||||
lines = file.readlines()
|
||||
assert len(lines) == 1
|
||||
assert (
|
||||
lines[0] == f"[{base_structure._current_timestamp()}]"
|
||||
f" [{event_type}] {event}\n"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(self):
|
||||
base_structure = BaseStructure()
|
||||
|
||||
async def async_function():
|
||||
return "Async Test Result"
|
||||
|
||||
result = await base_structure.run_async(async_function)
|
||||
assert result == "Async Test Result"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_metadata_async(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_metadata_path=tmp_dir)
|
||||
|
||||
metadata = {"name": "Test", "description": "Test metadata"}
|
||||
await base_structure.save_metadata_async(metadata)
|
||||
loaded_metadata = base_structure.load_metadata()
|
||||
|
||||
assert loaded_metadata == metadata
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_error_async(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_error_path=tmp_dir)
|
||||
|
||||
error_message = "Test error message"
|
||||
await base_structure.log_error_async(error_message)
|
||||
|
||||
log_file = os.path.join(tmp_dir, "TestStructure_errors.log")
|
||||
with open(log_file) as file:
|
||||
lines = file.readlines()
|
||||
assert len(lines) == 1
|
||||
assert lines[0] == f"{error_message}\n"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_artifact_async(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_artifact_path=tmp_dir)
|
||||
|
||||
artifact = {"key": "value"}
|
||||
artifact_name = "test_artifact"
|
||||
await base_structure.save_artifact_async(
|
||||
artifact, artifact_name
|
||||
)
|
||||
loaded_artifact = base_structure.load_artifact(artifact_name)
|
||||
|
||||
assert loaded_artifact == artifact
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_artifact_async(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_artifact_path=tmp_dir)
|
||||
|
||||
artifact = {"key": "value"}
|
||||
artifact_name = "test_artifact"
|
||||
base_structure.save_artifact(artifact, artifact_name)
|
||||
loaded_artifact = await base_structure.load_artifact_async(
|
||||
artifact_name
|
||||
)
|
||||
|
||||
assert loaded_artifact == artifact
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_event_async(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure(save_metadata_path=tmp_dir)
|
||||
|
||||
event = "Test event"
|
||||
event_type = "INFO"
|
||||
await base_structure.log_event_async(event, event_type)
|
||||
|
||||
log_file = os.path.join(tmp_dir, "TestStructure_events.log")
|
||||
with open(log_file) as file:
|
||||
lines = file.readlines()
|
||||
assert len(lines) == 1
|
||||
assert (
|
||||
lines[0] == f"[{base_structure._current_timestamp()}]"
|
||||
f" [{event_type}] {event}\n"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asave_to_file(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
file_path = os.path.join(tmp_dir, "test_file.json")
|
||||
data_to_save = {"key": "value"}
|
||||
base_structure = BaseStructure()
|
||||
|
||||
await base_structure.asave_to_file(data_to_save, file_path)
|
||||
loaded_data = base_structure.load_from_file(file_path)
|
||||
|
||||
assert loaded_data == data_to_save
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aload_from_file(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
file_path = os.path.join(tmp_dir, "test_file.json")
|
||||
data_to_save = {"key": "value"}
|
||||
base_structure = BaseStructure()
|
||||
base_structure.save_to_file(data_to_save, file_path)
|
||||
|
||||
loaded_data = await base_structure.aload_from_file(file_path)
|
||||
assert loaded_data == data_to_save
|
||||
|
||||
def test_run_in_thread(self):
|
||||
base_structure = BaseStructure()
|
||||
result = base_structure.run_in_thread(
|
||||
lambda: "Thread Test Result"
|
||||
)
|
||||
assert result.result() == "Thread Test Result"
|
||||
|
||||
def test_save_and_decompress_data(self):
|
||||
base_structure = BaseStructure()
|
||||
data = {"key": "value"}
|
||||
compressed_data = base_structure.compress_data(data)
|
||||
decompressed_data = base_structure.decompres_data(
|
||||
compressed_data
|
||||
)
|
||||
assert decompressed_data == data
|
||||
|
||||
def test_run_batched(self):
|
||||
base_structure = BaseStructure()
|
||||
|
||||
def run_function(data):
|
||||
return f"Processed {data}"
|
||||
|
||||
batched_data = list(range(10))
|
||||
result = base_structure.run_batched(
|
||||
batched_data, batch_size=5, func=run_function
|
||||
)
|
||||
|
||||
expected_result = [
|
||||
f"Processed {data}" for data in batched_data
|
||||
]
|
||||
assert result == expected_result
|
||||
|
||||
def test_load_config(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
config_file = os.path.join(tmp_dir, "config.json")
|
||||
config_data = {"key": "value"}
|
||||
base_structure = BaseStructure()
|
||||
|
||||
base_structure.save_to_file(config_data, config_file)
|
||||
loaded_config = base_structure.load_config(config_file)
|
||||
|
||||
assert loaded_config == config_data
|
||||
|
||||
def test_backup_data(self, tmpdir):
|
||||
tmp_dir = tmpdir.mkdir("test_dir")
|
||||
base_structure = BaseStructure()
|
||||
data_to_backup = {"key": "value"}
|
||||
base_structure.backup_data(
|
||||
data_to_backup, backup_path=tmp_dir
|
||||
)
|
||||
backup_files = os.listdir(tmp_dir)
|
||||
|
||||
assert len(backup_files) == 1
|
||||
loaded_data = base_structure.load_from_file(
|
||||
os.path.join(tmp_dir, backup_files[0])
|
||||
)
|
||||
assert loaded_data == data_to_backup
|
||||
|
||||
def test_monitor_resources(self):
|
||||
base_structure = BaseStructure()
|
||||
base_structure.monitor_resources()
|
||||
|
||||
def test_run_with_resources(self):
|
||||
base_structure = BaseStructure()
|
||||
|
||||
def run_function():
|
||||
base_structure.monitor_resources()
|
||||
return "Resource Test Result"
|
||||
|
||||
result = base_structure.run_with_resources(run_function)
|
||||
assert result == "Resource Test Result"
|
||||
|
||||
def test_run_with_resources_batched(self):
|
||||
base_structure = BaseStructure()
|
||||
|
||||
def run_function(data):
|
||||
base_structure.monitor_resources()
|
||||
return f"Processed {data}"
|
||||
|
||||
batched_data = list(range(10))
|
||||
result = base_structure.run_with_resources_batched(
|
||||
batched_data, batch_size=5, func=run_function
|
||||
)
|
||||
|
||||
expected_result = [
|
||||
f"Processed {data}" for data in batched_data
|
||||
]
|
||||
assert result == expected_result
|
||||
@ -1,67 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from swarm_models import OpenAIChat
|
||||
from swarms.structs import BaseWorkflow
|
||||
|
||||
load_dotenv()
|
||||
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
|
||||
|
||||
def setup_workflow():
|
||||
llm = OpenAIChat(openai_api_key=api_key)
|
||||
workflow = BaseWorkflow(max_loops=1)
|
||||
workflow.add("What's the weather in miami", llm)
|
||||
workflow.add("Create a report on these metrics", llm)
|
||||
workflow.save_workflow_state("workflow_state.json")
|
||||
return workflow
|
||||
|
||||
|
||||
def teardown_workflow():
|
||||
os.remove("workflow_state.json")
|
||||
|
||||
|
||||
def test_load_workflow_state():
|
||||
workflow = setup_workflow()
|
||||
workflow.load_workflow_state("workflow_state.json")
|
||||
assert workflow.max_loops == 1
|
||||
assert len(workflow.tasks) == 2
|
||||
assert (
|
||||
workflow.tasks[0].description == "What's the weather in miami"
|
||||
)
|
||||
assert (
|
||||
workflow.tasks[1].description
|
||||
== "Create a report on these metrics"
|
||||
)
|
||||
teardown_workflow()
|
||||
|
||||
|
||||
def test_load_workflow_state_with_missing_file():
|
||||
workflow = setup_workflow()
|
||||
with pytest.raises(FileNotFoundError):
|
||||
workflow.load_workflow_state("non_existent_file.json")
|
||||
teardown_workflow()
|
||||
|
||||
|
||||
def test_load_workflow_state_with_invalid_file():
|
||||
workflow = setup_workflow()
|
||||
with open("invalid_file.json", "w") as f:
|
||||
f.write("This is not valid JSON")
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
workflow.load_workflow_state("invalid_file.json")
|
||||
os.remove("invalid_file.json")
|
||||
teardown_workflow()
|
||||
|
||||
|
||||
def test_load_workflow_state_with_missing_keys():
|
||||
workflow = setup_workflow()
|
||||
with open("missing_keys.json", "w") as f:
|
||||
json.dump({"max_loops": 1}, f)
|
||||
with pytest.raises(KeyError):
|
||||
workflow.load_workflow_state("missing_keys.json")
|
||||
os.remove("missing_keys.json")
|
||||
teardown_workflow()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,354 @@
|
||||
from swarms import Agent
|
||||
from swarms.structs.hiearchical_swarm import HierarchicalSwarm
|
||||
|
||||
|
||||
def test_hierarchical_swarm_basic_initialization():
|
||||
"""Test basic HierarchicalSwarm initialization"""
|
||||
# Create worker agents
|
||||
research_agent = Agent(
|
||||
agent_name="Research-Specialist",
|
||||
agent_description="Specialist in research and data collection",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
analysis_agent = Agent(
|
||||
agent_name="Analysis-Expert",
|
||||
agent_description="Expert in data analysis and insights",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
implementation_agent = Agent(
|
||||
agent_name="Implementation-Manager",
|
||||
agent_description="Manager for implementation and execution",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create swarm with agents
|
||||
swarm = HierarchicalSwarm(
|
||||
name="Research-Analysis-Implementation-Swarm",
|
||||
description="Hierarchical swarm for comprehensive project execution",
|
||||
agents=[research_agent, analysis_agent, implementation_agent],
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Verify initialization
|
||||
assert swarm.name == "Research-Analysis-Implementation-Swarm"
|
||||
assert (
|
||||
swarm.description
|
||||
== "Hierarchical swarm for comprehensive project execution"
|
||||
)
|
||||
assert len(swarm.agents) == 3
|
||||
assert swarm.max_loops == 1
|
||||
assert swarm.director is not None
|
||||
|
||||
|
||||
def test_hierarchical_swarm_with_director():
|
||||
"""Test HierarchicalSwarm with custom director"""
|
||||
# Create a custom director
|
||||
director = Agent(
|
||||
agent_name="Project-Director",
|
||||
agent_description="Senior project director with extensive experience",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create worker agents
|
||||
developer = Agent(
|
||||
agent_name="Senior-Developer",
|
||||
agent_description="Senior software developer",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
tester = Agent(
|
||||
agent_name="QA-Lead",
|
||||
agent_description="Quality assurance lead",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create swarm with custom director
|
||||
swarm = HierarchicalSwarm(
|
||||
name="Software-Development-Swarm",
|
||||
description="Hierarchical swarm for software development projects",
|
||||
director=director,
|
||||
agents=[developer, tester],
|
||||
max_loops=2,
|
||||
)
|
||||
|
||||
assert swarm.director == director
|
||||
assert len(swarm.agents) == 2
|
||||
assert swarm.max_loops == 2
|
||||
|
||||
|
||||
def test_hierarchical_swarm_execution():
|
||||
"""Test HierarchicalSwarm execution with multiple agents"""
|
||||
# Create specialized agents
|
||||
market_researcher = Agent(
|
||||
agent_name="Market-Researcher",
|
||||
agent_description="Market research specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
product_strategist = Agent(
|
||||
agent_name="Product-Strategist",
|
||||
agent_description="Product strategy and planning expert",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
technical_architect = Agent(
|
||||
agent_name="Technical-Architect",
|
||||
agent_description="Technical architecture and design specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
risk_analyst = Agent(
|
||||
agent_name="Risk-Analyst",
|
||||
agent_description="Risk assessment and mitigation specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create hierarchical swarm
|
||||
swarm = HierarchicalSwarm(
|
||||
name="Product-Development-Swarm",
|
||||
description="Comprehensive product development hierarchical swarm",
|
||||
agents=[
|
||||
market_researcher,
|
||||
product_strategist,
|
||||
technical_architect,
|
||||
risk_analyst,
|
||||
],
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Execute swarm
|
||||
result = swarm.run(
|
||||
"Develop a comprehensive strategy for a new AI-powered healthcare platform"
|
||||
)
|
||||
|
||||
# Verify result structure
|
||||
assert result is not None
|
||||
# HierarchicalSwarm returns a SwarmSpec or conversation history, just ensure it's not None
|
||||
|
||||
|
||||
def test_hierarchical_swarm_multiple_loops():
|
||||
"""Test HierarchicalSwarm with multiple feedback loops"""
|
||||
# Create agents for iterative refinement
|
||||
planner = Agent(
|
||||
agent_name="Strategic-Planner",
|
||||
agent_description="Strategic planning and project management",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
executor = Agent(
|
||||
agent_name="Task-Executor",
|
||||
agent_description="Task execution and implementation",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
reviewer = Agent(
|
||||
agent_name="Quality-Reviewer",
|
||||
agent_description="Quality assurance and review specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create swarm with multiple loops for iterative refinement
|
||||
swarm = HierarchicalSwarm(
|
||||
name="Iterative-Development-Swarm",
|
||||
description="Hierarchical swarm with iterative feedback loops",
|
||||
agents=[planner, executor, reviewer],
|
||||
max_loops=3, # Allow multiple iterations
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Execute with multiple loops
|
||||
result = swarm.run(
|
||||
"Create a detailed project plan for implementing a machine learning recommendation system"
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_hierarchical_swarm_error_handling():
|
||||
"""Test HierarchicalSwarm error handling"""
|
||||
# Test with empty agents list
|
||||
try:
|
||||
swarm = HierarchicalSwarm(agents=[])
|
||||
assert (
|
||||
False
|
||||
), "Should have raised ValueError for empty agents list"
|
||||
except ValueError as e:
|
||||
assert "agents" in str(e).lower() or "empty" in str(e).lower()
|
||||
|
||||
# Test with invalid max_loops
|
||||
researcher = Agent(
|
||||
agent_name="Test-Researcher",
|
||||
agent_description="Test researcher",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
try:
|
||||
swarm = HierarchicalSwarm(agents=[researcher], max_loops=0)
|
||||
assert (
|
||||
False
|
||||
), "Should have raised ValueError for invalid max_loops"
|
||||
except ValueError as e:
|
||||
assert "max_loops" in str(e).lower() or "0" in str(e)
|
||||
|
||||
|
||||
def test_hierarchical_swarm_collaboration_prompts():
|
||||
"""Test HierarchicalSwarm with collaboration prompts enabled"""
|
||||
# Create agents
|
||||
data_analyst = Agent(
|
||||
agent_name="Data-Analyst",
|
||||
agent_description="Data analysis specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
business_analyst = Agent(
|
||||
agent_name="Business-Analyst",
|
||||
agent_description="Business analysis specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create swarm with collaboration prompts
|
||||
swarm = HierarchicalSwarm(
|
||||
name="Collaborative-Analysis-Swarm",
|
||||
description="Hierarchical swarm with enhanced collaboration",
|
||||
agents=[data_analyst, business_analyst],
|
||||
max_loops=1,
|
||||
add_collaboration_prompt=True,
|
||||
)
|
||||
|
||||
# Check that collaboration prompts were added to agents
|
||||
assert data_analyst.system_prompt is not None
|
||||
assert business_analyst.system_prompt is not None
|
||||
|
||||
# Execute swarm
|
||||
result = swarm.run(
|
||||
"Analyze customer behavior patterns and provide business recommendations"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_hierarchical_swarm_with_dashboard():
|
||||
"""Test HierarchicalSwarm with interactive dashboard"""
|
||||
# Create agents
|
||||
content_creator = Agent(
|
||||
agent_name="Content-Creator",
|
||||
agent_description="Content creation specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
editor = Agent(
|
||||
agent_name="Editor",
|
||||
agent_description="Content editor and proofreader",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
publisher = Agent(
|
||||
agent_name="Publisher",
|
||||
agent_description="Publishing and distribution specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create swarm with interactive dashboard
|
||||
swarm = HierarchicalSwarm(
|
||||
name="Content-Publishing-Swarm",
|
||||
description="Hierarchical swarm for content creation and publishing",
|
||||
agents=[content_creator, editor, publisher],
|
||||
max_loops=1,
|
||||
interactive=True,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Verify dashboard was created
|
||||
assert swarm.dashboard is not None
|
||||
assert swarm.interactive is True
|
||||
|
||||
# Execute swarm
|
||||
result = swarm.run(
|
||||
"Create a comprehensive guide on machine learning best practices"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_hierarchical_swarm_real_world_scenario():
|
||||
"""Test HierarchicalSwarm in a realistic business scenario"""
|
||||
# Create agents representing different business functions
|
||||
market_intelligence = Agent(
|
||||
agent_name="Market-Intelligence-Director",
|
||||
agent_description="Director of market intelligence and competitive analysis",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
product_strategy = Agent(
|
||||
agent_name="Product-Strategy-Manager",
|
||||
agent_description="Product strategy and roadmap manager",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
engineering_lead = Agent(
|
||||
agent_name="Engineering-Lead",
|
||||
agent_description="Senior engineering lead and technical architect",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
operations_manager = Agent(
|
||||
agent_name="Operations-Manager",
|
||||
agent_description="Operations and implementation manager",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
compliance_officer = Agent(
|
||||
agent_name="Compliance-Officer",
|
||||
agent_description="Legal compliance and regulatory specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create comprehensive hierarchical swarm
|
||||
swarm = HierarchicalSwarm(
|
||||
name="Enterprise-Strategy-Swarm",
|
||||
description="Enterprise-level strategic planning and execution swarm",
|
||||
agents=[
|
||||
market_intelligence,
|
||||
product_strategy,
|
||||
engineering_lead,
|
||||
operations_manager,
|
||||
compliance_officer,
|
||||
],
|
||||
max_loops=2,
|
||||
verbose=True,
|
||||
add_collaboration_prompt=True,
|
||||
)
|
||||
|
||||
# Test with complex enterprise scenario
|
||||
result = swarm.run(
|
||||
"Develop a comprehensive 5-year strategic plan for our company to become a leader in "
|
||||
"AI-powered enterprise solutions. Consider market opportunities, competitive landscape, "
|
||||
"technical requirements, operational capabilities, and regulatory compliance."
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
@ -1,84 +1,268 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from swarms.structs.mixture_of_agents import MixtureOfAgents
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms_memory import BaseVectorDatabase
|
||||
|
||||
|
||||
def test_init():
|
||||
with patch.object(
|
||||
MixtureOfAgents, "agent_check"
|
||||
) as mock_agent_check, patch.object(
|
||||
MixtureOfAgents, "final_agent_check"
|
||||
) as mock_final_agent_check, patch.object(
|
||||
MixtureOfAgents, "swarm_initialization"
|
||||
) as mock_swarm_initialization, patch.object(
|
||||
MixtureOfAgents, "communication_protocol"
|
||||
) as mock_communication_protocol:
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
scp = Mock(spec=BaseVectorDatabase)
|
||||
MixtureOfAgents(
|
||||
agents=agents, final_agent=final_agent, scp=scp
|
||||
)
|
||||
mock_agent_check.assert_called_once()
|
||||
mock_final_agent_check.assert_called_once()
|
||||
mock_swarm_initialization.assert_called_once()
|
||||
mock_communication_protocol.assert_called_once()
|
||||
|
||||
|
||||
def test_communication_protocol():
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
scp = Mock(spec=BaseVectorDatabase)
|
||||
swarm = MixtureOfAgents(
|
||||
agents=agents, final_agent=final_agent, scp=scp
|
||||
)
|
||||
swarm.communication_protocol()
|
||||
for agent in agents:
|
||||
agent.long_term_memory.assert_called_once_with(scp)
|
||||
|
||||
|
||||
def test_agent_check():
|
||||
final_agent = Mock(spec=Agent)
|
||||
with pytest.raises(TypeError):
|
||||
MixtureOfAgents(agents="not a list", final_agent=final_agent)
|
||||
with pytest.raises(TypeError):
|
||||
MixtureOfAgents(
|
||||
agents=["not an agent"], final_agent=final_agent
|
||||
)
|
||||
|
||||
|
||||
def test_final_agent_check():
|
||||
agents = [Mock(spec=Agent)]
|
||||
with pytest.raises(TypeError):
|
||||
MixtureOfAgents(agents=agents, final_agent="not an agent")
|
||||
def test_mixture_of_agents_basic_initialization():
|
||||
"""Test basic MixtureOfAgents initialization with multiple agents"""
|
||||
# Create multiple specialized agents
|
||||
research_agent = Agent(
|
||||
agent_name="Research-Specialist",
|
||||
agent_description="Specialist in research and data collection",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
analysis_agent = Agent(
|
||||
agent_name="Analysis-Expert",
|
||||
agent_description="Expert in data analysis and insights",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
strategy_agent = Agent(
|
||||
agent_name="Strategy-Consultant",
|
||||
agent_description="Strategy and planning consultant",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
def test_swarm_initialization():
|
||||
with patch(
|
||||
"swarms.structs.mixture_of_agents.logger"
|
||||
) as mock_logger:
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
swarm = MixtureOfAgents(
|
||||
agents=agents, final_agent=final_agent
|
||||
)
|
||||
swarm.swarm_initialization()
|
||||
assert mock_logger.info.call_count == 3
|
||||
|
||||
|
||||
def test_run():
|
||||
with patch("swarms.structs.mixture_of_agents.logger"), patch(
|
||||
"builtins.open", new_callable=Mock
|
||||
) as mock_open:
|
||||
agents = [Mock(spec=Agent)]
|
||||
final_agent = Mock(spec=Agent)
|
||||
swarm = MixtureOfAgents(
|
||||
agents=agents, final_agent=final_agent
|
||||
# Create aggregator agent
|
||||
aggregator = Agent(
|
||||
agent_name="Aggregator-Agent",
|
||||
agent_description="Agent that aggregates responses from other agents",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create mixture of agents
|
||||
moa = MixtureOfAgents(
|
||||
name="Business-Analysis-Mixture",
|
||||
description="Mixture of agents for comprehensive business analysis",
|
||||
agents=[research_agent, analysis_agent, strategy_agent],
|
||||
aggregator_agent=aggregator,
|
||||
layers=3,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Verify initialization
|
||||
assert moa.name == "Business-Analysis-Mixture"
|
||||
assert (
|
||||
moa.description
|
||||
== "Mixture of agents for comprehensive business analysis"
|
||||
)
|
||||
assert len(moa.agents) == 3
|
||||
assert moa.aggregator_agent == aggregator
|
||||
assert moa.layers == 3
|
||||
assert moa.max_loops == 1
|
||||
|
||||
|
||||
def test_mixture_of_agents_execution():
|
||||
"""Test MixtureOfAgents execution with multiple agents"""
|
||||
# Create diverse agents for different perspectives
|
||||
market_analyst = Agent(
|
||||
agent_name="Market-Analyst",
|
||||
agent_description="Market analysis and trend specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
technical_expert = Agent(
|
||||
agent_name="Technical-Expert",
|
||||
agent_description="Technical feasibility and implementation specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
financial_analyst = Agent(
|
||||
agent_name="Financial-Analyst",
|
||||
agent_description="Financial modeling and ROI specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
risk_assessor = Agent(
|
||||
agent_name="Risk-Assessor",
|
||||
agent_description="Risk assessment and mitigation specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create aggregator for synthesis
|
||||
aggregator = Agent(
|
||||
agent_name="Executive-Summary-Agent",
|
||||
agent_description="Executive summary and recommendation specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create mixture of agents
|
||||
moa = MixtureOfAgents(
|
||||
name="Comprehensive-Evaluation-Mixture",
|
||||
description="Mixture of agents for comprehensive business evaluation",
|
||||
agents=[
|
||||
market_analyst,
|
||||
technical_expert,
|
||||
financial_analyst,
|
||||
risk_assessor,
|
||||
],
|
||||
aggregator_agent=aggregator,
|
||||
layers=2,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Test execution
|
||||
result = moa.run(
|
||||
"Evaluate the feasibility of launching an AI-powered healthcare platform"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_mixture_of_agents_multiple_layers():
|
||||
"""Test MixtureOfAgents with multiple layers"""
|
||||
# Create agents for layered analysis
|
||||
data_collector = Agent(
|
||||
agent_name="Data-Collector",
|
||||
agent_description="Data collection and research specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
pattern_analyzer = Agent(
|
||||
agent_name="Pattern-Analyzer",
|
||||
agent_description="Pattern recognition and analysis specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
insight_generator = Agent(
|
||||
agent_name="Insight-Generator",
|
||||
agent_description="Insight generation and interpretation specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create aggregator
|
||||
final_aggregator = Agent(
|
||||
agent_name="Final-Aggregator",
|
||||
agent_description="Final aggregation and conclusion specialist",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create mixture with multiple layers for deeper analysis
|
||||
moa = MixtureOfAgents(
|
||||
name="Multi-Layer-Analysis-Mixture",
|
||||
description="Mixture of agents with multiple analysis layers",
|
||||
agents=[data_collector, pattern_analyzer, insight_generator],
|
||||
aggregator_agent=final_aggregator,
|
||||
layers=4,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Test multi-layer execution
|
||||
result = moa.run(
|
||||
"Analyze customer behavior patterns and provide strategic insights"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_mixture_of_agents_error_handling():
|
||||
"""Test MixtureOfAgents error handling and validation"""
|
||||
# Test with empty agents list
|
||||
try:
|
||||
moa = MixtureOfAgents(agents=[])
|
||||
assert (
|
||||
False
|
||||
), "Should have raised ValueError for empty agents list"
|
||||
except ValueError as e:
|
||||
assert "No agents provided" in str(e)
|
||||
|
||||
# Test with invalid aggregator system prompt
|
||||
analyst = Agent(
|
||||
agent_name="Test-Analyst",
|
||||
agent_description="Test analyst",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
try:
|
||||
moa = MixtureOfAgents(
|
||||
agents=[analyst], aggregator_system_prompt=""
|
||||
)
|
||||
swarm.run("task")
|
||||
for agent in agents:
|
||||
agent.run.assert_called_once()
|
||||
final_agent.run.assert_called_once()
|
||||
mock_open.assert_called_once_with(swarm.saved_file_name, "w")
|
||||
assert (
|
||||
False
|
||||
), "Should have raised ValueError for empty system prompt"
|
||||
except ValueError as e:
|
||||
assert "No aggregator system prompt" in str(e)
|
||||
|
||||
|
||||
def test_mixture_of_agents_real_world_scenario():
|
||||
"""Test MixtureOfAgents in a realistic business scenario"""
|
||||
# Create agents representing different business functions
|
||||
marketing_director = Agent(
|
||||
agent_name="Marketing-Director",
|
||||
agent_description="Senior marketing director with market expertise",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
product_manager = Agent(
|
||||
agent_name="Product-Manager",
|
||||
agent_description="Product strategy and development manager",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
engineering_lead = Agent(
|
||||
agent_name="Engineering-Lead",
|
||||
agent_description="Senior engineering and technical architecture lead",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
sales_executive = Agent(
|
||||
agent_name="Sales-Executive",
|
||||
agent_description="Enterprise sales and customer relationship executive",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
legal_counsel = Agent(
|
||||
agent_name="Legal-Counsel",
|
||||
agent_description="Legal compliance and regulatory counsel",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create aggregator for executive decision making
|
||||
executive_aggregator = Agent(
|
||||
agent_name="Executive-Decision-Maker",
|
||||
agent_description="Executive decision maker and strategic aggregator",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Create comprehensive mixture of agents
|
||||
moa = MixtureOfAgents(
|
||||
name="Executive-Board-Mixture",
|
||||
description="Mixture of agents representing executive board for strategic decisions",
|
||||
agents=[
|
||||
marketing_director,
|
||||
product_manager,
|
||||
engineering_lead,
|
||||
sales_executive,
|
||||
legal_counsel,
|
||||
],
|
||||
aggregator_agent=executive_aggregator,
|
||||
layers=3,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Test with complex business scenario
|
||||
result = moa.run(
|
||||
"Develop a comprehensive go-to-market strategy for our new AI-powered enterprise platform. "
|
||||
"Consider market positioning, technical requirements, competitive landscape, sales channels, "
|
||||
"and legal compliance requirements."
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
|
||||
@ -1,201 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from swarms import Agent
|
||||
from swarm_models import OpenAIChat
|
||||
from experimental.multi_agent_collab import MultiAgentCollaboration
|
||||
|
||||
# Initialize the director agent
|
||||
|
||||
director = Agent(
|
||||
agent_name="Director",
|
||||
system_prompt="Directs the tasks for the workers",
|
||||
llm=OpenAIChat(),
|
||||
max_loops=1,
|
||||
dashboard=False,
|
||||
streaming_on=True,
|
||||
verbose=True,
|
||||
stopping_token="<DONE>",
|
||||
state_save_file_type="json",
|
||||
saved_state_path="director.json",
|
||||
)
|
||||
|
||||
|
||||
# Initialize worker 1
|
||||
|
||||
worker1 = Agent(
|
||||
agent_name="Worker1",
|
||||
system_prompt="Generates a transcript for a youtube video on what swarms are",
|
||||
llm=OpenAIChat(),
|
||||
max_loops=1,
|
||||
dashboard=False,
|
||||
streaming_on=True,
|
||||
verbose=True,
|
||||
stopping_token="<DONE>",
|
||||
state_save_file_type="json",
|
||||
saved_state_path="worker1.json",
|
||||
)
|
||||
|
||||
|
||||
# Initialize worker 2
|
||||
worker2 = Agent(
|
||||
agent_name="Worker2",
|
||||
system_prompt="Summarizes the transcript generated by Worker1",
|
||||
llm=OpenAIChat(),
|
||||
max_loops=1,
|
||||
dashboard=False,
|
||||
streaming_on=True,
|
||||
verbose=True,
|
||||
stopping_token="<DONE>",
|
||||
state_save_file_type="json",
|
||||
saved_state_path="worker2.json",
|
||||
)
|
||||
|
||||
|
||||
# Create a list of agents
|
||||
agents = [director, worker1, worker2]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def collaboration():
|
||||
return MultiAgentCollaboration(agents)
|
||||
|
||||
|
||||
def test_collaboration_initialization(collaboration):
|
||||
assert len(collaboration.agents) == 2
|
||||
assert callable(collaboration.select_next_speaker)
|
||||
assert collaboration.max_loops == 10
|
||||
assert collaboration.results == []
|
||||
assert collaboration.logging is True
|
||||
|
||||
|
||||
def test_reset(collaboration):
|
||||
collaboration.reset()
|
||||
for agent in collaboration.agents:
|
||||
assert agent.step == 0
|
||||
|
||||
|
||||
def test_inject(collaboration):
|
||||
collaboration.inject("TestName", "TestMessage")
|
||||
for agent in collaboration.agents:
|
||||
assert "TestName" in agent.history[-1]
|
||||
assert "TestMessage" in agent.history[-1]
|
||||
|
||||
|
||||
def test_inject_agent(collaboration):
|
||||
agent3 = Agent(llm=OpenAIChat(), max_loops=2)
|
||||
collaboration.inject_agent(agent3)
|
||||
assert len(collaboration.agents) == 3
|
||||
assert agent3 in collaboration.agents
|
||||
|
||||
|
||||
def test_step(collaboration):
|
||||
collaboration.step()
|
||||
for agent in collaboration.agents:
|
||||
assert agent.step == 1
|
||||
|
||||
|
||||
def test_ask_for_bid(collaboration):
|
||||
agent = Mock()
|
||||
agent.bid.return_value = "<5>"
|
||||
bid = collaboration.ask_for_bid(agent)
|
||||
assert bid == 5
|
||||
|
||||
|
||||
def test_select_next_speaker(collaboration):
|
||||
collaboration.select_next_speaker = Mock(return_value=0)
|
||||
idx = collaboration.select_next_speaker(1, collaboration.agents)
|
||||
assert idx == 0
|
||||
|
||||
|
||||
def test_run(collaboration):
|
||||
collaboration.run()
|
||||
for agent in collaboration.agents:
|
||||
assert agent.step == collaboration.max_loops
|
||||
|
||||
|
||||
def test_format_results(collaboration):
|
||||
collaboration.results = [
|
||||
{"agent": "Agent1", "response": "Response1"}
|
||||
]
|
||||
formatted_results = collaboration.format_results(
|
||||
collaboration.results
|
||||
)
|
||||
assert "Agent1 responded: Response1" in formatted_results
|
||||
|
||||
|
||||
def test_save_and_load(collaboration):
|
||||
collaboration.save()
|
||||
loaded_state = collaboration.load()
|
||||
assert loaded_state["_step"] == collaboration._step
|
||||
assert loaded_state["results"] == collaboration.results
|
||||
|
||||
|
||||
def test_performance(collaboration):
|
||||
performance_data = collaboration.performance()
|
||||
for agent in collaboration.agents:
|
||||
assert agent.name in performance_data
|
||||
assert "metrics" in performance_data[agent.name]
|
||||
|
||||
|
||||
def test_set_interaction_rules(collaboration):
|
||||
rules = {"rule1": "action1", "rule2": "action2"}
|
||||
collaboration.set_interaction_rules(rules)
|
||||
assert hasattr(collaboration, "interaction_rules")
|
||||
assert collaboration.interaction_rules == rules
|
||||
|
||||
|
||||
def test_repr(collaboration):
|
||||
repr_str = repr(collaboration)
|
||||
assert isinstance(repr_str, str)
|
||||
assert "MultiAgentCollaboration" in repr_str
|
||||
|
||||
|
||||
def test_load(collaboration):
|
||||
state = {
|
||||
"step": 5,
|
||||
"results": [{"agent": "Agent1", "response": "Response1"}],
|
||||
}
|
||||
with open(collaboration.saved_file_path_name, "w") as file:
|
||||
json.dump(state, file)
|
||||
|
||||
loaded_state = collaboration.load()
|
||||
assert loaded_state["_step"] == state["step"]
|
||||
assert loaded_state["results"] == state["results"]
|
||||
|
||||
|
||||
def test_save(collaboration, tmp_path):
|
||||
collaboration.saved_file_path_name = tmp_path / "test_save.json"
|
||||
collaboration.save()
|
||||
|
||||
with open(collaboration.saved_file_path_name) as file:
|
||||
saved_data = json.load(file)
|
||||
|
||||
assert saved_data["_step"] == collaboration._step
|
||||
assert saved_data["results"] == collaboration.results
|
||||
|
||||
|
||||
# Add more tests here...
|
||||
|
||||
# Add more parameterized tests for different scenarios...
|
||||
|
||||
|
||||
# Example of exception testing
|
||||
def test_exception_handling(collaboration):
|
||||
agent = Mock()
|
||||
agent.bid.side_effect = ValueError("Invalid bid")
|
||||
with pytest.raises(ValueError):
|
||||
collaboration.ask_for_bid(agent)
|
||||
|
||||
|
||||
# Add more exception testing...
|
||||
|
||||
|
||||
# Example of environment variable testing (if applicable)
|
||||
@pytest.mark.parametrize("env_var", ["ENV_VAR_1", "ENV_VAR_2"])
|
||||
def test_environment_variables(collaboration, monkeypatch, env_var):
|
||||
monkeypatch.setenv(env_var, "test_value")
|
||||
assert os.getenv(env_var) == "test_value"
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue