feat(go/ai): added DefineMiddleware (Middleware V2)

genkit-tools/common/src/types/index.ts

@@ -23,6 +23,7 @@ export * from './document';
 export * from './env';
 export * from './eval';
 export * from './evaluator';
+export * from './middleware';
 export * from './model';
 export * from './prompt';
 export * from './reflection';

genkit-tools/common/src/types/middleware.ts

@@ -0,0 +1,38 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { z } from 'zod';
+
+/** Descriptor for a registered middleware, returned by reflection API. */
+export const MiddlewareDescSchema = z.object({
+  /** Unique name of the middleware. */
+  name: z.string(),
+  /** Human-readable description of what the middleware does. */
+  description: z.string().optional(),
+  /** JSON Schema for the middleware's configuration. */
+  configSchema: z.record(z.any()).nullish(),
+  /** User defined metadata for the middleware. */
+  metadata: z.record(z.any()).optional(),
+});
+export type MiddlewareDesc = z.infer<typeof MiddlewareDescSchema>;
+
+/** Reference to a registered middleware with optional configuration. */
+export const MiddlewareRefSchema = z.object({
+  /** Name of the registered middleware. */
+  name: z.string(),
+  /** Configuration for the middleware (schema defined by the middleware). */
+  config: z.any().optional(),
+});
+export type MiddlewareRef = z.infer<typeof MiddlewareRefSchema>;

genkit-tools/common/src/types/model.ts

@@ -15,6 +15,7 @@
  */
 import { z } from 'zod';
 import { DocumentDataSchema } from './document';
+import { MiddlewareRefSchema } from './middleware';
 import {
   CustomPartSchema,
   DataPartSchema,
@@ -64,32 +65,6 @@ export {
 // IMPORTANT: Keep this file in sync with genkit/ai/src/model-types.ts!
 //
 
-/** Descriptor for a registered middleware, returned by reflection API. */
-export const MiddlewareDescSchema = z.object({
-  /** Unique name of the middleware. */
-  name: z.string(),
-  /** Human-readable description of what the middleware does. */
-  description: z.string().optional(),
-  /** JSON Schema for the middleware's configuration. */
-  configSchema: z.record(z.any()).nullish(),
-  /** User defined metadata for the middleware. */
-  metadata: z.record(z.any()).nullish(),
-});
-export type MiddlewareDesc = z.infer<typeof MiddlewareDescSchema>;
-
-/**
- * Zod schema of middleware reference.
- */
-export const MiddlewareRefSchema = z.object({
-  name: z.string(),
-  config: z.any().optional(),
-});
-
-/**
- * Middleware reference.
- */
-export type MiddlewareRef = z.infer<typeof MiddlewareRefSchema>;
-
 /**
  * Zod schema of an opration representing a model reference.
  */

genkit-tools/genkit-schema.json

@@ -270,6 +270,49 @@
       ],
       "additionalProperties": false
     },
+    "MiddlewareDesc": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        },
+        "description": {
+          "type": "string"
+        },
+        "configSchema": {
+          "anyOf": [
+            {
+              "type": "object",
+              "additionalProperties": {}
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "metadata": {
+          "type": "object",
+          "additionalProperties": {}
+        }
+      },
+      "required": [
+        "name"
+      ],
+      "additionalProperties": false
+    },
+    "MiddlewareRef": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        },
+        "config": {}
+      },
+      "required": [
+        "name"
+      ],
+      "additionalProperties": false
+    },
     "CandidateError": {
       "type": "object",
       "properties": {
@@ -756,56 +799,6 @@
       ],
       "additionalProperties": false
     },
-    "MiddlewareDesc": {
-      "type": "object",
-      "properties": {
-        "name": {
-          "type": "string"
-        },
-        "description": {
-          "type": "string"
-        },
-        "configSchema": {
-          "anyOf": [
-            {
-              "type": "object",
-              "additionalProperties": {}
-            },
-            {
-              "type": "null"
-            }
-          ]
-        },
-        "metadata": {
-          "anyOf": [
-            {
-              "type": "object",
-              "additionalProperties": {}
-            },
-            {
-              "type": "null"
-            }
-          ]
-        }
-      },
-      "required": [
-        "name"
-      ],
-      "additionalProperties": false
-    },
-    "MiddlewareRef": {
-      "type": "object",
-      "properties": {
-        "name": {
-          "type": "string"
-        },
-        "config": {}
-      },
-      "required": [
-        "name"
-      ],
-      "additionalProperties": false
-    },
     "ModelInfo": {
       "type": "object",
       "properties": {

genkit-tools/scripts/schema-exporter.ts

@@ -26,6 +26,7 @@ const EXPORTED_TYPE_MODULES = [
   '../common/src/types/embedder.ts',
   '../common/src/types/evaluator.ts',
   '../common/src/types/error.ts',
+  '../common/src/types/middleware.ts',
   '../common/src/types/model.ts',
   '../common/src/types/parts.ts',
   '../common/src/types/reranker.ts',

go/ai/gen.go

@@ -111,8 +111,9 @@ type GenerateActionOptions struct {
 	// the model to choose a tool, and none forces the model not to use any tools. Defaults to auto.
 	ToolChoice ToolChoice `json:"toolChoice,omitempty"`
 	// Tools is a list of registered tool names for this generation if supported.
-	Tools []string         `json:"tools,omitempty"`
-	Use   []*MiddlewareRef `json:"use,omitempty"`
+	Tools []string `json:"tools,omitempty"`
+	// Use is middleware to apply to this generation, referenced by name with optional config.
+	Use []*MiddlewareRef `json:"use,omitempty"`
 }
 
 // GenerateActionResume holds options for resuming an interrupted generation.
@@ -225,16 +226,25 @@ type Message struct {
 	Role Role `json:"role,omitempty"`
 }
 
+// MiddlewareDesc is the registered descriptor for a middleware.
 type MiddlewareDesc struct {
-	ConfigSchema any    `json:"configSchema,omitempty"`
-	Description  string `json:"description,omitempty"`
-	Metadata     any    `json:"metadata,omitempty"`
-	Name         string `json:"name,omitempty"`
+	// ConfigSchema is a JSON Schema describing the middleware's configuration.
+	ConfigSchema map[string]any `json:"configSchema,omitempty"`
+	// Description explains what the middleware does.
+	Description string `json:"description,omitempty"`
+	// Metadata contains additional context for the middleware.
+	Metadata map[string]any `json:"metadata,omitempty"`
+	// Name is the middleware's unique identifier.
+	Name          string `json:"name,omitempty"`
+	buildFromJSON middlewareFactoryFunc
 }
 
+// MiddlewareRef is a serializable reference to a registered middleware with config.
 type MiddlewareRef struct {
-	Config any    `json:"config,omitempty"`
-	Name   string `json:"name,omitempty"`
+	// Config contains the middleware configuration.
+	Config any `json:"config,omitempty"`
+	// Name is the name of the registered middleware.
+	Name string `json:"name,omitempty"`
 }
 
 // ModelInfo contains metadata about a model's capabilities and characteristics.