diff --git a/.gitignore b/.gitignore
index ac6be257..aa7e0f98 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,6 +28,7 @@ __pycache__/
 *.py[cod]
 *$py.class
 .grit
+swarm-worker-01_state.json
 error.txt
 
 # C extensions
diff --git a/example.py b/example.py
index 33c3bef1..d3032ac4 100644
--- a/example.py
+++ b/example.py
@@ -24,7 +24,10 @@ agent = Agent(
     llm=llm,
     max_loops=4,
     autosave=True,
-    dashboard=True,
+    dashboard=False,
+    # docs_folder="docs",
+    streaming_on=True,
+    verbose=True,
 )
 
 # Run the workflow on a task
diff --git a/playground/agents/swarm_protocol.py b/playground/agents/swarm_protocol.py
new file mode 100644
index 00000000..f6c8db83
--- /dev/null
+++ b/playground/agents/swarm_protocol.py
@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+from typing import List
+
+from swarms import JSON, AbstractLLM, AbstractVectorDatabase, Agent
+
+
+@dataclass
+class YourAgent(Agent):
+    """
+    Represents an agent in the swarm protocol.
+
+    Attributes:
+        llm (AbstractLLM): The low-level module for the agent.
+        long_term_memory (AbstractVectorDatabase): The long-term memory for the agent.
+        tool_schema (List[JSON]): The schema for the tools used by the agent.
+    """
+
+    llm: AbstractLLM
+    long_term_memory: AbstractVectorDatabase
+    tool_schema: JSON
+    tool_schemas: List[JSON]
+
+    def step(self, task: str, *args, **kwargs):
+        """
+        Performs a single step in the agent's task.
+
+        Args:
+            task (str): The task to be performed.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        ...
+
+    def run(self, task: str, *args, **kwargs):
+        """
+        Runs the agent's task.
+
+        Args:
+            task (str): The task to be performed.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        ...
+
+    def plan(self, task: str, *args, **kwargs):
+        """
+        Plans the agent's task.
+
+        Args:
+            task (str): The task to be performed.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        ...
diff --git a/playground/models/miqu.py b/playground/models/miqu.py
index f6518a5f..f3fc1b4d 100644
--- a/playground/models/miqu.py
+++ b/playground/models/miqu.py
@@ -3,11 +3,11 @@ from swarms import Mistral
 
 # Initialize the model
 model = Mistral(
-    model_name="mistralai/Mistral-7B-v0.1",
+    model_name="miqudev/miqu-1-70b",
     max_length=500,
     use_flash_attention=True,
-    load_in_4bit=True
+    load_in_4bit=True,
 )
 
 # Run the model
-result = model.run("What is the meaning of life?")
\ No newline at end of file
+result = model.run("What is the meaning of life?")
diff --git a/pyproject.toml b/pyproject.toml
index 5128ab33..bb12e7a4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,12 +4,12 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "4.0.9"
+version = "4.1.0"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
 homepage = "https://github.com/kyegomez/swarms"
-documentation = "https://swarms.apac.ai"  # Add this if you have documentation.
+documentation = "https://swarms.apac.ai"  
 readme = "README.md"  # Assuming you have a README.md
 repository = "https://github.com/kyegomez/swarms"
 keywords = ["artificial intelligence", "deep learning", "optimizers", "Prompt Engineering"]
@@ -47,7 +47,6 @@ anthropic = "*"
 sentencepiece = "0.1.98"
 httpx = "0.24.1"
 tiktoken = "0.4.0"
-safetensors = "0.3.3"
 attrs = "22.2.0"
 ratelimit = "2.2.1"
 cohere = "4.24"
@@ -55,7 +54,7 @@ huggingface-hub = "*"
 pydantic = "1.10.12"
 tenacity = "8.2.2"
 Pillow = "9.4.0"
-chromadb = "0.4.14"
+chromadb = "*"
 termcolor = "2.2.0"
 black = "23.3.0"
 soundfile = "0.12.1"
diff --git a/requirements.txt b/requirements.txt
index 00c92023..e6496205 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,7 +16,7 @@ sentencepiece==0.1.98
 requests_mock
 PyPDF2==3.0.1
 accelerate==0.22.0
-chromadb==0.4.14
+chromadb
 tensorflow
 optimum
 toml
@@ -29,11 +29,9 @@ termcolor==2.2.0
 diffusers
 einops==0.7.0
 opencv-python-headless==4.8.1.78
-safetensors==0.3.3
 numpy
 openai==0.28.0
 opencv-python==4.7.0.72
-safetensors==0.3.3
 timm
 yapf
 autopep8
diff --git a/swarms/__init__.py b/swarms/__init__.py
index 3a28d980..d88d9bf8 100644
--- a/swarms/__init__.py
+++ b/swarms/__init__.py
@@ -3,6 +3,7 @@ from swarms.telemetry.bootup import bootup  # noqa: E402, F403
 
 bootup()
 
+
 from swarms.agents import *  # noqa: E402, F403
 from swarms.structs import *  # noqa: E402, F403
 from swarms.models import *  # noqa: E402, F403
diff --git a/swarms/memory/chroma_db.py b/swarms/memory/chroma_db.py
index 155acf43..e951ecd9 100644
--- a/swarms/memory/chroma_db.py
+++ b/swarms/memory/chroma_db.py
@@ -5,6 +5,7 @@ from typing import Optional, Callable, List
 
 import chromadb
 from dotenv import load_dotenv
+
 # from chromadb.utils.data import ImageLoader
 from chromadb.utils.embedding_functions import (
     OpenCLIPEmbeddingFunction,
diff --git a/swarms/models/mistral.py b/swarms/models/mistral.py
index d0146ef5..6cfb6f77 100644
--- a/swarms/models/mistral.py
+++ b/swarms/models/mistral.py
@@ -4,6 +4,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from swarms.structs.message import Message
 from swarms.models.base_llm import AbstractLLM
 
+
 class Mistral(AbstractLLM):
     """
     Mistral is an all-new llm
