How to wait sink node after multiple Send to same node

[azidanit]

code:

import operator
from typing import Annotated
from typing_extensions import TypedDict

from langgraph.types import Send, Command
from langgraph.graph import END, StateGraph, START

from modules.graph.component.state.reducers import use_right


# This will be the overall state of the main graph.
# It will contain a topic (which we expect the user to provide)
# and then will generate a list of subjects, and then a joke for
# each subject
class OverallState(TypedDict):
    topic: str
    subjects: Annotated[list, operator.add]
    # Notice here we use the operator.add
    # This is because we want combine all the jokes we generate
    # from individual nodes back into one list - this is essentially
    # the "reduce" part
    jokes: Annotated[list, operator.add]
    best_selected_joke: Annotated[str, use_right]

    # for individual Send
    subject: Annotated[str, use_right]


# This is the function we will use to generate the subjects of the jokes.
# In general the length of the list generated by this node could vary each run.
def generate_topics(state: OverallState):
    # Simulate a LLM.
    return {"subjects": ["penguins", "lions", "elephants"]}


called_get_joke_tool = set()
called_generate_joke = set()
# Here we generate a joke, given a subject
def generate_joke(state: OverallState):
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)
    if state["subject"] not in called_get_joke_tool:
        print("from generate_joke calling get_joke: ", state)
        called_get_joke_tool.add(state["subject"])
        return state
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)


    print("from generate_joke: ", state)

    called_generate_joke.add(state["subject"])
    # Simulate a LLM.
    joke_map = {
        "lions": "Why don't lions like fast food? Because they can't catch it!",
        "elephants": "Why don't elephants use computers? They're afraid of the mouse!",
        "penguins": (
            "Why don't penguins like talking to strangers at parties? "
            "Because they find it hard to break the ice."
        ),
    }
    resp = {"jokes": [joke_map[state["subject"]]], "subject": state["subject"]}
    print("final from generate_joke: ", resp)
    return resp


def get_joke_tool(state: OverallState):
    # Dummy node Simulate a tool node.
    print("get_joke_tool called: ", state)
    return {}


# Here we define the logic to map out over the generated subjects
# We will use this as an edge in the graph
def continue_to_jokes(state: OverallState):
    return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]


def continue_to_get_joke_tool(state: OverallState):
    if state["subject"] not in called_generate_joke:
        return [Send("get_joke_tool", {"subject": state["subject"]})]

    return END


# Here we will judge the best joke
def best_joke(state: OverallState):
    print("BEST joke:", state)
    return {"best_selected_joke": "penguins"}


# Construct the graph: here we put everything together to construct our graph
builder = StateGraph(OverallState)
builder.add_node("generate_topics", generate_topics)
builder.add_node("generate_joke", generate_joke)
builder.add_node("get_joke_tool", get_joke_tool)
builder.add_node("best_joke", best_joke, defer=True)
builder.add_edge(START, "generate_topics")
builder.add_conditional_edges("generate_topics", continue_to_jokes, ["generate_joke"])
builder.add_conditional_edges("generate_joke", continue_to_get_joke_tool, {
    "get_joke_tool": "get_joke_tool",
    END: "best_joke"
})
builder.add_edge("get_joke_tool", "generate_joke")
builder.add_edge("best_joke", END)
graph = builder.compile()

# Call the graph: here we call it to generate a list of jokes
for step in graph.stream({"topic": "animals"}):
    print("=======STEP", step)

---
config:
  flowchart:
    curve: linear
---
graph TD;
	__start__([<p>__start__</p>]):::first
	generate_topics(generate_topics)
	generate_joke(generate_joke)
	get_joke_tool(get_joke_tool)
	best_joke(best_joke)
	__end__([<p>__end__</p>]):::last
	__start__ --> generate_topics;
	generate_joke -.-> best_joke;
	generate_joke -.-> get_joke_tool;
	generate_topics -.-> generate_joke;
	get_joke_tool --> generate_joke;
	best_joke --> __end__;
	classDef default fill:#f2f0ff,line-height:1.2
	classDef first fill-opacity:0
	classDef last fill:#bfb6fc

