docs + error handlong + logging

Former-commit-id: a76f10d4677d86a47850a1195324e4a67ad688ad

swarms/hivemind.py

@@ -0,0 +1,2 @@
+# many boss + workers in unison
+#kye gomez jul 13 4:01pm, can scale up the number of swarms working on a probkem with `hivemind(swarms=4, or swarms=auto which will scale the agents depending on the complexity)`

swarms/swarms.py

@@ -7,64 +7,141 @@ import logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
 class Swarms:
-    def __init__(self, openai_api_key):
+    def __init__(self, openai_api_key=""):
+        #openai_api_key: the openai key. Default is empty
+
+        if not openai_api_key:
+            logging.error("OpenAI key is not provided")
+            raise ValueError("OpenAI API key is required")
+        
         self.openai_api_key = openai_api_key
 
     def initialize_llm(self, llm_class, temperature=0.5):
-        # Initialize language model
-        return llm_class(openai_api_key=self.openai_api_key, temperature=temperature)
+        """
+        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 initialize_tools(self, llm_class):
-        llm = self.initialize_llm(llm_class)
-        # Initialize tools
-        web_search = DuckDuckGoSearchRun()
-        tools = [
-            web_search,
-            WriteFileTool(root_dir=ROOT_DIR),
-            ReadFileTool(root_dir=ROOT_DIR),
+        """
+        Init tools
+        
+        Params:
+            llm_class (class): The Language model class. Default is OpenAI
+        """
+        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)),
+
+            ]
+            assert tools is not None, "tools is not initialized"
+            return tools
+        except Exception as e:
+            logging.error(f"Failed to initialize tools: {e}")
+            raise
 
-            process_csv,
-            WebpageQATool(qa_chain=load_qa_with_sources_chain(llm)),
+    def initialize_vectorstore(self):
+        """
+        Init vector store
+        """
+
+        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 initialize_worker_node(self, worker_tools, vectorstore, llm_class=ChatOpenAI, ai_name="Swarm Worker AI Assistant"):
+        """
+        Init WorkerNode
+
+        Params:
+            worker_tools (list): The list of worker tools.
+            vectorstore (object): The vector store object
+            llm_class (class): The Language model class. Default is ChatOpenAI
+            ai_name (str): The AI name. Default is "Swarms worker AI assistant"        
+        """
+
+        try:
+                
+            # Initialize worker node
+            llm = self.initialize_llm(ChatOpenAI)
+            worker_node = WorkerNode(llm=llm, tools=worker_tools, vectorstore=vectorstore)
+            worker_node.create_agent(ai_name=ai_name, ai_role="Assistant", human_in_the_loop=False, search_kwargs={}) # add search kwargs
+
+            worker_node_tool = Tool(name="WorkerNode AI Agent", func=worker_node.run, description="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, max_iterations=5, verbose=True):
+        """
+        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)
+            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. 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. Output: a todo list for that objective. 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)
+            return BossNode(llm, vectorstore, agent_executor, max_iterations=max_iterations)
+        except Exception as e:
+            logging.error(f"Failed to initialize boss node: {e}")
+            raise
 
-        ]
-        assert tools is not None, "tools is not initialized"
-        return tools
 
-    def initialize_vectorstore(self):
-        # Initialize vector store
-        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({}), {})
-
-    def initialize_worker_node(self, worker_tools, vectorstore):
-        # Initialize worker node
-        llm = self.initialize_llm(ChatOpenAI)
-        worker_node = WorkerNode(llm=llm, tools=worker_tools, vectorstore=vectorstore)
-        worker_node.create_agent(ai_name="Swarm Worker AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={})
-        worker_node_tool = Tool(name="WorkerNode AI Agent", func=worker_node.run, description="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
-
-    def initialize_boss_node(self, vectorstore, worker_node):
-        # Initialize boss node
-        llm = self.initialize_llm(OpenAI)
-        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. 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. Output: a todo list for that objective. 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=True)
-        # return BossNode(return BossNode(llm, vectorstore, agent_executor, max_iterations=5)
-        return BossNode(llm, vectorstore, agent_executor, max_iterations=5)
 
 
     def run_swarms(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.initialize_tools(OpenAI)
@@ -81,15 +158,29 @@ class Swarms:
             logging.error(f"An error occurred in run_swarms: {e}")
             raise
 
-
 # usage
-def swarm(api_key, objective):
+def swarm(api_key="", objective=""):
     """
-    import swarm
-    api_key = "APIKEY"
-    objective = "What is the capital of the UK?"
-    result = swarm(api_key, objective)
-    print(result)  # Prints: "The capital of the UK is London."
+    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.
     """
-    swarms = Swarms(api_key)
-    return swarms.run_swarms(objective)
+
+    if not api_key:
+        logging.error("OpenAIkey is not provided")
+        raise ValueError("OpenAI API key is not provided")
+    if not objective:
+        logging.error("Objective is not provided")
+        raise ValueError("Objective is required")
+    try:
+            
+        swarms = Swarms(api_key)
+        return swarms.run_swarms(objective)
+    except Exception as e:
+        logging.error(f"An error occured in swarm: {e}")
+        raise