From 8be0fb675a48c6bcbbbaca2343960281bbc7cd46 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Fri, 21 Mar 2025 15:25:28 -0400 Subject: [PATCH] x --- libs/langgraph/bench/__main__.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py index ad6c7e65d..06ec43a28 100644 --- a/libs/langgraph/bench/__main__.py +++ b/libs/langgraph/bench/__main__.py @@ -35,14 +35,19 @@ async def arun_first_event_latency(graph: Pregel, input: dict) -> None: Run the graph until the first event is processed and then stop. """ - async for _ in graph.astream( + stream = graph.astream( input, { "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, - ): - break + ) + + try: + async for _ in stream: + break + finally: + await stream.aclose() def run(graph: Pregel, input: dict): @@ -65,14 +70,19 @@ def run_first_event_latency(graph: Pregel, input: dict) -> None: Run the graph until the first event is processed and then stop. """ - for _ in graph.stream( + stream = graph.stream( input, { "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, - ): - break + ) + + try: + for _ in stream: + break + finally: + stream.close() def compile_graph(graph: StateGraph) -> None: