feat: middleware V2

ai/src/main/java/com/google/genkit/ai/GenerateOptions.java

@@ -18,6 +18,7 @@
 
 package com.google.genkit.ai;
 
+import com.google.genkit.ai.middleware.GenerationMiddleware;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -43,6 +44,7 @@ public class GenerateOptions<T> {
   private final Integer maxTurns;
   private final ResumeOptions resume;
   private final Class<T> outputClass;
+  private final List<GenerationMiddleware> use;
 
   /**
    * Creates new GenerateOptions.
@@ -74,7 +76,8 @@ public GenerateOptions(
       Map<String, Object> context,
       Integer maxTurns,
       ResumeOptions resume,
-      Class<T> outputClass) {
+      Class<T> outputClass,
+      List<GenerationMiddleware> use) {
     this.model = model;
     this.prompt = prompt;
     this.messages = messages;
@@ -88,6 +91,7 @@ public GenerateOptions(
     this.maxTurns = maxTurns;
     this.resume = resume;
     this.outputClass = outputClass;
+    this.use = use;
   }
 
   /**
@@ -286,6 +290,15 @@ public Class<T> getOutputClass() {
     return outputClass;
   }
 
+  /**
+   * Gets the V2 middleware to apply to this generation.
+   *
+   * @return the middleware list, or null if not set
+   */
+  public List<GenerationMiddleware> getUse() {
+    return use;
+  }
+
   /**
    * Builder for GenerateOptions.
    *
@@ -305,6 +318,7 @@ public static class Builder<T> {
     private Integer maxTurns;
     private ResumeOptions resume;
     private Class<T> outputClass;
+    private List<GenerationMiddleware> use;
 
     public Builder<T> model(String model) {
       this.model = model;
@@ -407,6 +421,29 @@ public Builder<T> resume(ResumeOptions resume) {
       return this;
     }
 
+    /**
+     * Sets V2 middleware to apply to this generation. Middleware hooks wrap the generate loop,
+     * model calls, and tool executions.
+     *
+     * @param use the middleware to apply
+     * @return this builder
+     */
+    public Builder<T> use(List<GenerationMiddleware> use) {
+      this.use = use;
+      return this;
+    }
+
+    /**
+     * Sets V2 middleware to apply to this generation.
+     *
+     * @param middleware the middleware to apply
+     * @return this builder
+     */
+    public Builder<T> use(GenerationMiddleware... middleware) {
+      this.use = List.of(middleware);
+      return this;
+    }
+
     public GenerateOptions<T> build() {
       return new GenerateOptions<>(
           model,
@@ -421,7 +458,8 @@ public GenerateOptions<T> build() {
           context,
           maxTurns,
           resume,
-          outputClass);
+          outputClass,
+          use);
     }
   }
 }

ai/src/main/java/com/google/genkit/ai/middleware/BaseGenerationMiddleware.java

@@ -0,0 +1,78 @@
+/*
+ * 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.google.genkit.ai.middleware;
+
+import com.google.genkit.ai.ModelResponse;
+import com.google.genkit.ai.Tool;
+import com.google.genkit.ai.ToolResponse;
+import com.google.genkit.core.ActionContext;
+import com.google.genkit.core.GenkitException;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * BaseGenerationMiddleware provides default pass-through implementations for all three hooks.
+ * Extend this class and override only the hooks you need.
+ *
+ * <p>Example:
+ *
+ * <pre>{@code
+ * public class TimingMiddleware extends BaseGenerationMiddleware {
+ *   @Override
+ *   public String name() { return "timing"; }
+ *
+ *   @Override
+ *   public GenerationMiddleware newInstance() { return new TimingMiddleware(); }
+ *
+ *   @Override
+ *   public ModelResponse wrapModel(ActionContext ctx, ModelParams params, ModelNext next)
+ *       throws GenkitException {
+ *     long start = System.currentTimeMillis();
+ *     ModelResponse resp = next.apply(ctx, params);
+ *     System.out.println("Model call took " + (System.currentTimeMillis() - start) + "ms");
+ *     return resp;
+ *   }
+ * }
+ * }</pre>
+ */
+public abstract class BaseGenerationMiddleware implements GenerationMiddleware {
+
+  @Override
+  public ModelResponse wrapGenerate(ActionContext ctx, GenerateParams params, GenerateNext next)
+      throws GenkitException {
+    return next.apply(ctx, params);
+  }
+
+  @Override
+  public ModelResponse wrapModel(ActionContext ctx, ModelParams params, ModelNext next)
+      throws GenkitException {
+    return next.apply(ctx, params);
+  }
+
+  @Override
+  public ToolResponse wrapTool(ActionContext ctx, ToolParams params, ToolNext next)
+      throws GenkitException {
+    return next.apply(ctx, params);
+  }
+
+  @Override
+  public List<Tool<?, ?>> tools() {
+    return Collections.emptyList();
+  }
+}

ai/src/main/java/com/google/genkit/ai/middleware/GenerateNext.java

@@ -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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.google.genkit.ai.middleware;
+
+import com.google.genkit.ai.ModelResponse;
+import com.google.genkit.core.ActionContext;
+import com.google.genkit.core.GenkitException;
+
+/** Next function in the {@link GenerationMiddleware#wrapGenerate} hook chain. */
+@FunctionalInterface
+public interface GenerateNext {
+
+  /**
+   * Calls the next handler in the generate chain.
+   *
+   * @param ctx the action context
+   * @param params the generate parameters
+   * @return the model response
+   * @throws GenkitException if processing fails
+   */
+  ModelResponse apply(ActionContext ctx, GenerateParams params) throws GenkitException;
+}

