NeuroDock/Plugin-API/src/main/java/me/neurodock/plugin/tool/ToolParameters.java

package me.neurodock.plugin.tool;

import org.json.JSONObject;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

/**
 * Represents the parameters of a tool.
 * This is used by Ollama to determine the parameters of a tool.
 */
public class ToolParameters {

    /**
     * Creates a new instance of {@link ToolParameters}.
     * @param properties The properties of the parameters
     * @param required The required parameters
     */
    private ToolParameters(JSONObject properties, String[] required) {
        this.properties = properties;
        this.required = required;
    };

    public JSONObject getJSON() {
        JSONObject json = new JSONObject();

        json.put("type", "object");
        json.put("properties", properties);
        json.put("required", required);

        return json;
    }

    /**
     * the properties of the parameters
     */
    JSONObject properties;
    /**
     * the required parameters
     */
    String[] required;

    /**
     * Gets the properties of the {@link ToolParameters}
     * @return The properties of the {@link ToolParameters}
     */
    public JSONObject getProperties() {
        return properties;
    }

    /**
     * Gets the required parameters of the {@link ToolParameters}
     * @return The required parameters of the {@link ToolParameters}
     */
    public String[] getRequired() {
        return required;
    }

    /**
     * Creates a new instance of {@link ToolParametersBuilder}.
     * @return The {@link ToolParametersBuilder}
     */
    public static ToolParametersBuilder builder() {
        return new ToolParametersBuilder();
    }

    /**
     * Creates an empty {@link ToolParameters}
     * @return an empty {@link ToolParameters}
     * @apiNote This is equvalent to
     * <pre>{@code
     *    ToolParameters.builder().build();
     * }</pre>
     */
    public static ToolParameters empty() {
        return builder().build();
    }

    /**
     * Represents a builder for {@link ToolParameters}.
     */
    public static class ToolParametersBuilder {
        /**
         * The properties of the parameters.
         */
        Map<String, Property> propertyMap = new HashMap<>();
        /**
         * The required parameters.
         */
        ArrayList<String> required = new ArrayList<>();

        /**
         * Add an optinal perameter to this {@link ToolParametersBuilder}
         * @param name The name of the parameter
         * @param type The type of the parameter
         * @param description The description of the parameter
         * @return The {@link ToolParametersBuilder}
         * @apiNote Prefer {@link #addProperty(String, Type, String, boolean)} to be explicit about required state
         */
        public ToolParametersBuilder addProperty(String name, Type type, String description) {
            return ToolParametersBuilder.this.addProperty(name, type, description, false);
        }

        /**
         * Add a potentialy required peremeter to this {@link ToolParametersBuilder}.
         * @param name The name of the parameter
         * @param type The type of the parameter
         * @param description The description of the parameter
         * @param required The required state of the parameter
         * @return The {@link ToolParametersBuilder}
         */
        public ToolParametersBuilder addProperty(String name, Type type, String description, boolean required) {
            if(name == null || type == null || description == null) {
                return this;
            }
            propertyMap.put(name, new Property(type.getType(), description));
            if(required) {
                this.required.add(name);
            }
            return this;
        }

        /**
         * Makes a previusly optinal perameter required for this {@link ToolParametersBuilder}
         * @param name The name of the parameter
         * @return The {@link ToolParametersBuilder}
         */
        public ToolParametersBuilder required(String name) {
            required.add(name);
            return this;
        }

        /**
         * Removes a property from the parameters.
         * @param name The name of the property to remove
         * @return The {@link ToolParametersBuilder}
         */
        public ToolParametersBuilder removeProperty(String name) {
            propertyMap.remove(name);
            required.remove(name);
            return this;
        }

        /**
         * Builds the {@link ToolParameters}
         * @return The {@link ToolParameters}
         */
        public ToolParameters build() {
            JSONObject properties = new JSONObject();
            for(String name : propertyMap.keySet()) {
                properties.put(name, new JSONObject(propertyMap.get(name).toString()));
            }
            return new ToolParameters(properties, required.toArray(new String[0]));
        }

        /**
         * Represents a property of a parameter.
         */
        private static class Property {
            /**
             * The type of the property.
             */
            String type;
            /**
             * The description of the property.
             */
            String description;

            /**
             * Creates a new instance of {@link Property}.
             * @param type The type of the property
             * @param description The description of the property
             */
            public Property(String type, String description) {
                this.type = type;
                this.description = description;
            }

            @Override
            public String toString() {
                JSONObject json = new JSONObject();

                json.put("type", type);
                json.put("description", description);

                return json.toString();
            }
        }

        /**
         * Represents the type of parameter.
         */
        public enum Type {
            /**
             * Represents a string parameter.
             */
            STRING("string"),
            /**
             * Represents an integer parameter.
             */
            INT("int"),
            /**
             * Represents a boolean parameter.
             */
            BOOLEAN("boolean"),
            /**
             * Represents a enum parameter.
             */
            ENUM("enum"),
            /**
             * Represents a array parameter.
             */
            ARRAY("array"),
            /**
             * Represents a object parameter.
             */
            OBJECT("object");

            /**
             * The type of the parameter.
             */
            private final String type;

            /**
             * Gets the type of the parameter.
             * @return The type of the parameter
             */
            public String getType() {
                return type;
            }

            /**
             * Creates a new instance of {@link Type}.
             * @param type The type of the parameter
             */
            Type(String type) {
                this.type = type;
            }
        }
    }
}