[OPIK-218] Remove limitations in dataset items

comet-ml · Oct 11, 2024 · ab2f38f · ab2f38f
1 parent 557b6c7
commit ab2f38f
Show file tree

Hide file tree

Showing 13 changed files with 531 additions and 96 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
@@ -12,6 +12,8 @@
 
 import java.time.Instant;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
 @Builder(toBuilder = true)
@@ -21,12 +23,16 @@
 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) @NotNull JsonNode input,
+        @JsonView({DatasetItem.View.Public.class,
+                DatasetItem.View.Write.class}) @Schema(deprecated = true) 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 +48,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 {

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/DatasetItemInputValue.java b/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();
+}