diff --git a/examples/map-reduce.ipynb b/examples/map-reduce.ipynb index 5262a1d93..073aa9071 100644 --- a/examples/map-reduce.ipynb +++ b/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",