swarm class simplification boss node is pre configed and worker node is pre configed

Former-commit-id: cba0d4c759

swarms/boss/boss_node.py

@@ -20,7 +20,7 @@ class BossNodeInitializer:
     The BossNode class is responsible for creating and executing tasks using the BabyAGI model.
     It takes a language model (llm), a vectorstore for memory, an agent_executor for task execution, and a maximum number of iterations for the BabyAGI model.
     """
-    def __init__(self, llm, vectorstore, agent_executor, max_iterations):
+    def __init__(self, llm, vectorstore, agent_executor, max_iterations, human_in_the_loop, embedding_size):
         if not llm or not vectorstore or not agent_executor or not max_iterations:
             logging.error("llm, vectorstore, agent_executor, and max_iterations cannot be None.")
             raise ValueError("llm, vectorstore, agent_executor, and max_iterations cannot be None.")
@@ -28,6 +28,8 @@ class BossNodeInitializer:
         self.vectorstore = vectorstore
         self.agent_executor = agent_executor
         self.max_iterations = max_iterations
+        self.human_in_the_loop = human_in_the_loop
+        self.embedding_size = embedding_size
 
         try:
             self.baby_agi = BabyAGI.from_llm(
@@ -35,7 +37,7 @@ class BossNodeInitializer:
                 vectorstore=self.vectorstore,
                 task_execution_chain=self.agent_executor,
                 max_iterations=self.max_iterations,
-                human_in_the_loop=True
+                human_in_the_loop=self.human_in_the_loop
             )
         except ValidationError as e:
             logging.error(f"Validation Error while initializing BabyAGI: {e}")
@@ -50,7 +52,7 @@ class BossNodeInitializer:
         """
         try:     
             embeddings_model = OpenAIEmbeddings(openai_api_key=self.openai_api_key)
-            embedding_size = 1536
+            embedding_size = self.embedding_size
             index = faiss.IndexFlatL2(embedding_size)
             return FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
         except Exception as e:

swarms/swarms.py

@@ -16,8 +16,8 @@ from langchain.vectorstores import FAISS
 
 
 from swarms.agents.tools.main import WebpageQATool, process_csv
