weiqingy commented on code in PR #1097:
URL: https://github.com/apache/flink-agents/pull/1097#discussion_r4011678864
##########
integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java:
##########
@@ -173,19 +326,73 @@ public ChatMessage chat(
}
}
- 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. 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());
Review Comment:
#1117 was merged before the enum fix reached it, so the watsonx part is in
#1120 now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
