From bf070372cc1f5916b82292a5a6aa6c836912a0b5 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Wed, 2 Sep 2026 23:32:15 -0700 Subject: [PATCH 1/3] [integrations][bedrock] Extract buildRequest as a test seam The Bedrock connection built its ConverseRequest inline in chat(), and every request-building helper was private, so nothing could assert what the connection actually sends. Move the construction into a package-private buildRequest() and cover it. The extraction is behavior-preserving: the moved block is unchanged, and resolveModel() still runs first so a missing model still fails before anything else happens. Five tests, none of which had an equivalent before: model-id resolution across both the configured default and a per-call override, tool config, the system/conversation split, inference config, and message merging. Generated-by: Claude Code 2.1.259 (Claude Opus 5) --- .../bedrock/BedrockChatModelConnection.java | 50 +++++-- .../BedrockChatModelConnectionTest.java | 125 ++++++++++++++++++ 2 files changed, 162 insertions(+), 13 deletions(-) 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 86105572b..25bdf7c53 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 @@ -120,6 +120,42 @@ public BedrockChatModelConnection( @Override public ChatMessage chat( List messages, List tools, Map modelParams) { + ConverseRequest request = buildRequest(messages, tools, modelParams); + 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, and the inference + * configuration. + * + *

Package-private so a test can assert the request body without issuing a live call through + * the Bedrock runtime client. + * + *

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 + * @return the request to send to Converse + * @throws IllegalArgumentException if neither the call nor the connection supplies a model id + */ + ConverseRequest buildRequest( + List messages, List tools, Map modelParams) { String modelId = resolveModel(modelParams); List systemMsgs = @@ -173,19 +209,7 @@ public ChatMessage chat( } } - ConverseRequest request = requestBuilder.build(); - - 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; + return requestBuilder.build(); } 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 5e570f59d..7bb3f9d7c 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 @@ -23,8 +23,17 @@ 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 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 java.util.*; @@ -45,6 +54,34 @@ private static ResourceDescriptor descriptor(String region, String model) { 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 extraArgs = new HashMap<>(); + extraArgs.put("externalId", externalId); + return new ChatMessage(MessageRole.TOOL, content, extraArgs); + } + @Test @DisplayName("Constructor creates client with default region") void testConstructorDefaultRegion() { @@ -119,4 +156,92 @@ void testStripMarkdownFencesPlainBlock() { void testStripMarkdownFencesNull() { assertThat(BedrockChatModelConnection.stripMarkdownFences(null)).isNull(); } + + @Test + @DisplayName("buildRequest: the effective model id lands in the request") + void testBuildRequestResolvesModelId() { + List messages = List.of(ChatMessage.user("hello")); + + ConverseRequest fromConnection = connection().buildRequest(messages, null, Map.of()); + assertThat(fromConnection.modelId()) + .isEqualTo("us.anthropic.claude-sonnet-4-20250514-v1:0"); + + ConverseRequest fromCall = + connection().buildRequest(messages, null, Map.of("model", "per-call-model")); + 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()); + + 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()); + + 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 messages = List.of(ChatMessage.user("hello")); + + ConverseRequest configured = + connection() + .buildRequest(messages, null, Map.of("temperature", 0.7, "max_tokens", 64)); + 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()); + 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()); + + 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"); + } } From 0c56a3b456624650a4d3fb8f7aac59ea995e436c Mon Sep 17 00:00:00 2001 From: weiqingy Date: Thu, 3 Sep 2026 00:06:39 -0700 Subject: [PATCH 2/3] [integrations][bedrock] Bump the AWS SDK to 2.41.22 Converse exposes outputConfig, its native structured-output surface, from 2.41.22 onward; 2.41.21 does not have it. Pin the minimum version that carries the surface rather than the latest, to keep the change reviewable. The property is shared, so this also moves the Bedrock embedding model, the OpenSearch vector store and the S3 Vectors vector store. All four modules compile and test unchanged, with no source edits. NOTICE follows the resulting set: 27 AWS entries and 10 netty entries, which move to 4.1.130.Final transitively through netty-nio-client, plus a new entry for utils-lite. Nothing leaves the set. Generated-by: Claude Code 2.1.259 (Claude Opus 5) --- dist/src/main/resources/META-INF/NOTICE | 75 +++++++++++++------------ integrations/pom.xml | 2 +- 2 files changed, 39 insertions(+), 38 deletions(-) diff --git a/dist/src/main/resources/META-INF/NOTICE b/dist/src/main/resources/META-INF/NOTICE index aa53540ab..39b97f2d5 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/pom.xml b/integrations/pom.xml index 8c7b4cbae..624fef8d3 100644 --- a/integrations/pom.xml +++ b/integrations/pom.xml @@ -36,7 +36,7 @@ under the License. 2.6.18 4.8.0 2.12.0 - 2.32.16 + 2.41.22 1.56.0 From 5fd5a7a8530f4fa1a46f3e9def821ecb7d6001b7 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Thu, 3 Sep 2026 01:13:11 -0700 Subject: [PATCH 3/3] [integrations][bedrock] Apply Bedrock native structured output Bedrock inherited the base 4-arg chat(), which rejects any non-null output schema, so it was the last Java chat model with no native path. Wire Converse's outputConfig so the provider constrains the response, falling back to prompt engineering when the schema, the derivation or the model cannot support it. Capability is an exact match against the model ids AWS documents as supporting structured output, retried once after stripping a leading inference-profile segment. Support is per model rather than per family, so a prefix match would claim a capability the provider denies for one Qwen model while granting it for eight siblings. Anything unrecognised, ARNs included, answers false and takes the fallback: an ARN carries no model information, and a prompt router does not choose its model until the request runs. The AWS SDK ships no schema generator, so the schema is derived with victools and serialised, JsonSchemaDefinition taking a string rather than a JSON value. JacksonModule keeps @JsonProperty names and drops @JsonIgnore fields; without a required check victools emits no required key at all, which would let an empty document satisfy every schema. additionalProperties is never emitted with a non-false value because Bedrock rejects that, which leaves a map's values undescribed. No collision guard: this connection reads only model, temperature and max_tokens, so there is no channel for a caller to have set outputConfig already. Generated-by: Claude Code 2.1.259 (Claude Opus 5) --- integrations/chat-models/bedrock/pom.xml | 11 + .../bedrock/BedrockChatModelConnection.java | 191 ++++++++++- .../BedrockChatModelConnectionTest.java | 313 +++++++++++++++++- 3 files changed, 504 insertions(+), 11 deletions(-) diff --git a/integrations/chat-models/bedrock/pom.xml b/integrations/chat-models/bedrock/pom.xml index b2fb44888..eb88c88a0 100644 --- a/integrations/chat-models/bedrock/pom.xml +++ b/integrations/chat-models/bedrock/pom.xml @@ -43,6 +43,17 @@ under the License. bedrockruntime ${aws.sdk.version} + + + + com.github.victools + jsonschema-generator + + + + com.github.victools + jsonschema-module-jackson + 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 25bdf7c53..6421bd529 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,13 @@ 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 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 +45,12 @@ 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 +65,10 @@ 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 +103,51 @@ 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 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", + "mistral.mistral-large-3-675b-instruct", + "openai.gpt-oss-120b-1:0", + "openai.gpt-oss-20b-1:0", + "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"); + + // 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, and the strip is attempted only after the id + // itself fails to match, so a bare model id is never shortened. + 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,10 +177,61 @@ public BedrockChatModelConnection( .build(); } + /** + * Whether AWS documents structured-output support for {@code effectiveModel}. + * + *

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. + * + *

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. + * + *

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. + * + *

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 messages, List tools, Map modelParams) { - ConverseRequest request = buildRequest(messages, tools, 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 messages, + List tools, + Map modelParams, + Object outputSchema) { + ConverseRequest request = buildRequest(messages, tools, modelParams, outputSchema); String modelId = request.modelId(); ConverseResponse response = @@ -138,8 +249,8 @@ public ChatMessage chat( /** * Translate the flink-agents call arguments into a Converse request: the effective model id, - * the SYSTEM/conversation message split, the tool configuration, and the inference - * configuration. + * 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. * *

Package-private so a test can assert the request body without issuing a live call through * the Bedrock runtime client. @@ -151,11 +262,17 @@ public ChatMessage chat( * @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 messages, List tools, Map modelParams) { + List messages, + List tools, + Map modelParams, + Object outputSchema) { String modelId = resolveModel(modelParams); List systemMsgs = @@ -209,9 +326,75 @@ ConverseRequest buildRequest( } } + if (outputSchema instanceof Class && supportsNativeStructuredOutput(modelId)) { + requestBuilder.outputConfig(nativeOutputConfig((Class) outputSchema)); + } + return requestBuilder.build(); } + /** + * Wraps the schema derived from {@code schemaClass} in the request element Converse reads it + * from. + * + *

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. 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 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()); + configBuilder + .forFields() + .withRequiredCheck(field -> !Optional.class.equals(field.getRawMember().getType())); + return new SchemaGenerator(configBuilder.build()).generateSchema(schemaClass); + } + private static boolean isRetryable(Exception e) { String msg = e.toString(); return msg.contains("ThrottlingException") 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 7bb3f9d7c..b39a67b97 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,6 +18,10 @@ package org.apache.flink.agents.integrations.chatmodels.bedrock; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; +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; @@ -30,12 +34,19 @@ 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; @@ -119,6 +130,22 @@ void testChatThrowsWithoutModel() { .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 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() { @@ -162,12 +189,12 @@ void testStripMarkdownFencesNull() { void testBuildRequestResolvesModelId() { List messages = List.of(ChatMessage.user("hello")); - ConverseRequest fromConnection = connection().buildRequest(messages, null, Map.of()); + 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")); + connection().buildRequest(messages, null, Map.of("model", "per-call-model"), null); assertThat(fromCall.modelId()).isEqualTo("per-call-model"); } @@ -179,7 +206,8 @@ void testBuildRequestPreservesToolConfig() { .buildRequest( List.of(ChatMessage.user("hello")), List.of(new SchemaOnlyTool("{\"type\": \"object\"}")), - Map.of()); + Map.of(), + null); assertThat(request.toolConfig()).isNotNull(); assertThat(request.toolConfig().tools()).hasSize(1); @@ -198,7 +226,8 @@ void testBuildRequestPreservesSystemMessages() { .buildRequest( List.of(ChatMessage.system("be terse"), ChatMessage.user("hello")), null, - Map.of()); + Map.of(), + null); assertThat(request.system()).hasSize(1); assertThat(request.system().get(0).text()).isEqualTo("be terse"); @@ -214,12 +243,13 @@ void testBuildRequestPreservesInferenceConfig() { ConverseRequest configured = connection() - .buildRequest(messages, null, Map.of("temperature", 0.7, "max_tokens", 64)); + .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()); + ConverseRequest bare = connection().buildRequest(messages, null, Map.of(), null); assertThat(bare.inferenceConfig()).isNull(); } @@ -234,7 +264,8 @@ void testBuildRequestMergesConsecutiveToolMessages() { toolMessage("call-1", "first result"), toolMessage("call-2", "second result")), null, - Map.of()); + Map.of(), + null); assertThat(request.messages()).hasSize(2); Message merged = request.messages().get(1); @@ -244,4 +275,272 @@ void testBuildRequestMergesConsecutiveToolMessages() { .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. + * + *

{@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 note; + + public Map counts; + } + + /** + * Every model id the connection reports capable. + * + *

The list is the whole allowlist, so an entry dropped or mistyped fails here rather than + * narrowing capability silently. The four suffix shapes are part of what is under test: three + * entries carry no version suffix at all, and {@code openai.gpt-oss-120b-1:0} carries {@code + * -1:0} rather than {@code -v1:0}, so a rule that assumes one shape fails on the others. + */ + private static Stream 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", + "mistral.mistral-large-3-675b-instruct", + "openai.gpt-oss-120b-1:0", + "openai.gpt-oss-20b-1:0", + "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"); + } + + private static Map params(String model) { + Map params = new HashMap<>(); + params.put("model", model); + return params; + } + + /** + * Reads the schema the request carries. + * + *

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. + assertThat(connection().supportsNativeStructuredOutput(prefix + CAPABLE_MODEL)).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 prefix + // shared by four capable entries. 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 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 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 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")); + } + + @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(); + } }