-from swarms.boss.boss_node import BossNodeInitializer as BossNode
+from swarms.boss.boss_node import BossNode
-from swarms.workers.worker_node import WorkerNodeInitializer
+from swarms.workers.worker_node import WorkerNode
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
@@ -36,16 +36,13 @@ class HierarchicalSwarm:
     def __init__(
         self, 
         openai_api_key: Optional[str] = "", 
-
         use_vectorstore: Optional[bool] = True, 
         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",
-
         worker_prompt: Optional[str] = None,
         temperature: Optional[float] = 0.5,
         max_iterations: Optional[int] = None,
@@ -53,144 +50,38 @@ class HierarchicalSwarm:
             
         self.openai_api_key = openai_api_key
         self.use_vectorstore = use_vectorstore
-
         self.use_async = use_async
         self.human_in_the_loop = human_in_the_loop
-
         self.embedding_size = embedding_size
         self.boss_prompt = boss_prompt
         self.worker_prompt = worker_prompt
-
         self.temperature = temperature
         self.max_iterations = max_iterations
         self.logging_enabled = logging_enabled
-
         self.verbose = verbose
 
+        self.worker_node = WorkerNode(
+            openai_api_key=self.openai_api_key,
+            use_vectorstore=self.use_vectorstore,
+            embedding_size=self.embedding_size,
+            worker_name=self.worker_name,
+            worker_prompt=self.worker_prompt,
+            temperature=self.temperature,
+            human_in_the_loop=self.human_in_the_loop,
+            verbose=self.verbose
+        )
+
+        self.boss_node = BossNode(
+            worker_node=self.worker_node,
+            max_iterations=self.max_iterations,
+            human_in_the_loop=self.human_in_the_loop,
+            embedding_size=self.embedding_size
+        )
 
         self.logger = logging.getLogger()
         if not logging_enabled:
             self.logger.disabled = True
 
-
-
-    def initialize_llm(self, llm_class):
-        """
-        Init LLM 
-
-        Params:
-            llm_class(class): The Language model class. Default is OpenAI.
-            temperature (float): The Temperature for the language model. Default is 0.5
-        """
-        try: 
-            # Initialize language model
-            return llm_class(openai_api_key=self.openai_api_key, temperature=self.temperature)
-        except Exception as e:
-            logging.error(f"Failed to initialize language model: {e}")
-
-    def initialize_tools(self, llm_class, extra_tools=None):
-        """
-        Init tools
-        
-        Params:
-            llm_class (class): The Language model class. Default is OpenAI
-
-        extra_tools = [CustomTool()]
-            worker_tools = swarms.initialize_tools(OpenAI, extra_tools)
-        """
-        try:
-            llm = self.initialize_llm(llm_class)
-            # Initialize tools
-            web_search = DuckDuckGoSearchRun()
-            tools = [
-                web_search,
-                WriteFileTool(root_dir=ROOT_DIR),
-                ReadFileTool(root_dir=ROOT_DIR),
-
-                process_csv,
-                WebpageQATool(qa_chain=load_qa_with_sources_chain(llm)),
-            ]
-
-            if extra_tools:
-                tools.extend(extra_tools)
-
-            assert tools is not None, "tools is not initialized"
-            return tools
-
-        except Exception as e:
-            logging.error(f"Failed to initialize tools: {e}")
-            raise
-
-    def initialize_vectorstore(self):
-        """
-        Init vector store
-        """
-        try:     
-            embeddings_model = OpenAIEmbeddings(openai_api_key=self.openai_api_key)
-            embedding_size = self.embedding_size
-            index = faiss.IndexFlatL2(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 initialize_worker_node(self, worker_tools, vectorstore, llm_class=ChatOpenAI):
-        try:    
-            llm = self.initialize_llm(ChatOpenAI)
-            worker_node = WorkerNodeInitializer(llm=llm, tools=worker_tools, vectorstore=vectorstore)
-            worker_node.create_agent(ai_name=self.worker_name, ai_role="Assistant", search_kwargs={}, human_in_the_loop=self.human_in_the_loop)
-            worker_description = self.worker_prompt
-            worker_node_tool = Tool(name="WorkerNode AI Agent", func=worker_node.run, description= worker_description or "Input: an objective with a todo list for that objective. Output: your task completed: Please be very clear what the objective and task instructions are. The Swarm worker agent is Useful for when you need to spawn an autonomous agent instance as a worker to accomplish any complex tasks, it can search the internet or write code or spawn child multi-modality models to process and generate images and text or audio and so on")
-            return worker_node_tool
-        except Exception as e:
-            logging.error(f"Failed to initialize worker node: {e}")
-            raise
-
-    def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI):
-        """
-        Init BossNode
-
-        Params:
-            vectorstore (object): the vector store object.
-            worker_node (object): the worker node object
-            llm_class (class): the language model class. Default is OpenAI
-            max_iterations(int): The number of max iterations. Default is 5
-            verbose(bool): Debug mode. Default is False
-        
-        """
-        try:
-
-            # Initialize boss node
-            llm = self.initialize_llm(llm_class)
-            
-            # prompt = self.boss_prompt
-            todo_prompt = PromptTemplate.from_template(self.boss_prompt)
-            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 your objective. Note create a todo list then assign a ranking from 0.0 to 1.0 to each task, then sort the tasks based on the tasks most likely to achieve the objective. The Output: a todo list for that objective with rankings for each step from 0.1 Please be very clear what the objective is!"),
-                worker_node,
-            ]
-
-            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 """
-            
-            prompt = ZeroShotAgent.create_prompt(tools, prefix=prefix, suffix=suffix, input_variables=["objective", "task", "context", "agent_scratchpad"],)
-            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=self.verbose)
-            return BossNode(llm, vectorstore, 
-            agent_executor, max_iterations=self.max_iterations)
-        except Exception as e:
-            logging.error(f"Failed to initialize boss node: {e}")
-            raise
-
-
-
-
     def run(self, objective):
         """
         Run the swarm with the given objective