Loading

Output:

=======STEP {'generate_topics': {'subjects': ['penguins', 'lions', 'elephants']}}
from generate_joke calling get_joke:  {'subject': 'penguins'}
from generate_joke calling get_joke:  {'subject': 'lions'}
from generate_joke calling get_joke:  {'subject': 'elephants'}
=======STEP {'generate_joke': {'subject': 'penguins'}}
=======STEP {'generate_joke': {'subject': 'elephants'}}
=======STEP {'generate_joke': {'subject': 'lions'}}
get_joke_tool called:  {'subject': 'penguins'}
get_joke_tool called:  {'subject': 'lions'}
get_joke_tool called:  {'subject': 'elephants'}
=======STEP {'get_joke_tool': None}
=======STEP {'get_joke_tool': None}
=======STEP {'get_joke_tool': None}
from generate_joke:  {'topic': 'animals', 'subjects': ['penguins', 'lions', 'elephants'], 'jokes': [], 'best_selected_joke': '', 'subject': 'elephants'}
final from generate_joke:  {'jokes': ["Why don't elephants use computers? They're afraid of the mouse!"], 'subject': 'elephants'}
=======STEP {'generate_joke': {'jokes': ["Why don't elephants use computers? They're afraid of the mouse!"], 'subject': 'elephants'}}
BEST joke: {'topic': 'animals', 'subjects': ['penguins', 'lions', 'elephants'], 'jokes': ["Why don't elephants use computers? They're afraid of the mouse!"], 'best_selected_joke': '', 'subject': 'elephants'}
=======STEP {'best_joke': {'best_selected_joke': 'penguins'}}

As we can see in here, the best_joke node directly executed after 1 generate_joke finished (even tho there is other 2 that not yet finish). How I can make best_joke wait until all generate_joke finish?

Need your help in here. Thanks a lot!

[azidanit]

If I comment this part to call another node within generate_joke node it works as expected

def generate_joke(state: OverallState):
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)
    # if state["subject"] not in called_get_joke_tool:
    #     print("from generate_joke calling get_joke: ", state)
    #     called_get_joke_tool.add(state["subject"])
    #     return state
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)
    # ===== SIMULATE CALLING SOME RANDOM NODE (tool)
...

giving me output like this:

=======STEP {'generate_topics': {'subjects': ['penguins', 'lions', 'elephants']}}
from generate_joke:  {'subject': 'penguins'}
final from generate_joke:  {'jokes': ["Why don't penguins like talking to strangers at parties? Because they find it hard to break the ice."], 'subject': 'penguins'}
from generate_joke:  {'subject': 'lions'}
final from generate_joke:  {'jokes': ["Why don't lions like fast food? Because they can't catch it!"], 'subject': 'lions'}
from generate_joke:  {'subject': 'elephants'}
final from generate_joke:  {'jokes': ["Why don't elephants use computers? They're afraid of the mouse!"], 'subject': 'elephants'}
=======STEP {'generate_joke': {'jokes': ["Why don't penguins like talking to strangers at parties? Because they find it hard to break the ice."], 'subject': 'penguins'}}
=======STEP {'generate_joke': {'jokes': ["Why don't lions like fast food? Because they can't catch it!"], 'subject': 'lions'}}
=======STEP {'generate_joke': {'jokes': ["Why don't elephants use computers? They're afraid of the mouse!"], 'subject': 'elephants'}}
BEST joke: {'topic': 'animals', 'subjects': ['penguins', 'lions', 'elephants'], 'jokes': ["Why don't penguins like talking to strangers at parties? Because they find it hard to break the ice.", "Why don't lions like fast food? Because they can't catch it!", "Why don't elephants use computers? They're afraid of the mouse!"], 'best_selected_joke': '', 'subject': 'elephants'}
=======STEP {'best_joke': {'best_selected_joke': 'penguins'}}

where the best_joke can get all the jokes value