You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
269 lines
8.4 KiB
269 lines
8.4 KiB
from swarms.structs.agent import Agent
|
|
from swarms.structs.majority_voting import MajorityVoting
|
|
|
|
|
|
def test_majority_voting_basic_execution():
|
|
"""Test basic MajorityVoting execution with multiple agents"""
|
|
# Create specialized agents with different perspectives
|
|
geographer = Agent(
|
|
agent_name="Geography-Expert",
|
|
agent_description="Expert in geography and world capitals",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
historian = Agent(
|
|
agent_name="History-Scholar",
|
|
agent_description="Historical and cultural context specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
political_analyst = Agent(
|
|
agent_name="Political-Analyst",
|
|
agent_description="Political and administrative specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
# Create majority voting system
|
|
mv = MajorityVoting(
|
|
name="Geography-Consensus-System",
|
|
description="Majority voting system for geographical questions",
|
|
agents=[geographer, historian, political_analyst],
|
|
max_loops=1,
|
|
verbose=True,
|
|
)
|
|
|
|
# Test execution
|
|
result = mv.run("What is the capital city of France?")
|
|
assert result is not None
|
|
|
|
|
|
def test_majority_voting_multiple_loops():
|
|
"""Test MajorityVoting with multiple loops for consensus refinement"""
|
|
# Create agents with different knowledge bases
|
|
trivia_expert = Agent(
|
|
agent_name="Trivia-Expert",
|
|
agent_description="General knowledge and trivia specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
research_analyst = Agent(
|
|
agent_name="Research-Analyst",
|
|
agent_description="Research and fact-checking specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
subject_matter_expert = Agent(
|
|
agent_name="Subject-Matter-Expert",
|
|
agent_description="Deep subject matter expertise specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
# Create majority voting with multiple loops for iterative refinement
|
|
mv = MajorityVoting(
|
|
name="Multi-Loop-Consensus-System",
|
|
description="Majority voting with iterative consensus refinement",
|
|
agents=[
|
|
trivia_expert,
|
|
research_analyst,
|
|
subject_matter_expert,
|
|
],
|
|
max_loops=3, # Allow multiple iterations
|
|
verbose=True,
|
|
)
|
|
|
|
# Test multi-loop execution
|
|
result = mv.run(
|
|
"What are the main causes of climate change and what can be done to mitigate them?"
|
|
)
|
|
assert result is not None
|
|
|
|
|
|
def test_majority_voting_business_scenario():
|
|
"""Test MajorityVoting in a realistic business scenario"""
|
|
# Create agents representing different business perspectives
|
|
market_strategist = Agent(
|
|
agent_name="Market-Strategist",
|
|
agent_description="Market strategy and competitive analysis specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
financial_analyst = Agent(
|
|
agent_name="Financial-Analyst",
|
|
agent_description="Financial modeling and ROI analysis specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
technical_architect = Agent(
|
|
agent_name="Technical-Architect",
|
|
agent_description="Technical feasibility and implementation specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
risk_manager = Agent(
|
|
agent_name="Risk-Manager",
|
|
agent_description="Risk assessment and compliance specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
operations_expert = Agent(
|
|
agent_name="Operations-Expert",
|
|
agent_description="Operations and implementation specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
# Create majority voting for business decisions
|
|
mv = MajorityVoting(
|
|
name="Business-Decision-Consensus",
|
|
description="Majority voting system for business strategic decisions",
|
|
agents=[
|
|
market_strategist,
|
|
financial_analyst,
|
|
technical_architect,
|
|
risk_manager,
|
|
operations_expert,
|
|
],
|
|
max_loops=2,
|
|
verbose=True,
|
|
)
|
|
|
|
# Test with complex business decision
|
|
result = mv.run(
|
|
"Should our company invest in developing an AI-powered customer service platform? "
|
|
"Consider market demand, financial implications, technical feasibility, risk factors, "
|
|
"and operational requirements."
|
|
)
|
|
|
|
assert result is not None
|
|
|
|
|
|
def test_majority_voting_error_handling():
|
|
"""Test MajorityVoting error handling and validation"""
|
|
# Test with empty agents list
|
|
try:
|
|
MajorityVoting(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
|
|
analyst = Agent(
|
|
agent_name="Test-Analyst",
|
|
agent_description="Test analyst",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
try:
|
|
MajorityVoting(agents=[analyst], 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_majority_voting_different_output_types():
|
|
"""Test MajorityVoting with different output types"""
|
|
# Create agents for technical analysis
|
|
Agent(
|
|
agent_name="Security-Expert",
|
|
agent_description="Cybersecurity and data protection specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
Agent(
|
|
agent_name="Compliance-Officer",
|
|
agent_description="Regulatory compliance and legal specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
Agent(
|
|
agent_name="Privacy-Advocate",
|
|
agent_description="Privacy protection and data rights specialist",
|
|
model_name="gpt-4o",
|
|
max_loops=1,
|
|
)
|
|
|
|
# Assert majority vote is correct
|
|
assert majority_vote is not None
|
|
|
|
|
|
def test_streaming_majority_voting():
|
|
"""
|
|
Test the streaming_majority_voting with logging/try-except and assertion.
|
|
"""
|
|
logs = []
|
|
|
|
def streaming_callback(
|
|
agent_name: str, chunk: str, is_final: bool
|
|
):
|
|
# Chunk buffer static per call (reset each session)
|
|
if not hasattr(streaming_callback, "_buffer"):
|
|
streaming_callback._buffer = ""
|
|
streaming_callback._buffer_size = 0
|
|
|
|
min_chunk_size = 512 # or any large chunk size you want
|
|
|
|
if chunk:
|
|
streaming_callback._buffer += chunk
|
|
streaming_callback._buffer_size += len(chunk)
|
|
if (
|
|
streaming_callback._buffer_size >= min_chunk_size
|
|
or is_final
|
|
):
|
|
if streaming_callback._buffer:
|
|
print(streaming_callback._buffer, end="", flush=True)
|
|
logs.append(streaming_callback._buffer)
|
|
streaming_callback._buffer = ""
|
|
streaming_callback._buffer_size = 0
|
|
if is_final:
|
|
print()
|
|
|
|
try:
|
|
# Initialize the agent
|
|
agent = Agent(
|
|
agent_name="Financial-Analysis-Agent",
|
|
agent_description="Personal finance advisor agent",
|
|
system_prompt="You are a financial analysis agent.", # replaced missing const
|
|
max_loops=1,
|
|
model_name="gpt-4.1",
|
|
dynamic_temperature_enabled=True,
|
|
user_name="swarms_corp",
|
|
retry_attempts=3,
|
|
context_length=8192,
|
|
return_step_meta=False,
|
|
output_type="str", # "json", "dict", "csv" OR "string" "yaml" and
|
|
auto_generate_prompt=False, # Auto generate prompt for the agent based on name, description, and system prompt, task
|
|
max_tokens=4000, # max output tokens
|
|
saved_state_path="agent_00.json",
|
|
interactive=False,
|
|
streaming_on=True, # if concurrent agents want to be streamed
|
|
)
|
|
|
|
swarm = MajorityVoting(agents=[agent, agent, agent])
|
|
|
|
result = swarm.run(
|
|
"Create a table of super high growth opportunities for AI. I have $40k to invest in ETFs, index funds, and more. Please create a table in markdown.",
|
|
streaming_callback=streaming_callback,
|
|
)
|
|
assert result is not None
|
|
except Exception as e:
|
|
print("Error in test_streaming_majority_voting:", e)
|
|
print("Logs so far:", logs)
|
|
raise
|