{
 "cells": [
  {
   "attachments": {
    "51f122de-b2ce-4c21-a5a7-c3be70c28a91.png": {
     "image/png": "[encoded data, 52544 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "710dc4f0-1c88-4386-9e9d-fec3de6bb774",
   "metadata": {},
   "source": [
    "# How to create branches for parallel node execution\n",
    "\n",
    "Parallel execution of nodes is essential to speed up overall graph operation. LangGraph offers native support for parallel execution of nodes, which can significantly enhance the performance of graph-based workflows. This parallelization is achieved through fan-out and fan-in mechanisms, utilizing both standard edges and [conditional_edges](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.MessageGraph.add_conditional_edges). Below are some examples showing how to add create branching dataflows that work for you. \n",
    "\n",
    "![Screenshot 2024-07-09 at 2.55.56 PM.png](attachment:51f122de-b2ce-4c21-a5a7-c3be70c28a91.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bb54e2d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6c05fc4-ecd8-483f-a9fd-b1a055f922d9",
   "metadata": {},
   "source": [
    "## Parallel node fan-out and fan-in\n",
    "\n",
    "In this example, we fan out from `Node A` to `B and C` and then fan in to `D`. With our state, [we specify the reducer add operation](https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers). This will combine or accumulate values for the specific key in the State, rather than simply overwriting the existing value. For lists, this means concatenating the new list with the existing list. \n",
    "\n",
    "Note that LangGraph uses `Annotated` type to specify reducer functions for specific keys in the State: it maintains the original type (`list`) for type checking, but allows attaching the reducer function (`add`) to the type without changing the type itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "09372b8b-edea-4b9d-9ec3-3d93ce1ba819",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, Any\n",
    "\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    # The operator.add reducer fn makes this append-only\n",
    "    aggregate: Annotated[list, operator.add]\n",
    "\n",
    "\n",
    "class ReturnNodeValue:\n",
    "    def __init__(self, node_secret: str):\n",
    "        self._value = node_secret\n",
    "\n",
    "    def __call__(self, state: State) -> Any:\n",
    "        print(f\"Adding {self._value} to {state['aggregate']}\")\n",
    "        return {\"aggregate\": [self._value]}\n",
    "\n",
    "\n",
    "builder = StateGraph(State)\n",
    "builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
    "builder.add_edge(START, \"a\")\n",
    "builder.add_node(\"b\", ReturnNodeValue(\"I'm B\"))\n",
    "builder.add_node(\"c\", ReturnNodeValue(\"I'm C\"))\n",
    "builder.add_node(\"d\", ReturnNodeValue(\"I'm D\"))\n",
    "builder.add_edge(\"a\", \"b\")\n",
    "builder.add_edge(\"a\", \"c\")\n",
    "builder.add_edge(\"b\", \"d\")\n",
    "builder.add_edge(\"c\", \"d\")\n",
    "builder.add_edge(\"d\", END)\n",
    "graph = builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "66f52a20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 12056 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74dd577b-0474-44c4-b4bc-9113090e3121",
   "metadata": {},
   "source": [
    "With the reducer, you can see that the values added in each node are accumulated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "38846b01",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding I'm A to []\n",
      "Adding I'm B to [\"I'm A\"]\n",
      "Adding I'm C to [\"I'm A\"]\n",
      "Adding I'm D to [\"I'm A\", \"I'm B\", \"I'm C\"]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm D\"]}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.invoke({\"aggregate\": []}, {\"configurable\": {\"thread_id\": \"foo\"}})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c392b3d2",
   "metadata": {},
   "source": [
    "<details class=\"note\"> <summary>Exception handling?</summary>\n",
    " <p>LangGraph executes nodes within <a href=\"https://langchain-ai.github.io/langgraph/concepts/#core-design\">\"supersteps\"</a>, meaning that while parallel branches are executed in parallel, the entire superstep is <b>transactional</b>. If any of these branches raises an exception, <b>none</b> of the updates are applied to the state (the entire superstep errors).<br><br>\n",
    " If you have error-prone (perhaps want to handle flakey API calls), LangGraph provides two ways to address this:<br>\n",
    " <ol>\n",
    "  <li>You can write regular python code within your node to catch and handle exceptions.</li>\n",
    "  <li>You can set a <b><a href=\"https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.graph.CompiledGraph.retry_policy\">retry_policy</a></b> to direct the graph to retry nodes that raise certain types of exceptions. Only failing branches are retried, so you needn't worry about performing redundant work.</li>\n",
    "</ol></p>\n",
    "Together, these let you perform parallel execution and fully control exception handling.\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08d8162e-1785-4ae1-993f-6d2ed48c22ae",
   "metadata": {},
   "source": [
    "## Parallel node fan-out and fan-in with extra steps\n",
    "\n",
    "The above example showed how to fan-out and fan-in when each path was only one step. But what if one path had more than one step?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "259a7704-5aa0-4e4c-aeef-cca04e8be0ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated\n",
    "\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import StateGraph\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    # The operator.add reducer fn makes this append-only\n",
    "    aggregate: Annotated[list, operator.add]\n",
    "\n",
    "\n",
    "builder = StateGraph(State)\n",
    "builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
    "builder.add_edge(START, \"a\")\n",
    "builder.add_node(\"b\", ReturnNodeValue(\"I'm B\"))\n",
    "builder.add_node(\"b2\", ReturnNodeValue(\"I'm B2\"))\n",
    "builder.add_node(\"c\", ReturnNodeValue(\"I'm C\"))\n",
    "builder.add_node(\"d\", ReturnNodeValue(\"I'm D\"))\n",
    "builder.add_edge(\"a\", \"b\")\n",
    "builder.add_edge(\"a\", \"c\")\n",
    "builder.add_edge(\"b\", \"b2\")\n",
    "builder.add_edge([\"b2\", \"c\"], \"d\")\n",
    "builder.add_edge(\"d\", END)\n",
    "graph = builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "83320227-8ab3-44c0-b6cf-064a7a425b9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 13592 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3f971fa3-29e4-466f-a85e-2863bfecf7fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding I'm A to []\n",
      "Adding I'm B to [\"I'm A\"]\n",
      "Adding I'm C to [\"I'm A\"]\n",
      "Adding I'm B2 to [\"I'm A\", \"I'm B\", \"I'm C\"]\n",
      "Adding I'm D to [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm B2\"]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm B2\", \"I'm D\"]}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.invoke({\"aggregate\": []})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d45f4477",
   "metadata": {},
   "source": [
    "## Conditional Branching\n",
    "\n",
    "If your fan-out is not deterministic, you can use [add_conditional_edges](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.StateGraph.add_conditional_edges) directly.\n",
    "\n",
    "If you have a known \"sink\" node that the conditional branches will route to afterwards, you can provide `then=<final-node-name>` when creating the conditional edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "95f5e026",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, Sequence\n",
    "\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import END, START, StateGraph\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    # The operator.add reducer fn makes this append-only\n",
    "    aggregate: Annotated[list, operator.add]\n",
    "    which: str\n",
    "\n",
    "\n",
    "builder = StateGraph(State)\n",
    "builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
    "builder.add_edge(START, \"a\")\n",
    "builder.add_node(\"b\", ReturnNodeValue(\"I'm B\"))\n",
    "builder.add_node(\"c\", ReturnNodeValue(\"I'm C\"))\n",
    "builder.add_node(\"d\", ReturnNodeValue(\"I'm D\"))\n",
    "builder.add_node(\"e\", ReturnNodeValue(\"I'm E\"))\n",
    "\n",
    "\n",
    "def route_bc_or_cd(state: State) -> Sequence[str]:\n",
    "    if state[\"which\"] == \"cd\":\n",
    "        return [\"c\", \"d\"]\n",
    "    return [\"b\", \"c\"]\n",
    "\n",
    "\n",
    "intermediates = [\"b\", \"c\", \"d\"]\n",
    "builder.add_conditional_edges(\n",
    "    \"a\",\n",
    "    route_bc_or_cd,\n",
    "    intermediates,\n",
    ")\n",
    "for node in intermediates:\n",
    "    builder.add_edge(node, \"e\")\n",
    "\n",
    "\n",
    "builder.add_edge(\"e\", END)\n",
    "graph = builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1d0e6c56",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 15180 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7134f652",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding I'm A to []\n",
      "Adding I'm B to [\"I'm A\"]\n",
      "Adding I'm C to [\"I'm A\"]\n",
      "Adding I'm E to [\"I'm A\", \"I'm B\", \"I'm C\"]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'aggregate': [\"I'm A\", \"I'm B\", \"I'm C\", \"I'm E\"], 'which': 'bc'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.invoke({\"aggregate\": [], \"which\": \"bc\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b130e694",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding I'm A to []\n",
      "Adding I'm D to [\"I'm A\"]\n",
      "Adding I'm C to [\"I'm A\"]\n",
      "Adding I'm E to [\"I'm A\", \"I'm C\", \"I'm D\"]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'aggregate': [\"I'm A\", \"I'm C\", \"I'm D\", \"I'm E\"], 'which': 'cd'}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.invoke({\"aggregate\": [], \"which\": \"cd\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "952cd6f3",
   "metadata": {},
   "source": [
    "## Stable Sorting\n",
    "\n",
    "When fanned out, nodes are run in parallel as a single \"superstep\". The updates from each superstep are all applied to the state in sequence once the superstep has completed. \n",
    "\n",
    "If you need consistent, predetermined ordering of updates from a parallel superstep, you should write the outputs (along with an identifying key) to a separate field in your state, then combine them in the \"sink\" node by adding regular `edge`'s from each of the fanout nodes to the rendezvous point.\n",
    "\n",
    "For instance, suppose I want to order the outputs of the parallel step by \"reliability\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "836bc12d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, Sequence\n",
    "\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import StateGraph\n",
    "\n",
    "\n",
    "def reduce_fanouts(left, right):\n",
    "    if left is None:\n",
    "        left = []\n",
    "    if not right:\n",
    "        # Overwrite\n",
    "        return []\n",
    "    return left + right\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    # The operator.add reducer fn makes this append-only\n",
    "    aggregate: Annotated[list, operator.add]\n",
    "    fanout_values: Annotated[list, reduce_fanouts]\n",
    "    which: str\n",
    "\n",
    "\n",
    "builder = StateGraph(State)\n",
    "builder.add_node(\"a\", ReturnNodeValue(\"I'm A\"))\n",
    "builder.add_edge(START, \"a\")\n",
    "\n",
    "\n",
    "class ParallelReturnNodeValue:\n",
    "    def __init__(\n",
    "        self,\n",
    "        node_secret: str,\n",
    "        reliability: float,\n",
    "    ):\n",
    "        self._value = node_secret\n",
    "        self._reliability = reliability\n",
    "\n",
    "    def __call__(self, state: State) -> Any:\n",
    "        print(f\"Adding {self._value} to {state['aggregate']} in parallel.\")\n",
    "        return {\n",
    "            \"fanout_values\": [\n",
    "                {\n",
    "                    \"value\": [self._value],\n",
    "                    \"reliability\": self._reliability,\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "\n",
    "\n",
    "builder.add_node(\"b\", ParallelReturnNodeValue(\"I'm B\", reliability=0.9))\n",
    "\n",
    "builder.add_node(\"c\", ParallelReturnNodeValue(\"I'm C\", reliability=0.1))\n",
    "builder.add_node(\"d\", ParallelReturnNodeValue(\"I'm D\", reliability=0.3))\n",
    "\n",
    "\n",
    "def aggregate_fanout_values(state: State) -> Any:\n",
    "    # Sort by reliability\n",
    "    ranked_values = sorted(\n",
    "        state[\"fanout_values\"], key=lambda x: x[\"reliability\"], reverse=True\n",
    "    )\n",
    "    return {\n",
    "        \"aggregate\": [x[\"value\"] for x in ranked_values] + [\"I'm E\"],\n",
    "        \"fanout_values\": [],\n",
    "    }\n",
    "\n",
    "\n",
    "builder.add_node(\"e\", aggregate_fanout_values)\n",
    "\n",
    "\n",
    "def route_bc_or_cd(state: State) -> Sequence[str]:\n",
    "    if state[\"which\"] == \"cd\":\n",
    "        return [\"c\", \"d\"]\n",
    "    return [\"b\", \"c\"]\n",
    "\n",
    "\n",
    "intermediates = [\"b\", \"c\", \"d\"]\n",
    "builder.add_conditional_edges(\"a\", route_bc_or_cd, intermediates)\n",
    "\n",
    "for node in intermediates:\n",
    "    builder.add_edge(node, \"e\")\n",
    "\n",
    "builder.add_edge(\"e\", END)\n",
    "graph = builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "932c497e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 15180 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "933b3afd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding I'm A to []\n",
      "Adding I'm B to [\"I'm A\"] in parallel.\n",
      "Adding I'm C to [\"I'm A\"] in parallel.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'aggregate': [\"I'm A\", [\"I'm B\"], [\"I'm C\"], \"I'm E\"],\n",
       " 'fanout_values': [],\n",
       " 'which': 'bc'}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.invoke({\"aggregate\": [], \"which\": \"bc\", \"fanout_values\": []})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e30531bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding I'm A to []\n",
      "Adding I'm C to [\"I'm A\"] in parallel.\n",
      "Adding I'm D to [\"I'm A\"] in parallel.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'aggregate': [\"I'm A\", [\"I'm D\"], [\"I'm C\"], \"I'm E\"],\n",
       " 'fanout_values': [],\n",
       " 'which': 'cd'}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.invoke({\"aggregate\": [], \"which\": \"cd\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
