hiearchical swarm cleanup

swarms/swarms/swarms.py

@@ -9,23 +9,17 @@ from swarms.workers.worker_node import WorkerNode
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# TODO: Pass in abstract LLM class that can utilize Hf or Anthropic models, Move away from OPENAI
-# TODO: ADD Universal Communication Layer, a ocean vectorstore instance
-# TODO: BE MORE EXPLICIT ON TOOL USE, TASK DECOMPOSITION AND TASK COMPLETETION AND ALLOCATION
-# TODO: Add RLHF Data collection, ask user how the swarm is performing
-# TODO: Create an onboarding process if not settings are preconfigured like `from swarms import Swarm, Swarm()` => then initiate onboarding name your swarm + provide purpose + etc 
-
 
 # ---------- Constants ----------
 ROOT_DIR = "./data/"
 
 
+
 class HierarchicalSwarm:
     def __init__(
         self, 
         openai_api_key: Optional[str] = "", 
         use_vectorstore: Optional[bool] = True, 
-
         use_async: Optional[bool] = True, 
         worker_name: Optional[str] = "Swarm Worker AI Assistant",
         verbose: Optional[bool] = False,
@@ -34,14 +28,13 @@ class HierarchicalSwarm:
         worker_prompt: Optional[str] = None,
         temperature: Optional[float] = 0.5,
         max_iterations: Optional[int] = None,
-        logging_enabled: Optional[bool] = True):
+        logging_enabled: Optional[bool] = True
-            
+    ):
         self.openai_api_key = openai_api_key
         self.use_vectorstore = use_vectorstore
         self.use_async = use_async
         self.worker_name = worker_name
         self.human_in_the_loop = human_in_the_loop
-
         self.boss_prompt = boss_prompt
         self.temperature = temperature
         self.max_iterations = max_iterations
@@ -50,7 +43,6 @@ class HierarchicalSwarm:
 
         self.worker_node = WorkerNode(
             openai_api_key=self.openai_api_key,
-            # worker_name=self.worker_name,
             temperature=self.temperature,
             human_in_the_loop=self.human_in_the_loop,
             verbose=self.verbose
@@ -85,9 +77,8 @@ class HierarchicalSwarm:
 
 def swarm(
     api_key: Optional[str]="", 
-    objective: Optional[str]="", 
+    objective: Optional[str]=""
-    ):
+):
-
     if not api_key or not isinstance(api_key, str):
         logging.error("Invalid OpenAI key")
         raise ValueError("A valid OpenAI API key is required")
