weiqingy commented on code in PR #1097:
URL: https://github.com/apache/flink-agents/pull/1097#discussion_r4011473762
##########
integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnection.java:
##########
@@ -173,19 +356,78 @@ 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, 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
Review Comment:
Yes, that works. Once the remaining provider PRs land, I'll add a structured
output section to the chat model docs that covers every provider, including
that an enum used in an output schema needs `@JsonProperty` on every constant.
--
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]
