clean up worker_node

swarms/agents/workers/worker_agent.py

@@ -36,21 +36,18 @@ class WorkerNode:
         self.agent.chain.verbose = True
 
     def add_tool(self, tool: Tool):
-        """adds a new tool to the agents toolset"""
         self.tools.append(tool)
 
-    def run(self, tool_input: Dict[str, Any]) -> str:
-        if not isinstance(tool_input, dict):
-            raise TypeError("tool_input must be a dictionary")
-        if 'prompt' not in tool_input:
-            raise ValueError("tool_input must contain the key 'prompt'")
-        
-        prompt = tool_input['prompt']
-        if prompt is None:
-            raise ValueError("Prompt not found in tool_input")
+    def run(self, prompt: str) -> str:
+        if not isinstance(prompt, str):
+            raise TypeError("Prompt must be a string")
         
+        if not prompt:
+            raise ValueError("Prompt is empty")
+
         self.agent.run([f"{prompt}"])
         return "Task completed by WorkerNode"
+    
 
 worker_tool = Tool(
     name="WorkerNode AI Agent",
@@ -59,7 +56,6 @@ worker_tool = Tool(
 )
 
 
-
 class WorkerNodeInitializer:
     def __init__(self, openai_api_key):
         self.openai_api_key = openai_api_key
@@ -92,9 +88,7 @@ class WorkerNodeInitializer:
         worker_node.create_agent(ai_name="Swarm Worker AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={})
         return worker_node
 
-
-# usage
-def worker_node(api_key):
-    initializer = WorkerNodeInitializer(api_key)
-    worker = initializer.create_worker_node()
-    return worker
+def worker_node(openai_api_key):
+    initializer = WorkerNodeInitializer(openai_api_key)
+    worker_node = initializer.create_worker_node()
+    return worker_node