-
-
Notifications
You must be signed in to change notification settings - Fork 495
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
adding unit test for spreadsheet_swarm
- Loading branch information
1 parent
6aa2522
commit a359c05
Showing
2 changed files
with
87 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,78 @@ | ||
import os | ||
from datetime import datetime | ||
from uuid import uuid4 | ||
|
||
# Import necessary classes from your swarm module | ||
from swarms.structs.agent import Agent | ||
from swarms.structs.base_swarm import BaseSwarm | ||
from swarms.telemetry.capture_sys_data import log_agent_data | ||
from swarms.utils.file_processing import create_file_in_folder | ||
from swarms import SpreadSheetSwarm | ||
|
||
# Ensure you have an environment variable or default workspace dir | ||
workspace_dir = os.getenv("WORKSPACE_DIR", "./workspace") | ||
|
||
def create_agents(num_agents: int): | ||
""" | ||
Create a list of agent instances. | ||
Args: | ||
num_agents (int): The number of agents to create. | ||
Returns: | ||
List[Agent]: List of created Agent objects. | ||
""" | ||
agents = [] | ||
for i in range(num_agents): | ||
agent_name = f"Agent-{i + 1}" | ||
agents.append(Agent(agent_name=agent_name)) | ||
return agents | ||
|
||
def main(): | ||
# Number of agents to create | ||
num_agents = 5 | ||
|
||
# Create the agents | ||
agents = create_agents(num_agents) | ||
|
||
# Initialize the swarm with agents and other configurations | ||
swarm = SpreadSheetSwarm( | ||
name="Test-Swarm", | ||
description="A swarm for testing purposes.", | ||
agents=agents, | ||
autosave_on=True, | ||
max_loops=2, | ||
workspace_dir=workspace_dir | ||
) | ||
|
||
# Run a sample task in the swarm (synchronously) | ||
task = "process_data" | ||
|
||
# Ensure the run method is synchronous | ||
swarm_metadata = swarm.run(task) # Assuming this is made synchronous | ||
|
||
# Print swarm metadata after task completion | ||
print("Swarm Metadata:") | ||
print(swarm_metadata) | ||
|
||
# Check if CSV file has been created and saved | ||
if os.path.exists(swarm.save_file_path): | ||
print(f"Metadata saved to: {swarm.save_file_path}") | ||
else: | ||
print(f"Metadata not saved correctly. Check the save path.") | ||
|
||
# Test saving metadata to JSON file | ||
swarm.data_to_json_file() | ||
|
||
# Test exporting metadata to JSON | ||
swarm_json = swarm.export_to_json() | ||
print("Exported JSON metadata:") | ||
print(swarm_json) | ||
|
||
# Log agent data | ||
print("Logging agent data:") | ||
print(log_agent_data(swarm.metadata.model_dump())) | ||
|
||
# Run the synchronous main function | ||
if __name__ == "__main__": | ||
main() |