@@ -199,24 +90,13 @@ class HierarchicalSwarm:
             objective(str): The task
         """
         try:
-            # Run the swarm with the given objective
+            task = self.boss_node.create_task(objective)
-            worker_tools = self.initialize_tools(OpenAI)
-            assert worker_tools is not None, "worker_tools is not initialized"
-
-            vectorstore = self.initialize_vectorstore() if self.use_vectorstore else None
-            assert vectorstore is not None, "vectorstore is not initialized"
-
-            worker_node = self.initialize_worker_node(worker_tools, vectorstore)
-
-            boss_node = self.initialize_boss_node(vectorstore, worker_node)
-
-            task = boss_node.create_task(objective)
             logging.info(f"Running task: {task}")
             if self.use_async:
                 loop = asyncio.get_event_loop()
-                result = loop.run_until_complete(boss_node.run(task))
+                result = loop.run_until_complete(self.boss_node.run(task))
             else:
-                result = boss_node.run(task)
+                result = self.boss_node.run(task)
             logging.info(f"Completed tasks: {task}")
             return result
         except Exception as e:
@@ -257,3 +137,161 @@ def swarm(
         logging.error(f"An error occured in swarm: {e}")
         return None
 
+
+
+
+# class HierarchicalSwarm:
+#     def __init__(
+#         self, 
+#         openai_api_key: Optional[str] = "", 
+
+#         use_vectorstore: Optional[bool] = True, 
+#         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",
+
+#         worker_prompt: Optional[str] = None,
+#         temperature: Optional[float] = 0.5,
+#         max_iterations: Optional[int] = None,
+#         logging_enabled: Optional[bool] = True):
+            
+#         self.openai_api_key = openai_api_key
+#         self.use_vectorstore = use_vectorstore
+
+#         self.use_async = use_async
+#         self.human_in_the_loop = human_in_the_loop
+
+#         self.embedding_size = embedding_size
+#         self.boss_prompt = boss_prompt
+#         self.worker_prompt = worker_prompt
+
+#         self.temperature = temperature
+#         self.max_iterations = max_iterations
+#         self.logging_enabled = logging_enabled
+
+#         self.verbose = verbose
+
+#         self.worker_node = WorkerNode(openai_api_key)
+
+#         self.logger = logging.getLogger()
+#         if not logging_enabled:
+#             self.logger.disabled = True
+
+#     def initialize_worker_node(self, worker_tools, vectorstore, llm_class=ChatOpenAI):
+#         try:    
+#             worker_node = self.worker_node.create_worker_node(llm_class=llm_class, ai_name=self.worker_name, ai_role="Assistant", human_in_the_loop=self.human_in_the_loop, search_kwargs={}, verbose=self.verbose)
+#             worker_description = self.worker_prompt
+#             worker_node_tool = Tool(name="WorkerNode AI Agent", func=worker_node.run, description= worker_description or "Input: an objective with a todo list for that objective. Output: your task completed: Please be very clear what the objective and task instructions are. The Swarm worker agent is Useful for when you need to spawn an autonomous agent instance as a worker to accomplish any complex tasks, it can search the internet or write code or spawn child multi-modality models to process and generate images and text or audio and so on")
+#             return worker_node_tool
+#         except Exception as e:
+#             logging.error(f"Failed to initialize worker node: {e}")
+#             raise
+
+#     def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI):
+#         """
+#         Init BossNode
+
+#         Params:
+#             vectorstore (object): the vector store object.
+#             worker_node (object): the worker node object
+#             llm_class (class): the language model class. Default is OpenAI
+#             max_iterations(int): The number of max iterations. Default is 5
+#             verbose(bool): Debug mode. Default is False
+        
+#         """
+#         try:
+
+#             # Initialize boss node
+#             llm = self.worker_node.initialize_llm(llm_class)
+            
+#             # prompt = self.boss_prompt
+#             todo_prompt = PromptTemplate.from_template(self.boss_prompt)
+#             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 your objective. Note create a todo list then assign a ranking from 0.0 to 1.0 to each task, then sort the tasks based on the tasks most likely to achieve the objective. The Output: a todo list for that objective with rankings for each step from 0.1 Please be very clear what the objective is!"),
+#                 worker_node,
+#             ]
+
+#             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 """
+            
+#             prompt = ZeroShotAgent.create_prompt(tools, prefix=prefix, suffix=suffix, input_variables=["objective", "task", "context", "agent_scratchpad"],)
+#             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=self.verbose)
+#             return BossNode(llm, vectorstore, agent_executor, self.max_iterations)
+#         except Exception as e:
+#             logging.error(f"Failed to initialize boss node: {e}")
+#             raise
+
+#     def run(self, objective):
+#         """
+#         Run the swarm with the given objective
+
+#         Params:
+#             objective(str): The task
+#         """
+#         try:
+#             # Run the swarm with the given objective
+#             worker_tools = self.worker_node.initialize_tools(OpenAI)
+#             assert worker_tools is not None, "worker_tools is not initialized"
+
+#             vectorstore = self.worker_node.initialize_vectorstore() if self.use_vectorstore else None
+#             assert vectorstore is not None, "vectorstore is not initialized"
+
+#             worker_node = self.initialize_worker_node(worker_tools, vectorstore)
+
+#             boss_node = self.initialize_boss_node(vectorstore, worker_node)
+
+#             task = boss_node.create_task(objective)
+#             logging.info(f"Running task: {task}")
+#             if self.use_async:
+#                 loop = asyncio.get_event_loop()
+#                 result = loop.run_until_complete(boss_node.run(task))
+#             else:
+#                 result = boss_node.run(task)
+#             logging.info(f"Completed tasks: {task}")
+#             return result
+#         except Exception as e:
+#             logging.error(f"An error occurred in run: {e}")
+#             return None
+        
+# # usage-# usage-
+# def swarm(
+#     api_key: Optional[str]="", 
+#     objective: Optional[str]="", 
+#     ):
+#     """
+#     Run the swarm with the given API key and objective.
+
+#     Parameters:
+#     api_key (str): The OpenAI API key. Default is an empty string.
+#     objective (str): The objective. Default is an empty string.
+
+#     Returns:
+#     The result of the swarm.
+#     """
+
+#     if not api_key or not isinstance(api_key, str):
+#         logging.error("Invalid OpenAI key")
+#         raise ValueError("A valid OpenAI API key is required")
+#     if not objective or not isinstance(objective, str):
+#         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
+#         result = swarms.run(objective)
+#         if result is None:
+#             logging.error("Failed to run swarms")
+#         else:
+#             logging.info(f"Successfully ran swarms with results: {result}")
+#         return result
+#     except Exception as e:
+#         logging.error(f"An error occured in swarm: {e}")
+#         return None

