import streamlit as st
from assist.chat import chat as chat_plugin
from assist.heatmap import get_heatmap_figure

st.title("RepoSage Chatbot Demo")

# Tab or expander for heatmap
with st.expander("🔢 View Token Similarity Heatmap"):
    if st.button("Show Heatmap"):
        fig = get_heatmap_figure("tensor.pt")
        st.pyplot(fig)

# Chat UI
question = st.text_input("Ask RepoSage a question:", "")
if st.button("Ask RepoSage"):
    response = chat_plugin(question)
    st.write(response)