From 04a5151343abaf92dec6e3ad92074c3a12715675 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Thu, 6 Jul 2023 13:43:32 -0400
Subject: [PATCH] worker class clean up

---
 swarms/agents/workers/worker.py | 38 ++++++++++++++-------------------
 1 file changed, 16 insertions(+), 22 deletions(-)

diff --git a/swarms/agents/workers/worker.py b/swarms/agents/workers/worker.py
index f15e5adf..cb097446 100644
--- a/swarms/agents/workers/worker.py
+++ b/swarms/agents/workers/worker.py
@@ -6,35 +6,30 @@ from langchain.callbacks.manager import (
     AsyncCallbackManagerForToolRun,
     CallbackManagerForToolRun,
 )
-from typing import List, Any
+from typing import List, Any, Dict
 from langchain.memory.chat_message_histories import FileChatMessageHistory
 
 import logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
+class WorkerNodeArgs(BaseModel):
+    prompt: str
+    run_manager: Optional[CallbackManagerForToolRun] = None
+
 @tool("WorkerNode")
-class 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"""
-    # name = "WorkerNode"
-    # 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"
-    # llm: ChatOpenAI  # add this line
-    # tools: List[Tool]
-    # vectorstore: VectorStore
-    llm: Any
-    tools: List[Any]
-    vectorstore: Any
-    agent: Any = None
-
     
-    def __init__(self, llm: Any, tools: List[Any], vectorstore: Any):
-        logging.info("Initializing WorkerNode")
-        super().__init__()
-        self.llm = llm
-        self.tools = tools
-        self.vectorstore = vectorstore
+    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)
+        self.llm = kwargs.get('llm')
+        self.tools = kwargs.get('tools')
+        self.vectorstore = kwargs.get('vectorstore')
         self.agent = None
 
-
     def create_agent(self, ai_name, ai_role, human_in_the_loop, search_kwargs):
         logging.info("Creating agent in WorkerNode")
         self.agent = AutoGPT.from_llm_and_tools(
@@ -48,11 +43,11 @@ class WorkerNode:
         )
         self.agent.chain.verbose = True
 
-    # @tool("perform_task")
     def _run(
-        self, prompt: str, run_manager: Optional[CallbackManagerForToolRun] = None
+        self, tool_input: Dict[str, Any], run_manager: Optional[CallbackManagerForToolRun] = None
     ) -> str:
         """Use the tool."""
+        prompt = self._parse_input(tool_input)
         tree_of_thoughts_prompt = """
         Imagine three different experts are answering this question. All experts will write down each chain of thought of each step of their thinking, then share it with the group. Then all experts will go on to the next step, etc. If any expert realises they're wrong at any point then they leave. The question is...
         """
@@ -68,7 +63,6 @@ class 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 """
 #     name = "WorkerNode"