diff --git a/ai/src/main/java/com/google/genkit/ai/GenerateAction.java b/ai/src/main/java/com/google/genkit/ai/GenerateAction.java index fb193f689..85856461c 100644 --- a/ai/src/main/java/com/google/genkit/ai/GenerateAction.java +++ b/ai/src/main/java/com/google/genkit/ai/GenerateAction.java @@ -20,6 +20,14 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.middleware.GenerateNext; +import com.google.genkit.ai.middleware.GenerateParams; +import com.google.genkit.ai.middleware.GenerationMiddleware; +import com.google.genkit.ai.middleware.GenerationMiddlewareDesc; +import com.google.genkit.ai.middleware.ModelNext; +import com.google.genkit.ai.middleware.ModelParams; +import com.google.genkit.ai.middleware.ToolNext; +import com.google.genkit.ai.middleware.ToolParams; import com.google.genkit.ai.telemetry.ModelTelemetryHelper; import com.google.genkit.core.*; import com.google.genkit.core.tracing.SpanMetadata; @@ -94,42 +102,62 @@ public ModelResponse run( throw new GenkitException("GenerateActionOptions cannot be null"); } + // Resolve middleware names from the registry. The Dev UI's Middleware panel sends + // selected middleware as a list of names in `options.use`. Each name is looked up + // in the "middleware" value bucket (registered via Genkit.Builder.middleware(...)). + // Mirrors the JS SDK's resolveMiddleware() in js/ai/src/generate/action.ts. + final List middlewares = resolveMiddlewares(options.getUse()); + + // Core: run the full tool loop (possibly streaming). model.run() inside is wrapped + // by the wrapModel chain; tool execution is wrapped by the wrapTool chain. + GenerateNext core = + (gctx, gparams) -> + runIterations(gctx, gparams.getRequest(), gparams.getOnChunk(), middlewares); + + // Outermost: wrapGenerate chain + GenerateNext chain = chainGenerate(middlewares, core); + + int initialMsgIdx = options.getMessages() != null ? options.getMessages().size() : 0; + return chain.apply(ctx, new GenerateParams(options, 0, initialMsgIdx, streamCallback)); + } + + /** + * Executes the tool-call loop. Each iteration wraps the model call in the {@code wrapModel} + * middleware chain and tool execution in the {@code wrapTool} chain. + */ + private ModelResponse runIterations( + ActionContext ctx, + GenerateActionOptions options, + Consumer streamCallback, + List middlewares) + throws GenkitException { + String modelName = options.getModel(); if (modelName == null || modelName.isEmpty()) { throw new GenkitException("Model name is required"); } - // Resolve the model action key String modelKey = resolveModelKey(modelName); - - // Look up the model in the registry Action action = registry.lookupAction(modelKey); if (action == null) { throw new GenkitException("Model not found: " + modelName + " (key: " + modelKey + ")"); } - if (!(action instanceof Model)) { throw new GenkitException("Action is not a model: " + modelKey); } + final Model model = (Model) action; - Model model = (Model) action; - - // Build the model request from the options ModelRequest request = buildModelRequest(options); logger.debug("Generating with model: {}", modelKey); - // Determine if we should return tool requests without executing them boolean returnToolRequests = Boolean.TRUE.equals(options.getReturnToolRequests()); - - // Get max turns for tool loop (default to 5) int maxTurns = options.getMaxTurns() != null ? options.getMaxTurns() : 5; int turn = 0; String flowName = ctx.getFlowName(); while (turn < maxTurns) { - // Create span metadata for the model call SpanMetadata spanMetadata = SpanMetadata.builder() .name(modelName) @@ -144,7 +172,7 @@ public ModelResponse run( final ModelRequest currentRequest = request; final String spanPath = "/generate/" + modelName; - // Run the model wrapped in a span + // Run the model wrapped in a span and through the wrapModel middleware chain. ModelResponse response = Tracer.runInNewSpan( ctx, @@ -152,49 +180,46 @@ public ModelResponse run( request, (spanCtx, req) -> { ActionContext newCtx = ctx.withSpanContext(spanCtx); - if (streamCallback != null && model.supportsStreaming()) { - return ModelTelemetryHelper.runWithTelemetryStreaming( - modelName, - flowName, - spanPath, - currentRequest, - r -> model.run(newCtx, r, streamCallback)); - } else { - return ModelTelemetryHelper.runWithTelemetry( - modelName, flowName, spanPath, currentRequest, r -> model.run(newCtx, r)); - } + ModelNext modelCore = + (mctx, mparams) -> { + ModelRequest mreq = mparams.getRequest(); + Consumer sc = mparams.getStreamCallback(); + if (sc != null && model.supportsStreaming()) { + return ModelTelemetryHelper.runWithTelemetryStreaming( + modelName, flowName, spanPath, mreq, r -> model.run(mctx, r, sc)); + } else { + return ModelTelemetryHelper.runWithTelemetry( + modelName, flowName, spanPath, mreq, r -> model.run(mctx, r)); + } + }; + ModelNext wrappedModel = chainModel(middlewares, modelCore); + return wrappedModel.apply(newCtx, new ModelParams(currentRequest, streamCallback)); }); // Check if the model requested tool calls List toolRequestParts = extractToolRequestParts(response); - // If no tool requests or we should return them without executing, return - // response if (toolRequestParts.isEmpty() || returnToolRequests) { return response; } - // Check if we have tools to execute if (options.getTools() == null || options.getTools().isEmpty()) { - // No tools available, return response with tool requests return response; } - // Execute tools - List toolResponseParts = executeTools(ctx, toolRequestParts, options.getTools()); + // Execute tools through the wrapTool chain + List toolResponseParts = + executeTools(ctx, toolRequestParts, options.getTools(), middlewares); - // Add the assistant message with tool requests Message assistantMessage = response.getMessage(); List updatedMessages = new ArrayList<>(request.getMessages()); updatedMessages.add(assistantMessage); - // Add tool response message Message toolResponseMessage = new Message(); toolResponseMessage.setRole(Role.TOOL); toolResponseMessage.setContent(toolResponseParts); updatedMessages.add(toolResponseMessage); - // Update request with new messages for next turn request = ModelRequest.builder() .messages(updatedMessages) @@ -209,6 +234,89 @@ public ModelResponse run( throw new GenkitException("Max tool execution turns (" + maxTurns + ") exceeded"); } + /** + * Resolves middleware references to fresh per-call middleware instances by looking them up in the + * registry's {@code "middleware"} value bucket. Each reference is a {@code {name, config?}} + * object (the {@code MiddlewareRef} shape the Dev UI's Middleware panel sends); a bare JSON + * string is also accepted as a name-only reference. The optional {@code config} is passed + * opaquely to the descriptor's {@link GenerationMiddlewareDesc#instantiate(JsonNode)} — no + * server-side validation is performed; defaults live inside the middleware. Unknown names are + * logged and skipped. + */ + private List resolveMiddlewares(List refs) + throws GenkitException { + if (refs == null || refs.isEmpty()) { + return List.of(); + } + List resolved = new ArrayList<>(refs.size()); + for (JsonNode ref : refs) { + if (ref == null || ref.isNull()) continue; + String name; + JsonNode config = null; + if (ref.isTextual()) { + name = ref.asText(); + } else if (ref.isObject() && ref.hasNonNull("name")) { + name = ref.get("name").asText(); + config = ref.get("config"); // may be absent (null) or JSON null + } else { + logger.warn("Unrecognized middleware reference shape: {}", ref); + continue; + } + if (name == null || name.isEmpty()) continue; + Object value = registry.lookupValue("middleware", name); + if (value instanceof GenerationMiddlewareDesc) { + // Bind config -> a fresh hooks instance, so config and per-call state are per-invocation. + resolved.add(((GenerationMiddlewareDesc) value).instantiate(config)); + } else if (value instanceof GenerationMiddleware) { + // Legacy: a bare middleware registered without a descriptor. Fresh instance per call. + resolved.add(((GenerationMiddleware) value).newInstance()); + } else { + logger.warn( + "Middleware '{}' was requested but is not registered. " + + "Register it via a MiddlewarePlugin or Genkit.Builder.middleware(...).", + name); + } + } + return resolved; + } + + /** Chains wrapGenerate hooks. First middleware is outermost. */ + private static GenerateNext chainGenerate( + List middlewares, GenerateNext core) { + if (middlewares.isEmpty()) return core; + GenerateNext current = core; + for (int i = middlewares.size() - 1; i >= 0; i--) { + final GenerationMiddleware mw = middlewares.get(i); + final GenerateNext next = current; + current = (ctx, params) -> mw.wrapGenerate(ctx, params, next); + } + return current; + } + + /** Chains wrapModel hooks. First middleware is outermost. */ + private static ModelNext chainModel(List middlewares, ModelNext core) { + if (middlewares.isEmpty()) return core; + ModelNext current = core; + for (int i = middlewares.size() - 1; i >= 0; i--) { + final GenerationMiddleware mw = middlewares.get(i); + final ModelNext next = current; + current = (ctx, params) -> mw.wrapModel(ctx, params, next); + } + return current; + } + + /** Chains wrapTool hooks. First middleware is outermost. */ + private static ToolNext chainTool(List middlewares, ToolNext core) { + if (middlewares.isEmpty()) return core; + ToolNext current = core; + for (int i = middlewares.size() - 1; i >= 0; i--) { + final GenerationMiddleware mw = middlewares.get(i); + final ToolNext next = current; + current = (ctx, params) -> mw.wrapTool(ctx, params, next); + } + return current; + } + /** Extracts tool request parts from a model response. */ private List extractToolRequestParts(ModelResponse response) { List toolRequestParts = new ArrayList<>(); @@ -224,20 +332,47 @@ private List extractToolRequestParts(ModelResponse response) { return toolRequestParts; } - /** Executes tools and returns the response parts. */ + /** Executes tools through the wrapTool middleware chain and returns the response parts. */ private List executeTools( - ActionContext ctx, List toolRequestParts, List toolNames) { + ActionContext ctx, + List toolRequestParts, + List toolNames, + List middlewares) { List responseParts = new ArrayList<>(); + // Core tool invocation — runs after all wrapTool middleware + ToolNext toolCore = + (tctx, tparams) -> { + Tool tool = tparams.getTool(); + ToolRequest toolReq = tparams.getRequest(); + Object toolInput = toolReq.getInput(); + + // Convert input if necessary. Use JsonUtils.convert (the centrally configured mapper) + // rather than a local ObjectMapper, so custom (de)serializers, date formats, and naming + // strategies registered globally are honored — consistent with Genkit.java. + if (toolInput instanceof Map + && tool.getInputClass() != null + && !Map.class.isAssignableFrom(tool.getInputClass())) { + toolInput = JsonUtils.convert(toolInput, tool.getInputClass()); + } + + @SuppressWarnings("unchecked") + Tool typedTool = (Tool) tool; + Object result = typedTool.run(tctx, toolInput); + + Part responsePart = new Part(); + responsePart.setToolResponse( + new ToolResponse(toolReq.getRef(), toolReq.getName(), result)); + return responsePart; + }; + ToolNext wrappedTool = chainTool(middlewares, toolCore); + for (Part toolRequestPart : toolRequestParts) { ToolRequest toolRequest = toolRequestPart.getToolRequest(); String toolName = toolRequest.getName(); - Object toolInput = toolRequest.getInput(); - // Find the tool Tool tool = findTool(toolName, toolNames); if (tool == null) { - // Tool not found, create an error response Part errorPart = new Part(); ToolResponse errorResponse = new ToolResponse( @@ -249,26 +384,8 @@ private List executeTools( } try { - // Execute the tool - @SuppressWarnings("unchecked") - Tool typedTool = (Tool) tool; - - // Convert input if necessary - Object convertedInput = toolInput; - if (toolInput instanceof Map - && tool.getInputClass() != null - && !Map.class.isAssignableFrom(tool.getInputClass())) { - convertedInput = objectMapper.convertValue(toolInput, tool.getInputClass()); - } - - Object result = typedTool.run(ctx, convertedInput); - - // Create tool response part - Part responsePart = new Part(); - ToolResponse toolResponse = new ToolResponse(toolRequest.getRef(), toolName, result); - responsePart.setToolResponse(toolResponse); + Part responsePart = wrappedTool.apply(ctx, new ToolParams(toolRequestPart, tool)); responseParts.add(responsePart); - logger.debug("Executed tool '{}' successfully", toolName); } catch (Exception e) { logger.error("Tool execution failed for '{}': {}", toolName, e.getMessage()); diff --git a/ai/src/main/java/com/google/genkit/ai/GenerateActionOptions.java b/ai/src/main/java/com/google/genkit/ai/GenerateActionOptions.java index e3ebd52fb..f301a37f2 100644 --- a/ai/src/main/java/com/google/genkit/ai/GenerateActionOptions.java +++ b/ai/src/main/java/com/google/genkit/ai/GenerateActionOptions.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.JsonNode; import java.util.ArrayList; import java.util.List; @@ -81,6 +82,23 @@ public class GenerateActionOptions { @JsonProperty("stepName") private String stepName; + /** + * Middleware references. In the JS SDK, this is a list of {@code ModelMiddleware} functions or + * name strings. The Dev UI populates this field with the names of middlewares the user has + * selected in the Middleware panel (which correspond to middlewares registered via {@code + * Genkit.Builder.middleware(...)} under the {@code "middleware"} value bucket). + * + *

