From 03ba50356949025d681786cd4a7038dafd6604e7 Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 5 Jun 2024 07:40:58 -0700
Subject: [PATCH] Rm END check in STORM (#600)

---
 examples/storm/storm.ipynb | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/examples/storm/storm.ipynb b/examples/storm/storm.ipynb
index cd13007ad..32c060627 100644
--- a/examples/storm/storm.ipynb
+++ b/examples/storm/storm.ipynb
@@ -1489,6 +1489,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from langgraph.checkpoint.memory import MemorySaver\n",
+    "\n",
     "builder_of_storm = StateGraph(ResearchState)\n",
     "\n",
     "nodes = [\n",
@@ -1507,7 +1509,7 @@
     "\n",
     "builder_of_storm.set_entry_point(nodes[0][0])\n",
     "builder_of_storm.set_finish_point(nodes[-1][0])\n",
-    "storm = builder_of_storm.compile()"
+    "storm = builder_of_storm.compile(checkpointer=MemorySaver())"
    ]
   },
   {
@@ -1558,16 +1560,16 @@
     }
    ],
    "source": [
+    "config = {\"configurable\": {\"thread_id\": \"my-thread\"}}\n",
     "async for step in storm.astream(\n",
     "    {\n",
     "        \"topic\": \"Groq, NVIDIA, Llamma.cpp and the future of LLM Inference\",\n",
-    "    }\n",
+    "    },\n",
+    "    config\n",
     "):\n",
     "    name = next(iter(step))\n",
     "    print(name)\n",
-    "    print(\"-- \", str(step[name])[:300])\n",
-    "    if END in step:\n",
-    "        results = step"
+    "    print(\"-- \", str(step[name])[:300])"
    ]
   },
   {
@@ -1576,7 +1578,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "article = results[END][\"article\"]"
+    "checkpoint = storm.get_state(config)\n",
+    "article = checkpoint.values[\"article\"]"
    ]
   },
   {
@@ -1687,7 +1690,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,