adispatch_custom_event doesn't work with LangGraph

[lucasalvarezlacasa]

Hello. I'm trying to dispatch custom events from certain nodes of my graph.

I'm currently using this code, but custom events are never shown. Any ideas why?

class State(MessagesState):
    docs: list[str]

async def retrieve_docs(state: State):
    await adispatch_custom_event("custom_retrieve_docs", 0)
    await asyncio.sleep(1)
    await adispatch_custom_event("custom_retrieve_docs", 1)
    return {"docs": ["doc2", "doc1", "doc3"]}

async def sort_docs(state: State):
    await adispatch_custom_event("custom_sort_docs", 0)
    await asyncio.sleep(1)
    await adispatch_custom_event("custom_sort_docs", 1)
    return {"docs": ["doc1", "doc2", "doc3"]}

async def call_model(state: MessagesState):
    response = await llm.ainvoke(state['messages'])
    return {"messages": response}


workflow = StateGraph(MessagesState)

workflow.add_node(retrieve_docs)
workflow.add_node(sort_docs)
workflow.add_node(call_model)

workflow.add_edge(START, "retrieve_docs")
workflow.add_edge("retrieve_docs", "sort_docs")
workflow.add_edge("sort_docs", "call_model")
workflow.add_edge("call_model", END)
compiled_graph = workflow.compile()
display(Image(compiled_graph.get_graph().draw_mermaid_png()))

inputs = [{"role": "user", "content": "hi!"}]
async for event in compiled_graph.astream_events({"messages": inputs}, version="v1"):
    kind = event["event"]
    print(f"{kind}: {event['name']}")

[vbarda]

i believe you need to use version="v2"

[lucasalvarezlacasa]

Ahh very nice, I didn't realize this. Thank you!