src/04-orchestration-lc/chat-conversation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with memory \n",
    "When implementing a long running chat or a chat that goes over multiple sessions you might want to persist the messages of a conversation in an external history store. This will give you the ability to load previous conversations in the model chat prompt to provide context, trim old messages to reduce the amount of distracting information and leverage summaries to keep really long running conversations on point.\n",
    "\n",
    "Overview: <br>\n",
    "We will be using Azure CosmosDB for MongoDB (vCore) to persist chat, use langchain to store, load and trim conversation flows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "session_id = \"session1\" # your session id\n",
    "\n",
    "# Load environment variables\n",
    "if load_dotenv():\n",
    "    print(\"Found OpenAPI Base Endpoint: \" + os.getenv(\"AZURE_OPENAI_ENDPOINT\"))\n",
    "else: \n",
    "    print(\"No file .env found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets create a simple conversation and use the SDK to persist the messages in CosmosDB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#pip install -U langchain-mongodb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "connection_string=os.getenv(\"MONGO_DB_CONNECTION_STRING\")\n",
    "database_name=os.getenv(\"MONGO_DB_database_name\")\n",
    "collection_name=os.getenv(\"MONGO_DB_chathistory_collection_name\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_mongodb import MongoDBChatMessageHistory\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "\n",
    "chat_message_history = MongoDBChatMessageHistory(\n",
    "    session_id=session_id,\n",
    "    connection_string=connection_string,\n",
    "    database_name=database_name,\n",
    "    collection_name=collection_name,\n",
    ")\n",
    "\n",
    "conversational_memory = ConversationBufferMemory(\n",
    "    chat_memory=chat_message_history,\n",
    "    memory_key='chat_history',\n",
    "    return_messages=True)\n",
    "\n",
    "# Prompt: Create a conversation flow of 10 questions and answers between Luke and Yoda about the power of the force for creating cloud born applications on Azure.\n",
    "chat_message_history.add_user_message(\"Master, how can Azure help me in creating cloud-born applications?\")\n",
    "chat_message_history.add_ai_message(\"Azure, a powerful platform it is. Build, deploy, and manage applications across a global network of Microsoft-managed data centers, you can. With your preferred tools and frameworks, it works.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"What kind of applications can I create using Azure?\")\n",
    "chat_message_history.add_ai_message(\"Limitless, the possibilities are. Web apps, mobile apps, AI apps, IoT apps, all can be built. Even microservices or docker-based apps, Azure can handle.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"Can Azure help me in scaling these applications?\")\n",
    "chat_message_history.add_ai_message(\"Indeed, young Padawan. Automatically, Azure scales. Adapts to demand, it does. Worry about infrastructure, you need not. Focus on code, you should.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"How secure is Azure for building applications?\")\n",
    "chat_message_history.add_ai_message(\"Strong, Azure's security is. Advanced security, compliance, and privacy features built-in, it has. Trust Azure, you can.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"What if my application needs to be available all the time?\")\n",
    "chat_message_history.add_ai_message(\"Worry not. Azure promises 99.99% availability. Disaster recovery, backup, and migration options, it provides. Always on, your applications will be.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"Azure help me analyze data from my applications?\")\n",
    "chat_message_history.add_ai_message(\"Yes, young one. Powerful analytics tools Azure has. Insight into your data, it will give. Make informed decisions, you can.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"What about the costs? Is Azure affordable?\")\n",
    "chat_message_history.add_ai_message(\"Flexible, Azure's pricing is. Pay for what you use, you do. Even offer free services, they do.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"Can Azure support open-source technologies?\")\n",
    "chat_message_history.add_ai_message(\"Indeed, Luke. Strong supporter of open source, Azure is. Many languages, tools, and frameworks, it supports.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"Is it possible to automate tasks in Azure?\")\n",
    "chat_message_history.add_ai_message(\"Mmm, automate tasks, you can. Azure DevOps and Azure Automation, use you can. Increase productivity, you will.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"Finally, what if I need help? Does Azure provide support?\")\n",
    "chat_message_history.add_ai_message(\"Fear not, Luke. Strong support, Azure provides. Community support, documentation, tutorials, all available they are. Even professional support options, they have.\")\n",
    "\n",
    "print(\"This is what has been persisted:\")\n",
    "chat_message_history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now go into the Data Explorer of CosmosDB or Mongo Compass to check the documents that have been created.\n",
    "There should be several documents each for a single message that will be connected by the SessionId:\n",
    "\n",
    "```json\n",
    "{\n",
    "\t\"_id\" : ObjectId(\"65d620e61875ca4a55e09fae\"),\n",
    "\t\"SessionId\" : \"session1\",\n",
    "\t\"History\" : \"{\\\"type\\\": \\\"human\\\", \\\"data\\\": {\\\"content\\\": \\\"Can Azure support open-source technologies?\\\", \\\"additional_kwargs\\\": {}, \\\"type\\\": \\\"human\\\", \\\"example\\\": false}}\"\n",
    "}\n",
    "```\n",
    "\n",
    "If you want to delete all items in the database use the MongoDB shell and execute the following command:\n",
    "```\n",
    "use my_db\n",
    "db.chat_histories.deleteMany({})\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets create a conversation flow that includes the history to also use questions on the chat history itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import AzureChatOpenAI\n",
    "import os\n",
    "chat = AzureChatOpenAI(\n",
    "    azure_deployment = os.getenv(\"AZURE_OPENAI_COMPLETION_DEPLOYMENT_NAME\")\n",
    ")\n",
    "\n",
    "from langchain_core.messages import AIMessage, HumanMessage\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant. Answer all questions to the best of your ability.\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"messages\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = prompt | chat\n",
    "\n",
    "follow_up_question = \"Can you summarize the last two answers I got from you?\"\n",
    "\n",
    "chat_message_history.add_user_message(follow_up_question)\n",
    "\n",
    "response = chain.invoke(\n",
    "    {\n",
    "        \"messages\": chat_message_history.messages,\n",
    "    }\n",
    ")\n",
    "\n",
    "print(response)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trim the message history\n",
    "The downside of this approach is that we always have to pass the messages to the chain explicitly. That approach is valid but requires us to keep the message history in sync manually. To overcome that we can use the RunnableWithMessageHistory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ChatMessageHistory\n",
    "demo_ephemeral_chat_history = ChatMessageHistory()\n",
    "\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",
    "\n",
    "demo_ephemeral_chat_history.messages\n",
    "\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant. Answer all questions to the best of your ability.\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = prompt | chat\n",
    "\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",
    ")\n",
    "\n",
    "chain_with_message_history.invoke(\n",
    "    {\"input\": \"What was my name again?\"},\n",
    "    {\"configurable\": {\"session_id\": \"unused\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More interesting is the idea of regularly trimming the history to a fixed size to keep the context window of messages (and the amount of prompt tokens low and relevant). In this case we will only keep the two last messages in the history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "\n",
    "def trim_messages(chain_input):\n",
    "    stored_messages = demo_ephemeral_chat_history.messages\n",
    "    if len(stored_messages) <= 2:\n",
    "        return False\n",
    "\n",
    "    demo_ephemeral_chat_history.clear()\n",
    "\n",
    "    for message in stored_messages[-2:]:\n",
    "        demo_ephemeral_chat_history.add_message(message)\n",
    "\n",
    "    return True\n",
    "\n",
    "\n",
    "chain_with_trimming = (\n",
    "    RunnablePassthrough.assign(messages_trimmed=trim_messages)\n",
    "    | chain_with_message_history\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we now ask the same question as before we will see that the first question in the context is a different one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"These are all the messages before the questions\")\n",
    "\n",
    "demo_ephemeral_chat_history.messages\n",
    "\n",
    "chain_with_trimming.invoke(\n",
    "    {\"input\": \"What is my name again?\"},\n",
    "    {\"configurable\": {\"session_id\": \"unused\"}},\n",
    ")\n",
    "\n",
    "print(\"These are all the messages after the questions\")\n",
    "demo_ephemeral_chat_history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatically create summaries\n",
    "\n",
    "This approach keeps the amount of chat history low but the history can loose relevance very quickly. An approach to solve that is to ask the LLM to create a summary of a fixed size of the previous conversation and only keep that summary in the context while updating the summary continously in CosmosDB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.chat_history import BaseChatMessageHistory\n",
    "from langchain_mongodb import MongoDBChatMessageHistory\n",
    "\n",
    "chat_message_history = MongoDBChatMessageHistory(\n",
    "    session_id=session_id,\n",
    "    connection_string=connection_string,\n",
    "    database_name=database_name,\n",
    "    collection_name=collection_name,\n",
    ")\n",
    "\n",
    "# Prompt: Create a conversation flow of 10 questions and answers between Luke and Yoda about the power of the force for creating cloud born applications on Azure.\n",
    "chat_message_history.add_user_message(\"Master, how can Azure help me in creating cloud-born applications?\")\n",
    "chat_message_history.add_ai_message(\"Azure, a powerful platform it is. Build, deploy, and manage applications across a global network of Microsoft-managed data centers, you can. With your preferred tools and frameworks, it works.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"What kind of applications can I create using Azure?\")\n",
    "chat_message_history.add_ai_message(\"Limitless, the possibilities are. Web apps, mobile apps, AI apps, IoT apps, all can be built. Even microservices or docker-based apps, Azure can handle.\")\n",
    "\n",
    "chat_message_history.add_user_message(\"Can Azure help me in scaling these applications?\")\n",
    "chat_message_history.add_ai_message(\"Indeed, young Padawan. Automatically, Azure scales. Adapts to demand, it does. Worry about infrastructure, you need not. Focus on code, you should.\")\n",
    "\n",
    "print(\"These messages have been persisted:\")\n",
    "chat_message_history.messages\n",
    "\n",
    "def get_session_history(session_id: str) -> BaseChatMessageHistory:\n",
    "    return chat_message_history\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"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",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"user\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = prompt | chat\n",
    "\n",
    "chain_with_saved_message_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    get_session_history,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"chat_history\",\n",
    ")\n",
    "\n",
    "def summarize_messages(chain_input):\n",
    "    stored_messages = chat_message_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 | chat\n",
    "\n",
    "    summary_message = summarization_chain.invoke({\"chat_history\": stored_messages})\n",
    "\n",
    "    chat_message_history.clear()\n",
    "\n",
    "    chat_message_history.add_message(summary_message)\n",
    "\n",
    "    return True\n",
    "\n",
    "\n",
    "chain_with_summarization = (\n",
    "    RunnablePassthrough.assign(messages_summarized=summarize_messages)\n",
    "    | chain_with_saved_message_history\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now if we ask a question the items in the CosmosDB will be updated and you will no longer see every message but instead only a summary of the details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_with_summarization.invoke(\n",
    "    {\"input\": \"What did we talk about Azure?\"},\n",
    "    {\"configurable\": {\"session_id\": session_id}},\n",
    ")\n",
    "\n",
    "print(\"These messages have been persisted after the summary:\")\n",
    "chat_message_history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This demonstrates the basic principles for persisting a chat history, trim the contents and use automatic summaries. Of course we have not implemented session handling or multiple users yet but maybe you can extend the scenario?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}