Your Name
3 months ago
parent
38a6e6d5a5
commit
989e35aaa8
@ -0,0 +1,38 @@
|
||||
agents:
|
||||
- agent_name: "Financial-Analysis-Agent"
|
||||
model:
|
||||
openai_api_key: "your_openai_api_key"
|
||||
model_name: "gpt-4o-mini"
|
||||
temperature: 0.1
|
||||
max_tokens: 2000
|
||||
system_prompt: "financial_agent_sys_prompt" # Reference to system prompt file
|
||||
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" # Can be "json" or any other format
|
||||
|
||||
# - agent_name: "Stock-Analysis-Agent"
|
||||
# model:
|
||||
# openai_api_key: "your_openai_api_key"
|
||||
# model_name: "gpt-4o-mini"
|
||||
# temperature: 0.2
|
||||
# max_tokens: 1500
|
||||
# system_prompt: "stock_agent_sys_prompt" # Reference to system prompt file
|
||||
# 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: "str"
|
||||
@ -0,0 +1,20 @@
|
||||
from loguru import logger
|
||||
from dotenv import load_dotenv
|
||||
from swarms import create_agents_from_yaml
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Path to your YAML file
|
||||
yaml_file = 'agents_config.yaml'
|
||||
|
||||
try:
|
||||
# Create agents and run tasks (using 'both' to return agents and task results)
|
||||
agents, task_results = create_agents_from_yaml(yaml_file, return_type="both")
|
||||
|
||||
# Print the results of the tasks
|
||||
for result in task_results:
|
||||
print(f"Agent: {result['agent_name']} | Task: {result['task']} | Output: {result.get('output', 'Error encountered')}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"An error occurred: {e}")
|
||||
@ -0,0 +1,208 @@
|
||||
# `create_agents_from_yaml`
|
||||
|
||||
The `create_agents_from_yaml` function enables the dynamic creation and execution of agents based on configurations defined in a YAML file. This function is designed to support enterprise use-cases, offering flexibility, reliability, and scalability for various agent-based workflows.
|
||||
|
||||
By allowing the user to define multiple agents and tasks in a YAML configuration file, this function streamlines the process of initializing and executing tasks through agents while supporting advanced features such as multi-agent orchestration, logging, error handling, and flexible return values.
|
||||
|
||||
---
|
||||
|
||||
# Key Features
|
||||
- **Multi-Agent Creation**: Automatically create multiple agents based on a single YAML configuration file.
|
||||
- **Task Execution**: Each agent can execute a predefined task if specified in the YAML configuration.
|
||||
- **Logging with Loguru**: Integrated logging using `loguru` for robust, real-time tracking and error reporting.
|
||||
- **Dynamic Return Types**: Offers flexible return values (agents, tasks, or both) based on user needs.
|
||||
- **Error Handling**: Gracefully handles missing configurations, invalid inputs, and runtime errors.
|
||||
- **Extensibility**: Supports additional positional (`*args`) and keyword arguments (`**kwargs`) to customize agent behavior.
|
||||
|
||||
---
|
||||
|
||||
# Function Signature
|
||||
|
||||
```python
|
||||
def create_agents_from_yaml(yaml_file: str, return_type: str = "agents", *args, **kwargs)
|
||||
```
|
||||
|
||||
### Parameters:
|
||||
|
||||
- **`yaml_file: str`**
|
||||
- Description: The path to the YAML file containing agent configurations.
|
||||
- Required: Yes
|
||||
- Example: `'agents_config.yaml'`
|
||||
|
||||
- **`return_type: str`**
|
||||
- Description: Determines the type of data the function should return.
|
||||
- Options:
|
||||
- `"agents"`: Returns a list of the created agents.
|
||||
- `"tasks"`: Returns a list of task results (outputs or errors).
|
||||
- `"both"`: Returns both the list of agents and the task results as a tuple.
|
||||
- Default: `"agents"`
|
||||
- Required: No
|
||||
|
||||
- **`*args` and `**kwargs`**:
|
||||
- Description: Additional arguments to customize agent behavior. These can be passed through to the underlying `Agent` or `OpenAIChat` class constructors.
|
||||
- Required: No
|
||||
- Example: Can be used to modify model configurations or agent behavior dynamically.
|
||||
|
||||
### Returns:
|
||||
- **Based on `return_type`:**
|
||||
- `return_type="agents"`: Returns a list of initialized `Agent` objects.
|
||||
- `return_type="tasks"`: Returns a list of task results (success or error).
|
||||
- `return_type="both"`: Returns a tuple containing both the list of agents and task results.
|
||||
|
||||
---
|
||||
|
||||
# YAML Configuration Structure
|
||||
|
||||
The function relies on a YAML file for defining agents and tasks. Below is an example YAML configuration:
|
||||
|
||||
### Example YAML (agents_config.yaml):
|
||||
```yaml
|
||||
agents:
|
||||
- agent_name: "Financial-Analysis-Agent"
|
||||
model:
|
||||
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:
|
||||
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?"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
# Enterprise Use Cases
|
||||
|
||||
### 1. **Automating Financial Analysis**
|
||||
- An enterprise can use this function to create agents that analyze financial data in real-time. For example, an agent can be configured to provide financial advice based on the latest market trends, using predefined tasks in YAML to query the agent.
|
||||
|
||||
### 2. **Scalable Stock Analysis**
|
||||
- Multiple stock analysis agents can be created, each tasked with analyzing specific stocks or investment strategies. This setup can help enterprises handle large-scale financial modeling and stock analysis without manual intervention.
|
||||
|
||||
### 3. **Task Scheduling and Execution**
|
||||
- In enterprise operations, agents can be pre-configured with tasks such as risk assessment, regulatory compliance checks, or financial forecasting. The function automatically runs these tasks and returns actionable results or alerts.
|
||||
|
||||
---
|
||||
|
||||
# Example: Creating Agents and Running Tasks
|
||||
|
||||
### Example 1: Creating and Returning Agents
|
||||
```python
|
||||
from swarms import create_agents_from_yaml
|
||||
|
||||
yaml_file = 'agents_config.yaml'
|
||||
agents = create_agents_from_yaml(yaml_file, return_type="agents")
|
||||
|
||||
for agent in agents:
|
||||
print(f"Agent {agent.agent_name} created.")
|
||||
```
|
||||
|
||||
### Example 2: Creating Agents and Returning Task Results
|
||||
```python
|
||||
from swarms import create_agents_from_yaml
|
||||
|
||||
yaml_file = 'agents_config.yaml'
|
||||
task_results = create_agents_from_yaml(yaml_file, return_type="tasks")
|
||||
|
||||
for result in task_results:
|
||||
print(f"Agent {result['agent_name']} executed task '{result['task']}': {result['output']}")
|
||||
```
|
||||
|
||||
### Example 3: Returning Both Agents and Task Results
|
||||
```python
|
||||
from swarms import create_agents_from_yaml
|
||||
|
||||
yaml_file = 'agents_config.yaml'
|
||||
agents, task_results = create_agents_from_yaml(yaml_file, return_type="both")
|
||||
|
||||
# Handling agents
|
||||
for agent in agents:
|
||||
print(f"Agent {agent.agent_name} created.")
|
||||
|
||||
# Handling task results
|
||||
for result in task_results:
|
||||
print(f"Agent {result['agent_name']} executed task '{result['task']}': {result['output']}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
# Error Handling
|
||||
|
||||
### Common Errors:
|
||||
1. **Missing API Key**:
|
||||
- Error: `API key is missing for agent: <agent_name>`
|
||||
- Cause: The API key is either not provided in the YAML or not available as an environment variable.
|
||||
- Solution: Ensure the API key is either defined in the YAML configuration or set as an environment variable.
|
||||
|
||||
2. **Missing System Prompt**:
|
||||
- Error: `System prompt is missing for agent: <agent_name>`
|
||||
- Cause: The `system_prompt` field is not defined in the YAML configuration.
|
||||
- Solution: Define the system prompt field for each agent.
|
||||
|
||||
3. **Invalid `return_type`**:
|
||||
- Error: `Invalid return_type: <return_type>`
|
||||
- Cause: The `return_type` provided is not one of `"agents"`, `"tasks"`, or `"both"`.
|
||||
- Solution: Ensure that `return_type` is set to one of the valid options.
|
||||
|
||||
### Logging:
|
||||
- The function integrates `loguru` logging to track all key actions:
|
||||
- File loading, agent creation, task execution, and errors are all logged.
|
||||
- Use this logging output to monitor operations and diagnose issues in production environments.
|
||||
|
||||
---
|
||||
|
||||
# Scalability and Extensibility
|
||||
|
||||
### Scalability:
|
||||
The `create_agents_from_yaml` function is designed for use in high-scale enterprise environments:
|
||||
- **Multi-Agent Creation**: Create and manage large numbers of agents simultaneously.
|
||||
- **Parallel Task Execution**: Tasks can be executed in parallel for real-time analysis and decision-making across multiple business units.
|
||||
|
||||
### Extensibility:
|
||||
- **Customizable Behavior**: Through `*args` and `**kwargs`, users can extend the functionality of the agents or models without altering the core YAML configuration.
|
||||
- **Seamless Integration**: The function can be easily integrated with larger multi-agent systems and workflows, enabling rapid scaling across departments.
|
||||
|
||||
---
|
||||
|
||||
# Security Considerations
|
||||
|
||||
For enterprise deployments, consider the following security best practices:
|
||||
1. **API Key Management**: Ensure that API keys are stored securely (e.g., using environment variables or secret management tools).
|
||||
2. **Data Handling**: Be mindful of sensitive information within tasks or agent responses. Implement data sanitization where necessary.
|
||||
3. **Task Validation**: Validate tasks in the YAML file to ensure they meet your organization's security and operational policies before execution.
|
||||
|
||||
---
|
||||
|
||||
# Conclusion
|
||||
|
||||
The `create_agents_from_yaml` function is a powerful tool for enterprises to dynamically create, manage, and execute tasks using AI-powered agents. With its flexible configuration, logging, and error handling, this function is ideal for scaling agent-based systems and automating complex workflows across industries.
|
||||
|
||||
Integrating this function into your enterprise workflow will enhance efficiency, provide real-time insights, and reduce operational overhead.
|
||||
@ -0,0 +1,148 @@
|
||||
import os
|
||||
import yaml
|
||||
from loguru import logger
|
||||
from swarms.structs.agent import Agent
|
||||
from swarm_models import OpenAIChat
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Function to create and optionally run agents from a YAML file
|
||||
def create_agents_from_yaml(yaml_file: str, return_type: str = "agents", *args, **kwargs):
|
||||
"""
|
||||
Create agents based on configurations defined in a YAML file.
|
||||
If a 'task' is provided in the YAML, the agent will execute the task after creation.
|
||||
|
||||
Args:
|
||||
yaml_file (str): Path to the YAML file containing agent configurations.
|
||||
return_type (str): Determines the return value. "agents" to return agent list,
|
||||
"tasks" to return task results, "both" to return both agents and tasks.
|
||||
*args: Additional positional arguments for agent customization.
|
||||
**kwargs: Additional keyword arguments for agent customization.
|
||||
|
||||
Returns:
|
||||
List[Agent] or List[Task Results] or Tuple(List[Agent], List[Task Results])
|
||||
"""
|
||||
logger.info(f"Checking if the YAML file {yaml_file} exists...")
|
||||
|
||||
# Check if the YAML file exists
|
||||
if not os.path.exists(yaml_file):
|
||||
logger.error(f"YAML file {yaml_file} not found.")
|
||||
raise FileNotFoundError(f"YAML file {yaml_file} not found.")
|
||||
|
||||
# Load the YAML configuration
|
||||
logger.info(f"Loading YAML file {yaml_file}")
|
||||
with open(yaml_file, 'r') as file:
|
||||
config = yaml.safe_load(file)
|
||||
|
||||
# Ensure agents key exists
|
||||
if "agents" not in config:
|
||||
logger.error("The YAML configuration does not contain 'agents'.")
|
||||
raise ValueError("The YAML configuration does not contain 'agents'.")
|
||||
|
||||
# List to store created agents and task results
|
||||
agents = []
|
||||
task_results = []
|
||||
|
||||
# Iterate over each agent configuration and create agents
|
||||
for agent_config in config["agents"]:
|
||||
logger.info(f"Creating agent: {agent_config['agent_name']}")
|
||||
|
||||
# Get the OpenAI API key from environment or YAML config
|
||||
api_key = os.getenv("OPENAI_API_KEY") or agent_config["model"].get("openai_api_key")
|
||||
if not api_key:
|
||||
logger.error(f"API key is missing for agent: {agent_config['agent_name']}")
|
||||
raise ValueError(f"API key is missing for agent: {agent_config['agent_name']}")
|
||||
|
||||
# Create an instance of OpenAIChat model
|
||||
model = OpenAIChat(
|
||||
openai_api_key=api_key,
|
||||
model_name=agent_config["model"]["model_name"],
|
||||
temperature=agent_config["model"]["temperature"],
|
||||
max_tokens=agent_config["model"]["max_tokens"],
|
||||
*args, **kwargs # Pass any additional arguments to the model
|
||||
)
|
||||
|
||||
# Ensure the system prompt is provided
|
||||
if "system_prompt" not in agent_config:
|
||||
logger.error(f"System prompt is missing for agent: {agent_config['agent_name']}")
|
||||
raise ValueError(f"System prompt is missing for agent: {agent_config['agent_name']}")
|
||||
|
||||
# Dynamically choose the system prompt based on the agent config
|
||||
try:
|
||||
system_prompt = globals().get(agent_config["system_prompt"])
|
||||
if not system_prompt:
|
||||
logger.error(f"System prompt {agent_config['system_prompt']} not found.")
|
||||
raise ValueError(f"System prompt {agent_config['system_prompt']} not found.")
|
||||
except KeyError:
|
||||
logger.error(f"System prompt {agent_config['system_prompt']} is not valid.")
|
||||
raise ValueError(f"System prompt {agent_config['system_prompt']} is not valid.")
|
||||
|
||||
# Initialize the agent using the configuration
|
||||
agent = Agent(
|
||||
agent_name=agent_config["agent_name"],
|
||||
system_prompt=system_prompt,
|
||||
llm=model,
|
||||
max_loops=agent_config.get("max_loops", 1),
|
||||
autosave=agent_config.get("autosave", True),
|
||||
dashboard=agent_config.get("dashboard", False),
|
||||
verbose=agent_config.get("verbose", False),
|
||||
dynamic_temperature_enabled=agent_config.get("dynamic_temperature_enabled", False),
|
||||
saved_state_path=agent_config.get("saved_state_path"),
|
||||
user_name=agent_config.get("user_name", "default_user"),
|
||||
retry_attempts=agent_config.get("retry_attempts", 1),
|
||||
context_length=agent_config.get("context_length", 100000),
|
||||
return_step_meta=agent_config.get("return_step_meta", False),
|
||||
output_type=agent_config.get("output_type", "str"),
|
||||
*args, **kwargs # Pass any additional arguments to the agent
|
||||
)
|
||||
|
||||
logger.info(f"Agent {agent_config['agent_name']} created successfully.")
|
||||
agents.append(agent)
|
||||
|
||||
# Check if a task is provided, and if so, run the agent
|
||||
task = agent_config.get("task")
|
||||
if task:
|
||||
logger.info(f"Running task '{task}' with agent {agent_config['agent_name']}")
|
||||
try:
|
||||
output = agent.run(task)
|
||||
logger.info(f"Output for agent {agent_config['agent_name']}: {output}")
|
||||
task_results.append({
|
||||
"agent_name": agent_config["agent_name"],
|
||||
"task": task,
|
||||
"output": output
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Error running task for agent {agent_config['agent_name']}: {e}")
|
||||
task_results.append({
|
||||
"agent_name": agent_config["agent_name"],
|
||||
"task": task,
|
||||
"error": str(e)
|
||||
})
|
||||
|
||||
# Return results based on the `return_type`
|
||||
if return_type == "agents":
|
||||
return agents
|
||||
elif return_type == "tasks":
|
||||
return task_results
|
||||
elif return_type == "both":
|
||||
return agents, task_results
|
||||
else:
|
||||
logger.error(f"Invalid return_type: {return_type}")
|
||||
raise ValueError(f"Invalid return_type: {return_type}")
|
||||
|
||||
# # Usage example
|
||||
# yaml_file = 'agents_config.yaml'
|
||||
|
||||
# try:
|
||||
# # Auto-create agents from the YAML file and return both agents and task results
|
||||
# agents, task_results = create_agents_from_yaml(yaml_file, return_type="tasks")
|
||||
|
||||
# # Example: Print task results
|
||||
# for result in task_results:
|
||||
# print(f"Agent: {result['agent_name']} | Task: {result['task']} | Output: {result.get('output', 'Error encountered')}")
|
||||
|
||||
# except FileNotFoundError as e:
|
||||
# logger.error(f"Error: {e}")
|
||||
# except ValueError as e:
|
||||
# logger.error(f"Error: {e}")
|
||||
@ -0,0 +1,216 @@
|
||||
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()
|
||||
Loading…
Reference in new issue