{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Memory-Enhanced Conversational Agent\n",
    "\n",
    "## Overview\n",
    "This tutorial outlines the process of creating a conversational AI agent with enhanced memory capabilities. The agent incorporates both short-term and long-term memory to maintain context and improve the quality of interactions over time.\n",
    "\n",
    "## Motivation\n",
    "Traditional chatbots often struggle with maintaining context beyond immediate interactions. This limitation can lead to disjointed conversations and a lack of personalization. By implementing both short-term and long-term memory, we aim to create an agent that can:\n",
    "- Maintain context within a single conversation\n",
    "- Remember important information across multiple sessions\n",
    "- Provide more coherent and personalized responses\n",
    "\n",
    "## Key Components\n",
    "1. **Language Model**: The core AI component for understanding and generating responses.\n",
    "2. **Short-term Memory**: Stores the immediate conversation history.\n",
    "3. **Long-term Memory**: Retains important information across multiple conversations.\n",
    "4. **Prompt Template**: Structures the input for the language model, incorporating both types of memory.\n",
    "5. **Memory Manager**: Handles the storage and retrieval of information in both memory types.\n",
    "\n",
    "## Method Details\n",
    "\n",
    "### Setting Up the Environment\n",
    "1. Import necessary libraries for the language model, memory management, and prompt handling.\n",
    "2. Initialize the language model with desired parameters (e.g., model type, token limit).\n",
    "\n",
    "### Implementing Memory Systems\n",
    "1. Create a store for short-term memory (conversation history):\n",
    "   - Use a dictionary to manage multiple conversation sessions.\n",
    "   - Implement a function to retrieve or create new conversation histories.\n",
    "\n",
    "2. Develop a long-term memory system:\n",
    "   - Create a separate store for persistent information.\n",
    "   - Implement functions to update and retrieve long-term memories.\n",
    "   - Define simple criteria for what information to store long-term (e.g., longer user inputs).\n",
    "\n",
    "### Designing the Conversation Structure\n",
    "1. Create a prompt template that includes:\n",
    "   - A system message defining the AI's role.\n",
    "   - A section for long-term memory context.\n",
    "   - A placeholder for the conversation history (short-term memory).\n",
    "   - The current user input.\n",
    "\n",
    "### Building the Conversational Chain\n",
    "1. Combine the prompt template with the language model.\n",
    "2. Wrap this combination with a component that manages message history.\n",
    "3. Ensure the chain can access and update both short-term and long-term memory.\n",
    "\n",
    "### Creating the Interaction Loop\n",
    "1. Develop a main chat function that:\n",
    "   - Retrieves relevant long-term memories.\n",
    "   - Invokes the conversational chain with the current input and memories.\n",
    "   - Updates the long-term memory based on the interaction.\n",
    "   - Returns the AI's response.\n",
    "\n",
    "### Testing and Refinement\n",
    "1. Run example conversations to test the agent's ability to maintain context.\n",
    "2. Review both short-term and long-term memories after interactions.\n",
    "3. Adjust memory management criteria and prompt structure as needed.\n",
    "\n",
    "## Conclusion\n",
    "The Memory-Enhanced Conversational Agent offers several advantages over traditional chatbots:\n",
    "\n",
    "- **Improved Context Awareness**: By utilizing both short-term and long-term memory, the agent can maintain context within and across conversations.\n",
    "- **Personalization**: Long-term memory allows the agent to remember user preferences and past interactions, enabling more personalized responses.\n",
    "- **Flexible Memory Management**: The implementation allows for easy adjustment of what information is stored long-term and how it's used in conversations.\n",
    "- **Scalability**: The session-based approach allows for managing multiple independent conversations.\n",
    "\n",
    "This implementation provides a foundation for creating more sophisticated AI agents. Future enhancements could include:\n",
    "- More advanced criteria for long-term memory storage\n",
    "- Implementation of memory consolidation or summarization techniques\n",
    "- Integration with external knowledge bases\n",
    "- Emotional or sentiment tracking across interactions\n",
    "\n",
    "By focusing on memory enhancement, this conversational agent design significantly improves upon basic chatbot functionality, paving the way for more engaging, context-aware, and intelligent AI assistants."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Imports\n",
    "\n",
    "First, we'll import the necessary modules and set up our environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "from langchain_community.chat_message_histories import ChatMessageHistory\n",
    "\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from dotenv import load_dotenv\n",
    "import os\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv('OPENAI_API_KEY')\n",
    "# Initialize the language model\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", max_tokens=1000, temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Memory Stores\n",
    "\n",
    "We'll create stores for both short-term (chat history) and long-term memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_store = {}\n",
    "long_term_memory = {}\n",
    "\n",
    "def get_chat_history(session_id: str):\n",
    "    if session_id not in chat_store:\n",
    "        chat_store[session_id] = ChatMessageHistory()\n",
    "    return chat_store[session_id]\n",
    "\n",
    "def update_long_term_memory(session_id: str, input: str, output: str):\n",
    "    if session_id not in long_term_memory:\n",
    "        long_term_memory[session_id] = []\n",
    "    if len(input) > 20:  # Simple logic: store inputs longer than 20 characters\n",
    "        long_term_memory[session_id].append(f\"User said: {input}\")\n",
    "    if len(long_term_memory[session_id]) > 5:  # Keep only last 5 memories\n",
    "        long_term_memory[session_id] = long_term_memory[session_id][-5:]\n",
    "\n",
    "def get_long_term_memory(session_id: str):\n",
    "    return \". \".join(long_term_memory.get(session_id, []))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Template\n",
    "\n",
    "We'll create a prompt template that includes both short-term and long-term memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"You are a helpful AI assistant. Use the information from long-term memory if relevant.\"),\n",
    "    (\"system\", \"Long-term memory: {long_term_memory}\"),\n",
    "    MessagesPlaceholder(variable_name=\"history\"),\n",
    "    (\"human\", \"{input}\")\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conversational Chain\n",
    "\n",
    "Now, we'll set up the conversational chain with message history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = prompt | llm\n",
    "chain_with_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    get_chat_history,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"history\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat Function\n",
    "\n",
    "We'll create a function to handle chat interactions, including updating long-term memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(input_text: str, session_id: str):\n",
    "    long_term_mem = get_long_term_memory(session_id)\n",
    "    response = chain_with_history.invoke(\n",
    "        {\"input\": input_text, \"long_term_memory\": long_term_mem},\n",
    "        config={\"configurable\": {\"session_id\": session_id}}\n",
    "    )\n",
    "    update_long_term_memory(session_id, input_text, response.content)\n",
    "    return response.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example Usage\n",
    "\n",
    "Let's test our memory-enhanced conversational agent with a series of interactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AI: Hello, Alice! How can I assist you today?\n",
      "AI: I don't have real-time weather data, but you can check a weather website or app for the most accurate and up-to-date information. If you tell me your location, I can suggest what to look for!\n",
      "AI: Sunny days are wonderful! They can really lift your mood and are perfect for outdoor activities. Do you have any favorite things you like to do on sunny days?\n",
      "AI: Yes, your name is Alice! How can I assist you further today?\n"
     ]
    }
   ],
   "source": [
    "session_id = \"user_123\"\n",
    "\n",
    "print(\"AI:\", chat(\"Hello! My name is Alice.\", session_id))\n",
    "print(\"AI:\", chat(\"What's the weather like today?\", session_id))\n",
    "print(\"AI:\", chat(\"I love sunny days.\", session_id))\n",
    "print(\"AI:\", chat(\"Do you remember my name?\", session_id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Review Memory\n",
    "\n",
    "Let's review the conversation history and long-term memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conversation History:\n",
      "human: Hello! My name is Alice.\n",
      "ai: Hello, Alice! How can I assist you today?\n",
      "human: What's the weather like today?\n",
      "ai: I don't have real-time weather data, but you can check a weather website or app for the most accurate and up-to-date information. If you tell me your location, I can suggest what to look for!\n",
      "human: I love sunny days.\n",
      "ai: Sunny days are wonderful! They can really lift your mood and are perfect for outdoor activities. Do you have any favorite things you like to do on sunny days?\n",
      "human: Do you remember my name?\n",
      "ai: Yes, your name is Alice! How can I assist you further today?\n",
      "\n",
      "Long-term Memory:\n",
      "User said: Hello! My name is Alice.. User said: What's the weather like today?. User said: Do you remember my name?\n"
     ]
    }
   ],
   "source": [
    "print(\"Conversation History:\")\n",
    "for message in chat_store[session_id].messages:\n",
    "    print(f\"{message.type}: {message.content}\")\n",
    "\n",
    "print(\"\\nLong-term Memory:\")\n",
    "print(get_long_term_memory(session_id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://europe-west1-genai-agents-views-tracker.cloudfunctions.net/genai-agents-tracker?notebook=all-agents-tutorials--memory-enhanced-conversational-agent)"
   ]
  }
 ],
 "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}