Kye Gomez
2 months ago
parent
d8e06900a4
commit
43c5845bbd
@ -0,0 +1,36 @@
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.groupchat import GroupChat
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
# Example agents
|
||||
agent1 = Agent(
|
||||
agent_name="Financial-Analysis-Agent",
|
||||
system_prompt="You are a financial analyst specializing in investment strategies.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
dynamic_temperature_enabled=True,
|
||||
)
|
||||
|
||||
agent2 = Agent(
|
||||
agent_name="Tax-Adviser-Agent",
|
||||
system_prompt="You are a tax adviser who provides clear and concise guidance on tax-related queries.",
|
||||
model_name="gpt-4o",
|
||||
max_loops=1,
|
||||
dynamic_temperature_enabled=True,
|
||||
)
|
||||
|
||||
agents = [agent1, agent2]
|
||||
|
||||
chat = GroupChat(
|
||||
name="Investment Advisory",
|
||||
description="Financial and tax analysis group",
|
||||
agents=agents,
|
||||
max_loops=1,
|
||||
)
|
||||
|
||||
history = chat.run(
|
||||
"How to optimize tax strategy for investments?"
|
||||
)
|
||||
print(history)
|
||||
@ -1,396 +0,0 @@
|
||||
import json
|
||||
import sched
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Any, Callable, ClassVar, Dict, List, Union
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.conversation import Conversation
|
||||
from swarms.structs.omni_agent_types import AgentType
|
||||
from typing import Optional
|
||||
from swarms.utils.loguru_logger import initialize_logger
|
||||
|
||||
logger = initialize_logger(log_folder="task")
|
||||
|
||||
|
||||
class Task(BaseModel):
|
||||
"""
|
||||
Task class for running a task in a sequential workflow.
|
||||
|
||||
Attributes:
|
||||
description (str): Description of the task.
|
||||
agent (Union[Callable, Agent]): Agent or callable object to run the task.
|
||||
args (List[Any]): Arguments to pass to the agent or callable object.
|
||||
kwargs (Dict[str, Any]): Keyword arguments to pass to the agent or callable object.
|
||||
result (Any): Result of the task.
|
||||
history (List[Any]): History of the task.
|
||||
schedule_time (datetime): Time to schedule the task.
|
||||
scheduler (sched.scheduler): Scheduler to schedule the task.
|
||||
trigger (Callable): Trigger to run the task.
|
||||
action (Callable): Action to run the task.
|
||||
condition (Callable): Condition to run the task.
|
||||
priority (int): Priority of the task.
|
||||
dependencies (List[Task]): List of tasks that need to be completed before this task can be executed.
|
||||
|
||||
Methods:
|
||||
execute: Execute the task by calling the agent or model with the arguments and keyword arguments.
|
||||
handle_scheduled_task: Handles the execution of a scheduled task.
|
||||
set_trigger: Sets the trigger for the task.
|
||||
set_action: Sets the action for the task.
|
||||
set_condition: Sets the condition for the task.
|
||||
is_completed: Checks whether the task has been completed.
|
||||
add_dependency: Adds a task to the list of dependencies.
|
||||
set_priority: Sets the priority of the task.
|
||||
check_dependency_completion: Checks whether all the dependencies have been completed.
|
||||
|
||||
|
||||
Examples:
|
||||
>>> from swarms.structs import Task, Agent
|
||||
>>> from swarm_models import OpenAIChat
|
||||
>>> agent = Agent(llm=OpenAIChat(openai_api_key=""), max_loops=1, dashboard=False)
|
||||
>>> task = Task(description="What's the weather in miami", agent=agent)
|
||||
>>> task.run()
|
||||
|
||||
>>> task.result
|
||||
|
||||
"""
|
||||
|
||||
name: Optional[str] = "Task"
|
||||
description: Optional[str] = (
|
||||
"A task is a unit of work that needs to be completed for a workflow to progress."
|
||||
)
|
||||
agent: Optional[Union[Callable, Agent, AgentType]] = Field(
|
||||
None,
|
||||
description="Agent or callable object to run the task",
|
||||
)
|
||||
result: Optional[Any] = None
|
||||
history: List[Any] = Field(default_factory=list)
|
||||
schedule_time: Optional[datetime] = Field(
|
||||
None,
|
||||
description="Time to schedule the task",
|
||||
)
|
||||
scheduler: ClassVar[sched.scheduler] = sched.scheduler(
|
||||
time.time, time.sleep
|
||||
)
|
||||
trigger: Optional[Callable] = Field(
|
||||
None,
|
||||
description="Trigger to run the task",
|
||||
)
|
||||
action: Optional[Callable] = Field(
|
||||
None,
|
||||
description="Action to run the task",
|
||||
)
|
||||
condition: Optional[Callable] = Field(
|
||||
None,
|
||||
description="Condition to run the task",
|
||||
)
|
||||
priority: Optional[int] = Field(
|
||||
0.4,
|
||||
description="Priority of the task",
|
||||
)
|
||||
dependencies: List["Task"] = Field(default_factory=list)
|
||||
args: List[Any] = Field(default_factory=list)
|
||||
kwargs: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
# We need to check that the agent exists
|
||||
|
||||
def step(self, task: str = None, *args, **kwargs):
|
||||
"""
|
||||
Execute the task by calling the agent or model with the arguments and
|
||||
keyword arguments. You can add images to the agent by passing the
|
||||
path to the image as a keyword argument.
|
||||
|
||||
|
||||
Examples:
|
||||
>>> from swarms.structs import Task, Agent
|
||||
>>> from swarm_models import OpenAIChat
|
||||
>>> agent = Agent(llm=OpenAIChat(openai_api_key=""), max_loops=1, dashboard=False)
|
||||
>>> task = Task(description="What's the weather in miami", agent=agent)
|
||||
>>> task.run()
|
||||
>>> task.result
|
||||
|
||||
"""
|
||||
|
||||
logger.info(f"Running task: {task}")
|
||||
|
||||
# Check dependencies
|
||||
if not self.check_dependency_completion():
|
||||
logger.info(
|
||||
f"Task {self.description} is waiting for dependencies to complete"
|
||||
)
|
||||
return None
|
||||
|
||||
# Check the condition before executing the task
|
||||
if self.condition is not None:
|
||||
try:
|
||||
condition_result = self.condition()
|
||||
if not condition_result:
|
||||
logger.info(
|
||||
f"Completion not met for the task: {task} Skipping execution"
|
||||
)
|
||||
return None
|
||||
except Exception as error:
|
||||
logger.error(f"[ERROR][Task] {error}")
|
||||
return None
|
||||
|
||||
# Execute the task
|
||||
if self.trigger is None or self.trigger():
|
||||
try:
|
||||
logger.info(f"Executing task: {task}")
|
||||
self.result = self.agent.run(task, *args, **kwargs)
|
||||
|
||||
# Ensure the result is either a string or a dict
|
||||
if isinstance(self.result, str):
|
||||
logger.info(f"Task result: {self.result}")
|
||||
elif isinstance(self.result, dict):
|
||||
logger.info(f"Task result: {self.result}")
|
||||
else:
|
||||
logger.error(
|
||||
"Task result must be either a string or a dict"
|
||||
)
|
||||
|
||||
# Add the result to the history
|
||||
self.history.append(self.result)
|
||||
|
||||
# If an action is specified, execute it
|
||||
if self.action is not None:
|
||||
try:
|
||||
logger.info(
|
||||
f"Executing action for task: {task}"
|
||||
)
|
||||
self.action()
|
||||
except Exception as error:
|
||||
logger.error(f"[ERROR][Task] {error}")
|
||||
except Exception as error:
|
||||
logger.error(f"[ERROR][Task] {error}")
|
||||
else:
|
||||
logger.info(f"Task {task} is not triggered")
|
||||
|
||||
def run(self, task: str = None, *args, **kwargs):
|
||||
now = datetime.now()
|
||||
|
||||
# If the task is scheduled for the future, schedule it
|
||||
if self.schedule_time and self.schedule_time > now:
|
||||
delay = (self.schedule_time - now).total_seconds()
|
||||
logger.info(
|
||||
f"Scheduling task: {self.description} for {self.schedule_time}"
|
||||
)
|
||||
self.scheduler.enter(
|
||||
delay,
|
||||
1,
|
||||
self.step,
|
||||
argument=(task, args, kwargs),
|
||||
)
|
||||
self.scheduler.run()
|
||||
|
||||
# We need to return the result
|
||||
else:
|
||||
# If no scheduling or the time has already passed run the task
|
||||
return self.step(task, *args, **kwargs)
|
||||
|
||||
def handle_scheduled_task(self):
|
||||
"""
|
||||
Handles the execution of a scheduled task.
|
||||
|
||||
If the schedule time is not set or has already passed, the task is executed immediately.
|
||||
Otherwise, the task is scheduled to be executed at the specified schedule time.
|
||||
"""
|
||||
logger.info(
|
||||
f"[INFO][Task] Handling scheduled task: {self.description}"
|
||||
)
|
||||
try:
|
||||
if (
|
||||
self.schedule_time is None
|
||||
or self.schedule_time <= datetime.now()
|
||||
):
|
||||
self.execute()
|
||||
|
||||
else:
|
||||
delay = (
|
||||
self.schedule_time - datetime.now()
|
||||
).total_seconds()
|
||||
self.scheduler.enter(delay, 1, self.execute)
|
||||
self.scheduler_run()
|
||||
except Exception as error:
|
||||
logger.error(f"[ERROR][Task] {error}")
|
||||
|
||||
def set_trigger(self, trigger: Callable):
|
||||
"""
|
||||
Sets the trigger for the task.
|
||||
|
||||
Args:
|
||||
trigger (Callable): The trigger to set.
|
||||
"""
|
||||
self.trigger = trigger
|
||||
|
||||
def set_action(self, action: Callable):
|
||||
"""
|
||||
Sets the action for the task.
|
||||
|
||||
Args:
|
||||
action (Callable): The action to set.
|
||||
"""
|
||||
self.action = action
|
||||
|
||||
def set_condition(self, condition: Callable):
|
||||
"""
|
||||
Sets the condition for the task.
|
||||
|
||||
Args:
|
||||
condition (Callable): The condition to set.
|
||||
"""
|
||||
self.condition = condition
|
||||
|
||||
def is_completed(self):
|
||||
"""Is the task completed?
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
return self.result is not None
|
||||
|
||||
def add_dependency(self, task):
|
||||
"""Adds a task to the list of dependencies.
|
||||
|
||||
Args:
|
||||
task (_type_): _description_
|
||||
"""
|
||||
self.dependencies.append(task)
|
||||
|
||||
def set_priority(self, priority: int):
|
||||
"""Sets the priority of the task.
|
||||
|
||||
Args:
|
||||
priority (int): _description_
|
||||
"""
|
||||
self.priority = priority
|
||||
|
||||
def check_dependency_completion(self):
|
||||
"""
|
||||
Checks whether all the dependencies have been completed.
|
||||
|
||||
Returns:
|
||||
bool: True if all the dependencies have been completed, False otherwise.
|
||||
"""
|
||||
logger.info("[INFO][Task] Checking dependency completion")
|
||||
try:
|
||||
for task in self.dependencies:
|
||||
if not task.is_completed():
|
||||
return False
|
||||
except Exception as error:
|
||||
logger.error(
|
||||
f"[ERROR][Task][check_dependency_completion] {error}"
|
||||
)
|
||||
|
||||
def context(
|
||||
self,
|
||||
task: "Task" = None,
|
||||
context: List["Task"] = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Set the context for the task.
|
||||
|
||||
Args:
|
||||
context (str): The context to set.
|
||||
"""
|
||||
# For sequential workflow, sequentially add the context of the previous task in the list
|
||||
new_context = Conversation(time_enabled=True, *args, **kwargs)
|
||||
|
||||
if context:
|
||||
for task in context:
|
||||
description = (
|
||||
task.description
|
||||
if task.description is not None
|
||||
else ""
|
||||
)
|
||||
|
||||
result = (
|
||||
task.result if task.result is not None else ""
|
||||
)
|
||||
|
||||
# Add the context of the task to the conversation
|
||||
new_context.add(
|
||||
task.agent.agent_name, f"{description} {result}"
|
||||
)
|
||||
|
||||
elif task:
|
||||
description = (
|
||||
task.description
|
||||
if task.description is not None
|
||||
else ""
|
||||
)
|
||||
result = task.result if task.result is not None else ""
|
||||
new_context.add(
|
||||
task.agent.agent_name, f"{description} {result}"
|
||||
)
|
||||
|
||||
prompt = new_context.return_history_as_string()
|
||||
|
||||
# Add to history
|
||||
return self.history.append(prompt)
|
||||
|
||||
def to_dict(self):
|
||||
"""
|
||||
Convert the task to a dictionary.
|
||||
|
||||
Returns:
|
||||
dict: The task as a dictionary.
|
||||
"""
|
||||
return self.model_dump_json(indent=4)
|
||||
|
||||
def save_to_file(self, file_path: str):
|
||||
"""
|
||||
Save the task to a file.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the file to save the task to.
|
||||
"""
|
||||
with open(file_path, "w") as file:
|
||||
file.write(self.to_json(indent=4))
|
||||
|
||||
@classmethod
|
||||
def load_from_file(cls, file_path: str):
|
||||
"""
|
||||
Load a task from a file.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the file to load the task from.
|
||||
|
||||
Returns:
|
||||
Task: The task loaded from the file.
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
task_dict = json.load(file)
|
||||
return Task(**task_dict)
|
||||
|
||||
def schedule_task_with_sched(
|
||||
function: Callable, run_date: datetime
|
||||
) -> None:
|
||||
now = datetime.now()
|
||||
|
||||
if run_date <= now:
|
||||
raise ValueError("run_date must be in the future")
|
||||
|
||||
# Calculate the delay in seconds
|
||||
delay = (run_date - now).total_seconds()
|
||||
|
||||
scheduler = sched.scheduler(time.time, time.sleep)
|
||||
|
||||
# Schedule the function
|
||||
scheduler.enter(delay, 1, function)
|
||||
|
||||
# Start the scheduler
|
||||
scheduler.run(delay, 1, function)
|
||||
|
||||
# Start the scheduler
|
||||
logger.info(f"Task scheduled for {run_date}")
|
||||
scheduler.run()
|
||||
|
||||
return None
|
||||
@ -1,147 +1,388 @@
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from swarm_models import OpenAIChat
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.groupchat import GroupChat, expertise_based
|
||||
|
||||
from swarms.structs.groupchat import (
|
||||
GroupChat,
|
||||
round_robin,
|
||||
expertise_based,
|
||||
random_selection,
|
||||
sentiment_based,
|
||||
length_based,
|
||||
question_based,
|
||||
topic_based,
|
||||
)
|
||||
from datetime import datetime
|
||||
import time
|
||||
|
||||
|
||||
class TestReport:
|
||||
def __init__(self):
|
||||
self.results = []
|
||||
self.start_time = None
|
||||
self.end_time = None
|
||||
|
||||
def add_result(self, test_name, passed, message="", duration=0):
|
||||
self.results.append(
|
||||
{
|
||||
"test_name": test_name,
|
||||
"passed": passed,
|
||||
"message": message,
|
||||
"duration": duration,
|
||||
}
|
||||
)
|
||||
|
||||
def setup_test_agents():
|
||||
model = OpenAIChat(
|
||||
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
||||
model_name="gpt-4",
|
||||
temperature=0.1,
|
||||
def start(self):
|
||||
self.start_time = datetime.now()
|
||||
|
||||
def end(self):
|
||||
self.end_time = datetime.now()
|
||||
|
||||
def generate_report(self):
|
||||
total_tests = len(self.results)
|
||||
passed_tests = sum(1 for r in self.results if r["passed"])
|
||||
failed_tests = total_tests - passed_tests
|
||||
duration = (self.end_time - self.start_time).total_seconds()
|
||||
|
||||
report = "\n" + "=" * 50 + "\n"
|
||||
report += "GROUP CHAT TEST SUITE REPORT\n"
|
||||
report += "=" * 50 + "\n\n"
|
||||
report += f"Test Run: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}\n"
|
||||
report += f"Duration: {duration:.2f} seconds\n"
|
||||
report += f"Total Tests: {total_tests}\n"
|
||||
report += f"Passed: {passed_tests}\n"
|
||||
report += f"Failed: {failed_tests}\n"
|
||||
report += (
|
||||
f"Success Rate: {(passed_tests/total_tests)*100:.1f}%\n\n"
|
||||
)
|
||||
|
||||
return [
|
||||
Agent(
|
||||
agent_name="Agent1",
|
||||
system_prompt="You only respond with 'A'",
|
||||
llm=model,
|
||||
report += "Detailed Test Results:\n"
|
||||
report += "-" * 50 + "\n"
|
||||
|
||||
for result in self.results:
|
||||
status = "✓" if result["passed"] else "✗"
|
||||
report += f"{status} {result['test_name']} ({result['duration']:.2f}s)\n"
|
||||
if result["message"]:
|
||||
report += f" {result['message']}\n"
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def create_test_agents(num_agents, diverse_prompts=False):
|
||||
"""Helper function to create test agents with diverse prompts"""
|
||||
agents = []
|
||||
specialties = [
|
||||
(
|
||||
"Finance",
|
||||
"You are a financial expert focusing on investment strategies and market analysis. Be concise and data-driven in your responses.",
|
||||
),
|
||||
Agent(
|
||||
agent_name="Agent2",
|
||||
system_prompt="You only respond with 'B'",
|
||||
llm=model,
|
||||
(
|
||||
"Tech",
|
||||
"You are a technology expert specializing in AI and cybersecurity. Use technical terms and provide practical examples.",
|
||||
),
|
||||
Agent(
|
||||
agent_name="Agent3",
|
||||
system_prompt="You only respond with 'C'",
|
||||
llm=model,
|
||||
(
|
||||
"Healthcare",
|
||||
"You are a healthcare professional with expertise in public health. Focus on evidence-based information and patient care.",
|
||||
),
|
||||
(
|
||||
"Marketing",
|
||||
"You are a marketing strategist focusing on digital trends. Be creative and audience-focused in your responses.",
|
||||
),
|
||||
(
|
||||
"Legal",
|
||||
"You are a legal expert specializing in corporate law. Be precise and reference relevant regulations.",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_round_robin_speaking():
|
||||
chat = GroupChat(agents=setup_test_agents())
|
||||
history = chat.run("Say your letter")
|
||||
|
||||
# Verify agents speak in order
|
||||
responses = [
|
||||
r.message for t in history.turns for r in t.responses
|
||||
for i in range(num_agents):
|
||||
specialty, base_prompt = specialties[i % len(specialties)]
|
||||
if diverse_prompts:
|
||||
# Add personality traits and communication style to make responses more diverse
|
||||
traits = [
|
||||
"Be analytical and data-focused",
|
||||
"Use analogies and examples",
|
||||
"Be concise and direct",
|
||||
"Ask thought-provoking questions",
|
||||
"Provide practical applications",
|
||||
]
|
||||
assert responses == ["A", "B", "C"] * (len(history.turns))
|
||||
prompt = f"{base_prompt} {traits[i % len(traits)]}"
|
||||
else:
|
||||
prompt = base_prompt
|
||||
|
||||
agents.append(
|
||||
Agent(
|
||||
agent_name=f"{specialty}-Agent-{i+1}",
|
||||
system_prompt=prompt,
|
||||
model_name="gpt-4",
|
||||
max_loops=1,
|
||||
temperature=0.7, # Add temperature to increase response variety
|
||||
)
|
||||
)
|
||||
return agents
|
||||
|
||||
def test_concurrent_processing():
|
||||
chat = GroupChat(agents=setup_test_agents())
|
||||
tasks = ["Task1", "Task2", "Task3"]
|
||||
histories = chat.concurrent_run(tasks)
|
||||
|
||||
assert len(histories) == len(tasks)
|
||||
for history in histories:
|
||||
assert history.total_messages > 0
|
||||
def test_basic_groupchat(report):
|
||||
"""Test basic GroupChat initialization and conversation"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
agents = create_test_agents(2)
|
||||
chat = GroupChat(
|
||||
name="Test Chat",
|
||||
description="A test group chat",
|
||||
agents=agents,
|
||||
max_loops=2,
|
||||
)
|
||||
|
||||
def test_expertise_based_speaking():
|
||||
agents = setup_test_agents()
|
||||
chat = GroupChat(agents=agents, speaker_fn=expertise_based)
|
||||
result = chat.run("Say hello!")
|
||||
report.add_result(
|
||||
"Basic GroupChat Test",
|
||||
True,
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
# Test each agent's expertise trigger
|
||||
for agent in agents:
|
||||
history = chat.run(f"Trigger {agent.system_prompt}")
|
||||
first_response = history.turns[0].responses[0]
|
||||
assert first_response.agent_name == agent.agent_name
|
||||
except Exception as e:
|
||||
report.add_result(
|
||||
"Basic GroupChat Test",
|
||||
False,
|
||||
message=str(e),
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
|
||||
def test_max_loops_limit():
|
||||
max_loops = 3
|
||||
chat = GroupChat(agents=setup_test_agents(), max_loops=max_loops)
|
||||
history = chat.run("Test message")
|
||||
def test_speaker_functions(report):
|
||||
"""Test all available speaker functions with enhanced prompts"""
|
||||
speaker_functions = {
|
||||
"round_robin": (
|
||||
round_robin,
|
||||
"What are your thoughts on sustainable practices?",
|
||||
),
|
||||
"expertise_based": (
|
||||
expertise_based,
|
||||
"Discuss the impact of AI on your field.",
|
||||
),
|
||||
"random_selection": (
|
||||
random_selection,
|
||||
"How do you approach problem-solving?",
|
||||
),
|
||||
"sentiment_based": (
|
||||
sentiment_based,
|
||||
"Share your positive outlook on future trends.",
|
||||
),
|
||||
"length_based": (
|
||||
length_based,
|
||||
"Provide a detailed analysis of recent developments.",
|
||||
),
|
||||
"question_based": (
|
||||
question_based,
|
||||
"What challenges do you foresee in your industry?",
|
||||
),
|
||||
"topic_based": (
|
||||
topic_based,
|
||||
"How does digital transformation affect your sector?",
|
||||
),
|
||||
}
|
||||
|
||||
assert len(history.turns) == max_loops
|
||||
for name, (func, prompt) in speaker_functions.items():
|
||||
start_time = time.time()
|
||||
try:
|
||||
# Create agents with diverse prompts for this test
|
||||
agents = create_test_agents(3, diverse_prompts=True)
|
||||
chat = GroupChat(
|
||||
name=f"{name.title()} Test",
|
||||
description=f"Testing {name} speaker function with diverse responses",
|
||||
agents=agents,
|
||||
speaker_fn=func,
|
||||
max_loops=2,
|
||||
rules="1. Be unique in your responses\n2. Build on others' points\n3. Stay relevant to your expertise",
|
||||
)
|
||||
|
||||
result = chat.run(prompt)
|
||||
report.add_result(
|
||||
f"Speaker Function - {name}",
|
||||
True,
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
def test_error_handling():
|
||||
broken_agent = Agent(
|
||||
agent_name="BrokenAgent",
|
||||
system_prompt="You raise errors",
|
||||
llm=None,
|
||||
except Exception as e:
|
||||
report.add_result(
|
||||
f"Speaker Function - {name}",
|
||||
False,
|
||||
message=str(e),
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
chat = GroupChat(agents=[broken_agent])
|
||||
history = chat.run("Trigger error")
|
||||
|
||||
assert "Error" in history.turns[0].responses[0].message
|
||||
def test_varying_agent_counts(report):
|
||||
"""Test GroupChat with different numbers of agents"""
|
||||
agent_counts = [1, 3, 5, 7]
|
||||
|
||||
for count in agent_counts:
|
||||
start_time = time.time()
|
||||
try:
|
||||
agents = create_test_agents(count)
|
||||
chat = GroupChat(
|
||||
name=f"{count}-Agent Test", agents=agents, max_loops=2
|
||||
)
|
||||
|
||||
def test_conversation_context():
|
||||
agents = setup_test_agents()
|
||||
complex_prompt = "Previous message refers to A. Now trigger B. Finally discuss C."
|
||||
result = chat.run("Introduce yourselves briefly.")
|
||||
report.add_result(
|
||||
f"Agent Count Test - {count} agents",
|
||||
True,
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
report.add_result(
|
||||
f"Agent Count Test - {count} agents",
|
||||
False,
|
||||
message=str(e),
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
chat = GroupChat(agents=agents, speaker_fn=expertise_based)
|
||||
history = chat.run(complex_prompt)
|
||||
|
||||
responses = [
|
||||
r.agent_name for t in history.turns for r in t.responses
|
||||
def test_error_cases(report):
|
||||
"""Test error handling with expanded cases"""
|
||||
test_cases = [
|
||||
("Empty Agents List", lambda: GroupChat(agents=[])),
|
||||
(
|
||||
"Invalid Max Loops",
|
||||
lambda: GroupChat(
|
||||
agents=[create_test_agents(1)[0]], max_loops=0
|
||||
),
|
||||
),
|
||||
(
|
||||
"Empty Task",
|
||||
lambda: GroupChat(agents=[create_test_agents(1)[0]]).run(
|
||||
""
|
||||
),
|
||||
),
|
||||
(
|
||||
"None Task",
|
||||
lambda: GroupChat(agents=[create_test_agents(1)[0]]).run(
|
||||
None
|
||||
),
|
||||
),
|
||||
(
|
||||
"Invalid Speaker Function",
|
||||
lambda: GroupChat(
|
||||
agents=[create_test_agents(1)[0]],
|
||||
speaker_fn=lambda x, y: "not a boolean", # This should raise ValueError
|
||||
),
|
||||
),
|
||||
]
|
||||
assert all(agent.agent_name in responses for agent in agents)
|
||||
|
||||
for name, test_func in test_cases:
|
||||
start_time = time.time()
|
||||
try:
|
||||
test_func()
|
||||
report.add_result(
|
||||
f"Error Case - {name}",
|
||||
False,
|
||||
message="Expected ValueError not raised",
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
except (
|
||||
ValueError,
|
||||
TypeError,
|
||||
): # Include TypeError for invalid speaker function
|
||||
report.add_result(
|
||||
f"Error Case - {name}",
|
||||
True,
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
except Exception as e:
|
||||
report.add_result(
|
||||
f"Error Case - {name}",
|
||||
False,
|
||||
message=f"Unexpected error: {str(e)}",
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
|
||||
def test_concurrent_execution(report):
|
||||
"""Test concurrent execution with various task counts"""
|
||||
start_time = time.time()
|
||||
|
||||
def test_large_agent_group():
|
||||
large_group = setup_test_agents() * 5 # 15 agents
|
||||
chat = GroupChat(agents=large_group)
|
||||
history = chat.run("Test scaling")
|
||||
try:
|
||||
agents = create_test_agents(3)
|
||||
chat = GroupChat(
|
||||
name="Concurrent Test", agents=agents, max_loops=1
|
||||
)
|
||||
|
||||
assert history.total_messages > len(large_group)
|
||||
tasks = [
|
||||
"Task 1: Introduce yourself",
|
||||
"Task 2: What's your specialty?",
|
||||
"Task 3: How can you help?",
|
||||
"Task 4: What are your limitations?",
|
||||
"Task 5: Give an example of your expertise",
|
||||
]
|
||||
|
||||
results = chat.concurrent_run(tasks)
|
||||
report.add_result(
|
||||
"Concurrent Execution Test",
|
||||
True,
|
||||
message=f"Successfully completed {len(results)} tasks",
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
report.add_result(
|
||||
"Concurrent Execution Test",
|
||||
False,
|
||||
message=str(e),
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
def test_long_conversations():
|
||||
chat = GroupChat(agents=setup_test_agents(), max_loops=50)
|
||||
history = chat.run("Long conversation test")
|
||||
|
||||
assert len(history.turns) == 50
|
||||
assert history.total_messages > 100
|
||||
def test_conversation_rules(report):
|
||||
"""Test GroupChat with different conversation rules"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
agents = create_test_agents(3, diverse_prompts=True)
|
||||
chat = GroupChat(
|
||||
name="Rules Test",
|
||||
description="Testing conversation with specific rules",
|
||||
agents=agents,
|
||||
max_loops=2,
|
||||
rules="""
|
||||
1. Keep responses under 50 words
|
||||
2. Always be professional
|
||||
3. Stay on topic
|
||||
4. Provide unique perspectives
|
||||
5. Build on previous responses
|
||||
""",
|
||||
)
|
||||
|
||||
def test_stress_batched_runs():
|
||||
chat = GroupChat(agents=setup_test_agents())
|
||||
tasks = ["Task"] * 100
|
||||
histories = chat.batched_run(tasks)
|
||||
result = chat.run(
|
||||
"How can we ensure ethical AI development across different sectors?"
|
||||
)
|
||||
report.add_result(
|
||||
"Conversation Rules Test",
|
||||
True,
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
assert len(histories) == len(tasks)
|
||||
total_messages = sum(h.total_messages for h in histories)
|
||||
assert total_messages > len(tasks) * 3
|
||||
except Exception as e:
|
||||
report.add_result(
|
||||
"Conversation Rules Test",
|
||||
False,
|
||||
message=str(e),
|
||||
duration=time.time() - start_time,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
functions = [
|
||||
test_round_robin_speaking,
|
||||
test_concurrent_processing,
|
||||
test_expertise_based_speaking,
|
||||
test_max_loops_limit,
|
||||
test_error_handling,
|
||||
test_conversation_context,
|
||||
test_large_agent_group,
|
||||
test_long_conversations,
|
||||
test_stress_batched_runs,
|
||||
]
|
||||
report = TestReport()
|
||||
report.start()
|
||||
|
||||
for func in functions:
|
||||
try:
|
||||
print(f"Running {func.__name__}...")
|
||||
func()
|
||||
print("✓ Passed")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed: {str(e)}")
|
||||
print("Starting Enhanced GroupChat Test Suite...\n")
|
||||
|
||||
# Run all tests
|
||||
test_basic_groupchat(report)
|
||||
test_speaker_functions(report)
|
||||
test_varying_agent_counts(report)
|
||||
test_error_cases(report)
|
||||
test_concurrent_execution(report)
|
||||
test_conversation_rules(report)
|
||||
|
||||
report.end()
|
||||
print(report.generate_report())
|
||||
|
||||
Loading…
Reference in new issue