import gradio as gr
import os
from crewai import Agent, Task, Crew, Process
from langchain_groq import ChatGroq
from langchain.schema import HumanMessage

# Initialize the GROQ language model
groq_llm = ChatGroq(
    groq_api_key=os.environ["GROQ_API_KEY"],
    model_name="mixtral-8x7b-32768"
)

# Create agents
def create_agent(role, goal, backstory):
    return Agent(
        role=role,
        goal=goal,
        backstory=backstory,
        verbose=True,
        allow_delegation=False,
        llm=groq_llm
    )

researcher = create_agent(
    'Senior Researcher',
    'Conduct thorough research on given topics',
    'You are an experienced researcher with a keen eye for detail and the ability to find relevant information quickly.'
)

writer = create_agent(
    'Content Writer',
    'Create engaging and informative content based on research',
    'You are a skilled writer capable of turning complex information into easily understandable and engaging content.'
)

editor = create_agent(
    'Editor',
    'Refine and improve the written content',
    'You are a meticulous editor with a strong command of language and an eye for clarity and coherence.'
)

def create_crew(query):
    # Create tasks
    research_task = Task(
        description=f"Research the following topic thoroughly: {query}",
        agent=researcher
    )

    writing_task = Task(
        description="Write an informative article based on the research conducted",
        agent=writer
    )

    editing_task = Task(
        description="Review and refine the written article, ensuring clarity, coherence, and engagement",
        agent=editor
    )

    # Create the crew
    crew = Crew(
        agents=[researcher, writer, editor],
        tasks=[research_task, writing_task, editing_task],
        verbose=2,
        process=Process.sequential
    )

    return crew

def process_query(query, max_tokens=500):
    crew = create_crew(query)
    result = crew.kickoff()
    return result

# Create the Gradio interface
iface = gr.Interface(
    fn=process_query,
    inputs=[
        gr.Textbox(lines=5, label="Enter your query"),
        gr.Slider(minimum=100, maximum=1000, value=500, step=50, label="Max Tokens")
    ],
    outputs=gr.Textbox(lines=10, label="AI Agent Response"),
    title="CrewAI-powered Chatbot using GROQ and LangChain",
    description="Enter a query and get a researched, written, and edited response using CrewAI, GROQ API, and LangChain."
)

iface.launch()