map_reduce.py

import operator
from typing import Annotated
from typing_extensions import TypedDict

from pydantic import BaseModel

from langchain_openai import ChatOpenAI 

from langgraph.constants import Send
from langgraph.graph import END, StateGraph, START
from dotenv import load_dotenv
from utils import save_graph

load_dotenv()

# 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 = ChatOpenAI(model="gpt-4o", temperature=0) 

# Define the state
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

def generate_topics(state: OverallState)-> Subjects:
    prompt = subjects_prompt.format(topic=state["topic"])
    response = model.with_structured_output(Subjects).invoke(prompt)
    return {"subjects": response.subjects}

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]}

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]}

def continue_to_jokes(state: OverallState):
    return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]

# Construct the graph: here we put everything together to construct our graph
graph_builder = StateGraph(OverallState)
graph_builder.add_node("generate_topics", generate_topics)
graph_builder.add_node("generate_joke", generate_joke)
graph_builder.add_node("best_joke", best_joke)
graph_builder.add_edge(START, "generate_topics")
# CRITICAL PART !
graph_builder.add_conditional_edges("generate_topics", continue_to_jokes, ["generate_joke"])
graph_builder.add_edge("generate_joke", "best_joke")

graph_builder.add_edge("best_joke", END)

# Compile the graph
graph = graph_builder.compile()


def main():
    """
    Executes the map-reduce joke generation pipeline:
    - Generates sub-topics based on an overall topic.
    - For each sub-topic, generates a joke.
    - Selects the best joke from the generated jokes.
    """
    # Define the overall state with an example topic
    initial_state = {
        "topic": "technology",
        "subjects": [],
        "jokes": [],
        "best_selected_joke": ""
    }
    
    # Execute the graph pipeline
    result = graph.invoke(initial_state, debug=True)
    
    # Print the best selected joke from the pipeline
    print("\n=== Best Joke Selected ===")
    if "best_selected_joke" in result and result["best_selected_joke"]:
        print(result["best_selected_joke"])
    else:
        print("No joke was selected.")
        
    # Save the graph visualization
    save_graph(graph, "./images/map_reduce.png")
    
    
if __name__ == "__main__":
    main()