{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 13 - Agent Memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "This notebook demonstrates how to build a travel booking agent with **persistent memory** using the **Microsoft Agent Framework** (MAF).\n",
    "\n",
    "You will learn how different types of agent memory — working, short-term, and long-term — affect how an agent retains and uses information across conversations.\n",
    "\n",
    "**Prerequisites:**\n",
    "- A Microsoft Foundry project with a deployed chat model (e.g. `gpt-5-mini`).\n",
    "- Logged in with the Azure CLI — run `az login` in your terminal.\n",
    "- `AZURE_AI_PROJECT_ENDPOINT` — your Microsoft Foundry project endpoint.\n",
    "- `AZURE_AI_MODEL_DEPLOYMENT_NAME` — the name of your deployed model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install agent-framework azure-ai-projects azure-identity python-dotenv -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logging.getLogger(\"agent_framework.foundry\").setLevel(logging.ERROR)\n",
    "\n",
    "import os\n",
    "import json\n",
    "import dotenv\n",
    "from typing import Annotated\n",
    "from datetime import datetime\n",
    "\n",
    "from agent_framework import tool\n",
    "from agent_framework.foundry import FoundryChatClient\n",
    "from azure.identity import DefaultAzureCredential\n",
    "\n",
    "dotenv.load_dotenv()\n",
    "\n",
    "endpoint = os.getenv(\"AZURE_AI_PROJECT_ENDPOINT\")\n",
    "deployment_name = os.getenv(\"AZURE_AI_MODEL_DEPLOYMENT_NAME\")\n",
    "\n",
    "missing = [k for k, v in {\n",
    "    \"AZURE_AI_PROJECT_ENDPOINT\": endpoint,\n",
    "    \"AZURE_AI_MODEL_DEPLOYMENT_NAME\": deployment_name\n",
    "}.items() if not v]\n",
    "\n",
    "if missing:\n",
    "    raise ValueError(\n",
    "        f\"Missing required environment variables: {', '.join(missing)}. \"\n",
    "        \"Please set them as environment variables (e.g., in your .env file or shell environment).\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the Microsoft Foundry client\n",
    "client = FoundryChatClient(\n",
    "    project_endpoint=endpoint,\n",
    "    model=deployment_name,\n",
    "    credential=DefaultAzureCredential()\n",
    ")\n",
    "\n",
    "print(\"Microsoft Foundry client configured\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Types of Agent Memory\n",
    "\n",
    "AI agents can leverage different kinds of memory, each serving a distinct purpose:\n",
    "\n",
    "### Working Memory\n",
    "The conversation thread itself — the messages exchanged in a single session. The agent can refer back to earlier messages in the same thread to maintain coherence. In MAF this is created via **`agent.create_session()`**, which returns an `AgentSession`.\n",
    "\n",
    "### Short-Term Memory\n",
    "Information that persists for the duration of a task or session but is not stored permanently. For example, the agent may accumulate facts during a multi-turn planning conversation and use them to produce a final itinerary.\n",
    "\n",
    "### Long-Term Memory\n",
    "Preferences and facts that persist **across sessions**. A returning user should not have to repeat their dietary restrictions or travel style. Long-term memory is typically backed by an external store — a database, file, or vector index — and surfaced to the agent via tools."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working Memory with Sessions\n",
    "\n",
    "The simplest form of memory is the conversation session. When you pass the same session (created via `agent.create_session()`) to successive `agent.run()` calls, the agent sees the full history of that conversation and can recall earlier details.\n",
    "\n",
    "Let's create a travel agent and demonstrate working memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = client.as_agent(\n",
    "    name=\"TravelMemoryAgent\",\n",
    "    instructions=(\n",
    "        \"You are a travel agent who remembers user preferences across conversations. \"\n",
    "        \"Track destinations mentioned, budget constraints, and travel dates.\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "session = agent.create_session()\n",
    "\n",
    "# First message — the user shares preferences\n",
    "response = await agent.run(\n",
    "    \"I love beach destinations and my budget is $3000\",\n",
    "    session=session,\n",
    ")\n",
    "print(\"Agent:\", response)\n",
    "\n",
    "# Second message — the agent should recall the budget from the thread\n",
    "response = await agent.run(\n",
    "    \"What did I say my budget was?\",\n",
    "    session=session,\n",
    ")\n",
    "print(\"Agent:\", response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The agent correctly recalled the budget because both messages share the same session. This is **working memory** — it exists only for the lifetime of the session.\n",
    "\n",
    "### What happens with a new thread?\n",
    "\n",
    "If we create a **new** session, the agent has no memory of the previous conversation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_session = agent.create_session()\n",
    "\n",
    "response = await agent.run(\n",
    "    \"What is my budget?\",\n",
    "    session=new_session,\n",
    ")\n",
    "print(\"Agent:\", response)\n",
    "print(\"\\n💡 The agent has no memory of the previous conversation — it's a fresh session.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Long-Term Memory Pattern\n",
    "\n",
    "To remember user preferences **across sessions**, we need a persistent store that lives outside the conversation thread. The agent accesses this store through **tools** — functions it can call to save and retrieve information.\n",
    "\n",
    "Below we implement a simple in-memory preference store (in production you would back this with a database or vector index) and expose it as tools the agent can use.\n",
    "\n",
    "### Architecture\n",
    "```\n",
    "┌─────────────────┐     ┌──────────────────┐     ┌─────────────────┐\n",
    "│  MAF Agent      │────▶│  @tool functions  │────▶│  Preference     │\n",
    "│  (LLM)          │     │  save / retrieve  │     │  Store (dict)   │\n",
    "└─────────────────┘     └──────────────────┘     └─────────────────┘\n",
    "         │                                                 │\n",
    "    AgentSession                                   Persists across\n",
    "    (working memory)                               sessions\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Persistent preference store (simulated) ---\n",
    "preference_store: dict[str, list[str]] = {}\n",
    "\n",
    "\n",
    "@tool(approval_mode=\"never_require\")\n",
    "def save_preference(\n",
    "    user_id: Annotated[str, \"User identifier\"],\n",
    "    preference: Annotated[str, \"A travel preference to remember\"],\n",
    ") -> str:\n",
    "    \"\"\"Save a user travel preference to long-term memory.\"\"\"\n",
    "    preference_store.setdefault(user_id, []).append(preference)\n",
    "    return f\"✅ Stored: {preference}\"\n",
    "\n",
    "\n",
    "@tool(approval_mode=\"never_require\")\n",
    "def get_preferences(\n",
    "    user_id: Annotated[str, \"User identifier\"],\n",
    ") -> str:\n",
    "    \"\"\"Retrieve all saved travel preferences for a user.\"\"\"\n",
    "    prefs = preference_store.get(user_id, [])\n",
    "    if not prefs:\n",
    "        return f\"No saved preferences for {user_id}.\"\n",
    "    return \"Saved preferences:\\n- \" + \"\\n- \".join(prefs)\n",
    "\n",
    "\n",
    "@tool(approval_mode=\"never_require\")\n",
    "def search_hotels(\n",
    "    query: Annotated[str, \"Search query — location, amenities, or tags\"],\n",
    ") -> str:\n",
    "    \"\"\"Search the hotel database for matching properties.\"\"\"\n",
    "    hotels = [\n",
    "        {\"name\": \"Le Meurice Paris\", \"location\": \"Paris, France\", \"price\": 850, \"tags\": [\"luxury\", \"romantic\", \"spa\"]},\n",
    "        {\"name\": \"Four Seasons Maui\", \"location\": \"Maui, Hawaii\", \"price\": 695, \"tags\": [\"beach\", \"family\", \"resort\"]},\n",
    "        {\"name\": \"Aman Tokyo\", \"location\": \"Tokyo, Japan\", \"price\": 780, \"tags\": [\"luxury\", \"city\", \"spa\"]},\n",
    "        {\"name\": \"Hotel Sacher Vienna\", \"location\": \"Vienna, Austria\", \"price\": 420, \"tags\": [\"historic\", \"accessible\", \"cultural\"]},\n",
    "        {\"name\": \"Fairmont Whistler\", \"location\": \"Whistler, Canada\", \"price\": 380, \"tags\": [\"ski\", \"family\", \"mountain\"]},\n",
    "    ]\n",
    "    q = query.lower()\n",
    "    matches = [\n",
    "        h for h in hotels\n",
    "        if q in h[\"name\"].lower()\n",
    "        or q in h[\"location\"].lower()\n",
    "        or any(q in t for t in h[\"tags\"])\n",
    "    ]\n",
    "    if not matches:\n",
    "        matches = hotels[:3]\n",
    "    return json.dumps(matches, indent=2)\n",
    "\n",
    "\n",
    "print(\"✅ Tools defined: save_preference, get_preferences, search_hotels\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 1 — First-time user books an anniversary trip\n",
    "\n",
    "Sarah visits for the first time. The agent should store her preferences via the tools and use them to recommend hotels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "travel_agent = client.as_agent(\n",
    "    tools=[save_preference, get_preferences],\n",
    "    name=\"TravelBookingAssistant\",\n",
    "    instructions=(\n",
    "        \"You are a personalized travel booking assistant with long-term memory.\\n\"\n",
    "        \"WORKFLOW:\\n\"\n",
    "        \"1. When a user starts a conversation, call get_preferences() to check for saved information.\\n\"\n",
    "        \"2. Store any new preferences the user mentions using save_preference().\\n\"\n",
    "        \"3. Use search_hotels() to find suitable options that match their preferences and budget.\\n\"\n",
    "        \"4. Do NOT recommend hotels that exceed the user's budget.\\n\\n\"\n",
    "        \"IMPORTANT: Always use user_id='sarah_johnson_123' for all memory operations.\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "session_1 = travel_agent.create_session()\n",
    "\n",
    "response = await travel_agent.run(\n",
    "    \"Hi! I'm Sarah and I'm planning a trip for my 10th wedding anniversary. \"\n",
    "    \"We love romantic destinations, fine dining, and spa experiences. \"\n",
    "    \"My husband has mobility issues, so we need accessible accommodations. \"\n",
    "    \"Our budget is around $700-800 per night.\",\n",
    "    session=session_1,\n",
    ")\n",
    "print(\"🤖 Agent:\", response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await travel_agent.run(\n",
    "    \"The Hotel Sacher sounds perfect! We're both vegetarian and I have a \"\n",
    "    \"severe nut allergy. Can you note that for future trips?\",\n",
    "    session=session_1,\n",
    ")\n",
    "print(\"🤖 Agent:\", response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify what was stored\n",
    "print(\"📋 Preference store contents:\")\n",
    "for uid, prefs in preference_store.items():\n",
    "    print(f\"\\n  User: {uid}\")\n",
    "    for p in prefs:\n",
    "        print(f\"    - {p}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 2 — Sarah returns weeks later\n",
    "\n",
    "Sarah starts a **brand-new thread** (simulating a new session). The working memory is empty, but the long-term preference store still has her information. The agent should retrieve it and use it to personalize recommendations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session_2 = travel_agent.create_session()  # New session — no working memory\n",
    "\n",
    "response = await travel_agent.run(\n",
    "    \"Hi, my husband and I are planning another trip. Can you recommend a good hotel?\",\n",
    "    session=session_2,\n",
    ")\n",
    "print(\"🤖 Agent:\", response)\n",
    "print(\"\\n💡 The agent retrieved Sarah's saved preferences from long-term memory \"\n",
    "      \"even though this is a completely new conversation thread.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await travel_agent.run(\n",
    "    \"Great suggestions! For the Maui option, what activities would you recommend for the kids?\",\n",
    "    session=session_2,\n",
    ")\n",
    "print(\"🤖 Agent:\", response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this lesson you learned three types of agent memory and how to implement them with the Microsoft Agent Framework:\n",
    "\n",
    "| Memory Type | MAF Mechanism | Lifetime |\n",
    "|---|---|---|\n",
    "| **Working** | `agent.create_session()` | Single conversation |\n",
    "| **Short-term** | Accumulated context within a thread | Single task / session |\n",
    "| **Long-term** | External store accessed via `@tool` functions | Across sessions |\n",
    "\n",
    "### Key Takeaways\n",
    "1. **`agent.create_session()`** provides working memory — the agent sees the full conversation history within a session.\n",
    "2. **New sessions lose context** — without long-term memory the agent cannot recall past conversations.\n",
    "3. **`@tool` functions** bridge the gap — they let the agent save and retrieve information from a persistent store.\n",
    "4. **Personalization improves over time** — the more preferences are stored, the better the agent's recommendations.\n",
    "\n",
    "### Real-World Applications\n",
    "- **Customer Service**: Remember customer history and preferences\n",
    "- **Personal Assistants**: Maintain context across days or weeks\n",
    "- **Healthcare**: Track patient information and preferences\n",
    "- **E-commerce**: Personalized shopping based on history\n",
    "\n",
    "### Next Steps\n",
    "- Replace the in-memory dict with a database or vector store (e.g. Azure AI Search)\n",
    "- Add memory expiration for time-sensitive information\n",
    "- Build multi-agent systems with shared memory\n",
    "- Explore the Cognee notebook for knowledge-graph-backed memory"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