swarms/workers/worker_node.py

@@ -34,8 +34,9 @@ class WorkerNodeInitializer:
                  llm: Optional[Union[InMemoryDocstore, ChatOpenAI]] = None, 
                  tools: Optional[List[Tool]] = None, 
                  vectorstore: Optional[FAISS] = None,
-                 ai_name: str = "Swarm Worker AI Assistant", 
+                 embedding_size: Optional[int] = 1926,
-                 ai_role: str = "Assistant", 
+                 worker_name: str = "Swarm Worker AI Assistant", 
+                 worker_role: str = "Assistant", 
                  human_in_the_loop: bool = False, 
                  search_kwargs: dict = {}, 
                  verbose: bool = False,
@@ -46,8 +47,8 @@ class WorkerNodeInitializer:
         self.vectorstore = vectorstore #if vectorstore is not None else #FAISS(faiss.IndexFlatIP(512))
 
         # Initializing agent in the constructor
-        self.ai_name = ai_name
+        self.worker_name = worker_name
-        self.ai_role = ai_role
+        self.worker_role = worker_role
         self.human_in_the_loop = human_in_the_loop
         self.search_kwargs = search_kwargs
         self.verbose = verbose
@@ -62,8 +63,8 @@ class WorkerNodeInitializer:
             if self.vectorstore is None:
                 raise ValueError("Vectorstore is not initialized in WorkerNodeInitializer")
             self.agent = AutoGPT.from_llm_and_tools(
-                ai_name=self.ai_name,
+                worker_name=self.worker_name,
-                ai_role=self.ai_role,
+                worker_role=self.worker_role,
                 tools=self.tools,
                 llm=self.llm,
                 memory=self.vectorstore.as_retriever(search_kwargs=self.search_kwargs),
@@ -101,7 +102,6 @@ class WorkerNodeInitializer:
             logging.error(f"While running the agent: {str(e)}")
             raise e
 
-
 class WorkerNode:
     def __init__(self, openai_api_key):
         if not openai_api_key:
@@ -109,6 +109,7 @@ class WorkerNode:
             raise ValueError("openai_api_key cannot be None")
         
         self.openai_api_key = openai_api_key
+        self.worker_node_initializer = WorkerNodeInitializer(openai_api_key)
 
     def initialize_llm(self, llm_class, temperature=0.5):
         if not llm_class:
@@ -126,7 +127,6 @@ class WorkerNode:
             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)
             web_search = DuckDuckGoSearchRun()
