[integrations][gemini] Apply Gemini native structured output

integrations/chat-models/gemini/pom.xml

@@ -43,6 +43,17 @@ under the License.
             <artifactId>google-genai</artifactId>
             <version>${google.genai.version}</version>
         </dependency>
+
+        <!-- Versions managed by the victools BOM imported in the root pom. -->
+        <dependency>
+            <groupId>com.github.victools</groupId>
+            <artifactId>jsonschema-generator</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>com.github.victools</groupId>
+            <artifactId>jsonschema-module-jackson</artifactId>
+        </dependency>
     </dependencies>
 
 </project>

integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java

@@ -19,7 +19,17 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.github.victools.jsonschema.generator.Option;
+import com.github.victools.jsonschema.generator.OptionPreset;
+import com.github.victools.jsonschema.generator.SchemaGenerator;
+import com.github.victools.jsonschema.generator.SchemaGeneratorConfigBuilder;
+import com.github.victools.jsonschema.generator.SchemaVersion;
+import com.github.victools.jsonschema.generator.impl.PropertySortUtils;
+import com.github.victools.jsonschema.module.jackson.JacksonModule;
+import com.github.victools.jsonschema.module.jackson.JacksonOption;
 import com.google.genai.Client;
 import com.google.genai.types.Candidate;
 import com.google.genai.types.Content;
@@ -48,6 +58,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import java.util.stream.Collectors;
 