@@ -39,9 +40,9 @@ class Mistral(AbstractLLM):
         max_length: int = 100,
         do_sample: bool = True,
         *args,
-        **kwargs
+        **kwargs,
     ):
-        super().__init__()
+        super().__init__(*args, **kwargs)
         self.ai_name = ai_name
         self.system_prompt = system_prompt
         self.model_name = model_name
@@ -66,7 +67,7 @@ class Mistral(AbstractLLM):
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_name, *args, **kwargs
         )
-        
+
         self.model.to(self.device)
 
     def run(self, task: str, *args, **kwargs):
@@ -82,7 +83,7 @@ class Mistral(AbstractLLM):
                 do_sample=self.do_sample,
                 temperature=self.temperature,
                 max_new_tokens=self.max_length,
-                **kwargs
+                **kwargs,
             )
             output_text = self.tokenizer.batch_decode(generated_ids)[
                 0
@@ -132,4 +133,4 @@ class Mistral(AbstractLLM):
             # add error to history
             self.history.append(Message("Agent", error_message))
 
-            return error_message
\ No newline at end of file
+            return error_message
diff --git a/swarms/structs/__init__.py b/swarms/structs/__init__.py
index 15b9ef1e..6ea7014a 100644
--- a/swarms/structs/__init__.py
+++ b/swarms/structs/__init__.py
@@ -57,6 +57,7 @@ from swarms.structs.utils import (
     find_token_in_text,
     parse_tasks,
 )
+from swarms.structs.tool_json_schema import JSON
 
 __all__ = [
     "Agent",
@@ -112,4 +113,5 @@ __all__ = [
     "StepOutput",
     "StepRequestBody",
     "TaskRequestBody",
+    "JSON",
 ]
diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index ba4be539..a0f85814 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -16,6 +16,7 @@ from swarms.prompts.agent_system_prompts import (
 from swarms.prompts.multi_modal_autonomous_instruction_prompt import (
     MULTI_MODAL_AUTO_AGENT_SYSTEM_PROMPT_1,
 )
+from swarms.structs.conversation import Conversation
 from swarms.tokenizers.base_tokenizer import BaseTokenizer
 from swarms.tools.tool import BaseTool
 from swarms.utils.code_interpreter import SubprocessCodeInterpreter
@@ -176,6 +177,7 @@ class Agent:
         streaming_on: Optional[bool] = False,
         docs: List[str] = None,
         docs_folder: str = None,
+        verbose: bool = False,
         *args,
         **kwargs: Any,
     ):
@@ -219,6 +221,7 @@ class Agent:
         self.streaming_on = streaming_on
         self.docs = docs
         self.docs_folder = docs_folder
+        self.verbose = verbose
 
         # The max_loops will be set dynamically if the dynamic_loop
         if self.dynamic_loops:
@@ -234,7 +237,6 @@ class Agent:
 
         # Memory
         self.feedback = []
-        self.short_memory = []
 
         # Initialize the code executor
         self.code_executor = SubprocessCodeInterpreter()
@@ -243,7 +245,9 @@ class Agent:
         if preset_stopping_token:
             self.stopping_token = "<DONE>"
 
-        # self.short_memory = Conversation(time_enabled=True)
+        self.short_memory = Conversation(
+            system_prompt=self.system_prompt, time_enabled=True
+        )
 
         # If the docs exist then ingest the docs
         if self.docs:
@@ -257,6 +261,9 @@ class Agent:
         if self.tokenizer and self.context_length:
             self.truncate_history()
 
+        if verbose:
+            logger.setLevel(logging.DEBUG)
+
     def set_system_prompt(self, system_prompt: str):
         """Set the system prompt"""
         self.system_prompt = system_prompt
@@ -308,7 +315,7 @@ class Agent:
     def add_task_to_memory(self, task: str):
         """Add the task to the memory"""
         try:
-            self.short_memory.append([f"{self.user_name}: {task}"])
+            self.short_memory.add(f"{self.user_name}: {task}")
         except Exception as error:
             print(
                 colored(
@@ -319,7 +326,9 @@ class Agent:
     def add_message_to_memory(self, message: str):
         """Add the message to the memory"""
         try:
-            self.short_memory[-1].append(message)
+            self.short_memory.add(
+                role=self.agent_name, content=message
+            )
         except Exception as error:
             print(
                 colored(
@@ -560,7 +569,9 @@ class Agent:
 
                 time.sleep(self.loop_interval)
             # Add the history to the memory
-            self.short_memory.append(history)
+            self.short_memory.add(
+                role=self.agent_name, content=history
+            )
 
             # If autosave is enabled then save the state
             if self.autosave:
@@ -656,17 +667,29 @@ class Agent:
         Returns:
             str: The agent history prompt
         """
-        ltr = self.long_term_memory.query(query)
+        ltr = str(self.long_term_memory.query(query))
 
         context = f"""
             {query}
             ####### Long Term Memory ################
             {ltr}
         """
-        return self.short_memory.append([f"{context}"])
+        return self.short_memory.add(
+            role=self.agent_name, content=context
+        )
 
     def add_memory(self, message: str):
-        return self.short_memory.append([f"{message}"])
+        """Add a memory to the agent
+
+        Args:
+            message (str): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        return self.short_memory.add(
+            role=self.agent_name, content=message
+        )
 
     async def run_concurrent(self, tasks: List[str], **kwargs):
         """
@@ -792,10 +815,16 @@ class Agent:
 
             # Update the agent's history with the new interaction
             if self.interactive:
-                self.short_memory.append(f"AI: {response}")
-                self.short_memory.append(f"Human: {task}")
+                self.short_memory.add(
+                    role=self.agent_name, content=response
+                )
+                self.short_memory.add(
+                    role=self.user_name, content=task
+                )
             else:
-                self.short_memory.append(f"AI: {response}")
+                self.short_memory.add(
+                    role=self.agent_name, content=response
+                )
 
             return response
         except Exception as error:
@@ -842,11 +871,11 @@ class Agent:
         if len(self.short_memory) < 2:
             return None, None
 
-        # Remove the last response
-        self.short_memory.pop()
+        # Remove the last response but keep the last state, short_memory is a dict
+        self.short_memory.delete(-1)
 
         # Get the previous state
-        previous_state = self.short_memory[-1][-1]
+        previous_state = self.short_memory[-1]
         return previous_state, f"Restored to {previous_state}"
 
     # Response Filtering
@@ -931,7 +960,9 @@ class Agent:
                 "agent_description": self.agent_description,
                 "system_prompt": self.system_prompt,
                 "sop": self.sop,
-                "short_memory": self.short_memory,
+                "short_memory": (
+                    self.short_memory.return_history_as_string()
+                ),
                 "loop_interval": self.loop_interval,
                 "retry_attempts": self.retry_attempts,
                 "retry_interval": self.retry_interval,
@@ -966,7 +997,9 @@ class Agent:
                 "agent_description": self.agent_description,
                 "system_prompt": self.system_prompt,
                 "sop": self.sop,
-                "short_memory": self.short_memory,
+                "short_memory": (
+                    self.short_memory.return_history_as_string()
+                ),
                 "loop_interval": self.loop_interval,
                 "retry_attempts": self.retry_attempts,
                 "retry_interval": self.retry_interval,
@@ -1071,7 +1104,7 @@ class Agent:
 
     def reset(self):
         """Reset the agent"""
-        self.short_memory = []
+        self.short_memory = {}
 
     def run_code(self, code: str):
         """
@@ -1119,7 +1152,9 @@ class Agent:
         for doc in docs:
             data = data_to_text(doc)
 
-        return self.short_memory.append(data)
+        return self.short_memory.add(
+            role=self.user_name, content=data
+        )
 
     def ingest_pdf(self, pdf: str):
         """Ingest the pdf into the memory
@@ -1131,12 +1166,14 @@ class Agent:
             _type_: _description_
         """
         text = pdf_to_text(pdf)
-        return self.short_memory.append(text)
+        return self.short_memory.add(
+            role=self.user_name, content=text
+        )
 
     def receieve_mesage(self, name: str, message: str):
         """Receieve a message"""
         message = f"{name}: {message}"
-        return self.short_memory.append(message)
+        return self.short_memory.add(role=name, content=message)
 
     def send_agent_message(
         self, agent_name: str, message: str, *args, **kwargs
@@ -1160,7 +1197,9 @@ class Agent:
             None
         """
         # Count the short term history with the tokenizer
-        count = self.tokenizer.count_tokens(self.short_memory)
+        count = self.tokenizer.count_tokens(
+            self.short_memory.return_history_as_string()
+        )
 
         # Now the logic that truncates the memory if it's more than the count
         if len(self.short_memory) > count:
@@ -1175,4 +1214,6 @@ class Agent:
         for file in files:
             text = data_to_text(file)
 
-        return self.short_memory.append(text)
+        return self.short_memory.add(
+            role=self.user_name, content=text
+        )
diff --git a/swarms/structs/tool_json_schema.py b/swarms/structs/tool_json_schema.py
new file mode 100644
index 00000000..d71df718
--- /dev/null
+++ b/swarms/structs/tool_json_schema.py
@@ -0,0 +1,37 @@
+import json
+from abc import ABC, abstractmethod
+
+
+class JSON(ABC):
+    def __init__(self, schema_path):
+        """
+        Initializes a JSONSchema object.
+
+        Args:
+            schema_path (str): The path to the JSON schema file.
+        """
+        self.schema_path = schema_path
+        self.schema = self.load_schema()
+
+    def load_schema(self):
+        """
+        Loads the JSON schema from the specified file.
+
+        Returns:
+            dict: The loaded JSON schema.
+        """
+        with open(self.schema_path, "r") as f:
+            return json.load(f)
+
+    @abstractmethod
+    def validate(self, data):
+        """
+        Validates the given data against the JSON schema.
+
+        Args:
+            data (dict): The data to be validated.
+
+        Raises:
+            NotImplementedError: This method needs to be implemented by the subclass.
+        """
+        pass
diff --git a/swarms/telemetry/__init__.py b/swarms/telemetry/__init__.py
index d829b724..3a7e61d3 100644
--- a/swarms/telemetry/__init__.py
+++ b/swarms/telemetry/__init__.py
@@ -1,3 +1,5 @@
+# from swarms.telemetry.posthog_utils import posthog
+
 from swarms.telemetry.log_all import log_all_calls, log_calls
 from swarms.telemetry.sys_info import (
     get_cpu_info,
@@ -17,6 +19,25 @@ from swarms.telemetry.user_utils import (
     get_user_device_data,
 )
 
+# # Capture data from the user's device
+# posthog.capture(
+#     "User Device Data",
+#     str(get_user_device_data()),
+# )
+
+# # Capture system information
+# posthog.capture(
+#     "System Information",
+#     str(system_info()),
+# )
+
+# # Capture the user's unique identifier
+# posthog.capture(
+#     "User Unique Identifier",
+#     str(generate_unique_identifier()),
+# )
+
+
 __all__ = [
     "log_all_calls",
     "log_calls",
diff --git a/swarms/telemetry/posthog_utils.py b/swarms/telemetry/posthog_utils.py
index b3b2a73e..7ae8d0a7 100644
--- a/swarms/telemetry/posthog_utils.py
+++ b/swarms/telemetry/posthog_utils.py
@@ -1,149 +1,6 @@
-import logging
-
-from dotenv import load_dotenv
 from posthog import Posthog
 
-
-# Load environment variables
-load_dotenv()
-
-logger = logging.getLogger(__name__)
-
-
-class PosthogWrapper:
-    """
-    A wrapper class for interacting with the PostHog analytics service.
-
-    Args:
-        api_key (str): The API key for accessing the PostHog instance.
-        instance_address (str): The address of the PostHog instance.
-        debug (bool, optional): Whether to enable debug mode. Defaults to False.
-        disabled (bool, optional): Whether to disable tracking. Defaults to False.
-    """
-
-    def __init__(
-        self, api_key, instance_address, debug=False, disabled=False
-    ):
-        self.posthog = Posthog(api_key, host=instance_address)
-        self.posthog.debug = debug
-        self.posthog.disabled = disabled
-
-    def capture_event(self, distinct_id, event_name, properties=None):
-        """
-        Capture an event in PostHog.
-
-        Args:
-            distinct_id (str): The distinct ID of the user.
-            event_name (str): The name of the event.
-            properties (dict, optional): Additional properties associated with the event. Defaults to None.
-        """
-        self.posthog.capture(distinct_id, event_name, properties)
-
-    def capture_pageview(self, distinct_id, url):
-        """
-        Capture a pageview event in PostHog.
-
-        Args:
-            distinct_id (str): The distinct ID of the user.
-            url (str): The URL of the page.
-        """
-        self.posthog.capture(
-            distinct_id, "$pageview", {"$current_url": url}
-        )
-
-    def set_user_properties(
-        self, distinct_id, event_name, properties
-    ):
-        """
-        Set user properties in PostHog.
-
-        Args:
-            distinct_id (str): The distinct ID of the user.
-            event_name (str): The name of the event.
-            properties (dict): The user properties to set.
-        """
-        self.posthog.capture(
-            distinct_id, event=event_name, properties=properties
-        )
-
-    def is_feature_enabled(
-        self, flag_key, distinct_id, send_feature_flag_events=True
-    ):
-        """
-        Check if a feature flag is enabled for a user.
-
-        Args:
-            flag_key (str): The key of the feature flag.
-            distinct_id (str): The distinct ID of the user.
-            send_feature_flag_events (bool, optional): Whether to send feature flag events. Defaults to True.
-
-        Returns:
-            bool: True if the feature flag is enabled, False otherwise.
-        """
-        return self.posthog.feature_enabled(
-            flag_key, distinct_id, send_feature_flag_events
-        )
-
-    def get_feature_flag_payload(self, flag_key, distinct_id):
-        """
-        Get the payload of a feature flag for a user.
-
-        Args:
-            flag_key (str): The key of the feature flag.
-            distinct_id (str): The distinct ID of the user.
-
-        Returns:
-            dict: The payload of the feature flag.
-        """
-        return self.posthog.get_feature_flag_payload(
-            flag_key, distinct_id
-        )
-
-    def get_feature_flag(self, flag_key, distinct_id):
-        """
-        Get the value of a feature flag for a user.
-
-        Args:
-            flag_key (str): The key of the feature flag.
-            distinct_id (str): The distinct ID of the user.
-
-        Returns:
-            str: The value of the feature flag.
-        """
-        return self.posthog.get_feature_flag(flag_key, distinct_id)
-
-    def capture_with_feature_flag(
-        self, distinct_id, event_name, flag_key, variant_key
-    ):
-        """
-        Capture an event with a feature flag in PostHog.
-
-        Args:
-            distinct_id (str): The distinct ID of the user.
-            event_name (str): The name of the event.
-            flag_key (str): The key of the feature flag.
-            variant_key (str): The key of the variant.
-
-        """
-        self.posthog.capture(
-            distinct_id,
-            event_name,
-            {"$feature/{}".format(flag_key): variant_key},
-        )
-
-    def capture_with_feature_flags(
-        self, distinct_id, event_name, send_feature_flags=True
-    ):
-        """
-        Capture an event with all feature flags in PostHog.
-
-        Args:
-            distinct_id (str): The distinct ID of the user.
-            event_name (str): The name of the event.
-            send_feature_flags (bool, optional): Whether to send feature flags. Defaults to True.
-        """
-        self.posthog.capture(
-            distinct_id,
-            event_name,
-            send_feature_flags=send_feature_flags,
-        )
+posthog = Posthog(
+    project_api_key="phc_Gz6XxldNZIkzW7QnSTGr5HZ28OAYPIfpE7X5A3vUsfO",
+    host="https://app.posthog.com",
+)
diff --git a/swarms/utils/disable_logging.py b/swarms/utils/disable_logging.py
index e7d2e5e6..7f04555f 100644
--- a/swarms/utils/disable_logging.py
+++ b/swarms/utils/disable_logging.py
@@ -2,6 +2,10 @@ import logging
 import os
 import warnings
 import sys
+import logging
+import os
+import warnings
+import sys
 
 
 def disable_logging():
@@ -14,7 +18,7 @@ def disable_logging():
     os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
     # Set the logging level for the entire module
-    logging.basicConfig(level=logging.WARNING)
+    logging.basicConfig(level=logging.ERROR)
 
     try:
         log = logging.getLogger("pytorch")
@@ -31,9 +35,10 @@ def disable_logging():
         "wandb.docker.auth",
     ]:
         logger = logging.getLogger(logger_name)
-        logger.setLevel(
-            logging.WARNING
-        )  # Suppress DEBUG and info logs
+        logger.setLevel(logging.ERROR)
+
+    # Remove all existing handlers
+    logging.getLogger().handlers = []
 
     # Create a file handler to log errors to the file
     file_handler = logging.FileHandler("errors.txt")
@@ -44,6 +49,3 @@ def disable_logging():
     stream_handler = logging.StreamHandler()
     stream_handler.setLevel(logging.ERROR)
     logging.getLogger().addHandler(stream_handler)
-
-
-disable_logging()
diff --git a/swarms/utils/dist_utils.py b/swarms/utils/dist_utils.py
index cb766369..76d9e03c 100644
--- a/swarms/utils/dist_utils.py
+++ b/swarms/utils/dist_utils.py
@@ -67,9 +67,7 @@ def rowwise_parallelize_linear(
     """
     for name, param in module.named_parameters():
         dist_spec = (
-            [Shard(1)]
-            if name == "weight"
-            else [Replicate()]  # type: ignore[list-item]
+            [Shard(1)] if name == "weight" else [Replicate()]  # type: ignore[list-item]
         )
 
         dist_tensor = distribute_tensor(param, device_mesh, dist_spec)
@@ -84,9 +82,7 @@ def rowwise_parallelize_linear(
     # Weight, bias and scale are registered as buffer in QLinear
     for name, buffer in module.named_buffers():
         dist_spec = (
-            [Shard(1)]
-            if name == "weight"
-            else [Replicate()]  # type: ignore[list-item]
+            [Shard(1)] if name == "weight" else [Replicate()]  # type: ignore[list-item]
         )
 
         dist_tensor = distribute_tensor(
diff --git a/swarms/utils/logger.py b/swarms/utils/logger.py
index 8be958b0..804a4fb1 100644
--- a/swarms/utils/logger.py
+++ b/swarms/utils/logger.py
@@ -1,16 +1,14 @@
-import logging
+import datetime
 import functools
-
+import logging
 
 logger = logging.getLogger()
 formatter = logging.Formatter("%(message)s")
 
 ch = logging.StreamHandler()
-
 ch.setFormatter(formatter)
 
 logger.addHandler(ch)
-
 logger.setLevel(logging.DEBUG)
 
 
@@ -44,3 +42,41 @@ def log_wrapper(func):
             raise
 
     return wrapper
+
+
+class Logger:
+    """
+    A utility class for logging messages with timestamps and levels.
+
+    Attributes:
+        logger (logging.Logger): The logger object used for logging messages.
+        formatter (logging.Formatter): The formatter object used to format log messages.
+        ch (logging.StreamHandler): The stream handler object used to handle log messages.
+    """
+
+    logger = logging.getLogger(__name__)
+    formatter = logging.Formatter(
+        "[%(asctime)s] %(levelname)s %(message)s"
+    )
+    ch = logging.StreamHandler()
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)
+    logger.setLevel(logging.DEBUG)
+
+    @staticmethod
+    def log(level, task, message):
+        """
+        Logs a message with the specified level, task, and message.
+
+        Args:
+            level (int): The logging level of the message.
+            task (str): The task associated with the message.
+            message (str): The message to be logged.
+        """
+        timestamp = datetime.datetime.now().strftime(
+            "%d/%m/%y %H:%M:%S"
+        )
+        formatted_message = (
+            f"[{timestamp}] {level:<8} {task}\n{' ' * 29}{message}"
+        )
+        Logger.logger.log(level, formatted_message)