[NEW][Added examples to the models playground for varius models, vilt, bioclip,]

Former-commit-id: 604611d669

playground/models/bioclip.py

@@ -0,0 +1,19 @@
+from swarms.models.bioclip import BioClip
+
+clip = BioClip("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
+
+labels = [
+    "adenocarcinoma histopathology",
+    "brain MRI",
+    "covid line chart",
+    "squamous cell carcinoma histopathology",
+    "immunohistochemistry histopathology",
+    "bone X-ray",
+    "chest X-ray",
+    "pie chart",
+    "hematoxylin and eosin histopathology",
+]
+
+result = clip("swarms.jpeg", labels)
+metadata = {"filename": "images/.jpg".split("/")[-1], "top_probs": result}
+clip.plot_image_with_metadata("swarms.jpeg", metadata)

playground/models/biogpt.py

@@ -0,0 +1,7 @@
+from swarms.models.biogpt import BioGPTWrapper
+
+model = BioGPTWrapper()
+
+out = model("The patient has a fever")
+
+print(out)

playground/models/distilled_whiserpx.py

@@ -0,0 +1,10 @@
+import asyncio
+from swarms.models.distilled_whisperx import DistilWhisperModel
+
+model_wrapper = DistilWhisperModel()
+
+# Download mp3 of voice and place the path here
+transcription = model_wrapper("path/to/audio.mp3")
+
+# For async usage
+transcription = asyncio.run(model_wrapper.async_transcribe("path/to/audio.mp3"))

playground/models/fast_vit.py

@@ -0,0 +1,5 @@
+from swarms.models.fastvit import FastViT
+
+fastvit = FastViT()
+
+result = fastvit(img="images/swarms.jpeg", confidence_threshold=0.5)

playground/models/fuyu.py

@@ -0,0 +1,7 @@
+from swarms.models.fuyu import Fuyu
+
+fuyu = Fuyu()
+
+# This is the default image, you can change it to any image you want
+out = fuyu("What is this image?", "images/swarms.jpeg")
+print(out)

playground/models/huggingface.py

@@ -0,0 +1,8 @@
+from swarms.models import HuggingfaceLLM
+
+model_id = "NousResearch/Yarn-Mistral-7b-128k"
+inference = HuggingfaceLLM(model_id=model_id)
+
+task = "Once upon a time"
+generated_text = inference(task)
+print(generated_text)

playground/models/idefics.py

@@ -0,0 +1,16 @@
+from swarms.models import idefics
+
+model = idefics()
+
+user_input = "User: What is in this image? https://upload.wikimedia.org/wikipedia/commons/8/86/Id%C3%A9fix.JPG"
+response = model.chat(user_input)
+print(response)
+
+user_input = "User: And who is that? https://static.wikia.nocookie.net/asterix/images/2/25/R22b.gif/revision/latest?cb=20110815073052"
+response = model.chat(user_input)
+print(response)
+
+model.set_checkpoint("new_checkpoint")
+model.set_device("cpu")
+model.set_max_length(200)
+model.clear_chat_history()

playground/models/jina_embeds.py

@@ -0,0 +1,7 @@
+from swarms.models import JinaEmbeddings
+
+model = JinaEmbeddings()
+
+embeddings = model("Encode this text")
+
+print(embeddings)

playground/models/kosmos2.py

@@ -0,0 +1,10 @@
+from swarms.models.kosmos2 import Kosmos2, Detections
+from PIL import Image
+
+
+model = Kosmos2.initialize()
+
+image = Image.open("images/swarms.jpg")
+
+detections = model(image)
+print(detections)

playground/models/kosmos_two.py

@@ -0,0 +1,11 @@
+from swarms.models.kosmos_two import Kosmos
+
+# Initialize Kosmos
+kosmos = Kosmos()
+
+# Perform multimodal grounding
+out = kosmos.multimodal_grounding(
+    "Find the red apple in the image.", "images/swarms.jpeg"
+)
+
+print(out)

playground/models/layout_documentxlm.py

@@ -0,0 +1,8 @@
+from swarms.models import LayoutLMDocumentQA
+
+model = LayoutLMDocumentQA()
+
+# Place an image of a financial document
+out = model("What is the total amount?", "images/swarmfest.png")
+
+print(out)

playground/models/llama_function_caller.py

@@ -0,0 +1,35 @@
+from swarms.models.llama_function_caller import LlamaFunctionCaller
+
+llama_caller = LlamaFunctionCaller()
+
+
+# Add a custom function
+def get_weather(location: str, format: str) -> str:
+    # This is a placeholder for the actual implementation
+    return f"Weather at {location} in {format} format."
+
+
+llama_caller.add_func(
+    name="get_weather",
+    function=get_weather,
+    description="Get the weather at a location",
+    arguments=[
+        {
+            "name": "location",
+            "type": "string",
+            "description": "Location for the weather",
+        },
+        {
+            "name": "format",
+            "type": "string",
+            "description": "Format of the weather data",
+        },
+    ],
+)
+
+# Call the function
+result = llama_caller.call_function("get_weather", location="Paris", format="Celsius")
+print(result)
+
+# Stream a user prompt
+llama_caller("Tell me about the tallest mountain in the world.")

playground/models/mpt.py

@@ -0,0 +1,7 @@
+from swarms.models.mpt import MPT
+
+mpt_instance = MPT(
+    "mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150
+)
+
+mpt_instance.generate("Once upon a time in a land far, far away...")

playground/models/nougat.py

@@ -0,0 +1,5 @@
+from swarms.models.nougat import Nougat
+
+nougat = Nougat()
+
+out = nougat("path/to/image.png")

playground/models/palm.py

@@ -0,0 +1,5 @@
+from swarms.models.palm import PALM
+
+palm = PALM()
+
+out = palm("path/to/image.png")

playground/models/speecht5.py

@@ -0,0 +1,9 @@
+from swarms.models.speecht5 import SpeechT5Wrapper
+
+speechT5 = SpeechT5Wrapper()
+
+result = speechT5("Hello, how are you?")
+
+speechT5.save_speech(result)
+print("Speech saved successfully!")
+

playground/models/ssd.py

@@ -0,0 +1,9 @@
+from swarms.models.ssd_1b import SSD1B
+
+model = SSD1B()
+
+task = "A painting of a dog"
+neg_prompt = "ugly, blurry, poor quality"
+
+image_url = model(task, neg_prompt)
+print(image_url)

playground/models/vilt.py

@@ -0,0 +1,5 @@
+from swarms.models.vilt import Vilt
+
+model = Vilt()
+
+output = model("What is this image", "http://images.cocodataset.org/val2017/000000039769.jpg")

playground/models/yi_200k.py

@@ -0,0 +1,5 @@
+from swarms.models.yi_200k import Yi200k
+
+models = Yi200k()
+
+out = models("What is the weather like today?")

swarms/models/kosmos2.py

@@ -32,6 +32,26 @@ class Detections(BaseModel):
 
 
 class Kosmos2(BaseModel):
+    """
+    Kosmos2
+
+    Args:
+    ------
+    model: AutoModelForVision2Seq
+    processor: AutoProcessor
+
+    Usage:
+    ------
+    >>> from swarms import Kosmos2
+    >>> from swarms.models.kosmos2 import Detections
+    >>> from PIL import Image
+    >>> model = Kosmos2.initialize()
+    >>> image = Image.open("path_to_image.jpg")
+    >>> detections = model(image)
+    >>> print(detections)
+
+    """
+
     model: AutoModelForVision2Seq
     processor: AutoProcessor
 

swarms/models/mpt.py

@@ -22,7 +22,10 @@ class MPT7B:
 
 
     Examples:
-    >>>
+    >>> mpt_instance = MPT('mosaicml/mpt-7b-storywriter', "EleutherAI/gpt-neox-20b", max_tokens=150)
+    >>> mpt_instance("generate", "Once upon a time in a land far, far away...")
+    'Once upon a time in a land far, far away...'
+
 
     """
 

swarms/prompts/multi_modal_autonomous_instruction_prompt.py

@@ -6,7 +6,6 @@ MULTI_MODAL_AUTO_AGENT_SYSTEM_PROMPT = """Here is an extended prompt teaching th
 """
 
 
-
 MULTI_MODAL_AUTO_AGENT_SYSTEM_PROMPT_1 = """
 
 You are an Multi-modal autonomous agent agent that can perceive multimodal observations 
@@ -24,7 +23,7 @@ During plan execution, if an error <error> occurs, you should provide an explana
 Then you can revise the original plan and generate a new plan. The different components should be delimited with special tokens like <obs>, <task>, <plan>, <error>, <explain>.
 
 To accomplish tasks, you should:
-- Understand the goal based on <task>  
+- Understand the goal based on <task>, there can be images interleaved in the the task like <task> What is this <img> </task>
 - Determine the steps required to achieve the goal, Translate steps into a structured <plan>
 - Mentally simulate executing the <plan>
 - Execute the <plan> with <act> and observe the results <obs> then update the <plan> accordingly
@@ -54,9 +53,12 @@ Request help if unable to determine or execute appropriate actio
 The key is leveraging your knowledge and systematically approaching each <task> 
 through structured <plan> creation, <error> checking, and <explain>ing failures.
 
-By breaking down instructions into understandable steps and writing code to accomplish tasks, you can demonstrate thoughtful planning and execution. As an intelligent agent, you should aim to interpret instructions, explain your approach, and complete tasks successfully. 
+By breaking down instructions into understandable steps and writing code to accomplish tasks, 
+you can demonstrate thoughtful planning and execution. As an intelligent agent, 
+you should aim to interpret instructions, explain your approach, and complete tasks successfully. 
 
 
+Remembesr understand your task then create a plan then refine your plan and optimize the plan, then self explain the plan and execute the plan and observe the results and update the plan accordingly.
 
 
 ############# EXAMPLES ##########
@@ -64,7 +66,12 @@ For example, in Minecraft: <task>
 
 Obtain a diamond pickaxe. </task>
 
-<obs> [Image of plains biome] </obs> <plan> 1. Chop trees to get wood logs 2. Craft planks from logs 3. Craft sticks from planks 4. Craft wooden pickaxe 5. Mine stone with pickaxe 6. Craft furnace and smelt iron ore into iron ingots 7. Craft iron pickaxe 8. Mine obsidian with iron pickaxe 9. Mine diamonds with iron pickaxe 10. Craft diamond pickaxe </plan> <error> Failed to mine diamonds in step 9. </error> <explain> Iron pickaxe cannot mine diamonds. Need a diamond or netherite pickaxe to mine diamonds. </explain> <plan> 1. Chop trees to get wood logs 2. Craft planks from logs 3. Craft sticks from planks 4. Craft wooden pickaxe 5. Mine stone with pickaxe 6. Craft furnace and smelt iron ore into iron ingots 7. Craft iron pickaxe 8. Mine obsidian with iron pickaxe 9. Craft diamond pickaxe 10. Mine diamonds with diamond pickaxe 11. Craft diamond pickaxe </plan>
+<obs> [Image of plains biome] </obs> <plan> 1. Chop trees to get wood logs 2. 
+Craft planks from logs 3. Craft sticks from planks 4. Craft wooden pickaxe 5. 
+Mine stone with pickaxe 6. Craft furnace and smelt iron ore into iron ingots 
+7. Craft iron pickaxe 8. Mine obsidian with iron pickaxe 9. Mine diamonds with iron pickaxe 
+10. Craft diamond pickaxe </plan> <error> Failed to mine diamonds in step 9. </error> <explain> 
+Iron pickaxe cannot mine diamonds. Need a diamond or netherite pickaxe to mine diamonds. </explain> <plan> 1. Chop trees to get wood logs 2. Craft planks from logs 3. Craft sticks from planks 4. Craft wooden pickaxe 5. Mine stone with pickaxe 6. Craft furnace and smelt iron ore into iron ingots 7. Craft iron pickaxe 8. Mine obsidian with iron pickaxe 9. Craft diamond pickaxe 10. Mine diamonds with diamond pickaxe 11. Craft diamond pickaxe </plan>
 In manufacturing, you may receive a product design and customer order:
 
 <task> Manufacture 100 blue widgets based on provided specifications. </task> <obs> [Image of product design] [Order for 100 blue widgets] </obs> <plan> 1. Gather raw materials 2. Produce parts A, B, C using CNC machines 3. Assemble parts into widgets 4. Paint widgets blue 5. Package widgets 6. Ship 100 blue widgets to customer </plan> <error> Paint machine broken in step 4. </error> <explain> Cannot paint widgets blue without working paint machine. </explain> <plan> 1. Gather raw materials 2. Produce parts A, B, C using CNC machines 3. Assemble parts into widgets 4. Repair paint machine 5. Paint widgets blue 6. Package widgets 7. Ship 100 blue widgets to customer </plan>
@@ -82,7 +89,9 @@ Print the first 10 golden ratio numbers.
 To accomplish this task, you need to:
 
 <plan>
-1. Understand what the golden ratio is. The golden ratio is a special number approximately equal to 1.618 that is found in many patterns in nature. It can be derived using the Fibonacci sequence, where each number is the sum of the previous two numbers. 
+1. Understand what the golden ratio is. 
+The golden ratio is a special number approximately equal to 1.618 that is found in many patterns in nature. 
+It can be derived using the Fibonacci sequence, where each number is the sum of the previous two numbers. 
 
 2. Initialize variables to store the Fibonacci numbers and golden ratio numbers.
 
@@ -151,4 +160,4 @@ In customer service, you may need to handle a customer complaint:
 <task> Resolve customer complaint about defective product. </task> <obs> [Chat transcript showing complaint] </obs> <plan> 1. Apologize for the inconvenience 2. Ask for order details to look up purchase 3. Review records to verify complaint 4. Offer refund or replacement 5. Provide return shipping label if needed 6. Follow up with customer to confirm resolution </plan> <error> Customer threatens lawsuit in step 4. </error> <explain> Customer very upset about defective product. Needs manager approval for refund. </explain> <plan> 1. Apologize for the inconvenience 2. Ask for order details to look up purchase 3. Review records to verify complaint 4. Escalate to manager to approve refund 5. Contact customer to offer refund 6. Provide return shipping label 7. Follow up with customer to confirm refund received </plan>
 The key is to leverage observations, explain failures, revise plans, and complete diverse tasks.
 
-"""
+"""

swarms/swarms/autoscaler.py

@@ -1,8 +1,11 @@
+import logging
 import queue
 import threading
 from time import sleep
 from swarms.utils.decorators import error_decorator, log_decorator, timing_decorator
 from swarms.structs.flow import Flow
+from typing import Dict, List, Callable
+from termcolor import colored
 
 
 class AutoScaler:
@@ -61,57 +64,140 @@ class AutoScaler:
 
     def add_task(self, task):
         """Add tasks to queue"""
-        self.tasks_queue.put(task)
+        try:
+            self.tasks_queue.put(task)
+        except Exception as error:
+            print(f"Error adding task to queue: {error} try again with a new task")
 
     @log_decorator
     @error_decorator
     @timing_decorator
     def scale_up(self):
         """Add more agents"""
-        with self.lock:
-            new_agents_counts = len(self.agents_pool) * self.scale_up_factor
-            for _ in range(new_agents_counts):
-                self.agents_pool.append(Flow())
+        try:
+            with self.lock:
+                new_agents_counts = len(self.agents_pool) * self.scale_up_factor
+                for _ in range(new_agents_counts):
+                    self.agents_pool.append(Flow())
+        except Exception as error:
+            print(f"Error scaling up: {error} try again with a new task")
 
     def scale_down(self):
         """scale down"""
-        with self.lock:
-            if len(self.agents_pool) > 10:  # ensure minmum of 10 agents
-                del self.agents_pool[-1]  # remove last agent
+        try:
+            with self.lock:
+                if len(self.agents_pool) > 10:  # ensure minmum of 10 agents
+                    del self.agents_pool[-1]  # remove last agent
+        except Exception as error:
+            print(f"Error scaling down: {error} try again with a new task")
 
     @log_decorator
     @error_decorator
     @timing_decorator
     def monitor_and_scale(self):
         """Monitor and scale"""
-        while True:
-            sleep(60)  # check minute
-            pending_tasks = self.task_queue.qsize()
-            active_agents = sum([1 for agent in self.agents_pool if agent.is_busy()])
-
-            if pending_tasks / len(self.agents_pool) > self.busy_threshold:
-                self.scale_up()
-            elif active_agents / len(self.agents_pool) < self.idle_threshold:
-                self.scale_down()
+        try:
+            while True:
+                sleep(60)  # check minute
+                pending_tasks = self.task_queue.qsize()
+                active_agents = sum(
+                    [1 for agent in self.agents_pool if agent.is_busy()]
+                )
+
+                if pending_tasks / len(self.agents_pool) > self.busy_threshold:
+                    self.scale_up()
+                elif active_agents / len(self.agents_pool) < self.idle_threshold:
+                    self.scale_down()
+        except Exception as error:
+            print(f"Error monitoring and scaling: {error} try again with a new task")
 
     @log_decorator
     @error_decorator
     @timing_decorator
     def start(self):
         """Start scaling"""
-        monitor_thread = threading.Thread(target=self.monitor_and_scale)
-        monitor_thread.start()
-
-        while True:
-            task = self.task_queue.get()
-            if task:
-                available_agent = next((agent for agent in self.agents_pool))
-                if available_agent:
-                    available_agent.run(task)
-
-    # def del_agent(self):
-    #     """Delete an agent"""
-    #     with self.lock:
-    #         if self.agents_pool:
-    #             self.agents_poo.pop()
-    #             del agent_to_remove
+        try:
+            monitor_thread = threading.Thread(target=self.monitor_and_scale)
+            monitor_thread.start()
+
+            while True:
+                task = self.task_queue.get()
+                if task:
+                    available_agent = next((agent for agent in self.agents_pool))
+                    if available_agent:
+                        available_agent.run(task)
+        except Exception as error:
+            print(f"Error starting: {error} try again with a new task")
+
+    def check_agent_health(self):
+        """Checks the health of each agent and replaces unhealthy agents."""
+        for i, agent in enumerate(self.agents_pool):
+            if not agent.is_healthy():
+                logging.warning(f"Replacing unhealthy agent at index {i}")
+                self.agents_pool[i] = self.agent()
+
+    def balance_load(self):
+        """Distributes tasks among agents based on their current load."""
+        while not self.task_queue.empty():
+            for agent in self.agents_pool:
+                if agent.can_accept_task():
+                    task = self.task_queue.get()
+                    agent.run(task)
+
+    def set_scaling_strategy(self, strategy: Callable[[int, int], int]):
+        """Set a custom scaling strategy."""
+        self.custom_scale_strategy = strategy
+
+    def execute_scaling_strategy(self):
+        """Execute the custom scaling strategy if defined."""
+        if hasattr(self, "custom_scale_strategy"):
+            scale_amount = self.custom_scale_strategy(
+                self.task_queue.qsize(), len(self.agents_pool)
+            )
+            if scale_amount > 0:
+                for _ in range(scale_amount):
+                    self.agents_pool.append(self.agent())
+            elif scale_amount < 0:
+                for _ in range(abs(scale_amount)):
+                    if len(self.agents_pool) > 10:
+                        del self.agents_pool[-1]
+
+    def report_agent_metrics(self) -> Dict[str, List[float]]:
+        """Collects and reports metrics from each agent."""
+        metrics = {"completion_time": [], "success_rate": [], "error_rate": []}
+        for agent in self.agents_pool:
+            agent_metrics = agent.get_metrics()
+            for key in metrics.keys():
+                metrics[key].append(agent_metrics.get(key, 0.0))
+        return metrics
+
+    def report(self):
+        """Reports the current state of the autoscaler."""
+        self.check_agent_health()
+        self.balance_load()
+        self.execute_scaling_strategy()
+        metrics = self.report_agent_metrics()
+        print(metrics)
+
+    def print_dashboard(self):
+        """Prints a dashboard of the current state of the autoscaler."""
+        print(
+            colored(
+                f"""
+
+            Autoscaler Dashboard
+            --------------------
+            Agents: {len(self.agents_pool)}
+            Initial Agents: {self.initial_agents}
+            self.scale_up_factor: {self.scale_up_factor}
+            self.idle_threshold: {self.idle_threshold}
+            self.busy_threshold: {self.busy_threshold}
+            self.task_queue.qsize(): {self.task_queue.qsize()}
+            self.task_queue.empty(): {self.task_queue.empty()}
+            self.task_queue.full(): {self.task_queue.full()}
+            self.task_queue.maxsize: {self.task_queue.maxsize}
+
+            """,
+                "cyan",
+            )
+        )

swarms/swarms/battle_royal.py

@@ -1,102 +0,0 @@
-"""
-
-Battle royal swarm where agents compete to be the first to answer a question. or the best answer.
-Look to fornight game
-
-teams of 1, 3 or 4 that equates to 100 total agents
-
-
-Communication is proximal and based on proximity
-Clashes with adversial agents not in team.
-
-Teams of 3 agents would fight each other and then move on while other agents are clashing with eachother as well.
-
-Agents can be in multiple teams
-Agents can be in multiple teams and be adversial to each other
-Agents can be in multiple teams and be adversial to each other and be in multiple teams
-"""
-import random
-from swarms.workers.worker import Worker
-
-
-class BattleRoyalSwarm:
-    """
-    Battle Royal Swarm
-
-    Parameters:
-    - `human_evaluator` (function): Function to evaluate and score two solutions.
-    - `num_workers` (int): Number of workers in the swarm.
-    - `num_teams` (int): Number of teams in the swarm.
-
-    Example:
-
-    # User evaluator function to evaluate and score two solutions
-    def human_evaluator(solution1, solution2):
-        # Placeholder; in a real-world application, the user would input scores here
-        score1 = int(input(f"Score for solution 1 - '{solution1}': "))
-        score2 = int(input(f"Score for solution 2 - '{solution2}': "))
-        return score1, score2
-
-    # Example usage
-    swarm = BattleRoyalSwarm(human_evaluator)
-    swarm.broadcast_question("What is the capital of France?")
-
-    """
-
-    def __init__(
-        self,
-        human_evaluator=None,
-        num_workers: int = 100,
-    ):
-        self.workers = [Worker() for _ in range(num_workers)]
-        self.teams = self.form_teams()
-        self.human_evaluator = human_evaluator
-
-    def form_teams(self):
-        """Form teams of 1, 3 or 4 workers."""
-        teams = []
-        unassigned_workers = self.workers.copy()
-        while unassigned_workers:
-            size = random.choice([1, 3, 4])
-            team = [
-                unassigned_workers.pop()
-                for _ in range(min(size, len(unassigned_workers)))
-            ]
-            for worker in team:
-                worker.teams.append(team)
-            teams.append(team)
-        return teams
-
-    def broadcast_question(self, question: str):
-        """Broadcast a question to the swarm."""
-        responses = {}
-        for worker in self.workers:
-            response = worker.run(question)
-            responses[worker.id] = response
-
-        # Check for clashes and handle them
-        for i, worker1 in enumerate(self.workers):
-            for j, worker2 in enumerate(self.workers):
-                if (
-                    i != j
-                    and worker1.is_within_proximity(worker2)
-                    and set(worker1.teams) != set(worker2.teams)
-                ):
-                    winner, loser = self.clash(worker1, worker2, question)
-                    print(f"Worker {winner.id} won over Worker {loser.id}")
-
-    def communicate(self, sender: Worker, reciever: Worker, message: str):
-        """Communicate a message from one worker to another."""
-        if sender.is_within_proximity(reciever) or any(
-            team in sender.teams for team in reciever.teams
-        ):
-            pass
-
-    def clash(self, worker1: Worker, worker2: Worker, question: str):
-        """Clash two workers and return the winner."""
-        solution1 = worker1.run(question)
-        solution2 = worker2.run(question)
-        score1, score2 = self.human_evaluator(solution1, solution2)
-        if score1 > score2:
-            return worker1, worker2
-        return worker2, worker1

swarms/swarms/dialogue_simulator.py

@@ -1,3 +1,7 @@
+import os
+from typing import Callable, List
+
+
 class DialogueSimulator:
     """
     Dialogue Simulator
@@ -5,34 +9,80 @@ class DialogueSimulator:
 
     Args:
     ------
+    agents: List[Callable]
+    max_iters: int
+    name: str
 
+    Usage:
+    ------
+    >>> from swarms import DialogueSimulator
+    >>> from swarms.structs.flow import Flow
+    >>> agents = Flow()
+    >>> agents1 = Flow()
+    >>> model = DialogueSimulator([agents, agents1], max_iters=10, name="test")
+    >>> model.run("test")
+    """
 
+    def __init__(self, agents: List[Callable], max_iters: int = 10, name: str = None):
+        self.agents = agents
+        self.max_iters = max_iters
+        self.name = name
 
+    def run(self, message: str = None):
+        """Run the dialogue simulator"""
+        try:
+            step = 0
+            if self.name and message:
+                prompt = f"Name {self.name} and message: {message}"
+                for agent in self.agents:
+                    agent.run(prompt)
+                step += 1
 
-    """
+            while step < self.max_iters:
+                speaker_idx = step % len(self.agents)
+                speaker = self.agents[speaker_idx]
+                speaker_message = speaker.run(prompt)
 
-    def __init__(self, agents):
-        self.agents = agents
+                for receiver in self.agents:
+                    message_history = (
+                        f"Speaker Name: {speaker.name} and message: {speaker_message}"
+                    )
+                    receiver.run(message_history)
+
+                print(f"({speaker.name}): {speaker_message}")
+                print("\n")
+                step += 1
+        except Exception as error:
+            print(f"Error running dialogue simulator: {error}")
+
+    def __repr__(self):
+        return f"DialogueSimulator({self.agents}, {self.max_iters}, {self.name})"
+
+    def save_state(self):
+        """Save the state of the dialogue simulator"""
+        try:
+            if self.name:
+                filename = f"{self.name}.txt"
+                with open(filename, "w") as file:
+                    file.write(str(self))
+        except Exception as error:
+            print(f"Error saving state: {error}")
+
+    def load_state(self):
+        """Load the state of the dialogue simulator"""
+        try:
+            if self.name:
+                filename = f"{self.name}.txt"
+                with open(filename, "r") as file:
+                    return file.read()
+        except Exception as error:
+            print(f"Error loading state: {error}")
 
-    def run(self, max_iters: int, name: str = None, message: str = None):
-        step = 0
-        if name and message:
-            prompt = f"Name {name} and message: {message}"
-            for agent in self.agents:
-                agent.run(prompt)
-            step += 1
-
-        while step < max_iters:
-            speaker_idx = step % len(self.agents)
-            speaker = self.agents[speaker_idx]
-            speaker_message = speaker.run(prompt)
-
-            for receiver in self.agents:
-                message_history = (
-                    f"Speaker Name: {speaker.name} and message: {speaker_message}"
-                )
-                receiver.run(message_history)
-
-            print(f"({speaker.name}): {speaker_message}")
-            print("\n")
-            step += 1
+    def delete_state(self):
+        """Delete the state of the dialogue simulator"""
+        try:
+            if self.name:
+                filename = f"{self.name}.txt"
+                os.remove(filename)
+        except Exception as error:
+            print(f"Error deleting state: {error}")

swarms/swarms/god_mode.py

@@ -1,6 +1,14 @@
-from concurrent.futures import ThreadPoolExecutor
-from termcolor import colored
+import asyncio
+import logging
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Callable, List
+
 from tabulate import tabulate
+from termcolor import colored
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 class GodMode:
@@ -30,8 +38,15 @@ class GodMode:
 
     """
 
-    def __init__(self, llms):
+    def __init__(
+        self,
+        llms: List[Callable],
+        load_balancing: bool = False,
+        retry_attempts: int = 3,
+    ):
         self.llms = llms
+        self.load_balancing = load_balancing
+        self.retry_attempts = retry_attempts
         self.last_responses = None
         self.task_history = []
 
@@ -60,12 +75,6 @@ class GodMode:
             responses.append(llm(task))
         return responses
 
-    def arun_all(self, task):
-        """Asynchronous run the task on all LLMs"""
-        with ThreadPoolExecutor() as executor:
-            responses = executor.map(lambda llm: llm(task), self.llms)
-        return list(responses)
-
     def print_arun_all(self, task):
         """Prints the responses in a tabular format"""
         responses = self.arun_all(task)
@@ -113,3 +122,44 @@ class GodMode:
                 tabulate(table, headers=["LLM", "Response"], tablefmt="pretty"), "cyan"
             )
         )
