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):