From e6895a9e769c6935664324554707def702a9bf21 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 20 Sep 2024 14:34:46 -0700 Subject: [PATCH] core: In astream_events propagate cancellation reason to inner task --- libs/core/langchain_core/tracers/event_stream.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/libs/core/langchain_core/tracers/event_stream.py b/libs/core/langchain_core/tracers/event_stream.py index 3b2f512560459..d374a3cd3f4a0 100644 --- a/libs/core/langchain_core/tracers/event_stream.py +++ b/libs/core/langchain_core/tracers/event_stream.py @@ -994,6 +994,10 @@ async def consume_astream() -> None: del event["data"]["input"] yield event + except asyncio.CancelledError as exc: + # Cancel the task if it's still running + task.cancel(exc.args[0] if exc.args else None) + raise finally: # Cancel the task if it's still running task.cancel()