[OPIK-593] OnlineScoring sampling and preparing traces for Automation…

… Rules (#1002)

* Adapting Trace ingestion to also trigger OnlineScoring.
* Changing LlmAsJudge typing from a free JsonNode into a proper defined object.
* Adapting methods to the defined evaluator payload
* Adapting contract to the new agreement and adding tests
* navigating using to JsonPath
* removing constructors from records

apps/opik-backend/pom.xml

@@ -32,6 +32,7 @@
         <redisson.version>3.41.0</redisson.version>
         <opentelmetry.version>2.10.0</opentelmetry.version>
         <aws.java.sdk.version>2.29.9</aws.java.sdk.version>
+        <json-path.version>2.9.0</json-path.version>
         <mainClass>com.comet.opik.OpikApplication</mainClass>
     </properties>
 
@@ -207,6 +208,11 @@
             <artifactId>java-uuid-generator</artifactId>
             <version>${uuid.java.generator.version}</version>
         </dependency>
+        <dependency>
+            <groupId>com.jayway.jsonpath</groupId>
+            <artifactId>json-path</artifactId>
+            <version>${json-path.version}</version>
+        </dependency>
         <dependency>
             <groupId>dev.langchain4j</groupId>
             <artifactId>langchain4j-open-ai</artifactId>

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

@@ -3,19 +3,15 @@
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.annotation.JsonView;
-import com.fasterxml.jackson.databind.JsonNode;
 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.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
-import lombok.EqualsAndHashCode;
-import lombok.ToString;
 import lombok.experimental.SuperBuilder;
 
-import java.beans.ConstructorProperties;
 import java.time.Instant;
 import java.util.List;
 import java.util.UUID;
@@ -24,38 +20,16 @@
 @SuperBuilder(toBuilder = true)
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "type", visible = true)
 @JsonSubTypes({
-        @JsonSubTypes.Type(value = AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class, name = "llm_as_judge")
+        @JsonSubTypes.Type(value = AutomationRuleEvaluatorLlmAsJudge.class, name = "llm_as_judge")
 })
 @Schema(name = "AutomationRuleEvaluator", discriminatorProperty = "type", discriminatorMapping = {
-        @DiscriminatorMapping(value = "llm_as_judge", schema = AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class)
+        @DiscriminatorMapping(value = "llm_as_judge", schema = AutomationRuleEvaluatorLlmAsJudge.class)
 })
 @AllArgsConstructor
