from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import ChatMessage
from langchain_openai import ChatOpenAI
import streamlit as st

st.set_page_config(page_title="Streamlit + Langchain")

st.title("Basic Chatbot with Streamlit and Langchain")
st.caption("Features text streaming")


class StreamHandler(BaseCallbackHandler):
    def __init__(self, container, text=""):
        self.container = container
        self.text = text

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.text += token
        self.container.markdown(self.text)


# Text input to enter OpenAI API key
with st.sidebar:
    OPENAI_API_KEY = st.text_input("Enter your OpenAI API Key", type="password")

# Streamlit session state
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        ChatMessage(role="assistant", content="How can I help you?")
    ]
# Display all chat messages from session state
for message in st.session_state.messages:
    st.chat_message(message.role).write(message.content)
# If user submits a prompt in the text input, continue
if prompt := st.chat_input():
    if not OPENAI_API_KEY:
        st.error("Please add your OpenAI API key to continue.")
        st.stop()
    # Add user's prompt to the chat messages
    st.session_state.messages.append(ChatMessage(role="user", content=prompt))
    st.chat_message("user").write(prompt)
    # Display the assistant's response with langchain query
    with st.chat_message("assistant"):
        stream_handler = StreamHandler(st.empty())
        llm = ChatOpenAI(
            model="gpt-4o-mini",
            openai_api_key=OPENAI_API_KEY,
            streaming=True,
            callbacks=[stream_handler],
        )
        response = llm.invoke(st.session_state.messages)
        st.session_state.messages.append(
            ChatMessage(role="assistant", content=response.content)
        )