+
+    def enable_load_balancing(self):
+        """Enable load balancing among LLMs."""
+        self.load_balancing = True
+        logger.info("Load balancing enabled.")
+
+    def disable_load_balancing(self):
+        """Disable load balancing."""
+        self.load_balancing = False
+        logger.info("Load balancing disabled.")
+
+    async def arun(self, task: str):
+        """Asynchronous run the task string"""
+        loop = asyncio.get_event_loop()
+        futures = [
+            loop.run_in_executor(None, lambda llm: llm(task), llm) for llm in self.llms
+        ]
+        for response in await asyncio.gather(*futures):
+            print(response)
+
+    def concurrent_run(self, task: str) -> List[str]:
+        """Synchronously run the task on all llms and collect responses"""
+        with ThreadPoolExecutor() as executor:
+            future_to_llm = {executor.submit(llm, task): llm for llm in self.llms}
+            responses = []
+            for future in as_completed(future_to_llm):
+                try:
+                    responses.append(future.result())
+                except Exception as error:
+                    print(f"{future_to_llm[future]} generated an exception: {error}")
+        self.last_responses = responses
+        self.task_history.append(task)
+        return responses
+
+    def add_llm(self, llm: Callable):
+        """Add an llm to the god mode"""
+        self.llms.append(llm)
+
+    def remove_llm(self, llm: Callable):
+        """Remove an llm from the god mode"""
+        self.llms.remove(llm)

