diff --git a/examples/notebooks/Customizing memory management.ipynb b/examples/notebooks/Customizing memory management.ipynb
index af167dce..95d3ed9a 100644
--- a/examples/notebooks/Customizing memory management.ipynb	
+++ b/examples/notebooks/Customizing memory management.ipynb	
@@ -253,15 +253,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "id": "7808912f-831b-4cdc-8606-40052eb809b4",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from typing import Optional, List\n",
+    "from typing import Optional, List, TYPE_CHECKING\n",
     "import json\n",
     "\n",
-    "def task_queue_push(self: \"Agent\", task_description: str):\n",
+    "if TYPE_CHECKING:\n",
+    "    from letta import AgentState\n",
+    "\n",
+    "def task_queue_push(agent_state: \"AgentState\", task_description: str):\n",
     "    \"\"\"\n",
     "    Push to a task queue stored in core memory. \n",
     "\n",
@@ -273,12 +276,12 @@
     "        does not produce a response.\n",
     "    \"\"\"\n",
     "    import json\n",
-    "    tasks = json.loads(self.memory.get_block(\"tasks\").value)\n",
+    "    tasks = json.loads(agent_state.memory.get_block(\"tasks\").value)\n",
     "    tasks.append(task_description)\n",
-    "    self.memory.update_block_value(\"tasks\", json.dumps(tasks))\n",
+    "    agent_state.memory.update_block_value(\"tasks\", json.dumps(tasks))\n",
     "    return None\n",
     "\n",
-    "def task_queue_pop(self: \"Agent\"):\n",
+    "def task_queue_pop(agent_state: \"AgentState\"):\n",
     "    \"\"\"\n",
     "    Get the next task from the task queue \n",
     "\n",
@@ -288,12 +291,12 @@
     "        None (the task queue is empty)\n",
     "    \"\"\"\n",
     "    import json\n",
-    "    tasks = json.loads(self.memory.get_block(\"tasks\").value)\n",
+    "    tasks = json.loads(agent_state.memory.get_block(\"tasks\").value)\n",
     "    if len(tasks) == 0: \n",
     "        return None\n",
     "    task = tasks[0]\n",
     "    print(\"CURRENT TASKS: \", tasks)\n",
-    "    self.memory.update_block_value(\"tasks\", json.dumps(tasks[1:]))\n",
+    "    agent_state.memory.update_block_value(\"tasks\", json.dumps(tasks[1:]))\n",
     "    return task\n",
     "\n",
     "push_task_tool = client.tools.upsert_from_function(func=task_queue_push)\n",
@@ -310,7 +313,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "id": "135fcf3e-59c4-4da3-b86b-dbffb21aa343",
    "metadata": {},
    "outputs": [],
@@ -340,6 +343,8 @@
     "        ),\n",
     "    ],\n",
     "    tool_ids=[push_task_tool.id, pop_task_tool.id],\n",
+    "    model=\"letta/letta-free\",\n",
+    "    embedding=\"letta/letta-free\",\n",
     ")"
    ]
   },