ai/src/main/java/com/google/genkit/ai/middleware/GenerateParams.java

@@ -0,0 +1,54 @@
+/*
+ * 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.google.genkit.ai.middleware;
+
+import com.google.genkit.ai.ModelRequest;
+
+/** Holds parameters for the {@link GenerationMiddleware#wrapGenerate} hook. */
+public class GenerateParams {
+
+  private final ModelRequest request;
+  private final int iteration;
+
+  /**
+   * Creates GenerateParams.
+   *
+   * @param request the current model request for this iteration
+   * @param iteration the current tool-loop iteration (0-indexed)
+   */
+  public GenerateParams(ModelRequest request, int iteration) {
+    this.request = request;
+    this.iteration = iteration;
+  }
+
+  /** Returns the current model request with accumulated messages. */
+  public ModelRequest getRequest() {
+    return request;
+  }
+
+  /** Returns the current tool-loop iteration (0-indexed). */
+  public int getIteration() {
+    return iteration;
+  }
+
+  /** Returns a new GenerateParams with the given request, preserving the iteration. */
+  public GenerateParams withRequest(ModelRequest request) {
+    return new GenerateParams(request, this.iteration);
+  }
+}

ai/src/main/java/com/google/genkit/ai/middleware/GenerationMiddleware.java

@@ -0,0 +1,126 @@
+/*
+ * 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.google.genkit.ai.middleware;
+
+import com.google.genkit.ai.ModelResponse;
+import com.google.genkit.ai.Tool;
+import com.google.genkit.ai.ToolResponse;
+import com.google.genkit.core.ActionContext;
+import com.google.genkit.core.GenkitException;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * GenerationMiddleware provides hooks for different stages of the generation pipeline.
+ *
+ * <p>This is the V2 middleware interface that replaces the generic {@code Middleware<I, O>}. It
+ * provides three distinct hooks:
+ *
+ * <ul>
+ *   <li>{@link #wrapGenerate} - wraps each iteration of the tool loop
+ *   <li>{@link #wrapModel} - wraps each model API call
+ *   <li>{@link #wrapTool} - wraps each tool execution
+ * </ul>
+ *
+ * <p>Each {@code generate()} call creates a fresh instance via {@link #newInstance()}, enabling
+ * per-invocation state (e.g., counters, timers) without shared mutable state across requests.
+ *
+ * <p>Example:
+ *
+ * <pre>{@code
+ * public class LoggingMiddleware extends BaseGenerationMiddleware {
+ *   private int modelCalls = 0;
+ *
+ *   @Override
+ *   public String name() { return "logging"; }
+ *
+ *   @Override
+ *   public GenerationMiddleware newInstance() { return new LoggingMiddleware(); }
+ *
+ *   @Override
+ *   public ModelResponse wrapModel(ActionContext ctx, ModelParams params, ModelNext next)
+ *       throws GenkitException {
+ *     modelCalls++;
+ *     System.out.println("Model call #" + modelCalls);
+ *     ModelResponse resp = next.apply(ctx, params);
+ *     System.out.println("Model responded with " + resp.getText());
+ *     return resp;
+ *   }
+ * }
+ * }</pre>
+ */
+public interface GenerationMiddleware {
+
+  /** Returns the middleware's unique identifier. */
+  String name();
+
+  /**
+   * Returns a fresh instance for each {@code generate()} call, enabling per-invocation state.
+   *
+   * <p>Stable state (e.g., API keys, configuration) should be preserved. Per-request state (e.g.,
+   * counters) should be reset.
+   */
+  GenerationMiddleware newInstance();
+
+  /**
+   * Wraps each iteration of the generate tool loop.
+   *
+   * @param ctx the action context
+   * @param params the generate parameters including the current request and iteration
+   * @param next the next function in the chain
+   * @return the model response
+   * @throws GenkitException if processing fails
+   */
+  ModelResponse wrapGenerate(ActionContext ctx, GenerateParams params, GenerateNext next)
+      throws GenkitException;
+
+  /**
+   * Wraps each model API call.
+   *
+   * @param ctx the action context
+   * @param params the model parameters including the request
+   * @param next the next function in the chain
+   * @return the model response
+   * @throws GenkitException if processing fails
+   */
+  ModelResponse wrapModel(ActionContext ctx, ModelParams params, ModelNext next)
+      throws GenkitException;
+
+  /**
+   * Wraps each tool execution. May be called concurrently when multiple tools execute in parallel.
+   * Implementations must be safe for concurrent use.
+   *
+   * @param ctx the action context
+   * @param params the tool parameters including the request and resolved tool
+   * @param next the next function in the chain
+   * @return the tool response
+   * @throws GenkitException if processing fails
+   */
+  ToolResponse wrapTool(ActionContext ctx, ToolParams params, ToolNext next) throws GenkitException;
+
+  /**
+   * Returns additional tools to make available during generation. These tools are dynamically added
+   * when the middleware is used.
+   *
+   * @return the list of additional tools, or empty list if none
+   */
+  default List<Tool<?, ?>> tools() {
+    return Collections.emptyList();
+  }
+}