all_agents_tutorials/memory-agent-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Memory-Enhanced Email Agent with LangGraph\n",
    "\n",
    "This tutorial demonstrates how to build an advanced AI agent with three types of memory using LangGraph and LangMem. We'll create an email assistant that can remember important facts, learn from past examples, and improve its behavior based on feedback.\n",
    "\n",
    "## Key Memory Types:\n",
    "- **Semantic Memory**: Stores facts and knowledge about contacts, preferences, and contexts\n",
    "- **Episodic Memory**: Remembers specific past interactions and examples\n",
    "- **Procedural Memory**: Learns and improves behavioral patterns over time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tutorial Overview: Email Assistant with Memory\n",
    "\n",
    "In this tutorial, we'll build an email agent that can:\n",
    "\n",
    "1. **Triage emails**: Classify incoming emails as 'ignore', 'notify', or 'respond'\n",
    "2. **Draft responses**: Compose contextually appropriate replies using stored knowledge\n",
    "3. **Learn from feedback**: Improve its performance based on user corrections\n",
    "\n",
    "The agent will leverage all three memory types to create a system that becomes more helpful and personalized over time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## System Workflow\n",
    "\n",
    "<div style=\"text-align: center;\">\n",
    "\n",
    "<img src=\"../images/memory-enhanced-email-agent.svg\" alt=\"essay grading system langgraph\" style=\"width:80%; height:auto;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Setting the Stage: Imports and Setup\n",
    "First, the imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from typing import TypedDict, Literal, Annotated, List\n",
    "from langgraph.graph import StateGraph, START, END, add_messages\n",
    "from langchain.chat_models import init_chat_model\n",
    "from langchain_core.tools import tool\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.schema import HumanMessage\n",
    "from pydantic import BaseModel, Field\n",
    "from langgraph.store.memory import InMemoryStore  # For storing memories\n",
    "from langmem import create_manage_memory_tool, create_search_memory_tool # LangMem!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load environment variables (your OpenAI API key)\n",
    "load_dotenv()\n",
    "\n",
    "# Initialize the LLM\n",
    "llm = init_chat_model(\"openai:gpt-4o-mini\")\n",
    "\n",
    "# Initialize the memory store (we'll use an in-memory store for simplicity)\n",
    "store = InMemoryStore(index={\"embed\": \"openai:text-embedding-3-small\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Defining Our Agent's \"Brain\": The State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class State(TypedDict):\n",
    "    email_input: dict  # The incoming email\n",
    "    messages: Annotated[list, add_messages]  # The conversation history\n",
    "    triage_result: str # The result of the triage (ignore, notify, respond)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. The Triage Center: Deciding What to Do (with Episodic Memory!)\n",
    "\n",
    "We'll enhance the triage step with episodic memory.\n",
    "\n",
    "First, let's define the Router:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Router(BaseModel):\n",
    "    reasoning: str = Field(description=\"Step-by-step reasoning behind the classification.\")\n",
    "    classification: Literal[\"ignore\", \"respond\", \"notify\"] = Field(\n",
    "        description=\"The classification of an email: 'ignore', 'notify', or 'respond'.\"\n",
    "    )\n",
    "\n",
    "llm_router = llm.with_structured_output(Router)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, the enhanced triage_email function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_few_shot_examples(examples):\n",
    "    formatted_examples = []\n",
    "    for eg in examples:\n",
    "        email = eg.value['email']\n",
    "        label = eg.value['label']\n",
    "        formatted_examples.append(\n",
    "            f\"From: {email['author']}\\nSubject: {email['subject']}\\nBody: {email['email_thread'][:300]}...\\n\\nClassification: {label}\"\n",
    "        )\n",
    "    return \"\\n\\n\".join(formatted_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def triage_email(state: State, config: dict, store: InMemoryStore) -> dict:\n",
    "    email = state[\"email_input\"]\n",
    "    user_id = config[\"configurable\"][\"langgraph_user_id\"]\n",
    "    namespace = (\"email_assistant\", user_id, \"examples\")  # Namespace for episodic memory\n",
    "\n",
    "    # Retrieve relevant examples from memory\n",
    "    examples = store.search(namespace, query=str(email))\n",
    "    formatted_examples = format_few_shot_examples(examples)\n",
    "\n",
    "    prompt_template = PromptTemplate.from_template(\"\"\"You are an email triage assistant.  Classify the following email:\n",
    "    From: {author}\n",
    "    To: {to}\n",
    "    Subject: {subject}\n",
    "    Body: {email_thread}\n",
    "\n",
    "    Classify as 'ignore', 'notify', or 'respond'.\n",
    "\n",
    "    Here are some examples of previous classifications:\n",
    "    {examples}\n",
    "    \"\"\")\n",
    "\n",
    "    prompt = prompt_template.format(examples=formatted_examples, **email)\n",
    "    messages = [HumanMessage(content=prompt)]\n",
    "    result = llm_router.invoke(messages)\n",
    "    return {\"triage_result\": result.classification}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Action Time: Defining Tools (with Semantic Memory!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def write_email(to: str, subject: str, content: str) -> str:\n",
    "    \"\"\"Write and send an email.\"\"\"\n",
    "    print(f\"Sending email to {to} with subject '{subject}'\\nContent:\\n{content}\\n\")\n",
    "    return f\"Email sent to {to} with subject '{subject}'\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def check_calendar_availability(day: str) -> str:\n",
    "    \"\"\"Check calendar availability for a given day.\"\"\"\n",
    "    return f\"Available times on {day}: 9:00 AM, 2:00 PM, 4:00 PM\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create LangMem memory tools (using the configured user ID)\n",
    "manage_memory_tool = create_manage_memory_tool(namespace=(\"email_assistant\", \"{langgraph_user_id}\", \"collection\"))\n",
    "search_memory_tool = create_search_memory_tool(namespace=(\"email_assistant\", \"{langgraph_user_id}\", \"collection\"))\n",
    "\n",
    "tools = [write_email, check_calendar_availability, manage_memory_tool, search_memory_tool]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. The Response Agent: Putting It All Together (with Semantic Memory!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "def create_agent_prompt(state, config, store):\n",
    "    messages = state['messages']\n",
    "    user_id = config[\"configurable\"][\"langgraph_user_id\"]\n",
    "    \n",
    "    # Get the current response prompt from procedural memory\n",
    "    system_prompt = store.get((\"email_assistant\", user_id, \"prompts\"), \"response_prompt\").value\n",
    "    \n",
    "    return [{\"role\": \"system\", \"content\": system_prompt}] + messages\n",
    "\n",
    "# Try using the current API signature\n",
    "response_agent = create_react_agent(\n",
    "    tools=tools,\n",
    "    prompt=create_agent_prompt,\n",
    "    store=store,\n",
    "    model=llm  # Using 'model' instead of 'llm'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Building the Graph: Connecting the Pieces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = StateGraph(State)\n",
    "\n",
    "# Update this line to pass the store to the node\n",
    "workflow.add_node(\"triage\", lambda state, config: triage_email(state, config, store))\n",
    "workflow.add_node(\"response_agent\", response_agent)\n",
    "\n",
    "def route_based_on_triage(state):\n",
    "  if state[\"triage_result\"] == \"respond\":\n",
    "    return \"response_agent\"\n",
    "  else:\n",
    "    return END\n",
    "\n",
    "workflow.add_edge(START, \"triage\")\n",
    "workflow.add_conditional_edges(\"triage\", route_based_on_triage,\n",
    "                              {\n",
    "                                  \"response_agent\": \"response_agent\",\n",
    "                                  END: END\n",
    "                              })\n",
    "\n",
    "# Compile the graph\n",
    "agent = workflow.compile(store=store)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Show the current agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAS8AAAD5CAIAAAD5rCQJAAAAAXNSR0IArs4c6QAAIABJREFUeJzt3WdcU+ffBvA7JyEJJGHLkrAEBQRBwYmrTrQu3HthraJVcdfaqo+2WncVrVZa98JVZx20dbVq3YoiG5lhQ8iez4v4T6kNqJTkPuH8vh9fYM5JuJJw5T4r59C0Wi0CAJAAgTsAAOANaCMAZAFtBIAsoI0AkAW0EQCygDYCQBYM3AHAPxTlyCRCtaRapVJo5TIN7jjvZsEk6AxkZc2w4tGd+CyGBXy+1x8N9jeSQfoTUcZzUXaS2CPQSq3UWvEY9s5MhTm0kckmqsqUEqFKUq0uLZA7e7B9gjnNw3hsKzruaOYH2ojZq/vCP8+X8ZtbegZwvII4TJZ5jy25qZLM5+Ki1zJ+C6uOHzvgjmNmoI3YCMuV1w4VWTtYdBrowLFubKsMD66V37tc3mucc4s2PNxZzAa0EY+sJPHN0yWDZrjaObFwZzEWjVp76+dShgUtYpAj7izmAdqIQUGm9NFvFQOmueEOYgqPfquQVKs7D4ZCvpt5r6WYoxd3qx4mUqWKCKE2PezYHOKXvYW4g5gBaKNJFb2WvbgjHDidKlXUCe9lb+/C/OtKOe4gZAdtNB2lXHPnYtnIWD7uIBi07+cgl6izXohxByE1aKPp3D5X6hvCxZ0Cm5ButjdPleBOQWrQRhMRlitzkiVBETa4g2BjbW/hEWCV9EcV7iDkBW00kWc3q7oOpfp2xc6DHDOei3CnIC9oo4k8u1XpEcDBnQIzCxahUaO8NAnuICQFbTSF18li9+ZWdDrNlL80ISFh5cqV9bjjkiVLzp8/b4RECCHkE8zJfA7bcgyDNppCXrq0eRtTb79JTk428R3fR7NWnLJChfEe36xBG02hOEfOtTXWkaiPHz+eNm1a9+7du3TpEh0d/ejRI4TQ9OnTz58/f+HChfDw8JSUFITQ5cuXx40b16VLl549e8bGxubl5enunpCQ0Lt37xs3bvTu3Xvr1q3h4eEFBQWrVq3q3r27MdJybS0Ks6RqFRwBZgC00RQk1SornlHaKJVK582b5+Pjs3fv3v379/v5+c2ZM0coFG7evNnf379Pnz6JiYm+vr4vXrxYvnx5RETEwYMHt23bJpVKFy1apHsECwsLqVR67NixlStXjhgx4tKlSwihRYsWnT171hiBEUIca4ZYqDLSg5u1xvbVAXKSVKuteEb5vp9AIBCLxf379/f29kYILVy4sHfv3kwmk81mMxgMJpNpa2uLEPL09Dx48KCfnx+DwUAIjR07dv78+eXl5fb29jQaTSaTjR07NiIiAiEkl8sRQlZWVjY2xtoZo2ujtb2FkR7ffEEbTYHJJugMo2zC8fDw8PT0XL58+fDhwzt06NCiRYuwsLB/z8blcvPz8+Pi4nJzc2UymVKpRAgJhUJ7e3vdDMHBwcaIZxDLitCawfeoMYAlVVMgCJqRls3odHp8fHyvXr3OnDkzfvz4gQMHXrx48d+zXb16denSpUFBQdu2bTty5MgXX3zx1gxcruk2MlWWKI20pGDuoI2mYGVNlwjVRnpwOzu7efPmnT17NiEhoV27ditWrPj3RtEzZ86Eh4fPnDnTy8vL0dFRJpMZKcz7EAtVje/b1Q0C2mgKzh4smdgobczPz79+/bruZx8fn2XLlhEEkZGRobtF/+VVhUKhW4HUuXz5cs2p/2a8b73KpGpXb7aFmZ9wxEjgRTEFZw926mOjHBEmEAgWL1586NCh7Ozs169fx8fHEwShWwnk8XgpKSkpKSmVlZVBQUF3795NSkoqLCxcu3ato6MjQujly5f/HiRZLBaLxXr06FFKSopK1fBL11nPxTAw1gbaaAreQZysJKMcgBIWFrZixYqLFy+OHz9+4sSJ9+7d27hxo6enJ0Jo9OjRJSUl0dHRycnJU6dODQsLmzlz5pQpUxwcHL766qv27duvWbNGP67WNHny5MTExJiYGKlU2uCBM5+LfYKpfoRgbeBMHCby2/GiFuG8ps2scAfB7NT2vKhZTQnCpAcJmgsYG00ksL3Nn+fLcKfA7K8r5e6+llDF2sASvIm4eLGteIzM5yKfYMP7EpYuXXr37l2Dk9RqNZ1ueJfAqlWrunXr1qBJ/1bbwXFqtVq3c8Xg1MTERN0xBm9RKTUPEytmbmjW0DEbD1hSNZ2KYsXdS2X9JrsanCqVSmvbaqJSqQz+fSOELC0ta5v031VXV9eWByFU2+/l8QyfQPX+tXIrLr1lR+p+3/qdoI0mlfqwOuuluO8EF9xBTO3VA2HuK2nv8c64g5AarDeaVPMwHs+WcfvnUtxBTCovTfL490qo4jvB2IhB0p9VlSVKipzwN/ul+Mn1yiExTXEHMQMwNmIQ1MmGbUVc2FOAO4jRPb1Z+fx2FVTxPcHYiE3WC/HvCcWh3Wzb9LDDnaXhZT4X/Xm+rHkbbrtIuFLV+4I24qTRaO9cKEu+Jwz9yNYrkOPoZvZXyBELVVlJ4txUiUaNOg10sHNi4k5kTqCN+EnF6ue3KzOeihUyjV9rLo2gcWzo1vYWGnP4EiCdQRNVKiVCtbhKVZInr65QeQdx/NvxXL0scUczP9BGEhGWKwszZdUVSnGVmkag6ooGPmg7KSnJ19eXzWY34GNybRhqldbKms6xYTi5M509oYT1B22kkCFDhmzfvp3Pp+KFQMwCbFMFgCygjQCQBbSRQpo1gyO2SQ3aSCH6M3QAcoI2Uoi1tTXuCKAu0EYKEQqFuCOAukAbKcTJyQl3BFAXaCOFFBcX444A6gJtpJDmzZvTaHBOGvKCNlJIamoqHHpFZtBGAMgC2kgh+itSAXKCNlJIeXk57gigLtBGCtFdOxV3ClAraCOFlJeXw1YcMoM2AkAW0EYK8fT0hCVVMoM2Usjr169hSZXMoI0AkAW0kUJ8fX1xRwB1gTZSSHp6Ou4IoC7QRgDIAtpIIfAdDpKDNlIIfIeD5KCNAJAFtJFC4AyOJAdtpBA4gyPJQRsBIAtoI4XA+VRJDtpIIXA+VZKDNlKIl5cX7G8kM2gjhWRnZ8P+RjKDNgJAFtBGCnF0dIQlVTKDNlJIaWkpLKmSGbSRQvz8/AgC3nHygveGQtLS0jQaDe4UoFbQRgqBsZHk4L2hEBgbSQ7aSCGurq64I4C60GAjW6PXt29fJpNJEERpaamNjQ2dTqfRaBwO5+jRo7ijgX9g4A4AjI5OpxcWFup+1l3emMlkRkdH484F3gZLqo1fx44d37qFz+cPHjwYUxxQK2hj4zdp0qQmTZro/8tkMseMGYM1ETAM2tj4eXh4tG/fXr+BwMPDY8iQIbhDAQOgjZQwefJkNzc33cA4cuRI3HGAYdBGSvDy8tKtPfL5/KFDh+KOAwyDbarkIq5SlQkUKmXD73bq2XFsyuPKfn37ZSaJG/zBCQLZOFrYOTEb/JEpBfY3kkV1hfLGyZLiXLlHAFciVOGO82E4NoyCDAnHmt6qq61vCBd3HHMFYyMpiCpVZ78v6D7K1cbRjIcXjUb76+ECGg01awWFrA9YbySFfauyB8V4mHUVEUIEQes9oenj61WvX0lwZzFL0Eb87lwq6zS4SaP5Vn7EYKcn1ytxpzBL0Eb8CjNlPDvzHhVr4tpa5KdL1CrYHvHBoI34adRanq0F7hQNycXLsrJUiTuF+YE24icWqhrZlw4lQhXRWBa8TQnaCABZQBsBIAtoIwBkAW0EgCygjQCQBbQRALKANgJAFtBGAMgC2ggAWUAbASALaCMAZAFtbGxOnznes3c73ClAfUAbzdLKVUsuXzlvcFLr0PB5c5eaOhBoCNBGs5SamlzbJG/vZgMHwFnhzBK00fx81DO8UFDw7fpVAwd3RwgNGdrr5KkjSz6f0yeyo0gkqrmkWlFR/s26r4aPjOzbr9P4iVGnTx/TP8jz508+mT62T2THyVNH3Pvrz8/mRm/9bp1uUmVlxTfrvho15uPI/hExsyc/fvIA0xOlHDhLlflJOHZp5Oj+n81e1LNnJEKIwWCcv3C6U8euE8dPY7PZNedcv/H/cnOyv/ziG3t7h+dJTzZt/trJ2aVzRHe5XL78qwVeXj474vaJRaIdOzdVVJb7NmuOENJoNEuWfiYSi5YsXulg73j23Imln8/5fscBHx9ffM+YKqCN5sfa2gYhZGVlZWNtgxCi0WhsFvvT6XP+PeesmAUEQbi5NkUI8fmeZ8+eePDgbueI7nfu3hIKq2Lnfu7l5YMQmvPZ4jnzpunu8uDhvdS0V5s37WodGo4Qmj1r4YOH906fObZwwXKTP1HKgTY2Bi1btjJ4uyXb8sixfU+ePKiqqtRoNNXVwqZN+QihnJxsLoerqyJCKDg41MbGVvdzcnKShYVFaEiY7r8EQbQKbp2enmKqp0Jp0MbGgMMxcP5SlUq1eOlstVo9e9ZCD74XnU5f/tUC3SShsMqKw6k5s268RQhJJGKlUtm3Xyf9JLVabW/vYORnABC0sTFLTk7KzEz/bsueVq1a626pqqxwdXFDCLFYLJlMVnNmobBK9wOHw2UymXt2H6k5lSBga58pwKtsrt55yQa5Ql5z0Hvx4lmhoEB3r6ZN+UJhVX5Bnm7S8+dPqqrenAHV37+lQqFQq9UeHl66f0wmy9HRycjPBiBoo1lisVgsFuvps0dp6SkqVa1X7PBt1pzJZJ4+c6ysrPT+g7vbtq9vG94hN+91RUV5h/adWSxW3I6NOTnZz58/+X73VgcHR929wtq08/Nt8c3aL588eVgoKEj89fL0T8eePXfChM+PuqCNZmnM6Mk3biQuXBQjlUlrm8fW1m7xohX3798ZN2HwwUPxSxavHDZsrEBQMH/hDHt7hxVfrsvNfT1t+pgdOzfFzIjlcLhMJgshRKfTv1233dvHd8WqxZOnDD94KH7ChGmjRk4w7fOjKLhGFX77V2f3nujOszXpOnyVsIrNYrNYLISQQqEYHNVj+idzooY0zIVWz+54/XG0m51zozplswnAVhwqEolE4ycMbtO63cQJn9BotOMnDhIE0bVLD9y5qA7aSEVcLvfbdXF79myfMy+aoBHNfJtv+HaHftUR4AJtpKjAgKAtm3fjTgH+AbbiAEAW0EYAyALaCABZQBsBIAtoIwBkAW0EgCygjQCQBbQRALKANgJAFtBGAMgC2oifnbMF0jSqb9LwHBh0OObyw0EbMTtw4EBe3uuyQtl7zGseFHJNQYZ007Y1arUadxYzA23ERiKRVFZWVlRU9BwUVFYgxx2nwRRlSwLa2bRt2/bKlSu4s5gZaCMGIpEoNja2qqqKx+PNnTu3Rbi1Uq5+eqMMd64GUFWquHuhpPvwJgMGDOjfvz9CKDo6+tmzZ7hzmQf47j8GP/74o5+fX9euXWvemHikiMGi27uwHJuyCYKGL1190AhtuUAhqlS++LNy/FIPBvPvT/nCwsIff/xx+fLlMpnsrVOhg7dAG03nzp07p0+f3rBhQ20zpDyqznouViq1ZflmtuBq58yk0ZC7n2WbHna1zXPy5EmxWDxp0iTTRjMn0EZTKC8vt7e3j4uLmzJlCuefpxWmlL1793bq1MnJycnOrtbSUhm00ejWr1/v7u4+duxY3EFIQa1Wv379ev369Rs3buRyDZwincqgjUakUqkEAsEff/wxatQo3FnI5f79+yUlJbrNPEAP2mgUOTk5K1eu3LlzJ2y3qFtMTMyQIUP69OmDOwgpwB4Oo9iyZcvcuXOhiu+0adOm33//vaysTKlU4s6CH4yNDenatWs5OTnR0dG4g5gZuVyekZFx8+bNGTNm4M6CE4yNDUOj0QgEgl9//XXCBDhJ/gdjsViBgYF0Ov3nn3/GnQUnGBsbwOHDh7t27dqkSRNYNP2PdEcIfPfdd3PnzsWdBQMYG/+rQ4cOFRUV8fl8qOJ/p3sNg4KCRo5smCuCmBcYG+vv8uXLkZGRxcXFTk5wecMGptFoCIL4448/wsLCqPMxB2NjPUVFRdFoNIQQVNEYdFdT9vb27tmzZ1FREe44JgJj4wfLysry8PDIz8/38PDAnYUS0tLSvL29EUIMRiP/CjOMjR9AKpWOGjWKIAg6nQ5VNBk/Pz+CICIiIpKTk3FnMS4YG9+XVqu9ffu2q6urr68v7iwUdfz48cZ9jCGMje8lNjZWq9V26dIFqoiRropLly7Nzc3FncUooI3vtnnz5qioKN12BYDdsmXL6viOqFmDJdW6nDt3btCgQSqVqtFvPzBHFy5cGDBgAO4UDQk+72u1YcOG6upqKmzKM1P+/v6RkZGN6cx0MDYaoFAomEzmvXv32rdvjzsLqEtJSYlGo2GxWLa2trizNAAYG9+Wnp4eHx+PEIIqkl+TJk2cnZ2Li4uPHTuGO0sDgDa+bfXq1TExMbhTgA/QvHnz3NzcRrChFZZU//bs2bNWrVrhTgHqqbi4uLi4OCgoCHeQ+oOx8Y3MzMyMjAzcKUD9OTk5OTg47N+/H3eQ+oM2IoSQQCB4/PhxVFQU7iDgP3F1deXz+UlJSbiD1BMsqSKtVqvVamHnfqNhvm+o+SVuWIWFhQMHDjTHdw7UhkajrVu37tSpU7iDfDCq/xWeOnXq5MmTuFOABrZs2TK5XC4QCHAH+TCwpAoAWVB3bLx169aaNWtwpwBGdOXKlc2bN+NO8QGo28aEhIRFixbhTgGMqG/fviqVKj09HXeQ9wVLqgCQBRXHRo1Gc/HiRdwpgIncv3/fXDbnULGNCQkJL1++xJ0CmIhCofjmm29wp3gvVPzmnoODA1yrjDoiIiKEQqFUKrW0tMSd5R1gvREAsqDckuqNGzfOnz+POwUwqdTU1B9++AF3inejXBsvX77MYrFwpwAm5ebmdvjwYdwp3o1yS6r3798PDg6mzqUdgM7Dhw8DAwNJvupIuTYCQFrUWlKVyWSzZ8/GnQJgsGvXrgcPHuBO8Q7UaqNIJEpLS8OdAmAgEAgKCwtxp3gHSiypzpgxQywWEwShVColEgmPxyMIQiaTHT9+HHc0YFyjR48mCEKr1eouCEmjvfmDP3r0KO5oBlBi7394ePju3bvf+tyhwscQ0Gq1qampNW/RaDSkPTcnJZZUx40b5+rqWvMWrVYbERGBLxEwkaFDh761Q8vW1nbq1Kn4EtWFEm20tLQcMmQInU7X38Lj8SZNmoQ1FDCFYcOG1bzSplarbdGiRdu2bbGGqhUl2ogQGjNmjLu7u/6/rVq1CgsLw5oImAKDwYiKimIymbr/Wltbk/lTmCptrDk8Ojg4TJkyBXciYCJDhw7l8/m6nwMCAjp06IA7Ua2o0kaE0PDhw/l8vlarDQgIaN26Ne44wEQYDMawYcNYLJa1tfWECRNwx6nLe21TVSk1UpHG+GGMzWJgv5EJCQljRkytrlDhDvNfaTVaawcL3Ck+jFSkVikxbMru/dGg0wm/uLq6Bvm3xfLWW3IIBvPdI9879jcm/yV8dquqXKCw5NLrmA2YnrWDRWGm1DuIE9bLztmD7Ifd3r1UmvxXNdfWQlxl9p+D9aBWa1mWRGhX26AImzpmq6uNf10tLy1Qhnaz59mb2WcwRWg0WmGZ4tbpoq5RTdz9SHo8tFajPburwM2P4+HP5VhTYv+2QdXlyqTb5RxbRsRAh9rmqbWN9y6XC8tUHQY4GTMhaBgX9+R2HuLo7kvGQp7Zke8Tau0TxMMdhBQeXCu1YKDOQxwNTjW8LFtRrCjNl0MVzUXPsa6Pfq3AncKAlIdCx6ZsqKJeeG9HUaWqOE9mcKrhNpbmy7VampGDgQbD5jBK8uRiIelWyQTZcpYVbHH4BxqdKMmTG5xkuI2iKnUTPtk3DICaPPw5FQIF7hRvUyq0ds5wpoV/aOLOEleoDU4yvFatlGuUhsdSQFLVFUotIt3iTHWFUtMIdo01KJVCK5cZflEotPcfAJKDNgJAFtBGAMgC2ggAWUAbASALaCMAZAFtBIAsoI0AkAW0EQCygDYCQBbQRgDIAtoIAEIIXb+R+FHP8KqqSowZoI2kk5WVMXrsANwpAAbQRtJJTU3GHQHg0WDnKRkytNf4cVPvP7j7+PH90yevcbncX3+7cuLEodc5WZaWVj0+6jstepbuGqZFRYJdu7c+efpQIhG7uLgNHzZ24IChCKEvvpxPJ+gtW7Y6feZYZWWFl6dPbOwy/xaBuse/eOnnhBOHCgryLC2t2rfrNHNGrL29A0IoaljvCeOii4oFv/1+RSqVBAe3Xjh/uYODI0Lo2bPH8T/tyMpKV6vVzZo1nzZ1VkhIG4SQSqU6dPjH336/WlRU2KSJ84jh4wYPGv7OJ/gq5WV8fFxaeopCIffy9ImOnhUe9uZyDucvnD585KeKivLAgODYeZ9PmjL8qy/XftS9N0IoNe1VfHxcSmqySqVs07rdrJgFLi6uCKGz507u3bdr7ddbt8VtyM3NtubZjB8f3b/f4H37d+8/sAch9FHP8Fkx84cPG9tQb1Dj9qGvs+7PYMfOTYmJv2i0mo4durRujf8E5A02NjIYjPMXTvt4+27ZtJvNZt++fX3N11+EhbXf88PRxYtW3Lz166YtX+vmXL9hVWlZyTdfb/3px4ShUaO3frfu/oO7CCEGnfH48f2CgrwD+06fPHHFxsZ25arFGo0GIXT16sWNm9b06f3xT/HH/2/lhtS0V58vm6s7ow+DwTh6fL+Xl8/Rw+d/ik9IS3t18FA8QkgqlS5bPs/L0ydu296dcfub+fgtXTZHWC1ECO3a/d3xhIPjxkz5Mf74iOHj4nZsvHjp57qfnVwuX7L0Mwsmc+OGnd/vOBDYstWXXy0oKSlGCCW/erF5yzedOnXbs/tIv8hBq9csQwjRaDTd5878BZ/SCGLLpt2bNu4SVlctWDRToVDoYovFogOH4letWH/+7PU+fT7esnVtSUnx6FGThg4d7eTk/PPpxIEDhjXUu9O41eN1RggdObrvwsUzMTHzd+86HBzcWvdng1eDtZFGo7FZ7E+nz2nZshWDwThybF9ISJtPps12b8rv0D7ik2mfJSb+UlxchBDKzEpvG94xwL9lUzf3wYOGx237qZmPn+5B1Bp1zMz5LBaLx+VNnPBJUZHgydOHCKETJw9HRHQbN3YKn+8ZGhr22exFqWmvkpKe6u7l6eHdL3IQg8FwcnJu17ZTSspLhFBxsUAsFvfu1d/T09vLy2f2rIVrv/6OacEUiURnz50YNXJC374D3JvyBw8a3rfPgCNH99X97Oh0+pZNu5cuXunn28LLy2fq5JkymSzpxVOE0NWrF+zs7GfNnO/h4dWnz8dduvTQ3+vc+ZM0Gm35F1/7+Pj6twhctnR1YWH+jZu/6qaqVKqxoyc7OTnTaLR+kYNVKlVGRiqbzWYxWTQazcbG9q3LuYDa1ON1RghdvXaxc0T3fpGDdH8G4WH4z0HekOuNLVu20v2g0WhSU5NrPr3QkDCEUGZmGkKoU8euR4/t2/n9loeP/lIqlQEBQbplTl2v9H+CXl7NEEL5+bkqlSojMy0wIFj/aC1aBCKE0jPeXAnM539lRgjxeNa6AdDd3YPP9/x67fIjR/elpr2i0+mhoWFsNjsjI1WlUtXMFhISVlCQJ5FI6nhqDAZDqVJu275+0pThw0b0nTApCiEkFFYhhHJyslsGttJfcqdL54/090pOTvJv0ZLHfXOOJmdnF1fXpunpKfoZ9Ml5PGuEULWo+sNfdVCf11mpVObn5/r7t9TPExAQZPLgb2vI81tyOFzdDzKZTK1W79u/+8DBPTVnKCsvRQjFzvvcx9v3WuKlEycPczicQQOHT50yk8FgIIQsLa30M+tWMkWiaqlMqtVqraw4+klWllYIIan0TX/eGkN0p6Og0+nbtsYfPbb/4sUze+LjnJ1dpk6e2afPxxKJGCEUu+BT3cKk/kKO5RVlVlZWqBZ5eTkLFs5oHdp22eerHR2aaDSakaP76yYJhVUOjk30c1pb/336WrFYlJae0ieyo/4WpVKpexEMJkdwScl6qcfrLJVJEUJM5t+31/zbw8UoZ5tls9kMBmNo1OiP+w+pebutnf3/roswZtiwMeXlZVevXfzxp522tnYjR4xHCOmqoiOWiHUfZpZsS4Ig/j1JX/7a2NrazZwxb+aMednZmQknDq39doWnl4/uXl8sW+Pj7VtzZqcmznU81G+/X1Wr1cu/+Fr3vhYVCfSTLJhMuezvkwhVVwv1P3M43ODg0AWxX9R8KDK8641MPV5nNoutq7H+FhEJFkyMsoeDIAg/P/+iokIPDy/dP1fXpnQGw5pnLRKJriX+olKpEEL29g6jR00MDAzOzEzX3TErO6NKWKX7Wbeh34PvxWAwfJs1f570RP/4L1880y+v1qagMP/27eu6n728fObHLiMIIjsrw8fHz8LCoqKiXJ/N2trGxsZWf1Exg5RKBYvF1n/EXku8pJ/k7u6RkvpSf5LoW7d/108KCAjKz891c3PX/y4ajabb3gsaUD1eZyaT6eLsmpHx92WPHz68Z5KwdTHW/sbRoybevPXbkaP7cnNfp6WnfLP2yzlzo8ViMY1G27b9242b1qSlpxQU5if+ejk1NTk09M2lFHk8640bV2dnZ6akJu/+4bumTfnBwaEIoREjxt+9ezvhxCGBoPDxkwfbd2wMCWnjX2cbi4sEK1YtTjhxKCcnOzf39cFD8QRBBAYGc7ncAQOG7tu/+7ffrxYU5j9+8mDh4ph161fW/XQC/IOqqip/uXyurKz057MnXqW8sLW1y8hIFYlE3bv2KioS7N23S/d0/rxzU3+vgQOGSaWSb9evTEtPycvLOXAwfkr0yFevXtT9u7hcXllZ6bNnjwWCwg94xSmsfq9zjx59b/9x/cLFM5mZ6QknDtVcz8TFWNdF6Nqlx7LPVx89tm/vvl0cDjcoKGTLpt0cDgch9O26uPj4uPkLPlUoFC4ublMmz4jsO1B3Ly9Pn/btIz5fNre0rMTXt8WqlRt0a3e9ekbK5bKEE4f2xMdxONzOEd0//XSApOpcAAADjElEQVRu3QFCQ8OWLFqRcPLQ3n276HS6p6fP6lUb+XxPhFDMjFgel/fDnm1lZaX29g6dOnaNnjqr7kfr1KnrqJETdv+wbef3m9u3i1i6eNXJU4ePHttPEMS8uUunTpl5+syxk6eOhISEzY9dNv3TcSwmCyHk4uK6edPuH37YNmduNJ1O9/Jqtmb15sDA4Lp/V88ekVeuXliwaObYMZOnTJ7xgS88FdXvdZ40cXpVVeWu3Vs1Gk2H9p2nT5+zctUSDdYTThq+DsdfV8oVMhTS3d6UUVasXCwSVW/a+L0pf+l/p9Vqy8vL9MtFz549nhv7yU/xx729m5kyxrWD+W372PObk+tSHGd25gd2tHfzIVcqvJLvVsolqi5RBhak4ci4/+rp00fDR0YeOBifl5eTlPR05/eb/f1benn54M4FzA91r+D1liNH9x09ZvgYAA8P7x3b99Z2x9DQsM+XrDp+4uCRo3u5XF5oSNin0+fqd5+A91dcXBT9ySiDk6ysuBKJyOCkut+dehg4uHttk1QqNYNh4LIiAf5B67+N+++/mkRLqnhVi6pr28ZtwbBwrLFHkZwawZKqSqUqKS02OEkhlzNrOTKpwd+dQkFBbZPkcrnBA6SYFsz331Rex5IqjI1v8Lg8/cEcAAsGg+Hq4oY7BcKYAdYbASALaCMAZAFtBIAsoI0AkAW0EQCygDYCQBbQRgDIAtoIAFlAGwEgC8PH4jDZNA2CIy3NCc/Ogka+j1aevQVh4LhOSmMwCVoto6DhW3l2FiWvpUZOBRpS9kuRg0tdpy/AgsWmlRfIcacgl+IcKdfO8ChouI1OfBZ8CcGMiCuVbt6WllzSDUOu3pYysQp3CnLRqLXOHmyDk2odG5v6sm+eEhicCsgm8XBB20g73CkM8A3hiiqVyfdwXmqGVP48V+TYlGlfy1KM4W9U6by4U5X2RBTSzcHOmUlnkG+lhPJkEnVVifz2meIBn7g6upH3VMiXDxTy7Jjuzbn2LuQNaVQatbZMIH/xRwW/uWVoN9vaZqurjQihrBfiJzcqBVkyOgOWXMnFztmiqkTpHcRp28fe2sECd5x3eHKjIvleNUKouoKKC64EQbNztgjpZusbUtdpR9/RRj25FOfZe8C/aTWIzTGzBRaNWqtUUPEMziw28T77KN63jQAAYzOzD1cAGjFoIwBkAW0EgCygjQCQBbQRALKANgJAFv8PnKsyxPTeyTkAAAAASUVORK5CYII=",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from langchain_core.runnables.graph import MermaidDrawMethod\n",
    "from IPython.display import display, Image\n",
    " \n",
    "display(\n",
    "    Image(\n",
    "        agent.get_graph().draw_mermaid_png(\n",
    "            draw_method=MermaidDrawMethod.API,\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7. Adding Procedural Memory (Updating Instructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Procedural memory allows the agent to learn and improve its instructions. This requires a separate agent (an \"optimizer\") to update the prompts based on feedback:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_triage_prompt = \"\"\"You are an email triage assistant. Classify the following email:\n",
    "From: {author}\n",
    "To: {to}\n",
    "Subject: {subject}\n",
    "Body: {email_thread}\n",
    "\n",
    "Classify as 'ignore', 'notify', or 'respond'.\n",
    "\n",
    "Here are some examples of previous classifications:\n",
    "{examples}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_response_prompt = \"\"\"You are a helpful assistant. Use the tools available, including memory tools, to assist the user.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store these prompts in the memory store\n",
    "store.put((\"email_assistant\", \"test_user\", \"prompts\"), \"triage_prompt\", initial_triage_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "store.put((\"email_assistant\", \"test_user\", \"prompts\"), \"response_prompt\", initial_response_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def triage_email_with_procedural_memory(state: State, config: dict, store: InMemoryStore) -> dict:\n",
    "    email = state[\"email_input\"]\n",
    "    user_id = config[\"configurable\"][\"langgraph_user_id\"]\n",
    "    \n",
    "    # Retrieve the current triage prompt (procedural memory)\n",
    "    current_prompt_template = store.get((\"email_assistant\", user_id, \"prompts\"), \"triage_prompt\").value\n",
    "    \n",
    "    # Retrieve relevant examples from memory (episodic memory)\n",
    "    namespace = (\"email_assistant\", user_id, \"examples\")\n",
    "    examples = store.search(namespace, query=str(email))\n",
    "    formatted_examples = format_few_shot_examples(examples)\n",
    "    \n",
    "    # Format the prompt\n",
    "    prompt = PromptTemplate.from_template(current_prompt_template).format(examples=formatted_examples, **email)\n",
    "    messages = [HumanMessage(content=prompt)]\n",
    "    result = llm_router.invoke(messages)\n",
    "    return {\"triage_result\": result.classification}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langmem import create_multi_prompt_optimizer\n",
    "\n",
    "def optimize_prompts(feedback: str, config: dict, store: InMemoryStore):\n",
    "    \"\"\"Improve our prompts based on feedback.\"\"\"\n",
    "    user_id = config[\"configurable\"][\"langgraph_user_id\"]\n",
    "    \n",
    "    # Get current prompts\n",
    "    triage_prompt = store.get((\"email_assistant\", user_id, \"prompts\"), \"triage_prompt\").value\n",
    "    response_prompt = store.get((\"email_assistant\", user_id, \"prompts\"), \"response_prompt\").value\n",
    "    \n",
    "    # Create a more relevant test example based on our actual email\n",
    "    sample_email = {\n",
    "        \"author\": \"Alice Smith <alice.smith@company.com>\",\n",
    "        \"to\": \"John Doe <john.doe@company.com>\",\n",
    "        \"subject\": \"Quick question about API documentation\",\n",
    "        \"email_thread\": \"Hi John, I was reviewing the API documentation and noticed a few endpoints are missing. Could you help? Thanks, Alice\",\n",
    "    }\n",
    "    \n",
    "    # Create the optimizer\n",
    "    optimizer = create_multi_prompt_optimizer(llm)\n",
    "    \n",
    "    # Create a more relevant conversation trajectory with feedback\n",
    "    conversation = [\n",
    "        {\"role\": \"system\", \"content\": response_prompt},\n",
    "        {\"role\": \"user\", \"content\": f\"I received this email: {sample_email}\"},\n",
    "        {\"role\": \"assistant\", \"content\": \"How can I assist you today?\"}\n",
    "    ]\n",
    "    \n",
    "    # Format prompts\n",
    "    prompts = [\n",
    "        {\"name\": \"triage\", \"prompt\": triage_prompt},\n",
    "        {\"name\": \"response\", \"prompt\": response_prompt}\n",
    "    ]\n",
    "    \n",
    "    try:\n",
    "        # More relevant trajectories \n",
    "        trajectories = [(conversation, {\"feedback\": feedback})]\n",
    "        result = optimizer.invoke({\"trajectories\": trajectories, \"prompts\": prompts})\n",
    "        \n",
    "        # Extract the improved prompts\n",
    "        improved_triage_prompt = next(p[\"prompt\"] for p in result if p[\"name\"] == \"triage\")\n",
    "        improved_response_prompt = next(p[\"prompt\"] for p in result if p[\"name\"] == \"response\")\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"API error: {e}\")\n",
    "        print(\"Using manual prompt improvement as fallback\")\n",
    "        \n",
    "        # More specific manual improvements\n",
    "        improved_triage_prompt = triage_prompt + \"\\n\\nNote: Emails about API documentation or missing endpoints are high priority and should ALWAYS be classified as 'respond'.\"\n",
    "        improved_response_prompt = response_prompt + \"\\n\\nWhen responding to emails about documentation or API issues, acknowledge the specific issue mentioned and offer specific assistance rather than generic responses.\"\n",
    "    \n",
    "    # Store the improved prompts\n",
    "    store.put((\"email_assistant\", user_id, \"prompts\"), \"triage_prompt\", improved_triage_prompt)\n",
    "    store.put((\"email_assistant\", user_id, \"prompts\"), \"response_prompt\", improved_response_prompt)\n",
    "    \n",
    "    print(f\"Triage prompt improved: {improved_triage_prompt[:100]}...\")\n",
    "    print(f\"Response prompt improved: {improved_response_prompt[:100]}...\")\n",
    "    \n",
    "    return \"Prompts improved based on feedback!\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 8. Let's Run It! (and Store Some Memories!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----\n",
      "triage:\n",
      "{'triage_result': 'respond'}\n",
      "-----\n",
      "-----\n",
      "response_agent:\n",
      "{'messages': [AIMessage(content='How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 306, 'total_tokens': 315, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_0392822090', 'id': 'chatcmpl-BO76Dza4MYUkrJzJkCrVzE4LxU4Bo', 'finish_reason': 'stop', 'logprobs': None}, id='run-d2670fbb-0370-4735-baa2-da204b44c38a-0', usage_metadata={'input_tokens': 306, 'output_tokens': 9, 'total_tokens': 315, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}\n",
      "-----\n"
     ]
    }
   ],
   "source": [
    "email_input = {\n",
    "    \"author\": \"Alice Smith <alice.smith@company.com>\",\n",
    "    \"to\": \"John Doe <john.doe@company.com>\",\n",
    "    \"subject\": \"Quick question about API documentation\",\n",
    "    \"email_thread\": \"\"\"Hi John,\n",
    "\n",
    "I was reviewing the API documentation and noticed a few endpoints are missing. Could you help?\n",
    "\n",
    "Thanks,\n",
    "Alice\"\"\",\n",
    "}\n",
    "\n",
    "config = {\"configurable\": {\"langgraph_user_id\": \"test_user\"}} # Set the user ID!\n",
    "inputs = {\"email_input\": email_input, \"messages\": []}\n",
    "\n",
    "for output in agent.stream(inputs, config=config): # Pass the config\n",
    "    for key, value in output.items():\n",
    "        print(f\"-----\\n{key}:\")\n",
    "        print(value)\n",
    "    print(\"-----\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "#add few shot examples to memory\n",
    "example1 = {\n",
    "    \"email\": {\n",
    "        \"author\": \"Spammy Marketer <spam@example.com>\",\n",
    "        \"to\": \"John Doe <john.doe@company.com>\",\n",
    "        \"subject\": \"BIG SALE!!!\",\n",
    "        \"email_thread\": \"Buy our product now and get 50% off!\",\n",
    "    },\n",
    "    \"label\": \"ignore\",\n",
    "}\n",
    "store.put((\"email_assistant\", \"test_user\", \"examples\"), \"spam_example\", example1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 9. Let's Run Our Complete Memory-Enhanced Agent!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_email_agent(store):\n",
    "    # Define the workflow\n",
    "    workflow = StateGraph(State)\n",
    "    workflow.add_node(\"triage\", lambda state, config: triage_email_with_procedural_memory(state, config, store))\n",
    "    \n",
    "    # Create a fresh response agent that will use the latest prompts\n",
    "    response_agent = create_react_agent(\n",
    "        tools=tools,\n",
    "        prompt=create_agent_prompt,\n",
    "        store=store,\n",
    "        model=llm\n",
    "    )\n",
    "    \n",
    "    workflow.add_node(\"response_agent\", response_agent)\n",
    "    \n",
    "    # The routing logic remains the same\n",
    "    workflow.add_edge(START, \"triage\")\n",
    "    workflow.add_conditional_edges(\"triage\", route_based_on_triage,\n",
    "                                {\n",
    "                                    \"response_agent\": \"response_agent\",\n",
    "                                    END: END\n",
    "                                })\n",
    "    \n",
    "    # Compile and return the graph\n",
    "    return workflow.compile(store=store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Processing original email BEFORE optimization...\n",
      "\n",
      "\n",
      "-----\n",
      "triage:\n",
      "{'triage_result': 'respond'}\n",
      "-----\n",
      "-----\n",
      "response_agent:\n",
      "{'messages': [AIMessage(content='How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 306, 'total_tokens': 315, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_0392822090', 'id': 'chatcmpl-BO76HT96rDP21FAxHnCpqodEfVVek', 'finish_reason': 'stop', 'logprobs': None}, id='run-93f88517-139d-4ef2-9cd7-c111cd0016b8-0', usage_metadata={'input_tokens': 306, 'output_tokens': 9, 'total_tokens': 315, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}\n",
      "-----\n",
      "Added API documentation example to episodic memory\n",
      "Triage prompt improved: You are an email triage assistant. Classify the following email:\n",
      "From: {author}\n",
      "To: {to}\n",
      "Subject: {s...\n",
      "Response prompt improved: You are a helpful assistant. Use the tools available, including memory tools, to assist the user. Pr...\n",
      "\n",
      "\n",
      "Processing the SAME email AFTER optimization with a fresh agent...\n",
      "\n",
      "\n",
      "-----\n",
      "triage:\n",
      "{'triage_result': 'respond'}\n",
      "-----\n",
      "-----\n",
      "response_agent:\n",
      "{'messages': [AIMessage(content='How can I assist you today? If you have any technical documentation inquiries or urgent topics to address, feel free to let me know!', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 29, 'prompt_tokens': 330, 'total_tokens': 359, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_dbaca60df0', 'id': 'chatcmpl-BO76alZhgDR4EBWFauOygbdDniQk1', 'finish_reason': 'stop', 'logprobs': None}, id='run-85c63574-6af5-42ff-b3ed-4fc8c164c3aa-0', usage_metadata={'input_tokens': 330, 'output_tokens': 29, 'total_tokens': 359, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}\n",
      "-----\n"
     ]
    }
   ],
   "source": [
    "# First process the original email to capture \"before\" behavior\n",
    "print(\"\\n\\nProcessing original email BEFORE optimization...\\n\\n\")\n",
    "agent = create_email_agent(store)  # Create a fresh agent\n",
    "for output in agent.stream(inputs, config=config):\n",
    "    for key, value in output.items():\n",
    "        print(f\"-----\\n{key}:\")\n",
    "        print(value)\n",
    "    print(\"-----\")\n",
    "\n",
    "# Add a specific example to episodic memory\n",
    "api_doc_example = {\n",
    "    \"email\": {\n",
    "        \"author\": \"Developer <dev@company.com>\",\n",
    "        \"to\": \"John Doe <john.doe@company.com>\",\n",
    "        \"subject\": \"API Documentation Issue\", \n",
    "        \"email_thread\": \"Found missing endpoints in the API docs. Need urgent update.\",\n",
    "    },\n",
    "    \"label\": \"respond\",\n",
    "}\n",
    "store.put((\"email_assistant\", \"test_user\", \"examples\"), \"api_doc_example\", api_doc_example)\n",
    "print(\"Added API documentation example to episodic memory\")\n",
    "\n",
    "# Provide feedback\n",
    "feedback = \"\"\"The agent didn't properly recognize that emails about API documentation issues \n",
    "are high priority and require immediate attention. When an email mentions \n",
    "'API documentation', it should always be classified as 'respond' with a helpful tone.\n",
    "Also, instead of just responding with 'How can I assist you today?', the agent should \n",
    "acknowledge the specific documentation issue mentioned and offer assistance.\"\"\"\n",
    "\n",
    "# Optimize prompts\n",
    "optimize_prompts(feedback, config, store)\n",
    "\n",
    "# Process the SAME email after optimization with a FRESH agent\n",
    "print(\"\\n\\nProcessing the SAME email AFTER optimization with a fresh agent...\\n\\n\")\n",
    "new_agent = create_email_agent(store)  # Create a fresh agent with updated prompts\n",
    "for output in new_agent.stream(inputs, config=config):\n",
    "    for key, value in output.items():\n",
    "        print(f\"-----\\n{key}:\")\n",
    "        print(value)\n",
    "    print(\"-----\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion: From Simple Scripts to Truly Intelligent Assistants\n",
    "\n",
    "We've now built an email agent that's far more than a simple script. Like a skilled human assistant who grows more valuable over time, our agent builds a multi-faceted memory system:\n",
    "\n",
    "1. **Semantic Memory**: A knowledge base of facts about your work context, contacts, and preferences\n",
    "2. **Episodic Memory**: A collection of specific examples that guide decision-making through pattern recognition\n",
    "3. **Procedural Memory**: The ability to improve its own processes based on feedback and experience\n",
    "\n",
    "This agent demonstrates how combining different types of memory creates an assistant that actually learns from interactions and gets better over time."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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
}