From e2d3188ba8fbb1e7a4a986ae3081998693787678 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Sun, 16 Jul 2023 14:16:03 -0400
Subject: [PATCH] clean up initialize

Former-commit-id: 82b5a4103f89ecacc15397ebe4900d20d080d2c9
---
 swarms/agents/boss/BossNode.py | 57 ++++++++++++++++++++++++++--------
 1 file changed, 44 insertions(+), 13 deletions(-)

diff --git a/swarms/agents/boss/BossNode.py b/swarms/agents/boss/BossNode.py
index 92c877fc..5b85b05e 100644
--- a/swarms/agents/boss/BossNode.py
+++ b/swarms/agents/boss/BossNode.py
@@ -1,7 +1,8 @@
-from swarms.tools.agent_tools import *
 from pydantic import ValidationError
 import logging
 
+from swarms.tools.agent_tools import *
+
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # ---------- Boss Node ----------
@@ -45,6 +46,21 @@ class BossNode:
         except Exception as e:
             logging.error(f"Failed to initialize vector store: {e}")
             return None
+        
+    def initialize_llm(self, llm_class, temperature=0.5):
+        """
+        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=temperature)
+        except Exception as e:
+            logging.error(f"Failed to initialize language model: {e}")
+
 
 
     def create_task(self, objective):
@@ -71,21 +87,36 @@ class BossNode:
 
 
 
-def boss_node(objective, api_key=None, llm=None, vectorstore=None, agent_executor=None, max_iterations=10):
-#wrapper function to initialize and use Bossnode with given parameters
-    #api keys can be passed as an argument or set as an env
-    api_key = api_key or os.getenv("API_KEY")
+# from swarms import BossNode, OpenAI, LLMChain, Tool, ZeroShotAgent, AgentExecutor, PromptTemplate
+
+def boss_node(objective, api_key=None, vectorstore=None, worker_node=None, llm_class=OpenAI, max_iterations=5, verbose=False):
+    """
+    Wrapper function to initialize and use BossNode with given parameters.
+    API key can be passed as argument or set as an environment variable.
+    """
+    api_key = api_key or os.getenv('API_KEY')
 
     if not api_key:
-        raise ValueError("API key must be providef either as argument as an env named 'api_key'")
-    
-    if not llm:
-        raise ValueError("Language model must be provided")
-    if not vectorstore:
-        raise ValueError("Vectorstore must be provided")
-    if not agent_executor:
-        raise ValueError('Agent Executor must be provided')
+        raise ValueError("API key must be provided either as argument or as an environment variable named 'API_KEY'.")
+
+    llm = BossNode.initialize_llm(llm_class)  # This function should be defined elsewhere
+
+    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 a 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 a 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
+    ]
+    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=verbose)
+
     boss = BossNode(llm, vectorstore, agent_executor, max_iterations)
     task = boss.create_task(objective)
     boss.execute_task(task)