from functions import *

# set the title
st.sidebar.title(DASHBOARD_TITLE)
info_section = st.empty()

# add an explanation of what is NER and why it is important for medical tasks
st.sidebar.markdown(
    f"""
    Meta Llama 3 8B Instruct is part of a family of large language models (LLMs) optimized for dialogue tasks.

    This project uses Streamlit to create a simple chatbot interface that allows you to chat with the model using the Hugging Face Inference API.
    
    Ask the model marketing-related questions and see how it responds. Have fun!

    Model used: [{MODEL_PATH}]({MODEL_LINK})
    """
)

first_assistant_message = "Hello! I am Marketing expert. What can I help you with today?"

# clear conversation
if st.sidebar.button("Clear conversation"):
    chat_history = [{'role':'assistant', 'content':first_assistant_message}]
    st.session_state['chat_history'] = chat_history
    st.rerun()

# Get the chat history
if "chat_history" not in st.session_state:
    chat_history = [{'role':'assistant', 'content':first_assistant_message}]
    st.session_state['chat_history'] = chat_history
else:
    chat_history = st.session_state['chat_history']

# print the conversation
for message in chat_history:
    with st.chat_message(message['role']):
        st.write(message['content'])

# keep only last 10 messages
shorter_history = [message for message in chat_history[-10:] if 'content' in message]

# include a system prompt to explain the bot what to do
system_prompt = """For this task, you are a Marketer specialized in E-commerce helping a user with marketing-related questions. Provide insights and recommendations based on the user's questions. Don't write more than 3-4 sentences per response."""
shorter_history = [{'role': 'system', 'content': system_prompt}] + shorter_history

# get the input from user
user_input = st.chat_input("Write something...")

if user_input:

    with st.chat_message("user"):
        st.write(user_input)

    # make the request
    with st.spinner("Generating the response..."):

        client = InferenceClient(
            "meta-llama/Meta-Llama-3-8B-Instruct",
            token=HUGGING_FACE_API_KEY,
            )
        
        messages = shorter_history + [{'role': 'user', 'content': user_input}]

        # query the model
        try:
            response = client.chat_completion(
                messages=messages,
                max_tokens = 500,
                stream = False,
                )
            
            # get the response
            message = response.choices[0].message['content']
            
            # append to the history
            chat_history.append({'content':user_input, 'role':'user'})
            chat_history.append(response.choices[0].message) # append the response

        except Exception as e:
            st.error(f"An error occurred: {e}")
            st.stop()

    st.session_state['chat_history'] = chat_history
    st.rerun()