{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Redis](https://redis.io/wp-content/uploads/2024/04/Logotype.svg?auto=webp&quality=85,75&width=120)\n",
    "\n",
    "# Redis Agent Memory with the NVIDIA NeMo Agent Toolkit\n",
    "\n",
    "## Introduction\n",
    "\n",
    "AI agents need **memory** to feel coherent: they should remember who the user is, what they prefer, and what was said earlier — across turns and across sessions. The [**NVIDIA NeMo Agent Toolkit**](https://github.com/NVIDIA/NeMo-Agent-Toolkit) (NAT) is a framework for building and orchestrating agents, and [**nemo-agent-toolkit-redis**](https://github.com/redis-developer/nemo-agent-toolkit-redis) plugs Redis-backed memory into it.\n",
    "\n",
    "In this recipe we wire NAT up to the open-source, self-hosted [**Redis Agent Memory Server**](https://github.com/redis/agent-memory-server) — a memory service that automatically extracts, stores, and semantically retrieves facts from conversations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What we'll build\n",
    "\n",
    "A chat agent whose memory is fully managed for us:\n",
    "\n",
    "- **Working (short-term) memory** — the live conversation, scoped to a session, with a TTL.\n",
    "- **Long-term memory** — durable facts (preferences, names, etc.) the server extracts from the conversation in the background and recalls later via semantic search.\n",
    "\n",
    "We use the toolkit's `redis_agent_memory_auto_memory` workflow wrapper, which on every turn hydrates the prompt with relevant memory and captures the new turn automatically — no manual save/load calls in our agent code.\n",
    "\n",
    "### Architecture\n",
    "\n",
    "```\n",
    "  your notebook ──run_workflow()──▶ NAT auto-memory workflow\n",
    "                                          │\n",
    "                                          ▼\n",
    "                              Redis Agent Memory Server  (localhost:8000)\n",
    "                                          │\n",
    "                                          ▼\n",
    "                                    Redis Stack  (localhost:6379)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's Begin\n",
    "\n",
    "> **NOTE:** This notebook drives real services (Redis + the Agent Memory Server via Docker) and calls OpenAI, so it is intended to run locally rather than in Google Colab or a CI pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "- **Docker** (to run Redis Stack and the Agent Memory Server locally).\n",
    "- An **OpenAI API key** — the Agent Memory Server uses it for extraction/embeddings, and NAT uses it for the chat LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "89a8927e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_SKIP\n",
    "# nvidia-nat-langchain registers the OpenAI->LangChain LLM client that the\n",
    "# chat_completion function uses; without it you get a KeyError at run time.\n",
    "%pip install -q nemo-agent-toolkit-redis nvidia-nat-langchain requests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efa45565",
   "metadata": {},
   "source": [
    "## Set environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6863fe72",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, getpass\n",
    "\n",
    "if not os.environ.get(\"OPENAI_API_KEY\"):\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API key: \")\n",
    "\n",
    "# The self-hosted Agent Memory Server we start below\n",
    "os.environ[\"REDIS_AGENT_MEMORY_URL\"] = \"http://localhost:8000\"\n",
    "os.environ[\"REDIS_AGENT_MEMORY_NAMESPACE\"] = \"nat-auto-memory\"\n",
    "os.environ[\"NAT_OPENAI_MODEL\"] = \"gpt-4o-mini\"\n",
    "\n",
    "# The open-source server runs with auth disabled, so no auth header is needed.\n",
    "AUTH_HEADERS: dict = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c543681b",
   "metadata": {},
   "source": [
    "## Start the Agent Memory Server (self-hosted)\n",
    "\n",
    "We start Redis Stack and the open-source Agent Memory Server with Docker. The server runs with `DISABLE_AUTH=true` for local development and enables background long-term-memory extraction with the `discrete` strategy.\n",
    "\n",
    "> The alternative is the project's `compose.yml`: `docker compose up -d`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "51d2cdbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nat-redis\n",
      "9ba75e15bb0cf789882a91e7e007d286d64f6b4428682aede02771c3f2ba39c4\n",
      "nat-agent-memory\n",
      "12cfc122f7cafb70bcb5eb93270501987f18d50ed0db03fa2be50085e2f2f3b7\n",
      "Agent Memory Server status: 200 {'now': 1784228341997}\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_SKIP\n",
    "import subprocess\n",
    "import time, requests\n",
    "\n",
    "def sh(cmd):\n",
    "    p = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
    "    if p.returncode != 0:\n",
    "        raise RuntimeError(f\"Command failed ({p.returncode}): {cmd}\\n{p.stderr.strip()}\")\n",
    "    print(p.stdout.strip() or 'ok')\n",
    "\n",
    "# Redis Stack (search + JSON) for the memory server to use as its store\n",
    "sh('docker rm -f nat-redis 2>/dev/null; '\n",
    "   'docker run -d --name nat-redis -p 6379:6379 redis/redis-stack:7.4.0-v8')\n",
    "\n",
    "# Agent Memory Server, talking to the Redis container over the host network.\n",
    "# `-e OPENAI_API_KEY` with no value tells docker to pass it through from our\n",
    "# environment, so the key never lands on the command line (ps / shell history).\n",
    "sh('docker rm -f nat-agent-memory 2>/dev/null; '\n",
    "   'docker run -d --name nat-agent-memory -p 8000:8000 '\n",
    "   '-e REDIS_URL=redis://host.docker.internal:6379 '\n",
    "   '-e OPENAI_API_KEY '\n",
    "   '-e DISABLE_AUTH=true '\n",
    "   '-e LONG_TERM_MEMORY=true '\n",
    "   '--add-host=host.docker.internal:host-gateway '\n",
    "   'redislabs/agent-memory-server:0.14.0 '\n",
    "   'agent-memory api --host 0.0.0.0 --port 8000 --task-backend=asyncio')\n",
    "\n",
    "# First run pulls images, so poll until the server accepts connections.\n",
    "for _ in range(30):\n",
    "    try:\n",
    "        r = requests.get('http://localhost:8000/v1/health', timeout=2)\n",
    "        if r.ok:\n",
    "            print('Agent Memory Server status:', r.status_code, r.json())\n",
    "            break\n",
    "    except requests.exceptions.RequestException:\n",
    "        pass\n",
    "    time.sleep(1)\n",
    "else:\n",
    "    raise RuntimeError(\"Agent Memory Server didn't come up. Check: docker logs nat-agent-memory\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0915691d",
   "metadata": {},
   "source": [
    "## Define the NAT workflow\n",
    "\n",
    "NAT is a full-featured agent toolkit and it manages its agents through workflow config yaml files. You can\n",
    "learn more about config files on Nvidia's docs site [here](https://docs.nvidia.com/nemo/agent-toolkit/1.2/workflows/about/index.html).\n",
    "\n",
    "The workflow config ties three things together: the chat LLM (`openai_llm`), the memory backend pointed at our server (`redis_agent_memory_backend`), and the `redis_agent_memory_auto_memory` wrapper that does hydration + capture. Environment variables (`${...}`) are resolved from the values we set above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2c615c38",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wrote nat_config.yml\n"
     ]
    }
   ],
   "source": [
    "config_yaml = \"\"\"general:\n",
    "  telemetry:\n",
    "    enabled: false\n",
    "\n",
    "llms:\n",
    "  openai_llm:\n",
    "    _type: openai\n",
    "    model_name: ${NAT_OPENAI_MODEL:-gpt-4o-mini}\n",
    "\n",
    "functions:\n",
    "  assistant_chat:\n",
    "    _type: chat_completion\n",
    "    llm_name: openai_llm\n",
    "    system_prompt: >-\n",
    "      You are a helpful assistant. When memory context is provided, use it to\n",
    "      answer consistently about the user's preferences and prior facts.\n",
    "\n",
    "memory:\n",
    "  redis_ltm:\n",
    "    _type: redis_agent_memory_backend\n",
    "    base_url: ${REDIS_AGENT_MEMORY_URL:-http://localhost:8000}\n",
    "    default_namespace: ${REDIS_AGENT_MEMORY_NAMESPACE:-nat-auto-memory}\n",
    "\n",
    "workflow:\n",
    "  _type: redis_agent_memory_auto_memory\n",
    "  description: >-\n",
    "    A chat agent that uses Redis Agent Memory working memory plus memory_prompt\n",
    "    hydration on every turn.\n",
    "  inner_agent_name: assistant_chat\n",
    "  memory_name: redis_ltm\n",
    "  default_user_id: demo-user\n",
    "  default_session_id: demo-session\n",
    "  memory_prompt:\n",
    "    optimize_query: false\n",
    "    long_term_search:\n",
    "      limit: 5\n",
    "  working_memory:\n",
    "    namespace: ${REDIS_AGENT_MEMORY_NAMESPACE:-nat-auto-memory}\n",
    "    model_name: ${NAT_OPENAI_MODEL:-gpt-4o-mini}\n",
    "    ttl_seconds: 86400\n",
    "    long_term_memory_strategy:\n",
    "      strategy: discrete\n",
    "\"\"\"\n",
    "\n",
    "with open(\"nat_config.yml\", \"w\") as f:\n",
    "    f.write(config_yaml)\n",
    "print(\"wrote nat_config.yml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c15353f5",
   "metadata": {},
   "source": [
    "## Run the agent\n",
    "\n",
    "`nat.utils.run_workflow` runs a single turn. We pass `session_kwargs` so NAT knows which `user_id` / `conversation_id` (session) the memory belongs to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "36cd903c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "import warnings\n",
    "from pathlib import Path\n",
    "\n",
    "from nat.utils import run_workflow\n",
    "\n",
    "CONFIG_FILE = Path(\"nat_config.yml\").resolve()\n",
    "\n",
    "\n",
    "async def chat(prompt: str, user_id: str = \"demo-user\", conversation_id: str = \"demo-session\") -> str:\n",
    "    \"\"\"Run one turn through the NAT auto-memory workflow.\n",
    "\n",
    "    NAT maps user_id -> Redis Agent Memory user_id and conversation_id -> session_id,\n",
    "    so working memory is hydrated and turns are captured automatically.\n",
    "    \"\"\"\n",
    "    warnings.filterwarnings(\"ignore\", message=r\".*\\.text\\(\\) as a method is deprecated.*\")\n",
    "    warnings.filterwarnings(\"ignore\", message=r\".*get_working_memory is deprecated.*\")\n",
    "    result = await run_workflow(\n",
    "        config_file=CONFIG_FILE,\n",
    "        prompt=prompt,\n",
    "        to_type=str,\n",
    "        session_kwargs={\"conversation_id\": conversation_id, \"user_id\": user_id},\n",
    "    )\n",
    "    print(f\"User: {prompt}\")\n",
    "    print(f\"Assistant: {result}\\n\")\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "41ba7dc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User: Hi! My name is Justin and my favorite city is Lisbon.\n",
      "Assistant: Hi Justin! Lisbon is such a beautiful city, known for its rich history, stunning architecture, and vibrant culture. What do you love most about it?\n",
      "\n",
      "User: I'm a vegetarian, by the way.\n",
      "Assistant: That's great to know, Justin! Lisbon has some fantastic vegetarian options and a growing food scene. Have you discovered any favorite vegetarian restaurants or dishes in the city?\n",
      "\n",
      "User: Where should I plan a food trip, and what should I keep in mind?\n",
      "Assistant: Planning a food trip is a great way to explore Lisbon! Here are some tips and suggestions:\n",
      "\n",
      "### Where to Eat:\n",
      "1. **Vegetarian and Vegan Restaurants**:\n",
      "   - **Jardim das Cerejas**: A popular buffet-style vegetarian restaurant.\n",
      "   - **The Food Temple**: A cozy spot with creative vegetarian dishes.\n",
      "   - **O Botanista**: Offers a range of delicious plant-based options.\n",
      "\n",
      "2. **Local Specialties**:\n",
      "   - **Bacalhau à Brás** (try a vegetarian version): A classic dish made with shredded cod, but you can often find creative plant-based interpretations.\n",
      "   - **Piri-piri Chickpeas**: A spicy and flavorful dish that you can find at many places.\n",
      "\n",
      "3. **Markets**:\n",
      "   - **Time Out Market**: A great place to sample a variety of local dishes, including vegetarian options. You can find several food stalls featuring fresh ingredients and innovative vegetarian cuisine.\n",
      "\n",
      "4. **Tascas (Traditional Taverns)**: Explore small local taverns that may offer vegetarian options, like grilled veggies and salads.\n",
      "\n",
      "### Tips to Keep in Mind:\n",
      "- **Make Reservations**: Some popular places can get busy, so it’s a good idea to reserve a table.\n",
      "- **Learn Basic Portuguese Phrases**: It can enhance your experience. Simple phrases like \"Sou vegetariano\" (I’m a vegetarian) can help communicate your dietary needs.\n",
      "- **Be Open to Exploration**: Try local foods, even if they're not vegetarian; many places might have creative and flavorful vegetable dishes.\n",
      "- **Stay Hydrated**: November in Lisbon can be mild, but it's always good to drink plenty of water, especially when trying rich foods.\n",
      "\n",
      "### Explore Neighborhoods:\n",
      "- **Alfama**: Famous for its narrow streets and traditional Portuguese eateries.\n",
      "- **Bairro Alto**: Known for vibrant nightlife and diverse food options.\n",
      "\n",
      "With these tips, you’ll be all set for a delicious adventure in Lisbon! Enjoy your food trip!\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Planning a food trip is a great way to explore Lisbon! Here are some tips and suggestions:\\n\\n### Where to Eat:\\n1. **Vegetarian and Vegan Restaurants**:\\n   - **Jardim das Cerejas**: A popular buffet-style vegetarian restaurant.\\n   - **The Food Temple**: A cozy spot with creative vegetarian dishes.\\n   - **O Botanista**: Offers a range of delicious plant-based options.\\n\\n2. **Local Specialties**:\\n   - **Bacalhau à Brás** (try a vegetarian version): A classic dish made with shredded cod, but you can often find creative plant-based interpretations.\\n   - **Piri-piri Chickpeas**: A spicy and flavorful dish that you can find at many places.\\n\\n3. **Markets**:\\n   - **Time Out Market**: A great place to sample a variety of local dishes, including vegetarian options. You can find several food stalls featuring fresh ingredients and innovative vegetarian cuisine.\\n\\n4. **Tascas (Traditional Taverns)**: Explore small local taverns that may offer vegetarian options, like grilled veggies and salads.\\n\\n### Tips to Keep in Mind:\\n- **Make Reservations**: Some popular places can get busy, so it’s a good idea to reserve a table.\\n- **Learn Basic Portuguese Phrases**: It can enhance your experience. Simple phrases like \"Sou vegetariano\" (I’m a vegetarian) can help communicate your dietary needs.\\n- **Be Open to Exploration**: Try local foods, even if they\\'re not vegetarian; many places might have creative and flavorful vegetable dishes.\\n- **Stay Hydrated**: November in Lisbon can be mild, but it\\'s always good to drink plenty of water, especially when trying rich foods.\\n\\n### Explore Neighborhoods:\\n- **Alfama**: Famous for its narrow streets and traditional Portuguese eateries.\\n- **Bairro Alto**: Known for vibrant nightlife and diverse food options.\\n\\nWith these tips, you’ll be all set for a delicious adventure in Lisbon! Enjoy your food trip!'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# NBVAL_SKIP\n",
    "# A multi-turn conversation. Turn 3 relies on memory captured in turns 1-2.\n",
    "await chat(\"Hi! My name is Justin and my favorite city is Lisbon.\")\n",
    "await chat(\"I'm a vegetarian, by the way.\")\n",
    "# New turn -> the auto-memory wrapper hydrates prior facts from Redis Agent Memory\n",
    "await chat(\"Where should I plan a food trip, and what should I keep in mind?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e56cec",
   "metadata": {},
   "source": [
    "Notice the third answer reflects facts (name, favorite city, vegetarian) from earlier turns even though we never passed them back in — the auto-memory wrapper retrieved them from Redis Agent Memory.\n",
    "\n",
    "## Inspect long-term memory\n",
    "\n",
    "Facts are promoted to long-term memory in the background. We can query them directly through the server's REST API.\n",
    "\n",
    "**NOTE:** Memory extraction is an async background task so these memories may not populate right away.\n",
    "If the request in the cell below isn't returning any memories try waiting a few seconds and running it again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8b2b64ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- [semantic] User's favorite city is Lisbon\n",
      "- [semantic] User is a vegetarian\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_SKIP\n",
    "# Inspect what got promoted to long-term memory via the Agent Memory REST API.\n",
    "base = os.environ[\"REDIS_AGENT_MEMORY_URL\"].rstrip(\"/\")\n",
    "namespace = os.environ.get(\"REDIS_AGENT_MEMORY_NAMESPACE\", \"nat-auto-memory\")\n",
    "\n",
    "resp = requests.post(\n",
    "    f\"{base}/v1/long-term-memory/search\",\n",
    "    headers={\"Content-Type\": \"application/json\", **AUTH_HEADERS},\n",
    "    json={\"text\": \"favorite city and diet\", \"namespace\": {\"eq\": namespace}, \"limit\": 5},\n",
    "    timeout=30,\n",
    ")\n",
    "resp.raise_for_status()\n",
    "for m in resp.json().get(\"memories\", []):\n",
    "    print(f\"- [{m.get('memory_type')}] {m.get('text')}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe6912c0",
   "metadata": {},
   "source": [
    "## Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d639c09a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nat-agent-memory\n",
      "nat-redis\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_SKIP\n",
    "# `|| true` keeps this idempotent — re-running when the containers are already\n",
    "# gone shouldn't raise.\n",
    "sh('docker rm -f nat-agent-memory nat-redis 2>/dev/null || true')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "We gave a NeMo Agent Toolkit agent persistent memory with **~30 lines of config and no memory plumbing** in the agent itself. The self-hosted [Redis Agent Memory Server](https://github.com/redis/agent-memory-server) handled extraction, storage, and semantic recall.\n",
    "\n",
    "Ready for production? The next notebook, `07_nemo_agent_toolkit_redis_cloud.ipynb`, runs the same agent against **managed [Redis Cloud Agent Memory](https://redis.io/agent-memory/)** — no server to operate."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "redis-ai-res",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}