From 32388a1f679b447199a8d020e5f2ab30c4c9c18c Mon Sep 17 00:00:00 2001 From: Kye Date: Sat, 29 Jul 2023 11:47:48 -0400 Subject: [PATCH] clean up --- swarms/swarms.py | 8 +++-- swarms/workers/worker_node.py | 57 ++++++++++++++++++++++------------- 2 files changed, 42 insertions(+), 23 deletions(-) diff --git a/swarms/swarms.py b/swarms/swarms.py index 11b37fc2..5392b9cd 100644 --- a/swarms/swarms.py +++ b/swarms/swarms.py @@ -41,6 +41,7 @@ class HierarchicalSwarm: 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", @@ -64,6 +65,9 @@ class HierarchicalSwarm: self.max_iterations = max_iterations self.logging_enabled = logging_enabled + self.verbose = verbose + + self.logger = logging.getLogger() if not logging_enabled: self.logger.disabled = True @@ -143,7 +147,7 @@ class HierarchicalSwarm: logging.error(f"Failed to initialize worker node: {e}") raise - def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI, max_iterations=None, verbose=False): + def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI): """ Init BossNode @@ -176,7 +180,7 @@ class HierarchicalSwarm: 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) + 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: diff --git a/swarms/workers/worker_node.py b/swarms/workers/worker_node.py index b07bce72..b306ea49 100644 --- a/swarms/workers/worker_node.py +++ b/swarms/workers/worker_node.py @@ -1,4 +1,5 @@ import logging +from typing import Optional, List, Union import faiss from langchain.agents import Tool @@ -25,46 +26,63 @@ ROOT_DIR = "./data/" logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') - class WorkerNodeInitializer: """Useful for when you need to spawn an autonomous agent instance as a worker to accomplish complex tasks, it can search the internet or spawn child multi-modality models to process and generate images and text or audio and so on""" - def __init__(self, llm, tools, vectorstore): - if not llm or not tools or not vectorstore: - logging.error("llm, tools, and vectorstore cannot be None.") - raise ValueError("llm, tools, and vectorstore cannot be None.") + def __init__(self, + llm: Optional[Union[InMemoryDocstore, ChatOpenAI]] = None, + tools: Optional[List[Tool]] = None, + vectorstore: Optional[FAISS] = None, + ai_name: str = "Swarm Worker AI Assistant", + ai_role: str = "Assistant", + human_in_the_loop: bool = False, + search_kwargs: dict = {}, + verbose: bool = False, + chat_history_file: str = "chat_history.txt"): - self.llm = llm - self.tools = tools - self.vectorstore = vectorstore - self.agent = None + self.llm = llm if llm is not None else ChatOpenAI() + self.tools = tools if tools is not None else [ReadFileTool(), WriteFileTool()] + self.vectorstore = vectorstore if vectorstore is not None else FAISS(faiss.IndexFlatIP(512)) + + # Initializing agent in the constructor + self.ai_name = ai_name + self.ai_role = ai_role + self.human_in_the_loop = human_in_the_loop + self.search_kwargs = search_kwargs + self.verbose = verbose + self.chat_history_file = chat_history_file + + self.create_agent() - def create_agent(self, ai_name="Swarm Worker AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={}, verbose=False): + def create_agent(self): + logging.info("Creating agent in WorkerNode") try: self.agent = AutoGPT.from_llm_and_tools( - ai_name=ai_name, - ai_role=ai_role, + ai_name=self.ai_name, + ai_role=self.ai_role, tools=self.tools, llm=self.llm, - memory=self.vectorstore.as_retriever(search_kwargs=search_kwargs), - human_in_the_loop=human_in_the_loop, - chat_history_memory=FileChatMessageHistory("chat_history.txt"), + memory=self.vectorstore.as_retriever(search_kwargs=self.search_kwargs), + human_in_the_loop=self.human_in_the_loop, + chat_history_memory=FileChatMessageHistory(self.chat_history_file), ) # self.agent.chain.verbose = verbose except Exception as e: logging.error(f"Error while creating agent: {str(e)}") raise e - - def add_tool(self, tool: Tool): + def add_tool(self, tool: Optional[Tool] = None): + if tool is None: + tool = DuckDuckGoSearchRun() + if not isinstance(tool, Tool): logging.error("Tool must be an instance of Tool.") raise TypeError("Tool must be an instance of Tool.") self.tools.append(tool) - def run(self, prompt: str) -> str: + def run(self, prompt) -> str: if not isinstance(prompt, str): logging.error("Prompt must be a string.") raise TypeError("Prompt must be a string.") @@ -80,9 +98,6 @@ class WorkerNodeInitializer: logging.error(f"While running the agent: {str(e)}") raise e - - - class WorkerNode: def __init__(self, openai_api_key):