From 7f5175de24525ebb4a7dfcae90c7b206ac04f07d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Cre=C3=A3o?= Date: Fri, 3 Jan 2025 11:01:08 +0000 Subject: [PATCH] [Online Scoring] Automation Rules endpoints (#945) * ## Details Data model, DAO, service layer and endpoints for Automation Rule Evaluators. We're using a polymorphic type for the evaluator code, which at least for now for LLM-as-Judge, we will use a JsonNode type ## Issues OPIK-590 OPIK-591 --- .../com/comet/opik/api/AutomationRule.java | 52 ++ .../opik/api/AutomationRuleEvaluator.java | 119 +++ .../opik/api/AutomationRuleEvaluatorType.java | 24 + .../api/AutomationRuleEvaluatorUpdate.java | 17 + .../AutomationRuleEvaluatorsResource.java | 161 ++++ .../AutomationModelEvaluatorMapper.java | 18 + .../comet/opik/domain/AutomationRuleDAO.java | 65 ++ .../domain/AutomationRuleEvaluatorDAO.java | 77 ++ .../domain/AutomationRuleEvaluatorModel.java | 17 + .../AutomationRuleEvaluatorRowMapper.java | 24 + .../AutomationRuleEvaluatorService.java | 196 +++++ .../opik/domain/AutomationRuleModel.java | 20 + .../opik/domain/AutomationRuleRowMapper.java | 25 + ...lmAsJudgeAutomationRuleEvaluatorModel.java | 30 + .../db/JsonNodeArgumentFactory.java | 36 + .../000009_add_automation_rule_tables.sql | 28 + .../AutomationRuleEvaluatorsResourceTest.java | 717 ++++++++++++++++++ 17 files changed, 1626 insertions(+) create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRule.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluator.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorType.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorUpdate.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResource.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationModelEvaluatorMapper.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleDAO.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorDAO.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorModel.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorRowMapper.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorService.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleModel.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleRowMapper.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/domain/LlmAsJudgeAutomationRuleEvaluatorModel.java create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/JsonNodeArgumentFactory.java create mode 100644 apps/opik-backend/src/main/resources/liquibase/db-app-state/migrations/000009_add_automation_rule_tables.sql create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRule.java b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRule.java new file mode 100644 index 0000000000..1cb3b5b052 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRule.java @@ -0,0 +1,52 @@ +package com.comet.opik.api; + +import com.fasterxml.jackson.annotation.JsonSubTypes; +import com.fasterxml.jackson.annotation.JsonTypeInfo; +import com.fasterxml.jackson.annotation.JsonValue; +import io.swagger.v3.oas.annotations.media.DiscriminatorMapping; +import io.swagger.v3.oas.annotations.media.Schema; +import lombok.AccessLevel; +import lombok.Getter; +import lombok.RequiredArgsConstructor; + +import java.time.Instant; +import java.util.Arrays; +import java.util.UUID; + +@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "action", visible = true) +@JsonSubTypes({ + @JsonSubTypes.Type(value = AutomationRuleEvaluator.class, name = "evaluator") +}) +@Schema(name = "AutomationRule", discriminatorProperty = "action", discriminatorMapping = { + @DiscriminatorMapping(value = "evaluator", schema = AutomationRuleEvaluator.class) +}) +public sealed interface AutomationRule permits AutomationRuleEvaluator { + + UUID getId(); + UUID getProjectId(); + String getName(); + + AutomationRuleAction getAction(); + Float getSamplingRate(); + + Instant getCreatedAt(); + String getCreatedBy(); + Instant getLastUpdatedAt(); + String getLastUpdatedBy(); + + @Getter + @RequiredArgsConstructor(access = AccessLevel.PRIVATE) + enum AutomationRuleAction { + + EVALUATOR("evaluator"); + + @JsonValue + private final String action; + + public static AutomationRule.AutomationRuleAction fromString(String action) { + return Arrays.stream(values()) + .filter(v -> v.action.equals(action)).findFirst() + .orElseThrow(() -> new IllegalArgumentException("Unknown rule type: " + action)); + } + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluator.java b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluator.java new file mode 100644 index 0000000000..5f3f90ed09 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluator.java @@ -0,0 +1,119 @@ +package com.comet.opik.api; + +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; + +@Data +@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") +}) +@Schema(name = "AutomationRuleEvaluator", discriminatorProperty = "type", discriminatorMapping = { + @DiscriminatorMapping(value = "llm_as_judge", schema = AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class) +}) +@AllArgsConstructor +public abstract sealed class AutomationRuleEvaluator implements AutomationRule { + + @EqualsAndHashCode(callSuper = true) + @Data + @SuperBuilder(toBuilder = true) + @ToString(callSuper = true) + public static final class AutomationRuleEvaluatorLlmAsJudge extends AutomationRuleEvaluator { + + @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; + } + } + + @JsonView({View.Public.class}) + @Schema(accessMode = Schema.AccessMode.READ_ONLY) + UUID id; + + @JsonView({View.Public.class, View.Write.class}) + @NotNull + UUID projectId; + + @JsonView({View.Public.class, View.Write.class}) + @Schema(accessMode = Schema.AccessMode.READ_WRITE) + @NotBlank + String name; + + @JsonView({View.Public.class, View.Write.class}) + @Schema(accessMode = Schema.AccessMode.READ_WRITE) + Float samplingRate; + + @JsonView({View.Public.class}) + @Schema(accessMode = Schema.AccessMode.READ_ONLY) + Instant createdAt; + + @JsonView({View.Public.class}) + @Schema(accessMode = Schema.AccessMode.READ_ONLY) + String createdBy; + + @JsonView({View.Public.class}) + @Schema(accessMode = Schema.AccessMode.READ_ONLY) + Instant lastUpdatedAt; + + @JsonView({View.Public.class}) + @Schema(accessMode = Schema.AccessMode.READ_ONLY) + String lastUpdatedBy; + + @JsonView({View.Public.class}) + public abstract AutomationRuleEvaluatorType type(); + + @JsonView({View.Public.class, View.Write.class}) + public abstract T getCode(); + + @Override + public AutomationRuleAction getAction() { + return AutomationRuleAction.EVALUATOR; + } + + public static class View { + public static class Write {} + public static class Public {} + } + + @Builder(toBuilder = true) + public record AutomationRuleEvaluatorPage( + @JsonView({View.Public.class}) int page, + @JsonView({View.Public.class}) int size, + @JsonView({View.Public.class}) long total, + @JsonView({View.Public.class}) List content) + implements Page{ + + public static AutomationRuleEvaluator.AutomationRuleEvaluatorPage empty(int page) { + return new AutomationRuleEvaluator.AutomationRuleEvaluatorPage(page, 0, 0, List.of()); + } + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorType.java b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorType.java new file mode 100644 index 0000000000..73fe68bc70 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorType.java @@ -0,0 +1,24 @@ +package com.comet.opik.api; + +import com.fasterxml.jackson.annotation.JsonValue; +import lombok.AccessLevel; +import lombok.Getter; +import lombok.RequiredArgsConstructor; + +import java.util.Arrays; + +@Getter +@RequiredArgsConstructor(access = AccessLevel.PRIVATE) +public enum AutomationRuleEvaluatorType { + + LLM_AS_JUDGE("llm_as_judge"); + + @JsonValue + private final String type; + + public static AutomationRuleEvaluatorType fromString(String type) { + return Arrays.stream(values()) + .filter(v -> v.type.equals(type)).findFirst() + .orElseThrow(() -> new IllegalArgumentException("Unknown evaluator type: " + type)); + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorUpdate.java b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorUpdate.java new file mode 100644 index 0000000000..b81e16d08b --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/AutomationRuleEvaluatorUpdate.java @@ -0,0 +1,17 @@ +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; +import lombok.Builder; + +@Builder(toBuilder = true) +@JsonIgnoreProperties(ignoreUnknown = true) +@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class) +public record AutomationRuleEvaluatorUpdate( + @NotNull String name, + @NotNull JsonNode code, + @NotNull Float samplingRate) { +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResource.java new file mode 100644 index 0000000000..b2a5e81493 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResource.java @@ -0,0 +1,161 @@ +package com.comet.opik.api.resources.v1.priv; + +import com.codahale.metrics.annotation.Timed; +import com.comet.opik.api.AutomationRuleEvaluator; +import com.comet.opik.api.AutomationRuleEvaluatorUpdate; +import com.comet.opik.api.BatchDelete; +import com.comet.opik.api.Page; +import com.comet.opik.domain.AutomationRuleEvaluatorService; +import com.comet.opik.infrastructure.auth.RequestContext; +import com.comet.opik.infrastructure.ratelimit.RateLimited; +import com.fasterxml.jackson.annotation.JsonView; +import io.swagger.v3.oas.annotations.Operation; +import io.swagger.v3.oas.annotations.headers.Header; +import io.swagger.v3.oas.annotations.media.Content; +import io.swagger.v3.oas.annotations.media.Schema; +import io.swagger.v3.oas.annotations.parameters.RequestBody; +import io.swagger.v3.oas.annotations.responses.ApiResponse; +import io.swagger.v3.oas.annotations.tags.Tag; +import jakarta.inject.Inject; +import jakarta.inject.Provider; +import jakarta.validation.Valid; +import jakarta.validation.constraints.Min; +import jakarta.validation.constraints.NotNull; +import jakarta.ws.rs.Consumes; +import jakarta.ws.rs.DefaultValue; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.PATCH; +import jakarta.ws.rs.POST; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.QueryParam; +import jakarta.ws.rs.core.Context; +import jakarta.ws.rs.core.MediaType; +import jakarta.ws.rs.core.Response; +import jakarta.ws.rs.core.UriInfo; +import lombok.NonNull; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +import java.net.URI; +import java.util.UUID; + +@Path("/v1/private/automations/projects/{projectId}/evaluators/") +@Produces(MediaType.APPLICATION_JSON) +@Consumes(MediaType.APPLICATION_JSON) +@Timed +@Slf4j +@RequiredArgsConstructor(onConstructor_ = @Inject) +@Tag(name = "Automation rule evaluators", description = "Automation rule evaluators resource") +public class AutomationRuleEvaluatorsResource { + + private final @NonNull AutomationRuleEvaluatorService service; + private final @NonNull Provider requestContext; + + @GET + @Operation(operationId = "findEvaluators", summary = "Find project Evaluators", description = "Find project Evaluators", responses = { + @ApiResponse(responseCode = "200", description = "Evaluators resource", content = @Content(schema = @Schema(implementation = AutomationRuleEvaluator.AutomationRuleEvaluatorPage.class))) + }) + @JsonView(AutomationRuleEvaluator.View.Public.class) + public Response find(@PathParam("projectId") UUID projectId, + @QueryParam("name") String name, + @QueryParam("page") @Min(1) @DefaultValue("1") int page, + @QueryParam("size") @Min(1) @DefaultValue("10") int size) { + + String workspaceId = requestContext.get().getWorkspaceId(); + log.info("Looking for automated evaluators for project id '{}' on workspaceId '{}' (page {})", projectId, + workspaceId, page); + Page definitionPage = service.find(projectId, workspaceId, name, page, size); + log.info("Found {} automated evaluators for project id '{}' on workspaceId '{}' (page {}, total {})", + definitionPage.size(), projectId, workspaceId, page, definitionPage.total()); + + return Response.ok() + .entity(definitionPage) + .build(); + } + + @GET + @Path("/{id}") + @Operation(operationId = "getEvaluatorById", summary = "Get automation rule evaluator by id", description = "Get automation rule by id", responses = { + @ApiResponse(responseCode = "200", description = "Automation Rule resource", content = @Content(schema = @Schema(implementation = AutomationRuleEvaluator.class))) + }) + @JsonView(AutomationRuleEvaluator.View.Public.class) + public Response getEvaluator(@PathParam("projectId") UUID projectId, @PathParam("id") UUID evaluatorId) { + String workspaceId = requestContext.get().getWorkspaceId(); + + log.info("Looking for automated evaluator: id '{}' on project_id '{}'", projectId, workspaceId); + AutomationRuleEvaluator evaluator = service.findById(evaluatorId, projectId, workspaceId); + log.info("Found automated evaluator: id '{}' on project_id '{}'", projectId, workspaceId); + + return Response.ok().entity(evaluator).build(); + } + + @POST + @Operation(operationId = "createAutomationRuleEvaluator", summary = "Create automation rule evaluator", description = "Create automation rule evaluator", responses = { + @ApiResponse(responseCode = "201", description = "Created", headers = { + @Header(name = "Location", required = true, example = "${basePath}/v1/private/automations/projects/{projectId}/evaluators/{evaluatorId}", schema = @Schema(implementation = String.class)) + }) + }) + @RateLimited + public Response createEvaluator( + @RequestBody(content = @Content(schema = @Schema(implementation = AutomationRuleEvaluator.class))) + @JsonView(AutomationRuleEvaluator.View.Write.class) @NotNull @Valid AutomationRuleEvaluator evaluator, + @Context UriInfo uriInfo) { + + String workspaceId = requestContext.get().getWorkspaceId(); + String userName = requestContext.get().getUserName(); + + log.info("Creating {} evaluator for project_id '{}' on workspace_id '{}'", evaluator.type(), + evaluator.getProjectId(), workspaceId); + AutomationRuleEvaluator savedEvaluator = service.save(evaluator, workspaceId, userName); + log.info("Created {} evaluator '{}' for project_id '{}' on workspace_id '{}'", evaluator.type(), + savedEvaluator.getId(), evaluator.getProjectId(), workspaceId); + + URI uri = uriInfo.getBaseUriBuilder() + .path("v1/private/automations/projects/{projectId}/evaluators/{id}") + .resolveTemplate("projectId", savedEvaluator.getProjectId().toString()) + .resolveTemplate("id", savedEvaluator.getId().toString()) + .build(); + return Response.created(uri).build(); + } + + @PATCH + @Path("/{id}") + @Operation(operationId = "updateAutomationRuleEvaluator", summary = "update Automation Rule Evaluator by id", description = "update Automation Rule Evaluator by id", responses = { + @ApiResponse(responseCode = "204", description = "No content"), + }) + @RateLimited + public Response updateEvaluator(@PathParam("id") UUID id, + @PathParam("projectId") UUID projectId, + @RequestBody(content = @Content(schema = @Schema(implementation = AutomationRuleEvaluatorUpdate.class))) @NotNull @Valid AutomationRuleEvaluatorUpdate evaluatorUpdate) { + + String workspaceId = requestContext.get().getWorkspaceId(); + String userName = requestContext.get().getUserName(); + + log.info("Updating automation rule evaluator by id '{}' and project_id '{}' on workspace_id '{}'", id, + projectId, workspaceId); + service.update(id, projectId, workspaceId, userName, evaluatorUpdate); + log.info("Updated automation rule evaluator by id '{}' and project_id '{}' on workspace_id '{}'", id, projectId, + workspaceId); + + return Response.noContent().build(); + } + + @POST + @Path("/delete") + @Operation(operationId = "deleteAutomationRuleEvaluatorBatch", summary = "Delete automation rule evaluators", description = "Delete automation rule evaluators batch", responses = { + @ApiResponse(responseCode = "204", description = "No Content"), + }) + public Response deleteEvaluators( + @NotNull @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @Valid BatchDelete batchDelete, @PathParam("projectId") UUID projectId) { + String workspaceId = requestContext.get().getWorkspaceId(); + log.info("Deleting automation rule evaluators by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(), + workspaceId); + service.delete(batchDelete.ids(), projectId, workspaceId); + log.info("Deleted automation rule evaluators by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(), + workspaceId); + return Response.noContent().build(); + } + +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationModelEvaluatorMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationModelEvaluatorMapper.java new file mode 100644 index 0000000000..a7ef7ebcaf --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationModelEvaluatorMapper.java @@ -0,0 +1,18 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRuleEvaluator; +import org.mapstruct.Mapper; +import org.mapstruct.factory.Mappers; + +import java.time.Instant; + +@Mapper(imports = Instant.class) +interface AutomationModelEvaluatorMapper { + + AutomationModelEvaluatorMapper INSTANCE = Mappers.getMapper(AutomationModelEvaluatorMapper.class); + + AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge map(LlmAsJudgeAutomationRuleEvaluatorModel model); + + LlmAsJudgeAutomationRuleEvaluatorModel map(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge dto); + +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleDAO.java new file mode 100644 index 0000000000..9b0dc2c8eb --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleDAO.java @@ -0,0 +1,65 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRule; +import com.comet.opik.api.AutomationRuleEvaluator; +import com.comet.opik.infrastructure.db.UUIDArgumentFactory; +import org.jdbi.v3.sqlobject.config.RegisterArgumentFactory; +import org.jdbi.v3.sqlobject.config.RegisterConstructorMapper; +import org.jdbi.v3.sqlobject.config.RegisterRowMapper; +import org.jdbi.v3.sqlobject.customizer.AllowUnusedBindings; +import org.jdbi.v3.sqlobject.customizer.Bind; +import org.jdbi.v3.sqlobject.customizer.BindList; +import org.jdbi.v3.sqlobject.customizer.BindMethods; +import org.jdbi.v3.sqlobject.customizer.Define; +import org.jdbi.v3.sqlobject.statement.SqlQuery; +import org.jdbi.v3.sqlobject.statement.SqlUpdate; +import org.jdbi.v3.stringtemplate4.UseStringTemplateEngine; + +import java.util.Set; +import java.util.UUID; + +@RegisterArgumentFactory(UUIDArgumentFactory.class) +@RegisterRowMapper(AutomationRuleRowMapper.class) +@RegisterConstructorMapper(AutomationRuleEvaluator.class) +interface AutomationRuleDAO { + + @SqlUpdate("INSERT INTO automation_rules(id, project_id, workspace_id, `action`, name, sampling_rate) "+ + "VALUES (:rule.id, :rule.projectId, :workspaceId, :rule.action, :rule.name, :rule.samplingRate)") + void saveBaseRule(@BindMethods("rule") AutomationRuleModel rule, @Bind("workspaceId") String workspaceId); + + @SqlUpdate(""" + UPDATE automation_rules + SET name = :name, + sampling_rate = :samplingRate + WHERE id = :id AND project_id = :projectId AND workspace_id = :workspaceId + """) + int updateBaseRule(@Bind("id") UUID id, + @Bind("projectId") UUID projectId, + @Bind("workspaceId") String workspaceId, + @Bind("name") String name, + @Bind("samplingRate") float samplingRate, + @Bind("lastUpdatedBy") String lastUpdatedBy); + + @SqlUpdate(""" + DELETE FROM automation_rules + WHERE project_id = :projectId AND workspace_id = :workspaceId + AND id IN () + """) + @UseStringTemplateEngine + @AllowUnusedBindings + void deleteBaseRules(@Define("ids") @BindList(onEmpty = BindList.EmptyHandling.NULL_VALUE, value = "ids") Set ids, + @Bind("projectId") UUID projectId, + @Bind("workspaceId") String workspaceId); + + @SqlQuery(""" + SELECT COUNT(*) + FROM automation_rules + WHERE project_id = :projectId AND workspace_id = :workspaceId + AND `action` = :action + """) + @UseStringTemplateEngine + @AllowUnusedBindings + long findCount(@Bind("projectId") UUID projectId, + @Bind("workspaceId") String workspaceId, + @Define("action") @Bind("action") AutomationRule.AutomationRuleAction action); +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorDAO.java new file mode 100644 index 0000000000..3cc439fb6b --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorDAO.java @@ -0,0 +1,77 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRule; +import com.comet.opik.api.AutomationRuleEvaluatorUpdate; +import com.comet.opik.infrastructure.db.JsonNodeArgumentFactory; +import com.comet.opik.infrastructure.db.UUIDArgumentFactory; +import org.jdbi.v3.sqlobject.config.RegisterArgumentFactory; +import org.jdbi.v3.sqlobject.config.RegisterConstructorMapper; +import org.jdbi.v3.sqlobject.config.RegisterRowMapper; +import org.jdbi.v3.sqlobject.customizer.AllowUnusedBindings; +import org.jdbi.v3.sqlobject.customizer.Bind; +import org.jdbi.v3.sqlobject.customizer.BindList; +import org.jdbi.v3.sqlobject.customizer.BindMethods; +import org.jdbi.v3.sqlobject.customizer.Define; +import org.jdbi.v3.sqlobject.statement.SqlQuery; +import org.jdbi.v3.sqlobject.statement.SqlUpdate; +import org.jdbi.v3.stringtemplate4.UseStringTemplateEngine; + +import java.util.List; +import java.util.Set; +import java.util.UUID; + +@RegisterArgumentFactory(UUIDArgumentFactory.class) +@RegisterArgumentFactory(JsonNodeArgumentFactory.class) +@RegisterConstructorMapper(LlmAsJudgeAutomationRuleEvaluatorModel.class) +@RegisterRowMapper(AutomationRuleEvaluatorRowMapper.class) +public interface AutomationRuleEvaluatorDAO extends AutomationRuleDAO { + + @SqlUpdate("INSERT INTO automation_rule_evaluators(id, `type`, code, created_by, last_updated_by) "+ + "VALUES (:rule.id, :rule.type, :rule.code, :rule.createdBy, :rule.lastUpdatedBy)") + void saveEvaluator(@BindMethods("rule") AutomationRuleEvaluatorModel rule); + + @SqlUpdate(""" + UPDATE automation_rule_evaluators + SET code = :rule.code, + last_updated_by = :userName + WHERE id = :id + """) + int updateEvaluator(@Bind("id") UUID id, @BindMethods("rule") AutomationRuleEvaluatorUpdate ruleUpdate, @Bind("userName") String userName); + + @SqlQuery(""" + SELECT rule.id, rule.project_id, rule.action, rule.name, rule.sampling_rate, evaluator.type, evaluator.code, + evaluator.created_at, evaluator.created_by, evaluator.last_updated_at, evaluator.last_updated_by + FROM automation_rules rule + JOIN automation_rule_evaluators evaluator + ON rule.id = evaluator.id + WHERE workspace_id = :workspaceId AND project_id = :projectId + AND `action` = :action + AND rule.id IN () + AND name like concat('%', :name, '%') + LIMIT :limit OFFSET :offset + """) + @UseStringTemplateEngine + @AllowUnusedBindings + List> find(@Bind("workspaceId") String workspaceId, + @Bind("projectId") UUID projectId, + @Define("ids") @BindList(onEmpty = BindList.EmptyHandling.NULL_VALUE, value = "ids") Set ids, + @Define("name") @Bind("name") String name, + @Bind("action") AutomationRule.AutomationRuleAction action, + @Bind("offset") int offset, + @Bind("limit") int limit); + + @SqlUpdate(""" + DELETE FROM automation_rule_evaluators + WHERE id IN ( + SELECT id + FROM automation_rules + WHERE workspace_id = :workspaceId AND project_id = :projectId + AND id IN () + ) + """) + @UseStringTemplateEngine + @AllowUnusedBindings + void deleteEvaluatorsByIds(@Bind("workspaceId") String workspaceId, + @Bind("projectId") UUID projectId, + @Define("ids") @BindList(onEmpty = BindList.EmptyHandling.NULL_VALUE, value = "ids") Set ids); +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorModel.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorModel.java new file mode 100644 index 0000000000..c8bcfc56c4 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorModel.java @@ -0,0 +1,17 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRule; +import com.comet.opik.api.AutomationRuleEvaluatorType; +import org.jdbi.v3.json.Json; + +public sealed interface AutomationRuleEvaluatorModel extends AutomationRuleModel + permits LlmAsJudgeAutomationRuleEvaluatorModel { + + @Json T code(); + AutomationRuleEvaluatorType type(); + + @Override + default AutomationRule.AutomationRuleAction action() { + return AutomationRule.AutomationRuleAction.EVALUATOR; + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorRowMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorRowMapper.java new file mode 100644 index 0000000000..0a7890ff77 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorRowMapper.java @@ -0,0 +1,24 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRuleEvaluatorType; +import org.jdbi.v3.core.mapper.RowMapper; +import org.jdbi.v3.core.statement.StatementContext; + +import java.sql.ResultSet; +import java.sql.SQLException; + +public class AutomationRuleEvaluatorRowMapper implements RowMapper> { + + @Override + public AutomationRuleEvaluatorModel map(ResultSet rs, StatementContext ctx) throws SQLException { + + var type = AutomationRuleEvaluatorType.fromString(rs.getString("type")); + + return switch (type) { + case LLM_AS_JUDGE -> ctx.findMapperFor(LlmAsJudgeAutomationRuleEvaluatorModel.class) + .orElseThrow(() -> new IllegalStateException( + "No mapper found for Automation Rule Evaluator type: %s".formatted(type))) + .map(rs, ctx); + }; + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorService.java new file mode 100644 index 0000000000..613f04c8d3 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleEvaluatorService.java @@ -0,0 +1,196 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRule; +import com.comet.opik.api.AutomationRuleEvaluator; +import com.comet.opik.api.AutomationRuleEvaluatorUpdate; +import com.comet.opik.api.error.EntityAlreadyExistsException; +import com.comet.opik.api.error.ErrorMessage; +import com.google.inject.ImplementedBy; +import jakarta.inject.Inject; +import jakarta.inject.Singleton; +import jakarta.ws.rs.NotFoundException; +import jakarta.ws.rs.core.Response; +import lombok.NonNull; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.jdbi.v3.core.statement.UnableToExecuteStatementException; +import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate; + +import java.sql.SQLIntegrityConstraintViolationException; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.UUID; + +import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY; +import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.WRITE; + +@ImplementedBy(AutomationRuleEvaluatorServiceImpl.class) +public interface AutomationRuleEvaluatorService { + + > T save(T automationRuleEvaluator, @NonNull String workspaceId, @NonNull String userName); + + void update(@NonNull UUID id, @NonNull UUID projectId, @NonNull String workspaceId, @NonNull String userName, + AutomationRuleEvaluatorUpdate automationRuleEvaluator); + + > T findById(@NonNull UUID id, @NonNull UUID projectId, @NonNull String workspaceId); + + void delete(@NonNull Set ids, @NonNull UUID projectId, @NonNull String workspaceId); + + AutomationRuleEvaluator.AutomationRuleEvaluatorPage find(@NonNull UUID projectId, @NonNull String workspaceId, String name, int page, int size); +}@NonNull + +@Singleton +@RequiredArgsConstructor(onConstructor_ = @Inject) +@Slf4j +class AutomationRuleEvaluatorServiceImpl implements AutomationRuleEvaluatorService { + + private static final String EVALUATOR_ALREADY_EXISTS = "AutomationRuleEvaluator already exists"; + + private final @NonNull IdGenerator idGenerator; + private final @NonNull TransactionTemplate template; + private final int DEFAULT_PAGE_LIMIT = 10; + + @Override + public > T save(T inputRuleEvaluator, + @NonNull String workspaceId, + @NonNull String userName) { + + UUID id = idGenerator.generateId(); + IdGenerator.validateVersion(id, "AutomationRuleEvaluator"); + + return template.inTransaction(WRITE, handle -> { + var evaluatorsDAO = handle.attach(AutomationRuleEvaluatorDAO.class); + + AutomationRuleEvaluatorModel evaluator = switch (inputRuleEvaluator) { + case AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge llmAsJudge -> { + var definition = llmAsJudge.toBuilder() + .id(id) + .createdBy(userName) + .lastUpdatedBy(userName) + .build(); + + yield AutomationModelEvaluatorMapper.INSTANCE.map(definition); + } + + }; + + try { + log.debug("Creating {} AutomationRuleEvaluator with id '{}' in projectId '{}' and workspaceId '{}'", + evaluator.type(), id, evaluator.projectId(), workspaceId); + + evaluatorsDAO.saveBaseRule(evaluator, workspaceId); + evaluatorsDAO.saveEvaluator(evaluator); + + return findById(evaluator.id(), evaluator.projectId(), workspaceId); + + } catch (UnableToExecuteStatementException e) { + if (e.getCause() instanceof SQLIntegrityConstraintViolationException) { + log.info(EVALUATOR_ALREADY_EXISTS, e); + throw new EntityAlreadyExistsException(new ErrorMessage(List.of(EVALUATOR_ALREADY_EXISTS))); + } else { + throw e; + } + } + }); + } + + + @Override + public void update(@NonNull UUID id, @NonNull UUID projectId, @NonNull String workspaceId, + @NonNull String userName, @NonNull AutomationRuleEvaluatorUpdate evaluatorUpdate) { + + log.debug("Updating AutomationRuleEvaluator with id '{}' in projectId '{}' and workspaceId '{}'", id, projectId, workspaceId); + template.inTransaction(WRITE, handle -> { + var dao = handle.attach(AutomationRuleEvaluatorDAO.class); + + try { + int resultBase = dao.updateBaseRule(id, projectId, workspaceId, evaluatorUpdate.name(), evaluatorUpdate.samplingRate(), userName); + int resultEval = dao.updateEvaluator(id, evaluatorUpdate, userName); + + if (resultEval == 0 || resultBase == 0) { + throw newNotFoundException(); + } + } catch (UnableToExecuteStatementException e) { + if (e.getCause() instanceof SQLIntegrityConstraintViolationException) { + log.info(EVALUATOR_ALREADY_EXISTS); + throw new EntityAlreadyExistsException(new ErrorMessage(List.of(EVALUATOR_ALREADY_EXISTS))); + } else { + throw e; + } + } + + return null; + }); + } + + @Override + public > T findById(@NonNull UUID id, @NonNull UUID projectId, @NonNull String workspaceId) { + log.debug("Finding AutomationRuleEvaluator with id '{}' in projectId '{}' and workspaceId '{}'", id, projectId, workspaceId); + + return (T) template.inTransaction(READ_ONLY, handle -> { + var dao = handle.attach(AutomationRuleEvaluatorDAO.class); + var singleIdSet = Collections.singleton(id); + return dao.find(workspaceId, projectId, singleIdSet, null, AutomationRule.AutomationRuleAction.EVALUATOR, 0, DEFAULT_PAGE_LIMIT) + .stream() + .findFirst() + .map(ruleEvaluator -> switch (ruleEvaluator) { + case LlmAsJudgeAutomationRuleEvaluatorModel llmAsJudge -> AutomationModelEvaluatorMapper.INSTANCE.map(llmAsJudge); + }) + .orElseThrow(this::newNotFoundException); + }); + } + + @Override + public void delete(@NonNull Set ids, @NonNull UUID projectId, @NonNull String workspaceId) { + if (ids.isEmpty()) { + log.info("Delete AutomationRuleEvaluator: ids list is empty, returning"); + return; + } + + log.debug("Deleting AutomationRuleEvaluators with ids {} in projectId '{}' and workspaceId '{}'", ids, projectId, workspaceId); + + template.inTransaction(WRITE, handle -> { + var dao = handle.attach(AutomationRuleEvaluatorDAO.class); + dao.deleteEvaluatorsByIds(workspaceId, projectId, ids); + dao.deleteBaseRules(ids, projectId, workspaceId); + return null; + }); + } + + private NotFoundException newNotFoundException() { + String message = "AutomationRuleEvaluator not found"; + log.info(message); + return new NotFoundException(message, + Response.status(Response.Status.NOT_FOUND).entity(new ErrorMessage(List.of(message))).build()); + } + + @Override + public AutomationRuleEvaluator.AutomationRuleEvaluatorPage find(@NonNull UUID projectId, + @NonNull String workspaceId, + String name, + int pageNum, int size) { + + log.debug("Finding AutomationRuleEvaluators with name pattern '{}' in projectId '{}' and workspaceId '{}'", name, projectId, workspaceId); + + return template.inTransaction(READ_ONLY, handle -> { + var dao = handle.attach(AutomationRuleEvaluatorDAO.class); + var total = dao.findCount(projectId, workspaceId, AutomationRule.AutomationRuleAction.EVALUATOR); + var offset = (pageNum - 1) * size; + + var automationRuleEvaluators = dao.find(workspaceId, projectId, Collections.emptySet(), name, AutomationRule.AutomationRuleAction.EVALUATOR, offset, size) + .stream() + .map(evaluator -> switch (evaluator) { + case LlmAsJudgeAutomationRuleEvaluatorModel llmAsJudge -> + AutomationModelEvaluatorMapper.INSTANCE.map(llmAsJudge); + }) + .toList(); + log.info("Found {} AutomationRuleEvaluators for projectId '{}'", automationRuleEvaluators.size(), projectId); + + return new AutomationRuleEvaluator.AutomationRuleEvaluatorPage(pageNum, automationRuleEvaluators.size(), total, + automationRuleEvaluators); + + }); + } + +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleModel.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleModel.java new file mode 100644 index 0000000000..ef92968db3 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleModel.java @@ -0,0 +1,20 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRule; + +import java.util.UUID; + +public sealed interface AutomationRuleModel + permits AutomationRuleEvaluatorModel { + + UUID id(); + UUID projectId(); + String name(); + + Float samplingRate(); + + String createdBy(); + String lastUpdatedBy(); + + AutomationRule.AutomationRuleAction action(); +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleRowMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleRowMapper.java new file mode 100644 index 0000000000..277d4900c4 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/AutomationRuleRowMapper.java @@ -0,0 +1,25 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRule; +import com.comet.opik.api.AutomationRuleEvaluator; +import org.jdbi.v3.core.mapper.RowMapper; +import org.jdbi.v3.core.statement.StatementContext; + +import java.sql.ResultSet; +import java.sql.SQLException; + +public class AutomationRuleRowMapper implements RowMapper> { + + @Override + public AutomationRule map(ResultSet rs, StatementContext ctx) throws SQLException { + + var action = AutomationRule.AutomationRuleAction.fromString(rs.getString("action")); + + return switch (action) { + case EVALUATOR -> ctx.findMapperFor(AutomationRuleEvaluator.class) + .orElseThrow(() -> new IllegalStateException( + "No mapper found for Automation Rule Action type: %s".formatted(action))) + .map(rs, ctx); + }; + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmAsJudgeAutomationRuleEvaluatorModel.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmAsJudgeAutomationRuleEvaluatorModel.java new file mode 100644 index 0000000000..d8f51460c1 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmAsJudgeAutomationRuleEvaluatorModel.java @@ -0,0 +1,30 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.AutomationRuleEvaluatorType; +import com.fasterxml.jackson.databind.JsonNode; +import lombok.Builder; +import org.jdbi.v3.json.Json; + +import java.time.Instant; +import java.util.UUID; + + +@Builder(toBuilder = true) +public record LlmAsJudgeAutomationRuleEvaluatorModel ( + UUID id, + UUID projectId, + String name, + Float samplingRate, + @Json JsonNode code, + Instant createdAt, + String createdBy, + Instant lastUpdatedAt, + String lastUpdatedBy +) implements AutomationRuleEvaluatorModel { + + @Override + public AutomationRuleEvaluatorType type() { + return AutomationRuleEvaluatorType.LLM_AS_JUDGE; + } + +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/JsonNodeArgumentFactory.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/JsonNodeArgumentFactory.java new file mode 100644 index 0000000000..52d479387b --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/JsonNodeArgumentFactory.java @@ -0,0 +1,36 @@ +package com.comet.opik.infrastructure.db; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.jdbi.v3.core.argument.AbstractArgumentFactory; +import org.jdbi.v3.core.argument.Argument; +import org.jdbi.v3.core.config.ConfigRegistry; + +import java.sql.Types; + +public class JsonNodeArgumentFactory extends AbstractArgumentFactory { + + private final ObjectMapper objectMapper = new ObjectMapper(); + + public JsonNodeArgumentFactory() { + super(Types.OTHER); + } + + @Override + protected Argument build(JsonNode value, ConfigRegistry config) { + return (position, statement, ctx) -> { + try { + if (value == null) { + statement.setNull(position, Types.NULL); + } else { + // Convert JsonNode to a JSON string before binding + statement.setObject(position, objectMapper.writeValueAsString(value)); + } + } catch (JsonProcessingException e) { + // Wrap the exception in a RuntimeException to comply with the signature + throw new IllegalArgumentException("Failed to serialize JsonNode to JSON string", e); + } + }; + } +} diff --git a/apps/opik-backend/src/main/resources/liquibase/db-app-state/migrations/000009_add_automation_rule_tables.sql b/apps/opik-backend/src/main/resources/liquibase/db-app-state/migrations/000009_add_automation_rule_tables.sql new file mode 100644 index 0000000000..6c97908792 --- /dev/null +++ b/apps/opik-backend/src/main/resources/liquibase/db-app-state/migrations/000009_add_automation_rule_tables.sql @@ -0,0 +1,28 @@ +--liquibase formatted sql +--changeset DanielAugusto:000009_add_automation_rule_tables +CREATE TABLE IF NOT EXISTS automation_rules ( + id CHAR(36), + project_id CHAR(36) NOT NULL, + workspace_id VARCHAR(150) NOT NULL, + + `action` ENUM('evaluator') NOT NULL, + name VARCHAR(150) NOT NULL, + sampling_rate FLOAT NOT NULL CHECK (sampling_rate >= 0 AND sampling_rate <= 1), + + CONSTRAINT `automation_rules_pk` PRIMARY KEY (id), + INDEX `automation_rules_idx` (workspace_id, project_id, id) +); + +CREATE TABLE IF NOT EXISTS automation_rule_evaluators ( + id CHAR(36), + + `type` ENUM('llm_as_judge') NOT NULL, + code JSON NOT NULL, + + created_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + created_by VARCHAR(100) NOT NULL DEFAULT 'admin', + last_updated_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6), + last_updated_by VARCHAR(100) NOT NULL DEFAULT 'admin', + + CONSTRAINT `automation_rules_evaluators_pk` PRIMARY KEY (id) +); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java new file mode 100644 index 0000000000..66399bb4f6 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/AutomationRuleEvaluatorsResourceTest.java @@ -0,0 +1,717 @@ +package com.comet.opik.api.resources.v1.priv; + +import com.comet.opik.api.AutomationRuleEvaluator; +import com.comet.opik.api.AutomationRuleEvaluatorUpdate; +import com.comet.opik.api.BatchDelete; +import com.comet.opik.api.FeedbackDefinition; +import com.comet.opik.api.resources.utils.AuthTestUtils; +import com.comet.opik.api.resources.utils.ClientSupportUtils; +import com.comet.opik.api.resources.utils.MigrationUtils; +import com.comet.opik.api.resources.utils.MySQLContainerUtils; +import com.comet.opik.api.resources.utils.RedisContainerUtils; +import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; +import com.comet.opik.api.resources.utils.WireMockUtils; +import com.comet.opik.podam.PodamFactoryUtils; +import com.github.tomakehurst.wiremock.client.WireMock; +import com.redis.testcontainers.RedisContainer; +import jakarta.ws.rs.HttpMethod; +import jakarta.ws.rs.client.Entity; +import jakarta.ws.rs.core.HttpHeaders; +import jakarta.ws.rs.core.MediaType; +import org.jdbi.v3.core.Jdbi; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.testcontainers.containers.MySQLContainer; +import org.testcontainers.lifecycle.Startables; +import ru.vyarus.dropwizard.guice.test.ClientSupport; +import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; +import uk.co.jemos.podam.api.PodamFactory; + +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.UUID; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import static com.comet.opik.infrastructure.auth.RequestContext.SESSION_COOKIE; +import static com.comet.opik.infrastructure.auth.RequestContext.WORKSPACE_HEADER; +import static com.comet.opik.infrastructure.auth.TestHttpClientUtils.UNAUTHORIZED_RESPONSE; +import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; +import static com.github.tomakehurst.wiremock.client.WireMock.matching; +import static com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath; +import static com.github.tomakehurst.wiremock.client.WireMock.okJson; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo; +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@DisplayName("Automation Rule Evaluators Resource Test") +class AutomationRuleEvaluatorsResourceTest { + + private static final String URL_TEMPLATE = "%s/v1/private/automations/projects/%s/evaluators/"; + + private static final String USER = UUID.randomUUID().toString(); + private static final String API_KEY = UUID.randomUUID().toString(); + private static final String WORKSPACE_ID = UUID.randomUUID().toString(); + private static final String TEST_WORKSPACE = UUID.randomUUID().toString(); + + private static final RedisContainer REDIS = RedisContainerUtils.newRedisContainer(); + + private static final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer(); + + @RegisterExtension + private static final TestDropwizardAppExtension app; + + private static final WireMockUtils.WireMockRuntime wireMock; + + static { + Startables.deepStart(REDIS, MYSQL).join(); + + wireMock = WireMockUtils.startWireMock(); + + app = TestDropwizardAppExtensionUtils.newTestDropwizardAppExtension(MYSQL.getJdbcUrl(), null, + wireMock.runtimeInfo(), REDIS.getRedisURI()); + } + + private final PodamFactory factory = PodamFactoryUtils.newPodamFactory(); + + private String baseURI; + private ClientSupport client; + + @BeforeAll + void setUpAll(ClientSupport client, Jdbi jdbi) { + + MigrationUtils.runDbMigration(jdbi, MySQLContainerUtils.migrationParameters()); + + this.baseURI = "http://localhost:%d".formatted(client.getPort()); + this.client = client; + + ClientSupportUtils.config(client); + + mockTargetWorkspace(API_KEY, TEST_WORKSPACE, WORKSPACE_ID); + } + + private static void mockTargetWorkspace(String apiKey, String workspaceName, String workspaceId) { + AuthTestUtils.mockTargetWorkspace(wireMock.server(), apiKey, workspaceName, workspaceId, USER); + } + + @AfterAll + void tearDownAll() { + wireMock.server().stop(); + } + + private UUID create(AutomationRuleEvaluator evaluator, String apiKey, String workspaceName) { + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, evaluator.getProjectId())) + .request() + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(evaluator))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); + + return TestUtils.getIdFromLocation(actualResponse.getLocation()); + } + } + + @Nested + @DisplayName("Api Key Authentication:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class ApiKey { + + private final String fakeApikey = UUID.randomUUID().toString(); + private final String okApikey = UUID.randomUUID().toString(); + + Stream credentials() { + return Stream.of( + arguments(okApikey, true), + arguments(fakeApikey, false), + arguments("", false)); + } + + @BeforeEach + void setUp() { + + wireMock.server().stubFor( + post(urlPathEqualTo("/opik/auth")) + .withHeader(HttpHeaders.AUTHORIZATION, equalTo(fakeApikey)) + .withRequestBody(matchingJsonPath("$.workspaceName", matching(".+"))) + .willReturn(WireMock.unauthorized())); + + wireMock.server().stubFor( + post(urlPathEqualTo("/opik/auth")) + .withHeader(HttpHeaders.AUTHORIZATION, equalTo("")) + .withRequestBody(matchingJsonPath("$.workspaceName", matching(".+"))) + .willReturn(WireMock.unauthorized())); + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("create evaluator definition: when api key is present, then return proper response") + void createAutomationRuleEvaluator__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, + boolean isAuthorized) { + + var ruleEvaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build(); + + mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, ruleEvaluator.getProjectId())) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .post(Entity.json(ruleEvaluator))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("get evaluators by project id: when api key is present, then return proper response") + void getProjectAutomationRuleEvaluators__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, + boolean isAuthorized) { + + final String workspaceName = UUID.randomUUID().toString(); + final String workspaceId = UUID.randomUUID().toString(); + final UUID projectId = UUID.randomUUID(); + + mockTargetWorkspace(okApikey, workspaceName, workspaceId); + + int samplesToCreate = 15; + + IntStream.range(0, samplesToCreate).forEach(i -> { + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class) + .toBuilder().id(null).projectId(projectId).build(); + create(evaluator, okApikey, workspaceName); + }); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, projectId)) + .queryParam("size", samplesToCreate) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var actualEntity = actualResponse + .readEntity(AutomationRuleEvaluator.AutomationRuleEvaluatorPage.class); + assertThat(actualEntity.content()).hasSize(samplesToCreate); + assertThat(actualEntity.total()).isEqualTo(samplesToCreate); + + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @Test + @DisplayName("search project evaluators: when searching by name, then return evaluators") + void find__whenSearchingByName__thenReturnEvaluators() { + + var workspaceName = UUID.randomUUID().toString(); + var workspaceId = UUID.randomUUID().toString(); + var projectId = UUID.randomUUID(); + var apiKey = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + var name = "Evaluator Name: " + UUID.randomUUID(); + + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class) + .toBuilder().id(null) + .projectId(projectId) + .name(name) + .build(); + + create(evaluator, apiKey, workspaceName); + + var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, projectId)) + .queryParam("name", "aluator") + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get(); + + var actualEntity = actualResponse.readEntity(AutomationRuleEvaluator.AutomationRuleEvaluatorPage.class); + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualEntity.page()).isEqualTo(1); + assertThat(actualEntity.size()).isEqualTo(1); + assertThat(actualEntity.total()).isEqualTo(1); + + List content = actualEntity.content(); + assertThat(content.stream().map(AutomationRuleEvaluator::getName).toList()).contains(name); + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("get evaluator by id: when api key is present, then return proper response") + void getAutomationRuleEvaluatorById__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, + boolean isAuthorized) { + + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build(); + + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + + mockTargetWorkspace(okApikey, workspaceName, workspaceId); + + UUID id = create(evaluator, okApikey, workspaceName); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, evaluator.getProjectId())) + .path(id.toString()) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var ruleEvaluator = actualResponse.readEntity(AutomationRuleEvaluator.class); + assertThat(ruleEvaluator.getId()).isEqualTo(id); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("update evaluator: when api key is present, then return proper response") + void updateAutomationRuleEvaluator__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, + boolean isAuthorized) { + + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build(); + + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + + mockTargetWorkspace(okApikey, workspaceName, workspaceId); + + UUID id = create(evaluator, okApikey, workspaceName); + + var updatedEvaluator = factory.manufacturePojo(AutomationRuleEvaluatorUpdate.class); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, evaluator.getProjectId())) + .path(id.toString()) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .method(HttpMethod.PATCH, Entity.json(updatedEvaluator))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("delete evaluator by id: when api key is present, then return proper response") + void deleteAutomationRuleEvaluator__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, + boolean isAuthorized) { + + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build();; + + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + + mockTargetWorkspace(okApikey, workspaceName, workspaceId); + + UUID id = create(evaluator, okApikey, workspaceName); + + var deleteMethod = BatchDelete.builder().ids(Collections.singleton(id)).build(); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, evaluator.getProjectId())) + .path("delete") + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(deleteMethod))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("batch delete evaluators by id: when api key is present, then return proper response") + void deleteProjectAutomationRuleEvaluators__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, + boolean isAuthorized) { + var projectId = UUID.randomUUID(); + var workspaceName = UUID.randomUUID().toString(); + var workspaceId = UUID.randomUUID().toString(); + + mockTargetWorkspace(okApikey, workspaceName, workspaceId); + + var evaluator1 = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().projectId(projectId).build(); + var evalId1 = create(evaluator1, okApikey, workspaceName); + + var evaluator2 = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().projectId(projectId).build(); + var evalId2 = create(evaluator2, okApikey, workspaceName); + + var evaluator3 = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().projectId(projectId).build(); + create(evaluator3, okApikey, workspaceName); + + var evalIds1and2 = Set.of(evalId1, evalId2); + var deleteMethod = BatchDelete.builder().ids(evalIds1and2).build(); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, projectId)) + .path("delete") + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(deleteMethod))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + + // we shall see a single evaluators for the project now + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, projectId)) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var actualEntity = actualResponse + .readEntity(AutomationRuleEvaluator.AutomationRuleEvaluatorPage.class); + assertThat(actualEntity.content()).hasSize(1); + assertThat(actualEntity.total()).isEqualTo(1); + + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + } + + @Nested + @DisplayName("Session Token Authentication:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class SessionTokenCookie { + + private final String sessionToken = UUID.randomUUID().toString(); + private final String fakeSessionToken = UUID.randomUUID().toString(); + + Stream credentials() { + return Stream.of( + arguments(sessionToken, true, "OK_" + UUID.randomUUID()), + arguments(fakeSessionToken, false, UUID.randomUUID().toString())); + } + + @BeforeEach + void setUp() { + wireMock.server().stubFor( + post(urlPathEqualTo("/opik/auth-session")) + .withCookie(SESSION_COOKIE, equalTo(sessionToken)) + .withRequestBody(matchingJsonPath("$.workspaceName", matching("OK_.+"))) + .willReturn(okJson(AuthTestUtils.newWorkspaceAuthResponse(USER, WORKSPACE_ID)))); + + wireMock.server().stubFor( + post(urlPathEqualTo("/opik/auth-session")) + .withCookie(SESSION_COOKIE, equalTo(fakeSessionToken)) + .withRequestBody(matchingJsonPath("$.workspaceName", matching(".+"))) + .willReturn(WireMock.unauthorized())); + } + + // .cookie(SESSION_COOKIE, sessionToken) + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("create evaluator definition: when api key is present, then return proper response") + void createAutomationRuleEvaluator__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken, + boolean isAuthorized, + String workspaceName) { + + var ruleEvaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build(); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, ruleEvaluator.getProjectId())) + .request() + .cookie(SESSION_COOKIE, sessionToken) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(ruleEvaluator))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("get evaluators by project id: when api key is present, then return proper response") + void getProjectAutomationRuleEvaluators__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken, + boolean isAuthorized, + String workspaceName) { + + var projectId = UUID.randomUUID(); + + int samplesToCreate = 15; + var newWorkspaceName = UUID.randomUUID().toString(); + var newWorkspaceId = UUID.randomUUID().toString(); + + wireMock.server().stubFor( + post(urlPathEqualTo("/opik/auth-session")) + .withCookie(SESSION_COOKIE, equalTo(sessionToken)) + .withRequestBody(matchingJsonPath("$.workspaceName", equalTo(newWorkspaceName))) + .willReturn(okJson(AuthTestUtils.newWorkspaceAuthResponse(USER, newWorkspaceId)))); + + + IntStream.range(0, samplesToCreate).forEach(i -> { + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class) + .toBuilder().id(null).projectId(projectId).build(); + create(evaluator, API_KEY, TEST_WORKSPACE); + }); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, projectId)) + .queryParam("size", samplesToCreate) + .request() + .cookie(SESSION_COOKIE, sessionToken) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var actualEntity = actualResponse + .readEntity(AutomationRuleEvaluator.AutomationRuleEvaluatorPage.class); + assertThat(actualEntity.content()).hasSize(samplesToCreate); + assertThat(actualEntity.total()).isEqualTo(samplesToCreate); + + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("get evaluator by id: when api key is present, then return proper response") + void getAutomationRuleEvaluatorById__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken, + boolean isAuthorized, + String workspaceName) { + + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build(); + + UUID id = create(evaluator, API_KEY, TEST_WORKSPACE); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, evaluator.getProjectId())) + .path(id.toString()) + .request() + .cookie(SESSION_COOKIE, sessionToken) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var ruleEvaluator = actualResponse.readEntity(AutomationRuleEvaluator.class); + assertThat(ruleEvaluator.getId()).isEqualTo(id); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("update evaluator: when api key is present, then return proper response") + void updateAutomationRuleEvaluator__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken, + boolean isAuthorized, + String workspaceName) { + + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build(); + + UUID id = create(evaluator, API_KEY, TEST_WORKSPACE); + + var updatedEvaluator = factory.manufacturePojo(AutomationRuleEvaluatorUpdate.class); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, evaluator.getProjectId())) + .path(id.toString()) + .request() + .cookie(SESSION_COOKIE, sessionToken) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .method(HttpMethod.PATCH, Entity.json(updatedEvaluator))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("delete evaluator by id: when api key is present, then return proper response") + void deleteAutomationRuleEvaluator__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken, + boolean isAuthorized, + String workspaceName) { + + var evaluator = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().id(null).build();; + + var id = create(evaluator, API_KEY, TEST_WORKSPACE); + var deleteMethod = BatchDelete.builder().ids(Collections.singleton(id)).build(); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, evaluator.getProjectId())) + .path("delete") + .request() + .cookie(SESSION_COOKIE, sessionToken) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(deleteMethod))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + + @ParameterizedTest + @MethodSource("credentials") + @DisplayName("batch delete evaluators by id: when api key is present, then return proper response") + void deleteProjectAutomationRuleEvaluators__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken, + boolean isAuthorized, + String workspaceName) { + + var projectId = UUID.randomUUID(); + + var evaluator1 = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().projectId(projectId).build(); + var evalId1 = create(evaluator1, API_KEY, TEST_WORKSPACE); + + var evaluator2 = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().projectId(projectId).build(); + var evalId2 = create(evaluator2, API_KEY, TEST_WORKSPACE); + + var evaluator3 = factory.manufacturePojo(AutomationRuleEvaluator.AutomationRuleEvaluatorLlmAsJudge.class).toBuilder().projectId(projectId).build(); + create(evaluator3, API_KEY, TEST_WORKSPACE); + + var evalIds1and2 = Set.of(evalId1, evalId2); + var deleteMethod = BatchDelete.builder().ids(evalIds1and2).build(); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, projectId)) + .path("delete") + .request() + .cookie(SESSION_COOKIE, sessionToken) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(deleteMethod))) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + + // we shall see a single evaluators for the project now + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI, projectId)) + .request() + .cookie(SESSION_COOKIE, sessionToken) + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + if (isAuthorized) { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var actualEntity = actualResponse + .readEntity(AutomationRuleEvaluator.AutomationRuleEvaluatorPage.class); + assertThat(actualEntity.content()).hasSize(1); + assertThat(actualEntity.total()).isEqualTo(1); + + } else { + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401); + assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class)) + .isEqualTo(UNAUTHORIZED_RESPONSE); + } + } + } + } + +}