Chap4-Memory.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 4. Memory: Enabling Your Chatbot to Learn from Interactions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple version of this memory system using LangChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import modules\n",
    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
    "from langchain.memory import ChatMessageHistory\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.prompts.chat import MessagesPlaceholder\n",
    "from langchain_core.runnables import chain\n",
    "from langchain_core.chat_history import InMemoryChatMessageHistory\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, trim_messages, filter_messages\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from operator import itemgetter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='No problem! My name is John. How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 54, 'total_tokens': 69, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-06006072-8841-4bdf-8360-8f93653e3ff8-0', usage_metadata={'input_tokens': 54, 'output_tokens': 15, 'total_tokens': 69, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a prompt template\n",
    "prompt_temp = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"You are a helpful assistant. Answer all questions to the best of your ability.\"),\n",
    "    (\"placeholder\", \"{messages}\"),\n",
    "])\n",
    "\n",
    "# Obtain a chat model\n",
    "llm = ChatOpenAI()\n",
    "\n",
    "\n",
    "# Create a simple chain\n",
    "chain = prompt_temp | llm\n",
    "\n",
    "# Invoke the chain.\n",
    "# Note how the incorporation of the previous conversation in the chain enabled the model to answer the follow-up question in a context-aware manner.\n",
    "chain.invoke({\n",
    "    \"messages\": [\n",
    "        (\"human\",\"Wha is your name?\"),\n",
    "        (\"ai\", \"My name is John.\"),\n",
    "        (\"human\", \"Sorry. What is your name again?\"),\n",
    "    ],\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Whilst this may work for demo purposes, it won’t scale in a production environment because the list of conversation messages can grow significantly. Fortunately, LangChain provides a core utility class called *ChatMessageHistory*, which makes it easier to implement this memory system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='What is your name?', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='My name is John.', additional_kwargs={}, response_metadata={})]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a chat history object that can store messages in memory\n",
    "chat_history = InMemoryChatMessageHistory()\n",
    "\n",
    "# Add a user message to the chat history\n",
    "chat_history.add_user_message(\"What is your name?\")\n",
    "\n",
    "# Add an AI message to the chat history\n",
    "chat_history.add_ai_message(\"My name is John.\")\n",
    "\n",
    "# Print the chat history\n",
    "chat_history.messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='My name is John.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 5, 'prompt_tokens': 53, 'total_tokens': 58, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-f5f05bc3-f720-463c-b289-950e1a984496-0', usage_metadata={'input_tokens': 53, 'output_tokens': 5, 'total_tokens': 58, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We can then integrate the stored chat messages into our chain and send a final prompt to the model\n",
    "response = chain.invoke({\n",
    "    \"messages\": chat_history.messages,\n",
    "})\n",
    "input = \"Sorry, what is your name again?\"\n",
    "chat_history.add_user_message(input)\n",
    "chain.invoke({\n",
    "    \"messages\": chat_history.messages,\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the previous example, we integrated the chat messages into the chain explicitly but this requires the tedious manual management of each new message. In a production setting, we need a way to persist chat history and automate the insertion and updating of it.\n",
    "\n",
    "To solve this problem, we can utilize LangChain’s RunnableWithMessageHistory class to automatically insert and update chat messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, let’s modify our prompt template to incorporate a chat_history parameter which will later contain all prior chat messages\n",
    "prompt_temp = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"You are a helpful assistant. Answer all questions to the best of your ability.\"),\n",
    "    (\"placeholder\", \"{history}\"),\n",
    "    (\"human\", \"{input}\"),\n",
    "])\n",
    "chain = prompt_temp | llm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, let’s use the RunnableWithMessageHistory class to wrap our chain and incorporate the latest user input and chat history.\n",
    "chat_history_for_chain = InMemoryChatMessageHistory()\n",
    "chain_with_message_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    # Session_id is an identifier for the session (conversation) thread that the input messages correspond to. \n",
    "    # This allows you to maintain several conversations or threads with the same chain at the same time.\n",
    "    lambda session_id: chat_history_for_chain,\n",
    "    # An input_messages_key that specifies which part of the input should be tracked and stored in the chat history.\n",
    "    # In this example, we want to track the string passed in as input (match with the \"input\" key in the prompt).\n",
    "    input_messages_key=\"input\",\n",
    "    # A history_messages_key that specifies what the previous messages should be injected into the prompt as. \n",
    "    # Our prompt has a placeholder named \"history\", so we specify this property to match.\n",
    "    history_messages_key=\"history\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s look at an example where we return a chat history corresponding to each session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the chain we used before\n",
    "prompt_temp = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"You are a helpful assistant. Answer all questions to the best of your ability.\"),\n",
    "    (\"placeholder\", \"{history}\"),\n",
    "    (\"human\", \"{input}\"),\n",
    "])\n",
    "llm = ChatOpenAI()\n",
    "chain = prompt_temp | llm\n",
    "\n",
    "# Keep track of the history for each combination of user_id and conversation_id.\n",
    "# Note: This line uses type hinting in Python 3.9+, which indicates that histories is a dictionary that will map session_id strings to chat history objects.\n",
    "# The code implies that when a new session starts, it can be stored in this dictionary like so: histories[session_id] = InMemoryChatMessageHistory()\n",
    "histories: dict[str, InMemoryChatMessageHistory] = {}\n",
    "\n",
    "# Define a function that takes a session_id as an argument and returns a chat history object.\n",
    "# Note: This line also uses type hinting in Python 3.9+. Denoting that session_id is a string, and the default value is an empty string.\n",
    "def get_session_history(session_id: str = ''):\n",
    "    if session_id not in histories:\n",
    "        histories[session_id] = InMemoryChatMessageHistory()\n",
    "    return histories[session_id]\n",
    "\n",
    "# Chain with history\n",
    "with_message_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    get_session_history,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Hello Bob! How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 31, 'total_tokens': 41, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-5e4e40e8-bc1a-45df-bb9b-314db6711696-0', usage_metadata={'input_tokens': 31, 'output_tokens': 10, 'total_tokens': 41, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In action by providing the input and session id\n",
    "with_message_history.invoke(\n",
    "    {\"input\": \"hi im bob!\"},\n",
    "    config={\"configurable\": {\"session_id\": \"123\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Your name is Bob.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 5, 'prompt_tokens': 54, 'total_tokens': 59, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-bcf33259-6689-4d5f-9cbf-5c83e9db6895-0', usage_metadata={'input_tokens': 54, 'output_tokens': 5, 'total_tokens': 59, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Continue the conversation\n",
    "with_message_history.invoke(\n",
    "    {\"input\": \"whats my name?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"123\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"I'm sorry, but I don't have access to your personal information. However, if there's anything specific you'd like help with, feel free to ask!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 33, 'prompt_tokens': 32, 'total_tokens': 65, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-e37b132b-139a-47d9-8bcc-cda5e837be1a-0', usage_metadata={'input_tokens': 32, 'output_tokens': 33, 'total_tokens': 65, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# New session_id --> does not remember\n",
    "with_message_history.invoke(\n",
    "    {\"input\": \"whats my name?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"456\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'123': InMemoryChatMessageHistory(messages=[HumanMessage(content='hi im bob!', additional_kwargs={}, response_metadata={}), AIMessage(content='Hello Bob! How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 31, 'total_tokens': 41, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-5e4e40e8-bc1a-45df-bb9b-314db6711696-0', usage_metadata={'input_tokens': 31, 'output_tokens': 10, 'total_tokens': 41, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), HumanMessage(content='whats my name?', additional_kwargs={}, response_metadata={}), AIMessage(content='Your name is Bob.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 5, 'prompt_tokens': 54, 'total_tokens': 59, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-bcf33259-6689-4d5f-9cbf-5c83e9db6895-0', usage_metadata={'input_tokens': 54, 'output_tokens': 5, 'total_tokens': 59, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]),\n",
       " '456': InMemoryChatMessageHistory(messages=[HumanMessage(content='whats my name?', additional_kwargs={}, response_metadata={}), AIMessage(content=\"I'm sorry, but I don't have access to your personal information. However, if there's anything specific you'd like help with, feel free to ask!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 33, 'prompt_tokens': 32, 'total_tokens': 65, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-e37b132b-139a-47d9-8bcc-cda5e837be1a-0', usage_metadata={'input_tokens': 32, 'output_tokens': 33, 'total_tokens': 65, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})])}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# You can print the dictioinary\n",
    "histories"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Modify Chat History\n",
    "In many cases, the chat history messages aren’t in the best state or format to generate an accurate response from the model. To overcome this problem, we can modify the chat history in a variety of ways."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trimming messages​\n",
    "LLMs have limited context windows, therefore, the final prompt sent to the model can’t exceed the model’s input token limits. In addition, excessive prompt information can distract the model and lead to hallucination.\n",
    "\n",
    "An effective solution to this problem is to limit the number of messages retrieved from chat history and appended to the prompt. In practice, we need only to load and store the most recent chat n history messages. Let’s use an example chat history with some preloaded messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the LangChain's trim_messages function\n",
    "trimmer = trim_messages(\n",
    "    max_tokens=65,\n",
    "    # Maintain the last messages \n",
    "    strategy=\"last\",\n",
    "    token_counter=ChatOpenAI(model=\"gpt-4o\"),\n",
    "    # Include the system message\n",
    "    include_system=True,\n",
    "    # Do not allow partial messages\n",
    "    allow_partial=False,\n",
    "    # start_on=”human” ensures that we never remove an AIMessage (that is a response from the model) \n",
    "    # without also removing corresponding HumanMessage (ie the question for that response).\n",
    "    start_on=\"human\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SystemMessage(content=\"you're a good assistant\", additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='whats 2 + 2', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='4', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='thanks', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='no problem!', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='having fun?', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='yes!', additional_kwargs={}, response_metadata={})]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a long message\n",
    "messages = [\n",
    "    SystemMessage(content=\"you're a good assistant\"),\n",
    "    HumanMessage(content=\"hi! I'm bob\"),\n",
    "    AIMessage(content=\"hi!\"),\n",
    "    HumanMessage(content=\"I like vanilla ice cream\"),\n",
    "    AIMessage(content=\"nice\"),\n",
    "    HumanMessage(content=\"whats 2 + 2\"),\n",
    "    AIMessage(content=\"4\"),\n",
    "    HumanMessage(content=\"thanks\"),\n",
    "    AIMessage(content=\"no problem!\"),\n",
    "    HumanMessage(content=\"having fun?\"),\n",
    "    AIMessage(content=\"yes!\"),\n",
    "]\n",
    "\n",
    "# Trim the message\n",
    "trimmer.invoke(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let’s incorporate the trimmer into a chain and RunnableWithMessageHistory. To use it in the chain, we need to ensure that the trimmer is run before the messages input to our prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a prompt template\n",
    "prompt_temp = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"You are a helpful assistant. Answer all questions to the best of your ability.\"),\n",
    "    (\"placeholder\", \"{messages}\"),\n",
    "])\n",
    "\n",
    "# Obtain an LLM\n",
    "llm = ChatOpenAI()\n",
    "\n",
    "# This makes a \"messages\" key available to prompt template,\n",
    "# after passing the input messages list through the trimmer \n",
    "chain = {\"messages\": trimmer} | prompt_temp | llm\n",
    "\n",
    "# Tracking history\n",
    "history = InMemoryChatMessageHistory()\n",
    "with_message_history = RunnableWithMessageHistory(\n",
    "    chain, \n",
    "    lambda: history\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"Of course! Here's the joke again:\\n\\nWhy did the scarecrow win an award?\\nBecause he was outstanding in his field!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 26, 'prompt_tokens': 78, 'total_tokens': 104, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-38153048-f1ab-41a9-8a4c-76faa27a02b4-0', usage_metadata={'input_tokens': 78, 'output_tokens': 26, 'total_tokens': 104, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Using it\n",
    "with_message_history.invoke(\n",
    "[HumanMessage(content=\"Today is a good day to learn about LangChain. Do you agree?\")]\n",
    ")\n",
    "with_message_history.invoke(\n",
    "[HumanMessage(content=\"Why is sky blue?\")]\n",
    ")\n",
    "with_message_history.invoke(\n",
    "[HumanMessage(content=\"What is the capital of France?\")]\n",
    ")\n",
    "with_message_history.invoke(\n",
    "[HumanMessage(content=\"Tell me a joke.\")]\n",
    ")\n",
    "with_message_history.invoke(\n",
    "[HumanMessage(content=\"What joke did you tell me? Could you repeat it?\")]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "InMemoryChatMessageHistory(messages=[HumanMessage(content='Today is a good day to learn about LangChain. Do you agree?', additional_kwargs={}, response_metadata={}), AIMessage(content=\"Absolutely! Learning about LangChain can be a great way to expand your knowledge and stay informed about new technologies and developments in the field. Let me know if you have any specific questions about LangChain, and I'll do my best to provide you with the information you need.\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 55, 'prompt_tokens': 42, 'total_tokens': 97, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-dc0b5e2c-4982-41e4-86e2-a5fabeb79f01-0', usage_metadata={'input_tokens': 42, 'output_tokens': 55, 'total_tokens': 97, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), HumanMessage(content='Why is sky blue?', additional_kwargs={}, response_metadata={}), AIMessage(content=\"The sky appears blue because of the way Earth's atmosphere scatters sunlight. The short-wavelength blue and violet colors are scattered more easily than other colors in the spectrum, which is why we see the sky as blue during the day.\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 47, 'prompt_tokens': 32, 'total_tokens': 79, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-aa2fb9cb-58d9-4616-bd39-b9b61db67451-0', usage_metadata={'input_tokens': 32, 'output_tokens': 47, 'total_tokens': 79, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), HumanMessage(content='What is the capital of France?', additional_kwargs={}, response_metadata={}), AIMessage(content='The capital of France is Paris.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 7, 'prompt_tokens': 34, 'total_tokens': 41, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-a69b8e91-4952-411d-829f-af6bdd80f399-0', usage_metadata={'input_tokens': 34, 'output_tokens': 7, 'total_tokens': 41, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), HumanMessage(content='Tell me a joke.', additional_kwargs={}, response_metadata={}), AIMessage(content=\"Sure, here's a joke for you:\\n\\nWhy did the scarecrow win an award?\\nBecause he was outstanding in his field!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 26, 'prompt_tokens': 54, 'total_tokens': 80, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-a26ae3c9-3742-48f7-a56f-75c6de847f16-0', usage_metadata={'input_tokens': 54, 'output_tokens': 26, 'total_tokens': 80, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), HumanMessage(content='What joke did you tell me? Could you repeat it?', additional_kwargs={}, response_metadata={}), AIMessage(content=\"Of course! Here's the joke again:\\n\\nWhy did the scarecrow win an award?\\nBecause he was outstanding in his field!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 26, 'prompt_tokens': 78, 'total_tokens': 104, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-38153048-f1ab-41a9-8a4c-76faa27a02b4-0', usage_metadata={'input_tokens': 78, 'output_tokens': 26, 'total_tokens': 104, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Print the history of the conversation\n",
    "history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary memory\n",
    "Aside from trimming messages, we can utilize the LLM to generate a summary of the conversation and then incorporate this summary into the prompt sent to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content=\"Hey there! I'm Nemo.\", additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='Hello!', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='How are you today?', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='Fine thanks!', additional_kwargs={}, response_metadata={})]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use ChatMessageHistory() to save the chat history\n",
    "# Note: while ChatMessageHistory serves as a base class for managing chat histories with potential for various storage implementations, \n",
    "# InMemoryChatMessageHistory is a concrete subclass that handles storage in memory.\n",
    "demo_ephemeral_chat_history = ChatMessageHistory()\n",
    "demo_ephemeral_chat_history.add_user_message(\"Hey there! I'm Nemo.\")\n",
    "demo_ephemeral_chat_history.add_ai_message(\"Hello!\")\n",
    "demo_ephemeral_chat_history.add_user_message(\"How are you today?\")\n",
    "demo_ephemeral_chat_history.add_ai_message(\"Fine thanks!\")\n",
    "demo_ephemeral_chat_history.messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a prompt template\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are a helpful assistant. Answer all questions to the best of your ability. The provided chat history includes facts about the user you are speaking with.\",),\n",
    "    MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"user\", \"{input}\"),\n",
    "    ])\n",
    "\n",
    "# Create a chain that uses the prompt template\n",
    "chain = prompt | llm\n",
    "\n",
    "# Chain with history\n",
    "chain_with_message_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    lambda session_id: demo_ephemeral_chat_history,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"chat_history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, let’s create a function that will distill previous interactions into a summary. We can add this one to the front of the chain too.\n",
    "def summarize_messages(chain_input):\n",
    "    stored_messages = demo_ephemeral_chat_history.messages\n",
    "    if len(stored_messages) == 0:\n",
    "        return False\n",
    "    summarization_prompt = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "            (\n",
    "                \"user\",\n",
    "                \"Distill the above chat messages into a single summary message. Include as many specific details as you can.\",\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "    summarization_chain = summarization_prompt | llm\n",
    "    summary_message = summarization_chain.invoke({\"chat_history\": stored_messages})\n",
    "    demo_ephemeral_chat_history.clear()\n",
    "    demo_ephemeral_chat_history.add_message(summary_message)\n",
    "    return True\n",
    "\n",
    "# Finally, we can add this function to the chain with the message history.\n",
    "chain_with_summarization = (\n",
    "    RunnablePassthrough.assign(messages_summarized=summarize_messages)\n",
    "    | chain_with_message_history\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='You introduced yourself as Nemo.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 7, 'prompt_tokens': 88, 'total_tokens': 95, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-3bd66d22-f57b-4fb5-a2a8-9232c19e2cc1-0', usage_metadata={'input_tokens': 88, 'output_tokens': 7, 'total_tokens': 95, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Now, let’s invoke the chain and see if it remembers the chat history.\n",
    "chain_with_summarization.invoke(\n",
    "    {\"input\": \"What did I say my name was?\"},\n",
    "    {\"configurable\": {\"session_id\": \"unused\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filtering messages\n",
    "As the list of chat history messages grows, a wider variety of types, sub-chains, and models may be utilized. LangChain provides a filter_messages helper that makes it easier to filter the chat history messages by type, id, or name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='example input', additional_kwargs={}, response_metadata={}, name='example_user', id='2'),\n",
       " HumanMessage(content='real input', additional_kwargs={}, response_metadata={}, name='bob', id='4')]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Filtering messages\n",
    "messages = [\n",
    "    SystemMessage(\"you are a good assistant\", id=\"1\"),\n",
    "    HumanMessage(\"example input\", id=\"2\", name=\"example_user\"),\n",
    "    AIMessage(\"example output\", id=\"3\", name=\"example_assistant\"),\n",
    "    HumanMessage(\"real input\", id=\"4\", name=\"bob\"),\n",
    "    AIMessage(\"real output\", id=\"5\", name=\"alice\"),\n",
    "]\n",
    "filter_messages(messages, include_types=\"human\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SystemMessage(content='you are a good assistant', additional_kwargs={}, response_metadata={}, id='1'),\n",
       " HumanMessage(content='real input', additional_kwargs={}, response_metadata={}, name='bob', id='4'),\n",
       " AIMessage(content='real output', additional_kwargs={}, response_metadata={}, name='alice', id='5')]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Another filtering example\n",
    "filter_messages(messages, exclude_names=[\"example_user\", \"example_assistant\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='example input', additional_kwargs={}, response_metadata={}, name='example_user', id='2'),\n",
       " HumanMessage(content='real input', additional_kwargs={}, response_metadata={}, name='bob', id='4'),\n",
       " AIMessage(content='real output', additional_kwargs={}, response_metadata={}, name='alice', id='5')]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Another way to filter messages\n",
    "filter_messages(messages, include_types=[HumanMessage, AIMessage], exclude_ids=[\"3\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The filter_messages helper can also be used imperatively (as above) or declaratively (as below), \n",
    "# making it easy to compose with other components in a chain\n",
    "filter_ = filter_messages(exclude_names=[\"example_user\", \"example_assistant\"])\n",
    "chain = filter_ | llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat history with retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load the document \n",
    "loader = TextLoader(\"TeachingwithGenerativeAI.txt\")\n",
    "doc = loader.load()\n",
    "\n",
    "## Split the document into chunks\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=1000,\n",
    "    chunk_overlap=20,\n",
    ")\n",
    "chunks = text_splitter.split_documents(doc)\n",
    "\n",
    "## Define the embedding model\n",
    "embed_model = OpenAIEmbeddings()\n",
    "\n",
    "# Create the vector store\n",
    "vector_db = FAISS.from_documents(\n",
    "    documents = chunks, \n",
    "    embedding = embed_model)\n",
    "\n",
    "# Create the retriever\n",
    "retriever = vector_db.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let’s define a sub-chain that takes historical chat messages and the latest user question, and reformulates the question if it makes reference to any information in the historical information. We’ll then use this sub-chain inside the final RAG chain, which will, in order,\n",
    "\n",
    "1. Rephrase the user’s question given the conversation history (if there is history)\n",
    "\n",
    "2. Pass the rephrased question to the retriever (see above) to get the most relevant documents\n",
    "\n",
    "3. Pass the original question, chat history and documents to the final prompt to generate an answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to extract the content of a message\n",
    "def get_msg_content(msg):\n",
    "    return msg.content\n",
    "\n",
    "# Define the SYSTEM prompt for contextualizing the chat history to come up with a standalone question\n",
    "contextualize_system_prompt = (\n",
    "    \"Given a chat history and the latest user question \"\n",
    "    \"which might reference context in the chat history, \"\n",
    "    \"formulate a standalone question which can be understood \"\n",
    "    \"without the chat history. Do NOT answer the question, \"\n",
    "    \"just reformulate it if needed and otherwise return it as is.\"\n",
    ")\n",
    "\n",
    "# Define the prompt for contextualizing the chat history to come up with a standalone question\n",
    "contextualize_prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", contextualize_system_prompt),\n",
    "    (\"placeholder\", \"{chat_history}\"),\n",
    "    (\"human\", \"{input}\"),\n",
    "])\n",
    "\n",
    "# Define the chain for contextualizing the chat history to come up with a standalone question\n",
    "contextualize_chain = (\n",
    "    contextualize_prompt\n",
    "    | ChatOpenAI()\n",
    "    | get_msg_content\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the question-answering SYSTEM prompt to generate the final answer\n",
    "qa_system_prompt = (\n",
    "    \"You are an assistant for question-answering tasks. \"\n",
    "    \"Use the following pieces of retrieved context to answer \"\n",
    "    \"the question. If you don't know the answer, say that you \"\n",
    "    \"don't know. Use three sentences maximum and keep the \"\n",
    "    \"answer concise.\"\n",
    "    \"\\n\\n\"\n",
    "    \"{context}\"\n",
    ")\n",
    "\n",
    "# Define the question-answering prompt to generate the final answer\n",
    "qa_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", qa_system_prompt),\n",
    "        (\"placeholder\", \"{chat_history}\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Define the chain to generate the final answer\n",
    "qa_chain = (\n",
    "    qa_prompt\n",
    "    | ChatOpenAI()\n",
    "    | get_msg_content\n",
    ") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the overall chain the uses both the retrieved documents and the chat history to answer the question\n",
    "@chain\n",
    "def history_aware_qa(input):\n",
    "     # Rephrase the question if needed\n",
    "     if input.get('chat_history'):\n",
    "         question = contextualize_chain.invoke(input)\n",
    "     else:\n",
    "         question = input['input']\n",
    "    \n",
    "     # Get context from the retriever\n",
    "     context = retriever.invoke(question)\n",
    "\n",
    "     # Get the final answer\n",
    "     return qa_chain.invoke({\n",
    "         **input,\n",
    "         \"context\": context\n",
    "     })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, let’s incorporate stateful management of chat history and send the final prompt, \n",
    "# including chat history and retrieved context to the model for an output.\n",
    "chat_history_for_chain = InMemoryChatMessageHistory()\n",
    "qa_with_history = RunnableWithMessageHistory(\n",
    "    history_aware_qa,\n",
    "    lambda _: chat_history_for_chain,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"chat_history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Yes, faculty should declare their AI policy in their classes to help students understand what is and is not allowed regarding AI use. It is recommended that instructors explain their AI policy in the syllabus and discuss the reasons for adopting it in class. Providing clear guidelines about Dos and Don’ts, such as acknowledging AI use and specifying its limitations, can help prevent misuse or overuse of AI tools.'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Finally, let's invoke the chain\n",
    "qa_with_history.invoke(\n",
    "    {\"input\": \"Should faculty declare their AI policy in their classes?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"123\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'You asked whether faculty should declare their AI policy in their classes.'"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Try ask a related question that refers to the previous question\n",
    "qa_with_history.invoke(\n",
    "    {\"input\": \"What question did I just ask?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"123\"}},\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}