From 932bc1488edb8e940cfc61aabfa73ef5dcbf3235 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Tue, 3 Oct 2023 20:38:36 -0400
Subject: [PATCH] clean up

---
 swarms/agents/omni_modal_agent.py | 51 ++++++++++++++++---------------
 1 file changed, 26 insertions(+), 25 deletions(-)

diff --git a/swarms/agents/omni_modal_agent.py b/swarms/agents/omni_modal_agent.py
index 48c4b8a4..11079014 100644
--- a/swarms/agents/omni_modal_agent.py
+++ b/swarms/agents/omni_modal_agent.py
@@ -41,33 +41,33 @@ from transformers import load_tool
 # #         response = chat_huggingface(messages, api_key, api_type, api_endpoint)
 # #         return response
 
-# class Step:
-#     def __init__(
-#         self,
-#         task: str,
-#         id: int,
-#         dep: List[int],
-#         args: Dict[str, str],
-#         tool: BaseTool
-#     ):
-#         self.task = task
-#         self.id = id
-#         self.dep = dep
-#         self.args = args
-#         self.tool = tool
+class Step:
+    def __init__(
+        self,
+        task: str,
+        id: int,
+        dep: List[int],
+        args: Dict[str, str],
+        tool: BaseTool
+    ):
+        self.task = task
+        self.id = id
+        self.dep = dep
+        self.args = args
+        self.tool = tool
 
-# class Plan:
-#     def __init__(
-#         self,
-#         steps: List[Step]
-#     ):
-#         self.steps = steps
+class Plan:
+    def __init__(
+        self,
+        steps: List[Step]
+    ):
+        self.steps = steps
     
-#     def __str__(self) -> str:
-#         return str([str(step) for step in self.steps])
+    def __str__(self) -> str:
+        return str([str(step) for step in self.steps])
     
-#     def __repr(self) -> str:
-#         return str(self)
+    def __repr(self) -> str:
+        return str(self)
 
 
 
@@ -128,10 +128,11 @@ class OmniModalAgent:
         
         self.chat_planner = load_chat_planner(llm)
         self.response_generator = load_response_generator(llm)
-        self.task_executor: TaskExecutor
+        self.task_executor = TaskExecutor
     
 
     def run(self, input: str) -> str:
+        """Run the OmniAgent"""
         plan = self.chat_planner.plan(
             inputs={
                 "input": input,