fix(py/genkit): ai.run() -> ai.run_main() and fix docstrings/commentary. (#2731)

fix(py/genkit): `ai.run()` -> `ai.run_main()` and fix docstrings/commentary.

CHANGELOG:
- [ ] Update remaining instances of `ai.run()` -> `ai.run_main()` throughout.
- [ ] Fix docstrings and commentary.
- [ ] Fix types for `run_main()`.

Author: yesudeep

py/packages/genkit/src/genkit/ai/_base.py

@@ -20,7 +20,7 @@
 import threading
 from collections.abc import Coroutine
 from http.server import HTTPServer
-from typing import Any
+from typing import Any, TypeVar
 
 import structlog
 
@@ -36,6 +36,8 @@
 
 logger = structlog.get_logger(__name__)
 
+T = TypeVar('T')
+
 
 class GenkitBase(GenkitRegistry):
     """Base class with shared infra for Genkit instances (sync and async)."""
@@ -58,7 +60,7 @@ def __init__(
         self._initialize_server(reflection_server_spec)
         self._initialize_registry(model, plugins)
 
-    def run_main(self, coro: Coroutine[Any, Any, Any] | None = None) -> Any:
+    def run_main(self, coro: Coroutine[Any, Any, T] | None = None) -> T:
         """Runs the provided coroutine on an event loop.
 
         Args:
@@ -67,7 +69,6 @@ def run_main(self, coro: Coroutine[Any, Any, Any] | None = None) -> Any:
         Returns:
             The result of the coroutine.
         """
-
         if not coro:
 
             async def blank_coro():
@@ -78,7 +79,7 @@ async def blank_coro():
         result = None
         if self._loop:
 
-            async def run() -> Any:
+            async def run() -> T:
                 return await coro
 
             result = run_async(self._loop, run)

py/packages/genkit/src/genkit/ai/_base_async.py

@@ -14,7 +14,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Base/shared implementation for Genkit user-facing API."""
+"""Asynchronous server gateway interface implementation for Genkit."""
 
 from collections.abc import Coroutine
 from typing import Any, TypeVar
@@ -91,15 +91,17 @@ def resolver(kind, name, plugin=plugin):
                 else:
                     raise ValueError(f'Invalid {plugin=} provided to Genkit: must be of type `genkit.ai.Plugin`')
 
-    def run(self, coro: Coroutine[Any, Any, T]) -> T:
+    def run_main(self, coro: Coroutine[Any, Any, T]) -> T:
         """Run the user's main coroutine.
 
-        In development mode (`GENKIT_ENV=dev`), this starts the Genkit reflection
-        server and runs the user's coroutine concurrently within the same event loop,
-        blocking until the server is stopped (e.g., via Ctrl+C).
+        In development mode (`GENKIT_ENV=dev`), this starts the Genkit
+        reflection server and runs the user's coroutine concurrently within the
+        same event loop, blocking until the server is stopped (e.g., via
+        Ctrl+C).
 
         In production mode, this simply runs the user's coroutine to completion
-        using `asyncio.run()`.
+        using `uvloop.run()` for performance if available, otherwise
+        `asyncio.run()`.
 
         Args:
             coro: The main coroutine provided by the user.
@@ -111,7 +113,6 @@ def run(self, coro: Coroutine[Any, Any, T]) -> T:
             logger.info('Running in production mode.')
             return run_loop(coro)
 
-        # Development mode: Start reflection server and user coro concurrently.
         logger.info('Running in development mode.')
 
         spec = self._reflection_server_spec
@@ -138,24 +139,23 @@ async def run_user_coro_wrapper():
                 finally:
                     user_task_finished_event.set()
 
-            reflection_server = make_reflection_server(self.registry, spec)
+            reflection_server = _make_reflection_server(self.registry, spec)
 
             try:
                 async with RuntimeManager(spec):
-                    # We use anyio's task group because it's compatible with
+                    # We use anyio.TaskGroup because it is compatible with
                     # asyncio's event loop and works with Python 3.10
-                    # (asyncio.create_task_group was added in 3.11, and we can switch
-                    # to that if we drop support for 3.10).
+                    # (asyncio.TaskGroup was added in 3.11, and we can switch to
+                    # that when we drop support for 3.10).
                     async with anyio.create_task_group() as tg:
                         # Start reflection server in the background.
                         tg.start_soon(reflection_server.serve, name='genkit-reflection-server')
-                        await anyio.sleep(0.2)
-                        logger.info(f'Started Genkit reflection server at {spec.scheme}://{spec.host}:{spec.port}')
+                        logger.info(f'Started Genkit reflection server at {spec.url}')
 
                         # Start the (potentially short-lived) user coroutine wrapper
                         tg.start_soon(run_user_coro_wrapper, name='genkit-user-coroutine')
 
-                        # Block here until the TaskGroup is cancelled (e.g. Ctrl+C)
+                        # Block here until the task group is canceled (e.g. Ctrl+C)
                         # or a task raises an unhandled exception. It should not
                         # exit just because the user coroutine finishes.
 
@@ -166,7 +166,7 @@ async def run_user_coro_wrapper():
                 logger.exception(e)
                 raise
 
-            # After the TaskGroup finishes (error or cancellation)
+            # After the TaskGroup finishes (error or cancelation).
             if user_task_finished_event.is_set():
                 logger.debug('User coroutine finished before TaskGroup exit.')
                 return user_result
@@ -177,8 +177,19 @@ async def run_user_coro_wrapper():
         return anyio.run(dev_runner)
 
 
-def make_reflection_server(registry: GenkitRegistry, spec: ServerSpec) -> uvicorn.Server:
-    """Make a reflection server for the given registry and spec."""
+def _make_reflection_server(registry: GenkitRegistry, spec: ServerSpec) -> uvicorn.Server:
+    """Make a reflection server for the given registry and spec.
+
+    This is a helper function to make it easier to test the reflection server
+    in isolation.
+
+    Args:
+        registry: The registry to use for the reflection server.
+        spec: The spec to use for the reflection server.
+
+    Returns:
+        A uvicorn server instance.
+    """
     app = create_reflection_asgi_app(registry=registry)
     config = uvicorn.Config(app, host=spec.host, port=spec.port, loop='asyncio')
     return uvicorn.Server(config)

py/packages/genkit/src/genkit/ai/_runtime.py

@@ -145,18 +145,18 @@ def __init__(self, spec: ServerSpec, runtime_dir: str | Path | None = None):
         """
         self.spec = spec
         if runtime_dir is None:
-            self.runtime_dir = Path(os.getcwd()) / DEFAULT_RUNTIME_DIR_NAME
+            self._runtime_dir = Path(os.getcwd()) / DEFAULT_RUNTIME_DIR_NAME
         else:
-            self.runtime_dir = Path(runtime_dir)
+            self._runtime_dir = Path(runtime_dir)
 
         self._runtime_file_path: Path | None = None
 
     async def __aenter__(self) -> RuntimeManager:
         """Create the runtime directory and file."""
         try:
-            await logger.adebug(f'Ensuring runtime directory exists: {self.runtime_dir}')
-            self.runtime_dir.mkdir(parents=True, exist_ok=True)
-            runtime_file_path = _create_and_write_runtime_file(self.runtime_dir, self.spec)
+            await logger.adebug(f'Ensuring runtime directory exists: {self._runtime_dir}')
+            self._runtime_dir.mkdir(parents=True, exist_ok=True)
+            runtime_file_path = _create_and_write_runtime_file(self._runtime_dir, self.spec)
             _register_atexit_cleanup_handler(runtime_file_path)
 
         except Exception as e:
@@ -186,9 +186,9 @@ async def __aexit__(
     def __enter__(self) -> RuntimeManager:
         """Synchronous entry point: Create the runtime directory and file."""
         try:
-            logger.debug(f'[sync] Ensuring runtime directory exists: {self.runtime_dir}')
-            self.runtime_dir.mkdir(parents=True, exist_ok=True)
-            self._runtime_file_path = _create_and_write_runtime_file(self.runtime_dir, self.spec)
+            logger.debug(f'[sync] Ensuring runtime directory exists: {self._runtime_dir}')
+            self._runtime_dir.mkdir(parents=True, exist_ok=True)
+            self._runtime_file_path = _create_and_write_runtime_file(self._runtime_dir, self.spec)
             _register_atexit_cleanup_handler(self._runtime_file_path)
 
         except Exception as e:

py/samples/firestore-retreiver/src/main.py

@@ -93,6 +93,7 @@ async def index_documents() -> None:
 
 @ai.flow()
 async def retreive_documents():
+    """Retrieves the film documents from Firestore."""
     return await ai.retrieve(
         query=Document.from_text('sci-fi film'),
         retriever=firestore_action_name('filmsretriever'),

py/samples/google-genai-context-caching/src/context_caching.py

@@ -101,4 +101,4 @@ async def main() -> None:
 
 
 if __name__ == '__main__':
-    ai.run(main())
+    ai.run_main(main())

py/samples/short-n-long/README.md

@@ -39,3 +39,9 @@ To start as a server in dev mode:
 ```bash
 genkit start -- uv run src/short_n_long/main.py --server
 ```
+
+## Running with a specific version of Python
+
+```bash
+genkit start -- uv run --python python3.10 src/short_n_long/main.py
+```

py/samples/short-n-long/src/short_n_long/main.py

@@ -42,6 +42,7 @@
 """
 
 import argparse
+import asyncio
 
 import structlog
 import uvicorn
@@ -114,8 +115,8 @@ async def simple_generate_with_tools_flow(value: int) -> str:
     return response.text
 
 
-@ai.tool(name='gablorkenTool2')
-def gablorken_tool2(input_: GablorkenInput, ctx: ToolRunContext):
+@ai.tool(name='interruptingTool')
+def interrupting_tool(input_: GablorkenInput, ctx: ToolRunContext):
     """The user-defined tool function.
 
     Args:
@@ -146,7 +147,7 @@ async def simple_generate_with_interrupts(value: int) -> str:
                 content=[TextPart(text=f'what is a gablorken of {value}')],
             ),
         ],
-        tools=['gablorkenTool2'],
+        tools=['interruptingTool'],
     )
     await logger.ainfo(f'len(response.tool_requests)={len(response1.tool_requests)}')
     if len(response1.interrupts) == 0:
@@ -233,6 +234,29 @@ async def say_hi_stream(name: str, ctx):
     return result
 
 
+@ai.flow()
+async def stream_greeting(name: str, ctx) -> str:
+    """Stream a greeting for the given name.
+
+    Args:
+        name: the name to send to test function
+        ctx: the context of the tool
+
+    Returns:
+        The generated response with a function.
+    """
+    chunks = [
+        'hello',
+        name,
+        'how are you?',
+    ]
+    for data in chunks:
+        await asyncio.sleep(1)
+        ctx.send_chunk(data)
+
+    return 'test streaming response'
+
+
 class Skills(BaseModel):
     """Skills for an RPG character."""
 
@@ -358,4 +382,4 @@ async def main(ai: Genkit) -> None:
 if __name__ == '__main__':
     config: argparse.Namespace = parse_args()
     runner = server_main if config.server else main
-    ai.run(runner(ai))
+    ai.run_main(runner(ai))