diff --git a/examples/code_assistant/lcel-teacher-langgraph.ipynb b/examples/code_assistant/lcel-teacher-langgraph.ipynb
index f0caf96a9..f4c7170bb 100644
--- a/examples/code_assistant/lcel-teacher-langgraph.ipynb
+++ b/examples/code_assistant/lcel-teacher-langgraph.ipynb
@@ -335,6 +335,7 @@
     "    question = state_dict[\"question\"]\n",
     "    docs = state_dict[\"docs\"]\n",
     "    code_solution = state_dict[\"generation\"]\n",
+    "    prefix = code_solution[0].prefix\n",
     "    imports = code_solution[0].imports\n",
     "    code = code_solution[0].code\n",
     "    code_block = imports +\"\\n\"+ code\n",
@@ -354,8 +355,13 @@
     "        # No errors occurred\n",
     "        error = \"None\"\n",
     "\n",
-    "    return {\"keys\": {\"generation\": code_solution, \"question\": question, \"error\": error, \"docs\": docs}}\n",
-    "\n",
+    "    return {\"keys\": {\"generation\": code_solution, \n",
+    "                     \"question\": question, \n",
+    "                     \"error\": error, \n",
+    "                     \"docs\": docs,\n",
+    "                     \"prefix\":prefix,\n",
+    "                     \"imports\":imports,\n",
+    "                     \"code\":code}}\n",
     "\n",
     "### Edges\n",
     "\n",
@@ -364,7 +370,7 @@
     "    Determines whether to test code execution, or re-try answer generation.\n",
     "\n",
     "    Args:\n",
-    "        state (dict): The current state of the agent, including all keys.\n",
+    "       state (dict): The current graph state\n",
     "\n",
     "    Returns:\n",
     "        str: Next node to call\n",
@@ -391,7 +397,7 @@
     "    Determines whether to finish.\n",
     "\n",
     "    Args:\n",
-    "        state (dict): The current state of the agent, including all keys.\n",
+    "        state (dict): The current graph state\n",
     "\n",
     "    Returns:\n",
     "        str: Next node to call\n",
@@ -480,30 +486,29 @@
     "@run_evaluator\n",
     "def check_import(run: Run, example: Union[Example, None] = None):\n",
     "    model_outputs = run.outputs[\"keys\"]\n",
-    "    code_block = model_outputs['generation'][0]['repr']\n",
-    "    print(\"CODE BLOCK!\")\n",
-    "    print(code_block)\n",
+    "    imports = model_outputs['imports']\n",
     "    try:\n",
-    "        exec(code_block.imports)\n",
+    "        exec(imports)\n",
     "        score = 1\n",
+    "        print(\"Score: 1!\")\n",
     "    except:\n",
     "        score = 0\n",
+    "        print(\"Score: 0!\")\n",
     "    return EvaluationResult(key=\"check_import\", score=score)\n",
     "\n",
     "@run_evaluator\n",
     "def check_execution(run: Run, example: Union[Example, None] = None):\n",
     "    model_outputs = run.outputs[\"keys\"]\n",
-    "    code_block = model_outputs['generation'][0]['repr']\n",
-    "    print(\"CODE BLOCK!\")\n",
-    "    print(code_block)\n",
-    "    imports = code_block.imports\n",
-    "    code = code_block.code\n",
+    "    imports = model_outputs['imports']\n",
+    "    code = model_outputs['code']\n",
     "    code_to_execute = imports +\"\\n\"+ code\n",
     "    try:\n",
     "        exec(code_to_execute)\n",
     "        score = 1\n",
+    "        print(\"Score: 1!\")\n",
     "    except:\n",
     "        score = 0\n",
+    "        print(\"Score: 0!\")\n",
     "    return EvaluationResult(key=\"check_execution\", score=score)"
    ]
   },