import weave

import streamlit as st
from rag.rag import SimpleRAGPipeline

WANDB_PROJECT = "paper_reader"

weave.init(f"{WANDB_PROJECT}")

st.set_page_config(page_title="Chat with the Llama 3 paper!", page_icon="🦙", layout="centered", initial_sidebar_state="auto", menu_items=None)
st.title("Chat with the Llama 3 paper 💬🦙")

@st.cache_resource(show_spinner=False)
def load_rag_pipeline():
    rag_pipeline = SimpleRAGPipeline()
    rag_pipeline.build_query_engine()

    return rag_pipeline

if "rag_pipeline" not in st.session_state.keys():
    st.session_state.rag_pipeline = load_rag_pipeline()

rag_pipeline = st.session_state["rag_pipeline"]

# openai_api_key = st.sidebar.text_input('OpenAI API Key', type='password')

def generate_response(query):
    response = rag_pipeline.predict(query)
    st.write_stream(response.response_gen)

with st.form('my_form'):
    query = st.text_area('Ask your question about the Llama 3 paper here:')
    submitted = st.form_submit_button('Submit')
    if submitted:
        generate_response(query)