refactor(py/aio): don't require loop.run_forever, run until completion as necessary (#2483)

Author: pavelgj

py/packages/genkit-ai/src/genkit/ai/veneer.py

@@ -163,6 +163,7 @@ def __init__(
             self.thread = threading.Thread(
                 target=self.start_server,
                 args=[reflection_server_spec, self.loop],
+                daemon=True,
             )
             self.thread.start()
         else:
@@ -193,7 +194,6 @@ def resolver(kind, name, plugin=plugin):
 
     def join(self):
         if self.thread and self.loop:
-            self.loop.run_forever()
             self.thread.join()
 
     def start_server(

py/packages/genkit-ai/src/genkit/aio/loop.py

@@ -29,33 +29,37 @@ def create_loop():
 
 def run_async(loop: asyncio.AbstractEventLoop, fn: Callable):
     """Schedules the callable on the even loop and blocks until it completes."""
-    output = None
-    error = None
-    lock = threading.Lock()
-    lock.acquire()
-
-    async def run_fn():
-        nonlocal lock
-        nonlocal output
-        nonlocal error
-        try:
-            output = await fn()
-        except Exception as e:
-            error = e
-        finally:
-            lock.release()
-
-    asyncio.run_coroutine_threadsafe(run_fn(), loop=loop)
-
-    def wait_for_done():
-        nonlocal lock
+    if loop.is_running():
+        output = None
+        error = None
+        lock = threading.Lock()
         lock.acquire()
 
-    thread = threading.Thread(target=wait_for_done)
-    thread.start()
-    thread.join()
+        async def run_fn():
+            nonlocal lock
+            nonlocal output
+            nonlocal error
+            try:
+                output = await fn()
+                return output
+            except Exception as e:
+                error = e
+            finally:
+                lock.release()
 
-    if error:
-        raise error
+        asyncio.run_coroutine_threadsafe(run_fn(), loop=loop)
 
-    return output
+        def wait_for_done():
+            nonlocal lock
+            lock.acquire()
+
+        thread = threading.Thread(target=wait_for_done)
+        thread.start()
+        thread.join()
+
+        if error:
+            raise error
+
+        return output
+    else:
+        return loop.run_until_complete(fn())

py/packages/genkit-ai/src/genkit/core/reflection.py

@@ -74,7 +74,10 @@
 
 
 def make_reflection_server(
-    registry: Registry, loop: asyncio.AbstractEventLoop, encoding='utf-8'
+    registry: Registry,
+    loop: asyncio.AbstractEventLoop,
+    encoding='utf-8',
+    quiet=True,
 ):
     """Create and return a ReflectionServer class with the given registry.
 
@@ -93,6 +96,18 @@ class ReflectionServer(BaseHTTPRequestHandler):
         registered Genkit actions during development.
         """
 
+        def log_message(self, format, *args):
+            if not quiet:
+                message = format % args
+                logger.debug(
+                    '%s - - [%s] %s'
+                    % (
+                        self.address_string(),
+                        self.log_date_time_string(),
+                        message.translate(self._control_char_table),
+                    )
+                )
+
         def do_GET(self) -> None:  # noqa: N802
             """Handle GET requests to the reflection API.
 

py/samples/google_genai/src/hello.py

@@ -15,9 +15,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import asyncio
-import os
 
-import anyio
 from pydantic import BaseModel, Field
 
 from genkit.ai import Document, Genkit