import streamlit as st
from modules import model_selection, retrieve_query


# Streamlit functions to select the model
model, api_key = model_selection()

# RAG CHATBOT Interface
st.header("Rag Chatbot")
st.divider()
#-----------------------------------------------------
query = st.chat_input(placeholder="Enter Your Message")
if query:
    user_message = st.chat_message("User").write(query)
    response = retrieve_query(
        query=query, 
        df="Data/products.csv",
        column="product_info",
        api_key=api_key,
        model=model
        )
    ai_message = st.chat_message("AI").write(response)