[OPIK-218] Remove limitations in dataset items

apps/opik-backend/src/main/java/com/comet/opik/api/DatasetItem.java

@@ -1,5 +1,6 @@
 package com.comet.opik.api;
 
+import com.comet.opik.api.validate.DatasetItemInputValidation;
 import com.comet.opik.api.validate.SourceValidation;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonView;
@@ -12,21 +13,28 @@
 
 import java.time.Instant;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
 @Builder(toBuilder = true)
 @JsonIgnoreProperties(ignoreUnknown = true)
 @JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
 @SourceValidation
+@DatasetItemInputValidation
 public record DatasetItem(
         @JsonView( {
                 DatasetItem.View.Public.class, DatasetItem.View.Write.class}) UUID id,
-        @JsonView({DatasetItem.View.Public.class, DatasetItem.View.Write.class}) @NotNull JsonNode input,
-        @JsonView({DatasetItem.View.Public.class, DatasetItem.View.Write.class}) JsonNode expectedOutput,
+        @JsonView({DatasetItem.View.Public.class,
+                DatasetItem.View.Write.class}) @Schema(deprecated = true, description = "to be deprecated soon, please use input_data field") JsonNode input,
+        @JsonView({DatasetItem.View.Public.class,
+                DatasetItem.View.Write.class}) @Schema(deprecated = true, description = "to be deprecated soon, please use input_data field") JsonNode expectedOutput,
         @JsonView({DatasetItem.View.Public.class, DatasetItem.View.Write.class}) JsonNode metadata,
         @JsonView({DatasetItem.View.Public.class, DatasetItem.View.Write.class}) UUID traceId,
         @JsonView({DatasetItem.View.Public.class, DatasetItem.View.Write.class}) UUID spanId,
         @JsonView({DatasetItem.View.Public.class, DatasetItem.View.Write.class}) @NotNull DatasetItemSource source,
+        @JsonView({DatasetItem.View.Public.class,
+                DatasetItem.View.Write.class}) Map<String, DatasetItemInputValue<?>> inputData,
         @JsonView({
                 DatasetItem.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) List<ExperimentItem> experimentItems,
         @JsonView({DatasetItem.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant createdAt,
@@ -42,7 +50,8 @@ public record DatasetItemPage(
                     DatasetItem.View.Public.class}) List<DatasetItem> content,
             @JsonView({DatasetItem.View.Public.class}) int page,
             @JsonView({DatasetItem.View.Public.class}) int size,
-            @JsonView({DatasetItem.View.Public.class}) long total) implements Page<DatasetItem>{
+            @JsonView({DatasetItem.View.Public.class}) long total,
+            @JsonView({DatasetItem.View.Public.class}) Set<String> columns) implements Page<DatasetItem>{
     }
 
     public static class View {

apps/opik-backend/src/main/java/com/comet/opik/api/DatasetItemInputValue.java

@@ -0,0 +1,93 @@
+package com.comet.opik.api;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.PropertyNamingStrategies;
+import com.fasterxml.jackson.databind.annotation.JsonNaming;
+import io.swagger.v3.oas.annotations.media.DiscriminatorMapping;
+import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.validation.constraints.NotBlank;
+import jakarta.validation.constraints.NotNull;
+import lombok.Data;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.experimental.SuperBuilder;
+
+import java.beans.ConstructorProperties;
+
+import static com.comet.opik.api.DatasetItemInputValue.InputValueType;
+import static com.comet.opik.api.DatasetItemInputValue.JsonValue;
+import static com.comet.opik.api.DatasetItemInputValue.StringValue;
+
+@Data
+@SuperBuilder(toBuilder = true)
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "type", visible = true)
+@JsonSubTypes({
+        @JsonSubTypes.Type(value = JsonValue.class, name = InputValueType.JSON_TYPE),
+        @JsonSubTypes.Type(value = StringValue.class, name = InputValueType.STRING_TYPE),
+})
+@Schema(name = "InputData", discriminatorProperty = "type", discriminatorMapping = {
+        @DiscriminatorMapping(value = InputValueType.JSON_TYPE, schema = JsonValue.class),
+        @DiscriminatorMapping(value = InputValueType.STRING_TYPE, schema = StringValue.class)
+})
+@RequiredArgsConstructor
+public abstract sealed class DatasetItemInputValue<T> {
+
+    @Getter
+    @RequiredArgsConstructor
+    public enum InputValueType {
+
+        JSON(InputValueType.JSON_TYPE),
+        STRING(InputValueType.STRING_TYPE);
+
+        public static final String JSON_TYPE = "json";
+        public static final String STRING_TYPE = "string";
+
+        public static InputValueType fromString(String value) {
+            return InputValueType.valueOf(value);
+        }
+
+        @com.fasterxml.jackson.annotation.JsonValue
+        private final String value;
+    }
+
+    @Getter
+    @SuperBuilder(toBuilder = true)
+    @JsonIgnoreProperties(ignoreUnknown = true)
+    @JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
+    public static final class JsonValue extends DatasetItemInputValue<JsonNode> {
+
+        @ConstructorProperties({"value"})
+        public JsonValue(@NotNull JsonNode value) {
+            super(value);
+        }
+
+        @Override
+        public InputValueType getType() {
+            return InputValueType.JSON;
+        }
+    }
+
+    @Getter
+    @SuperBuilder(toBuilder = true)
+    @JsonIgnoreProperties(ignoreUnknown = true)
+    @JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
+    public static final class StringValue extends DatasetItemInputValue<String> {
+
+        @ConstructorProperties({"value"})
+        public StringValue(@NotBlank String value) {
+            super(value);
+        }
+
+        @Override
+        public InputValueType getType() {
+            return InputValueType.STRING;
+        }
+    }
+
+    protected final T value;
+
+    public abstract InputValueType getType();
+}

apps/opik-backend/src/main/java/com/comet/opik/api/validate/DatasetItemInputValidation.java

@@ -0,0 +1,22 @@
+package com.comet.opik.api.validate;
+
+import jakarta.validation.Constraint;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Target({ElementType.FIELD, ElementType.ANNOTATION_TYPE, ElementType.TYPE})
+@Retention(RetentionPolicy.RUNTIME)
+@Constraint(validatedBy = {DatasetItemInputValidator.class})
+@Documented
+public @interface DatasetItemInputValidation {
+
+    String message() default "must provide either input or input_data";
+
+    Class<?>[] groups() default {};
+
+    Class<?>[] payload() default {};
+}

apps/opik-backend/src/main/java/com/comet/opik/api/validate/DatasetItemInputValidator.java

@@ -0,0 +1,22 @@
+package com.comet.opik.api.validate;
+
+import com.comet.opik.api.DatasetItem;
+import jakarta.validation.ConstraintValidator;
+import jakarta.validation.ConstraintValidatorContext;
+
+public class DatasetItemInputValidator implements ConstraintValidator<DatasetItemInputValidation, DatasetItem> {
+
+    @Override
+    public boolean isValid(DatasetItem datasetItem, ConstraintValidatorContext context) {
+        boolean result = datasetItem.input() != null || (datasetItem.inputData() != null && !datasetItem.inputData().isEmpty());
+
+        if (!result) {
+            context.disableDefaultConstraintViolation();
+            context.buildConstraintViolationWithTemplate("must provide either input or input_data")
+                    .addPropertyNode("input")
+                    .addConstraintViolation();
+        }
+
+        return result;
+    }
+}