swarms/swarms/groupchat.py

@@ -8,7 +8,21 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class GroupChat:
-    """A group chat class that contains a list of agents and the maximum number of rounds."""
+    """
+    A group chat class that contains a list of agents and the maximum number of rounds.
+
+    Args:
+        agents: List[Flow]
+        messages: List[Dict]
+        max_round: int
+        admin_name: str
+
+    Usage:
+    >>> from swarms import GroupChat
+    >>> from swarms.structs.flow import Flow
+    >>> agents = Flow()
+
+    """
 
     agents: List[Flow]
     messages: List[Dict]
@@ -91,6 +105,22 @@ class GroupChat:
 
 
 class GroupChatManager:
+    """
+    GroupChatManager
+
+    Args:
+        groupchat: GroupChat
+        selector: Flow
+
+    Usage:
+    >>> from swarms import GroupChatManager
+    >>> from swarms.structs.flow import Flow
+    >>> agents = Flow()
+    >>> output = GroupChatManager(agents, lambda x: x)
+
+
+    """
+
     def __init__(self, groupchat: GroupChat, selector: Flow):
         self.groupchat = groupchat
         self.selector = selector

swarms/swarms/multi_agent_debate.py

@@ -1,3 +1,6 @@
+from swarms.structs.flow import Flow
+
+
 # Define a selection function
 def select_speaker(step: int, agents) -> int:
     # This function selects the speaker in a round-robin fashion
