LangChain with ConversationBufferMemory in Streamlit application does not work

I have a streamlit chatbot that works perfectly fine but does not remember previous chat history. I was trying to add it with langchain ConversationBufferMemory but it does not seem to work.

Here is a sample of chatbot I created:

import streamlit as st
from streamlit_chat import message

from langchain.chains import ConversationChain
from langchain.llms import OpenAI
from langchain.chat_models import AzureChatOpenAI

from langchain.memory import ConversationBufferMemory

from langchain.prompts import (
    ChatPromptTemplate, 
    MessagesPlaceholder, 
    SystemMessagePromptTemplate, 
    HumanMessagePromptTemplate
)


prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template("The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know."),
    MessagesPlaceholder(variable_name="history"),
    HumanMessagePromptTemplate.from_template("{input}")
])


def load_chain(prompt):
    """Logic for loading the chain you want to use should go here."""
    llm = AzureChatOpenAI(
                    deployment_name = 'gpt-35-turbo',
                    model_name = 'gpt-35-turbo',
                    temperature = 0,
                    openai_api_key = '.....',
                    openai_api_base = '.....',
                    openai_api_version = "2023-05-15", 
                    openai_api_type="azure"
                    )
    memory = ConversationBufferMemory(return_messages=True)
    chain = ConversationChain(
        llm=llm,
        verbose=True,
        prompt=prompt,
        memory=memory
    )
    return chain

chain = load_chain(prompt)

# From here down is all the StreamLit UI.
st.set_page_config(page_title="LangChain Demo", page_icon=":robot:")
st.header("LangChain Demo")

if "generated" not in st.session_state:
    st.session_state["generated"] = []

if "past" not in st.session_state:
    st.session_state["past"] = []

if "history" not in st.session_state:
    st.session_state["history"] = []

def get_text():
    input_text = st.text_input("You: ", "Hello, how are you?", key="input")
    return input_text


user_input = get_text()

if user_input:
    output = chain.run(input=user_input, history=st.session_state["history"])
    st.session_state["history"].append((user_input, output))
    st.session_state.past.append(user_input)
    st.session_state.generated.append(output)
    st.write(st.session_state["history"])

if st.session_state["generated"]:

    for i in range(len(st.session_state["generated"]) - 1, -1, -1):
        message(st.session_state["generated"][i], key=str(i))
        message(st.session_state["past"][i], is_user=True, key=str(i) + "_user")

It looks like bot ignores ConversationBufferMemory for some reason. Any help would be appreciated.

@hiper2d's answer said it clearly. I was facing the same issue for the past hour. In hindsight, I should have known.

Any time something must be updated on the screen, Streamlit reruns your entire Python script from top to bottom. This means, the ConversationalBufferMemory will be initiated every time we refresh. If we are going to store it in a session state, I don't see the use of Memory abstractions.

But, if you do want to do that, you can set up a session state:

if "messages" not in st.session_state:
    st.session_state.messages = []

# While processing
st.session_state.messages.append({"role" : "assistant", "content" : output_from_llm})