feat: PluginV2 refactor (#4132)

Author: huangjeff5

py/bin/sanitize_schema_typing.py

@@ -42,10 +42,9 @@
 
 import ast
 import sys
-from _ast import AST
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Type, cast
+from typing import Any, cast
 
 
 class ClassTransformer(ast.NodeTransformer):

py/packages/genkit/src/genkit/ai/__init__.py

@@ -35,9 +35,9 @@
 from genkit.core import GENKIT_CLIENT_HEADER, GENKIT_VERSION
 from genkit.core.action import ActionRunContext
 from genkit.core.action.types import ActionKind
+from genkit.core.plugin import Plugin
 
 from ._aio import Genkit
-from ._plugin import Plugin
 from ._registry import FlowWrapper, GenkitRegistry
 
 __all__ = [

py/packages/genkit/src/genkit/ai/_aio.py

@@ -43,9 +43,9 @@ class while customizing it with any plugins.
 from genkit.blocks.retriever import IndexerRef, IndexerRequest, RetrieverRef
 from genkit.core.action import ActionRunContext
 from genkit.core.action.types import ActionKind
+from genkit.core.plugin import Plugin
 from genkit.core.typing import (
     BaseDataPoint,
-    BaseEvalDataPoint,
     EmbedRequest,
     EmbedResponse,
     EvalRequest,
@@ -61,8 +61,7 @@ class while customizing it with any plugins.
     ToolChoice,
 )
 
-from ._base import GenkitBase
-from ._plugin import Plugin
+from ._base_async import GenkitBase
 from ._server import ServerSpec
 
 
@@ -328,7 +327,7 @@ async def retrieve(
         Args:
             retriever: Optional retriever name or reference to use.
             query: Text query or a DocumentData containing query text.
-            options: retriever options
+            options: Optional retriever-specific options.
 
         Returns:
             The generated response with documents.
@@ -349,11 +348,20 @@ async def retrieve(
         if isinstance(query, str):
             query = Document.from_text(query)
 
-        final_options = {**(retriever_config or {}), **(options or {})}
+        request_options = {**(retriever_config or {}), **(options or {})}
 
-        retrieve_action = self.registry.lookup_action(ActionKind.RETRIEVER, retriever_name)
+        retrieve_action = await self.registry.resolve_action(ActionKind.RETRIEVER, retriever_name)
+        if retrieve_action is None:
+            raise ValueError(f'Retriever "{retriever_name}" not found')
 
-        return (await retrieve_action.arun(RetrieverRequest(query=query, options=final_options))).response
+        return (
+            await retrieve_action.arun(
+                RetrieverRequest(
+                    query=query,
+                    options=request_options if request_options else None,
+                )
+            )
+        ).response
 
     async def index(
         self,
@@ -366,7 +374,7 @@ async def index(
         Args:
             indexer: Optional indexer name or reference to use.
             documents: Documents to index.
-            options: indexer options
+            options: Optional indexer-specific options.
         """
         indexer_name: str
         indexer_config: dict[str, Any] = {}
@@ -381,11 +389,18 @@ async def index(
         else:
             raise ValueError('Indexer must be specified as a string name or an IndexerRef.')
 
-        final_options = {**(indexer_config or {}), **(options or {})}
+        req_options = {**(indexer_config or {}), **(options or {})}
 
-        index_action = self.registry.lookup_action(ActionKind.INDEXER, indexer_name)
+        index_action = await self.registry.resolve_action(ActionKind.INDEXER, indexer_name)
+        if index_action is None:
+            raise ValueError(f'Indexer "{indexer_name}" not found')
 
-        await index_action.arun(IndexerRequest(documents=documents, options=final_options))
+        await index_action.arun(
+            IndexerRequest(
+                documents=documents,
+                options=req_options if req_options else None,
+            )
+        )
 
     async def embed(
         self,
@@ -410,7 +425,9 @@ async def embed(
         # Merge options passed to embed() with config from EmbedderRef
         final_options = {**(embedder_config or {}), **(options or {})}
 
-        embed_action = self.registry.lookup_action(ActionKind.EMBEDDER, embedder_name)
+        embed_action = await self.registry.resolve_action(ActionKind.EMBEDDER, embedder_name)
+        if embed_action is None:
+            raise ValueError(f'Embedder "{embedder_name}" not found')
 
         return (await embed_action.arun(EmbedRequest(input=documents, options=final_options))).response
 
@@ -445,7 +462,9 @@ async def evaluate(
 
         final_options = {**(evaluator_config or {}), **(options or {})}
 
-        eval_action = self.registry.lookup_action(ActionKind.EVALUATOR, evaluator_name)
+        eval_action = await self.registry.resolve_action(ActionKind.EVALUATOR, evaluator_name)
+        if eval_action is None:
+            raise ValueError(f'Evaluator "{evaluator_name}" not found')
 
         if not eval_run_id:
             eval_run_id = str(uuid.uuid4())

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

@@ -29,10 +29,10 @@
 from genkit.blocks.formats import built_in_formats
 from genkit.blocks.generate import define_generate_action
 from genkit.core.environment import is_dev_environment
+from genkit.core.plugin import Plugin
 from genkit.core.reflection import make_reflection_server
 from genkit.web.manager import find_free_port_sync
 
-from ._plugin import Plugin
 from ._registry import GenkitRegistry
 from ._server import ServerSpec, init_default_runtime
 
@@ -121,19 +121,7 @@ def _initialize_registry(self, model: str | None, plugins: list[Plugin] | None)
         else:
             for plugin in plugins:
                 if isinstance(plugin, Plugin):
-                    plugin.initialize(ai=self)
-
-                    def resolver(kind, name, plugin=plugin):
-                        return plugin.resolve_action(self, kind, name)
-
-                    def action_resolver(plugin=plugin):
-                        if isinstance(plugin.list_actions, list):
-                            return plugin.list_actions
-                        else:
-                            return plugin.list_actions()
-
-                    self.registry.register_action_resolver(plugin.plugin_name(), resolver)
-                    self.registry.register_list_actions_resolver(plugin.plugin_name(), action_resolver)
+                    self.registry.register_plugin(plugin)
                 else:
                     raise ValueError(f'Invalid {plugin=} provided to Genkit: must be of type `genkit.ai.Plugin`')
 

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

@@ -25,11 +25,12 @@
 
 from genkit.aio.loop import run_loop
 from genkit.blocks.formats import built_in_formats
+from genkit.blocks.generate import define_generate_action
 from genkit.core.environment import is_dev_environment
+from genkit.core.plugin import Plugin
 from genkit.core.reflection import create_reflection_asgi_app
 from genkit.web.manager import find_free_port_sync
 
-from ._plugin import Plugin
 from ._registry import GenkitRegistry
 from ._runtime import RuntimeManager
 from ._server import ServerSpec
@@ -59,6 +60,8 @@ def __init__(
         super().__init__()
         self._reflection_server_spec = reflection_server_spec
         self._initialize_registry(model, plugins)
+        # Ensure the default generate action is registered for async usage.
+        define_generate_action(self.registry)
 
     def _initialize_registry(self, model: str | None, plugins: list[Plugin] | None) -> None:
         """Initialize the registry for the Genkit instance.
@@ -82,12 +85,7 @@ def _initialize_registry(self, model: str | None, plugins: list[Plugin] | None)
         else:
             for plugin in plugins:
                 if isinstance(plugin, Plugin):
-                    plugin.initialize(ai=self)
-
-                    def resolver(kind, name, plugin=plugin):
-                        return plugin.resolve_action(self, kind, name)
-
-                    self.registry.register_action_resolver(plugin.plugin_name(), resolver)
+                    self.registry.register_plugin(plugin)
                 else:
                     raise ValueError(f'Invalid {plugin=} provided to Genkit: must be of type `genkit.ai.Plugin`')
 

py/packages/genkit/src/genkit/ai/_plugin.py

@@ -43,7 +43,7 @@
 import uuid
 from collections.abc import AsyncIterator, Callable
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Type
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from genkit.blocks.resource import ResourceFn, ResourceOptions
@@ -429,6 +429,7 @@ def define_reranker(
                 config_schema=reranker_meta['reranker'].get('customOptions'),
                 label=reranker_meta['reranker'].get('label'),
             ),
+            description=reranker_description,
         )
 
     async def rerank(
@@ -636,7 +637,7 @@ def define_model(
         self,
         name: str,
         fn: ModelFn,
-        config_schema: Type[BaseModel] | dict[str, Any] | None = None,
+        config_schema: type[BaseModel] | dict[str, Any] | None = None,
         metadata: dict[str, Any] | None = None,
         info: ModelInfo | None = None,
         description: str | None = None,
@@ -820,7 +821,6 @@ async def prompt(
         Raises:
             GenkitError: If the prompt is not found.
         """
-
         return await lookup_prompt(
             registry=self.registry,
             name=name,

py/packages/genkit/src/genkit/ai/_registry.py

@@ -28,6 +28,8 @@
 
 import structlog
 
+from genkit.core.constants import DEFAULT_GENKIT_VERSION
+
 from ._server import ServerSpec
 
 logger = structlog.get_logger(__name__)
@@ -87,15 +89,24 @@ def _create_and_write_runtime_file(runtime_dir: Path, spec: ServerSpec) -> Path:
         The Path object of the created file.
     """
     current_datetime = datetime.now()
-    runtime_file_name = f'{current_datetime.isoformat()}.json'
+    timestamp_ms = int(current_datetime.timestamp() * 1000)
+    pid = os.getpid()
+
+    # Build a unique runtime ID from the process ID and port
+    port = spec.port if spec.port else ''
+    runtime_id = f'{pid}-{port}' if port else f'{pid}'
+
+    # Include timestamp in filename to avoid collisions across restarts
+    runtime_file_name = f'{runtime_id}-{timestamp_ms}.json'
     runtime_file_path = runtime_dir / runtime_file_name
 
     metadata = json.dumps({
         'reflectionApiSpecVersion': 1,
-        'id': f'{os.getpid()}',
-        'pid': os.getpid(),
+        'id': runtime_id,
+        'pid': pid,
+        'genkitVersion': 'py/' + DEFAULT_GENKIT_VERSION,
         'reflectionServerUrl': spec.url,
-        'timestamp': f'{current_datetime.isoformat()}',
+        'timestamp': current_datetime.isoformat(),
     })
 
     logger.debug(f'Writing runtime file: {runtime_file_path}')
@@ -156,8 +167,8 @@ async def __aenter__(self) -> RuntimeManager:
         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)
-            _register_atexit_cleanup_handler(runtime_file_path)
+            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:
             logger.error(f'Failed to initialize runtime file: {e}', exc_info=True)