From 2de1853609254735c6acdbac87e8d2ef1832001b Mon Sep 17 00:00:00 2001 From: Kye Date: Tue, 25 Jul 2023 17:54:19 -0400 Subject: [PATCH] clean up --- swarms/swarms.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/swarms/swarms.py b/swarms/swarms.py index 710d8f73..708e429b 100644 --- a/swarms/swarms.py +++ b/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,