from crewai import Agent, Task, Crew
from crewai_tools import SerperDevTool, ScrapeWebsiteTool, SeleniumScrapingTool
from langchain_community.chat_models import ChatOllama
from langchain_openai import ChatOpenAI
from flask import Flask, jsonify
import os

# export OPENAI_API_BASE=http://the.mk:11434/v1
# export OPENAI_MODEL_NAME=crewai-llama3
# export SERPER_API_KEY=

# class MyAgent(Agent):
#     def __init__(self, role="My Math Professor", goal=None, backstory=None, allow_delegation=False, verbose=True, llm=None):
#         super().__init__(role=role, goal=goal, backstory=backstory, allow_delegation=allow_delegation, verbose=verbose, llm=llm)
#     def to_dict(self):
#         return {
#             'role': self.role,
#             'goal': self.goal,
#             'backstory': self.backstory,
#             'allow_delegation': self.allow_delegation,
#             'verbose': self.verbose
#         }


os.environ["OPENAI_API_KEY"] = "NA"

llm = ChatOpenAI(
    model = "crewai-llama3",
    base_url = "http://chat.the.mk:11434/v1",
    temperature=0.1)

# general_agent = MyAgent(role = "Math Professor",
#                       goal = """Provide the solution to the students that are asking mathematical questions and give them the answer.""",
#                       backstory = """You are an excellent math professor that likes to solve math questions in a way that everyone can understand your solution""",
#                       allow_delegation = False,
#                       verbose = True,
#                       llm = llm)
# task = Task (description="""what is 3 + 5""",
#              agent = general_agent,
#              expected_output="A numerical answer.")

# crew = Crew(
#             agents=[general_agent],
#             tasks=[task],
#             verbose=2
#         )



# Loading Tools
search_tool = SerperDevTool()
scrape_tool = SeleniumScrapingTool()


# Define your agents with roles, goals, tools, and additional attributes
researcher = Agent(
  role='Senior Research Analyst',
  goal='Uncover cutting-edge developments in AI and data science',
  backstory=(
    "You are a Senior Research Analyst at a leading tech think tank."
    "Your expertise lies in identifying emerging trends and technologies in AI and data science."
    "You have a knack for dissecting complex data and presenting actionable insights."
    "Always search the web first and make the determination for the best 4 Links"
    "if you are going to use Read website content tool replace the search positional argument: to 'website_url'"
    # "Don't include single or double quotes in any or the searches or parameters for the tools including website_url or css_element"
    "For any web searches, be sure to scrape the website content from the Link in the search"
  ),
  verbose=True,
  allow_delegation=False,
  tools=[search_tool, scrape_tool],
  max_rpm=100
)

writer = Agent(
  role='Tech Content Strategist',
  goal='Craft compelling content on tech advancements',
  backstory=(
    "You are a renowned Tech Content Strategist, known for your insightful and engaging articles on technology and innovation."
    "With a deep understanding of the tech industry, you transform complex concepts into compelling narratives."
    "Always search the web first and make the determination for the best 4 Links"
    "if you are going to use Read website content tool replace the search positional argument: to 'website_url'"
    # "Don't include single or double quotes in any or the searches or parameters for the tools including website_url or css_element"
    "For any web searches, be sure to scrape the website content from the Link in the search"
  ),
  verbose=True,
  allow_delegation=True,
  tools=[search_tool, scrape_tool],
  cache=False, # Disable cache for this agent
)

# Create tasks for your agents
task1 = Task(
  description=(
    "Conduct a comprehensive analysis of the latest advancements in AI in 2024."
    "Identify key trends, breakthrough technologies, and potential industry impacts."
    "Compile your findings in a detailed report."
    "Make sure to check with a human if the draft is good before finalizing your answer."
  ),
  expected_output='A comprehensive full report on the latest AI advancements in 2024, leave nothing out',
  agent=researcher,
  human_input=True,
)

task2 = Task(
  description=(
    "Using the insights from the researcher's report, develop an engaging blog post that highlights the most significant AI advancements."
    "Your post should be informative yet accessible, catering to a tech-savvy audience."
    "Aim for a narrative that captures the essence of these breakthroughs and their implications for the future."
  ),
  expected_output='A compelling three paragraphs blog post formatted as markdown with headings, subheadings, and a main thesis about the latest AI advancements in 2024',
  agent=writer
)

# Instantiate your crew with a sequential process
crew = Crew(
  agents=[researcher, writer],
  tasks=[task1, task2],
  verbose=4
)

app = Flask(__name__)

def get_agents(cls):
    return [agent.__dict__ for agent in cls.agents]

def get_tasks(cls):
    return [{**task.__dict__, 'expected_output': task.expected_output} for task in cls.tasks]


@app.route('/agents', methods=['GET'])
def get_agents_route():
    return jsonify({'agents': get_agents(crew)})

@app.route('/tasks', methods=['GET'])
def get_tasks_route():
    return jsonify({'tasks': get_tasks(crew)})

@app.route('/kickoff', methods=['GET'])
def get_kickoff_route():
    result = crew.kickoff()
    print(result)
    return result


if __name__ == '__main__':
    app.run(debug=True, port=5001)