From 6b3fba50eacce6709f690e866365ccaab7c5b972 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Thu, 6 Jul 2023 13:50:28 -0400
Subject: [PATCH] new tool style

---
 swarms/agents/workers/worker.py | 32 +++++++++++++++++++-------------
 swarms/swarms.py                |  3 ++-
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/swarms/agents/workers/worker.py b/swarms/agents/workers/worker.py
index 89fedeaf..66a7a14e 100644
--- a/swarms/agents/workers/worker.py
+++ b/swarms/agents/workers/worker.py
@@ -13,20 +13,22 @@ import logging
 from pydantic import BaseModel, Extra
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-class WorkerNodeArgs(BaseModel):
-    prompt: str
-    run_manager: Optional[CallbackManagerForToolRun] = None
-
-    class Config:
-        arbitrary_types_allowed = True
-        extra = Extra.forbid
-
-@tool("WorkerNode")
-class WorkerNode(BaseTool):
+# class WorkerNodeArgs(BaseModel):
+#     prompt: str
+#     run_manager: Optional[CallbackManagerForToolRun] = None
+
+#     class Config:
+#         arbitrary_types_allowed = True
+#         extra = Extra.forbid
+
+# @tool("WorkerNode")
+class WorkerNode(
+    # BaseTool
+    ):
     """Useful for when you need to spawn an autonomous agent instance as a worker to accomplish complex tasks, it can search the internet or spawn child multi-modality models to process and generate images and text or audio and so on"""
     
-    args_schema: Optional[Type[BaseModel]] = WorkerNodeArgs
-    """Pydantic model class to validate and parse the tool's input arguments."""
+    # args_schema: Optional[Type[BaseModel]] = WorkerNodeArgs
+    # """Pydantic model class to validate and parse the tool's input arguments."""
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -66,7 +68,11 @@ class WorkerNode(BaseTool):
         raise NotImplementedError("WorkerNode does not support async")
 
 
-
+worker_tool = Tool(
+    name="WorkerNode AI Agent",
+    func=WorkerNode._run,
+    description="Useful for when you need to spawn an autonomous agent instance as a worker to accomplish complex tasks, it can search the internet or spawn child multi-modality models to process and generate images and text or audio and so on"
+)
 
 # class WorkerNode(BaseTool):
 #     """Useful for when you need to spawn an autonomous agent instance as a worker to accomplish complex tasks, it can search the internet or spawn child multi-modality models to process and generate images and text or audio and so on """
diff --git a/swarms/swarms.py b/swarms/swarms.py
index 753ba1db..9fc1fad3 100644
--- a/swarms/swarms.py
+++ b/swarms/swarms.py
@@ -3,6 +3,7 @@ from swarms.agents.workers.worker import WorkerNode
 from swarms.agents.boss.boss_agent import BossNode
 # from swarms.agents.workers.omni_worker import OmniWorkerAgent
 
+from swarms.agents.workers.worker import worker_tool
 import logging
 
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -50,7 +51,7 @@ class Swarms:
         todo_chain = LLMChain(llm=llm, prompt=todo_prompt)
         tools = [
             Tool(name="TODO", func=todo_chain.run, description="useful for when you need to come up with todo lists. Input: an objective to create a todo list for. Output: a todo list for that objective. Please be very clear what the objective is!"),
-            worker_node,
+            worker_tool
         ]
         suffix = """Question: {task}\n{agent_scratchpad}"""
         prefix = """You are an Boss in a swarm who performs one task based on the following objective: {objective}. Take into account these previously completed tasks: {context}.\n"""