At runtime, {@link GenerateAction#run} resolves each name via {@code + * registry.lookupValue("middleware", name)} and dispatches the {@code wrapGenerate}/{@code + * wrapModel}/{@code wrapTool} hooks around the model invocation. + */ + /** + * Each element may be a JSON string (middleware name) or a JSON object with a {@code name} field, + * matching the shape sent by the Dev UI's Middleware panel. + */ + @JsonProperty("use") + private List use; + /** Default constructor for JSON deserialization. */ public GenerateActionOptions() {} @@ -109,6 +127,7 @@ public GenerateActionOptions withMessages(List newMessages) { copy.returnToolRequests = this.returnToolRequests; copy.maxTurns = this.maxTurns; copy.stepName = this.stepName; + copy.use = this.use; return copy; } @@ -216,4 +235,12 @@ public String getStepName() { public void setStepName(String stepName) { this.stepName = stepName; } + + public List getUse() { + return use; + } + + public void setUse(List use) { + this.use = use; + } } diff --git a/ai/src/main/java/com/google/genkit/ai/middleware/GenerationMiddlewareDesc.java b/ai/src/main/java/com/google/genkit/ai/middleware/GenerationMiddlewareDesc.java new file mode 100644 index 000000000..e5009000a --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/middleware/GenerationMiddlewareDesc.java @@ -0,0 +1,103 @@ +/* + * 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.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.core.GenkitException; +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * A registered, self-describing, parameterized middleware factory. + * + *

This is the descriptor that gets stored in the registry's {@code "middleware"} value bucket + * and surfaced to the Genkit Dev UI via the reflection API. It mirrors the JS SDK's {@code + * GenerateMiddleware} descriptor (whose {@code toJson()} produces {@code MiddlewareDesc}) and the + * Go SDK's {@code MiddlewareDesc}. + * + *

Unlike a bare {@link GenerationMiddleware} (which is the runtime hooks bundle), a + * descriptor carries metadata for discovery and a {@link #configSchema()} that lets the Dev UI + * render a parameters form. When the user selects a middleware in the Dev UI (optionally filling in + * parameters), the selection is sent back as a {@code {name, config}} reference and resolved via + * {@link #instantiate(JsonNode)}. + * + *

Build descriptors with {@link GenerationMiddlewares#define}. Plugins expose them via {@link + * MiddlewarePlugin}. + */ +public interface GenerationMiddlewareDesc { + + /** Returns the middleware's unique name. Must equal the key it is registered under. */ + String name(); + + /** Returns a human-readable description, or {@code null} if none. */ + default String description() { + return null; + } + + /** + * Returns the JSON Schema describing this middleware's configuration parameters, or {@code null} + * if the middleware takes no parameters. The Dev UI renders a form from this schema. + */ + default Map configSchema() { + return null; + } + + /** Returns arbitrary metadata, or {@code null} if none. */ + default Map metadata() { + return null; + } + + /** + * Instantiates a fresh {@link GenerationMiddleware} (hooks bundle) bound to the given + * configuration. + * + *

Mirrors the JS {@code def.instantiate({config})} and Go {@code buildFromJSON(configJSON)}. + * The {@code config} is applied opaquely — no server-side validation is performed; defaults live + * in the middleware/config type. A fresh instance is returned per call so per-invocation state is + * isolated. + * + * @param config the configuration as a JSON node (from the {@code use[].config} field), or {@code + * null}/JSON null when no configuration was supplied + * @return a fresh middleware instance + * @throws GenkitException if the configuration cannot be bound + */ + GenerationMiddleware instantiate(JsonNode config) throws GenkitException; + + /** + * Serializes this descriptor to the JSON shape expected by the reflection API / Dev UI: {@code + * {name, description?, configSchema?, metadata?}} (null fields omitted). Matches the JS/Go {@code + * MiddlewareDesc} wire shape. + * + * @return the serialized descriptor + */ + default Map toJson() { + Map json = new LinkedHashMap<>(); + json.put("name", name()); + if (description() != null) { + json.put("description", description()); + } + if (configSchema() != null) { + json.put("configSchema", configSchema()); + } + if (metadata() != null) { + json.put("metadata", metadata()); + } + return json; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/middleware/GenerationMiddlewares.java b/ai/src/main/java/com/google/genkit/ai/middleware/GenerationMiddlewares.java new file mode 100644 index 000000000..5de02122c --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/middleware/GenerationMiddlewares.java @@ -0,0 +1,149 @@ +/* + * 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.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.SchemaUtils; +import java.util.Map; +import java.util.function.Function; +import java.util.function.Supplier; + +/** + * Factory helpers for building {@link GenerationMiddlewareDesc} descriptors. + * + *

These mirror the JS SDK's {@code generateMiddleware(...)} helper and the Go SDK's {@code + * NewMiddleware(description, prototype)}: a descriptor pairs discovery metadata (name, description) + * and a {@code configSchema} with a factory that binds config to a fresh {@link + * GenerationMiddleware} hooks bundle. + */ +public final class GenerationMiddlewares { + + private GenerationMiddlewares() {} + + /** + * Defines a parameterized middleware. The config JSON Schema is inferred from {@code configClass} + * (via {@link SchemaUtils#inferSchema}), so the Dev UI can render a parameters form. At resolve + * time the incoming config JSON is deserialized onto a fresh {@code configClass} instance + * (missing fields keep the class's field defaults, mirroring JS/Go where defaults live in the + * middleware), then handed to {@code factory}. + * + * @param name the unique middleware name + * @param description a human-readable description + * @param configClass the configuration POJO type (its fields are the parameters) + * @param factory builds a fresh middleware from a bound config instance + * @param the configuration type + * @return the descriptor + */ + public static GenerationMiddlewareDesc define( + String name, + String description, + Class configClass, + Function factory) { + Map schema = SchemaUtils.inferSchema(configClass); + return new GenerationMiddlewareDesc() { + @Override + public String name() { + return name; + } + + @Override + public String description() { + return description; + } + + @Override + public Map configSchema() { + return schema; + } + + @Override + public GenerationMiddleware instantiate(JsonNode config) throws GenkitException { + C cfg = + (config == null || config.isNull()) + ? newDefault(configClass) + : JsonUtils.fromJsonNode(config, configClass); + return factory.apply(cfg); + } + }; + } + + /** + * Defines a parameterless middleware (no {@code configSchema}). + * + * @param name the unique middleware name + * @param description a human-readable description + * @param factory builds a fresh middleware instance + * @return the descriptor + */ + public static GenerationMiddlewareDesc define( + String name, String description, Supplier factory) { + return new GenerationMiddlewareDesc() { + @Override + public String name() { + return name; + } + + @Override + public String description() { + return description; + } + + @Override + public GenerationMiddleware instantiate(JsonNode config) { + return factory.get(); + } + }; + } + + /** + * Wraps an already-instantiated {@link GenerationMiddleware} in a parameterless descriptor. Used + * for backward compatibility so live middleware passed to {@code Genkit.Builder.middleware(...)} + * or {@code GenerateOptions.use(...)} still shows up in the Dev UI (without a parameters form). + * + * @param mw the live middleware to wrap + * @return the descriptor + */ + public static GenerationMiddlewareDesc of(GenerationMiddleware mw) { + return new GenerationMiddlewareDesc() { + @Override + public String name() { + return mw.name(); + } + + @Override + public GenerationMiddleware instantiate(JsonNode config) { + return mw.newInstance(); + } + }; + } + + private static C newDefault(Class configClass) throws GenkitException { + try { + return configClass.getDeclaredConstructor().newInstance(); + } catch (ReflectiveOperationException e) { + throw new GenkitException( + "Middleware config type " + + configClass.getName() + + " must have a public no-arg constructor", + e); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/middleware/MiddlewarePlugin.java b/ai/src/main/java/com/google/genkit/ai/middleware/MiddlewarePlugin.java new file mode 100644 index 000000000..5f2f859de --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/middleware/MiddlewarePlugin.java @@ -0,0 +1,58 @@ +/* + * 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.core.Registry; +import java.util.List; + +/** + * Implemented by plugins that share reusable generation middleware. + * + *

This is how Middleware V2 is shared as plugins, mirroring the JS SDK's {@code + * GenkitPluginV2.middleware()} and the Go SDK's {@code MiddlewarePlugin.Middlewares()}. A plugin + * implements both {@link com.google.genkit.core.Plugin} (for its actions, if any) and this + * interface; during {@code Genkit} initialization each returned descriptor is registered into the + * registry's {@code "middleware"} value bucket, making it discoverable in the Genkit Dev UI and + * resolvable by name at generate time. + * + *

Example: + * + * 

{@code
+ * public class MyMiddlewarePlugin implements Plugin, MiddlewarePlugin {
+ *   public List middlewares(Registry registry) {
+ *     return List.of(
+ *         GenerationMiddlewares.define("retry", "Retry failed model calls",
+ *             RetryOptions.class, RetryMiddleware::new));
+ *   }
+ *   // ...Plugin methods...
+ * }
+ * }
+ */ +public interface MiddlewarePlugin { + + /** + * Returns the generation middleware this plugin provides. Called once during {@code Genkit} + * initialization. + * + * @param registry the Genkit registry, for middleware that needs to resolve dependencies (e.g. a + * fallback middleware that runs other models) + * @return the middleware descriptors to register (may be empty) + */ + List middlewares(Registry registry); +} diff --git a/core/src/main/java/com/google/genkit/core/DefaultRegistry.java b/core/src/main/java/com/google/genkit/core/DefaultRegistry.java index 88465de93..389b2bcbd 100644 --- a/core/src/main/java/com/google/genkit/core/DefaultRegistry.java +++ b/core/src/main/java/com/google/genkit/core/DefaultRegistry.java @@ -35,6 +35,7 @@ public class DefaultRegistry implements Registry { private final Map> actions = new ConcurrentHashMap<>(); private final Map plugins = new ConcurrentHashMap<>(); private final Map values = new ConcurrentHashMap<>(); + private final Map> valuesByType = new ConcurrentHashMap<>(); private final Map> schemas = new ConcurrentHashMap<>(); private final Map partials = new ConcurrentHashMap<>(); private final Map helpers = new ConcurrentHashMap<>(); @@ -90,6 +91,17 @@ public void registerValue(String name, Object value) { logger.debug("Registered value: {}", name); } + @Override + public void registerValue(String type, String name, Object value) { + Map bucket = valuesByType.computeIfAbsent(type, t -> new ConcurrentHashMap<>()); + // Atomic check-then-act: putIfAbsent avoids a race where two threads registering the same + // (type, name) concurrently could both pass a containsKey check and clobber each other. + if (bucket.putIfAbsent(name, value) != null) { + throw new IllegalStateException("Value already registered: " + type + "/" + name); + } + logger.debug("Registered value: {}/{}", type, name); + } + @Override public void registerSchema(String name, Map schema) { if (schemas.containsKey(name)) { @@ -126,6 +138,16 @@ public Object lookupValue(String name) { return value; } + @Override + public Object lookupValue(String type, String name) { + Map bucket = valuesByType.get(type); + Object value = bucket != null ? bucket.get(name) : null; + if (value == null && parent != null) { + value = parent.lookupValue(type, name); + } + return value; + } + @Override public Map lookupSchema(String name) { Map schema = schemas.get(name); @@ -252,6 +274,19 @@ public Map listValues() { return allValues; } + @Override + public Map listValues(String type) { + Map allValues = new LinkedHashMap<>(); + if (parent != null) { + allValues.putAll(parent.listValues(type)); + } + Map bucket = valuesByType.get(type); + if (bucket != null) { + allValues.putAll(bucket); + } + return allValues; + } + @Override public void registerPartial(String name, String source) { partials.put(name, source); diff --git a/core/src/main/java/com/google/genkit/core/Registry.java b/core/src/main/java/com/google/genkit/core/Registry.java index 431fdf8c2..3150f443c 100644 --- a/core/src/main/java/com/google/genkit/core/Registry.java +++ b/core/src/main/java/com/google/genkit/core/Registry.java @@ -72,7 +72,7 @@ public interface Registry { void registerAction(String key, Action action); /** - * Records an arbitrary value in the registry. + * Records an arbitrary value in the registry under the default type bucket. * * @param name the value name * @param value the value to register @@ -80,6 +80,20 @@ public interface Registry { */ void registerValue(String name, Object value); + /** + * Records an arbitrary value in the registry under the given type bucket. + * + *

This mirrors the JS reflection API which keys values by {@code (type, name)} (e.g. {@code + * type="middleware"}, {@code type="defaultModel"}). Values registered here are exposed via the + * {@code /api/values?type=...} reflection endpoint and surfaced in the Genkit Dev UI. + * + * @param type the value type bucket (e.g. {@code "middleware"}) + * @param name the value name + * @param value the value to register + * @throws IllegalStateException if a value with the same (type, name) is already registered + */ + void registerValue(String type, String name, Object value); + /** * Records a JSON schema in the registry. * @@ -127,6 +141,16 @@ public interface Registry { */ Object lookupValue(String name); + /** + * Returns the value for the given type bucket and name. It first checks the current registry, + * then falls back to the parent if not found. + * + * @param type the value type bucket (e.g. {@code "middleware"}) + * @param name the value name + * @return the value, or null if not found + */ + Object lookupValue(String type, String name); + /** * Returns a JSON schema for the given name. It first checks the current registry, then falls back * to the parent if not found. @@ -190,12 +214,21 @@ default void registerAction(ActionType type, Action action) { List listPlugins(); /** - * Returns a map of all registered values. + * Returns a map of all registered values in the default type bucket. * * @return map of all registered values */ Map listValues(); + /** + * Returns a map of all registered values under the given type bucket. This includes values from + * both the current registry and its parent hierarchy. + * + * @param type the value type bucket (e.g. {@code "middleware"}) + * @return map of values keyed by name, or empty map if none + */ + Map listValues(String type); + /** * Registers a partial template for use with prompts. * diff --git a/docs/src/content/docs/middleware.md b/docs/src/content/docs/middleware.md index 89c945724..30b419f8f 100644 --- a/docs/src/content/docs/middleware.md +++ b/docs/src/content/docs/middleware.md @@ -314,6 +314,114 @@ ModelResponse response = genkit.generate( Middleware order matters — the **first** middleware listed is **outermost** (runs first on the way in, last on the way out). +### Using middleware from the Dev UI + +To make middleware selectable from the **Middleware** panel in the Genkit Dev UI, register it with the `Genkit` builder via `.middleware(...)`: + +```java +GenerationMiddleware modelLogging = new ModelLoggingMiddleware(); +GenerationMiddleware timing = new GenerateTimingMiddleware(); +GenerationMiddleware toolMonitor = new ToolMonitorMiddleware(); + +Genkit genkit = Genkit.builder() + .plugin(OpenAIPlugin.create()) + .middleware(modelLogging, timing, toolMonitor) + .build(); +``` + +Each registered middleware appears in the Dev UI Middleware panel by `name()`. When you select one or more middlewares in the panel and run a model from the **Models** runner, the Dev UI invokes the `/util/generate` action with the selected middleware names in the `use` field. The action resolves each name back to the registered middleware via the registry and runs it through the same `wrapGenerate` / `wrapModel` / `wrapTool` chain that applies to programmatic `generate()` calls. + +A fresh middleware instance (via `newInstance()`) is created per Dev UI invocation, so per-request state (counters, timers) is isolated just as it is for code-driven calls. + +> **Note:** `.middleware(...)` only controls Dev UI visibility. Middleware attached programmatically with `GenerateOptions.builder().use(...)` does not need to be registered with the builder. + +### Sharing middleware as a plugin + +Beyond the ad-hoc `.middleware(...)` method above, middleware can be **shared as a plugin** — the same way it works in the Genkit JS and Go SDKs. A plugin that provides middleware implements `MiddlewarePlugin` (alongside `Plugin`) and returns a list of **middleware descriptors** built with `GenerationMiddlewares.define(...)`: + +```java +public class MyMiddlewarePlugin implements Plugin, MiddlewarePlugin { + + @Override + public String getName() { return "my-middleware"; } + + @Override + public List> init() { return List.of(); } + + @Override + public List middlewares(Registry registry) { + return List.of( + GenerationMiddlewares.define( + "retry", // name + "Retry failed model calls.", // description + RetryMiddleware.Options.class, // config type (schema auto-inferred) + RetryMiddleware::new)); // factory: config -> middleware + } +} +``` + +Add the plugin to the builder and every middleware it provides is registered automatically: + +```java +Genkit genkit = Genkit.builder() + .plugin(OpenAIPlugin.create()) + .plugin(new MyMiddlewarePlugin()) + .build(); +``` + +### Parameterized middleware (Dev UI parameters form) + +A middleware defined with a config type exposes a **parameters form** in the Dev UI. The config type is a plain POJO whose fields become the parameters; the JSON Schema is inferred automatically (use `@JsonPropertyDescription` to document each field): + +```java +public class RetryMiddleware extends BaseGenerationMiddleware { + + private final Options options; + + public RetryMiddleware(Options options) { this.options = options; } + + @Override public String name() { return "retry"; } + @Override public GenerationMiddleware newInstance() { return this; } + + @Override + public ModelResponse wrapModel(ActionContext ctx, ModelParams params, ModelNext next) + throws GenkitException { + // ...use options.maxRetries etc... + return next.apply(ctx, params); + } + + /** These fields render as a form in the Dev UI Middleware panel. */ + public static class Options { + @JsonProperty("maxRetries") + @JsonPropertyDescription("Maximum number of retries after the initial attempt.") + public int maxRetries = 3; + + public Options() {} + } +} +``` + +In the Dev UI, selecting `retry` shows a form for `maxRetries`; the values you enter are sent to the `/util/generate` action as `use: [{ "name": "retry", "config": { "maxRetries": 2 } }]`. The runtime resolves the descriptor by name, binds the `config` onto a fresh `Options` instance (missing fields keep their defaults), and runs the resulting middleware. Config is applied as-is — defaults live in the middleware, so there is no separate validation step. + +### Built-in middleware plugin + +The `genkit-plugin-middleware` module ships a ready-to-use plugin, `GenerationMiddlewarePlugin`, with parameterized middleware (the Java equivalent of the JS `@genkit-ai/middleware` package): + +```java +Genkit genkit = Genkit.builder() + .plugin(OpenAIPlugin.create()) + .plugin(GenerationMiddlewarePlugin.create()) + .build(); +``` + +| Name | Hook | Parameters | +|------|------|------------| +| `retry` | `wrapModel` | `maxRetries`, `initialDelayMs`, `maxDelayMs`, `backoffFactor` | +| `fallback` | `wrapGenerate` | `models` (ordered list of fallback model names) | +| `simulateSystemPrompt` | `wrapGenerate` | `preface`, `acknowledgement` | + +Each appears in the Dev UI Middleware panel with a parameters form, and can also be used programmatically, e.g. `new RetryMiddleware(opts)` passed to `GenerateOptions.builder().use(...)`. + ### Multi-hook middleware A single middleware can implement all three hooks to observe every stage: diff --git a/genkit/src/main/java/com/google/genkit/Genkit.java b/genkit/src/main/java/com/google/genkit/Genkit.java index b1eec3c6c..9e639b010 100644 --- a/genkit/src/main/java/com/google/genkit/Genkit.java +++ b/genkit/src/main/java/com/google/genkit/Genkit.java @@ -679,6 +679,10 @@ private ModelResponse generateInternal( int maxTurns = options.getMaxTurns() != null ? options.getMaxTurns() : 5; + // Auto-register any middleware passed via .use(...) so it shows up in the Dev UI + // Middleware panel. Registration is idempotent (last write wins for a given name). + registerMiddlewareForDevUi(options.getUse()); + // Create fresh middleware instances for this invocation List middlewares = createMiddlewareInstances(options.getUse()); @@ -941,6 +945,65 @@ private List createMiddlewareInstances(List use) { + if (use == null || use.isEmpty()) { + return; + } + for (GenerationMiddleware mw : use) { + if (mw == null) continue; + String name = mw.name(); + if (name == null || name.isEmpty()) continue; + registerMiddlewareDesc(GenerationMiddlewares.of(mw)); + } + } + + /** + * Registers a middleware descriptor into the {@code "middleware"} value bucket so the Dev UI can + * list it (with a parameters form derived from its {@code configSchema}) and resolve it by name + * at generate time. + * + *

Idempotent — a descriptor already registered under the same name is kept (first registration + * wins). Concurrency-safe — the {@code lookupValue}/{@code registerValue} pair is not atomic, so + * a concurrent registration of the same name is caught and ignored rather than crashing the + * generation request with the registry's duplicate-key {@link IllegalStateException}. + */ + private void registerMiddlewareDesc(GenerationMiddlewareDesc desc) { + if (desc == null) return; + String name = desc.name(); + if (name == null || name.isEmpty()) return; + if (registry.lookupValue("middleware", name) != null) { + return; + } + try { + registry.registerValue("middleware", name, desc); + } catch (IllegalStateException e) { + // Another thread registered the same middleware concurrently — safe to ignore. + } + } + + /** + * Registers middleware shared by plugins implementing {@link MiddlewarePlugin} into the {@code + * "middleware"} value bucket. Called once during builder {@code build()} after plugin + * initialization, mirroring JS {@code GenkitPluginV2.middleware()} and Go {@code + * MiddlewarePlugin.Middlewares()}. + */ + private void registerPluginMiddlewares() { + for (Plugin plugin : plugins) { + if (plugin instanceof MiddlewarePlugin) { + List descs = ((MiddlewarePlugin) plugin).middlewares(registry); + if (descs == null) continue; + for (GenerationMiddlewareDesc desc : descs) { + registerMiddlewareDesc(desc); + } + } + } + } + /** * Converts {@link GenerateOptions} to a high-level {@link GenerateActionOptions}. * @@ -2335,6 +2398,8 @@ public EvalStore getEvalStore() { /** Builder for Genkit. */ public static class Builder { private final List plugins = new ArrayList<>(); + private final List middlewares = new ArrayList<>(); + private final List middlewareDescs = new ArrayList<>(); private GenkitOptions options = GenkitOptions.builder().build(); /** @@ -2359,6 +2424,39 @@ public Builder plugin(Plugin plugin) { return this; } + /** + * Optional: pre-registers one or more generation middlewares so they show up in the Genkit Dev + * UI Middleware panel before any flow has executed. This is a UX convenience only — + * middlewares are also auto-registered the first time they appear in a {@code + * GenerateOptions.use(...)} call, so production code does not need to declare them here. + * + * @param middlewares the middlewares to pre-register + * @return this builder + */ + public Builder middleware(GenerationMiddleware... middlewares) { + for (GenerationMiddleware mw : middlewares) { + this.middlewares.add(mw); + } + return this; + } + + /** + * Optional: pre-registers one or more middleware descriptors so they show up in the + * Genkit Dev UI Middleware panel — including a parameters form derived from each descriptor's + * {@code configSchema}. Use this for parameterized middleware defined via {@link + * com.google.genkit.ai.middleware.GenerationMiddlewares#define}. For middleware shared by a + * plugin, prefer implementing {@link com.google.genkit.ai.middleware.MiddlewarePlugin} instead. + * + * @param descriptors the middleware descriptors to pre-register + * @return this builder + */ + public Builder middleware(GenerationMiddlewareDesc... descriptors) { + for (GenerationMiddlewareDesc desc : descriptors) { + this.middlewareDescs.add(desc); + } + return this; + } + /** * Enables dev mode. * @@ -2390,6 +2488,16 @@ public Genkit build() { Genkit genkit = new Genkit(options); genkit.plugins.addAll(plugins); genkit.init(); + // Register middleware shared by plugins (those implementing MiddlewarePlugin) so it shows up + // in the Dev UI Middleware panel and is resolvable by name at generate time. Mirrors the JS + // GenkitPluginV2.middleware() / Go MiddlewarePlugin.Middlewares() registration during init. + genkit.registerPluginMiddlewares(); + // Pre-register any middleware declared directly via .middleware(...) so the Dev UI Middleware + // panel can list them before any generate() call runs. + genkit.registerMiddlewareForDevUi(middlewares); + for (GenerationMiddlewareDesc desc : middlewareDescs) { + genkit.registerMiddlewareDesc(desc); + } return genkit; } } diff --git a/genkit/src/main/java/com/google/genkit/ReflectionServer.java b/genkit/src/main/java/com/google/genkit/ReflectionServer.java index 6cc99bcb8..3b9e1b720 100644 --- a/genkit/src/main/java/com/google/genkit/ReflectionServer.java +++ b/genkit/src/main/java/com/google/genkit/ReflectionServer.java @@ -274,6 +274,25 @@ public boolean handle(Request request, Response response, Callback callback) thr } else if (target.startsWith("/api/actions/")) { String actionKey = target.substring("/api/actions/".length()); result = handleGetAction(actionKey); + } else if ("/api/values".equals(target) && "GET".equals(method)) { + String query = request.getHttpURI().getQuery(); + String type = parseQueryParam(query, "type"); + if (type == null) { + status = 400; + result = + createErrorResponse(3, "Query parameter \"type\" is required.", null); // 3=INVALID + } else if (!"middleware".equals(type) && !"defaultModel".equals(type)) { + status = 400; + result = + createErrorResponse( + 3, + "'type' " + + type + + " is not supported. Only 'defaultModel' and 'middleware' are supported", + null); // 3=INVALID + } else { + result = handleListValuesByType(type); + } } else if ("/api/notify".equals(target) && "POST".equals(method)) { String body = readRequestBody(request); result = handleNotify(body); @@ -466,6 +485,65 @@ private String handleGetAction(String actionKey) { return JsonUtils.toJson(actionInfo); } + /** + * Parses a query string and returns the value for the given parameter name, or null if absent. + */ + private String parseQueryParam(String query, String name) { + if (query == null || query.isEmpty()) return null; + for (String pair : query.split("&")) { + int eq = pair.indexOf('='); + String k = eq >= 0 ? pair.substring(0, eq) : pair; + String v = eq >= 0 ? pair.substring(eq + 1) : ""; + if (k.equals(name)) { + try { + return java.net.URLDecoder.decode(v, java.nio.charset.StandardCharsets.UTF_8); + } catch (Exception e) { + return v; + } + } + } + return null; + } + + /** + * Handles {@code GET /api/values?type=}. Returns a map of name -> value JSON + * representation for all values registered under the given type bucket. Mirrors the JS + * reflection API used by the Dev UI Middleware panel (type=middleware) and default-model + * indicator (type=defaultModel). + */ + private String handleListValuesByType(String type) { + Map values = registry.listValues(type); + Map mapped = new HashMap<>(); + if (values != null) { + for (Map.Entry e : values.entrySet()) { + mapped.put(e.getKey(), serializeValue(e.getValue(), e.getKey())); + } + } + return JsonUtils.toJson(mapped); + } + + /** + * Serializes a registered value for the reflection API. Middleware descriptors are serialized + * to their {@code {name, description?, configSchema?, metadata?}} shape (matching the JS/Go + * {@code MiddlewareDesc}), so the Dev UI can render a parameters form from {@code + * configSchema}. A bare {@link com.google.genkit.ai.middleware.GenerationMiddleware} (no + * descriptor) degrades to just its name. + */ + private Object serializeValue(Object value, String fallbackName) { + if (value == null) return null; + if (value instanceof com.google.genkit.ai.middleware.GenerationMiddlewareDesc) { + return ((com.google.genkit.ai.middleware.GenerationMiddlewareDesc) value).toJson(); + } + if (value instanceof com.google.genkit.ai.middleware.GenerationMiddleware) { + com.google.genkit.ai.middleware.GenerationMiddleware mw = + (com.google.genkit.ai.middleware.GenerationMiddleware) value; + Map json = new HashMap<>(); + json.put("name", mw.name() != null ? mw.name() : fallbackName); + return json; + } + return value; + } + private String handleRunAction(String body) throws GenkitException { JsonNode requestNode = JsonUtils.parseJson(body); diff --git a/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java b/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java index d6e145910..713fce969 100644 --- a/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java +++ b/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java @@ -428,9 +428,45 @@ private void handleListActions(String requestId) { private void handleListValues(String requestId, JsonNode params) { if (requestId == null) return; - // Currently no values to list for Java runtime + String type = + (params != null && params.hasNonNull("type")) ? params.get("type").asText() : null; + if (type == null) { + sendError(requestId, -32602, "Query parameter \"type\" is required.", null); + return; + } + if (!"middleware".equals(type) && !"defaultModel".equals(type)) { + sendError( + requestId, + -32602, + "'type' " + + type + + " is not supported. Only 'defaultModel' and 'middleware' are supported", + null); + return; + } + + Map values = new HashMap<>(); + Map raw = registry.listValues(type); + if (raw != null) { + for (Map.Entry e : raw.entrySet()) { + Object v = e.getValue(); + if (v instanceof com.google.genkit.ai.middleware.GenerationMiddlewareDesc) { + // {name, description?, configSchema?, metadata?} — configSchema drives the Dev UI form. + values.put( + e.getKey(), ((com.google.genkit.ai.middleware.GenerationMiddlewareDesc) v).toJson()); + } else if (v instanceof com.google.genkit.ai.middleware.GenerationMiddleware) { + com.google.genkit.ai.middleware.GenerationMiddleware mw = + (com.google.genkit.ai.middleware.GenerationMiddleware) v; + Map json = new HashMap<>(); + json.put("name", mw.name() != null ? mw.name() : e.getKey()); + values.put(e.getKey(), json); + } else { + values.put(e.getKey(), v); + } + } + } Map result = new HashMap<>(); - result.put("values", new HashMap<>()); + result.put("values", values); sendResponse(requestId, result); } diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/FallbackMiddleware.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/FallbackMiddleware.java new file mode 100644 index 000000000..566af66c3 --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/FallbackMiddleware.java @@ -0,0 +1,113 @@ +/* + * 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.plugins.middleware; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonPropertyDescription; +import com.google.genkit.ai.GenerateActionOptions; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.middleware.BaseGenerationMiddleware; +import com.google.genkit.ai.middleware.GenerateNext; +import com.google.genkit.ai.middleware.GenerateParams; +import com.google.genkit.ai.middleware.GenerationMiddleware; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.GenkitException; +import java.util.ArrayList; +import java.util.List; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Falls back to alternate models when the primary model fails. + * + *

Wraps the {@code wrapGenerate} hook: if generation with the primary model (the one on the + * incoming request) throws, each model in {@link Options#models} is tried in order by re-running + * the generate iteration with the request's model swapped. The first success is returned; if all + * fallbacks fail, the original error is rethrown. + * + *

Implemented at the generate level (rather than {@code wrapModel}) because switching model + * requires re-resolving the model action, which the core generate loop already does from {@code + * options.model}. As a consequence a failing iteration re-runs its tool calls under the fallback + * model; put side-effect-free tools before this middleware if that matters. + * + *

Mirrors the JS {@code fallback} middleware in {@code @genkit-ai/middleware} and the Go {@code + * Fallback} middleware. + */ +public class FallbackMiddleware extends BaseGenerationMiddleware { + + private static final Logger logger = LoggerFactory.getLogger(FallbackMiddleware.class); + + private final Options options; + + public FallbackMiddleware(Options options) { + this.options = options != null ? options : new Options(); + } + + @Override + public String name() { + return "fallback"; + } + + @Override + public GenerationMiddleware newInstance() { + return this; + } + + @Override + public ModelResponse wrapGenerate(ActionContext ctx, GenerateParams params, GenerateNext next) + throws GenkitException { + try { + return next.apply(ctx, params); + } catch (GenkitException primary) { + if (options.models == null || options.models.isEmpty()) { + throw primary; + } + GenerateActionOptions req = params.getRequest(); + for (String model : options.models) { + if (model == null || model.isEmpty()) continue; + // Full copy of the request with the model swapped (withMessages preserves all other + // fields). + GenerateActionOptions fallbackReq = req.withMessages(req.getMessages()); + fallbackReq.setModel(model); + try { + logger.warn( + "[fallback] primary model '{}' failed, trying fallback model '{}'", + req.getModel(), + model); + return next.apply(ctx, params.withRequest(fallbackReq)); + } catch (GenkitException e) { + logger.warn("[fallback] fallback model '{}' failed: {}", model, e.getMessage()); + } + } + throw primary; + } + } + + /** Configuration parameters for {@link FallbackMiddleware}. */ + public static class Options { + + @JsonProperty("models") + @JsonPropertyDescription( + "Ordered list of fallback model names to try when the primary model fails," + + " e.g. \"openai/gpt-4o-mini\".") + public List models = new ArrayList<>(); + + public Options() {} + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/GenerationMiddlewarePlugin.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/GenerationMiddlewarePlugin.java new file mode 100644 index 000000000..ee99e4568 --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/GenerationMiddlewarePlugin.java @@ -0,0 +1,107 @@ +/* + * 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.plugins.middleware; + +import com.google.genkit.ai.middleware.GenerationMiddlewareDesc; +import com.google.genkit.ai.middleware.GenerationMiddlewares; +import com.google.genkit.ai.middleware.MiddlewarePlugin; +import com.google.genkit.core.Action; +import com.google.genkit.core.Plugin; +import com.google.genkit.core.Registry; +import java.util.List; + +/** + * A plugin that ships a set of ready-to-use, parameterized generation middleware. + * + *

This is the Java analog of the JS {@code @genkit-ai/middleware} package and the Go {@code + * plugins/middleware} package. Adding it to the {@code Genkit} builder registers each middleware + * into the {@code "middleware"} registry bucket, so it appears in the Genkit Dev UI Middleware + * panel (with a parameters form derived from its {@code configSchema}) and can be attached to any + * {@code generate()} call by name. + * + *

Provided middleware: + * + *

    + *
  • {@code retry} — retry failed model calls with exponential backoff ({@link RetryMiddleware}) + *
  • {@code fallback} — fall back to alternate models on failure ({@link FallbackMiddleware}) + *
  • {@code simulateSystemPrompt} — rewrite the system message into a user/model exchange + * ({@link SimulateSystemPromptMiddleware}) + *
+ * + *

Usage: + * + * 

{@code
+ * Genkit genkit = Genkit.builder()
+ *     .plugin(OpenAIPlugin.create())
+ *     .plugin(GenerationMiddlewarePlugin.create())
+ *     .build();
+ *
+ * // Attach by name (config optional) from code or the Dev UI:
+ * genkit.generate(GenerateOptions.builder()
+ *     .model("openai/gpt-4o-mini")
+ *     .prompt("Hello")
+ *     .build());
+ * }
+ */ +public class GenerationMiddlewarePlugin implements Plugin, MiddlewarePlugin { + + /** The plugin name. */ + public static final String PLUGIN_NAME = "genkit-middleware"; + + /** + * Creates the plugin. + * + * @return a new plugin instance + */ + public static GenerationMiddlewarePlugin create() { + return new GenerationMiddlewarePlugin(); + } + + @Override + public String getName() { + return PLUGIN_NAME; + } + + @Override + public List> init() { + // This plugin provides middleware (via MiddlewarePlugin), not actions. + return List.of(); + } + + @Override + public List middlewares(Registry registry) { + return List.of( + GenerationMiddlewares.define( + "retry", + "Retries failed model calls with exponential backoff.", + RetryMiddleware.Options.class, + RetryMiddleware::new), + GenerationMiddlewares.define( + "fallback", + "Falls back to alternate models when the primary model fails.", + FallbackMiddleware.Options.class, + FallbackMiddleware::new), + GenerationMiddlewares.define( + "simulateSystemPrompt", + "Rewrites the system message into a user/model exchange for models without native" + + " system-prompt support.", + SimulateSystemPromptMiddleware.Options.class, + SimulateSystemPromptMiddleware::new)); + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/RetryMiddleware.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/RetryMiddleware.java new file mode 100644 index 000000000..c291d828c --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/RetryMiddleware.java @@ -0,0 +1,117 @@ +/* + * 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.plugins.middleware; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonPropertyDescription; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.middleware.BaseGenerationMiddleware; +import com.google.genkit.ai.middleware.GenerationMiddleware; +import com.google.genkit.ai.middleware.ModelNext; +import com.google.genkit.ai.middleware.ModelParams; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.GenkitException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Retries a failed model call with exponential backoff. + * + *

Wraps the {@code wrapModel} hook: if the underlying model call throws, it is retried up to + * {@link Options#maxRetries} times, sleeping {@code initialDelayMs} before the first retry and + * multiplying the delay by {@code backoffFactor} (capped at {@code maxDelayMs}) after each attempt. + * + *

Mirrors the JS {@code retry} middleware in {@code @genkit-ai/middleware} and the Go {@code + * Retry} middleware. + */ +public class RetryMiddleware extends BaseGenerationMiddleware { + + private static final Logger logger = LoggerFactory.getLogger(RetryMiddleware.class); + + private final Options options; + + public RetryMiddleware(Options options) { + this.options = options != null ? options : new Options(); + } + + @Override + public String name() { + return "retry"; + } + + @Override + public GenerationMiddleware newInstance() { + // Stateless: the backoff state is local to each wrapModel invocation. + return this; + } + + @Override + public ModelResponse wrapModel(ActionContext ctx, ModelParams params, ModelNext next) + throws GenkitException { + int attempts = 0; + long delay = Math.max(0, options.initialDelayMs); + while (true) { + try { + return next.apply(ctx, params); + } catch (GenkitException e) { + if (attempts >= options.maxRetries) { + throw e; + } + attempts++; + logger.warn( + "[retry] model call failed (retry {}/{}), retrying in {}ms: {}", + attempts, + options.maxRetries, + delay, + e.getMessage()); + if (delay > 0) { + try { + Thread.sleep(delay); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw e; + } + } + delay = Math.min((long) (delay * options.backoffFactor), options.maxDelayMs); + } + } + } + + /** Configuration parameters for {@link RetryMiddleware}. */ + public static class Options { + + @JsonProperty("maxRetries") + @JsonPropertyDescription("Maximum number of retries after the initial attempt.") + public int maxRetries = 3; + + @JsonProperty("initialDelayMs") + @JsonPropertyDescription("Delay before the first retry, in milliseconds.") + public long initialDelayMs = 1000; + + @JsonProperty("maxDelayMs") + @JsonPropertyDescription("Maximum delay between retries, in milliseconds.") + public long maxDelayMs = 60000; + + @JsonProperty("backoffFactor") + @JsonPropertyDescription("Multiplier applied to the delay after each retry.") + public double backoffFactor = 2.0; + + public Options() {} + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/SimulateSystemPromptMiddleware.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/SimulateSystemPromptMiddleware.java new file mode 100644 index 000000000..da8f0c0b6 --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/SimulateSystemPromptMiddleware.java @@ -0,0 +1,119 @@ +/* + * 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.plugins.middleware; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonPropertyDescription; +import com.google.genkit.ai.GenerateActionOptions; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.middleware.BaseGenerationMiddleware; +import com.google.genkit.ai.middleware.GenerateNext; +import com.google.genkit.ai.middleware.GenerateParams; +import com.google.genkit.ai.middleware.GenerationMiddleware; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.GenkitException; +import java.util.ArrayList; +import java.util.List; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Rewrites a {@code system} message into a user/model exchange, for models that do not natively + * support system prompts. + * + *

Wraps the {@code wrapGenerate} hook: each {@code system} message is replaced by a {@code user} + * message ({@link Options#preface} followed by the original system content) and a {@code model} + * message ({@link Options#acknowledgement}). Non-system messages are passed through unchanged. + * + *

Mirrors the JS {@code simulateSystemPrompt} model middleware. + */ +public class SimulateSystemPromptMiddleware extends BaseGenerationMiddleware { + + private static final Logger logger = + LoggerFactory.getLogger(SimulateSystemPromptMiddleware.class); + + private final Options options; + + public SimulateSystemPromptMiddleware(Options options) { + this.options = options != null ? options : new Options(); + } + + @Override + public String name() { + return "simulateSystemPrompt"; + } + + @Override + public GenerationMiddleware newInstance() { + return this; + } + + @Override + public ModelResponse wrapGenerate(ActionContext ctx, GenerateParams params, GenerateNext next) + throws GenkitException { + GenerateActionOptions req = params.getRequest(); + List messages = req.getMessages(); + if (messages == null || messages.isEmpty()) { + return next.apply(ctx, params); + } + + boolean hasSystem = messages.stream().anyMatch(m -> m != null && m.getRole() == Role.SYSTEM); + if (!hasSystem) { + return next.apply(ctx, params); + } + + List rewritten = new ArrayList<>(messages.size() + 2); + for (Message m : messages) { + if (m != null && m.getRole() == Role.SYSTEM) { + List userParts = new ArrayList<>(); + if (options.preface != null && !options.preface.isEmpty()) { + userParts.add(Part.text(options.preface)); + } + if (m.getContent() != null) { + userParts.addAll(m.getContent()); + } + rewritten.add(new Message(Role.USER, userParts)); + rewritten.add(new Message(Role.MODEL, List.of(Part.text(options.acknowledgement)))); + } else { + rewritten.add(m); + } + } + + logger.debug("[simulateSystemPrompt] rewrote system message(s) into a user/model exchange"); + return next.apply(ctx, params.withRequest(req.withMessages(rewritten))); + } + + /** Configuration parameters for {@link SimulateSystemPromptMiddleware}. */ + public static class Options { + + @JsonProperty("preface") + @JsonPropertyDescription( + "Text prepended to the system content when converting it into a user message.") + public String preface = "System Instructions:\n"; + + @JsonProperty("acknowledgement") + @JsonPropertyDescription("Assistant reply inserted after the converted system message.") + public String acknowledgement = "Understood."; + + public Options() {} + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java index 5e8d28ba2..4f8ce940b 100644 --- a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java @@ -19,9 +19,17 @@ /** * Middleware plugin for Genkit providing higher-level generation building blocks. * - *

This plugin provides: + *

This module provides: * *

    + *
  • {@link com.google.genkit.plugins.middleware.GenerationMiddlewarePlugin} — a {@code + * MiddlewarePlugin} that ships ready-to-use, parameterized Generation Middleware V2 (the Java + * equivalent of the JS {@code @genkit-ai/middleware} package): {@link + * com.google.genkit.plugins.middleware.RetryMiddleware retry}, {@link + * com.google.genkit.plugins.middleware.FallbackMiddleware fallback}, and {@link + * com.google.genkit.plugins.middleware.SimulateSystemPromptMiddleware simulateSystemPrompt}. + * Each is registered into the {@code "middleware"} bucket and appears in the Genkit Dev UI + * Middleware panel with a parameters form. *
  • {@link com.google.genkit.plugins.middleware.Agents} — sub-agent delegation, where * each configured sub-agent is exposed to the model as a {@code delegate_to_} tool that * runs the sub-agent for a single turn and returns its text (plus optional artifacts). @@ -29,11 +37,11 @@ * {@code write_artifact} tools operating on the active agent session's artifact store. *
* - *

Both are implemented as tool factories: they return {@code List>} (and, - * for {@code agents()}, a system-prompt fragment) that callers wire into an agent via {@code - * AgentConfig.tools(...)} and {@code AgentConfig.system(...)}. This avoids modifying the generate - * pipeline while still letting a model delegate to sub-agents and read/write artifacts. + *

{@code Agents} and {@code Artifacts} are implemented as tool factories: they return + * {@code List>} (and, for {@code agents()}, a system-prompt fragment) that callers wire + * into an agent via {@code AgentConfig.tools(...)} and {@code AgentConfig.system(...)}. * + * @see com.google.genkit.plugins.middleware.GenerationMiddlewarePlugin * @see com.google.genkit.plugins.middleware.Agents * @see com.google.genkit.plugins.middleware.Artifacts */ diff --git a/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/GenerationMiddlewarePluginTest.java b/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/GenerationMiddlewarePluginTest.java new file mode 100644 index 000000000..251dca6f6 --- /dev/null +++ b/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/GenerationMiddlewarePluginTest.java @@ -0,0 +1,150 @@ +/* + * 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.plugins.middleware; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.middleware.GenerationMiddleware; +import com.google.genkit.ai.middleware.GenerationMiddlewareDesc; +import com.google.genkit.ai.middleware.ModelNext; +import com.google.genkit.ai.middleware.ModelParams; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import org.junit.jupiter.api.Test; + +/** Tests the descriptor/plugin mechanism and the built-in parameterized middleware. */ +class GenerationMiddlewarePluginTest { + + private Map descriptors() { + return GenerationMiddlewarePlugin.create().middlewares(new DefaultRegistry()).stream() + .collect(Collectors.toMap(GenerationMiddlewareDesc::name, d -> d)); + } + + @Test + void listsThreeParameterizedMiddleware() { + Map byName = descriptors(); + assertEquals(3, byName.size()); + assertTrue(byName.containsKey("retry")); + assertTrue(byName.containsKey("fallback")); + assertTrue(byName.containsKey("simulateSystemPrompt")); + } + + @Test + void descriptorSerializesToDevUiShape() { + GenerationMiddlewareDesc retry = descriptors().get("retry"); + + Map json = retry.toJson(); + assertEquals("retry", json.get("name")); + assertNotNull(json.get("description")); + + @SuppressWarnings("unchecked") + Map configSchema = (Map) json.get("configSchema"); + assertNotNull(configSchema, "retry must expose a configSchema so the Dev UI renders a form"); + assertEquals("object", configSchema.get("type")); + + @SuppressWarnings("unchecked") + Map props = (Map) configSchema.get("properties"); + assertNotNull(props); + assertTrue(props.containsKey("maxRetries")); + assertTrue(props.containsKey("initialDelayMs")); + assertTrue(props.containsKey("backoffFactor")); + } + + @Test + void instantiateBindsConfigAndDefaults() throws Exception { + GenerationMiddlewareDesc retry = descriptors().get("retry"); + + JsonNode config = JsonUtils.parseJson("{\"maxRetries\":7,\"initialDelayMs\":250}"); + GenerationMiddleware bound = retry.instantiate(config); + assertInstanceOf(RetryMiddleware.class, bound); + assertEquals("retry", bound.name()); + + // Null config -> defaults (no exception). + GenerationMiddleware defaults = retry.instantiate(null); + assertNotNull(defaults); + } + + @Test + void retryRetriesThenSucceeds() throws Exception { + RetryMiddleware.Options opts = new RetryMiddleware.Options(); + opts.maxRetries = 2; + opts.initialDelayMs = 1; + opts.maxDelayMs = 1; + opts.backoffFactor = 1; + RetryMiddleware retry = new RetryMiddleware(opts); + + AtomicInteger calls = new AtomicInteger(); + ModelResponse success = new ModelResponse(); + ModelNext next = + (ctx, params) -> { + if (calls.incrementAndGet() < 3) { + throw new GenkitException("transient failure"); + } + return success; + }; + + ModelResponse result = retry.wrapModel(ctx(), new ModelParams(null, null), next); + assertEquals(success, result); + assertEquals(3, calls.get(), "1 initial attempt + 2 retries"); + } + + @Test + void retryExhaustsThenRethrows() { + RetryMiddleware.Options opts = new RetryMiddleware.Options(); + opts.maxRetries = 2; + opts.initialDelayMs = 1; + RetryMiddleware retry = new RetryMiddleware(opts); + + AtomicInteger calls = new AtomicInteger(); + ModelNext next = + (ctx, params) -> { + calls.incrementAndGet(); + throw new GenkitException("always fails"); + }; + + assertThrows( + GenkitException.class, () -> retry.wrapModel(ctx(), new ModelParams(null, null), next)); + assertEquals(3, calls.get(), "1 initial attempt + 2 retries before giving up"); + } + + @Test + void parameterlessDescriptorHasNoConfigSchema() { + // fallback/simulateSystemPrompt DO have schemas; verify a wrapped live middleware would not. + GenerationMiddlewareDesc simulate = descriptors().get("simulateSystemPrompt"); + assertNotNull(simulate.configSchema()); + // sanity: metadata is null by default (omitted from toJson) + assertNull(simulate.metadata()); + } + + private static ActionContext ctx() { + return ActionContext.builder().registry(new DefaultRegistry()).build(); + } +} diff --git a/samples/middleware-v2/README.md b/samples/middleware-v2/README.md index 5726e101c..a6f9e95ea 100644 --- a/samples/middleware-v2/README.md +++ b/samples/middleware-v2/README.md @@ -59,6 +59,12 @@ Logs tool execution name and duration. Stateless — `newInstance()` returns `th ### 4. FullObservabilityMiddleware (All 3 hooks) A single middleware that implements all three hooks, showing how one middleware can observe the entire pipeline with per-invocation counters. +### 5. TaggedLoggingMiddleware (parameterized, shared via a plugin) +A **parameterized** middleware with a `tag` and `logResponseLength` config. It is shared through the sample's own `SampleMiddlewarePlugin` (a `MiddlewarePlugin`), so it appears in the Dev UI Middleware panel **with a parameters form**. + +### Built-in plugin middleware (`GenerationMiddlewarePlugin`) +The sample also adds `GenerationMiddlewarePlugin.create()`, which contributes three ready-to-use, parameterized middleware — `retry`, `fallback`, and `simulateSystemPrompt` — the Java equivalent of the JS `@genkit-ai/middleware` package. The `v2-resilient` flow uses `retry` (imported from the plugin) programmatically. + ## Available Endpoints | Endpoint | Description | Middleware | @@ -67,6 +73,7 @@ A single middleware that implements all three hooks, showing how one middleware | `/v2-observable` | AI chat | Full observability (all 3 hooks) | | `/v2-stacked` | AI chat | Three separate middleware stacked | | `/v2-baseline` | AI chat | No middleware (baseline) | +| `/v2-resilient` | AI chat | `retry` (imported from the middleware plugin, configured in code) | ## Example Requests @@ -130,6 +137,36 @@ ModelResponse response = genkit.generate( .build()); ``` +## Using middleware from the Dev UI + +Register middleware with the `Genkit` builder so they appear in the Dev UI **Middleware** panel: + +```java +Genkit genkit = Genkit.builder() + .plugin(OpenAIPlugin.create()) + .middleware(new MyMiddleware(), new AnotherMiddleware()) + .build(); +``` + +In the Dev UI, open the Middleware panel, tick one or more middlewares, then run any model from the **Models** runner. The Dev UI sends the selected middleware in the `use` field of the `/util/generate` action as `{ "name": ..., "config": ... }` entries, which the runtime resolves from the registry and dispatches through the full `wrapGenerate` / `wrapModel` / `wrapTool` chain — middleware logs will appear in the server console. + +`.middleware(...)` only controls Dev UI visibility; programmatic `GenerateOptions.builder().use(...)` calls do not require registration. + +### Sharing middleware as a plugin (with parameters) + +To share reusable, **parameterized** middleware the JS/Go way, implement `MiddlewarePlugin` and return descriptors built with `GenerationMiddlewares.define(...)`. This sample does exactly that with `SampleMiddlewarePlugin` (providing the parameterized `tagged-logging`), and also adds the built-in `GenerationMiddlewarePlugin`: + +```java +Genkit genkit = Genkit.builder() + .plugin(OpenAIPlugin.create()) + .plugin(GenerationMiddlewarePlugin.create()) // retry, fallback, simulateSystemPrompt + .plugin(new SampleMiddlewarePlugin()) // this sample's tagged-logging + .middleware(modelLogging, generateTiming, ...) // ad-hoc, parameterless + .build(); +``` + +Middleware defined with a config type expose a **parameters form** in the Dev UI panel. When you fill it in and run a model, the values arrive as `use: [{ "name": "tagged-logging", "config": { "tag": "demo", "logResponseLength": true } }]` and are bound to a fresh middleware instance before it runs. + ## Architecture V2 middleware wraps the generation pipeline at three levels: diff --git a/samples/middleware-v2/pom.xml b/samples/middleware-v2/pom.xml index 0e4f7a077..bfc2b2069 100644 --- a/samples/middleware-v2/pom.xml +++ b/samples/middleware-v2/pom.xml @@ -61,6 +61,11 @@ genkit-plugin-jetty ${genkit.version} + + com.google.genkit + genkit-plugin-middleware + ${genkit.version} + ch.qos.logback logback-classic diff --git a/samples/middleware-v2/src/main/java/com/google/genkit/samples/MiddlewareV2Sample.java b/samples/middleware-v2/src/main/java/com/google/genkit/samples/MiddlewareV2Sample.java index 476eb2e04..e64e4c925 100644 --- a/samples/middleware-v2/src/main/java/com/google/genkit/samples/MiddlewareV2Sample.java +++ b/samples/middleware-v2/src/main/java/com/google/genkit/samples/MiddlewareV2Sample.java @@ -18,6 +18,8 @@ package com.google.genkit.samples; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonPropertyDescription; import com.google.genkit.Genkit; import com.google.genkit.GenkitOptions; import com.google.genkit.ai.GenerateOptions; @@ -29,15 +31,23 @@ import com.google.genkit.ai.middleware.GenerateNext; import com.google.genkit.ai.middleware.GenerateParams; import com.google.genkit.ai.middleware.GenerationMiddleware; +import com.google.genkit.ai.middleware.GenerationMiddlewareDesc; +import com.google.genkit.ai.middleware.GenerationMiddlewares; +import com.google.genkit.ai.middleware.MiddlewarePlugin; import com.google.genkit.ai.middleware.ModelNext; import com.google.genkit.ai.middleware.ModelParams; import com.google.genkit.ai.middleware.ToolNext; import com.google.genkit.ai.middleware.ToolParams; +import com.google.genkit.core.Action; import com.google.genkit.core.ActionContext; import com.google.genkit.core.Flow; import com.google.genkit.core.GenkitException; +import com.google.genkit.core.Plugin; +import com.google.genkit.core.Registry; import com.google.genkit.plugins.jetty.JettyPlugin; import com.google.genkit.plugins.jetty.JettyPluginOptions; +import com.google.genkit.plugins.middleware.GenerationMiddlewarePlugin; +import com.google.genkit.plugins.middleware.RetryMiddleware; import com.google.genkit.plugins.openai.OpenAIPlugin; import java.util.HashMap; import java.util.List; @@ -229,6 +239,91 @@ public Part wrapTool(ActionContext ctx, ToolParams params, ToolNext next) } } + // ========================================================================= + // Example 5: Custom PARAMETERIZED middleware shared via a plugin + // ========================================================================= + + /** + * A parameterized middleware: logs each model call tagged with a configurable {@code tag}, and + * optionally the response length. Because it is registered through a {@link MiddlewarePlugin} + * with a config schema (see {@link SampleMiddlewarePlugin}), the Dev UI renders its parameters + * ({@code tag}, {@code logResponseLength}) as a form. + */ + static class TaggedLoggingMiddleware extends BaseGenerationMiddleware { + + private final Options options; + + TaggedLoggingMiddleware(Options options) { + this.options = options != null ? options : new Options(); + } + + @Override + public String name() { + return "tagged-logging"; + } + + @Override + public GenerationMiddleware newInstance() { + return this; + } + + @Override + public ModelResponse wrapModel(ActionContext ctx, ModelParams params, ModelNext next) + throws GenkitException { + logger.info("[{}] model call", options.tag); + ModelResponse resp = next.apply(ctx, params); + if (options.logResponseLength) { + logger.info( + "[{}] response length: {}", + options.tag, + resp.getText() != null ? resp.getText().length() : 0); + } + return resp; + } + + /** Configuration parameters — rendered as a form in the Dev UI Middleware panel. */ + public static class Options { + @JsonProperty("tag") + @JsonPropertyDescription("Label prefixed to each log line.") + public String tag = "tagged"; + + @JsonProperty("logResponseLength") + @JsonPropertyDescription("Whether to log the model response length.") + public boolean logResponseLength = false; + + public Options() {} + } + } + + /** + * The sample's own middleware plugin. Demonstrates how an application (or a third-party library) + * shares middleware the JS/Go way: implement {@link MiddlewarePlugin} and return descriptors + * built with {@link GenerationMiddlewares#define}. Once added to the builder, {@code + * tagged-logging} shows up in the Dev UI Middleware panel with a parameters form and can be + * attached by name. + */ + static class SampleMiddlewarePlugin implements Plugin, MiddlewarePlugin { + @Override + public String getName() { + return "sample-middleware"; + } + + @Override + public List> init() { + return List.of(); + } + + @Override + public List middlewares(Registry registry) { + return List.of( + GenerationMiddlewares.define( + "tagged-logging", + "Logs each model call with a configurable tag.", + TaggedLoggingMiddleware.Options.class, + TaggedLoggingMiddleware::new)); + } + } + // ========================================================================= // Main // ========================================================================= @@ -236,19 +331,27 @@ public Part wrapTool(ActionContext ctx, ToolParams params, ToolNext next) public static void main(String[] args) throws Exception { JettyPlugin jetty = new JettyPlugin(JettyPluginOptions.builder().port(8080).build()); + // Instantiate middleware (templates — newInstance() is called per generate()) + GenerationMiddleware modelLogging = new ModelLoggingMiddleware(); + GenerationMiddleware generateTiming = new GenerateTimingMiddleware(); + GenerationMiddleware toolMonitor = new ToolMonitorMiddleware(); + GenerationMiddleware fullObservability = new FullObservabilityMiddleware(); + Genkit genkit = Genkit.builder() .options(GenkitOptions.builder().devMode(true).reflectionPort(3100).build()) .plugin(OpenAIPlugin.create()) .plugin(jetty) + // Shared, parameterized middleware imported from the built-in plugin (retry, fallback, + // simulateSystemPrompt). These appear in the Dev UI Middleware panel WITH a parameters + // form and can be attached to any generate() call by name. + .plugin(GenerationMiddlewarePlugin.create()) + // The sample's OWN middleware, shared as a plugin (parameterized "tagged-logging"). + .plugin(new SampleMiddlewarePlugin()) + // Ad-hoc live middleware registered directly (no parameters). + .middleware(modelLogging, generateTiming, toolMonitor, fullObservability) .build(); - // Instantiate middleware (templates — newInstance() is called per generate()) - GenerationMiddleware modelLogging = new ModelLoggingMiddleware(); - GenerationMiddleware generateTiming = new GenerateTimingMiddleware(); - GenerationMiddleware toolMonitor = new ToolMonitorMiddleware(); - GenerationMiddleware fullObservability = new FullObservabilityMiddleware(); - // Define a simple tool so the WrapTool hook gets exercised @SuppressWarnings("unchecked") Tool, Map> weatherTool = @@ -375,6 +478,37 @@ public static void main(String[] args) throws Exception { return response.getText(); }); + // ======================================================= + // Flow 5: Uses the imported plugin middleware programmatically + // ======================================================= + + // The plugin ships `retry` for Dev-UI use; its RetryMiddleware class can also be used directly. + RetryMiddleware.Options retryOpts = new RetryMiddleware.Options(); + retryOpts.maxRetries = 2; + retryOpts.initialDelayMs = 500; + + Flow resilientFlow = + genkit.defineFlow( + "v2-resilient", + String.class, + String.class, + (ctx, userMessage) -> { + ModelResponse response = + genkit.generate( + GenerateOptions.builder() + .model("openai/gpt-4o-mini") + .prompt(userMessage) + // Imported from the middleware plugin, configured in code. + .use(new RetryMiddleware(retryOpts), modelLogging) + .config( + GenerationConfig.builder() + .temperature(0.7) + .maxOutputTokens(200) + .build()) + .build()); + return response.getText(); + }); + logger.info("\n========================================"); logger.info("Genkit Middleware V2 Sample Started!"); logger.info("========================================\n"); @@ -383,7 +517,18 @@ public static void main(String[] args) throws Exception { logger.info(" - v2-chat: Model logging + generate timing middleware"); logger.info(" - v2-observable: Full observability (all 3 hooks in one middleware)"); logger.info(" - v2-stacked: Three separate middleware stacked together"); - logger.info(" - v2-baseline: No middleware (baseline comparison)\n"); + logger.info(" - v2-baseline: No middleware (baseline comparison)"); + logger.info( + " - v2-resilient: retry middleware imported from the plugin (configured in code)\n"); + + logger.info( + "Dev UI Middleware panel now lists (open http://localhost:4000 via `genkit start`):"); + logger.info(" Parameterized (from plugins, with a config form):"); + logger.info(" - retry, fallback, simulateSystemPrompt (genkit-middleware plugin)"); + logger.info( + " - tagged-logging (this sample's SampleMiddlewarePlugin)"); + logger.info(" Parameterless (registered via .middleware(...)):"); + logger.info(" - model-logging, generate-timing, tool-monitor, full-observability\n"); logger.info("Server running on http://localhost:8080"); logger.info("Reflection server running on http://localhost:3100");