agents.py

import os
from textwrap import dedent

from crewai import Agent
from crewai_tools import SerperDevTool
from custom_tools import ScrapeWebsiteTool
from custom_tools import SunoTool

from callback_handler import CustomStreamlitCallbackHandler


class MelodyAgents:
    
    def __init__(self, url: str, genre: str, llm):
        self.llm = llm
        self.url = url
        self.genre = genre
        
        # Define some useful tools
        self.search_tool = SerperDevTool()
        self.scrape_website_tool = ScrapeWebsiteTool(limit=1000)
        self.suno_tool = SunoTool(url=self.url, genre=self.genre)
        
        
    def web_researcher_agent(self):
        return Agent(
            role="Web Researcher",
            goal="Conducts a web search on a topic, generating a detailed report on the matter",
            tools = [self.search_tool, self.scrape_website_tool],
            backstory=dedent(
                "An expert in conducting web researchs about any topic"
            ),
            verbose=True,
            allow_delegation=False,
            callbacks=[CustomStreamlitCallbackHandler(color="green")],
            llm=self.llm,
            max_rpm=4000,       
        )
    
    def lyrics_creator_agent(self):
        return Agent(
            role="Lyrics Creator",
            goal=dedent("""Create the most amazing lyrics about a topic
                        adapting the writing style to the music genre."""),
            backstory="A creative lyricist who excels at creating high quality lyrics",
            verbose=True,
            allow_delegation=False,
            callbacks=[CustomStreamlitCallbackHandler(color="green")],
            llm=self.llm
        )

    def song_generator_agent(self):
        return Agent(
            role="Song Generator",
            goal=dedent("""Create a song using the Suno AI Tool. You must 
                        take the lyrics generated by the previous agent and
                        pass this lyrics to the Tool"""
            ),
            backstory="You are a song creator that loves creating songs using Suno AI",
            tools=[self.suno_tool],
            verbose=True,
            allow_delegation=False,
            callbacks=[CustomStreamlitCallbackHandler(color="green")],
            llm=self.llm,
        )