clean up

swarms/swarms.py

@@ -21,7 +21,9 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 # TODO: Add RLHF Data collection, ask user how the swarm is performing
 
 class HierarchicalSwarm:
-    def __init__(self, model_id: str = None, openai_api_key="", use_vectorstore=True, embedding_size: int = None, use_async=True, human_in_the_loop=True, model_type: str = None):
+    def __init__(self, model_id: str = None, openai_api_key="", 
+                 use_vectorstore=True, embedding_size: int = None, use_async=True, 
+    human_in_the_loop=True, model_type: str = None, boss_prompt: str = None):
         #openai_api_key: the openai key. Default is empty
         if not model_id:
             logging.error("Model ID is not provided")
@@ -33,8 +35,13 @@ class HierarchicalSwarm:
         self.model_id = model_id        
         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
+
 
     def initialize_llm(self, llm_class, temperature=0.5):
         """
@@ -79,8 +86,6 @@ class HierarchicalSwarm:
             if extra_tools:
                 tools.extend(extra_tools)
 
-
-
             assert tools is not None, "tools is not initialized"
             return tools
 
@@ -144,11 +149,11 @@ class HierarchicalSwarm:
 
             # Initialize boss node
             llm = self.initialize_llm(llm_class)
-            todo_prompt = PromptTemplate.from_template("You are a boss planer in a swarm who is an expert at coming up with a todo list for a given objective and then creating an worker to help you accomplish your task. Rate every task on the importance of it's probability to complete the main objective on a scale from 0 to 1, an integer. Come up with a todo list for this objective: {objective} and then spawn a worker agent to complete the task for you. Always spawn an worker agent after creating a plan and pass the objective and plan to the worker agent.")
+            
+            # prompt = self.boss_prompt
+            todo_prompt = PromptTemplate.from_template({self.boss_prompt} or "You are a boss planer in a swarm who is an expert at coming up with a todo list for a given objective and then creating an worker to help you accomplish your task. Rate every task on the importance of it's probability to complete the main objective on a scale from 0 to 1, an integer. Come up with a todo list for this objective: {objective} and then spawn a worker agent to complete the task for you. Always spawn an worker agent after creating a plan and pass the objective and plan to the worker agent.")
             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,