This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 2050b526 [integrations][java] Share POJO JSON Schema generation across
Ollama, Gemini, Bedrock and watsonx (#1120)
2050b526 is described below
commit 2050b5263e72541594cd798c2c7e52233b7ee98b
Author: Weiqing Yang <[email protected]>
AuthorDate: Tue Sep 15 02:31:31 2026 -0700
[integrations][java] Share POJO JSON Schema generation across Ollama,
Gemini, Bedrock and watsonx (#1120)
Generated-by: Claude Code 2.1.272 (Claude Opus 5)
---
dist/pom.xml | 5 +
integrations/chat-models/bedrock/pom.xml | 9 +-
.../bedrock/BedrockChatModelConnection.java | 46 +----
.../bedrock/BedrockChatModelConnectionTest.java | 56 ++----
.../chat-models/{ollama => common}/pom.xml | 17 +-
.../chatmodels/common/PojoJsonSchemaGenerator.java | 83 +++++++++
.../common/PojoJsonSchemaGeneratorTest.java | 196 +++++++++++++++++++++
integrations/chat-models/gemini/pom.xml | 9 +-
.../gemini/GeminiChatModelConnection.java | 65 ++-----
.../gemini/GeminiChatModelConnectionTest.java | 47 +----
integrations/chat-models/ollama/pom.xml | 9 +-
.../ollama/OllamaChatModelConnection.java | 48 ++---
.../ollama/OllamaChatModelConnectionTest.java | 109 ++++++------
integrations/chat-models/pom.xml | 1 +
integrations/chat-models/watsonx/pom.xml | 9 +-
.../watsonx/WatsonxChatModelConnection.java | 50 ++----
.../watsonx/WatsonxChatModelConnectionTest.java | 92 +++++-----
17 files changed, 473 insertions(+), 378 deletions(-)
diff --git a/dist/pom.xml b/dist/pom.xml
index 7cb74295..df52c7a3 100644
--- a/dist/pom.xml
+++ b/dist/pom.xml
@@ -62,6 +62,11 @@ under the License.
<artifactId>flink-agents-runtime</artifactId>
<version>${project.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+
<artifactId>flink-agents-integrations-chat-models-common</artifactId>
+ <version>${project.version}</version>
+ </dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-agents-integrations-chat-models-ollama</artifactId>
diff --git a/integrations/chat-models/bedrock/pom.xml
b/integrations/chat-models/bedrock/pom.xml
index eb88c88a..62bd09bf 100644
--- a/integrations/chat-models/bedrock/pom.xml
+++ b/integrations/chat-models/bedrock/pom.xml
@@ -44,15 +44,16 @@ under the License.
<version>${aws.sdk.version}</version>
</dependency>
- <!-- Versions managed by the victools BOM imported in the root pom. -->
<dependency>
- <groupId>com.github.victools</groupId>
- <artifactId>jsonschema-generator</artifactId>
+ <groupId>org.apache.flink</groupId>
+
<artifactId>flink-agents-integrations-chat-models-common</artifactId>
+ <version>${project.version}</version>
</dependency>
+ <!-- Versions managed by the victools BOM imported in the root pom. -->
<dependency>
<groupId>com.github.victools</groupId>
- <artifactId>jsonschema-module-jackson</artifactId>
+ <artifactId>jsonschema-generator</artifactId>
</dependency>
</dependencies>
diff --git
a/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java
b/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java
index 62cc250c..29989ab6 100644
---
a/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java
+++
b/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java
@@ -22,12 +22,6 @@ 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.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.module.jackson.JacksonModule;
-import com.github.victools.jsonschema.module.jackson.JacksonOption;
import org.apache.flink.agents.api.RetryExecutor;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
@@ -36,6 +30,7 @@ import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.tools.Tool;
import org.apache.flink.agents.api.tools.ToolMetadata;
+import
org.apache.flink.agents.integrations.chatmodels.common.PojoJsonSchemaGenerator;
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.core.SdkNumber;
import software.amazon.awssdk.core.document.Document;
@@ -66,7 +61,6 @@ import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
-import java.util.Optional;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@@ -388,24 +382,8 @@ public class BedrockChatModelConnection extends
BaseChatModelConnection {
.build();
}
- // Derives the JSON schema from a POJO class. Every setting below
addresses a concrete way the
- // generated schema otherwise fails to constrain generation:
- //
- // - DRAFT_2020_12 is the dialect Bedrock validates a schema against, so
the schema
- // declares it rather than the generator's older default.
- // - The PLAIN_JSON preset keeps generation to fields. Without a preset,
getters surface as
- // properties of their own, named after the accessor call, e.g.
"getSummary()".
- // - The required check marks every field required except an Optional
one. The default marks
- // nothing required, which lets a model omit fields at will, 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, and
- // list enum constants mapped by @JsonProperty or by a @JsonValue
method under the values
- // Jackson reads. The response is read back into the same class with
an ObjectMapper, so a
- // property that @JsonProperty renames or @JsonIgnore drops, and a
mapped enum constant,
- // have to appear in the schema as the mapper reads them, or a
response that satisfies the
- // schema still fails to deserialize. An enum annotating only some
constants falls back to
- // Java names for all of them, so its annotated constants do not read
back. The two enum
- // options change only the listed values: the required set stays the
one configured above.
+ // Derives the JSON schema from a POJO class with the shared generator,
adding no option. The
+ // shared schema declares draft 2020-12, the dialect Bedrock validates a
schema against.
//
// A Map's value schema is deliberately left underived. Bedrock accepts
additionalProperties
// only as false, and rejects a schema that carries it as a subschema, so
typing map values
@@ -413,21 +391,11 @@ public class BedrockChatModelConnection extends
BaseChatModelConnection {
// bare object.
//
// A self-referencing class derives its own field as a reference back to
the schema root,
- // whatever the required check says. Bedrock does not accept a recursive
schema and rejects the
- // request before the model runs, so declaring the field Optional does not
rescue it; only
- // flattening the recursion does.
+ // whatever the shared generator's required check says. Bedrock does not
accept a recursive
+ // schema and rejects the request before the model runs, so declaring the
field Optional does
+ // not rescue it; only flattening the recursion does.
private static JsonNode toNativeSchema(Class<?> schemaClass) {
- SchemaGeneratorConfigBuilder configBuilder =
- new SchemaGeneratorConfigBuilder(
- SchemaVersion.DRAFT_2020_12,
OptionPreset.PLAIN_JSON)
- .with(
- new JacksonModule(
-
JacksonOption.FLATTENED_ENUMS_FROM_JSONPROPERTY,
-
JacksonOption.FLATTENED_ENUMS_FROM_JSONVALUE));
- configBuilder
- .forFields()
- .withRequiredCheck(field ->
!Optional.class.equals(field.getRawMember().getType()));
- return new
SchemaGenerator(configBuilder.build()).generateSchema(schemaClass);
+ return PojoJsonSchemaGenerator.generate(schemaClass);
}
private static boolean isRetryable(Exception e) {
diff --git
a/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnectionTest.java
b/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnectionTest.java
index 59214174..334c94a7 100644
---
a/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnectionTest.java
+++
b/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnectionTest.java
@@ -18,7 +18,6 @@
package org.apache.flink.agents.integrations.chatmodels.bedrock;
-import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonValue;
import com.fasterxml.jackson.databind.JsonNode;
@@ -285,31 +284,23 @@ class BedrockChatModelConnectionTest {
private static final ObjectMapper SCHEMA_MAPPER = new ObjectMapper();
- /**
- * Output schema fixture shaped to expose Jackson's property model.
- *
- * <p>{@code name} is deserialized from {@code full_name} rather than from
the Java field name,
- * and {@code secret} is not deserialized at all.
- */
+ /** Output schema fixture for tests that need a POJO class schema of any
shape. */
public static class Profile {
- @JsonProperty("full_name")
public String name;
-
- @JsonIgnore public String secret;
-
- public int age;
}
- /**
- * Output schema fixture shaped to expose what the derived schema
constrains: which fields a
- * response must carry, and how a map field is rendered.
- */
+ /** Output schema fixture with a map field whose values carry a type. */
public static class Reading {
- public int score;
+ public Map<String, Integer> counts;
+ }
- public Optional<String> note;
+ /** Output schema fixture whose fields are declared in an order other than
alphabetical. */
+ public static class Task {
+ public String title;
- public Map<String, Integer> counts;
+ public int priority;
+
+ public boolean archived;
}
/**
@@ -499,23 +490,19 @@ class BedrockChatModelConnectionTest {
}
@Test
- @DisplayName("the derived schema names properties the way Jackson
deserializes them")
- void testDerivedSchemaFollowsJacksonPropertyNames() throws Exception {
+ @DisplayName("the derived schema lists properties in the order the class
declares them")
+ void testDerivedSchemaKeepsDeclarationOrder() throws Exception {
ConverseRequest request =
connection()
.buildRequest(
List.of(ChatMessage.user("hello")),
null,
params(CAPABLE_MODEL),
- Profile.class);
+ Task.class);
- // The response is read back into the same class, so a property that
@JsonProperty renames
- // or @JsonIgnore drops has to be stated under the name the mapper
reads. Derived without
- // Jackson's property model the schema instead names "name" and
demands the ignored
- // "secret", constraining the model to a document the mapper then
refuses.
assertThat(nativeSchema(request).path("properties").fieldNames())
.toIterable()
- .containsExactlyInAnyOrder("full_name", "age");
+ .containsExactly("title", "priority", "archived");
}
@Test
@@ -546,8 +533,8 @@ class BedrockChatModelConnectionTest {
}
@Test
- @DisplayName("the derived schema declares its dialect, requires fields,
and leaves maps bare")
- void testDerivedSchemaConstrainsTheResponse() throws Exception {
+ @DisplayName("the derived schema leaves map values untyped")
+ void testDerivedSchemaLeavesMapsBare() throws Exception {
ConverseRequest request =
connection()
.buildRequest(
@@ -557,23 +544,12 @@ class BedrockChatModelConnectionTest {
Reading.class);
JsonNode schema = nativeSchema(request);
- // Left to itself the generator marks nothing required, so an empty
document satisfies a
- // schema whose whole purpose is to constrain the response. An
Optional field is the one a
- // caller declared omissible, so it stays out of the required set.
- List<String> required = new ArrayList<>();
- schema.path("required").forEach(entry -> required.add(entry.asText()));
- assertThat(required).containsExactlyInAnyOrder("score", "counts");
-
// A map derives as a bare object. Typing its values renders them under
// additionalProperties, which Bedrock accepts only as false and
rejects as a subschema, so
// the value type is left off rather than putting the request outside
the accepted subset.
assertThat(schema.path("properties").path("counts").path("type").asText())
.isEqualTo("object");
assertThat(schema.findValues("additionalProperties")).isEmpty();
-
- // The dialect is stated rather than left on the generator's older
default.
- assertThat(schema.path("$schema").asText())
- .isEqualTo("https://json-schema.org/draft/2020-12/schema";);
}
@Test
diff --git a/integrations/chat-models/ollama/pom.xml
b/integrations/chat-models/common/pom.xml
similarity index 79%
copy from integrations/chat-models/ollama/pom.xml
copy to integrations/chat-models/common/pom.xml
index 49638efd..24c39026 100644
--- a/integrations/chat-models/ollama/pom.xml
+++ b/integrations/chat-models/common/pom.xml
@@ -27,21 +27,14 @@ under the License.
<relativePath>../pom.xml</relativePath>
</parent>
- <artifactId>flink-agents-integrations-chat-models-ollama</artifactId>
- <name>Flink Agents : Integrations: Chat Models: Ollama</name>
+ <artifactId>flink-agents-integrations-chat-models-common</artifactId>
+ <name>Flink Agents : Integrations: Chat Models: Common</name>
<packaging>jar</packaging>
<dependencies>
<dependency>
- <groupId>org.apache.flink</groupId>
- <artifactId>flink-agents-api</artifactId>
- <version>${project.version}</version>
- </dependency>
-
- <dependency>
- <groupId>io.github.ollama4j</groupId>
- <artifactId>ollama4j</artifactId>
- <version>${ollama4j.version}</version>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-databind</artifactId>
</dependency>
<!-- Versions managed by the victools BOM imported in the root pom. -->
@@ -56,4 +49,4 @@ under the License.
</dependency>
</dependencies>
-</project>
\ No newline at end of file
+</project>
diff --git
a/integrations/chat-models/common/src/main/java/org/apache/flink/agents/integrations/chatmodels/common/PojoJsonSchemaGenerator.java
b/integrations/chat-models/common/src/main/java/org/apache/flink/agents/integrations/chatmodels/common/PojoJsonSchemaGenerator.java
new file mode 100644
index 00000000..e36322e9
--- /dev/null
+++
b/integrations/chat-models/common/src/main/java/org/apache/flink/agents/integrations/chatmodels/common/PojoJsonSchemaGenerator.java
@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+package org.apache.flink.agents.integrations.chatmodels.common;
+
+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 java.util.Optional;
+
+/** Derives the JSON Schema that chat model connections send as a native
output schema. */
+public final class PojoJsonSchemaGenerator {
+
+ private PojoJsonSchemaGenerator() {}
+
+ // Every fixed setting below addresses a concrete way the generated schema
otherwise misstates
+ // the contract of the class a caller deserializes the response into:
+ //
+ // - DRAFT_2020_12 is the dialect pydantic generates on the Python side,
so a schema derived
+ // from a Java class states the same contract in the same dialect.
+ // - The PLAIN_JSON preset keeps generation to fields. Under the default
FULL_DOCUMENTATION
+ // preset, getters also surface as properties of their own, named
after the accessor call,
+ // e.g. "getDerived()".
+ // - Sorting fields before methods and applying no further comparison
leaves properties in
+ // declaration order, which is the order pydantic emits.
+ // - The required check marks every field required except an Optional
one. The default marks
+ // nothing required, which lets a model omit fields at will, while
marking everything
+ // required would force the fields a caller declared omissible.
+ // - The Jackson module names properties the way an ObjectMapper reads
them, honoring
+ // @JsonProperty renames and dropping @JsonIgnore fields, which the
required check would
+ // otherwise mark required. It lists enum constants mapped by
@JsonProperty or by a
+ // @JsonValue method under those mapped values. A property or enum
constant stated under
+ // any other name makes a response that satisfies the schema fail to
deserialize. An enum
+ // annotating only some constants falls back to Java names for all of
them, so its
+ // annotated constants do not read back. RESPECT_JSONPROPERTY_REQUIRED
stays off, so
+ // @JsonProperty(required = true) does not widen the required set
configured here.
+
+ /**
+ * Generates the JSON Schema for {@code type}.
+ *
+ * @param type the class a response is deserialized into
+ * @param options additional generator options enabled on top of the fixed
settings
+ * @return the schema document
+ */
+ public static ObjectNode generate(Class<?> type, Option... options) {
+ SchemaGeneratorConfigBuilder builder =
+ new SchemaGeneratorConfigBuilder(
+ SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON);
+ for (Option option : options) {
+ builder.with(option);
+ }
+ builder.with(
+ new JacksonModule(
+ JacksonOption.FLATTENED_ENUMS_FROM_JSONPROPERTY,
+ JacksonOption.FLATTENED_ENUMS_FROM_JSONVALUE));
+ builder.forTypesInGeneral()
+
.withPropertySorter(PropertySortUtils.SORT_PROPERTIES_FIELDS_BEFORE_METHODS);
+ builder.forFields()
+ .withRequiredCheck(field ->
!Optional.class.equals(field.getRawMember().getType()));
+ return new SchemaGenerator(builder.build()).generateSchema(type);
+ }
+}
diff --git
a/integrations/chat-models/common/src/test/java/org/apache/flink/agents/integrations/chatmodels/common/PojoJsonSchemaGeneratorTest.java
b/integrations/chat-models/common/src/test/java/org/apache/flink/agents/integrations/chatmodels/common/PojoJsonSchemaGeneratorTest.java
new file mode 100644
index 00000000..838336a7
--- /dev/null
+++
b/integrations/chat-models/common/src/test/java/org/apache/flink/agents/integrations/chatmodels/common/PojoJsonSchemaGeneratorTest.java
@@ -0,0 +1,196 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+package org.apache.flink.agents.integrations.chatmodels.common;
+
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonValue;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.github.victools.jsonschema.generator.Option;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Unit tests for {@link PojoJsonSchemaGenerator}. */
+class PojoJsonSchemaGeneratorTest {
+
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+
+ /**
+ * Output schema fixture shaped to expose the generation settings: fields
declared in a
+ * non-alphabetical order, one {@code Optional} field, and a getter with
no backing field.
+ */
+ public static class Report {
+ public String summary;
+ public Map<String, Integer> counts;
+ public Optional<String> note;
+ public int total;
+
+ public String getDerived() {
+ return summary + total;
+ }
+ }
+
+ /**
+ * Output schema fixture shaped to expose Jackson's property model.
+ *
+ * <p>{@code name} is deserialized from {@code full_name} rather than from
the Java field name,
+ * and {@code secret} is not deserialized at all.
+ */
+ public static class Profile {
+ @JsonProperty("full_name")
+ public String name;
+
+ @JsonIgnore public String secret;
+
+ public int age;
+ }
+
+ /**
+ * Output schema fixture whose enum constants are deserialized from values
other than their
+ * names, one through {@code @JsonProperty} on the constants and one
through a
+ * {@code @JsonValue} method.
+ */
+ public static class Ticket {
+ public Status status;
+
+ public Phase phase;
+ }
+
+ public enum Status {
+ @JsonProperty("in-progress")
+ IN_PROGRESS,
+ @JsonProperty("done")
+ DONE
+ }
+
+ public enum Phase {
+ STARTED("started"),
+ FINISHED("finished");
+
+ private final String wire;
+
+ Phase(String wire) {
+ this.wire = wire;
+ }
+
+ @JsonValue
+ public String wire() {
+ return wire;
+ }
+ }
+
+ @Test
+ @DisplayName("The generated schema names properties the way Jackson
deserializes them")
+ void namesPropertiesTheWayJacksonReadsThem() {
+ JsonNode schema = PojoJsonSchemaGenerator.generate(Profile.class);
+
+ // A caller deserializing the response into this class accepts the
renamed property and
+ // rejects the Java field name, and discards an ignored property that
the schema would
+ // otherwise state as required and so force a value for.
+ assertThat(fieldNames(schema.path("properties")))
+ .containsExactlyInAnyOrder("full_name", "age");
+ }
+
+ @Test
+ @DisplayName("The generated schema lists enum constants the way Jackson
deserializes them")
+ void listsEnumConstantsByTheirJacksonWireValues() throws Exception {
+ JsonNode properties =
PojoJsonSchemaGenerator.generate(Ticket.class).path("properties");
+
+ // Every listed value is one the model may emit, so each has to
deserialize into the enum.
+ // Listed by constant name instead, a plain mapper refuses every value
the schema allows.
+ List<Status> statuses = new ArrayList<>();
+ for (JsonNode value : properties.path("status").path("enum")) {
+ statuses.add(MAPPER.treeToValue(value, Status.class));
+ }
+ assertThat(statuses).containsExactlyInAnyOrder(Status.values());
+
+ List<Phase> phases = new ArrayList<>();
+ for (JsonNode value : properties.path("phase").path("enum")) {
+ phases.add(MAPPER.treeToValue(value, Phase.class));
+ }
+ assertThat(phases).containsExactlyInAnyOrder(Phase.values());
+ }
+
+ @Test
+ @DisplayName("The generated schema requires every field the caller did not
make optional")
+ void requiresEveryFieldExceptOptional() {
+ JsonNode schema = PojoJsonSchemaGenerator.generate(Report.class);
+
+ // Without a required set the model may omit fields the caller
declared, and with an
+ // all-inclusive one it must invent a value for the field the caller
made omissible.
+ assertThat(textValues(schema.path("required")))
+ .containsExactlyInAnyOrder("summary", "counts", "total");
+ }
+
+ @Test
+ @DisplayName("The generated schema keeps declaration order and emits no
getter property")
+ void keepsDeclarationOrderWithoutGetters() {
+ JsonNode schema = PojoJsonSchemaGenerator.generate(Report.class);
+
+ // The fixture declares its fields out of alphabetical order, so an
alphabetical sort
+ // reorders them, and a getter surfacing as a property adds an entry.
+ assertThat(fieldNames(schema.path("properties")))
+ .containsExactly("summary", "counts", "note", "total");
+ }
+
+ @Test
+ @DisplayName("The generated schema declares JSON Schema draft 2020-12")
+ void declaresDraft202012() {
+ JsonNode schema = PojoJsonSchemaGenerator.generate(Report.class);
+
+ assertThat(schema.path("$schema").asText())
+ .isEqualTo("https://json-schema.org/draft/2020-12/schema";);
+ }
+
+ @Test
+ @DisplayName("The generated schema applies an option only when it is
passed in")
+ void appliesAnOptionOnlyWhenPassedIn() {
+ JsonNode withOption =
+ PojoJsonSchemaGenerator.generate(
+ Report.class,
Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES)
+ .path("properties")
+ .path("counts");
+ JsonNode withoutOption =
+
PojoJsonSchemaGenerator.generate(Report.class).path("properties").path("counts");
+
+ // The map's value schema comes from the option: typed when it is
passed, absent when not.
+
assertThat(withOption.path("additionalProperties").path("type").asText())
+ .isEqualTo("integer");
+ assertThat(withoutOption.has("additionalProperties")).isFalse();
+ }
+
+ private static List<String> fieldNames(JsonNode node) {
+ List<String> names = new ArrayList<>();
+ node.fieldNames().forEachRemaining(names::add);
+ return names;
+ }
+
+ private static List<String> textValues(JsonNode array) {
+ List<String> values = new ArrayList<>();
+ array.forEach(value -> values.add(value.asText()));
+ return values;
+ }
+}
diff --git a/integrations/chat-models/gemini/pom.xml
b/integrations/chat-models/gemini/pom.xml
index 45516d3f..81587b63 100644
--- a/integrations/chat-models/gemini/pom.xml
+++ b/integrations/chat-models/gemini/pom.xml
@@ -44,15 +44,16 @@ under the License.
<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>
+ <groupId>org.apache.flink</groupId>
+
<artifactId>flink-agents-integrations-chat-models-common</artifactId>
+ <version>${project.version}</version>
</dependency>
+ <!-- Versions managed by the victools BOM imported in the root pom. -->
<dependency>
<groupId>com.github.victools</groupId>
- <artifactId>jsonschema-module-jackson</artifactId>
+ <artifactId>jsonschema-generator</artifactId>
</dependency>
</dependencies>
diff --git
a/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java
b/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java
index b4f128ab..3c34c431 100644
---
a/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java
+++
b/integrations/chat-models/gemini/src/main/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnection.java
@@ -23,13 +23,6 @@ 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;
@@ -47,6 +40,7 @@ import
org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.tools.ToolMetadata;
+import
org.apache.flink.agents.integrations.chatmodels.common.PojoJsonSchemaGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -401,17 +395,16 @@ public class GeminiChatModelConnection extends
BaseChatModelConnection {
// 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:
+ // without reporting which, so the shared schema and the options added to
it are checked
+ // 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()".
+ // - The shared schema is DRAFT_2020_12, 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.
// - 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.
+ // additionalProperties keyword carrying the declared value type.
Without it, 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
@@ -419,44 +412,20 @@ public class GeminiChatModelConnection extends
BaseChatModelConnection {
// 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.
+ // - Properties follow 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 the generator never emits.
//
// 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);
+ ObjectNode schema =
+ PojoJsonSchemaGenerator.generate(
+ schemaClass,
+ Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES,
+ Option.FORBIDDEN_ADDITIONAL_PROPERTIES_BY_DEFAULT);
stripRefSiblings(schema);
return schema;
}
diff --git
a/integrations/chat-models/gemini/src/test/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnectionTest.java
b/integrations/chat-models/gemini/src/test/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnectionTest.java
index 222db434..1cadda4e 100644
---
a/integrations/chat-models/gemini/src/test/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnectionTest.java
+++
b/integrations/chat-models/gemini/src/test/java/org/apache/flink/agents/integrations/chatmodels/gemini/GeminiChatModelConnectionTest.java
@@ -18,7 +18,6 @@
package org.apache.flink.agents.integrations.chatmodels.gemini;
-import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import com.fasterxml.jackson.annotation.JsonSubTypes;
@@ -113,8 +112,7 @@ class GeminiChatModelConnectionTest {
/**
* Output schema fixture shaped to expose the derivation settings: fields
are declared out of
- * alphabetical order, {@code counts} is a map whose values carry a type,
and {@code note} is
- * the only optional field.
+ * alphabetical order, and {@code counts} is a map whose values carry a
type.
*/
public static class Report {
public String summary;
@@ -123,21 +121,6 @@ class GeminiChatModelConnectionTest {
public int total;
}
- /**
- * Output schema fixture shaped to expose Jackson's property model.
- *
- * <p>{@code name} is deserialized from {@code full_name} rather than from
the Java field name,
- * and {@code secret} is not deserialized at all.
- */
- public static class Profile {
- @JsonProperty("full_name")
- public String name;
-
- @JsonIgnore public String secret;
-
- public int age;
- }
-
/** Nested type reused by two described fields of {@link Addresses}. */
public static class Address {
public String street;
@@ -800,20 +783,6 @@ class GeminiChatModelConnectionTest {
assertThat(config.responseMimeType()).isEmpty();
}
- @Test
- @DisplayName("The derived schema names properties the way Jackson
deserializes them")
- void derivedSchemaHonorsJacksonAnnotations() {
- // The response is read back with an ObjectMapper, which accepts the
renamed property and
- // rejects the Java field name, and which discards an ignored property
the schema would
- // otherwise force the model to fabricate.
- GenerateContentConfig config =
- connection()
- .buildConfig(userMessage(), null, params(),
CAPABLE_MODEL, Profile.class);
-
- assertThat(fieldNames(nativeSchema(config).path("properties")))
- .containsExactly("full_name", "age");
- }
-
@Test
@DisplayName("The derived schema lists enum constants by the values
Jackson deserializes")
void derivedSchemaListsEnumsByTheirJacksonWireValues() throws Exception {
@@ -861,20 +830,6 @@ class GeminiChatModelConnectionTest {
.isEqualTo("integer");
}
- @Test
- @DisplayName("The derived schema requires every field the caller did not
declare omissible")
- void derivedSchemaMarksNonOptionalFieldsRequired() {
- // Gemini treats a field the schema does not list as required as one
the model may skip,
- // so leaving `required` unset would let a response omit fields at
will.
- GenerateContentConfig config =
- connection()
- .buildConfig(userMessage(), null, params(),
CAPABLE_MODEL, Report.class);
-
- List<String> required = new ArrayList<>();
- nativeSchema(config).path("required").forEach(entry ->
required.add(entry.asText()));
- assertThat(required).containsExactlyInAnyOrder("summary", "counts",
"total");
- }
-
@Test
@DisplayName("A $ref carries no sibling that Gemini forbids beside it")
void derivedSchemaStripsRefSiblings() {
diff --git a/integrations/chat-models/ollama/pom.xml
b/integrations/chat-models/ollama/pom.xml
index 49638efd..6229e84a 100644
--- a/integrations/chat-models/ollama/pom.xml
+++ b/integrations/chat-models/ollama/pom.xml
@@ -44,15 +44,16 @@ under the License.
<version>${ollama4j.version}</version>
</dependency>
- <!-- Versions managed by the victools BOM imported in the root pom. -->
<dependency>
- <groupId>com.github.victools</groupId>
- <artifactId>jsonschema-generator</artifactId>
+ <groupId>org.apache.flink</groupId>
+
<artifactId>flink-agents-integrations-chat-models-common</artifactId>
+ <version>${project.version}</version>
</dependency>
+ <!-- Versions managed by the victools BOM imported in the root pom. -->
<dependency>
<groupId>com.github.victools</groupId>
- <artifactId>jsonschema-module-jackson</artifactId>
+ <artifactId>jsonschema-generator</artifactId>
</dependency>
</dependencies>
diff --git
a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
index be42f834..28857a3b 100644
---
a/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
+++
b/integrations/chat-models/ollama/src/main/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnection.java
@@ -22,12 +22,6 @@ import com.fasterxml.jackson.core.type.TypeReference;
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 io.github.ollama4j.exceptions.RoleNotFoundException;
import io.github.ollama4j.models.chat.*;
import io.github.ollama4j.models.request.OllamaChatEndpointCaller;
@@ -39,6 +33,7 @@ import
org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.tools.Tool;
+import
org.apache.flink.agents.integrations.chatmodels.common.PojoJsonSchemaGenerator;
import java.util.*;
import java.util.stream.Collectors;
@@ -329,29 +324,18 @@ public class OllamaChatModelConnection extends
BaseChatModelConnection {
return chatRequest;
}
- // Derives the JSON schema Ollama's format field expects from a POJO
class. Every setting below
- // addresses a concrete way the generated schema otherwise fails to
constrain generation:
+ // Derives the JSON schema Ollama's format field expects from a POJO
class. The schema comes
+ // from the shared generator with one option added:
//
- // - DRAFT_2020_12 is the draft pydantic generates on the Python side,
so a schema derived
- // from a Java class states the same contract in the same dialect.
- // - The PLAIN_JSON preset keeps generation to fields. Without a preset,
getters surface as
- // properties of their own, named after the accessor call, e.g.
"getSummary()".
// - MAP_VALUES_AS_ADDITIONAL_PROPERTIES gives a Map its value schema.
Without it the map
// admits any value, and a model does emit values that the declared
value type then fails
// to deserialize.
- // - Sorting fields before methods and applying no further comparison
leaves properties in
- // declaration order. Ollama's grammar fixes generation order to the
order the schema
- // declares its properties, so the default alphabetical order would
condition generation on
- // an order the class does not read in.
- // - The required check marks every field required except an Optional
one. The default marks
- // nothing required, which lets a model omit fields at will, 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.
- // It is applied with no JacksonOption, so it contributes property
naming and visibility
- // only: the required set and the property order stay the ones
configured below.
+ //
+ // When Ollama runs a model on llama.cpp, the grammar built from the
schema generates the
+ // required properties in the order the schema declares them, followed by
the optional ones,
+ // so the declaration order the shared schema keeps is the order the model
fills in a class's
+ // required fields. An alphabetical order would condition generation on an
order the class
+ // does not declare.
//
// Two settings are deliberately absent:
//
@@ -364,18 +348,8 @@ public class OllamaChatModelConnection extends
BaseChatModelConnection {
// $defs, so enabling it would silently unconstrain a common shape to
rescue a rare one. A
// recursive type instead fails loudly, with HTTP 400 from the server.
private static ObjectNode toNativeFormat(Class<?> schemaClass) {
- SchemaGeneratorConfigBuilder configBuilder =
- new SchemaGeneratorConfigBuilder(
- SchemaVersion.DRAFT_2020_12,
OptionPreset.PLAIN_JSON)
- .with(Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES)
- .with(new JacksonModule());
- configBuilder
- .forTypesInGeneral()
-
.withPropertySorter(PropertySortUtils.SORT_PROPERTIES_FIELDS_BEFORE_METHODS);
- configBuilder
- .forFields()
- .withRequiredCheck(field ->
!Optional.class.equals(field.getRawMember().getType()));
- return new
SchemaGenerator(configBuilder.build()).generateSchema(schemaClass);
+ return PojoJsonSchemaGenerator.generate(
+ schemaClass, Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES);
}
/**
diff --git
a/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
b/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
index f8b70633..3d8a0721 100644
---
a/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
+++
b/integrations/chat-models/ollama/src/test/java/org/apache/flink/agents/integrations/chatmodels/ollama/OllamaChatModelConnectionTest.java
@@ -17,9 +17,10 @@
*/
package org.apache.flink.agents.integrations.chatmodels.ollama;
-import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonValue;
import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
import io.github.ollama4j.models.chat.OllamaChatRequest;
import io.github.ollama4j.tools.Tools;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
@@ -41,7 +42,6 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.Optional;
import static org.assertj.core.api.Assertions.assertThat;
@@ -54,39 +54,46 @@ class OllamaChatModelConnectionTest {
private static final ResourceContext NOOP =
ResourceContext.fromGetResource((a, b) -> null);
- private static final String DRAFT_2020_12 =
"https://json-schema.org/draft/2020-12/schema";;
+ private static final ObjectMapper MAPPER = new ObjectMapper();
- /**
- * Output schema fixture shaped to expose the schema-generation settings.
- *
- * <p>Fields are declared out of alphabetical order, {@code counts} is a
map whose values carry
- * a type, {@code note} is the only optional field, and {@code getDerived}
is a getter backed by
- * no field.
- */
+ /** Output schema fixture with a plain field and a map whose values carry
a type. */
public static class Report {
public String summary;
public Map<String, Integer> counts;
- public Optional<String> note;
- public int total;
-
- public String getDerived() {
- return summary + total;
- }
}
/**
- * Output schema fixture shaped to expose Jackson's property model.
- *
- * <p>{@code name} is deserialized from {@code full_name} rather than from
the Java field name,
- * and {@code secret} is not deserialized at all.
+ * Output schema fixture whose enum constants are deserialized from values
other than their
+ * names, one through {@code @JsonProperty} on the constants and one
through a
+ * {@code @JsonValue} method.
*/
- public static class Profile {
- @JsonProperty("full_name")
- public String name;
+ public static class Ticket {
+ public Status status;
- @JsonIgnore public String secret;
+ public Phase phase;
+ }
- public int age;
+ public enum Status {
+ @JsonProperty("in-progress")
+ IN_PROGRESS,
+ @JsonProperty("done")
+ DONE
+ }
+
+ public enum Phase {
+ STARTED("started"),
+ FINISHED("finished");
+
+ private final String wire;
+
+ Phase(String wire) {
+ this.wire = wire;
+ }
+
+ @JsonValue
+ public String wire() {
+ return wire;
+ }
}
private static OllamaChatModelConnection connection() {
@@ -201,21 +208,13 @@ class OllamaChatModelConnectionTest {
}
@Test
- @DisplayName("The generated schema constrains draft, property order, map
values and required")
- void generatedSchemaShapeIsConstraining() {
+ @DisplayName("The generated schema gives map values their own schema")
+ void generatedSchemaGivesMapValuesTheirSchema() {
OllamaChatRequest request =
connection()
.buildRequest(userMessage(), List.of(),
params("qwen3:4b"), Report.class);
JsonNode schema = (JsonNode) request.getFormat();
- // Ollama fixes generation order to the order the schema declares its
properties, so the
- // emitted order has to follow the class rather than the alphabet. A
getter backed by no
- // field must not surface as a property of its own.
- assertThat(schema.path("$schema").asText()).isEqualTo(DRAFT_2020_12);
- assertThat(schema.path("properties").fieldNames())
- .toIterable()
- .containsExactly("summary", "counts", "note", "total");
-
// A map without a value schema admits any value, which the model does
take up and which
// then fails to deserialize into the declared type.
assertThat(schema.path("properties").path("counts").path("additionalProperties").isObject())
@@ -227,28 +226,30 @@ class OllamaChatModelConnectionTest {
.path("type")
.asText())
.isEqualTo("integer");
-
- // Every field is required except the one the caller declared
omissible.
- assertThat(textValues(schema.path("required")))
- .containsExactlyInAnyOrder("summary", "counts", "total");
}
@Test
- @DisplayName("The generated schema names properties the way Jackson
deserializes them")
- void generatedSchemaFollowsJacksonPropertyNames() {
+ @DisplayName("The generated schema lists enum constants the way Jackson
deserializes them")
+ void generatedSchemaFollowsJacksonEnumValues() throws Exception {
OllamaChatRequest request =
connection()
- .buildRequest(userMessage(), List.of(),
params("qwen3:4b"), Profile.class);
- JsonNode schema = (JsonNode) request.getFormat();
+ .buildRequest(userMessage(), List.of(),
params("qwen3:4b"), Ticket.class);
+ JsonNode properties = ((JsonNode)
request.getFormat()).path("properties");
+
+ // Every listed value is one the model may emit, so each has to
deserialize into the enum.
+ // Listed by constant name instead, the mapper reading the response
refuses every value
+ // the schema allows.
+ List<Status> statuses = new ArrayList<>();
+ for (JsonNode value : properties.path("status").path("enum")) {
+ statuses.add(MAPPER.treeToValue(value, Status.class));
+ }
+ assertThat(statuses).containsExactlyInAnyOrder(Status.values());
- // The response is read back with an ObjectMapper, which accepts the
renamed property and
- // rejects the Java field name, and which discards an ignored property
the schema would
- // otherwise force the model to fabricate.
- assertThat(schema.path("properties").fieldNames())
- .toIterable()
- .containsExactly("full_name", "age");
- assertThat(textValues(schema.path("required")))
- .containsExactlyInAnyOrder("full_name", "age");
+ List<Phase> phases = new ArrayList<>();
+ for (JsonNode value : properties.path("phase").path("enum")) {
+ phases.add(MAPPER.treeToValue(value, Phase.class));
+ }
+ assertThat(phases).containsExactlyInAnyOrder(Phase.values());
}
@ParameterizedTest
@@ -261,10 +262,4 @@ class OllamaChatModelConnectionTest {
// behavior change here.
assertThat(connection().supportsNativeStructuredOutput(model)).isTrue();
}
-
- private static List<String> textValues(JsonNode arrayNode) {
- List<String> values = new ArrayList<>();
- arrayNode.forEach(element -> values.add(element.asText()));
- return values;
- }
}
diff --git a/integrations/chat-models/pom.xml b/integrations/chat-models/pom.xml
index 37840ae7..2928abc6 100644
--- a/integrations/chat-models/pom.xml
+++ b/integrations/chat-models/pom.xml
@@ -33,6 +33,7 @@ under the License.
<modules>
<module>anthropic</module>
<module>bedrock</module>
+ <module>common</module>
<module>gemini</module>
<module>ollama</module>
<module>openai</module>
diff --git a/integrations/chat-models/watsonx/pom.xml
b/integrations/chat-models/watsonx/pom.xml
index c71baae3..6f8da02c 100644
--- a/integrations/chat-models/watsonx/pom.xml
+++ b/integrations/chat-models/watsonx/pom.xml
@@ -46,15 +46,16 @@ under the License.
<artifactId>jackson-databind</artifactId>
</dependency>
- <!-- Versions managed by the victools BOM imported in the root pom. -->
<dependency>
- <groupId>com.github.victools</groupId>
- <artifactId>jsonschema-generator</artifactId>
+ <groupId>org.apache.flink</groupId>
+
<artifactId>flink-agents-integrations-chat-models-common</artifactId>
+ <version>${project.version}</version>
</dependency>
+ <!-- Versions managed by the victools BOM imported in the root pom. -->
<dependency>
<groupId>com.github.victools</groupId>
- <artifactId>jsonschema-module-jackson</artifactId>
+ <artifactId>jsonschema-generator</artifactId>
</dependency>
<dependency>
diff --git
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
index 131d94cb..d8ddcc74 100644
---
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
+++
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
@@ -25,18 +25,13 @@ import com.fasterxml.jackson.databind.json.JsonMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
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 org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.tools.Tool;
+import
org.apache.flink.agents.integrations.chatmodels.common.PojoJsonSchemaGenerator;
import org.apache.flink.annotation.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -53,7 +48,6 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.regex.Matcher;
@@ -546,29 +540,17 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
return payload;
}
- // Derives the response_format value watsonx.ai expects from a POJO class.
Every setting below
- // addresses a concrete way the generated schema otherwise misstates the
contract:
+ // Derives the response_format value watsonx.ai expects from a POJO class.
The schema inside it
+ // comes from the shared generator with one option added:
//
- // - DRAFT_2020_12 is the draft pydantic generates on the Python side,
so a schema derived
- // from a Java class states the same contract in the same dialect.
- // - The PLAIN_JSON preset keeps generation to fields. Without a preset,
getters surface as
- // properties of their own, named after the accessor call, e.g.
"getDerived()".
// - MAP_VALUES_AS_ADDITIONAL_PROPERTIES gives a Map its value schema.
Without it the map
// admits any value, and a response that satisfies the schema can
still fail to deserialize
// into the declared value type at the caller.
- // - Sorting fields before methods and applying no further comparison
leaves properties in
- // declaration order, which is the order pydantic emits, so the
documents the two languages
- // derive from the same shape stay aligned.
- // - The required check marks every field required except an Optional
one. The default marks
- // nothing required, which lets a model omit fields at will, 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, because
- // a caller deserializing the response into this class honors
@JsonProperty and skips
- // @JsonIgnore. This connection returns the content as a string and
never deserializes into
- // the schema class, so a property stated under the wrong name
produces a response that
- // satisfies the schema and still fails to read back, at the caller
rather than here. It is
- // applied with no JacksonOption, so it contributes property naming
and visibility only:
- // the required set and the property order stay the ones configured
below.
+ //
+ // This connection returns the content as a string and never deserializes
into the schema
+ // class, so a property or enum constant the schema states under a name
Jackson does not read
+ // produces a response that satisfies the schema and still fails to read
back, at the caller
+ // rather than here.
//
// Two settings are deliberately absent:
//
@@ -584,24 +566,14 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
// recursive POJO reaches the provider unguarded either way, and this
setting cannot
// change that.
private static ObjectNode toNativeResponseFormat(Class<?> schemaClass) {
- SchemaGeneratorConfigBuilder configBuilder =
- new SchemaGeneratorConfigBuilder(
- SchemaVersion.DRAFT_2020_12,
OptionPreset.PLAIN_JSON)
- .with(Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES)
- .with(new JacksonModule());
- configBuilder
- .forTypesInGeneral()
-
.withPropertySorter(PropertySortUtils.SORT_PROPERTIES_FIELDS_BEFORE_METHODS);
- configBuilder
- .forFields()
- .withRequiredCheck(field ->
!Optional.class.equals(field.getRawMember().getType()));
-
final ObjectNode responseFormat = MAPPER.createObjectNode();
responseFormat.put("type", "json_schema");
final ObjectNode jsonSchema = responseFormat.putObject("json_schema");
jsonSchema.put("name", schemaClass.getSimpleName());
jsonSchema.set(
- "schema", new
SchemaGenerator(configBuilder.build()).generateSchema(schemaClass));
+ "schema",
+ PojoJsonSchemaGenerator.generate(
+ schemaClass,
Option.MAP_VALUES_AS_ADDITIONAL_PROPERTIES));
jsonSchema.put("strict", true);
return responseFormat;
}
diff --git
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
index 5dda6f7b..a1d1401e 100644
---
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
+++
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
@@ -17,8 +17,8 @@
*/
package org.apache.flink.agents.integrations.chatmodels.watsonx;
-import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonValue;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -65,16 +65,14 @@ class WatsonxChatModelConnectionTest {
private static final String CHAT_RESPONSE =
"{\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\","
+ "\"content\":\"Hello!\"},\"finish_reason\":\"stop\"}]}";
- private static final String DRAFT_2020_12 =
"https://json-schema.org/draft/2020-12/schema";;
private static final String MODEL = "ibm/granite-3-3-8b-instruct";
private static final Map<String, Object> CALLER_FORMAT = Map.of("type",
"json_object");
/**
* Output schema fixture shaped to expose the schema-generation settings.
*
- * <p>Fields are declared out of alphabetical order, {@code counts} is a
map whose values carry
- * a type, {@code note} is the only optional field, and {@code getDerived}
is a getter backed by
- * no field.
+ * <p>{@code counts} is a map whose values carry a type, and {@code
getDerived} is a getter
+ * backed by no field.
*/
public static class Report {
public String summary;
@@ -88,18 +86,37 @@ class WatsonxChatModelConnectionTest {
}
/**
- * Output schema fixture shaped to expose Jackson's property model.
- *
- * <p>{@code name} is deserialized from {@code full_name} rather than from
the Java field name,
- * and {@code secret} is not deserialized at all.
+ * Output schema fixture whose enum constants are deserialized from values
other than their
+ * names, one through {@code @JsonProperty} on the constants and one
through a
+ * {@code @JsonValue} method.
*/
- public static class Profile {
- @JsonProperty("full_name")
- public String name;
+ public static class Ticket {
+ public Status status;
+
+ public Phase phase;
+ }
+
+ public enum Status {
+ @JsonProperty("in-progress")
+ IN_PROGRESS,
+ @JsonProperty("done")
+ DONE
+ }
+
+ public enum Phase {
+ STARTED("started"),
+ FINISHED("finished");
- @JsonIgnore public String secret;
+ private final String wire;
+
+ Phase(String wire) {
+ this.wire = wire;
+ }
- public int age;
+ @JsonValue
+ public String wire() {
+ return wire;
+ }
}
private static ResourceDescriptor descriptor(String url, String apiKey,
String projectId) {
@@ -191,12 +208,6 @@ class WatsonxChatModelConnectionTest {
return names;
}
- private static List<String> textValues(JsonNode arrayNode) {
- List<String> values = new ArrayList<>();
- arrayNode.forEach(element -> values.add(element.asText()));
- return values;
- }
-
@ParameterizedTest(name = "{0}")
@MethodSource("missingRequiredConfiguration")
void testConstructorRejectsMissingRequiredConfiguration(
@@ -775,9 +786,6 @@ class WatsonxChatModelConnectionTest {
assertThat(jsonSchema.path("strict").booleanValue()).isTrue();
assertThat(fieldNames(jsonSchema.path("schema").path("properties")))
.containsExactlyInAnyOrder("summary", "counts", "note",
"total");
- // The draft is the one pydantic emits on the Python side, so both
languages state the
- // same contract in the same dialect. Nothing else in this suite pins
the draft.
-
assertThat(jsonSchema.path("schema").path("$schema").asText()).isEqualTo(DRAFT_2020_12);
}
@Test
@@ -805,28 +813,24 @@ class WatsonxChatModelConnectionTest {
}
@Test
- @DisplayName("The derived schema names properties the way Jackson
deserializes them")
- void derivedSchemaHonorsJacksonAnnotations() {
- JsonNode schema = derivedSchema(Profile.class);
-
- // A caller deserializing the response into this class accepts the
renamed property and
- // rejects the Java field name, and discards an ignored property that
the schema would
- // otherwise state as required and so force a value for.
- assertThat(fieldNames(schema.path("properties")))
- .containsExactlyInAnyOrder("full_name", "age");
- assertThat(textValues(schema.path("required")))
- .containsExactlyInAnyOrder("full_name", "age");
- }
+ @DisplayName("The derived schema lists enum constants the way Jackson
deserializes them")
+ void derivedSchemaFollowsJacksonEnumValues() throws Exception {
+ JsonNode properties = derivedSchema(Ticket.class).path("properties");
+
+ // Every listed value is one the model may emit, so each has to
deserialize into the enum.
+ // Listed by constant name instead, the caller's mapper refuses every
value the schema
+ // allows.
+ List<Status> statuses = new ArrayList<>();
+ for (JsonNode value : properties.path("status").path("enum")) {
+ statuses.add(MAPPER.treeToValue(value, Status.class));
+ }
+ assertThat(statuses).containsExactlyInAnyOrder(Status.values());
- @Test
- @DisplayName("The derived schema requires every field the caller did not
make optional")
- void derivedSchemaMarksNonOptionalFieldsRequired() {
- JsonNode schema = derivedSchema(Report.class);
-
- // Without a required set the model may omit fields the caller
declared, and with an
- // all-inclusive one it must invent a value for the field the caller
made omissible.
- assertThat(textValues(schema.path("required")))
- .containsExactlyInAnyOrder("summary", "counts", "total");
+ List<Phase> phases = new ArrayList<>();
+ for (JsonNode value : properties.path("phase").path("enum")) {
+ phases.add(MAPPER.treeToValue(value, Phase.class));
+ }
+ assertThat(phases).containsExactlyInAnyOrder(Phase.values());
}
@Test
