From e5b4cd2701a22f9fafc317836c02abb78ac06684 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 7 Oct 2024 10:16:36 -0700
Subject: [PATCH] scheduler-kafka: Pass loop arg to default async consumer and
 producer

- Some forks of aiokafka make this a required arg
---
 libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py     | 3 +++
 libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
index d8e150f5b..c2ad752e4 100644
--- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
@@ -69,6 +69,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
         self.retry_policy = retry_policy
 
     async def __aenter__(self) -> Self:
+        loop = asyncio.get_running_loop()
         self.subgraphs = {
             k: v async for k, v in self.graph.aget_subgraphs(recurse=True)
         }
@@ -81,6 +82,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
                     auto_offset_reset="earliest",
                     group_id="executor",
                     enable_auto_commit=False,
+                    loop=loop,
                     **self.kwargs,
                 )
             )
@@ -89,6 +91,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
 
             self.producer = await self.stack.enter_async_context(
                 DefaultAsyncProducer(
+                    loop=loop,
                     **self.kwargs,
                 )
             )
diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py
index 39e7b755b..493b02d42 100644
--- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py
+++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py
@@ -67,6 +67,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
         self.retry_policy = retry_policy
 
     async def __aenter__(self) -> Self:
+        loop = asyncio.get_running_loop()
         self.subgraphs = {
             k: v async for k, v in self.graph.aget_subgraphs(recurse=True)
         }
@@ -79,6 +80,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
                     auto_offset_reset="earliest",
                     group_id="orchestrator",
                     enable_auto_commit=False,
+                    loop=loop,
                     **self.kwargs,
                 )
             )
@@ -87,6 +89,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
 
             self.producer = await self.stack.enter_async_context(
                 DefaultAsyncProducer(
+                    loop=loop,
                     **self.kwargs,
                 )
             )