You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
swarms/examples/demos/ai_research_team/main_example.py

79 lines
1.7 KiB

import os
from dotenv import load_dotenv
from swarm_models import Anthropic, OpenAIChat
from swarms.prompts.ai_research_team import (
PAPER_IMPLEMENTOR_AGENT_PROMPT,
PAPER_SUMMARY_ANALYZER,
)
from swarms.structs import Agent
from swarms.utils.pdf_to_text import pdf_to_text
from swarms import rearrange
# Base llms
load_dotenv()
anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
openai_api_key = os.getenv("OPENAI_API_KEY")
PDF_PATH = "fasterffn.pdf"
# Base llms
llm1 = OpenAIChat(
openai_api_key=openai_api_key,
)
llm2 = Anthropic(
anthropic_api_key=anthropic_api_key,
)
# Agents
paper_summarizer_agent = Agent(
agent_name="paper_summarizer_agent",
llm=llm2,
sop=PAPER_SUMMARY_ANALYZER,
max_loops=1,
autosave=True,
saved_state_path="paper_summarizer.json",
)
paper_implementor_agent = Agent(
agent_name="paper_implementor_agent",
llm=llm1,
sop=PAPER_IMPLEMENTOR_AGENT_PROMPT,
max_loops=1,
autosave=True,
saved_state_path="paper_implementor.json",
code_interpreter=False,
)
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,
)
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)