diff --git a/examples/code_assistant/langgraph_code_assistant.ipynb b/examples/code_assistant/langgraph_code_assistant.ipynb
index 79450d604..37bc90db9 100644
--- a/examples/code_assistant/langgraph_code_assistant.ipynb
+++ b/examples/code_assistant/langgraph_code_assistant.ipynb
@@ -81,7 +81,9 @@
     "\n",
     "### Code solution\n",
     "\n",
-    "Try OpenAI and [Claude3](https://docs.anthropic.com/claude/docs/models-overview) with function calling."
+    "Try OpenAI and [Claude3](https://docs.anthropic.com/claude/docs/models-overview) with function calling.\n",
+    "\n",
+    "Create `code_gen_chain` w/ either OpenAI or Claude and test here."
    ]
   },
   {
@@ -179,7 +181,6 @@
     "\n",
     "structured_llm_claude = llm.with_structured_output(code, include_raw=True)\n",
     "\n",
-    "\n",
     "# Optional: Check for errors in case tool use is flaky\n",
     "def check_claude_output(tool_output):\n",
     "    \"\"\"Check for parse error or failure to call the tool\"\"\"\n",
@@ -202,13 +203,11 @@
     "        )\n",
     "    return tool_output\n",
     "\n",
-    "\n",
     "# Chain with output check\n",
     "code_chain_claude_raw = (\n",
     "    code_gen_prompt_claude | structured_llm_claude | check_claude_output\n",
     ")\n",
     "\n",
-    "\n",
     "def insert_errors(inputs):\n",
     "    \"\"\"Insert errors for tool parsing in the messages\"\"\"\n",
     "\n",
@@ -241,10 +240,7 @@
     "\n",
     "    return solution[\"parsed\"]\n",
     "\n",
-    "\n",
-    "# With re-try to correct for failure to invoke tool\n",
-    "# TODO: Annoying errors w/ \"user\" vs \"assistant\"\n",
-    "# Roles must alternate between \"user\" and \"assistant\", but found multiple \"user\" roles in a row\n",
+    "# Optional: With re-try to correct for failure to invoke tool\n",
     "code_gen_chain = code_gen_chain_re_try | parse_output\n",
     "\n",
     "# No re-try\n",
@@ -651,7 +647,7 @@
     "    solution = code_gen_chain.invoke(\n",
     "        {\"context\": concatenated_content, \"messages\": [(\"user\", example[\"question\"])]}\n",
     "    )\n",
-    "    solution_structured = structured_code_formatter.invoke([(\"code\", solution)])\n",
+    "    solution_structured = code_gen_chain.invoke([(\"code\", solution)])\n",
     "    return {\"imports\": solution_structured.imports, \"code\": solution_structured.code}\n",
     "\n",
     "\n",
@@ -758,7 +754,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,