From 95472ec6872532f105c2fb0cfcc597c94ac22df9 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Sun, 16 Jul 2023 15:13:17 -0400
Subject: [PATCH] clean up

Former-commit-id: f4d0e8e4b4bb2047e1fd8986b6bc0c35a4a61de9
---
 swarms/agents/boss/BossNode.py           |  2 +-
 swarms/agents/workers/WorkerUltraNode.py | 41 +++++++++++++++++-------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/swarms/agents/boss/BossNode.py b/swarms/agents/boss/BossNode.py
index 8ebf7e7e..d238e276 100644
--- a/swarms/agents/boss/BossNode.py
+++ b/swarms/agents/boss/BossNode.py
@@ -119,4 +119,4 @@ def BossNode(objective, api_key=None, vectorstore=None, worker_node=None, llm_cl
 
     boss = BossNode(llm, vectorstore, agent_executor, max_iterations)
     task = boss.create_task(objective)
-    boss.execute_task(task)
+    boss.run(task)
diff --git a/swarms/agents/workers/WorkerUltraNode.py b/swarms/agents/workers/WorkerUltraNode.py
index c37722a3..0f88c82d 100644
--- a/swarms/agents/workers/WorkerUltraNode.py
+++ b/swarms/agents/workers/WorkerUltraNode.py
@@ -107,15 +107,32 @@ class WorkerUltraNode:
 
 
 
-def WorkerUltra(objective: str, openai_api_key=None):
-    # If the openai_api_key parameter is not provided, try to get the API key from an environment variable # 
-    if openai_api_key is None:
-        openai_api_key = os.getenv('OPENAI_API_KEY')
-
-    if not openai_api_key or not isinstance(openai_api_key, str):
-        logging.error("Invalid OpenAI key")
-        raise ValueError("A valid OpenAI API key is required")
-
-    worker_node = WorkerUltraNode(objective, openai_api_key)
-    # Return the result of the execution
-    return worker_node.result
\ No newline at end of file
+# def WorkerUltra(objective: str, openai_api_key=None):
+#     # If the openai_api_key parameter is not provided, try to get the API key from an environment variable # 
+#     if openai_api_key is None:
+#         openai_api_key = os.getenv('OPENAI_API_KEY')
+
+#     if not openai_api_key or not isinstance(openai_api_key, str):
+#         logging.error("Invalid OpenAI key")
+#         raise ValueError("A valid OpenAI API key is required")
+
+#     worker_node = WorkerUltraNode(objective, openai_api_key)
+#     # Return the result of the execution
+#     return worker_node.result
+
+
+class WorkerUltra:
+    def __init__(self, objective, api_key=None):
+        self.api_key = api_key or os.getenv('OPENAI_API_KEY')
+        if not self.api_key:
+            raise ValueError("API key must be provided either as argument or as an environment variable named 'OPENAI_API_KEY'.")
+        self.worker_node = WorkerUltraNode(objective, self.api_key)
+
+    def execute(self):
+        try:
+            return self.worker_node.execute_task()
+        except Exception as e:
+            logging.error(f"Error while executing: {str(e)}")
+            raise e
+
+