import streamlit as st

from main import (
    load_dataset,
    build_vector_db,
    generate_response
)


st.title("Local RAG Chatbot for Cat Facts")

DATA_FILE = "cat-facts.txt"


@st.cache_resource
def initialize_rag():

    dataset = load_dataset(DATA_FILE)

    build_vector_db(dataset)

    return True


initialize_rag()

query = st.text_input(
    "Ask a question about cats:",
    placeholder="e.g. What is the average lifespan of a cat?"
)

if st.button("Answer"):

    if query.strip():

        with st.spinner("Retrieving relevant information..."):

            answer, results = generate_response(query)

        st.subheader("Answer")
        st.write(answer)