import os
import streamlit as st
from langchain_groq import ChatGroq

# Retrieve API key from environment variable
api_key = os.getenv('pack1st')
if not api_key:
    st.error("GROQ_API_KEY is not set. Please add it to your Hugging Face Space Secrets.")
    st.stop()

# Initialize ChatGroq model
llm = ChatGroq(model_name="gemma2-9b-it", api_key=api_key)

# Set up Streamlit page
st.set_page_config(page_title="Basic Langchain Model", page_icon=":robot:")
st.header("My  Chatbot")

# Initialize session state for chat history
if "messages" not in st.session_state:
    st.session_state.messages = []  # Stores chat history

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):  # "user" or "assistant"
        st.write(message["content"])

# Create new input box dynamically
if prompt := st.chat_input("Type your message..."):
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

    # Generate AI response
    with st.chat_message("assistant"):
        response = llm.invoke(prompt).content  # Extract only the text response
        st.write(response)

    # Add AI response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})