@@ -95,7 +86,7 @@ def swarm(
         logging.error("Invalid objective")
         raise ValueError("A valid objective is required")
     try:
-        swarms = HierarchicalSwarm(api_key, use_async=False) #logging_enabled=logging_enabled) # Turn off async
+        swarms = HierarchicalSwarm(api_key, use_async=False)
         result = swarms.run(objective)
         if result is None:
             logging.error("Failed to run swarms")
@@ -104,4 +95,12 @@ def swarm(
         return result
     except Exception as e:
         logging.error(f"An error occured in swarm: {e}")
-        return None
+        return None
+    
+
+
+# TODO: Pass in abstract LLM class that can utilize Hf or Anthropic models, Move away from OPENAI
+# TODO: ADD Universal Communication Layer, a ocean vectorstore instance
+# TODO: BE MORE EXPLICIT ON TOOL USE, TASK DECOMPOSITION AND TASK COMPLETETION AND ALLOCATION
+# TODO: Add RLHF Data collection, ask user how the swarm is performing
+# TODO: Create an onboarding process if not settings are preconfigured like `from swarms import Swarm, Swarm()` => then initiate onboarding name your swarm + provide purpose + etc 

swarms/workers/worker_node.py

@@ -109,30 +109,94 @@ class WorkerNodeInitializer:
             raise e
 
 class WorkerNode:
-    def __init__(self, 
+   def __init__(self, 
-            openai_api_key: str,
+                 openai_api_key: str,
-            temperature: int, 
+                 temperature: Optional[int] = None, 
-            llm: Optional[Union[InMemoryDocstore, ChatOpenAI]] = None, 
+                 llm: Optional[Union[InMemoryDocstore, ChatOpenAI]] = None, 
-            tools: Optional[List[Tool]] = None, 
+                 tools: Optional[List[Tool]] = None,
-            # vectorstore: Optional[FAISS] = None,
+                 embedding_size: Optional[int] = 8192, # Default value set as 8192 from WorkerNodeInitializer
-            embedding_size: Optional[int] = 4026,
+                 worker_name: Optional[str] = "Swarm Worker AI Assistant", 
-            worker_name: Optional[str] = "Swarm Worker AI Assistant", 
+                 worker_role: Optional[str] = "Assistant",
-            worker_role: Optional[str] = "Assistant", 
+                 human_in_the_loop: Optional[bool] = False, 
-            human_in_the_loop: Optional[bool] = False, 
+                 search_kwargs: dict = {},
-            search_kwargs: dict = {}, 
+                 verbose: Optional[bool] = False,
-            verbose: Optional[bool] = False,
+                 chat_history_file: str = "chat_history.txt"):
-            chat_history_file: str = "chat_history.txt"):       
+        
-    
         if not openai_api_key:
             raise ValueError("openai_api_key cannot be None")
-        
+
         self.openai_api_key = openai_api_key
-        self.worker_node_initializer = WorkerNodeInitializer(openai_api_key)
+        self.llm = llm if llm is not None else ChatOpenAI()
-        self.name = worker_name  # Added a name attribute
+        self.tools = tools if tools is not None else [ReadFileTool(), WriteFileTool()]
-        self.description = "A worker node that executes tasks"  # Added a description attribute
+        self.embedding_size = embedding_size
-        self.embedding_size = self.embedding_size
+        self.worker_name = worker_name
+        self.worker_role = worker_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
+
+        # Properties from WorkerNode
+        self.temperature = temperature
+        self.description = "A worker node that executes tasks"
+        self.create_agent()
+      
+    def create_agent(self):
+        logging.info("Creating agent in WorkerNode")
+        try:
+            vectorstore = self.initialize_vectorstore()
+            
+            self.agent = AutoGPT.from_llm_and_tools(
+                ai_name=self.worker_name,
+                ai_role=self.worker_role,
+                tools=self.tools,
+                llm=self.llm,
+                memory=vectorstore,
+                human_in_the_loop=self.human_in_the_loop,
+                chat_history_memory=FileChatMessageHistory(self.chat_history_file),
+            )
+        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()
+        
+        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 initialize_vectorstore(self):
+        try:
+            embeddings_model = OpenAIEmbeddings(openai_api_key=self.openai_api_key)
+            embedding_size = self.embedding_size
+            index = faiss.IndexFlatL2(embedding_size=embedding_size)
+            return FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
+        
+        except Exception as e:
+            logging.error(f"Failed to initialize vector store: {e}")
+            return None
 
+    def run(self, prompt) -> str:
+        if not isinstance(prompt, str):
+            logging.error("Prompt must be a string.")
+            raise TypeError("Prompt must be a string.")
+        
+        if not prompt:
+            logging.error("Prompt is empty.")
+            raise ValueError("Prompt is empty.")
+        
+        try:
+            self.agent.run([f"{prompt}"])
+            return "Task completed by WorkerNode"
+        except Exception as e:
+            logging.error(f"While running the agent: {str(e)}")
+            raise e
+
+    # Functions from WorkerNode
     def initialize_llm(self, llm_class, temperature):
         if not llm_class:
             logging.error("llm_class cannot be none")
@@ -144,14 +208,13 @@ class WorkerNode:
             logging.error(f"Failed to initialize language model: {e}")
             raise
 
-
     def initialize_tools(self, llm_class):
         if not llm_class:
             logging.error("llm_class not cannot be none")
             raise ValueError("llm_class cannot be none")
         try:
             logging.info('Creating WorkerNode')
-            llm = self.initialize_llm(llm_class)
+            llm = self.initialize_llm(llm_class, self.temperature)  # Passed self.temperature
 
             tools = [
                 web_search,
@@ -173,10 +236,10 @@ class WorkerNode:
             raise ValueError("llm_class cannot be None.")
         try:
             worker_tools = self.initialize_tools(llm_class)
-            vectorstore = self.worker_node_initializer.initialize_vectorstore()
+            vectorstore = self.initialize_vectorstore()
-            worker_node = WorkerNodeInitializer(
+            worker_node = UnifiedWorkerNode(
-                openai_api_key=self.openai_api_key, # pass the openai_api_key
+                openai_api_key=self.openai_api_key, 
-                llm=self.initialize_llm(llm_class), 
+                llm=self.initialize_llm(llm_class, self.temperature), 
                 tools=worker_tools, 
                 vectorstore=vectorstore,
                 ai_name=worker_name, 
@@ -190,8 +253,6 @@ class WorkerNode:
             logging.error(f"Failed to create worker node: {e}")
             raise
 
-
-
 def worker_node(openai_api_key):
     if not openai_api_key:
         logging.error("OpenAI API key is not provided")