Former-commit-id: 32388a1f67
pull/47/head
Kye 2 years ago
parent 832a7ae0d4
commit 43e014aac7

@ -41,6 +41,7 @@ class HierarchicalSwarm:
embedding_size: Optional[int] = None, embedding_size: Optional[int] = None,
use_async: Optional[bool] = True, use_async: Optional[bool] = True,
worker_name: Optional[str] = "Swarm Worker AI Assistant", worker_name: Optional[str] = "Swarm Worker AI Assistant",
verbose: Optional[bool] = False,
human_in_the_loop: Optional[bool] = True, 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", 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.max_iterations = max_iterations
self.logging_enabled = logging_enabled self.logging_enabled = logging_enabled
self.verbose = verbose
self.logger = logging.getLogger() self.logger = logging.getLogger()
if not logging_enabled: if not logging_enabled:
self.logger.disabled = True self.logger.disabled = True
@ -143,7 +147,7 @@ class HierarchicalSwarm:
logging.error(f"Failed to initialize worker node: {e}") logging.error(f"Failed to initialize worker node: {e}")
raise 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 Init BossNode
@ -176,7 +180,7 @@ class HierarchicalSwarm:
llm_chain = LLMChain(llm=llm, prompt=prompt) llm_chain = LLMChain(llm=llm, prompt=prompt)
agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=[tool.name for tool in tools]) 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, return BossNode(llm, vectorstore,
agent_executor, max_iterations=self.max_iterations) agent_executor, max_iterations=self.max_iterations)
except Exception as e: except Exception as e:

@ -1,4 +1,5 @@
import logging import logging
from typing import Optional, List, Union
import faiss import faiss
from langchain.agents import Tool from langchain.agents import Tool
@ -25,46 +26,63 @@ ROOT_DIR = "./data/"
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
class WorkerNodeInitializer: 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""" """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): def __init__(self,
if not llm or not tools or not vectorstore: llm: Optional[Union[InMemoryDocstore, ChatOpenAI]] = None,
logging.error("llm, tools, and vectorstore cannot be None.") tools: Optional[List[Tool]] = None,
raise ValueError("llm, tools, and vectorstore cannot be None.") vectorstore: Optional[FAISS] = None,
ai_name: str = "Swarm Worker AI Assistant",
self.llm = llm ai_role: str = "Assistant",
self.tools = tools human_in_the_loop: bool = False,
self.vectorstore = vectorstore search_kwargs: dict = {},
self.agent = None verbose: bool = False,
chat_history_file: str = "chat_history.txt"):
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):
def create_agent(self, ai_name="Swarm Worker AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={}, verbose=False):
logging.info("Creating agent in WorkerNode") logging.info("Creating agent in WorkerNode")
try: try:
self.agent = AutoGPT.from_llm_and_tools( self.agent = AutoGPT.from_llm_and_tools(
ai_name=ai_name, ai_name=self.ai_name,
ai_role=ai_role, ai_role=self.ai_role,
tools=self.tools, tools=self.tools,
llm=self.llm, llm=self.llm,
memory=self.vectorstore.as_retriever(search_kwargs=search_kwargs), memory=self.vectorstore.as_retriever(search_kwargs=self.search_kwargs),
human_in_the_loop=human_in_the_loop, human_in_the_loop=self.human_in_the_loop,
chat_history_memory=FileChatMessageHistory("chat_history.txt"), chat_history_memory=FileChatMessageHistory(self.chat_history_file),
) )
# self.agent.chain.verbose = verbose # self.agent.chain.verbose = verbose
except Exception as e: except Exception as e:
logging.error(f"Error while creating agent: {str(e)}") logging.error(f"Error while creating agent: {str(e)}")
raise e raise e
def add_tool(self, tool: Optional[Tool] = None):
if tool is None:
tool = DuckDuckGoSearchRun()
def add_tool(self, tool: Tool):
if not isinstance(tool, Tool): if not isinstance(tool, Tool):
logging.error("Tool must be an instance of Tool.") logging.error("Tool must be an instance of Tool.")
raise TypeError("Tool must be an instance of Tool.") raise TypeError("Tool must be an instance of Tool.")
self.tools.append(tool) self.tools.append(tool)
def run(self, prompt: str) -> str: def run(self, prompt) -> str:
if not isinstance(prompt, str): if not isinstance(prompt, str):
logging.error("Prompt must be a string.") logging.error("Prompt must be a string.")
raise TypeError("Prompt must be a string.") raise TypeError("Prompt must be a string.")
@ -81,9 +99,6 @@ class WorkerNodeInitializer:
raise e raise e
class WorkerNode: class WorkerNode:
def __init__(self, openai_api_key): def __init__(self, openai_api_key):
if not openai_api_key: if not openai_api_key:

Loading…
Cancel
Save