Merge pull request #336 from NVIDIA/fix/reuse-asyncio-loop-between-sy…

…nc-calls

Fix #320. Reuse the asyncio loop between sync calls

nemoguardrails/actions/llm/generation.py

@@ -100,12 +100,13 @@ def __init__(
 
         # There are still some edge cases not covered by nest_asyncio.
         # Using a separate thread always for now.
+        loop = asyncio.get_event_loop()
         if True or check_sync_call_from_async_loop():
             t = threading.Thread(target=asyncio.run, args=(self.init(),))
             t.start()
             t.join()
         else:
-            asyncio.run(self.init())
+            loop.run_until_complete(self.init())
 
         self.llm_task_manager = llm_task_manager
 

nemoguardrails/colang/v2_x/runtime/statemachine.py

@@ -698,9 +698,9 @@ def _resolve_action_conflicts(
                         index = competing_flow_state.action_uids.index(
                             competing_event.action_uid
                         )
-                        competing_flow_state.action_uids[index] = (
-                            winning_event.action_uid
-                        )
+                        competing_flow_state.action_uids[
+                            index
+                        ] = winning_event.action_uid
                         del state.actions[competing_event.action_uid]
 
                     advancing_heads.append(head)

nemoguardrails/rails/llm/llmrails.py

@@ -222,12 +222,13 @@ def __init__(
         # Next, we initialize the Knowledge Base
         # There are still some edge cases not covered by nest_asyncio.
         # Using a separate thread always for now.
+        loop = asyncio.get_event_loop()
         if True or check_sync_call_from_async_loop():
             t = threading.Thread(target=asyncio.run, args=(self._init_kb(),))
             t.start()
             t.join()
         else:
-            asyncio.run(self._init_kb())
+            loop.run_until_complete(self._init_kb())
 
         # We also register the kb as a parameter that can be passed to actions.
         self.runtime.register_action_param("kb", self.kb)
@@ -724,7 +725,8 @@ def generate(
                 "You should replace with `await generate_async(...)` or use `nest_asyncio.apply()`."
             )
 
-        return asyncio.run(
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(
             self.generate_async(
                 prompt=prompt,
                 messages=messages,
@@ -788,7 +790,8 @@ def generate_events(
                 "You should replace with `await generate_events_async(...)` or use `nest_asyncio.apply()`."
             )
 
-        return asyncio.run(self.generate_events_async(events=events))
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(self.generate_events_async(events=events))
 
     async def process_events_async(
         self, events: List[dict], state: Optional[dict] = None
@@ -835,7 +838,8 @@ def process_events(
                 "You should replace with `await generate_events_async(...)."
             )
 
-        return asyncio.run(self.process_events_async(events, state))
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(self.process_events_async(events, state))
 
     def register_action(self, action: callable, name: Optional[str] = None):
         """Register a custom action for the rails configuration."""

tests/test_nest_asyncio.py

@@ -54,11 +54,18 @@ async def test_async_api_error(monkeypatch):
 
     # Reload the module to re-run its top-level code with the new env var
     importlib.reload(nemoguardrails)
+    importlib.reload(nemoguardrails.patch_asyncio)
     importlib.reload(asyncio)
 
+    # Remove the patching marker
+    delattr(asyncio, "_nest_patched")
+
+    assert nemoguardrails.patch_asyncio.nest_asyncio_patch_applied is False
+    assert not hasattr(asyncio, "_nest_patched")
+
     with pytest.raises(
         RuntimeError,
-        match=r"asyncio.run\(\) cannot be called from a running event loop",
+        match=r"await generate_async",
     ):
         chat >> "Hi!"
         chat << "Hello there!"