-public abstract sealed class AutomationRuleEvaluator<T> implements AutomationRule<T> {
-
-    @EqualsAndHashCode(callSuper = true)
-    @Data
-    @SuperBuilder(toBuilder = true)
-    @ToString(callSuper = true)
-    public static final class AutomationRuleEvaluatorLlmAsJudge extends AutomationRuleEvaluator<JsonNode> {
-
-        @NotNull @JsonView({View.Public.class, View.Write.class})
-        @Schema(accessMode = Schema.AccessMode.READ_WRITE)
-        JsonNode code;
-
-        @ConstructorProperties({"id", "projectId", "name", "samplingRate", "code", "createdAt", "createdBy",
-                "lastUpdatedAt", "lastUpdatedBy"})
-        public AutomationRuleEvaluatorLlmAsJudge(UUID id, UUID projectId, @NotBlank String name, float samplingRate,
-                @NotNull JsonNode code,
-                Instant createdAt, String createdBy, Instant lastUpdatedAt, String lastUpdatedBy) {
-            super(id, projectId, name, samplingRate, createdAt, createdBy, lastUpdatedAt, lastUpdatedBy);
-            this.code = code;
-        }
-
-        @Override
-        public AutomationRuleEvaluatorType type() {
-            return AutomationRuleEvaluatorType.LLM_AS_JUDGE;
-        }
-    }
+public abstract sealed class AutomationRuleEvaluator<T>
+        implements
+            AutomationRule<T>
+        permits AutomationRuleEvaluatorLlmAsJudge {
 
     @JsonView({View.Public.class})
     @Schema(accessMode = Schema.AccessMode.READ_ONLY)

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

@@ -0,0 +1,19 @@
+package com.comet.opik.api;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import lombok.Builder;
+
+import java.util.Set;
+import java.util.UUID;
+
+@Builder(toBuilder = true)
+@JsonIgnoreProperties(ignoreUnknown = true)
+public record AutomationRuleEvaluatorCriteria(
+        AutomationRuleEvaluatorType type,
+        String name,
+        Set<UUID> ids) {
+
+    public AutomationRule.AutomationRuleAction action() {
+        return AutomationRule.AutomationRuleAction.EVALUATOR;
+    }
+}

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

@@ -0,0 +1,81 @@
+package com.comet.opik.api;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonView;
+import dev.langchain4j.data.message.ChatMessageType;
+import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.validation.constraints.NotBlank;
+import jakarta.validation.constraints.NotNull;
+import lombok.Builder;
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.ToString;
+import lombok.experimental.SuperBuilder;
+
+import java.beans.ConstructorProperties;
+import java.time.Instant;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+@EqualsAndHashCode(callSuper = true)
+@Data
+@SuperBuilder(toBuilder = true)
+@ToString(callSuper = true)
+public final class AutomationRuleEvaluatorLlmAsJudge
+        extends
+            AutomationRuleEvaluator<AutomationRuleEvaluatorLlmAsJudge.LlmAsJudgeCode> {
+
+    @NotNull @JsonView({View.Public.class, View.Write.class})
+    @Schema(accessMode = Schema.AccessMode.READ_WRITE)
+    private LlmAsJudgeCode code;
+
+    @Builder(toBuilder = true)
+    @JsonIgnoreProperties(ignoreUnknown = true)
+    public record LlmAsJudgeCode(
+            @JsonView( {
+                    View.Public.class, View.Write.class}) @NotNull LlmAsJudgeModelParameters model,
+            @JsonView({View.Public.class, View.Write.class}) @NotNull List<LlmAsJudgeMessage> messages,
+            @JsonView({View.Public.class, View.Write.class}) @NotNull Map<String, String> variables,
+            @JsonView({View.Public.class, View.Write.class}) @NotNull List<LlmAsJudgeOutputSchema> schema){
+    }
+
+    @Builder(toBuilder = true)
+    @JsonIgnoreProperties(ignoreUnknown = true)
+    public record LlmAsJudgeMessage(
+            @JsonView( {
+                    View.Public.class, View.Write.class}) @NotNull ChatMessageType role,
+            @JsonView({View.Public.class, View.Write.class}) @NotNull String content){
+    }
+
+    @Builder(toBuilder = true)
+    @JsonIgnoreProperties(ignoreUnknown = true)
+    public record LlmAsJudgeOutputSchema(
+            @JsonView( {
+                    View.Public.class, View.Write.class}) @NotNull String name,
+            @JsonView({View.Public.class, View.Write.class}) @NotNull LlmAsJudgeOutputSchemaType type,
+            @JsonView({View.Public.class, View.Write.class}) @NotNull String description){
+    }
+
+    @Builder(toBuilder = true)
+    @JsonIgnoreProperties(ignoreUnknown = true)
+    public record LlmAsJudgeModelParameters(
+            @JsonView( {
+                    View.Public.class, View.Write.class}) @NotNull String name,
+            @JsonView({View.Public.class, View.Write.class}) @NotNull Double temperature){
+    }
+
+    @ConstructorProperties({"id", "projectId", "name", "samplingRate", "code", "createdAt", "createdBy",
+            "lastUpdatedAt", "lastUpdatedBy"})
+    public AutomationRuleEvaluatorLlmAsJudge(UUID id, UUID projectId, @NotBlank String name, Float samplingRate,
+            @NotNull LlmAsJudgeCode code,
+            Instant createdAt, String createdBy, Instant lastUpdatedAt, String lastUpdatedBy) {
+        super(id, projectId, name, samplingRate, createdAt, createdBy, lastUpdatedAt, lastUpdatedBy);
+        this.code = code;
+    }
+
+    @Override
+    public AutomationRuleEvaluatorType type() {
+        return AutomationRuleEvaluatorType.LLM_AS_JUDGE;
+    }
+}

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

@@ -1,7 +1,6 @@
 package com.comet.opik.api;
 
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
-import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.PropertyNamingStrategies;
 import com.fasterxml.jackson.databind.annotation.JsonNaming;
 import jakarta.validation.constraints.NotNull;
@@ -12,6 +11,6 @@
 @JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
 public record AutomationRuleEvaluatorUpdate(
         @NotNull String name,
-        @NotNull JsonNode code,
+        @NotNull AutomationRuleEvaluatorLlmAsJudge.LlmAsJudgeCode code,
         @NotNull Float samplingRate) {
 }

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

@@ -0,0 +1,7 @@
+package com.comet.opik.api;
+
+public enum LlmAsJudgeOutputSchemaType {
+    BOOLEAN,
+    INTEGER,
+    DOUBLE
+}

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

@@ -1,20 +1,29 @@
 package com.comet.opik.api.events;
 
+import com.comet.opik.api.Trace;
 import com.comet.opik.infrastructure.events.BaseEvent;
 import lombok.Getter;
 import lombok.NonNull;
 import lombok.experimental.Accessors;
 
+import java.util.List;
 import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
 
 @Getter
 @Accessors(fluent = true)
 public class TracesCreated extends BaseEvent {
-    private final @NonNull Set<UUID> projectIds;
+    private final @NonNull List<Trace> traces;
 
-    public TracesCreated(@NonNull Set<UUID> projectIds, @NonNull String workspaceId, @NonNull String userName) {
+    public TracesCreated(@NonNull List<Trace> traces, @NonNull String workspaceId, @NonNull String userName) {
         super(workspaceId, userName);
-        this.projectIds = projectIds;
+        this.traces = traces;
+    }
+
+    public Set<UUID> projectIds() {
+        return traces.stream()
+                .map(Trace::projectId)
+                .collect(Collectors.toSet());
     }
 }

apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/events/LlmAsJudgeMessageRender.java

@@ -0,0 +1,136 @@
+package com.comet.opik.api.resources.v1.events;
+
+import com.comet.opik.api.AutomationRuleEvaluatorLlmAsJudge;
+import com.comet.opik.api.Trace;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.jayway.jsonpath.JsonPath;
+import dev.ai4j.openai4j.chat.Message;
+import dev.ai4j.openai4j.chat.SystemMessage;
+import dev.ai4j.openai4j.chat.UserMessage;
+import lombok.Builder;
+import lombok.experimental.UtilityClass;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.text.StringSubstitutor;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+@UtilityClass
+@Slf4j
+class LlmAsJudgeMessageRender {
+
+    /**
+     * Render the rule evaluator message template using the values from an actual trace.
+     *
+     * As the rule my consist in multiple messages, we check each one of them for variables to fill.
+     * Then we go through every variable template to replace them for the value from the trace.
+     *
+     * @param trace the trace with value to use to replace template variables
+     * @param evaluatorCode the evaluator
+     * @return a list of AI messages, with templates rendered
+     */
+    public static List<Message> renderMessages(Trace trace,
+            AutomationRuleEvaluatorLlmAsJudge.LlmAsJudgeCode evaluatorCode) {
+        // prepare the map of replacements to use in all messages
+        var parsedVariables = variableMapping(evaluatorCode.variables());
+
+        // extract the actual value from the Trace
+        var replacements = parsedVariables.stream().map(mapper -> {
+            var traceSection = switch (mapper.traceSection) {
+                case INPUT -> trace.input();
+                case OUTPUT -> trace.output();
+                case METADATA -> trace.metadata();
+            };
+
+            return mapper.toBuilder()
+                    .valueToReplace(extractFromJson(traceSection, mapper.jsonPath()))
+                    .build();
+        })
+                .filter(mapper -> mapper.valueToReplace() != null)
+                .collect(
+                        Collectors.toMap(LlmAsJudgeMessageRender.MessageVariableMapping::variableName,
+                                LlmAsJudgeMessageRender.MessageVariableMapping::valueToReplace));
+
+        // will convert all '{{key}}' into 'value'
+        // TODO: replace with Mustache Java to be in confirm with FE
+        var templateRenderer = new StringSubstitutor(replacements, "{{", "}}");
+
+        // render the message templates from evaluator rule
+        return evaluatorCode.messages().stream()
+                .map(templateMessage -> {
+                    var renderedMessage = templateRenderer.replace(templateMessage.content());
+
+                    return switch (templateMessage.role()) {
+                        case USER -> UserMessage.from(renderedMessage);
+                        case SYSTEM -> SystemMessage.from(renderedMessage);
+                        default -> {
+                            log.info("No mapping for message role type {}", templateMessage.role());
+                            yield null;
+                        }
+                    };
+                })
+                .filter(Objects::nonNull)
+                .toList();
+    }
+
+    /**
+     * Parse evaluator\'s variable mapper into an usable list of
+     *
+     * @param evaluatorVariables a map with variables and a path into a trace input/output/metadata to replace
+     * @return a parsed list of mappings, easier to use for the template rendering
+     */
+    public static List<MessageVariableMapping> variableMapping(Map<String, String> evaluatorVariables) {
+        return evaluatorVariables.entrySet().stream()
+                .map(mapper -> {
+                    var templateVariable = mapper.getKey();
+                    var tracePath = mapper.getValue();
+
+                    var builder = MessageVariableMapping.builder().variableName(templateVariable);
+
+                    if (tracePath.startsWith("input.")) {
+                        builder.traceSection(TraceSection.INPUT)
+                                .jsonPath("$" + tracePath.substring("input".length()));
+                    } else if (tracePath.startsWith("output.")) {
+                        builder.traceSection(TraceSection.OUTPUT)
+                                .jsonPath("$" + tracePath.substring("output".length()));
+                    } else if (tracePath.startsWith("metadata.")) {
+                        builder.traceSection(TraceSection.METADATA)
+                                .jsonPath("$" + tracePath.substring("metadata".length()));
+                    } else {
+                        log.info("Couldn't map trace path '{}' into a input/output/metadata path", tracePath);
+                        return null;
+                    }
+
+                    return builder.build();
+                })
+                .filter(Objects::nonNull)
+                .toList();
+    }
+
+    final ObjectMapper objectMapper = new ObjectMapper();
+
+    String extractFromJson(JsonNode json, String path) {
+        try {
+            // JsonPath didnt work with JsonNode, even explicitly using JacksonJsonProvider, so we convert to a Map
+            var forcedObject = objectMapper.convertValue(json, Map.class);
+            return JsonPath.parse(forcedObject).read(path);
+        } catch (Exception e) {
+            log.debug("Couldn't find path '{}' inside json {}: {}", path, json, e.getMessage());
+            return null;
+        }
+    }
+
+    public enum TraceSection {
+        INPUT,
+        OUTPUT,
+        METADATA
+    }
+
+    @Builder(toBuilder = true)
+    public record MessageVariableMapping(TraceSection traceSection, String variableName, String jsonPath,
+            String valueToReplace) {
+    }
+}