import utils
import streamlit as st
from streaming import StreamHandler

from langchain_openai import ChatOpenAI
from langchain.chains.conversation.base import ConversationChain
from langchain.memory.buffer import ConversationBufferMemory
from langchain.memory import ConversationSummaryMemory

st.set_page_config(page_title="Context aware chatbot", page_icon="🧠")
st.header('Context aware chatbot')
st.write('Enhancing Chatbot Interactions through Context Awareness')

class ContextChatbot:

    def __init__(self):
        self.openai_model = utils.configure_openai()

    @st.cache_resource
    def setup_chain(_self):
        summarizer = ChatOpenAI(model_name= "gpt-4o", temperature=0, streaming=True)
        memory = ConversationSummaryMemory(llm = summarizer)
        llm = ChatOpenAI(model_name=_self.openai_model, temperature=0, streaming=True)
        chain = ConversationChain(llm=llm, memory=memory, verbose=True)
        return chain

    @utils.enable_chat_history
    def main(self):
        chain = self.setup_chain()
        user_query = st.chat_input(placeholder="Ask me anything!")
        if user_query:
            utils.display_msg(user_query, 'user')
            with st.chat_message("assistant"):
                st_cb = StreamHandler(st.empty())
                result = chain.invoke(
                    {"input":user_query},
                    {"callbacks": [st_cb]}
                )
                response = result["response"]
                st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    obj = ContextChatbot()
    obj.main()