cr

examples/map-reduce.ipynb

@@ -11,12 +11,12 @@
     "\n",
     "LangGraph supports this via the `Send` api. This can be used to allow a conditional edge to `Send` multiple different states to multiple nodes. The state it sends can be different from the state of the core graph.\n",
     "\n",
-    "Let's see what this looks like in action! We'll put together a toy example of generating a list of words, and then writing a joke about each word."
+    "Let's see what this looks like in action! We'll put together a toy example of generating a list of words, and then writing a joke about each word, and then judging what the best joke is."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 10,
    "id": "0f0f78e4-423d-4e2d-aa1a-01efaec4715f",
    "metadata": {},
    "outputs": [
@@ -24,12 +24,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'generate_topics': {'subjects': ['cat', 'dog', 'elephant', 'lion', 'tiger']}}\n",
-      "{'generate_joke': {'jokes': ['Why did the tiger lose at poker? Because he was playing with a cheetah!']}}\n",
-      "{'generate_joke': {'jokes': [\"Why don't elephants use computers? Because they're afraid of the mouse!\"]}}\n",
-      "{'generate_joke': {'jokes': ['Why did the lion eat the tightrope walker? He wanted a well-balanced meal!']}}\n",
+      "{'generate_topics': {'subjects': ['cat', 'dog', 'rabbit', 'hamster', 'bird']}}\n",
+      "{'generate_joke': {'jokes': ['Why did the rabbit go to the barber shop? Because it needed a hare cut!']}}\n",
       "{'generate_joke': {'jokes': ['Why was the cat sitting on the computer? Because it wanted to keep an eye on the mouse!']}}\n",
-      "{'generate_joke': {'jokes': [\"Why do dogs run in circles before lying down? Because they're trying to make a 'ruff' impression!\"]}}\n"
+      "{'generate_joke': {'jokes': ['Why did the hamster join the band? Because it had great drumming skills!']}}\n",
+      "{'generate_joke': {'jokes': [\"Why did the dog sit in the shade? Because he didn't want to be a hot dog!\"]}}\n",
+      "{'generate_joke': {'jokes': ['Why did the bird join a band? Because it had the best tweet-talent!']}}\n",
+      "{'best_joke': {'best_selected_joke': \"Why did the dog sit in the shade? Because he didn't want to be a hot dog!\"}}\n"
      ]
     }
    ],
@@ -45,6 +46,9 @@
     "# Define model and prompts we will use\n",
     "subjects_prompt = \"\"\"Generate a comma separated list of between 2 and 5 {topic}.\"\"\"\n",
     "joke_prompt = \"\"\"Generate a joke about {subject}\"\"\"\n",
+    "best_joke_prompt = \"\"\"Below are a bunch of jokes about {topic}. Select the best one! Return the ID of the best one.\n",
+    "\n",
+    "{jokes}\"\"\"\n",
     "\n",
     "\n",
     "class Subjects(BaseModel):\n",
@@ -55,6 +59,10 @@
     "    joke: str\n",
     "\n",
     "\n",
+    "class BestJoke(BaseModel):\n",
+    "    id: int\n",
+    "\n",
+    "\n",
     "model = ChatOpenAI()\n",
     "\n",
     "# Graph components: define the components that will make up the graph\n",
@@ -71,6 +79,7 @@
     "    # from individual nodes back into one list - this is essentially\n",
     "    # the \"reduce\" part\n",
     "    jokes: Annotated[list, operator.add]\n",
+    "    best_selected_joke: str\n",
     "\n",
     "\n",
     "# This will be the state of the node that we will \"map\" all\n",
@@ -102,13 +111,25 @@
     "    return [Send(\"generate_joke\", {\"subject\": s}) for s in state['subjects']]\n",
     "\n",
     "\n",
+    "# Here we will judge the best joke\n",
+    "def best_joke(state: OverallState):\n",
+    "    jokes = \"\\n\\n\".format(\"Joke {i}: {j}\" for i, j in enumerate(state['jokes']))\n",
+    "    prompt = best_joke_prompt.format(topic=state['topic'], jokes=jokes)\n",
+    "    response = model.with_structured_output(BestJoke).invoke(prompt)\n",
+    "    return {\"best_selected_joke\": state['jokes'][response.id]}\n",
+    "    \n",
+    "\n",
+    "\n",
+    "\n",
     "# Construct the graph: here we put everything together to construct our graph\n",
     "graph = StateGraph(OverallState)\n",
     "graph.add_node(\"generate_topics\", generate_topics)\n",
     "graph.add_node(\"generate_joke\", generate_joke)\n",
+    "graph.add_node(\"best_joke\", best_joke)\n",
     "graph.set_entry_point(\"generate_topics\")\n",
     "graph.add_conditional_edges(\"generate_topics\", continue_to_jokes)\n",
-    "graph.add_edge(\"generate_joke\", END)\n",
+    "graph.add_edge(\"generate_joke\", \"best_joke\")\n",
+    "graph.add_edge(\"best_joke\", END)\n",
     "app = graph.compile()\n",
     "\n",
     "\n",