{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "91c6a7ef",
   "metadata": {},
   "source": [
    "# Streamlit Chat Message History\n",
    "\n",
    "This notebook goes over how to store and use chat message history in a Streamlit app. StreamlitChatMessageHistory will store messages in\n",
    "[Streamlit session state](https://docs.streamlit.io/library/api-reference/session-state)\n",
    "at the specified `key=`. The default key is `\"langchain_messages\"`.\n",
    "\n",
    "- Note, StreamlitChatMessageHistory only works when run in a Streamlit app.\n",
    "- You may also be interested in [StreamlitCallbackHandler](/docs/integrations/callbacks/streamlit) for LangChain.\n",
    "- For more on Streamlit check out their\n",
    "[getting started documentation](https://docs.streamlit.io/library/get-started).\n",
    "\n",
    "You can see the [full app example running here](https://langchain-st-memory.streamlit.app/), and more examples in\n",
    "[github.com/langchain-ai/streamlit-agent](https://github.com/langchain-ai/streamlit-agent)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d15e3302",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import StreamlitChatMessageHistory\n",
    "\n",
    "history = StreamlitChatMessageHistory(key=\"chat_messages\")\n",
    "\n",
    "history.add_user_message(\"hi!\")\n",
    "history.add_ai_message(\"whats up?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64fc465e",
   "metadata": {},
   "outputs": [],
   "source": [
    "history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b60dc735",
   "metadata": {},
   "source": [
    "You can integrate StreamlitChatMessageHistory into ConversationBufferMemory and chains or agents as usual. The history will be persisted across re-runs of the Streamlit app within a given user session. A given StreamlitChatMessageHistory will NOT be persisted or shared across user sessions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42ab5bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.memory.chat_message_histories import StreamlitChatMessageHistory\n",
    "\n",
    "# Optionally, specify your own session_state key for storing messages\n",
    "msgs = StreamlitChatMessageHistory(key=\"special_app_key\")\n",
    "\n",
    "memory = ConversationBufferMemory(memory_key=\"history\", chat_memory=msgs)\n",
    "if len(msgs.messages) == 0:\n",
    "    msgs.add_ai_message(\"How can I help you?\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a29252de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import LLMChain\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.prompts import PromptTemplate\n",
    "template = \"\"\"You are an AI chatbot having a conversation with a human.\n",
    "\n",
    "{history}\n",
    "Human: {human_input}\n",
    "AI: \"\"\"\n",
    "prompt = PromptTemplate(input_variables=[\"history\", \"human_input\"], template=template)\n",
    "\n",
    "# Add the memory to an LLMChain as usual\n",
    "llm_chain = LLMChain(llm=OpenAI(), prompt=prompt, memory=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cd99b4b",
   "metadata": {},
   "source": [
    "Conversational Streamlit apps will often re-draw each previous chat message on every re-run. This is easy to do by iterating through `StreamlitChatMessageHistory.messages`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bdb637b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import streamlit as st\n",
    "\n",
    "for msg in msgs.messages:\n",
    "    st.chat_message(msg.type).write(msg.content)\n",
    "\n",
    "if prompt := st.chat_input():\n",
    "    st.chat_message(\"human\").write(prompt)\n",
    "\n",
    "    # As usual, new messages are added to StreamlitChatMessageHistory when the Chain is called.\n",
    "    response = llm_chain.run(prompt)\n",
    "    st.chat_message(\"ai\").write(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7adaf3d6",
   "metadata": {},
   "source": [
    "**[View the final app](https://langchain-st-memory.streamlit.app/).**"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "poetry-venv",
   "language": "python",
   "name": "poetry-venv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}