@@ -102,6 +113,28 @@ public class GeminiChatModelConnection extends BaseChatModelConnection {
 
     private static final TypeReference<Map<String, Object>> MAP_TYPE = new TypeReference<>() {};
 
+    // Models for which Google documents native structured output.
+    // Source of truth: each model's Capabilities row at
+    // https://ai.google.dev/gemini-api/docs/models
+    //
+    // Google added JSON Schema support to every actively supported Gemini model, and every model
+    // predating that rollout has since been shut down, so the whole live text family is matched by
+    // the family prefix rather than enumerated. A new Flash generation has shipped roughly monthly,
+    // so an enumerated list would report not-capable for models that do work.
+    //
+    // Capability splits by output modality, not by generation: the image, speech, audio, Live,
+    // transcription, video and embedding variants share the family prefix but expose no schema
+    // parameter, so a modality marker is checked first and wins over the prefix. The markers also
+    // cover gemini-2.5-flash-image, whose published capability row claims support that the service
+    // rejects.
+    //
+    // A name outside the family — a Gemma model served by the same endpoint, a tuned model, or a
+    // path-qualified form such as models/gemini-2.5-flash — reports not-capable and degrades to the
+    // prompt fallback rather than failing at the provider.
+    private static final String NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX = "gemini-";
+    private static final Set<String> NON_TEXT_MODALITY_MARKERS =
+            Set.of("-image", "-tts", "-audio", "-live", "-transcribe", "-embedding", "-omni");
+
     private final ObjectMapper mapper = new ObjectMapper();
     private final Client client;
     private final String defaultModel;
@@ -173,11 +206,67 @@ public void close() {
         this.client.close();
     }
 
+    /**
+     * Whether Google documents native structured output for {@code effectiveModel}.
+     *
+     * <p>A {@code true} means the request carries the schema as a native {@code responseJsonSchema}
+     * and the model is one Google documents as accepting one. What that buys is structural
+     * conformance: the response is syntactically valid JSON whose object shape, key set and value
+     * types follow the schema as the service interpreted it.
+     *
+     * <p>It does not buy conformance to every constraint the schema expresses. Gemini supports a
+     * subset of JSON Schema and ignores the keywords outside that subset, server-side, without
+     * reporting which. A {@code pattern}, {@code minLength}, {@code maxLength}, {@code
+     * minProperties} or {@code maxProperties} in a derived schema is accepted by the request and
+     * has no effect on the response, so well-formed JSON of the right shape is not evidence that
+     * every declared constraint held. It does not buy semantically correct values either.
+     *
+     * <p>Capability is read from the model name alone, so a Gemini model released after this was
+     * written is treated as capable. One that turns out not to support structured output fails at
+     * the provider with {@code 400 INVALID_ARGUMENT}, "JSON mode is not enabled for this model",
+     * rather than degrading quietly. A name outside the family — a Gemma model on the same
+     * endpoint, a tuned model, or the path-qualified {@code models/gemini-2.5-flash} form the SDK
+     * also accepts — reports not-capable and keeps the prompt-engineering fallback.
+     */
+    @Override
+    protected boolean supportsNativeStructuredOutput(String effectiveModel) {
+        if (effectiveModel == null || effectiveModel.isBlank()) {
+            return false;
+        }
+        if (NON_TEXT_MODALITY_MARKERS.stream().anyMatch(effectiveModel::contains)) {
+            return false;
+        }
+        return effectiveModel.startsWith(NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX)
+                && effectiveModel.length() > NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX.length();
+    }
+
     @Override
     public ChatMessage chat(
             List<ChatMessage> messages,
             List<org.apache.flink.agents.api.tools.Tool> tools,
             Map<String, Object> arguments) {
+        return chat(messages, tools, arguments, null);
+    }
+
+    /**
+     * Translates {@code outputSchema} into Gemini's native {@code responseJsonSchema} when it is a
+     * POJO {@link Class}, the request carries no tools, and the effective model is one Google
+     * documents structured-output support for. Any other combination sends no derived schema, so a
+     * schema that cannot be sent natively degrades to the prompt-engineering fallback rather than
+     * failing at the provider.
+     *
+     * <p>The tools condition is a provider constraint rather than a preference: outside a
+     * documented preview, Gemini answers a request that combines function declarations with a JSON
+     * response mime type with {@code 400 INVALID_ARGUMENT}, "Function calling with a response mime
+     * type: 'application/json' is unsupported". The request therefore proceeds with its tools and
+     * without the schema.
+     */
+    @Override
+    public ChatMessage chat(
+            List<ChatMessage> messages,
+            List<org.apache.flink.agents.api.tools.Tool> tools,
+            Map<String, Object> arguments,
+            Object outputSchema) {
         Map<String, Object> args = arguments != null ? new HashMap<>(arguments) : new HashMap<>();
 
         Object modelObj = args.remove("model");
@@ -202,7 +291,8 @@ public ChatMessage chat(
                             .map(m -> convertToContent(m, toolCallIdToName))
                             .collect(Collectors.toList());
 
-            GenerateContentConfig config = buildConfig(messages, tools, args);
+            GenerateContentConfig config =
+                    buildConfig(messages, tools, args, modelName, outputSchema);
 
             GenerateContentResponse response =
                     client.models.generateContent(modelName, contents, config);
@@ -251,11 +341,15 @@ static Map<String, String> buildToolCallIdToNameMap(List<ChatMessage> messages)
         return map;
     }
 
-    // Package-visible for unit testing of the request-config assembly.
+    // Package-visible for unit testing of the request-config assembly. modelName is passed
+    // explicitly rather than read from arguments: chat() removes the model key before calling
+    // this method, so the map never carries it here.
     GenerateContentConfig buildConfig(
             List<ChatMessage> messages,
             List<org.apache.flink.agents.api.tools.Tool> tools,
-            Map<String, Object> arguments) {
+            Map<String, Object> arguments,
+            String modelName,
+            Object outputSchema) {
         GenerateContentConfig.Builder builder = GenerateContentConfig.builder();
 
         Content systemInstruction = extractSystemInstruction(messages);
@@ -284,9 +378,120 @@ GenerateContentConfig buildConfig(
             builder.tools(List.of(convertTools(tools)));
         }
 
+        // Native structured output applies only for a POJO Class schema; any other schema form,
+        // such as a RowTypeInfo wrapped in OutputSchema, keeps the prompt-engineering fallback.
+        // Nothing above writes either field this branch sets: the keys read directly are
+        // temperature and max_output_tokens, and applyAdditionalKwargs recognizes only top_k,
+        // top_p and stop_sequences, so there is no caller-supplied value to collide with.
+        //
+        // TODO(#912): the requested strategy is not visible here, so this re-check cannot tell an
+        // explicit NATIVE request apart from one that merely resolved to native. A caller asking
+        // for NATIVE on a schema form, a model, or a tool-carrying request this branch skips
+        // therefore gets an unconstrained response instead of an error. Once strategy resolution is
+        // wired up, NATIVE must either bypass this capability re-check or fail explicitly.
+        if (outputSchema instanceof Class
+                && (tools == null || tools.isEmpty())
+                && supportsNativeStructuredOutput(modelName)) {
+            builder.responseMimeType("application/json");
+            builder.responseJsonSchema(toNativeJsonSchema((Class<?>) outputSchema));
+        }
+
         return builder.build();
     }
 
+    // Derives the JSON Schema Gemini's responseJsonSchema field expects from a POJO class. Gemini
+    // supports a subset of JSON Schema and ignores the keywords outside that subset server-side
+    // without reporting which, so every setting below is chosen against that published subset:
+    //
+    //   - DRAFT_2020_12 is the draft whose keywords Gemini's supported list names: $defs and
+    //     prefixItems are listed, while the older drafts' definitions and tuple-form items are not.
+    //   - The PLAIN_JSON preset keeps generation to fields. A preset is mandatory, and the
+    //     generator's default one, FULL_DOCUMENTATION, surfaces getters as properties of their
+    //     own, named after the accessor call, e.g. "getSummary()".
+    //   - MAP_VALUES_AS_ADDITIONAL_PROPERTIES gives a Map its value schema, as an
+    //     additionalProperties keyword carrying the declared value type. Dropped from this recipe,
+    //     the map instead takes the additionalProperties:false of the option below and admits no
+    //     entries at all.
+    //   - FORBIDDEN_ADDITIONAL_PROPERTIES_BY_DEFAULT closes every object. No Gemini document states
+    //     that a schema omitting the keyword is closed, and under ordinary JSON Schema semantics it
+    //     is not, so without it a response may carry an undeclared key that the ObjectMapper
+    //     read-back then rejects. If the service honors the keyword the object is closed; if it
+    //     does not, the keyword is ignored like any other unsupported one, which is where omitting
+    //     it would have left us. It applies to the enclosing object and leaves a Map's declared
+    //     value schema alone.
+    //   - Sorting fields before methods and applying no further comparison leaves properties in
+    //     declaration order, which keeps the emitted document stable rather than alphabetized. It
+    //     is not an ordering guarantee: Gemini's ordering knob is the non-standard propertyOrdering
+    //     keyword, which this generator never emits.
+    //   - The required check marks every field required except an Optional one. Gemini treats a
+    //     field the schema does not list as required as optional and lets the model skip it, while
+    //     marking everything required would force the fields a caller declared omissible.
+    //   - The Jackson module makes the schema name properties the way Jackson names them. The
+    //     response is read back into the same class with an ObjectMapper, so a property that
+    //     @JsonProperty renames or @JsonIgnore drops has to be stated in the schema under the name
+    //     the mapper reads, or a response that satisfies the schema still fails to deserialize.
+    //     Enum constants carry the same hazard: the FLATTENED_ENUMS options list each constant by
+    //     its @JsonValue method or @JsonProperty value, as the mapper reads it. An enum annotating
+    //     only some constants falls back to Java names for all of them, so its annotated constants
+    //     do not read back. No other JacksonOption is enabled, so the required set and property
+    //     order stay as configured above.
+    //
+    // DEFINITION_FOR_MAIN_SCHEMA is deliberately absent. Without it a recursive type emits
+    // {"$ref": "#"} at the recursion point, which is the form Google's own recursion example uses.
+    // Enabling it instead produces a $defs entry referencing another $defs entry, a shape no
+    // published Gemini example demonstrates.
+    private static ObjectNode toNativeJsonSchema(Class<?> schemaClass) {
+        SchemaGeneratorConfigBuilder configBuilder =
+                new SchemaGeneratorConfigBuilder(
+                                SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON)
+                        .with(Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES)
+                        .with(Option.FORBIDDEN_ADDITIONAL_PROPERTIES_BY_DEFAULT)
+                        .with(
+                                new JacksonModule(
+                                        JacksonOption.FLATTENED_ENUMS_FROM_JSONVALUE,
+                                        JacksonOption.FLATTENED_ENUMS_FROM_JSONPROPERTY));
+        configBuilder
+                .forTypesInGeneral()
+                .withPropertySorter(PropertySortUtils.SORT_PROPERTIES_FIELDS_BEFORE_METHODS);
+        configBuilder
+                .forFields()
+                .withRequiredCheck(field -> !Optional.class.equals(field.getRawMember().getType()));
+        ObjectNode schema = new SchemaGenerator(configBuilder.build()).generateSchema(schemaClass);
+        stripRefSiblings(schema);
+        return schema;
+    }
+
+    // Gemini states that a sub-schema setting $ref may set no other property except those starting
+    // with a $, so every non-$ sibling is removed wherever a $ref appears. The generator writes a
+    // referencing field's own keywords beside the reference, and two shapes are measured to produce
+    // the pairing: a described field whose type is used more than once, which is extracted into
+    // $defs and referenced there, and a described field that recurses into its own type, which
+    // references the document root. What the strip removes is not only documentation: a Map's value
+    // schema is written beside a $ref the same way a description is.
+    //
+    // The test is on the $ref value rather than on the presence of a member named $ref, because a
+    // POJO may declare a property called $ref: its schema is then an object sitting under the
+    // enclosing properties map, which makes that map itself carry a $ref member while being an
+    // ordinary properties map rather than a reference. A reference's own $ref value is a URI
+    // string, so requiring a textual value tells the two apart. Matching on presence alone would
+    // delete every other declared property from such a map while required still listed them,
+    // producing a document nothing can satisfy.
+    private static void stripRefSiblings(JsonNode node) {
+        if (node instanceof ObjectNode && node.path("$ref").isTextual()) {
+            ObjectNode object = (ObjectNode) node;
+            List<String> siblings = new ArrayList<>();
+            object.fieldNames()
+                    .forEachRemaining(
+                            name -> {
+                                if (!name.startsWith("$")) {
+                                    siblings.add(name);
+                                }
+                            });
+            siblings.forEach(object::remove);
+        }
+        node.forEach(GeminiChatModelConnection::stripRefSiblings);
+    }
+
     // Package-visible for unit testing of the additional-kwargs forwarding.
     void applyAdditionalKwargs(GenerateContentConfig.Builder builder, Map<String, Object> kwargs) {
         for (Map.Entry<String, Object> entry : kwargs.entrySet()) {