import os
import streamlit as st
from groq import Groq
import streamlit as st
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
import time


# Sidebar for selecting the model and entering the API key
st.sidebar.title("Model Selection")
model_name = st.sidebar.selectbox("Select a model", ["llama3-70b-8192"])
api_key = st.sidebar.text_input("Enter your Groq API key", type="password")
st.sidebar.markdown("Get the API key from here: [https://console.groq.com/keys](https://console.groq.com/keys)")

# Initialize the Groq client
client = Groq(api_key=api_key)

if api_key:
    client = Groq(api_key=api_key)

# Chat interface
st.title("Chatbot")

msgs = StreamlitChatMessageHistory(key="special_app_key")

for msg in msgs.messages:
    st.chat_message(msg.type).write(msg.content)

if prompt := st.chat_input():
    start_time = time.time()
    st.chat_message("human").write(prompt)
    msgs.add_user_message(prompt)

    with st.spinner("Waiting for response..."):
        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model=model_name,
        )
        res = chat_completion.choices[0].message.content
    
    if res:
        st.chat_message("ai").write(res)
        end_time = time.time()
        print(f"Total time {end_time-start_time}")
        msgs.add_ai_message(res)
    else:
        st.error("No valid response received from the AI.")