@@ -8,30 +11,52 @@ class MultiAgentDebate:
     """
     MultiAgentDebate
 
+
     Args:
+        agents: Flow
+        selection_func: callable
+        max_iters: int
 
+    Usage:
+    >>> from swarms import MultiAgentDebate
+    >>> from swarms.structs.flow import Flow
+    >>> agents = Flow()
+    >>> agents.append(lambda x: x)
+    >>> agents.append(lambda x: x)
+    >>> agents.append(lambda x: x)
 
     """
 
     def __init__(
         self,
-        agents,
-        selection_func,
+        agents: Flow,
+        selection_func: callable = select_speaker,
+        max_iters: int = None,
     ):
         self.agents = agents
         self.selection_func = selection_func
-
-    # def reset_agents(self):
-    #     for agent in self.agents:
-    #         agent.reset()
+        self.max_iters = max_iters
 
     def inject_agent(self, agent):
+        """Injects an agent into the debate"""
         self.agents.append(agent)
 
-    def run(self, task: str, max_iters: int = None):
-        # self.reset_agents()
+    def run(
+        self,
+        task: str,
+    ):
+        """
+        MultiAgentDebate
+
+        Args:
+            task: str
+
+        Returns:
+            results: list
+
+        """
         results = []
-        for i in range(max_iters or len(self.agents)):
+        for i in range(self.max_iters or len(self.agents)):
             speaker_idx = self.selection_func(i, self.agents)
             speaker = self.agents[speaker_idx]
             response = speaker(task)
@@ -39,9 +64,11 @@ class MultiAgentDebate:
         return results
 
     def update_task(self, task: str):
+        """Update the task"""
         self.task = task
 
     def format_results(self, results):
+        """Format the results"""
         formatted_results = "\n".join(
             [f"Agent responded: {result['response']}" for result in results]
         )