Amazon Bedrock structured output using tool calls

AmazonBedrockClient.java

import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.core.document.Document;
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient;
import software.amazon.awssdk.services.bedrockruntime.model.*;

import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

public class AmazonBedrockClient {
    private final BedrockRuntimeClient client;
    private final Consumer<InferenceConfiguration.Builder> inferenceFactory;

    private final static String DEFAULT_MODEL = "eu.amazon.nova-lite-v1:0";

    private final String modelId;
    private final Float temperature;

    public AmazonBedrockClient(String modelId, Double temperature) {
        this.temperature = temperature == null ? 0f : temperature.floatValue();
        this.modelId = modelId == null ? DEFAULT_MODEL : modelId;
        this.client = BedrockRuntimeClient.builder()
                .credentialsProvider(DefaultCredentialsProvider.create())
                .build();
        this.inferenceFactory = config -> {
            config.temperature(this.temperature);
        };
    }
    
    public AmazonBedrockClient() {
        this(null, null);
    }

    private static final String FORMAT_OUTPUT_TOOL_NAME = "format_output";
    private ToolConfiguration getToolConfiguration(Document schema) {
        var schemaDocument = Document.mapBuilder()
                .putString("type", "object")
                .putMap("properties", Map.of("response", schema))
                .putList("required", List.of(Document.fromString("response")))
                .putBoolean("additionalProperties", false)
                .build();
        var toolInputSchema = ToolInputSchema.fromJson(schemaDocument);
        var toolSpec = ToolSpecification.builder()
                .name(FORMAT_OUTPUT_TOOL_NAME)
                .inputSchema(toolInputSchema)
                .build();
        var tool = Tool.fromToolSpec(toolSpec);
        var specificToolChoice = SpecificToolChoice.builder()
                .name(FORMAT_OUTPUT_TOOL_NAME)
                .build();
        var toolChoice = ToolChoice.fromTool(specificToolChoice);

        return ToolConfiguration.builder()
            .tools(tool)
            .toolChoice(toolChoice)
            .build();
    }

    public Document chat(String userInput, Document schema) {
        var systemPrompt = "You MUST pass your response through the provided tool \"" + FORMAT_OUTPUT_TOOL_NAME + "\"";
        var sys = SystemContentBlock.fromText(systemPrompt);

        var m = Message.builder()
                .role(ConversationRole.USER)
                .content(ContentBlock.fromText(userInput))
                .build();

        var request = ConverseRequest.builder()
                .modelId(modelId)
                .system(sys)
                .messages(m)
                .toolConfig(getToolConfiguration(schema))
                .inferenceConfig(inferenceFactory)
                .build();

        var converseResponse = client.converse(request);
        var assistantResponse = converseResponse.output().message().content().get(0);
        if (StopReason.TOOL_USE.equals(converseResponse.stopReason())) {
            return assistantResponse.toolUse().input().asMap().get("response");
        }
        throw new RuntimeException(assistantResponse.text());
    }
}

AmazonBedrockClientTest.java

import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import software.amazon.awssdk.core.document.Document;

import java.util.List;
import java.util.Map;

public class AmazonBedrockClientTest {
    private Document s(final String value) {
        return Document.fromString(value);
    }
    @Test
    void testUsingFormatTool() {
        var client = new AmazonBedrockClient();
        var result = client.chat("Suggest a list of face features",
              Document.mapBuilder()
                .putString("$schema", "http://json-schema.org/draft-07/schema#")
                .putString("type", "array")
                .putDocument("items", Document.mapBuilder()
                    .putString("type", "object")
                    .putDocument("properties", Document.mapBuilder()
                        .putMap("name", Map.of("type", s("string")))
                        .putMap("description", Map.of("type", s("string")))
                        .putDocument("position", Document.mapBuilder()
                            .putString("type", "object")
                            .putDocument("properties", Document.mapBuilder()
                                .putMap("horizontal", Map.of(
                                     "type", s("string"),
                                     "enum", Document.fromList(List.of(s("SIDE"), s("MIDDLE")))
                                ))
                                .putMap("vertical", Map.of(
                                    "type", s("string"),
                                    "enum", Document.fromList(List.of(s("TOP"), s("MIDDLE"), s("BOTTOM")))
                                ))
                                .build()
                            )
                            .build()
                        )
                        .build()
                    )
                    .build()
                )
                .build());

        // Debug here
        Assertions.assertNotNull(result);
    }
}

result.json

[ {
  "name" : "Eyes",
  "description" : "The eyes are the organs of vision, located on the face.",
  "position" : {
    "horizontal" : "SIDE",
    "vertical" : "MIDDLE"
  }
}, {
  "name" : "Nose",
  "description" : "The nose is the organ for breathing and smelling, located in the center of the face.",
  "position" : {
    "horizontal" : "MIDDLE",
    "vertical" : "MIDDLE"
  }
}, {
  "name" : "Mouth",
  "description" : "The mouth is the opening through which food is taken in and vocal sounds are produced.",
  "position" : {
    "horizontal" : "MIDDLE",
    "vertical" : "BOTTOM"
  }
}, {
  "name" : "Ears",
  "description" : "The ears are the organs of hearing, located on the sides of the head.",
  "position" : {
    "horizontal" : "SIDE",
    "vertical" : "TOP"
  }
}, {
  "name" : "Eyebrows",
  "description" : "The eyebrows are the arches of hair above the eyes, used for expression and protection.",
  "position" : {
    "horizontal" : "SIDE",
    "vertical" : "TOP"
  }
}, {
  "name" : "Cheeks",
  "description" : "The cheeks are the fleshy parts of the face below the eyes and around the nose.",
  "position" : {
    "horizontal" : "SIDE",
    "vertical" : "MIDDLE"
  }
}, {
  "name" : "Chin",
  "description" : "The chin is the projection at the bottom of the face, below the mouth.",
  "position" : {
    "horizontal" : "MIDDLE",
    "vertical" : "BOTTOM"
  }
} ]