From 38ac04771122865c5f7792432be0a18e309aeb96 Mon Sep 17 00:00:00 2001
From: Kye Gomez <swarms_wd@Kyes-MacBook-Pro.local>
Date: Tue, 9 Jul 2024 09:37:56 -0700
Subject: [PATCH] [DEMO][FIX]

---
 .../demos/ai_research_team/main_example.py    | 33 +++++++++++++++----
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/playground/demos/ai_research_team/main_example.py b/playground/demos/ai_research_team/main_example.py
index dc6e54ae..d77d560d 100644
--- a/playground/demos/ai_research_team/main_example.py
+++ b/playground/demos/ai_research_team/main_example.py
@@ -9,9 +9,9 @@ from swarms.prompts.ai_research_team import (
 )
 from swarms.structs import Agent
 from swarms.utils.pdf_to_text import pdf_to_text
+from swarms import rearrange
 
 # Base llms
-# Environment variables
 load_dotenv()
 anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
 openai_api_key = os.getenv("OPENAI_API_KEY")
@@ -30,6 +30,7 @@ llm2 = Anthropic(
 
 # Agents
 paper_summarizer_agent = Agent(
+    agent_name = "paper_summarizer_agent",
     llm=llm2,
     sop=PAPER_SUMMARY_ANALYZER,
     max_loops=1,
@@ -38,6 +39,7 @@ paper_summarizer_agent = Agent(
 )
 
 paper_implementor_agent = Agent(
+    agent_name = "paper_implementor_agent",
     llm=llm1,
     sop=PAPER_IMPLEMENTOR_AGENT_PROMPT,
     max_loops=1,
@@ -46,9 +48,28 @@ paper_implementor_agent = Agent(
     code_interpreter=False,
 )
 
-paper = pdf_to_text(PDF_PATH)
-algorithmic_psuedocode_agent = paper_summarizer_agent.run(
-    "Focus on creating the algorithmic pseudocode for the novel"
-    f" method in this paper: {paper}"
+pytorch_pseudocode_agent = Agent(
+    agent_name = "pytorch_pseudocode_agent",
+    llm=llm1,
+    sop=PAPER_IMPLEMENTOR_AGENT_PROMPT,
+    max_loops=1,
+    autosave=True,
+    saved_state_path="pytorch_pseudocode_agent.json",
+    code_interpreter=False,
 )
-pytorch_code = paper_implementor_agent.run(algorithmic_psuedocode_agent)
+
+
+paper = pdf_to_text(PDF_PATH)
+task = f"""
+    Focus on creating the algorithmic pseudocode for the novel 
+    f" method in this paper: {paper}
+"""
+
+
+agents = [paper_summarizer_agent, paper_implementor_agent, pytorch_pseudocode_agent]
+
+flow = "paper_summarizer_agent -> paper_implementor_agent -> pytorch_pseudocode_agent"
+
+swarm = rearrange(agents, flow, task)
+print(swarm)
+