clean up

Former-commit-id: 32388a1f67

swarms/swarms.py

@@ -41,6 +41,7 @@ class HierarchicalSwarm:
         embedding_size: Optional[int] = None, 
         use_async: Optional[bool] = True, 
         worker_name: Optional[str] = "Swarm Worker AI Assistant",
+        verbose: Optional[bool] = False,
 
         human_in_the_loop: Optional[bool] = True, 
         boss_prompt: Optional[str] = "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",
@@ -64,6 +65,9 @@ class HierarchicalSwarm:
         self.max_iterations = max_iterations
         self.logging_enabled = logging_enabled
 
+        self.verbose = verbose
+        
+
         self.logger = logging.getLogger()
         if not logging_enabled:
             self.logger.disabled = True
@@ -143,7 +147,7 @@ class HierarchicalSwarm:
             logging.error(f"Failed to initialize worker node: {e}")
             raise
 
-    def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI, max_iterations=None, verbose=False):
+    def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI):
         """
         Init BossNode
 
@@ -176,7 +180,7 @@ class HierarchicalSwarm:
             llm_chain = LLMChain(llm=llm, prompt=prompt)
             agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=[tool.name for tool in tools])
 
-            agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=verbose)
+            agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=self.verbose)
             return BossNode(llm, vectorstore, 
             agent_executor, max_iterations=self.max_iterations)
         except Exception as e:

swarms/workers/worker_node.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Optional, List, Union
 
 import faiss
 from langchain.agents import Tool
@@ -25,46 +26,63 @@ ROOT_DIR = "./data/"
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-
 class WorkerNodeInitializer:
     """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"""
 
-    def __init__(self, llm, tools, vectorstore):
-        if not llm or not tools or not vectorstore:
-            logging.error("llm, tools, and vectorstore cannot be None.")
-            raise ValueError("llm, tools, and vectorstore cannot be None.")
-        
-        self.llm = llm
-        self.tools = tools
-        self.vectorstore = vectorstore
-        self.agent = None
+    def __init__(self, 
+                 llm: Optional[Union[InMemoryDocstore, ChatOpenAI]] = None, 
+                 tools: Optional[List[Tool]] = None, 
+                 vectorstore: Optional[FAISS] = None,
+                 ai_name: str = "Swarm Worker AI Assistant", 
+                 ai_role: str = "Assistant", 
+                 human_in_the_loop: bool = False, 
+                 search_kwargs: dict = {}, 
+                 verbose: bool = False,
+                 chat_history_file: str = "chat_history.txt"):
+        
+        self.llm = llm if llm is not None else ChatOpenAI()
+        self.tools = tools if tools is not None else [ReadFileTool(), WriteFileTool()]
+        self.vectorstore = vectorstore if vectorstore is not None else FAISS(faiss.IndexFlatIP(512))
+
+        # Initializing agent in the constructor
+        self.ai_name = ai_name
+        self.ai_role = ai_role
+        self.human_in_the_loop = human_in_the_loop
+        self.search_kwargs = search_kwargs
+        self.verbose = verbose
+        self.chat_history_file = chat_history_file
+
+        self.create_agent()
+
+    def create_agent(self):
         
-    def create_agent(self, ai_name="Swarm Worker AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={}, verbose=False):
         logging.info("Creating agent in WorkerNode")
         try:
             self.agent = AutoGPT.from_llm_and_tools(
-                ai_name=ai_name,
-                ai_role=ai_role,
+                ai_name=self.ai_name,
+                ai_role=self.ai_role,
                 tools=self.tools,
                 llm=self.llm,
-                memory=self.vectorstore.as_retriever(search_kwargs=search_kwargs),
-                human_in_the_loop=human_in_the_loop,
-                chat_history_memory=FileChatMessageHistory("chat_history.txt"),
+                memory=self.vectorstore.as_retriever(search_kwargs=self.search_kwargs),
+                human_in_the_loop=self.human_in_the_loop,
+                chat_history_memory=FileChatMessageHistory(self.chat_history_file),
             )
             # self.agent.chain.verbose = verbose
         except Exception as e:
             logging.error(f"Error while creating agent: {str(e)}")
             raise e
 
+    def add_tool(self, tool: Optional[Tool] = None):
+        if tool is None:
+            tool = DuckDuckGoSearchRun()
         
-    def add_tool(self, tool: Tool):
         if not isinstance(tool, Tool):
             logging.error("Tool must be an instance of Tool.")
             raise TypeError("Tool must be an instance of Tool.")
         
         self.tools.append(tool)
 
-    def run(self, prompt: str) -> str:
+    def run(self, prompt) -> str:
         if not isinstance(prompt, str):
             logging.error("Prompt must be a string.")
             raise TypeError("Prompt must be a string.")
@@ -81,9 +99,6 @@ class WorkerNodeInitializer:
             raise e
 
 
-    
-
-
 class WorkerNode:
     def __init__(self, openai_api_key):
         if not openai_api_key: