JsonObjectSchema.java

package dev.langchain4j.model.chat.request.json;

import dev.langchain4j.Experimental;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import static dev.langchain4j.internal.Utils.copyIfNotNull;
import static dev.langchain4j.internal.Utils.quoted;
import static java.util.Arrays.asList;

@Experimental
public class JsonObjectSchema implements JsonSchemaElement {

    private final String description;
    private final Map<String, JsonSchemaElement> properties;
    private final List<String> required;
    private final Boolean additionalProperties;
    private final Map<String, JsonSchemaElement> definitions;

    public JsonObjectSchema(BuilderFields builder) {
        this.description = builder.description;
        this.properties = copyIfNotNull(builder.properties);
        this.required = copyIfNotNull(builder.required);
        this.additionalProperties = builder.additionalProperties;
        this.definitions = copyIfNotNull(builder.definitions);
    }

    public String description() {
        return description;
    }

    public Map<String, JsonSchemaElement> properties() {
        return properties;
    }

    public List<String> required() {
        return required;
    }

    public Boolean additionalProperties() {
        return additionalProperties;
    }

    /**
     * Used together with {@link JsonReferenceSchema} when recursion is required.
     */
    public Map<String, JsonSchemaElement> definitions() {
        return definitions;
    }

    public static Builder builder() {
        return new Builder();
    }

    public abstract static class BuilderFields {

        protected String description;
        protected Map<String, JsonSchemaElement> properties = new LinkedHashMap<>();
        protected List<String> required = new ArrayList<>();
        protected Boolean additionalProperties;
        protected Map<String, JsonSchemaElement> definitions;

        public JsonObjectSchema build() {
            return new JsonObjectSchema(this);
        }
    }

    public static class Builder extends BuilderFields {

        public Builder description(String description) {
            this.description = description;
            return this;
        }

        /**
         * Sets the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addProperty(String, JsonSchemaElement)
         * @see #addStringProperty(String)
         * @see #addStringProperty(String, String)
         * @see #addIntegerProperty(String)
         * @see #addIntegerProperty(String, String)
         * @see #addNumberProperty(String)
         * @see #addNumberProperty(String, String)
         * @see #addBooleanProperty(String)
         * @see #addBooleanProperty(String, String)
         * @see #addEnumProperty(String, List)
         * @see #addEnumProperty(String, List, String)
         */
        public AllPropertySetBuilder properties(Map<String, JsonSchemaElement> properties) {
            AllPropertySetBuilder allPropertySetBuilder = new AllPropertySetBuilder();
            allPropertySetBuilder.properties = properties;
            allPropertySetBuilder.description = this.description;
            allPropertySetBuilder.required = this.required;
            allPropertySetBuilder.additionalProperties = this.additionalProperties;
            allPropertySetBuilder.definitions = this.definitions;
            return allPropertySetBuilder;
        }

        /**
         * Adds a single property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #properties(Map)
         * @see #addStringProperty(String)
         * @see #addStringProperty(String, String)
         * @see #addIntegerProperty(String)
         * @see #addIntegerProperty(String, String)
         * @see #addNumberProperty(String)
         * @see #addNumberProperty(String, String)
         * @see #addBooleanProperty(String)
         * @see #addBooleanProperty(String, String)
         * @see #addEnumProperty(String, List)
         * @see #addEnumProperty(String, List, String)
         */
        public PropertyAdderBuilder addProperty(String name, JsonSchemaElement jsonSchemaElement) {
            this.properties.put(name, jsonSchemaElement);
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single string property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addStringProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addStringProperty(String name) {
            this.properties.put(name, new JsonStringSchema());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single string property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addStringProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addStringProperty(String name, String description) {
            this.properties.put(name, JsonStringSchema.builder().description(description).build());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single integer property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addIntegerProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addIntegerProperty(String name) {
            this.properties.put(name, new JsonIntegerSchema());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single integer property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addIntegerProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addIntegerProperty(String name, String description) {
            this.properties.put(name, JsonIntegerSchema.builder().description(description).build());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single number property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addNumberProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addNumberProperty(String name) {
            this.properties.put(name, new JsonNumberSchema());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single number property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addNumberProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addNumberProperty(String name, String description) {
            this.properties.put(name, JsonNumberSchema.builder().description(description).build());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single boolean property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addBooleanProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addBooleanProperty(String name) {
            this.properties.put(name, new JsonBooleanSchema());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single boolean property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addBooleanProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addBooleanProperty(String name, String description) {
            this.properties.put(name, JsonBooleanSchema.builder().description(description).build());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single enum property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addEnumProperty(String, List, String)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addEnumProperty(String name, List<String> enumValues) {
            this.properties.put(name, JsonEnumSchema.builder().enumValues(enumValues).build());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        /**
         * Adds a single enum property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addEnumProperty(String, List)
         * @see #addProperty(String, JsonSchemaElement)
         * @see #properties(Map)
         */
        public PropertyAdderBuilder addEnumProperty(String name, List<String> enumValues, String description) {
            this.properties.put(name, JsonEnumSchema.builder().enumValues(enumValues).description(description).build());
            return new PropertyAdderBuilder(description, properties, required, additionalProperties, definitions);
        }

        public Builder required(List<String> required) {
            this.required = required;
            return this;
        }

        public Builder required(String... required) {
            return required(asList(required));
        }

        public Builder additionalProperties(Boolean additionalProperties) {
            this.additionalProperties = additionalProperties;
            return this;
        }

        /**
         * Used together with {@link JsonReferenceSchema} when recursion is required.
         */
        public Builder definitions(Map<String, JsonSchemaElement> definitions) {
            this.definitions = definitions;
            return this;
        }
    }

    public static class AllPropertySetBuilder extends BuilderFields {

        public AllPropertySetBuilder description(String description) {
            this.description = description;
            return this;
        }

        public AllPropertySetBuilder required(List<String> required) {
            this.required = required;
            return this;
        }

        public AllPropertySetBuilder required(String... required) {
            return required(asList(required));
        }

        public AllPropertySetBuilder additionalProperties(Boolean additionalProperties) {
            this.additionalProperties = additionalProperties;
            return this;
        }

        /**
         * Used together with {@link JsonReferenceSchema} when recursion is required.
         */
        public AllPropertySetBuilder definitions(Map<String, JsonSchemaElement> definitions) {
            this.definitions = definitions;
            return this;
        }
    }

    public static class PropertyAdderBuilder extends BuilderFields {

        PropertyAdderBuilder(String description, Map<String, JsonSchemaElement> properties, List<String> required, Boolean additionalProperties, Map<String, JsonSchemaElement> definitions) {
            this.description = description;
            this.properties = properties;
            this.required = required;
            this.additionalProperties = additionalProperties;
            this.definitions = definitions;
        }

        public PropertyAdderBuilder description(String description) {
            this.description = description;
            return this;
        }

        /**
         * Adds a single property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addStringProperty(String)
         * @see #addStringProperty(String, String)
         * @see #addIntegerProperty(String)
         * @see #addIntegerProperty(String, String)
         * @see #addNumberProperty(String)
         * @see #addNumberProperty(String, String)
         * @see #addBooleanProperty(String)
         * @see #addBooleanProperty(String, String)
         * @see #addEnumProperty(String, List)
         * @see #addEnumProperty(String, List, String)
         */
        public PropertyAdderBuilder addProperty(String name, JsonSchemaElement jsonSchemaElement) {
            this.properties.put(name, jsonSchemaElement);
            return this;
        }

        /**
         * Adds a single string property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addStringProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addStringProperty(String name) {
            this.properties.put(name, new JsonStringSchema());
            return this;
        }

        /**
         * Adds a single string property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addStringProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addStringProperty(String name, String description) {
            this.properties.put(name, JsonStringSchema.builder().description(description).build());
            return this;
        }

        /**
         * Adds a single integer property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addIntegerProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addIntegerProperty(String name) {
            this.properties.put(name, new JsonIntegerSchema());
            return this;
        }

        /**
         * Adds a single integer property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addIntegerProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addIntegerProperty(String name, String description) {
            this.properties.put(name, JsonIntegerSchema.builder().description(description).build());
            return this;
        }

        /**
         * Adds a single number property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addNumberProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addNumberProperty(String name) {
            this.properties.put(name, new JsonNumberSchema());
            return this;
        }

        /**
         * Adds a single number property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addNumberProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addNumberProperty(String name, String description) {
            this.properties.put(name, JsonNumberSchema.builder().description(description).build());
            return this;
        }

        /**
         * Adds a single boolean property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addBooleanProperty(String, String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addBooleanProperty(String name) {
            this.properties.put(name, new JsonBooleanSchema());
            return this;
        }

        /**
         * Adds a single boolean property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addBooleanProperty(String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addBooleanProperty(String name, String description) {
            this.properties.put(name, JsonBooleanSchema.builder().description(description).build());
            return this;
        }

        /**
         * Adds a single enum property to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addEnumProperty(String, List, String)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addEnumProperty(String name, List<String> enumValues) {
            this.properties.put(name, JsonEnumSchema.builder().enumValues(enumValues).build());
            return this;
        }

        /**
         * Adds a single enum property with a description to the properties of this JSON object.
         * Please note that {@link #required(List)} should be set explicitly if you want the properties to be mandatory.
         *
         * @see #addEnumProperty(String, List)
         * @see #addProperty(String, JsonSchemaElement)
         */
        public PropertyAdderBuilder addEnumProperty(String name, List<String> enumValues, String description) {
            this.properties.put(name, JsonEnumSchema.builder().enumValues(enumValues).description(description).build());
            return this;
        }

        public PropertyAdderBuilder required(List<String> required) {
            this.required = required;
            return this;
        }

        public PropertyAdderBuilder required(String... required) {
            return required(asList(required));
        }

        public PropertyAdderBuilder additionalProperties(Boolean additionalProperties) {
            this.additionalProperties = additionalProperties;
            return this;
        }

        /**
         * Used together with {@link JsonReferenceSchema} when recursion is required.
         */
        public PropertyAdderBuilder definitions(Map<String, JsonSchemaElement> definitions) {
            this.definitions = definitions;
            return this;
        }
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        JsonObjectSchema that = (JsonObjectSchema) o;
        return Objects.equals(this.description, that.description)
                && Objects.equals(this.properties, that.properties)
                && Objects.equals(this.required, that.required)
                && Objects.equals(this.additionalProperties, that.additionalProperties)
                && Objects.equals(this.definitions, that.definitions);
    }

    @Override
    public int hashCode() {
        return Objects.hash(description, properties, required, additionalProperties, definitions);
    }

    @Override
    public String toString() {
        return "JsonObjectSchema {" +
                "description = " + quoted(description) +
                ", properties = " + properties +
                ", required = " + required +
                ", additionalProperties = " + additionalProperties +
                ", definitions = " + definitions +
                " }";
    }
}