[OPIK-218] Remove limitations in dataset items (#369)

* [OPIK-218] Remove limitations in dataset items * Add tests * Address PR review * Fixing contract * Fix fixture generation * Deprecate metadata * Add PR review comments * Remove dead code
comet-ml · Oct 16, 2024 · 6e5bb09 · 6e5bb09
1 parent a593751
commit 6e5bb09
Show file tree

Hide file tree

Showing 15 changed files with 446 additions and 110 deletions.
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/DatasetItem.java b/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,29 @@
 
 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}) JsonNode metadata,
+        @JsonView({DatasetItem.View.Public.class,
+                DatasetItem.View.Write.class}) @Schema(deprecated = true, description = "to be deprecated soon, please use data field") JsonNode input,
+        @JsonView({DatasetItem.View.Public.class,
+                DatasetItem.View.Write.class}) @Schema(deprecated = true, description = "to be deprecated soon, please use data field") JsonNode expectedOutput,
+        @JsonView({DatasetItem.View.Public.class,
+                DatasetItem.View.Write.class}) @Schema(deprecated = true, description = "to be deprecated soon, please use data field") 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, JsonNode> data,
         @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 +51,11 @@ 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<Column> columns) implements Page<DatasetItem>{
+
+        public record Column(String name, String type) {
+        }
     }
 
     public static class View {

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/DatasetsResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/DatasetsResource.java
@@ -301,7 +301,8 @@ public Response createDatasetItems(
                 return item.toBuilder().id(idGenerator.generateId()).build();
             }
             return item;
-        }).toList();
+        })
+                .toList();
 
         String workspaceId = requestContext.get().getWorkspaceId();
 

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/validate/DatasetItemInputValidation.java b/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 data field";
+
+    Class<?>[] groups() default {};
+
+    Class<?>[] payload() default {};
+}
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/validate/DatasetItemInputValidator.java b/apps/opik-backend/src/main/java/com/comet/opik/api/validate/DatasetItemInputValidator.java
@@ -0,0 +1,43 @@
+package com.comet.opik.api.validate;
+
+import com.comet.opik.api.DatasetItem;
+import com.fasterxml.jackson.databind.JsonNode;
+import jakarta.validation.ConstraintValidator;
+import jakarta.validation.ConstraintValidatorContext;
+import org.apache.commons.collections4.MapUtils;
+
+import java.util.Map;
+import java.util.Optional;
+
+public class DatasetItemInputValidator implements ConstraintValidator<DatasetItemInputValidation, DatasetItem> {
+
+    @Override
+    public boolean isValid(DatasetItem datasetItem, ConstraintValidatorContext context) {
+        boolean result = datasetItem.input() != null || MapUtils.isNotEmpty(datasetItem.data());
+
+        if (!result) {
+            context.disableDefaultConstraintViolation();
+            context.buildConstraintViolationWithTemplate("must provide either input or data field")
+                    .addPropertyNode("input")
+                    .addConstraintViolation();
+        }
+
+        if (result && datasetItem.data() != null) {
+            Optional<Map.Entry<String, JsonNode>> error = datasetItem.data().entrySet()
+                    .stream()
+                    .filter(entry -> entry.getValue() == null)
+                    .findAny();
+
+            if (error.isPresent()) {
+                context.disableDefaultConstraintViolation();
+                context.buildConstraintViolationWithTemplate("field must not contain null key or value")
+                        .addPropertyNode("data")
+                        .addConstraintViolation();
+            }
+
+            result = error.isEmpty();
+        }
+
+        return result;
+    }
+}