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,
@@ -157,6 +163,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.
@@ -425,7 +432,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],
@@ -453,7 +460,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,
@@ -481,7 +488,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],
@@ -509,7 +516,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,
@@ -535,7 +542,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,
@@ -725,6 +732,45 @@ 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."""
+        return new_middleware(middleware_cls)
+
+    def define_middleware(self, middleware_cls: type[BaseMiddleware]) -> MiddlewareDesc:
+        """Register a middleware class on this app and return the descriptor.
+
+        Registering a class:
+
+        * Makes it visible to the **Dev UI** through the reflection API.
+        * Allows it to be referenced by name via :class:`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 :class:`MiddlewareRef`, so the same
+        pydantic fields drive both:
+
+        * the inline path: ``use=[cls(...)]``
+        * the registered path: ``use=[MiddlewareRef(name=cls.name)]``
+
+        Returns:
+            The registered :class:`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():
@@ -799,7 +845,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]: ...
 
@@ -826,7 +872,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]: ...
 
@@ -851,7 +897,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.
@@ -860,6 +906,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,
@@ -879,13 +928,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]
         )
 
@@ -912,7 +960,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]: ...
@@ -940,7 +988,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]: ...
@@ -966,14 +1014,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,
@@ -993,14 +1044,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]
             )
 
@@ -1184,7 +1234,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."""