Kye Gomez
4 days ago
parent
7a7b4a4c07
commit
63fdc5cfb4
@ -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']}")
|
||||
@ -0,0 +1,64 @@
|
||||
import pytest
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.multi_agent_router import MultiAgentRouter
|
||||
|
||||
|
||||
# Agents fixture
|
||||
@pytest.fixture
|
||||
def agents():
|
||||
return [
|
||||
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,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
"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",
|
||||
],
|
||||
)
|
||||
def test_multiagentrouter_models(agents, model_name):
|
||||
"""
|
||||
Run MultiAgentRouter on a variety of models to ensure no errors are raised.
|
||||
"""
|
||||
task = "Use all the agents available to you to remake the Fibonacci function in Python, providing both an explanation and code."
|
||||
router_execute = MultiAgentRouter(
|
||||
agents=agents,
|
||||
temperature=0.5,
|
||||
model=model_name,
|
||||
)
|
||||
try:
|
||||
result = router_execute.run(task)
|
||||
assert result is not None
|
||||
except Exception as e:
|
||||
pytest.fail(f"Model {model_name} raised exception: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main(args=[__file__])
|
||||
@ -0,0 +1,897 @@
|
||||
"""
|
||||
Comprehensive unit test suite for SwarmRouter class.
|
||||
|
||||
This test suite covers all functions, swarm types, and ensures outputs are not None
|
||||
to verify successful execution.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
from typing import List, Dict, Any
|
||||
import concurrent.futures
|
||||
import json
|
||||
|
||||
from swarms.structs.swarm_router import (
|
||||
SwarmRouter,
|
||||
SwarmRouterConfig,
|
||||
SwarmRouterRunError,
|
||||
SwarmRouterConfigError,
|
||||
SwarmType,
|
||||
Document,
|
||||
)
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.utils.output_types import OutputType
|
||||
|
||||
|
||||
class TestSwarmRouterInitialization:
|
||||
"""Test SwarmRouter initialization with various configurations."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agents(self):
|
||||
"""Create sample agents for testing."""
|
||||
return [
|
||||
Agent(
|
||||
agent_name="ResearchAgent",
|
||||
agent_description="Specializes in researching topics",
|
||||
system_prompt="You are a research specialist.",
|
||||
max_loops=1,
|
||||
),
|
||||
Agent(
|
||||
agent_name="CodeAgent",
|
||||
agent_description="Expert in coding",
|
||||
system_prompt="You are a coding expert.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
def test_default_initialization(self):
|
||||
"""Test SwarmRouter with default parameters."""
|
||||
router = SwarmRouter()
|
||||
|
||||
assert router.name == "swarm-router"
|
||||
assert router.description == "Routes your task to the desired swarm"
|
||||
assert router.max_loops == 1
|
||||
assert router.agents == []
|
||||
assert router.swarm_type == "SequentialWorkflow"
|
||||
assert router.autosave is False
|
||||
assert router.return_json is False
|
||||
assert router.auto_generate_prompts is False
|
||||
assert router.shared_memory_system is None
|
||||
assert router.rules is None
|
||||
assert router.documents == []
|
||||
assert router.output_type == "dict-all-except-first"
|
||||
assert router.verbose is False
|
||||
assert router.telemetry_enabled is False
|
||||
|
||||
def test_custom_initialization(self, sample_agents):
|
||||
"""Test SwarmRouter with custom parameters."""
|
||||
router = SwarmRouter(
|
||||
name="test-router",
|
||||
description="Test router description",
|
||||
max_loops=3,
|
||||
agents=sample_agents,
|
||||
swarm_type="ConcurrentWorkflow",
|
||||
autosave=True,
|
||||
return_json=True,
|
||||
auto_generate_prompts=True,
|
||||
rules="Test rules",
|
||||
documents=["doc1.txt", "doc2.txt"],
|
||||
output_type="json",
|
||||
verbose=True,
|
||||
telemetry_enabled=True,
|
||||
)
|
||||
|
||||
assert router.name == "test-router"
|
||||
assert router.description == "Test router description"
|
||||
assert router.max_loops == 3
|
||||
assert router.agents == sample_agents
|
||||
assert router.swarm_type == "ConcurrentWorkflow"
|
||||
assert router.autosave is True
|
||||
assert router.return_json is True
|
||||
assert router.auto_generate_prompts is True
|
||||
assert router.rules == "Test rules"
|
||||
assert router.documents == ["doc1.txt", "doc2.txt"]
|
||||
assert router.output_type == "json"
|
||||
assert router.verbose is True
|
||||
assert router.telemetry_enabled is True
|
||||
|
||||
def test_initialization_with_heavy_swarm_config(self, sample_agents):
|
||||
"""Test SwarmRouter with HeavySwarm specific configuration."""
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type="HeavySwarm",
|
||||
heavy_swarm_loops_per_agent=2,
|
||||
heavy_swarm_question_agent_model_name="gpt-4",
|
||||
heavy_swarm_worker_model_name="gpt-3.5-turbo",
|
||||
heavy_swarm_swarm_show_output=False,
|
||||
)
|
||||
|
||||
assert router.swarm_type == "HeavySwarm"
|
||||
assert router.heavy_swarm_loops_per_agent == 2
|
||||
assert router.heavy_swarm_question_agent_model_name == "gpt-4"
|
||||
assert router.heavy_swarm_worker_model_name == "gpt-3.5-turbo"
|
||||
assert router.heavy_swarm_swarm_show_output is False
|
||||
|
||||
def test_initialization_with_council_judge_config(self):
|
||||
"""Test SwarmRouter with CouncilAsAJudge specific configuration."""
|
||||
router = SwarmRouter(
|
||||
swarm_type="CouncilAsAJudge",
|
||||
council_judge_model_name="gpt-4o",
|
||||
)
|
||||
|
||||
assert router.swarm_type == "CouncilAsAJudge"
|
||||
assert router.council_judge_model_name == "gpt-4o"
|
||||
|
||||
def test_initialization_with_agent_rearrange_flow(self, sample_agents):
|
||||
"""Test SwarmRouter with AgentRearrange and flow configuration."""
|
||||
flow = "agent1 -> agent2 -> agent1"
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type="AgentRearrange",
|
||||
rearrange_flow=flow,
|
||||
)
|
||||
|
||||
assert router.swarm_type == "AgentRearrange"
|
||||
assert router.rearrange_flow == flow
|
||||
|
||||
def test_initialization_with_shared_memory(self, sample_agents):
|
||||
"""Test SwarmRouter with shared memory system."""
|
||||
mock_memory = Mock()
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
shared_memory_system=mock_memory,
|
||||
)
|
||||
|
||||
assert router.shared_memory_system == mock_memory
|
||||
|
||||
def test_initialization_with_worker_tools(self, sample_agents):
|
||||
"""Test SwarmRouter with worker tools."""
|
||||
mock_tool = Mock()
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
worker_tools=[mock_tool],
|
||||
)
|
||||
|
||||
assert router.worker_tools == [mock_tool]
|
||||
|
||||
|
||||
class TestSwarmRouterConfigurationValidation:
|
||||
"""Test configuration validation and error handling."""
|
||||
|
||||
def test_invalid_swarm_type(self):
|
||||
"""Test error when invalid swarm type is provided."""
|
||||
with pytest.raises(ValueError):
|
||||
SwarmRouter(swarm_type="InvalidSwarmType")
|
||||
|
||||
def test_no_agents_for_swarm_requiring_agents(self):
|
||||
"""Test error when no agents provided for swarm requiring agents."""
|
||||
with pytest.raises(SwarmRouterConfigError):
|
||||
SwarmRouter(swarm_type="SequentialWorkflow", agents=None)
|
||||
|
||||
def test_no_rearrange_flow_for_agent_rearrange(self):
|
||||
"""Test error when no rearrange_flow provided for AgentRearrange."""
|
||||
agents = [Agent(agent_name="test", agent_description="test")]
|
||||
with pytest.raises(SwarmRouterConfigError):
|
||||
SwarmRouter(
|
||||
agents=agents,
|
||||
swarm_type="AgentRearrange",
|
||||
rearrange_flow=None,
|
||||
)
|
||||
|
||||
def test_zero_max_loops(self):
|
||||
"""Test error when max_loops is 0."""
|
||||
with pytest.raises(SwarmRouterConfigError):
|
||||
SwarmRouter(max_loops=0)
|
||||
|
||||
def test_heavy_swarm_without_agents(self):
|
||||
"""Test HeavySwarm can be created without agents."""
|
||||
router = SwarmRouter(swarm_type="HeavySwarm", agents=None)
|
||||
assert router.swarm_type == "HeavySwarm"
|
||||
|
||||
def test_council_judge_without_agents(self):
|
||||
"""Test CouncilAsAJudge can be created without agents."""
|
||||
router = SwarmRouter(swarm_type="CouncilAsAJudge", agents=None)
|
||||
assert router.swarm_type == "CouncilAsAJudge"
|
||||
|
||||
|
||||
class TestSwarmFactoryMethods:
|
||||
"""Test all factory methods for creating different swarm types."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agents(self):
|
||||
"""Create sample agents for testing."""
|
||||
return [
|
||||
Agent(
|
||||
agent_name="Agent1",
|
||||
agent_description="First agent",
|
||||
system_prompt="You are agent 1.",
|
||||
max_loops=1,
|
||||
),
|
||||
Agent(
|
||||
agent_name="Agent2",
|
||||
agent_description="Second agent",
|
||||
system_prompt="You are agent 2.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
def test_swarm_factory_initialization(self, sample_agents):
|
||||
"""Test that swarm factory is properly initialized."""
|
||||
router = SwarmRouter(agents=sample_agents)
|
||||
factory = router._initialize_swarm_factory()
|
||||
|
||||
expected_types = [
|
||||
"HeavySwarm", "AgentRearrange", "MALT", "CouncilAsAJudge",
|
||||
"InteractiveGroupChat", "HiearchicalSwarm", "MixtureOfAgents",
|
||||
"MajorityVoting", "GroupChat", "MultiAgentRouter",
|
||||
"SequentialWorkflow", "ConcurrentWorkflow", "BatchedGridWorkflow"
|
||||
]
|
||||
|
||||
for swarm_type in expected_types:
|
||||
assert swarm_type in factory
|
||||
assert callable(factory[swarm_type])
|
||||
|
||||
def test_create_heavy_swarm(self, sample_agents):
|
||||
"""Test HeavySwarm creation."""
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type="HeavySwarm",
|
||||
heavy_swarm_loops_per_agent=2,
|
||||
heavy_swarm_question_agent_model_name="gpt-4",
|
||||
heavy_swarm_worker_model_name="gpt-3.5-turbo",
|
||||
)
|
||||
|
||||
swarm = router._create_heavy_swarm()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_agent_rearrange(self, sample_agents):
|
||||
"""Test AgentRearrange creation."""
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type="AgentRearrange",
|
||||
rearrange_flow="agent1 -> agent2",
|
||||
)
|
||||
|
||||
swarm = router._create_agent_rearrange()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_sequential_workflow(self, sample_agents):
|
||||
"""Test SequentialWorkflow creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
swarm = router._create_sequential_workflow()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_concurrent_workflow(self, sample_agents):
|
||||
"""Test ConcurrentWorkflow creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="ConcurrentWorkflow")
|
||||
|
||||
swarm = router._create_concurrent_workflow()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_group_chat(self, sample_agents):
|
||||
"""Test GroupChat creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="GroupChat")
|
||||
|
||||
swarm = router._create_group_chat()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_multi_agent_router(self, sample_agents):
|
||||
"""Test MultiAgentRouter creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="MultiAgentRouter")
|
||||
|
||||
swarm = router._create_multi_agent_router()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_mixture_of_agents(self, sample_agents):
|
||||
"""Test MixtureOfAgents creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="MixtureOfAgents")
|
||||
|
||||
swarm = router._create_mixture_of_agents()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_majority_voting(self, sample_agents):
|
||||
"""Test MajorityVoting creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="MajorityVoting")
|
||||
|
||||
swarm = router._create_majority_voting()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_malt(self, sample_agents):
|
||||
"""Test MALT creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="MALT")
|
||||
|
||||
swarm = router._create_malt()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_council_as_judge(self):
|
||||
"""Test CouncilAsAJudge creation."""
|
||||
router = SwarmRouter(swarm_type="CouncilAsAJudge")
|
||||
|
||||
swarm = router._create_council_as_judge()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_interactive_group_chat(self, sample_agents):
|
||||
"""Test InteractiveGroupChat creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="InteractiveGroupChat")
|
||||
|
||||
swarm = router._create_interactive_group_chat()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_hierarchical_swarm(self, sample_agents):
|
||||
"""Test HierarchicalSwarm creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="HiearchicalSwarm")
|
||||
|
||||
swarm = router._create_hierarchical_swarm()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
def test_create_batched_grid_workflow(self, sample_agents):
|
||||
"""Test BatchedGridWorkflow creation."""
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="BatchedGridWorkflow")
|
||||
|
||||
swarm = router._create_batched_grid_workflow()
|
||||
assert swarm is not None
|
||||
assert hasattr(swarm, 'run')
|
||||
|
||||
|
||||
class TestSwarmRouterSwarmTypes:
|
||||
"""Test all available swarm types with actual execution."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agents(self):
|
||||
"""Create sample agents for testing."""
|
||||
return [
|
||||
Agent(
|
||||
agent_name="ResearchAgent",
|
||||
agent_description="Specializes in researching topics",
|
||||
system_prompt="You are a research specialist. Provide detailed, well-researched information.",
|
||||
max_loops=1,
|
||||
),
|
||||
Agent(
|
||||
agent_name="CodeAgent",
|
||||
agent_description="Expert in coding",
|
||||
system_prompt="You are a coding expert. Write clean, efficient code.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("swarm_type", [
|
||||
"SequentialWorkflow",
|
||||
"ConcurrentWorkflow",
|
||||
"GroupChat",
|
||||
"MultiAgentRouter",
|
||||
"MixtureOfAgents",
|
||||
"MajorityVoting",
|
||||
"MALT",
|
||||
"CouncilAsAJudge",
|
||||
"InteractiveGroupChat",
|
||||
"HiearchicalSwarm",
|
||||
"BatchedGridWorkflow",
|
||||
])
|
||||
def test_swarm_types_execution(self, sample_agents, swarm_type):
|
||||
"""Test execution of all swarm types with mock LLM."""
|
||||
with patch('swarms.structs.agent.LiteLLM') as mock_llm:
|
||||
# Mock the LLM to return a simple response
|
||||
mock_llm_instance = Mock()
|
||||
mock_llm_instance.agenerate.return_value = "Test response from agent"
|
||||
mock_llm.return_value = mock_llm_instance
|
||||
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type=swarm_type,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
# Test with a simple task
|
||||
task = "Write a simple Python function to add two numbers"
|
||||
|
||||
try:
|
||||
result = router.run(task)
|
||||
assert result is not None, f"Swarm type {swarm_type} returned None result"
|
||||
except Exception as e:
|
||||
# Some swarm types might have specific requirements
|
||||
if swarm_type in ["AgentRearrange"]:
|
||||
# AgentRearrange requires rearrange_flow
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type=swarm_type,
|
||||
rearrange_flow="agent1 -> agent2",
|
||||
max_loops=1,
|
||||
)
|
||||
result = router.run(task)
|
||||
assert result is not None, f"Swarm type {swarm_type} returned None result"
|
||||
|
||||
def test_heavy_swarm_execution(self):
|
||||
"""Test HeavySwarm execution."""
|
||||
with patch('swarms.structs.heavy_swarm.HeavySwarm') as mock_heavy_swarm:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "HeavySwarm response"
|
||||
mock_heavy_swarm.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(swarm_type="HeavySwarm")
|
||||
|
||||
result = router.run("Test task")
|
||||
assert result is not None
|
||||
assert result == "HeavySwarm response"
|
||||
|
||||
def test_agent_rearrange_execution(self, sample_agents):
|
||||
"""Test AgentRearrange execution with flow."""
|
||||
with patch('swarms.structs.agent_rearrange.AgentRearrange') as mock_agent_rearrange:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "AgentRearrange response"
|
||||
mock_agent_rearrange.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type="AgentRearrange",
|
||||
rearrange_flow="agent1 -> agent2",
|
||||
)
|
||||
|
||||
result = router.run("Test task")
|
||||
assert result is not None
|
||||
assert result == "AgentRearrange response"
|
||||
|
||||
|
||||
class TestSwarmRouterRunMethods:
|
||||
"""Test run, batch_run, concurrent_run methods."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agents(self):
|
||||
"""Create sample agents for testing."""
|
||||
return [
|
||||
Agent(
|
||||
agent_name="TestAgent",
|
||||
agent_description="Test agent",
|
||||
system_prompt="You are a test agent.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
def test_run_method(self, sample_agents):
|
||||
"""Test basic run method."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Test response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
result = router.run("Test task")
|
||||
assert result is not None
|
||||
assert result == "Test response"
|
||||
|
||||
def test_run_with_image(self, sample_agents):
|
||||
"""Test run method with image input."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Image analysis response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
result = router.run("Analyze this image", img="test_image.jpg")
|
||||
assert result is not None
|
||||
assert result == "Image analysis response"
|
||||
|
||||
def test_run_with_tasks_list(self, sample_agents):
|
||||
"""Test run method with tasks list."""
|
||||
with patch('swarms.structs.batched_grid_workflow.BatchedGridWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Batch response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="BatchedGridWorkflow")
|
||||
|
||||
tasks = ["Task 1", "Task 2", "Task 3"]
|
||||
result = router.run(tasks=tasks)
|
||||
assert result is not None
|
||||
assert result == "Batch response"
|
||||
|
||||
def test_batch_run_method(self, sample_agents):
|
||||
"""Test batch_run method."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Batch task response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
tasks = ["Task 1", "Task 2", "Task 3"]
|
||||
results = router.batch_run(tasks)
|
||||
|
||||
assert results is not None
|
||||
assert len(results) == 3
|
||||
assert all(result is not None for result in results)
|
||||
|
||||
def test_concurrent_run_method(self, sample_agents):
|
||||
"""Test concurrent_run method."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Concurrent response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
result = router.concurrent_run("Test task")
|
||||
assert result is not None
|
||||
assert result == "Concurrent response"
|
||||
|
||||
def test_call_method(self, sample_agents):
|
||||
"""Test __call__ method."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Call response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
result = router("Test task")
|
||||
assert result is not None
|
||||
assert result == "Call response"
|
||||
|
||||
def test_call_with_image(self, sample_agents):
|
||||
"""Test __call__ method with image."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Image call response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
result = router("Test task", img="test.jpg")
|
||||
assert result is not None
|
||||
assert result == "Image call response"
|
||||
|
||||
def test_call_with_images_list(self, sample_agents):
|
||||
"""Test __call__ method with images list."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Images call response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
result = router("Test task", imgs=["test1.jpg", "test2.jpg"])
|
||||
assert result is not None
|
||||
assert result == "Images call response"
|
||||
|
||||
|
||||
class TestSwarmRouterUtilityMethods:
|
||||
"""Test utility methods like to_dict, activate_ape, handle_rules, etc."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agents(self):
|
||||
"""Create sample agents for testing."""
|
||||
return [
|
||||
Agent(
|
||||
agent_name="TestAgent",
|
||||
agent_description="Test agent",
|
||||
system_prompt="You are a test agent.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
def test_to_dict_method(self, sample_agents):
|
||||
"""Test to_dict method serialization."""
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
name="test-router",
|
||||
description="Test description",
|
||||
swarm_type="SequentialWorkflow",
|
||||
)
|
||||
|
||||
result_dict = router.to_dict()
|
||||
|
||||
assert isinstance(result_dict, dict)
|
||||
assert result_dict["name"] == "test-router"
|
||||
assert result_dict["description"] == "Test description"
|
||||
assert result_dict["swarm_type"] == "SequentialWorkflow"
|
||||
assert "agents" in result_dict
|
||||
|
||||
def test_activate_ape(self, sample_agents):
|
||||
"""Test activate_ape method."""
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
auto_generate_prompts=True,
|
||||
)
|
||||
|
||||
# Mock the auto_generate_prompt attribute
|
||||
for agent in router.agents:
|
||||
agent.auto_generate_prompt = False
|
||||
|
||||
router.activate_ape()
|
||||
|
||||
# Check that APE was activated for agents that support it
|
||||
for agent in router.agents:
|
||||
if hasattr(agent, "auto_generate_prompt"):
|
||||
assert agent.auto_generate_prompt is True
|
||||
|
||||
def test_handle_rules(self, sample_agents):
|
||||
"""Test handle_rules method."""
|
||||
rules = "Always be helpful and accurate."
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
original_prompts = [agent.system_prompt for agent in router.agents]
|
||||
router.handle_rules()
|
||||
|
||||
# Check that rules were added to system prompts
|
||||
for i, agent in enumerate(router.agents):
|
||||
assert rules in agent.system_prompt
|
||||
assert agent.system_prompt == original_prompts[i] + f"### Swarm Rules ### {rules}"
|
||||
|
||||
def test_activate_shared_memory(self, sample_agents):
|
||||
"""Test activate_shared_memory method."""
|
||||
mock_memory = Mock()
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
shared_memory_system=mock_memory,
|
||||
)
|
||||
|
||||
router.activate_shared_memory()
|
||||
|
||||
# Check that all agents have the shared memory system
|
||||
for agent in router.agents:
|
||||
assert agent.long_term_memory == mock_memory
|
||||
|
||||
def test_update_system_prompt_for_agent_in_swarm(self, sample_agents):
|
||||
"""Test update_system_prompt_for_agent_in_swarm method."""
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
multi_agent_collab_prompt=True,
|
||||
)
|
||||
|
||||
original_prompts = [agent.system_prompt for agent in router.agents]
|
||||
router.update_system_prompt_for_agent_in_swarm()
|
||||
|
||||
# Check that collaboration prompt was added
|
||||
for i, agent in enumerate(router.agents):
|
||||
assert len(agent.system_prompt) > len(original_prompts[i])
|
||||
|
||||
def test_agent_config(self, sample_agents):
|
||||
"""Test agent_config method."""
|
||||
router = SwarmRouter(agents=sample_agents)
|
||||
|
||||
config = router.agent_config()
|
||||
|
||||
assert isinstance(config, dict)
|
||||
assert len(config) == len(sample_agents)
|
||||
|
||||
for agent in sample_agents:
|
||||
assert agent.agent_name in config
|
||||
|
||||
def test_fetch_message_history_as_string(self, sample_agents):
|
||||
"""Test fetch_message_history_as_string method."""
|
||||
router = SwarmRouter(agents=sample_agents)
|
||||
|
||||
# Mock the swarm and conversation
|
||||
mock_conversation = Mock()
|
||||
mock_conversation.return_all_except_first_string.return_value = "Test history"
|
||||
router.swarm = Mock()
|
||||
router.swarm.conversation = mock_conversation
|
||||
|
||||
result = router.fetch_message_history_as_string()
|
||||
assert result == "Test history"
|
||||
|
||||
def test_fetch_message_history_as_string_error(self, sample_agents):
|
||||
"""Test fetch_message_history_as_string method with error."""
|
||||
router = SwarmRouter(agents=sample_agents)
|
||||
|
||||
# Mock the swarm to raise an exception
|
||||
router.swarm = Mock()
|
||||
router.swarm.conversation.return_all_except_first_string.side_effect = Exception("Test error")
|
||||
|
||||
result = router.fetch_message_history_as_string()
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestSwarmRouterErrorHandling:
|
||||
"""Test error handling and exception scenarios."""
|
||||
|
||||
def test_swarm_creation_error(self):
|
||||
"""Test error handling when swarm creation fails."""
|
||||
router = SwarmRouter(swarm_type="SequentialWorkflow", agents=None)
|
||||
|
||||
with pytest.raises(SwarmRouterConfigError):
|
||||
router.run("Test task")
|
||||
|
||||
def test_run_error_handling(self):
|
||||
"""Test error handling during task execution."""
|
||||
agents = [Agent(agent_name="test", agent_description="test")]
|
||||
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.side_effect = Exception("Test execution error")
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
with pytest.raises(SwarmRouterRunError):
|
||||
router.run("Test task")
|
||||
|
||||
def test_batch_run_error_handling(self):
|
||||
"""Test error handling during batch execution."""
|
||||
agents = [Agent(agent_name="test", agent_description="test")]
|
||||
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.side_effect = Exception("Test batch error")
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
router.batch_run(["Task 1", "Task 2"])
|
||||
|
||||
def test_invalid_swarm_type_error(self):
|
||||
"""Test error when creating swarm with invalid type."""
|
||||
router = SwarmRouter(swarm_type="SequentialWorkflow")
|
||||
|
||||
# Manually set an invalid swarm type to test the factory
|
||||
router.swarm_type = "InvalidType"
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
router._create_swarm("Test task")
|
||||
|
||||
|
||||
class TestSwarmRouterCaching:
|
||||
"""Test swarm caching functionality."""
|
||||
|
||||
def test_swarm_caching(self):
|
||||
"""Test that swarms are cached for performance."""
|
||||
agents = [Agent(agent_name="test", agent_description="test")]
|
||||
router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
# Create swarm first time
|
||||
swarm1 = router._create_swarm("Task 1")
|
||||
|
||||
# Create swarm second time with same parameters
|
||||
swarm2 = router._create_swarm("Task 1")
|
||||
|
||||
# Should be the same cached instance
|
||||
assert swarm1 is swarm2
|
||||
|
||||
def test_swarm_cache_different_parameters(self):
|
||||
"""Test that different parameters create different cached swarms."""
|
||||
agents = [Agent(agent_name="test", agent_description="test")]
|
||||
router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
|
||||
|
||||
# Create swarms with different parameters
|
||||
swarm1 = router._create_swarm("Task 1", param1="value1")
|
||||
swarm2 = router._create_swarm("Task 2", param1="value2")
|
||||
|
||||
# Should be different instances
|
||||
assert swarm1 is not swarm2
|
||||
|
||||
|
||||
class TestSwarmRouterOutputTypes:
|
||||
"""Test different output types."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_agents(self):
|
||||
"""Create sample agents for testing."""
|
||||
return [
|
||||
Agent(
|
||||
agent_name="TestAgent",
|
||||
agent_description="Test agent",
|
||||
system_prompt="You are a test agent.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("output_type", [
|
||||
"string", "str", "list", "json", "dict", "yaml", "xml",
|
||||
"dict-all-except-first", "dict-first", "list-all-except-first"
|
||||
])
|
||||
def test_output_types(self, sample_agents, output_type):
|
||||
"""Test different output types."""
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = f"Response for {output_type}"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(
|
||||
agents=sample_agents,
|
||||
swarm_type="SequentialWorkflow",
|
||||
output_type=output_type,
|
||||
)
|
||||
|
||||
result = router.run("Test task")
|
||||
assert result is not None
|
||||
assert result == f"Response for {output_type}"
|
||||
|
||||
|
||||
class TestSwarmRouterIntegration:
|
||||
"""Integration tests for SwarmRouter."""
|
||||
|
||||
def test_full_workflow_with_sequential_workflow(self):
|
||||
"""Test complete workflow with SequentialWorkflow."""
|
||||
agents = [
|
||||
Agent(
|
||||
agent_name="ResearchAgent",
|
||||
agent_description="Research specialist",
|
||||
system_prompt="You are a research specialist.",
|
||||
max_loops=1,
|
||||
),
|
||||
Agent(
|
||||
agent_name="CodeAgent",
|
||||
agent_description="Code expert",
|
||||
system_prompt="You are a code expert.",
|
||||
max_loops=1,
|
||||
),
|
||||
]
|
||||
|
||||
with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
|
||||
mock_instance = Mock()
|
||||
mock_instance.run.return_value = "Complete workflow response"
|
||||
mock_workflow.return_value = mock_instance
|
||||
|
||||
router = SwarmRouter(
|
||||
agents=agents,
|
||||
swarm_type="SequentialWorkflow",
|
||||
max_loops=2,
|
||||
rules="Always provide accurate information",
|
||||
multi_agent_collab_prompt=True,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Test various methods
|
||||
result = router.run("Research and code a Python function")
|
||||
assert result is not None
|
||||
|
||||
batch_results = router.batch_run(["Task 1", "Task 2"])
|
||||
assert batch_results is not None
|
||||
assert len(batch_results) == 2
|
||||
|
||||
concurrent_result = router.concurrent_run("Concurrent task")
|
||||
assert concurrent_result is not None
|
||||
|
||||
# Test serialization
|
||||
config_dict = router.to_dict()
|
||||
assert isinstance(config_dict, dict)
|
||||
assert config_dict["swarm_type"] == "SequentialWorkflow"
|
||||
|
||||
def test_swarm_router_config_model(self):
|
||||
"""Test SwarmRouterConfig model."""
|
||||
config = SwarmRouterConfig(
|
||||
name="test-config",
|
||||
description="Test configuration",
|
||||
swarm_type="SequentialWorkflow",
|
||||
task="Test task",
|
||||
multi_agent_collab_prompt=True,
|
||||
)
|
||||
|
||||
assert config.name == "test-config"
|
||||
assert config.description == "Test configuration"
|
||||
assert config.swarm_type == "SequentialWorkflow"
|
||||
assert config.task == "Test task"
|
||||
assert config.multi_agent_collab_prompt is True
|
||||
|
||||
def test_document_model(self):
|
||||
"""Test Document model."""
|
||||
doc = Document(
|
||||
file_path="test.txt",
|
||||
data="Test document content"
|
||||
)
|
||||
|
||||
assert doc.file_path == "test.txt"
|
||||
assert doc.data == "Test document content"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@ -1,62 +0,0 @@
|
||||
from litellm import completion
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Retrieve detailed current weather information for a specified location, including temperature, humidity, wind speed, and atmospheric conditions.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state, e.g. San Francisco, CA, or a specific geographic coordinate in the format 'latitude,longitude'.",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit", "kelvin"],
|
||||
"description": "The unit of temperature measurement to be used in the response.",
|
||||
},
|
||||
"include_forecast": {
|
||||
"type": "boolean",
|
||||
"description": "Indicates whether to include a short-term weather forecast along with the current conditions.",
|
||||
},
|
||||
"time": {
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
"description": "Optional parameter to specify the time for which the weather data is requested, in ISO 8601 format.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"location",
|
||||
"unit",
|
||||
"include_forecast",
|
||||
"time",
|
||||
],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What's the weather like in Boston today?",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
response = completion(
|
||||
model="gpt-4o-mini",
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
parallel_tool_calls=True,
|
||||
)
|
||||
|
||||
print(response.choices[0].message.tool_calls[0].function.arguments)
|
||||
print(response.choices[0].message)
|
||||
@ -1,545 +0,0 @@
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
from swarms import (
|
||||
Agent,
|
||||
ConcurrentWorkflow,
|
||||
GroupChat,
|
||||
SequentialWorkflow,
|
||||
)
|
||||
|
||||
from swarms.utils.formatter import Formatter
|
||||
|
||||
|
||||
class MarkdownTestSwarm:
|
||||
"""A test swarm that demonstrates markdown output capabilities"""
|
||||
|
||||
def __init__(self):
|
||||
self.formatter = Formatter(markdown=True)
|
||||
self.setup_agents()
|
||||
self.setup_swarm()
|
||||
|
||||
def setup_agents(self):
|
||||
"""Setup specialized agents for markdown testing"""
|
||||
|
||||
# Research Agent - Generates structured markdown reports
|
||||
self.research_agent = Agent(
|
||||
agent_name="Research Agent",
|
||||
system_prompt="""You are a research specialist. When given a topic, create a comprehensive markdown report with:
|
||||
- Clear headers and subheaders
|
||||
- Code examples when relevant
|
||||
- Bullet points and numbered lists
|
||||
- Bold and italic text for emphasis
|
||||
- Tables for data comparison
|
||||
- Code blocks with syntax highlighting
|
||||
|
||||
Always format your response as clean markdown with proper structure.""",
|
||||
model_name="gpt-4o-mini", # Use a more capable model
|
||||
temperature=0.7,
|
||||
max_tokens=4000,
|
||||
max_loops=1,
|
||||
context_length=8000, # Limit context to prevent overflow
|
||||
return_history=False, # Don't return history to reduce context
|
||||
)
|
||||
|
||||
# Code Analysis Agent - Generates code-heavy markdown
|
||||
self.code_agent = Agent(
|
||||
agent_name="Code Analysis Agent",
|
||||
system_prompt="""You are a code analysis specialist. When given code or programming concepts, create markdown documentation with:
|
||||
- Syntax-highlighted code blocks
|
||||
- Function documentation
|
||||
- Code examples
|
||||
- Performance analysis
|
||||
- Best practices
|
||||
|
||||
Use proper markdown formatting with code blocks, inline code, and structured content.""",
|
||||
model_name="gpt-4o-mini", # Use a more capable model
|
||||
temperature=0.5,
|
||||
max_tokens=4000,
|
||||
max_loops=1,
|
||||
context_length=8000, # Limit context to prevent overflow
|
||||
return_history=False, # Don't return history to reduce context
|
||||
)
|
||||
|
||||
# Data Visualization Agent - Creates data-focused markdown
|
||||
self.data_agent = Agent(
|
||||
agent_name="Data Visualization Agent",
|
||||
system_prompt="""You are a data visualization specialist. When given data or analysis requests, create markdown reports with:
|
||||
- Data tables
|
||||
- Statistical analysis
|
||||
- Charts and graphs descriptions
|
||||
- Key insights with bold formatting
|
||||
- Recommendations in structured lists
|
||||
|
||||
Format everything as clean, readable markdown.""",
|
||||
model_name="gpt-4o-mini", # Use a more capable model
|
||||
temperature=0.6,
|
||||
max_tokens=4000,
|
||||
max_loops=1,
|
||||
context_length=8000, # Limit context to prevent overflow
|
||||
return_history=False, # Don't return history to reduce context
|
||||
)
|
||||
|
||||
def setup_swarm(self):
|
||||
"""Setup the swarm with the agents"""
|
||||
# Create different swarm types for testing
|
||||
self.sequential_swarm = SequentialWorkflow(
|
||||
name="Markdown Test Sequential",
|
||||
description="Sequential workflow for markdown testing",
|
||||
agents=[
|
||||
self.research_agent,
|
||||
self.code_agent,
|
||||
self.data_agent,
|
||||
],
|
||||
max_loops=1, # Reduce loops to prevent context overflow
|
||||
)
|
||||
|
||||
self.concurrent_swarm = ConcurrentWorkflow(
|
||||
name="Markdown Test Concurrent",
|
||||
description="Concurrent workflow for markdown testing",
|
||||
agents=[
|
||||
self.research_agent,
|
||||
self.code_agent,
|
||||
self.data_agent,
|
||||
],
|
||||
max_loops=1, # Reduce loops to prevent context overflow
|
||||
)
|
||||
|
||||
self.groupchat_swarm = GroupChat(
|
||||
name="Markdown Test Group Chat",
|
||||
description="A group chat for testing markdown output",
|
||||
agents=[
|
||||
self.research_agent,
|
||||
self.code_agent,
|
||||
self.data_agent,
|
||||
],
|
||||
max_loops=1, # Reduce loops to prevent context overflow
|
||||
)
|
||||
|
||||
# Default swarm for main tests
|
||||
self.swarm = self.sequential_swarm
|
||||
|
||||
def test_basic_markdown_output(self):
|
||||
"""Test basic markdown output with a simple topic"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 1: Basic Markdown Output")
|
||||
print("=" * 60)
|
||||
|
||||
topic = "Python Web Development with FastAPI"
|
||||
|
||||
self.formatter.print_panel(
|
||||
f"Starting research on: {topic}",
|
||||
title="Research Topic",
|
||||
style="bold blue",
|
||||
)
|
||||
|
||||
# Run the research agent
|
||||
result = self.research_agent.run(topic)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result, title="Research Report", border_style="green"
|
||||
)
|
||||
|
||||
def test_code_analysis_markdown(self):
|
||||
"""Test markdown output with code analysis"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 2: Code Analysis Markdown")
|
||||
print("=" * 60)
|
||||
|
||||
code_sample = """
|
||||
def fibonacci(n):
|
||||
if n <= 1:
|
||||
return n
|
||||
return fibonacci(n-1) + fibonacci(n-2)
|
||||
|
||||
# Test the function
|
||||
for i in range(10):
|
||||
print(fibonacci(i))
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Analyzing Python code sample",
|
||||
title="Code Analysis",
|
||||
style="bold cyan",
|
||||
)
|
||||
|
||||
# Run the code analysis agent
|
||||
result = self.code_agent.run(
|
||||
f"Analyze this Python code and provide improvements:\n\n{code_sample}"
|
||||
)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result,
|
||||
title="Code Analysis Report",
|
||||
border_style="yellow",
|
||||
)
|
||||
|
||||
def test_data_analysis_markdown(self):
|
||||
"""Test markdown output with data analysis"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 3: Data Analysis Markdown")
|
||||
print("=" * 60)
|
||||
|
||||
data_request = """
|
||||
Analyze the following dataset:
|
||||
- Sales: $1.2M (Q1), $1.5M (Q2), $1.8M (Q3), $2.1M (Q4)
|
||||
- Growth Rate: 8%, 12%, 15%, 18%
|
||||
- Customer Count: 1000, 1200, 1400, 1600
|
||||
|
||||
Provide insights and recommendations in markdown format.
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Analyzing quarterly business data",
|
||||
title="Data Analysis",
|
||||
style="bold magenta",
|
||||
)
|
||||
|
||||
# Run the data analysis agent
|
||||
result = self.data_agent.run(data_request)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result, title="Data Analysis Report", border_style="red"
|
||||
)
|
||||
|
||||
def test_swarm_collaboration_markdown(self):
|
||||
"""Test markdown output with swarm collaboration"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 4: Swarm Collaboration Markdown")
|
||||
print("=" * 60)
|
||||
|
||||
complex_topic = """
|
||||
Create a comprehensive guide on building a machine learning pipeline that includes:
|
||||
1. Data preprocessing techniques
|
||||
2. Model selection strategies
|
||||
3. Performance evaluation metrics
|
||||
4. Deployment considerations
|
||||
|
||||
Each agent should contribute their expertise and the final output should be well-formatted markdown.
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Swarm collaboration on ML pipeline guide",
|
||||
title="Swarm Task",
|
||||
style="bold green",
|
||||
)
|
||||
|
||||
# Run the swarm
|
||||
results = self.swarm.run(complex_topic)
|
||||
|
||||
# Display individual agent results
|
||||
# SequentialWorkflow returns a list of results, not a dict
|
||||
for i, result in enumerate(results, 1):
|
||||
agent_name = f"Agent {i}"
|
||||
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
# Extract the output from dict result
|
||||
result_content = result.get("output", str(result))
|
||||
else:
|
||||
result_content = str(result)
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title=f"Agent {i}: {agent_name}",
|
||||
border_style="blue",
|
||||
)
|
||||
|
||||
def test_markdown_toggle_functionality(self):
|
||||
"""Test the markdown enable/disable functionality"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 5: Markdown Toggle Functionality")
|
||||
print("=" * 60)
|
||||
|
||||
test_content = """
|
||||
# Test Content
|
||||
|
||||
This is a **bold** test with `inline code`.
|
||||
|
||||
## Code Block
|
||||
```python
|
||||
def test_function():
|
||||
return "Hello, World!"
|
||||
```
|
||||
|
||||
## List
|
||||
- Item 1
|
||||
- Item 2
|
||||
- Item 3
|
||||
"""
|
||||
|
||||
# Test with markdown enabled
|
||||
self.formatter.print_panel(
|
||||
"Testing with markdown ENABLED",
|
||||
title="Markdown Enabled",
|
||||
style="bold green",
|
||||
)
|
||||
self.formatter.print_markdown(test_content, "Markdown Output")
|
||||
|
||||
# Disable markdown
|
||||
self.formatter.disable_markdown()
|
||||
self.formatter.print_panel(
|
||||
"Testing with markdown DISABLED",
|
||||
title="Markdown Disabled",
|
||||
style="bold red",
|
||||
)
|
||||
self.formatter.print_panel(test_content, "Plain Text Output")
|
||||
|
||||
# Re-enable markdown
|
||||
self.formatter.enable_markdown()
|
||||
self.formatter.print_panel(
|
||||
"Testing with markdown RE-ENABLED",
|
||||
title="Markdown Re-enabled",
|
||||
style="bold blue",
|
||||
)
|
||||
self.formatter.print_markdown(
|
||||
test_content, "Markdown Output Again"
|
||||
)
|
||||
|
||||
def test_different_swarm_types(self):
|
||||
"""Test markdown output with different swarm types"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 6: Different Swarm Types")
|
||||
print("=" * 60)
|
||||
|
||||
simple_topic = (
|
||||
"Explain the benefits of using Python for data science"
|
||||
)
|
||||
|
||||
# Test Sequential Workflow
|
||||
print("\n--- Sequential Workflow ---")
|
||||
self.formatter.print_panel(
|
||||
"Testing Sequential Workflow (agents work in sequence)",
|
||||
title="Swarm Type Test",
|
||||
style="bold blue",
|
||||
)
|
||||
sequential_results = self.sequential_swarm.run(simple_topic)
|
||||
for i, result in enumerate(sequential_results, 1):
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
result_content = result.get("output", str(result))
|
||||
else:
|
||||
result_content = str(result)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title=f"Sequential Agent {i}",
|
||||
border_style="blue",
|
||||
)
|
||||
|
||||
# Test Concurrent Workflow
|
||||
print("\n--- Concurrent Workflow ---")
|
||||
self.formatter.print_panel(
|
||||
"Testing Concurrent Workflow (agents work in parallel)",
|
||||
title="Swarm Type Test",
|
||||
style="bold green",
|
||||
)
|
||||
concurrent_results = self.concurrent_swarm.run(simple_topic)
|
||||
for i, result in enumerate(concurrent_results, 1):
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
result_content = result.get("output", str(result))
|
||||
else:
|
||||
result_content = str(result)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title=f"Concurrent Agent {i}",
|
||||
border_style="green",
|
||||
)
|
||||
|
||||
# Test Group Chat
|
||||
print("\n--- Group Chat ---")
|
||||
self.formatter.print_panel(
|
||||
"Testing Group Chat (agents collaborate in conversation)",
|
||||
title="Swarm Type Test",
|
||||
style="bold magenta",
|
||||
)
|
||||
groupchat_results = self.groupchat_swarm.run(simple_topic)
|
||||
|
||||
# Handle different result types for GroupChat
|
||||
if isinstance(groupchat_results, dict):
|
||||
result_content = groupchat_results.get(
|
||||
"output", str(groupchat_results)
|
||||
)
|
||||
else:
|
||||
result_content = str(groupchat_results)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title="Group Chat Result",
|
||||
border_style="magenta",
|
||||
)
|
||||
|
||||
def test_simple_formatter_only(self):
|
||||
"""Test just the formatter functionality without agents"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 7: Simple Formatter Test (No Agents)")
|
||||
print("=" * 60)
|
||||
|
||||
# Test basic markdown rendering
|
||||
simple_markdown = """
|
||||
# Simple Test
|
||||
|
||||
This is a **bold** test with `inline code`.
|
||||
|
||||
## Code Block
|
||||
```python
|
||||
def hello_world():
|
||||
print("Hello, World!")
|
||||
return "Success"
|
||||
```
|
||||
|
||||
## List
|
||||
- Item 1
|
||||
- Item 2
|
||||
- Item 3
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Testing formatter without agents",
|
||||
title="Formatter Test",
|
||||
style="bold cyan",
|
||||
)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
simple_markdown,
|
||||
title="Simple Markdown Test",
|
||||
border_style="green",
|
||||
)
|
||||
|
||||
# Test toggle functionality
|
||||
self.formatter.disable_markdown()
|
||||
self.formatter.print_panel(
|
||||
"Markdown disabled - this should be plain text",
|
||||
title="Plain Text Test",
|
||||
style="bold red",
|
||||
)
|
||||
self.formatter.enable_markdown()
|
||||
|
||||
def test_error_handling_markdown(self):
|
||||
"""Test markdown output with error handling"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 8: Error Handling in Markdown")
|
||||
print("=" * 60)
|
||||
|
||||
# Test with malformed markdown
|
||||
malformed_content = """
|
||||
# Incomplete header
|
||||
**Unclosed bold
|
||||
```python
|
||||
def incomplete_code():
|
||||
# Missing closing backticks
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Testing error handling with malformed markdown",
|
||||
title="Error Handling Test",
|
||||
style="bold yellow",
|
||||
)
|
||||
|
||||
# This should handle the error gracefully
|
||||
self.formatter.print_markdown(
|
||||
malformed_content,
|
||||
title="Malformed Markdown Test",
|
||||
border_style="yellow",
|
||||
)
|
||||
|
||||
# Test with empty content
|
||||
self.formatter.print_markdown(
|
||||
"", title="Empty Content Test", border_style="cyan"
|
||||
)
|
||||
|
||||
# Test with None content
|
||||
self.formatter.print_markdown(
|
||||
None, title="None Content Test", border_style="magenta"
|
||||
)
|
||||
|
||||
def run_all_tests(self):
|
||||
"""Run all markdown output tests"""
|
||||
print(" Starting Swarm Markdown Output Tests")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
# Test 1: Basic markdown output
|
||||
self.test_basic_markdown_output()
|
||||
|
||||
# Test 2: Code analysis markdown
|
||||
self.test_code_analysis_markdown()
|
||||
|
||||
# Test 3: Data analysis markdown
|
||||
self.test_data_analysis_markdown()
|
||||
|
||||
# Test 4: Swarm collaboration
|
||||
self.test_swarm_collaboration_markdown()
|
||||
|
||||
# Test 5: Markdown toggle functionality
|
||||
self.test_markdown_toggle_functionality()
|
||||
|
||||
# Test 6: Different swarm types
|
||||
self.test_different_swarm_types()
|
||||
|
||||
# Test 7: Simple formatter test (no agents)
|
||||
self.test_simple_formatter_only()
|
||||
|
||||
# Test 8: Error handling
|
||||
self.test_error_handling_markdown()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(" All tests completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n Test failed with error: {str(e)}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to run the markdown output tests"""
|
||||
print("Swarms Markdown Output Test Suite")
|
||||
print(
|
||||
"Testing the current state of formatter.py with real swarm agents"
|
||||
)
|
||||
print("=" * 60)
|
||||
|
||||
# Check environment setup
|
||||
api_key = os.getenv("OPENAI_API_KEY") or os.getenv(
|
||||
"SWARMS_API_KEY"
|
||||
)
|
||||
if not api_key:
|
||||
print(
|
||||
"⚠ Warning: No API key found. Please set OPENAI_API_KEY or SWARMS_API_KEY environment variable."
|
||||
)
|
||||
print(
|
||||
" You can create a .env file with: OPENAI_API_KEY=your_api_key_here"
|
||||
)
|
||||
print(
|
||||
" Or set it in your environment: export OPENAI_API_KEY=your_api_key_here"
|
||||
)
|
||||
print()
|
||||
|
||||
try:
|
||||
# Create and run the test swarm
|
||||
test_swarm = MarkdownTestSwarm()
|
||||
test_swarm.run_all_tests()
|
||||
except Exception as e:
|
||||
print(f"\n Test failed with error: {str(e)}")
|
||||
print("\n Troubleshooting tips:")
|
||||
print(
|
||||
"1. Make sure you have set your API key (OPENAI_API_KEY or SWARMS_API_KEY)"
|
||||
)
|
||||
print("2. Check your internet connection")
|
||||
print("3. Verify you have sufficient API credits")
|
||||
print("4. Try running with a simpler test first")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,41 +0,0 @@
|
||||
import pypdf
|
||||
import pytest
|
||||
|
||||
from swarms.utils import pdf_to_text
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pdf_file(tmpdir):
|
||||
pdf_writer = pypdf.PdfWriter()
|
||||
pdf_page = pypdf.PageObject.create_blank_page(None, 200, 200)
|
||||
pdf_writer.add_page(pdf_page)
|
||||
pdf_file = tmpdir.join("temp.pdf")
|
||||
with open(pdf_file, "wb") as output:
|
||||
pdf_writer.write(output)
|
||||
return str(pdf_file)
|
||||
|
||||
|
||||
def test_valid_pdf_to_text(pdf_file):
|
||||
result = pdf_to_text(pdf_file)
|
||||
assert isinstance(result, str)
|
||||
|
||||
|
||||
def test_non_existing_file():
|
||||
with pytest.raises(FileNotFoundError):
|
||||
pdf_to_text("non_existing_file.pdf")
|
||||
|
||||
|
||||
def test_passing_non_pdf_file(tmpdir):
|
||||
file = tmpdir.join("temp.txt")
|
||||
file.write("This is a test")
|
||||
with pytest.raises(
|
||||
Exception,
|
||||
match=r"An error occurred while reading the PDF file",
|
||||
):
|
||||
pdf_to_text(str(file))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_pdf_file", [None, 123, {}, []])
|
||||
def test_invalid_pdf_to_text(invalid_pdf_file):
|
||||
with pytest.raises(Exception):
|
||||
pdf_to_text(invalid_pdf_file)
|
||||
Loading…
Reference in new issue