diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py index 968394f17..06ec43a28 100644 --- a/libs/langgraph/bench/__main__.py +++ b/libs/langgraph/bench/__main__.py @@ -30,6 +30,26 @@ async def arun(graph: Pregel, input: dict): ) +async def arun_first_event_latency(graph: Pregel, input: dict) -> None: + """Latency for the first event. + + Run the graph until the first event is processed and then stop. + """ + stream = graph.astream( + input, + { + "configurable": {"thread_id": str(uuid4())}, + "recursion_limit": 1000000000, + }, + ) + + try: + async for _ in stream: + break + finally: + await stream.aclose() + + def run(graph: Pregel, input: dict): len( [ @@ -45,6 +65,26 @@ def run(graph: Pregel, input: dict): ) +def run_first_event_latency(graph: Pregel, input: dict) -> None: + """Latency for the first event. + + Run the graph until the first event is processed and then stop. + """ + stream = graph.stream( + input, + { + "configurable": {"thread_id": str(uuid4())}, + "recursion_limit": 1000000000, + }, + ) + + try: + for _ in stream: + break + finally: + stream.close() + + def compile_graph(graph: StateGraph) -> None: """Compile the graph.""" graph.compile() @@ -342,6 +382,21 @@ for name, agraph, graph, input in benchmarks: if graph is not None: r.bench_func(name + "_sync", run, graph, input) + +# First event latency +for name, agraph, graph, input in benchmarks: + r.bench_async_func( + name + "_first_event_latency", + arun_first_event_latency, + agraph, + input, + loop_factory=new_event_loop, + ) + if graph is not None: + r.bench_func( + name + "_first_event_latency_sync", run_first_event_latency, graph, input + ) + # Graph compilation times compilation_benchmarks = ( ( diff --git a/libs/langgraph/bench/sequential.py b/libs/langgraph/bench/sequential.py index 3ab92912f..ef3ce6329 100644 --- a/libs/langgraph/bench/sequential.py +++ b/libs/langgraph/bench/sequential.py @@ -4,7 +4,7 @@ from langgraph.graph import MessagesState, StateGraph from langgraph.utils.runnable import RunnableCallable -def create_sequential(number_nodes) -> StateGraph: +def create_sequential(number_nodes: int) -> StateGraph: """Create a sequential no-op graph consisting of a few hundred nodes.""" builder = StateGraph(MessagesState)