{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "from typing import TypedDict\n",
    "from loguru import logger\n",
    "import time\n",
    "\n",
    "\n",
    "class OverallState(TypedDict):\n",
    "    initial_input: str\n",
    "    node_a1: str\n",
    "    node_a2: str\n",
    "    node_b: str\n",
    "\n",
    "\n",
    "def node_a1(state: OverallState) -> OverallState:\n",
    "    time.sleep(1)\n",
    "    logger.info(\"node_a1\")\n",
    "    return OverallState(node_a1=\"node_a1\")\n",
    "\n",
    "\n",
    "def node_a2(state: OverallState) -> OverallState:\n",
    "    time.sleep(3)\n",
    "    logger.info(\"node_a2\")\n",
    "\n",
    "    return OverallState(node_a2=\"node_a2\")\n",
    "\n",
    "\n",
    "def node_b(state: OverallState) -> OverallState:\n",
    "    time.sleep(1)\n",
    "    logger.info(\"node_b\")\n",
    "    return OverallState(node_b=\"node_b\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.checkpoint.memory import InMemorySaver\n",
    "\n",
    "builder = StateGraph(state_schema=OverallState)\n",
    "builder.add_node(\"node_a1\", node_a1)\n",
    "builder.add_node(\"node_a2\", node_a2)\n",
    "builder.add_node(\"node_b\", node_b)\n",
    "\n",
    "builder.add_edge(START, \"node_a1\")\n",
    "builder.add_edge(START, \"node_a2\")\n",
    "builder.add_edge([\"node_a1\", \"node_a2\"], \"node_b\")\n",
    "builder.add_edge(\"node_b\", END)\n",
    "\n",
    "checkpoint_saver = InMemorySaver()\n",
    "\n",
    "graph = builder.compile(checkpointer=checkpoint_saver)\n",
    "\n",
    "from IPython.display import display\n",
    "\n",
    "display(graph)"
   ],
   "id": "87e186e33e43ad5f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from rich import print as rp\n",
    "\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        \"thread_id\": \"12345\"\n",
    "    }\n",
    "}\n",
    "\n",
    "for chunk in graph.stream(\n",
    "        {\"initial_input\": \"init\"},\n",
    "        stream_mode=[\"checkpoints\"],\n",
    "        config=config\n",
    "):\n",
    "    print(\"=\" * 50)\n",
    "    # print(chunk)\n",
    "    rp(chunk)"
   ],
   "id": "80bb26c320de30c7",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}