@@ -145,26 +145,24 @@ class WorkerNode:
         except Exception as e:
             logging.error(f"Failed to initialize tools: {e}")
 
-    def initialize_vectorstore(self):
-        try:
-                
-            embeddings_model = OpenAIEmbeddings(openai_api_key=self.openai_api_key)
-            embedding_size = 1536
-            index = faiss.IndexFlatL2(embedding_size)
-            return FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
-        except Exception as e:
-            logging.error(f"Failed to initialize vector store: {e}")
-            raise
-
     def create_worker_node(self, llm_class=ChatOpenAI, ai_name="Swarm Worker AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={}, verbose=False):
         if not llm_class:
             logging.error("llm_class cannot be None.")
             raise ValueError("llm_class cannot be None.")
         try:
             worker_tools = self.initialize_tools(llm_class)
-            vectorstore = self.initialize_vectorstore()
+            vectorstore = self.worker_node_initializer.initialize_vectorstore()
-            worker_node = WorkerNode(llm=self.initialize_llm(llm_class), tools=worker_tools, vectorstore=vectorstore)
+            worker_node = WorkerNodeInitializer(
-            worker_node.create_agent(ai_name=ai_name, ai_role=ai_role, human_in_the_loop=human_in_the_loop, search_kwargs=search_kwargs, verbose=verbose)
+                openai_api_key=self.openai_api_key,
+                llm=self.initialize_llm(llm_class), 
+                tools=worker_tools, 
+                vectorstore=vectorstore,
+                ai_name=ai_name, 
+                ai_role=ai_role, 
+                human_in_the_loop=human_in_the_loop, 
+                search_kwargs=search_kwargs, 
+                verbose=verbose
+            )
             return worker_node
         except Exception as e:
             logging.error(f"Failed to create worker node: {e}")
@@ -176,12 +174,8 @@ def worker_node(openai_api_key):
         raise ValueError("OpenAI API key is required")
     
     try:
-
+        worker_node = WorkerNode(openai_api_key)
-        initializer = WorkerNodeInitializer(openai_api_key)
+        return worker_node.create_worker_node()
-        worker_node = initializer.create_worker_node()
-        return worker_node
     except Exception as e:
         logging.error(f"An error occured in worker_node: {e}")
         raise
-
-