Map-reduce¶
Review¶
We're building up to a multi-agent research assistant that ties together all of the modules from this course.
To build this multi-agent assistant, we've been introducing a few LangGraph controllability topics.
We just covered parallelization and sub-graphs.
Goals¶
Now, we're going to cover map reduce.
from dotenv import load_dotenv
load_dotenv()
%%capture --no-stderr
%pip install -U langchain-deepseek langgraph
import os
os.environ["LANGSMITH_TRACING"] = "true"
os.environ["LANGSMITH_PROJECT"] = "langchain-academy"
Problem¶
Map-reduce operations are essential for efficient task decomposition and parallel processing.
It has two phases:
(1) Map - Break a task into smaller sub-tasks, processing each sub-task in parallel.
(2) Reduce - Aggregate the results across all of the completed, parallelized sub-tasks.
Let's design a system that will do two things:
(1) Map - Create a set of jokes about a topic.
(2) Reduce - Pick the best joke from the list.
We'll use an LLM to do the job generation and selection.
from langchain_deepseek import ChatDeepSeek
# Prompts we will use
subjects_prompt = """Generate a list of 3 sub-topics that are all related to this overall topic: {topic}."""
joke_prompt = """Generate a joke about {subject}"""
best_joke_prompt = """Below are a bunch of jokes about {topic}. Select the best one! Return the ID of the best one, starting 0 as the ID for the first joke. Jokes: \n\n {jokes}"""
# LLM
model = ChatDeepSeek(model="deepseek-v4-pro", temperature=0)
State¶
Parallelizing joke generation¶
First, let's define the entry point of the graph that will:
- Take a user input topic
- Produce a list of joke topics from it
- Send each joke topic to our above joke generation node
Our state has a jokes key, which will accumulate jokes from parallelized joke generation
import operator
from typing import Annotated
from typing_extensions import TypedDict
from pydantic import BaseModel
class Subjects(BaseModel):
subjects: list[str]
class BestJoke(BaseModel):
id: int
class OverallState(TypedDict):
topic: str
subjects: list
jokes: Annotated[list, operator.add]
best_selected_joke: str
Generate subjects for jokes.
def generate_topics(state: OverallState):
prompt = subjects_prompt.format(topic=state["topic"])
response = model.with_structured_output(Subjects).invoke(prompt)
return {"subjects": response.subjects}
Here is the magic: we use the Send to create a joke for each subject.
This is very useful! It can automatically parallelize joke generation for any number of subjects.
generate_joke: the name of the node in the graph{"subject": s}: the state to send
Send allow you to pass any state that you want to generate_joke! It does not have to align with OverallState.
In this case, generate_joke is using its own internal state, and we can populate this via Send.
from langgraph.types import Send
def continue_to_jokes(state: OverallState):
return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
Joke generation (map)¶
Now, we just define a node that will create our jokes, generate_joke!
We write them back out to jokes in OverallState!
This key has a reducer that will combine lists.
class JokeState(TypedDict):
subject: str
class Joke(BaseModel):
joke: str
def generate_joke(state: JokeState):
prompt = joke_prompt.format(subject=state["subject"])
response = model.with_structured_output(Joke).invoke(prompt)
return {"jokes": [response.joke]}
Best joke selection (reduce)¶
Now, we add logic to pick the best joke.
def best_joke(state: OverallState):
jokes = "\n\n".join(state["jokes"])
prompt = best_joke_prompt.format(topic=state["topic"], jokes=jokes)
response = model.with_structured_output(BestJoke).invoke(prompt)
return {"best_selected_joke": state["jokes"][response.id]}
Compile¶
from IPython.display import Image
from langgraph.graph import END, StateGraph, START
# Construct the graph: here we put everything together to construct our graph
graph = StateGraph(OverallState)
graph.add_node("generate_topics", generate_topics)
graph.add_node("generate_joke", generate_joke)
graph.add_node("best_joke", best_joke)
graph.add_edge(START, "generate_topics")
graph.add_conditional_edges("generate_topics", continue_to_jokes, ["generate_joke"])
graph.add_edge("generate_joke", "best_joke")
graph.add_edge("best_joke", END)
# Compile the graph
app = graph.compile()
Image(app.get_graph().draw_mermaid_png())
# Call the graph: here we call it to generate a list of jokes
for s in app.stream({"topic": "animals"}):
print(s)
{'generate_topics': {'subjects': ['mammals', 'reptiles', 'birds']}}
{'generate_joke': {'jokes': ["Why don't mammals ever get lost? Because they always follow their 'instincts'!"]}}
{'generate_joke': {'jokes': ["Why don't alligators like fast food? Because they can't catch it!"]}}
{'generate_joke': {'jokes': ["Why do birds fly south for the winter? Because it's too far to walk!"]}}
{'best_joke': {'best_selected_joke': "Why don't alligators like fast food? Because they can't catch it!"}}
Studio¶
⚠️ Notice
Since filming these videos, we've updated Studio so that it can now be run locally and accessed through your browser. This is the preferred way to run Studio instead of using the Desktop App shown in the video. It is now called LangSmith Studio instead of LangGraph Studio. Detailed setup instructions are available in the "Getting Setup" guide at the start of the course. You can find a description of Studio here, and specific details for local deployment here.
To start the local development server, run the following command in your terminal in the /studio directory in this module:
langgraph dev
You should see the following output:
- 🚀 API: http://127.0.0.1:2024
- 🎨 Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024
- 📚 API Docs: http://127.0.0.1:2024/docs
Open your browser and navigate to the Studio UI URL shown above.
Let's load the above graph in the Studio UI, which uses module-4/studio/map_reduce.py set in module-4/studio/langgraph.json.