from langchain.schema import SystemMessage
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.agents import OpenAIFunctionsAgent
from langchain.prompts import MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI
from langchain.agents import AgentExecutor
from final_tools import custom_tools
import gradio as gr



define_agent = """
You are Apollo, an AI music-player assistant, designed to provide a personalized and engaging listening experience through thoughtful interaction and intelligent tool usage.

Your Main Responsibilities:

1. **Play Music:** Utilize your specialized toolkit to fulfill music requests.

2. **Mood Monitoring:** Constantly gauge the user's mood and adapt the music accordingly. For example, if the mood shifts from 'Happy' to 'more upbeat,' select 'Energetic' music. 

3. **Track and Artist Memory:** Be prepared to recall tracks and/or artists that the user has previously requested.

4. **Provide Guidance:** If the user appears indecisive or unsure about their selection, proactively offer suggestions based on their previous preferences or popular choices within the desired mood or genre.

5. **Seek Clarification:** If a user's request is ambiguous, don't hesitate to ask for more details.
"""


#global llm # None until create_agent() is called
# making global variable so explain_track() (and future functions that need an llm) can recognize it
llm_state = gr.State()


def create_agent(key):

    system_message = SystemMessage(content=define_agent)
    MEMORY_KEY = "chat_history"
    prompt = OpenAIFunctionsAgent.create_prompt(
        system_message=system_message,
        extra_prompt_messages=[MessagesPlaceholder(variable_name=MEMORY_KEY)]
    )
    memory = ConversationBufferMemory(memory_key=MEMORY_KEY, return_messages=True)

    llm = ChatOpenAI(openai_api_key=key, max_retries=3, temperature=0, model_name="gpt-4")
    llm_state.value = llm
    agent = OpenAIFunctionsAgent(llm=llm_state.value, tools=custom_tools, prompt=prompt)
    agent_executor = AgentExecutor(agent=agent, tools=custom_tools, memory=memory, verbose=True)
    return agent_executor