clean up

swarms/swarms.py

@@ -4,26 +4,14 @@ from swarms.agents.boss.boss_agent import BossNode
 # from swarms.agents.workers.omni_worker import OmniWorkerAgent
 
 
+
 class Swarms:
-    def __init__(self,
-                  openai_api_key,
-                #   omni_api_key=None, 
-                #   omni_api_endpoint=None,
-                #   omni_api_type=None
-                  ):
+    def __init__(self, openai_api_key):
         self.openai_api_key = openai_api_key
-        # self.omni_api_key = omni_api_key
-        # self.omni_api_endpoint = omni_api_endpoint
-        # self.omni_api_key = omni_api_type
-
-        # if omni_api_key and omni_api_endpoint and omni_api_type:
-        #     self.omni_worker_agent = OmniWorkerAgent(omni_api_key, omni_api_endpoint, omni_api_type)
-        # else:
-        #     self.omni_worker_agent = None
 
-    def initialize_llm(self):
+    def initialize_llm(self, llm_class, temperature=0):
         # Initialize language model
-        return ChatOpenAI(model_name="gpt-4", temperature=1.0, openai_api_key=self.openai_api_key)
+        return llm_class(temperature=temperature, openai_api_key=self.openai_api_key)
 
     def initialize_tools(self, llm):
         # Initialize tools
@@ -35,8 +23,6 @@ class Swarms:
             process_csv,
             WebpageQATool(qa_chain=load_qa_with_sources_chain(llm)),
         ]
-        # if self.omni_worker_agent:
-        #     tools.append(self.omni_worker_agent.chat) #add omniworker agent class
         return tools
 
     def initialize_vectorstore(self):
@@ -46,16 +32,18 @@ class Swarms:
         index = faiss.IndexFlatL2(embedding_size)
         return FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
 
-    def initialize_worker_node(self, llm, worker_tools, vectorstore):
+    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="AI Assistant", ai_role="Assistant", human_in_the_loop=False, search_kwargs={})
         return worker_node
 
-    def initialize_boss_node(self, llm, vectorstore, worker_node):
+    def initialize_boss_node(self, vectorstore, worker_node):
         # Initialize boss node
+        llm = self.initialize_llm(OpenAI)
         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)
+        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,
@@ -63,18 +51,17 @@ class Swarms:
         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)
+        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(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()
-        worker_tools = self.initialize_tools(llm)
+        worker_tools = self.initialize_tools(ChatOpenAI)
         vectorstore = self.initialize_vectorstore()
-        worker_node = self.initialize_worker_node(llm, worker_tools, vectorstore)
-        boss_node = self.initialize_boss_node(llm, vectorstore, worker_node)
+        worker_node = self.initialize_worker_node(worker_tools, vectorstore)
+        boss_node = self.initialize_boss_node(vectorstore, worker_node)
         task = boss_node.create_task(objective)
         boss_node.execute_task(task)
         worker_node.run_agent(objective)
@@ -82,6 +69,124 @@ class Swarms:
 
 
 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+#omni agent  ===> working
+# class Swarms:
+#     def __init__(self,
+#                   openai_api_key,
+#                 #   omni_api_key=None, 
+#                 #   omni_api_endpoint=None,
+#                 #   omni_api_type=None
+#                   ):
+#         self.openai_api_key = openai_api_key
+#         # self.omni_api_key = omni_api_key
+#         # self.omni_api_endpoint = omni_api_endpoint
+#         # self.omni_api_key = omni_api_type
+
+#         # if omni_api_key and omni_api_endpoint and omni_api_type:
+#         #     self.omni_worker_agent = OmniWorkerAgent(omni_api_key, omni_api_endpoint, omni_api_type)
+#         # else:
+#         #     self.omni_worker_agent = None
+
+#     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)),
+#         ]
+#         # if self.omni_worker_agent:
+#         #     tools.append(self.omni_worker_agent.chat) #add omniworker agent class
+#         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, worker_tools, vectorstore):
+#         # Initialize worker node
+#         worker_node = WorkerNode(llm=llm, tools=worker_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, 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)
+#         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=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()
+#         worker_tools = self.initialize_tools(llm)
+#         vectorstore = self.initialize_vectorstore()
+#         worker_node = self.initialize_worker_node(llm, worker_tools, 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)
+
+
+
+
 # usage
 def swarm(api_key, objective):
     swarms = Swarms(api_key)