memory

langchain-streamlit-demo/app.py

@@ -9,7 +9,7 @@ from langchain.schema.runnable import RunnableConfig
 from langsmith.client import Client
 
 from llm_stuff import (
-    _MEMORY,
+    _STMEMORY,
     _MODEL_DICT,
     _SUPPORTED_MODELS,
     _DEFAULT_MODEL,
@@ -37,6 +37,14 @@ if "trace_link" not in st.session_state:
     st.session_state.trace_link = None
 if "run_id" not in st.session_state:
     st.session_state.run_id = None
+if len(_STMEMORY.messages) == 0:
+    _STMEMORY.add_ai_message("Hello! I'm a helpful AI chatbot. Ask me a question!")
+
+for msg in _STMEMORY.messages:
+    st.chat_message(
+        msg.type,
+        avatar="🦜" if msg.type in ("ai", "assistant") else None,
+    ).write(msg.content)
 
 model = st.sidebar.selectbox(
     label="Chat Model",
@@ -76,6 +84,12 @@ system_prompt = (
     .replace("}", "}}")
 )
 
+if st.sidebar.button("Clear message history"):
+    print("Clearing message history")
+    _STMEMORY.clear()
+    st.session_state.trace_link = None
+    st.session_state.run_id = None
+
 temperature = st.sidebar.slider(
     "Temperature",
     min_value=_MIN_TEMPERATURE,
@@ -103,19 +117,6 @@ if provider_api_key:
 
 run_collector = RunCollectorCallbackHandler()
 
-if st.sidebar.button("Clear message history"):
-    print("Clearing message history")
-    st.session_state["langchain_messages"].memory.clear()
-    st.session_state.trace_link = None
-    st.session_state.run_id = None
-
-for msg in _MEMORY.messages:
-    with st.chat_message(
-        msg.type,
-        avatar="🦜" if msg.type in ("ai", "assistant") else None,
-    ):
-        st.markdown(msg.content)
-
 
 def _reset_feedback():
     st.session_state.feedback_update = None

langchain-streamlit-demo/llm_stuff.py

@@ -6,11 +6,15 @@ from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chat_models import ChatOpenAI, ChatAnyscale, ChatAnthropic
 from langchain.chat_models.base import BaseChatModel
 from langchain.memory import ConversationBufferMemory, StreamlitChatMessageHistory
-from langchain.memory.chat_memory import BaseChatMemory
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from streamlit_feedback import streamlit_feedback
 
-_MEMORY = StreamlitChatMessageHistory(key="langchain_messages")
+_STMEMORY = StreamlitChatMessageHistory(key="langchain_messages")
+_MEMORY = ConversationBufferMemory(
+    chat_memory=_STMEMORY,
+    return_messages=True,
+    memory_key="chat_history",
+)
 
 _DEFAULT_SYSTEM_PROMPT = "You are a helpful chatbot."
 
@@ -35,14 +39,6 @@ _MIN_TOKENS = 1
 _MAX_TOKENS = 100000
 
 
-def get_memory() -> BaseChatMemory:
-    return ConversationBufferMemory(
-        chat_memory=_MEMORY,
-        return_messages=True,
-        memory_key="chat_history",
-    )
-
-
 def get_llm(
     model: str,
     provider_api_key: str,
@@ -95,9 +91,8 @@ def get_llm_chain(
             ("human", "{input}"),
         ],
     ).partial(time=lambda: str(datetime.now()))
-    memory = get_memory()
     llm = get_llm(model, provider_api_key, temperature, max_tokens)
-    return LLMChain(prompt=prompt, llm=llm, memory=memory or get_memory())
+    return LLMChain(prompt=prompt, llm=llm, memory=_MEMORY)
 
 
 class StreamHandler(BaseCallbackHandler):