main
Kye 2 years ago
parent 920697c616
commit a588ff1245

@ -7,61 +7,57 @@ class Swarms:
self.openai_api_key = openai_api_key
def initialize_llm(self):
# Initialize language model
return ChatOpenAI(model_name="gpt-4", temperature=1.0, openai_api_key=self.openai_api_key)
def initialize_tools(self, llm):
# 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)),
# Tool(name='terminal', func=Terminal.execute, description='Operates a terminal'),
# Tool(name='code_writer', func=CodeWriter(), description='Writes code'),
# Tool(name='code_editor', func=CodeEditor(), description='Edits code'),#
]
return tools
def initialize_vectorstore(self):
# Initialize vector store
embeddings_model = OpenAIEmbeddings()
embedding_size = 1536
index = faiss.IndexFlatL2(embedding_size)
return FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
def initialize_worker_node(self, llm, tools, vectorstore):
# Initialize worker node
worker_node = WorkerNode(llm=llm, tools=tools, vectorstore=vectorstore)
worker_node.create_agent(ai_name="AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={})
return worker_node
def initialize_boss_node(self, llm, vectorstore):
todo_prompt = PromptTemplate.from_template("You are a planner who is an expert at coming up with a todo list for a given objective. Come up with a todo list for this objective: {objective}""")
def initialize_boss_node(self, llm, vectorstore, worker_node):
# Initialize boss node
todo_prompt = PromptTemplate.from_template("You are a planner who is an expert at coming up with a todo list for a given objective. Come up with a todo list for this objective: {objective}")
todo_chain = LLMChain(llm=OpenAI(temperature=0), prompt=todo_prompt)
# search = SerpAPIWrapper()
tools = [
# Tool(name="Search", func=search.run, description="useful for when you need to answer questions about current events"),
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!"),
# Tool(name="AUTONOMOUS Worker AGENT", func=self.worker_node, description="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")
self.worker_node,
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=OpenAI(temperature=0), 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(self.openai_api_key, llm, vectorstore, agent_executor, verbose=True, max_iterations=5)
def run_swarms(self, objective):
# Run the swarm with the given objective
llm = self.initialize_llm()
tools = self.initialize_tools(llm)
vectorstore = self.initialize_vectorstore()
worker_node = self.initialize_worker_node(llm, tools, vectorstore)
boss_node = self.initialize_boss_node(llm, vectorstore)
boss_node = self.initialize_boss_node(llm, vectorstore, worker_node)
task = boss_node.create_task(objective)
boss_node.execute_task(task)
worker_node.run_agent(objective)

@ -39,7 +39,7 @@ ROOT_DIR = "./data/"
openai_api_key = os.environ["OPENAI_API_KEY"]
llm = ChatOpenAI(model_name="gpt-4", temperature=1.0, openai_api_key=openai_api_key)
tools = [
worker_tools = [
DuckDuckGoSearchRun(),
WriteFileTool(root_dir=ROOT_DIR),
ReadFileTool(root_dir=ROOT_DIR),

Loading…
Cancel
Save