import streamlit as st
import os
import weaviate
from dotenv import load_dotenv,find_dotenv
from langchain.vectorstores import Weaviate
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

load_dotenv(find_dotenv())
weaviate_api_key = os.getenv('WEAVIATE_API_KEY')
weaviate_cluster = os.getenv('WEAVIATE_CLUSTER')

st.set_page_config(page_title="RAG-Roman Empire")
st.image("public/images/banner.png") 

st.write(
    """
    <div style='text-align: center;'>
        <p>Salvete! I'm here to help you learn about the Roman Empire.
        Feel free to ask me anything about Roman history and culture!</p>
    </div>
    """,
    unsafe_allow_html=True
)

with st.expander("Sample questions you can ask"):
    st.markdown(
        """
        <div>
            <ul>
                <li>What were the major achievements of the Roman Empire?</li>
                <li>How did the Roman Empire rise to power?</li>
               <li>How did the Roman Empire expand its territory over time? etc.</li>
            </ul>
            <p><b>Note:</b> you can also ask irrelevant questions to test how the RAG framework is working.</p>
        </div>
        """,
        unsafe_allow_html=True,
    )

st.sidebar.subheader("Your OpenAI Key Please")
openai_api_key = st.sidebar.text_input('Put Your Key Here: ', type="password")

with st.sidebar:
    if openai_api_key.startswith('sk-') and weaviate_api_key and weaviate_cluster:
        st.success('API keys and URL already provided!', icon='✅')
    else:
        st.warning('Please enter your credentials!', icon='⚠️')
        
        
with st.sidebar.expander("ℹ️ Disclaimer"):
    st.caption(
        """We appreciate your engagement! Please note, This demo app can be shut down 
        after the Weaviate cluster expires on 6/22/2024.
        """
    )
        
# Adding custom HTML content to the sidebar with your name and heading      
    st.sidebar.markdown(
        """
            <p style='text-align: center'> This RAG App is built using: </p>
            <div style='text-align: center;'>
            <a href='https://weaviate.io/' target='_blank'> Weaviate</a>
            <a href='https://www.langchain.com/' target='_blank'> LangChain</a>
                <a href='https://openai.com/' target='_blank'><img src="https://img.icons8.com/?size=100&id=ka3InxFU3QZa&format=png&color=000000" height="26"></a>
                <a href='https://huggingface.co/' target='_blank'> <img src="https://img.icons8.com/?size=100&id=sop9ROXku5bb&format=png&color=000000" height="25"></a>
            </div>
            </p>
        """,
        unsafe_allow_html=True,
    )
    
    st.sidebar.markdown(
        """
        <div style='margin-top: 20px;'>
            <p style='text-align: center; margin-bottom: 0; margin-top: 10;'>App created by</p>
            <p style='text-align: center; margin-top: 0;'><b>Prasad Mahamulkar</b></p>
            <p style='text-align: center'> GitHub
            <a href='https://github.com/prsdm/RAG-Application' target='_blank'> Repository</a>
            <p style='text-align: center'> Follow me for more! </p>
            <div style='text-align: center;'>
                <a href='https://github.com/prsdm' target='_blank'><img src="https://img.icons8.com/fluency/48/000000/github.png" height="26"></a>
                <a href='https://x.com/prsdm17' target='_blank'> <img src="https://img.icons8.com/?size=100&id=phOKFKYpe00C&format=png&color=000000" height="25"></a>
                <a href='https://www.linkedin.com/in/prasad-mahamulkar/' target='_blank'><img src="https://img.icons8.com/?size=100&id=8808&format=png&color=000000" height="25"></a>
                <a href='https://medium.com/@prasadmahamulkar' target='_blank'><img src="https://img.icons8.com/?size=100&id=XVNvUWCvvlD9&format=png&color=000000" height="25"></a>
            </div>
            </p>
        </div>
        """,
        unsafe_allow_html=True,
    )

def get_qa_chain():
    auth_config = weaviate.auth.AuthApiKey(api_key=weaviate_api_key)
    client = weaviate.Client(
        url=weaviate_cluster,
        additional_headers={"X-OpenAI-Api-key": openai_api_key},
        auth_client_secret=auth_config,
        startup_period=10
    )
    
    vectorstore = Weaviate(client, "RomanEmpire", "content", attributes=["source"])
    retriever = vectorstore.as_retriever()

    template = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question.
    If the answer is not in the context, just say that you I don't have access to this information. Use five sentences maximum and keep the answer concise.
    Question: {question}
    Context: {context}
    Answer:
    """
    prompt = ChatPromptTemplate.from_template(template)
    
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, openai_api_key=openai_api_key)
    
    chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )
    
    return chain

def generate_response(input_text):
    if openai_api_key and weaviate_api_key and weaviate_cluster:
        chain = get_qa_chain()
        answer = chain.invoke(input_text)
        st.info(answer)
    else:
        st.warning("Please provide all required API keys and URL!")

with st.form('my_form'):
    text = st.text_input('Enter text:', 'Tell me interesting facts about Roman Empire')
    submitted = st.form_submit_button('Submit')
    if not openai_api_key.startswith('sk-'):
        st.warning('Please enter your OpenAI API key!', icon='⚠')

        
    if submitted and openai_api_key.startswith('sk-'):
        generate_response(text)