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 4d119211 [integrations][bedrock] Apply Bedrock native structured
output (#1097)
4d119211 is described below
commit 4d119211e9a91c200ad23739ec898df6eb32a4ea
Author: Weiqing Yang <[email protected]>
AuthorDate: Mon Sep 14 19:38:42 2026 -0700
[integrations][bedrock] Apply Bedrock native structured output (#1097)
Generated-by: Claude Code 2.1.270 (Claude Opus 5)
---
dist/src/main/resources/META-INF/NOTICE | 75 +--
integrations/chat-models/bedrock/pom.xml | 11 +
.../bedrock/BedrockChatModelConnection.java | 264 ++++++++++-
.../bedrock/BedrockChatModelConnectionTest.java | 520 +++++++++++++++++++++
integrations/pom.xml | 2 +-
5 files changed, 823 insertions(+), 49 deletions(-)
diff --git a/dist/src/main/resources/META-INF/NOTICE
b/dist/src/main/resources/META-INF/NOTICE
index aa53540a..39b97f2d 100644
--- a/dist/src/main/resources/META-INF/NOTICE
+++ b/dist/src/main/resources/META-INF/NOTICE
@@ -51,16 +51,16 @@ This project bundles the following dependencies under the
Apache Software Licens
- org.jetbrains.kotlin:kotlin-stdlib-common:1.9.10
- org.jetbrains.kotlin:kotlin-reflect:1.8.10
- org.yaml:snakeyaml:2.3
-- io.netty:netty-handler:4.1.118.Final
-- io.netty:netty-resolver:4.1.118.Final
-- io.netty:netty-transport:4.1.118.Final
-- io.netty:netty-buffer:4.1.118.Final
-- io.netty:netty-codec:4.1.118.Final
-- io.netty:netty-codec-http:4.1.118.Final
-- io.netty:netty-codec-http2:4.1.118.Final
-- io.netty:netty-transport-native-unix-common:4.1.118.Final
-- io.netty:netty-transport-classes-epoll:4.1.118.Final
-- io.netty:netty-common:4.1.118.Final
+- io.netty:netty-handler:4.1.130.Final
+- io.netty:netty-resolver:4.1.130.Final
+- io.netty:netty-transport:4.1.130.Final
+- io.netty:netty-buffer:4.1.130.Final
+- io.netty:netty-codec:4.1.130.Final
+- io.netty:netty-codec-http:4.1.130.Final
+- io.netty:netty-codec-http2:4.1.130.Final
+- io.netty:netty-transport-native-unix-common:4.1.130.Final
+- io.netty:netty-transport-classes-epoll:4.1.130.Final
+- io.netty:netty-common:4.1.130.Final
- io.projectreactor:reactor-core:3.7.0
- org.elasticsearch.client:elasticsearch-rest-client:8.19.0
- org.apache.httpcomponents:httpclient:4.5.13
@@ -88,33 +88,34 @@ This project bundles the following dependencies under the
Apache Software Licens
- io.prometheus:simpleclient_tracer_common:0.16.0
- io.prometheus:simpleclient_tracer_otel:0.16.0
- io.prometheus:simpleclient_tracer_otel_agent:0.16.0
-- software.amazon.awssdk:bedrockruntime:2.32.16
-- software.amazon.awssdk:s3vectors:2.32.16
-- software.amazon.awssdk:sdk-core:2.32.16
-- software.amazon.awssdk:aws-core:2.32.16
-- software.amazon.awssdk:auth:2.32.16
-- software.amazon.awssdk:regions:2.32.16
-- software.amazon.awssdk:profiles:2.32.16
-- software.amazon.awssdk:utils:2.32.16
-- software.amazon.awssdk:annotations:2.32.16
-- software.amazon.awssdk:apache-client:2.32.16
-- software.amazon.awssdk:netty-nio-client:2.32.16
-- software.amazon.awssdk:http-client-spi:2.32.16
-- software.amazon.awssdk:http-auth:2.32.16
-- software.amazon.awssdk:http-auth-aws:2.32.16
-- software.amazon.awssdk:http-auth-aws-eventstream:2.32.16
-- software.amazon.awssdk:http-auth-spi:2.32.16
-- software.amazon.awssdk:identity-spi:2.32.16
-- software.amazon.awssdk:metrics-spi:2.32.16
-- software.amazon.awssdk:endpoints-spi:2.32.16
-- software.amazon.awssdk:checksums:2.32.16
-- software.amazon.awssdk:checksums-spi:2.32.16
-- software.amazon.awssdk:retries:2.32.16
-- software.amazon.awssdk:retries-spi:2.32.16
-- software.amazon.awssdk:aws-json-protocol:2.32.16
-- software.amazon.awssdk:protocol-core:2.32.16
-- software.amazon.awssdk:json-utils:2.32.16
-- software.amazon.awssdk:third-party-jackson-core:2.32.16
+- software.amazon.awssdk:bedrockruntime:2.41.22
+- software.amazon.awssdk:s3vectors:2.41.22
+- software.amazon.awssdk:sdk-core:2.41.22
+- software.amazon.awssdk:aws-core:2.41.22
+- software.amazon.awssdk:auth:2.41.22
+- software.amazon.awssdk:regions:2.41.22
+- software.amazon.awssdk:profiles:2.41.22
+- software.amazon.awssdk:utils:2.41.22
+- software.amazon.awssdk:utils-lite:2.41.22
+- software.amazon.awssdk:annotations:2.41.22
+- software.amazon.awssdk:apache-client:2.41.22
+- software.amazon.awssdk:netty-nio-client:2.41.22
+- software.amazon.awssdk:http-client-spi:2.41.22
+- software.amazon.awssdk:http-auth:2.41.22
+- software.amazon.awssdk:http-auth-aws:2.41.22
+- software.amazon.awssdk:http-auth-aws-eventstream:2.41.22
+- software.amazon.awssdk:http-auth-spi:2.41.22
+- software.amazon.awssdk:identity-spi:2.41.22
+- software.amazon.awssdk:metrics-spi:2.41.22
+- software.amazon.awssdk:endpoints-spi:2.41.22
+- software.amazon.awssdk:checksums:2.41.22
+- software.amazon.awssdk:checksums-spi:2.41.22
+- software.amazon.awssdk:retries:2.41.22
+- software.amazon.awssdk:retries-spi:2.41.22
+- software.amazon.awssdk:aws-json-protocol:2.41.22
+- software.amazon.awssdk:protocol-core:2.41.22
+- software.amazon.awssdk:json-utils:2.41.22
+- software.amazon.awssdk:third-party-jackson-core:2.41.22
- software.amazon.eventstream:eventstream:1.0.1
- com.google.genai:google-genai:1.56.0
- com.google.auto.value:auto-value-annotations:1.11.0
diff --git a/integrations/chat-models/bedrock/pom.xml
b/integrations/chat-models/bedrock/pom.xml
index b2fb4488..eb88c88a 100644
--- a/integrations/chat-models/bedrock/pom.xml
+++ b/integrations/chat-models/bedrock/pom.xml
@@ -43,6 +43,17 @@ under the License.
<artifactId>bedrockruntime</artifactId>
<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>
+ </dependency>
+
+ <dependency>
+ <groupId>com.github.victools</groupId>
+ <artifactId>jsonschema-module-jackson</artifactId>
+ </dependency>
</dependencies>
</project>
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 86105572..62cc250c 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
@@ -20,7 +20,14 @@ package
org.apache.flink.agents.integrations.chatmodels.bedrock;
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;
@@ -39,7 +46,12 @@ import
software.amazon.awssdk.services.bedrockruntime.model.ConversationRole;
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
import
software.amazon.awssdk.services.bedrockruntime.model.InferenceConfiguration;
+import
software.amazon.awssdk.services.bedrockruntime.model.JsonSchemaDefinition;
import software.amazon.awssdk.services.bedrockruntime.model.Message;
+import software.amazon.awssdk.services.bedrockruntime.model.OutputConfig;
+import software.amazon.awssdk.services.bedrockruntime.model.OutputFormat;
+import
software.amazon.awssdk.services.bedrockruntime.model.OutputFormatStructure;
+import software.amazon.awssdk.services.bedrockruntime.model.OutputFormatType;
import software.amazon.awssdk.services.bedrockruntime.model.SystemContentBlock;
import software.amazon.awssdk.services.bedrockruntime.model.ToolConfiguration;
import software.amazon.awssdk.services.bedrockruntime.model.ToolInputSchema;
@@ -54,6 +66,10 @@ 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;
import java.util.stream.Collectors;
/**
@@ -88,6 +104,80 @@ import java.util.stream.Collectors;
public class BedrockChatModelConnection extends BaseChatModelConnection {
private static final ObjectMapper MAPPER = new ObjectMapper();
+
+ // Models AWS documents structured-output support for on the
bedrock-runtime endpoint. There is
+ // no single list page: the feature page delegates the per-model answer to
the individual model
+ // cards, where each card carries it as a "Structured outputs" bullet in
the Supported or Not
+ // Supported column of its "Features supported using bedrock-runtime
endpoint" table.
+ //
+ // The ids are the Model ID column of each card's Programmatic Access
table, read from the
+ // bedrock-runtime row. A card commonly prints a different id for
bedrock-mantle and can carry
+ // opposite verdicts for the two, so the endpoint an id was read from is
part of what makes the
+ // entry correct. This connection calls Converse on bedrock-runtime.
+ //
+ // Matching is exact, never by prefix. A Bedrock id already pins the
vendor, the snapshot date
+ // and the version in one string, so there is no alias for a prefix to
cover, and a prefix would
+ // over-capture: "qwen.qwen3" admits qwen.qwen3-vl-235b-a22b, which AWS
documents as not
+ // supported, and "anthropic.claude-sonnet-4" admits
anthropic.claude-sonnet-4-20250514-v1:0,
+ // whose card carries no answer at all. Exact matching also keeps
irregular id shapes correct
+ // with no normalisation rule: mistral.mistral-large-3-675b-instruct
carries no version suffix,
+ // openai.gpt-oss-120b-1:0 carries "-1:0" rather than "-v1:0".
+ //
+ // A card whose capability table carries the bullet in neither column is
undocumented rather
+ // than negative, and is absent from this set for that reason.
+ private static final Set<String> NATIVE_STRUCTURED_OUTPUT_MODELS =
+ Set.of(
+ "anthropic.claude-sonnet-4-5-20250929-v1:0",
+ "anthropic.claude-opus-4-5-20251101-v1:0",
+ "anthropic.claude-haiku-4-5-20251001-v1:0",
+ "anthropic.claude-opus-4-6-v1",
+ "anthropic.claude-sonnet-4-6",
+ "deepseek.v3-v1:0",
+ "deepseek.v3.2",
+ "google.gemma-3-12b-it",
+ "google.gemma-3-27b-it",
+ "minimax.minimax-m2",
+ "minimax.minimax-m2.1",
+ "minimax.minimax-m2.5",
+ "mistral.mistral-large-3-675b-instruct",
+ "mistral.devstral-2-123b",
+ "mistral.magistral-small-2509",
+ "mistral.ministral-3-14b-instruct",
+ "mistral.ministral-3-3b-instruct",
+ "mistral.ministral-3-8b-instruct",
+ "mistral.voxtral-mini-3b-2507",
+ "mistral.voxtral-small-24b-2507",
+ "moonshot.kimi-k2-thinking",
+ "moonshotai.kimi-k2.5",
+ "nvidia.nemotron-nano-12b-v2",
+ "nvidia.nemotron-nano-3-30b",
+ "nvidia.nemotron-nano-9b-v2",
+ "nvidia.nemotron-super-3-120b",
+ "openai.gpt-oss-120b-1:0",
+ "openai.gpt-oss-20b-1:0",
+ "openai.gpt-5.6-luna",
+ "openai.gpt-oss-safeguard-120b",
+ "openai.gpt-oss-safeguard-20b",
+ "qwen.qwen3-235b-a22b-2507-v1:0",
+ "qwen.qwen3-32b-v1:0",
+ "qwen.qwen3-coder-30b-a3b-v1:0",
+ "qwen.qwen3-coder-480b-a35b-v1:0",
+ "qwen.qwen3-coder-next",
+ "qwen.qwen3-next-80b-a3b",
+ "writer.palmyra-vision-7b",
+ "zai.glm-4.7",
+ "zai.glm-4.7-flash",
+ "zai.glm-5");
+
+ // A cross-Region inference profile id is a model id behind a geographic
or global prefix, and
+ // AWS documents structured output as working through cross-Region
inference. The prefix set is
+ // open-ended — the documentation names members by example and states that
new profiles may
+ // be created — so a leading segment is matched by shape rather than
against a fixed list,
+ // which would already have missed the documented us-gov. profiles. The
charset excludes ":"
+ // and "/", so no ARN can be shortened this way. The strip is attempted
only after the id itself
+ // fails to match, so a listed model id always matches as itself.
+ private static final Pattern INFERENCE_PROFILE_PREFIX =
Pattern.compile("^[a-z0-9-]+\\.(.+)$");
+
private final BedrockRuntimeClient client;
private final String defaultModel;
private final RetryExecutor retryExecutor;
@@ -117,9 +207,102 @@ public class BedrockChatModelConnection extends
BaseChatModelConnection {
.build();
}
+ /**
+ * Whether AWS documents structured-output support for {@code
effectiveModel}.
+ *
+ * <p>See the allowlist above for the source of truth, for why the match
is exact, and for why a
+ * geographic or global inference-profile prefix is stripped before it.
+ *
+ * <p>Every ARN reports {@code false}. A provisioned-throughput,
imported-model,
+ * custom-model-deployment, application-inference-profile or
marketplace-endpoint ARN identifies
+ * a resource without naming the model behind it, and a prompt-router ARN
names a set whose
+ * member is chosen per request, so for none of them is an answer
derivable from the identifier
+ * the request carries. An unrecognized identifier reports {@code false}
so that it degrades to
+ * the prompt-engineering fallback rather than failing at the provider.
+ *
+ * <p>A null or blank model reports {@code false} rather than throwing:
{@code resolveModel}
+ * rejects one before a request is built, but this method is part of the
connection contract and
+ * answers for whatever it is given. Only the null case needs a guard of
its own, because the
+ * allowlist is an immutable Set whose {@code contains(null)} throws; a
blank model is merely
+ * absent from it.
+ *
+ * <p>Reads no instance state, so capability stays answerable
independently of how the
+ * connection was configured.
+ */
+ @Override
+ protected boolean supportsNativeStructuredOutput(String effectiveModel) {
+ // Load-bearing: the allowlist is an immutable Set, whose
contains(null) throws rather than
+ // reporting absence.
+ if (effectiveModel == null || effectiveModel.isBlank()) {
+ return false;
+ }
+ if (NATIVE_STRUCTURED_OUTPUT_MODELS.contains(effectiveModel)) {
+ return true;
+ }
+ Matcher profile = INFERENCE_PROFILE_PREFIX.matcher(effectiveModel);
+ return profile.matches() &&
NATIVE_STRUCTURED_OUTPUT_MODELS.contains(profile.group(1));
+ }
+
@Override
public ChatMessage chat(
List<ChatMessage> messages, List<Tool> tools, Map<String, Object>
modelParams) {
+ return chat(messages, tools, modelParams, null);
+ }
+
+ /**
+ * Translates {@code outputSchema} into Converse's native {@code
outputConfig} when it is a POJO
+ * {@link Class} and the effective model is one AWS documents as
supporting it. Any other schema
+ * form — notably a {@code RowTypeInfo} wrapped in {@code OutputSchema} —
and any other model
+ * leave the request unconstrained, so that the caller keeps the
prompt-engineering fallback.
+ */
+ @Override
+ public ChatMessage chat(
+ List<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
+ ConverseRequest request = buildRequest(messages, tools, modelParams,
outputSchema);
+ String modelId = request.modelId();
+
+ ConverseResponse response =
+ retryExecutor.execute(() -> client.converse(request),
"BedrockConverse");
+
+ ChatMessage result = convertResponse(response);
+ if (response.usage() != null) {
+ result.getExtraArgs().put("model_name", modelId);
+ result.getExtraArgs().put("promptTokens",
response.usage().inputTokens().longValue());
+ result.getExtraArgs()
+ .put("completionTokens",
response.usage().outputTokens().longValue());
+ }
+ return result;
+ }
+
+ /**
+ * Translate the flink-agents call arguments into a Converse request: the
effective model id,
+ * the SYSTEM/conversation message split, the tool configuration, the
inference configuration,
+ * and the native output configuration when the schema and the model both
admit one.
+ *
+ * <p>Package-private so a test can assert the request body without
issuing a live call through
+ * the Bedrock runtime client.
+ *
+ * <p>Resolving the model is the first step, so an absent model id is
rejected before any
+ * request state is built.
+ *
+ * @param messages the conversation, SYSTEM messages included; must not be
null
+ * @param tools the tools to advertise, or {@code null} / empty for none
+ * @param modelParams per-call parameters; {@code model}, {@code
temperature} and {@code
+ * max_tokens} are read, and {@code null} is accepted
+ * @param outputSchema the schema the response should conform to, or
{@code null} for an
+ * unconstrained response; applied natively only for a POJO {@link
Class} on a model that
+ * supports it, and otherwise left to the caller's prompt-engineering
fallback
+ * @return the request to send to Converse
+ * @throws IllegalArgumentException if neither the call nor the connection
supplies a model id
+ */
+ ConverseRequest buildRequest(
+ List<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
String modelId = resolveModel(modelParams);
List<ChatMessage> systemMsgs =
@@ -173,19 +356,78 @@ public class BedrockChatModelConnection extends
BaseChatModelConnection {
}
}
- ConverseRequest request = requestBuilder.build();
+ if (outputSchema instanceof Class &&
supportsNativeStructuredOutput(modelId)) {
+ requestBuilder.outputConfig(nativeOutputConfig((Class<?>)
outputSchema));
+ }
- ConverseResponse response =
- retryExecutor.execute(() -> client.converse(request),
"BedrockConverse");
+ return requestBuilder.build();
+ }
- ChatMessage result = convertResponse(response);
- if (response.usage() != null) {
- result.getExtraArgs().put("model_name", modelId);
- result.getExtraArgs().put("promptTokens",
response.usage().inputTokens().longValue());
- result.getExtraArgs()
- .put("completionTokens",
response.usage().outputTokens().longValue());
- }
- return result;
+ /**
+ * Wraps the schema derived from {@code schemaClass} in the request
element Converse reads it
+ * from.
+ *
+ * <p>Converse takes the schema as serialized text rather than as a
document, unlike the tool
+ * input schema on the same request, so the derived schema is written out
here.
+ */
+ private static OutputConfig nativeOutputConfig(Class<?> schemaClass) {
+ return OutputConfig.builder()
+ .textFormat(
+ OutputFormat.builder()
+ .type(OutputFormatType.JSON_SCHEMA)
+ .structure(
+ OutputFormatStructure.builder()
+ .jsonSchema(
+
JsonSchemaDefinition.builder()
+ .schema(
+
toNativeSchema(schemaClass)
+
.toString())
+ .build())
+ .build())
+ .build())
+ .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.
+ //
+ // 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
+ // would trade an unconstrained map for a rejected request. A Map field
reaches the model as a
+ // 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.
+ 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);
}
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 5e570f59..59214174 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,15 +18,36 @@
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;
+import com.fasterxml.jackson.databind.ObjectMapper;
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.api.tools.ToolMetadata;
+import org.apache.flink.agents.api.tools.ToolParameters;
+import org.apache.flink.agents.api.tools.ToolResponse;
+import org.apache.flink.agents.api.tools.ToolType;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.NullSource;
+import org.junit.jupiter.params.provider.ValueSource;
+import software.amazon.awssdk.core.document.Document;
+import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole;
+import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
+import software.amazon.awssdk.services.bedrockruntime.model.Message;
+import software.amazon.awssdk.services.bedrockruntime.model.OutputFormatType;
import java.util.*;
+import java.util.stream.Stream;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -45,6 +66,34 @@ class BedrockChatModelConnectionTest {
return b.build();
}
+ private static BedrockChatModelConnection connection() {
+ return new BedrockChatModelConnection(
+ descriptor("us-east-1",
"us.anthropic.claude-sonnet-4-20250514-v1:0"), NOOP);
+ }
+
+ /** Minimal tool carrying only metadata; never invoked in these tests. */
+ private static final class SchemaOnlyTool extends Tool {
+ SchemaOnlyTool(String inputSchema) {
+ super(new ToolMetadata("add", "Add two numbers.", inputSchema));
+ }
+
+ @Override
+ public ToolType getToolType() {
+ return ToolType.FUNCTION;
+ }
+
+ @Override
+ public ToolResponse call(ToolParameters parameters) {
+ throw new UnsupportedOperationException("not invoked in this
test");
+ }
+ }
+
+ private static ChatMessage toolMessage(String externalId, String content) {
+ Map<String, Object> extraArgs = new HashMap<>();
+ extraArgs.put("externalId", externalId);
+ return new ChatMessage(MessageRole.TOOL, content, extraArgs);
+ }
+
@Test
@DisplayName("Constructor creates client with default region")
void testConstructorDefaultRegion() {
@@ -82,6 +131,22 @@ class BedrockChatModelConnectionTest {
.isInstanceOf(RuntimeException.class);
}
+ @Test
+ @DisplayName("chat accepts a POJO output schema instead of rejecting it as
untranslatable")
+ void testChatAcceptsPojoOutputSchema() {
+ // A connection that does not translate schemas itself inherits a
four-argument chat that
+ // refuses every non-null schema outright, so the feature can be
removed at the call
+ // boundary while buildRequest still wires outputConfig correctly.
Configuring no model
+ // makes the overridden path fail while resolving the model, which
reports that the schema
+ // was accepted and a request was being built, without any call
reaching the provider.
+ BedrockChatModelConnection conn =
+ new BedrockChatModelConnection(descriptor("us-east-1", null),
NOOP);
+ List<ChatMessage> msgs = List.of(new ChatMessage(MessageRole.USER,
"hello"));
+ assertThatThrownBy(() -> conn.chat(msgs, null, Collections.emptyMap(),
Profile.class))
+ .isNotInstanceOf(UnsupportedOperationException.class)
+ .isInstanceOf(IllegalArgumentException.class);
+ }
+
@Test
@DisplayName("stripMarkdownFences: normal text with braces is not
modified")
void testStripMarkdownFencesPreservesTextWithBraces() {
@@ -119,4 +184,459 @@ class BedrockChatModelConnectionTest {
void testStripMarkdownFencesNull() {
assertThat(BedrockChatModelConnection.stripMarkdownFences(null)).isNull();
}
+
+ @Test
+ @DisplayName("buildRequest: the effective model id lands in the request")
+ void testBuildRequestResolvesModelId() {
+ List<ChatMessage> messages = List.of(ChatMessage.user("hello"));
+
+ ConverseRequest fromConnection = connection().buildRequest(messages,
null, Map.of(), null);
+ assertThat(fromConnection.modelId())
+ .isEqualTo("us.anthropic.claude-sonnet-4-20250514-v1:0");
+
+ ConverseRequest fromCall =
+ connection().buildRequest(messages, null, Map.of("model",
"per-call-model"), null);
+ assertThat(fromCall.modelId()).isEqualTo("per-call-model");
+ }
+
+ @Test
+ @DisplayName("buildRequest: tools land in toolConfig")
+ void testBuildRequestPreservesToolConfig() {
+ ConverseRequest request =
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.user("hello")),
+ List.of(new SchemaOnlyTool("{\"type\":
\"object\"}")),
+ Map.of(),
+ null);
+
+ assertThat(request.toolConfig()).isNotNull();
+ assertThat(request.toolConfig().tools()).hasSize(1);
+
assertThat(request.toolConfig().tools().get(0).toolSpec().name()).isEqualTo("add");
+
assertThat(request.toolConfig().tools().get(0).toolSpec().description())
+ .isEqualTo("Add two numbers.");
+
assertThat(request.toolConfig().tools().get(0).toolSpec().inputSchema().json().asMap())
+ .containsEntry("type", Document.fromString("object"));
+ }
+
+ @Test
+ @DisplayName("buildRequest: system messages land in system, the rest in
messages")
+ void testBuildRequestPreservesSystemMessages() {
+ ConverseRequest request =
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.system("be terse"),
ChatMessage.user("hello")),
+ null,
+ Map.of(),
+ null);
+
+ assertThat(request.system()).hasSize(1);
+ assertThat(request.system().get(0).text()).isEqualTo("be terse");
+ assertThat(request.messages()).hasSize(1);
+
assertThat(request.messages().get(0).role()).isEqualTo(ConversationRole.USER);
+
assertThat(request.messages().get(0).content().get(0).text()).isEqualTo("hello");
+ }
+
+ @Test
+ @DisplayName("buildRequest: temperature and max_tokens land in
inferenceConfig")
+ void testBuildRequestPreservesInferenceConfig() {
+ List<ChatMessage> messages = List.of(ChatMessage.user("hello"));
+
+ ConverseRequest configured =
+ connection()
+ .buildRequest(
+ messages, null, Map.of("temperature", 0.7,
"max_tokens", 64), null);
+ assertThat(configured.inferenceConfig()).isNotNull();
+ assertThat(configured.inferenceConfig().temperature()).isEqualTo(0.7f);
+ assertThat(configured.inferenceConfig().maxTokens()).isEqualTo(64);
+
+ ConverseRequest bare = connection().buildRequest(messages, null,
Map.of(), null);
+ assertThat(bare.inferenceConfig()).isNull();
+ }
+
+ @Test
+ @DisplayName("buildRequest: consecutive tool messages merge into one user
message")
+ void testBuildRequestMergesConsecutiveToolMessages() {
+ ConverseRequest request =
+ connection()
+ .buildRequest(
+ List.of(
+ ChatMessage.user("hello"),
+ toolMessage("call-1", "first result"),
+ toolMessage("call-2", "second
result")),
+ null,
+ Map.of(),
+ null);
+
+ assertThat(request.messages()).hasSize(2);
+ Message merged = request.messages().get(1);
+ assertThat(merged.role()).isEqualTo(ConversationRole.USER);
+ assertThat(merged.content()).hasSize(2);
+ assertThat(merged.content())
+ .extracting(block -> block.toolResult().toolUseId())
+ .containsExactly("call-1", "call-2");
+ }
+
+ /** Documented on its AWS model card as supporting structured output. */
+ private static final String CAPABLE_MODEL =
"anthropic.claude-sonnet-4-5-20250929-v1:0";
+
+ /** Documented on its AWS model card as not supporting structured output.
*/
+ private static final String INCAPABLE_MODEL = "amazon.nova-micro-v1:0";
+
+ 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.
+ */
+ 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.
+ */
+ public static class Reading {
+ public int score;
+
+ public Optional<String> note;
+
+ public Map<String, Integer> counts;
+ }
+
+ /**
+ * 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;
+ }
+ }
+
+ /**
+ * Every model id the connection reports capable.
+ *
+ * <p>The list is the whole allowlist, so an entry dropped or mistyped
fails here rather than
+ * narrowing capability silently. The differing suffix shapes are part of
what is under test:
+ * {@code mistral.mistral-large-3-675b-instruct} carries no version
suffix, {@code
+ * openai.gpt-oss-120b-1:0} carries {@code -1:0}, {@code
+ * anthropic.claude-sonnet-4-5-20250929-v1:0} carries {@code -v1:0}, and
{@code
+ * anthropic.claude-opus-4-6-v1} carries {@code -v1} with no {@code :0},
so a rule that assumes
+ * one shape fails on the others.
+ */
+ private static Stream<String> capableModels() {
+ return Stream.of(
+ "anthropic.claude-sonnet-4-5-20250929-v1:0",
+ "anthropic.claude-opus-4-5-20251101-v1:0",
+ "anthropic.claude-haiku-4-5-20251001-v1:0",
+ "anthropic.claude-opus-4-6-v1",
+ "anthropic.claude-sonnet-4-6",
+ "deepseek.v3-v1:0",
+ "deepseek.v3.2",
+ "google.gemma-3-12b-it",
+ "google.gemma-3-27b-it",
+ "minimax.minimax-m2",
+ "minimax.minimax-m2.1",
+ "minimax.minimax-m2.5",
+ "mistral.mistral-large-3-675b-instruct",
+ "mistral.devstral-2-123b",
+ "mistral.magistral-small-2509",
+ "mistral.ministral-3-14b-instruct",
+ "mistral.ministral-3-3b-instruct",
+ "mistral.ministral-3-8b-instruct",
+ "mistral.voxtral-mini-3b-2507",
+ "mistral.voxtral-small-24b-2507",
+ "moonshot.kimi-k2-thinking",
+ "moonshotai.kimi-k2.5",
+ "nvidia.nemotron-nano-12b-v2",
+ "nvidia.nemotron-nano-3-30b",
+ "nvidia.nemotron-nano-9b-v2",
+ "nvidia.nemotron-super-3-120b",
+ "openai.gpt-oss-120b-1:0",
+ "openai.gpt-oss-20b-1:0",
+ "openai.gpt-5.6-luna",
+ "openai.gpt-oss-safeguard-120b",
+ "openai.gpt-oss-safeguard-20b",
+ "qwen.qwen3-235b-a22b-2507-v1:0",
+ "qwen.qwen3-32b-v1:0",
+ "qwen.qwen3-coder-30b-a3b-v1:0",
+ "qwen.qwen3-coder-480b-a35b-v1:0",
+ "qwen.qwen3-coder-next",
+ "qwen.qwen3-next-80b-a3b",
+ "writer.palmyra-vision-7b",
+ "zai.glm-4.7",
+ "zai.glm-4.7-flash",
+ "zai.glm-5");
+ }
+
+ private static Map<String, Object> params(String model) {
+ Map<String, Object> params = new HashMap<>();
+ params.put("model", model);
+ return params;
+ }
+
+ /**
+ * Reads the schema the request carries.
+ *
+ * <p>Walks the typed accessors because {@code ConverseRequest.toString()}
prints the structure
+ * as redacted sensitive data.
+ */
+ private static JsonNode nativeSchema(ConverseRequest request) throws
Exception {
+ return SCHEMA_MAPPER.readTree(
+
request.outputConfig().textFormat().structure().jsonSchema().schema());
+ }
+
+ @ParameterizedTest
+ @MethodSource("capableModels")
+ @DisplayName("every documented model reports capable")
+ void testCapableModelsReportCapable(String model) {
+ // connection() is configured with a model that is not on the list, so
a predicate reading
+ // the configured model rather than its argument disagrees with itself
here.
+
assertThat(connection().supportsNativeStructuredOutput(model)).isTrue();
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"us.", "eu.", "apac.", "au.", "jp.", "global."})
+ @DisplayName("a geographic inference-profile prefix resolves to the model
it fronts")
+ void testGeoPrefixResolvesToTheModelItFronts(String prefix) {
+ // A cross-Region inference profile id is a model id behind a leading
segment, and the model
+ // behind it is the one whose capability the request gets.
+ String profile = prefix + "anthropic.claude-opus-4-6-v1";
+
assertThat(connection().supportsNativeStructuredOutput(profile)).isTrue();
+ }
+
+ @Test
+ @DisplayName("the us-gov inference-profile prefix resolves to the model it
fronts")
+ void testHyphenatedPrefixResolvesToTheModelItFronts() {
+ // us-gov. is a documented prefix that no other documented prefix
resembles, so a rule
+ // written as a fixed set of prefixes tends to omit it while a
leading-segment strip covers
+ // it without being told.
+
assertThat(connection().supportsNativeStructuredOutput("us-gov.openai.gpt-oss-120b-1:0"))
+ .isTrue();
+ }
+
+ @Test
+ @DisplayName("a model documented as unsupported reports not capable")
+ void testDocumentedUnsupportedModelReportsNotCapable() {
+ // AWS documents this model as not supporting structured output, and
it extends the
+ // qwen.qwen3- prefix that several capable entries share. Any prefix
match claims a
+ // capability the provider denies.
+
assertThat(connection().supportsNativeStructuredOutput("qwen.qwen3-vl-235b-a22b"))
+ .isFalse();
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ INCAPABLE_MODEL,
+ "anthropic.claude-sonnet-4-20250514-v1:0",
+ "mistral.mistral-large-2402-v1:0",
+ "us.anthropic.claude-sonnet-4-20250514-v1:0"
+ })
+ @DisplayName("a model with no documented answer reports not capable")
+ void testUndocumentedModelsReportNotCapable(String model) {
+ // An absent answer is not a positive one. The middle two each extend
a capable entry
+ // truncated at a version boundary; the last is the id this module's
own example uses, so
+ // its behavior is pinned here rather than discovered at the provider.
+
assertThat(connection().supportsNativeStructuredOutput(model)).isFalse();
+ }
+
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+
"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-sonnet-4-5-20250929-v1:0",
+
"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/anthropic.claude-sonnet-4-5-20250929-v1:0"
+ })
+ @DisplayName("an ARN reports not capable even when it spells out a capable
model")
+ void testArnFormsReportNotCapable(String model) {
+ // An ARN names a resource rather than a model, and an application
inference profile's
+ // trailing segment can spell an id it does not front. A substring
match reports both of
+ // these capable.
+
assertThat(connection().supportsNativeStructuredOutput(model)).isFalse();
+ }
+
+ @ParameterizedTest
+ @NullSource
+ @ValueSource(strings = {"", " "})
+ @DisplayName("a null or blank model reports not capable")
+ void testNullOrBlankModelReportsNotCapable(String model) {
+ // The guard is load-bearing rather than defensive: the allowlist is
an immutable Set, whose
+ // contains(null) throws instead of reporting absence.
+
assertThat(connection().supportsNativeStructuredOutput(model)).isFalse();
+ }
+
+ @Test
+ @DisplayName("the derived schema names properties the way Jackson
deserializes them")
+ void testDerivedSchemaFollowsJacksonPropertyNames() throws Exception {
+ ConverseRequest request =
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.user("hello")),
+ null,
+ params(CAPABLE_MODEL),
+ Profile.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");
+ }
+
+ @Test
+ @DisplayName("the derived schema lists enum constants the way Jackson
deserializes them")
+ void testDerivedSchemaFollowsJacksonEnumValues() throws Exception {
+ ConverseRequest request =
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.user("hello")),
+ null,
+ params(CAPABLE_MODEL),
+ Ticket.class);
+ JsonNode properties = nativeSchema(request).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 refuses every value the
schema allows.
+ List<Status> statuses = new ArrayList<>();
+ for (JsonNode value : properties.path("status").path("enum")) {
+ statuses.add(SCHEMA_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(SCHEMA_MAPPER.treeToValue(value, Phase.class));
+ }
+ assertThat(phases).containsExactlyInAnyOrder(Phase.values());
+ }
+
+ @Test
+ @DisplayName("the derived schema declares its dialect, requires fields,
and leaves maps bare")
+ void testDerivedSchemaConstrainsTheResponse() throws Exception {
+ ConverseRequest request =
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.user("hello")),
+ null,
+ params(CAPABLE_MODEL),
+ 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
+ @DisplayName("tool definitions and an output schema ride the same request")
+ void testToolConfigAndOutputConfigCoexist() {
+ ConverseRequest request =
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.user("hello")),
+ List.of(new SchemaOnlyTool("{\"type\":
\"object\"}")),
+ params(CAPABLE_MODEL),
+ Profile.class);
+
+ // Converse carries tool definitions and an output schema together, so
neither branch may
+ // suppress the other.
+ assertThat(request.toolConfig()).isNotNull();
+ assertThat(request.toolConfig().tools()).hasSize(1);
+ assertThat(request.outputConfig()).isNotNull();
+ assertThat(request.outputConfig().textFormat().type())
+ .isEqualTo(OutputFormatType.JSON_SCHEMA);
+ }
+
+ @Test
+ @DisplayName("the native path applies for a POJO class schema on a capable
model")
+ void testNativeSchemaAppliedWhenGateHolds() throws Exception {
+ ConverseRequest applied =
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.user("hello")),
+ null,
+ params(CAPABLE_MODEL),
+ Profile.class);
+
+ assertThat(applied.outputConfig().textFormat().type())
+ .isEqualTo(OutputFormatType.JSON_SCHEMA);
+ // Bedrock takes the schema as serialized text rather than as a typed
object, so what
+ // arrives has to parse back into the schema document itself.
+
assertThat(nativeSchema(applied).path("type").asText()).isEqualTo("object");
+ }
+
+ private static Stream<Arguments> gateFailures() {
+ return Stream.of(
+ Arguments.of(INCAPABLE_MODEL, Profile.class),
+ Arguments.of(CAPABLE_MODEL, null),
+ // A RowTypeInfo schema arrives wrapped rather than as a bare
Class and has no
+ // native translation here, so it degrades to the fallback
rather than failing.
+ // RowTypeInfo itself is not on this module's test classpath;
any non-Class object
+ // exercises the same gate.
+ Arguments.of(CAPABLE_MODEL, "row<name STRING>"));
+ }
+
+ @ParameterizedTest
+ @MethodSource("gateFailures")
+ @DisplayName("the native path is skipped for an incapable model or a
non-POJO schema")
+ void testNativeSchemaSkippedWhenGateFails(String model, Object
outputSchema) {
+ assertThat(
+ connection()
+ .buildRequest(
+ List.of(ChatMessage.user("hello")),
+ null,
+ params(model),
+ outputSchema)
+ .outputConfig())
+ .isNull();
+ }
}
diff --git a/integrations/pom.xml b/integrations/pom.xml
index 8c7b4cba..624fef8d 100644
--- a/integrations/pom.xml
+++ b/integrations/pom.xml
@@ -36,7 +36,7 @@ under the License.
<milvus.version>2.6.18</milvus.version>
<openai.version>4.8.0</openai.version>
<anthropic.version>2.12.0</anthropic.version>
- <aws.sdk.version>2.32.16</aws.sdk.version>
+ <aws.sdk.version>2.41.22</aws.sdk.version>
<google.genai.version>1.56.0</google.genai.version>
</properties>
