feat(py): generate middleware

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

@@ -41,6 +41,7 @@
     Media,
     MediaPart,
     Metadata,
+    MiddlewareRef,
     Part,
     ReasoningPart,
     Role,
@@ -131,6 +132,8 @@
     'DocumentPart',
     # Plugin interface
     'Plugin',
+    # Middleware references (wire form for use= parameter)
+    'MiddlewareRef',
     # AI runtime
     'ActionKind',
     'ActionRunContext',

py/packages/genkit/src/genkit/_ai/_aio.py

@@ -44,12 +44,16 @@
 )
 from genkit._ai._formats import built_in_formats
 from genkit._ai._formats._types import FormatDef
-from genkit._ai._generate import define_generate_action, generate_action, registry_with_inline_tools
+from genkit._ai._generate import (
+    define_generate_action,
+    generate_action,
+    registry_with_inline_middleware,
+    registry_with_inline_tools,
+)
 from genkit._ai._model import (
     Message,
     ModelConfig,
     ModelFn,
-    ModelMiddleware,
     ModelResponse,
     ModelResponseChunk,
     define_model,
@@ -90,6 +94,7 @@
 from genkit._core._environment import is_dev_environment
 from genkit._core._error import GenkitError
 from genkit._core._logger import get_logger
+from genkit._core._middleware import BaseMiddleware, MiddlewareDesc, new_middleware
 from genkit._core._model import Document
 from genkit._core._plugin import Plugin
 from genkit._core._reflection import ReflectionServer, ServerSpec, create_reflection_asgi_app
@@ -102,6 +107,7 @@
     EmbedRequest,
     EvalRequest,
     EvalResponse,
+    MiddlewareRef,
     ModelInfo,
     Operation,
     Part,
@@ -158,6 +164,7 @@ def __init__(
         self._initialize_registry(model, plugins)
         # Ensure the default generate action is registered for async usage.
         define_generate_action(self.registry)
+        self._register_plugin_middleware(plugins)
         # In dev mode, start the reflection server immediately in a background
         # daemon thread so it's available regardless of which web framework (or
         # none) the user chooses.
@@ -426,7 +433,7 @@ def define_prompt(
         metadata: dict[str, object] | None = None,
         tools: Sequence[str | Tool] | None = None,
         tool_choice: ToolChoice | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         input_schema: type[InputT],
         output_schema: type[OutputT],
@@ -454,7 +461,7 @@ def define_prompt(
         metadata: dict[str, object] | None = None,
         tools: Sequence[str | Tool] | None = None,
         tool_choice: ToolChoice | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         input_schema: type[InputT],
         output_schema: dict[str, object] | str | None = None,
@@ -482,7 +489,7 @@ def define_prompt(
         metadata: dict[str, object] | None = None,
         tools: Sequence[str | Tool] | None = None,
         tool_choice: ToolChoice | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         input_schema: dict[str, object] | str | None = None,
         output_schema: type[OutputT],
@@ -510,7 +517,7 @@ def define_prompt(
         metadata: dict[str, object] | None = None,
         tools: Sequence[str | Tool] | None = None,
         tool_choice: ToolChoice | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         input_schema: dict[str, object] | str | None = None,
         output_schema: dict[str, object] | str | None = None,
@@ -536,7 +543,7 @@ def define_prompt(
         metadata: dict[str, object] | None = None,
         tools: Sequence[str | Tool] | None = None,
         tool_choice: ToolChoice | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         input_schema: type | dict[str, object] | str | None = None,
         output_schema: type | dict[str, object] | str | None = None,
@@ -726,6 +733,47 @@ def _initialize_registry(self, model: str | None, plugins: list[Plugin] | None)
                 else:
                     raise ValueError(f'Invalid {plugin=} provided to Genkit: must be of type `genkit.ai.Plugin`')
 
+    def _register_plugin_middleware(self, plugins: list[Plugin] | None) -> None:
+        """Register middleware descriptors returned by ``Plugin.list_middleware``."""
+        if not plugins:
+            return
+        for plugin in plugins:
+            for desc in plugin.list_middleware():
+                self.registry.register_value('middleware', desc.name, desc)
+
+    def new_middleware(self, middleware_cls: type[BaseMiddleware]) -> MiddlewareDesc:
+        """Build a ``MiddlewareDesc`` from a class (same as ``genkit.middleware.new_middleware``).
+
+        Does not register on the registry. Pass the result to ``middleware_plugin([...])``
+        or return it from a custom ``Plugin.list_middleware`` so it is registered when the
+        app is constructed.
+
+        Returns:
+            The ``MiddlewareDesc`` instance.
+        """
+        return new_middleware(middleware_cls)
+
+    def define_middleware(self, middleware_cls: type[BaseMiddleware]) -> MiddlewareDesc:
+        """Register a middleware class on this app and return the resulting descriptor.
+
+        Registering a class makes it visible to the **Dev UI** (via the reflection
+        API) and allows it to be referenced by name using ``MiddlewareRef``.
+        Equivalent to building the descriptor with ``new_middleware(cls)`` and wiring
+        it through ``middleware_plugin([...])`` at construction time, but usable after
+        ``Genkit`` has already been built. The factory instantiates
+        ``middleware_cls(**config)`` each time a request resolves the name via
+        ``MiddlewareRef``, so the same pydantic fields drive both the inline
+        (``use=[cls(...)]``) and registered (``use=[MiddlewareRef(name=cls.name)]``)
+        paths.
+
+        Returns:
+            The registered ``MiddlewareDesc``; also available via
+            ``registry.lookup_value('middleware', cls.name)``.
+        """
+        desc = new_middleware(middleware_cls)
+        self.registry.register_value('middleware', desc.name, desc)
+        return desc
+
     def run_main(self, coro: Coroutine[Any, Any, T]) -> T | None:
         """Run the user's main coroutine, blocking in dev mode for the reflection server."""
         if not is_dev_environment():
@@ -800,7 +848,7 @@ async def generate(
         output_content_type: str | None = None,
         output_instructions: str | None = None,
         output_constrained: bool | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
     ) -> ModelResponse[OutputT]: ...
 
@@ -827,7 +875,7 @@ async def generate(
         output_content_type: str | None = None,
         output_instructions: str | None = None,
         output_constrained: bool | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
     ) -> ModelResponse[Any]: ...
 
@@ -852,7 +900,7 @@ async def generate(
         output_content_type: str | None = None,
         output_instructions: str | None = None,
         output_constrained: bool | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
     ) -> ModelResponse[Any]:
         """Generate text or structured data using a language model.
@@ -861,6 +909,9 @@ async def generate(
         is covariant: ``list[Tool]`` or ``list[str]`` are both assignable to
         ``Sequence[str | Tool]``, but not to ``list[str | Tool]``.
         """
+        registry = await registry_with_inline_tools(self.registry, tools)
+        child_registry = registry if registry.is_child else registry.new_child()
+        refs = registry_with_inline_middleware(child_registry, use) or None
         prompt_config = PromptConfig(
             model=model,
             prompt=prompt,
@@ -880,13 +931,12 @@ async def generate(
             output_schema=output_schema,
             output_constrained=output_constrained,
             docs=docs,
+            use=refs,
         )
-        registry = await registry_with_inline_tools(self.registry, prompt_config.tools)
         gen_options = await to_generate_action_options(registry, prompt_config)
         return await generate_action(
-            registry,
+            child_registry,
             gen_options,
-            middleware=use,
             context=context if context else ActionRunContext._current_context(),  # pyright: ignore[reportPrivateUsage]
         )
 
@@ -913,7 +963,7 @@ def generate_stream(
         output_content_type: str | None = None,
         output_instructions: str | None = None,
         output_constrained: bool | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         timeout: float | None = None,
     ) -> ModelStreamResponse[OutputT]: ...
@@ -941,7 +991,7 @@ def generate_stream(
         output_content_type: str | None = None,
         output_instructions: str | None = None,
         output_constrained: bool | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         timeout: float | None = None,
     ) -> ModelStreamResponse[Any]: ...
@@ -967,14 +1017,17 @@ def generate_stream(
         output_content_type: str | None = None,
         output_instructions: str | None = None,
         output_constrained: bool | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
         timeout: float | None = None,
     ) -> ModelStreamResponse[Any]:
         """Stream generated text, returning a ModelStreamResponse with .stream and .response."""
         channel: Channel[ModelResponseChunk, ModelResponse[Any]] = Channel(timeout=timeout)
 
         async def _run_generate() -> ModelResponse[Any]:
+            registry = await registry_with_inline_tools(self.registry, tools)
+            child_registry = registry if registry.is_child else registry.new_child()
+            refs = registry_with_inline_middleware(child_registry, use) or None
             prompt_config = PromptConfig(
                 model=model,
                 prompt=prompt,
@@ -994,14 +1047,13 @@ async def _run_generate() -> ModelResponse[Any]:
                 output_schema=output_schema,
                 output_constrained=output_constrained,
                 docs=docs,
+                use=refs,
             )
-            registry = await registry_with_inline_tools(self.registry, prompt_config.tools)
             gen_options = await to_generate_action_options(registry, prompt_config)
             return await generate_action(
-                registry,
+                child_registry,
                 gen_options,
                 on_chunk=lambda c: channel.send(c),
-                middleware=use,
                 context=context if context else ActionRunContext._current_context(),  # pyright: ignore[reportPrivateUsage]
             )
 
@@ -1185,7 +1237,7 @@ async def generate_operation(
         output_content_type: str | None = None,
         output_instructions: str | None = None,
         output_constrained: bool | None = None,
-        use: list[ModelMiddleware] | None = None,
+        use: list[BaseMiddleware | MiddlewareRef] | None = None,
         docs: list[Document] | None = None,
     ) -> Operation:
         """Generate content using a long-running model, returning an Operation to poll."""