Feature/trace span batch annotation

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

@@ -0,0 +1,24 @@
+package com.comet.opik.api;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.databind.PropertyNamingStrategies;
+import com.fasterxml.jackson.databind.annotation.JsonNaming;
+import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.validation.constraints.NotBlank;
+import jakarta.validation.constraints.NotEmpty;
+import jakarta.validation.constraints.NotNull;
+import jakarta.validation.constraints.Size;
+import lombok.Builder;
+
+import java.util.List;
+import java.util.Set;
+import java.util.UUID;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
+@Builder(toBuilder = true)
+public record CommentsBatchCreate(
+		@Schema(description = "IDs of entities to comment on (max 1K)") @NotNull @Size(min = 1, max = 1000) Set<@NotNull UUID> ids,
+		@Schema(description = "Comment text to apply to all entities") @NotBlank String text
+) {}
+

apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java

@@ -4,6 +4,7 @@
 import com.comet.opik.api.BatchDelete;
 import com.comet.opik.api.Comment;
 import com.comet.opik.api.DeleteFeedbackScore;
+import com.comet.opik.api.CommentsBatchCreate;
 import com.comet.opik.api.FeedbackDefinition;
 import com.comet.opik.api.FeedbackScore;
 import com.comet.opik.api.FeedbackScoreBatchContainer;
@@ -487,13 +488,36 @@ public Response deleteSpanComments(
 
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        log.info("Delete span comments with ids '{}' on workspaceId '{}'", batchDelete.ids(), workspaceId);
+        log.info("Delete span comments with size '{}' on workspaceId '{}'", batchDelete.ids().size(), workspaceId);
 
         commentService.delete(batchDelete)
                 .contextWrite(ctx -> setRequestContext(ctx, requestContext))
                 .block();
 
-        log.info("Deleted span comments with ids '{}' on workspaceId '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleted span comments with size '{}' on workspaceId '{}'", batchDelete.ids().size(), workspaceId);
+
+        return Response.noContent().build();
+    }
+
+    @POST
+    @Path("/comments/batch")
+    @Operation(operationId = "createSpanCommentsBatch", summary = "Create comments for multiple spans", description = "Create the same comment text for multiple span IDs (max 1K per request)", responses = {
+            @ApiResponse(responseCode = "204", description = "No Content"),
+            @ApiResponse(responseCode = "400", description = "Bad Request", content = @Content(schema = @Schema(implementation = ErrorMessage.class))),
+            @ApiResponse(responseCode = "404", description = "Not found", content = @Content(schema = @Schema(implementation = ErrorMessage.class)))})
+    @RateLimited
+    public Response createSpanCommentsBatch(
+            @RequestBody(content = @Content(schema = @Schema(implementation = CommentsBatchCreate.class))) @NotNull @Valid CommentsBatchCreate payload) {
+
+        String workspaceId = requestContext.get().getWorkspaceId();
+
+        log.info("Creating span comments batch with size '{}' on workspaceId '{}'", payload.ids().size(), workspaceId);
+
+        commentService.createBatchForSpans(payload.ids(), payload.text())
+                .contextWrite(ctx -> setRequestContext(ctx, requestContext))
+                .block();
+
+        log.info("Created span comments batch with size '{}' on workspaceId '{}'", payload.ids().size(), workspaceId);
 
         return Response.noContent().build();
     }

apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TracesResource.java

@@ -4,6 +4,7 @@
 import com.comet.opik.api.BatchDelete;
 import com.comet.opik.api.BatchDeleteByProject;
 import com.comet.opik.api.Comment;
+import com.comet.opik.api.CommentsBatchCreate;
 import com.comet.opik.api.DeleteFeedbackScore;
 import com.comet.opik.api.DeleteThreadFeedbackScores;
 import com.comet.opik.api.DeleteTraceThreads;
@@ -535,13 +536,36 @@ public Response deleteTraceComments(
 
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        log.info("Delete trace comments with ids '{}' on workspaceId '{}'", batchDelete.ids(), workspaceId);
+        log.info("Delete trace comments with size '{}' on workspaceId '{}'", batchDelete.ids().size(), workspaceId);
 
         commentService.delete(batchDelete)
                 .contextWrite(ctx -> setRequestContext(ctx, requestContext))
                 .block();
 
-        log.info("Deleted trace comments with ids '{}' on workspaceId '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleted trace comments with size '{}' on workspaceId '{}'", batchDelete.ids().size(), workspaceId);
+
+        return Response.noContent().build();
+    }
+
+    @POST
+    @Path("/comments/batch")
+    @Operation(operationId = "createTraceCommentsBatch", summary = "Create comments for multiple traces", description = "Create the same comment text for multiple trace IDs (max 10 per request)", responses = {
+            @ApiResponse(responseCode = "204", description = "No Content"),
+            @ApiResponse(responseCode = "400", description = "Bad Request", content = @Content(schema = @Schema(implementation = ErrorMessage.class))),
+            @ApiResponse(responseCode = "404", description = "Not found", content = @Content(schema = @Schema(implementation = ErrorMessage.class)))})
+    @RateLimited
+    public Response createTraceCommentsBatch(
+            @RequestBody(content = @Content(schema = @Schema(implementation = CommentsBatchCreate.class))) @NotNull @Valid CommentsBatchCreate payload) {
+
+        String workspaceId = requestContext.get().getWorkspaceId();
+
+        log.info("Creating trace comments batch with size '{}' on workspaceId '{}'", payload.ids().size(), workspaceId);
+
+        commentService.createBatchForTraces(payload.ids(), payload.text())
+                .contextWrite(ctx -> setRequestContext(ctx, requestContext))
+                .block();
+
+        log.info("Created trace comments batch with size '{}' on workspaceId '{}'", payload.ids().size(), workspaceId);
 
         return Response.noContent().build();
     }
@@ -696,7 +720,7 @@ public Response deleteTraceThreads(
                 .contextWrite(ctx -> setRequestContext(ctx, requestContext))
                 .block();
 
-        log.info("Deleted trace threads with ids '{}' on workspaceId '{}'", traceThreads.threadIds(), workspaceId);
+        log.info("Deleted trace threads with size '{}' on workspaceId '{}'", traceThreads.threadIds().size(), workspaceId);
 
         return Response.noContent().build();
     }
@@ -922,13 +946,13 @@ public Response deleteThreadComments(
 
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        log.info("Delete thread comments with ids '{}' on workspaceId '{}'", batchDelete.ids(), workspaceId);
+        log.info("Delete thread comments with size '{}' on workspaceId '{}'", batchDelete.ids().size(), workspaceId);
 
         commentService.delete(batchDelete)
                 .contextWrite(ctx -> setRequestContext(ctx, requestContext))
                 .block();
 
-        log.info("Deleted thread comments with ids '{}' on workspaceId '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleted thread comments with size '{}' on workspaceId '{}'", batchDelete.ids().size(), workspaceId);
 
         return Response.noContent().build();
     }

apps/opik-backend/src/main/java/com/comet/opik/domain/CommentDAO.java

@@ -1,6 +1,7 @@
 package com.comet.opik.domain;
 
 import com.comet.opik.api.Comment;
+import com.comet.opik.utils.TemplateUtils;
 import com.comet.opik.infrastructure.db.TransactionTemplateAsync;
 import com.google.inject.ImplementedBy;
 import io.r2dbc.spi.Result;
@@ -14,6 +15,7 @@
 import org.stringtemplate.v4.ST;
 import reactor.core.publisher.Mono;
 
+import java.util.List;
 import java.util.Set;
 import java.util.UUID;
 
@@ -39,6 +41,8 @@ enum EntityType {
 
     Mono<Long> addComment(UUID commentId, UUID entityId, EntityType entityType, UUID projectId, Comment comment);
 
+    Mono<Long> addCommentsBatch(EntityType entityType, List<UUID> entityIds, List<UUID> commentIds, List<UUID> projectIds, Comment comment);
+
     Mono<Comment> findById(UUID entityId, UUID commentId);
 
     Mono<Void> updateComment(UUID commentId, Comment comment);
@@ -80,6 +84,30 @@ INSERT INTO comments(
             ;
             """;
 
+    private static final String INSERT_COMMENTS_BATCH = """
+            INSERT INTO comments(
+                id,
+                entity_id,
+                entity_type,
+                project_id,
+                workspace_id,
+                text,
+                created_by,
+                last_updated_by
+            ) VALUES
+                <items:{item | (
+                    :id<item.index>,
+                    :entity_id<item.index>,
+                    :entity_type,
+                    :project_id<item.index>,
+                    :workspace_id,
+                    :text,
+                    :user_name,
+                    :user_name
+                )<if(item.hasNext)>,<endif>}>
+            ;
+            """;
+
     private static final String SELECT_COMMENT_BY_ID = """
             SELECT
                 *
@@ -145,6 +173,28 @@ public Mono<Long> addComment(@NonNull UUID commentId, @NonNull UUID entityId, @N
         });
     }
 
+    @Override
+    public Mono<Long> addCommentsBatch(@NonNull EntityType entityType, @NonNull List<UUID> entityIds,
+            @NonNull List<UUID> commentIds, @NonNull List<UUID> projectIds, @NonNull Comment comment) {
+        return asyncTemplate.nonTransaction(connection -> makeMonoContextAware((userName, workspaceId) -> {
+            var items = TemplateUtils.getQueryItemPlaceHolder(entityIds.size());
+            var template = new ST(INSERT_COMMENTS_BATCH).add("items", items);
+
+            var statement = connection.createStatement(template.render())
+                    .bind("entity_type", entityType.getType())
+                    .bind("text", comment.text());
+
+            for (int i = 0; i < entityIds.size(); i++) {
+                statement.bind("id" + i, commentIds.get(i))
+                        .bind("entity_id" + i, entityIds.get(i))
+                        .bind("project_id" + i, projectIds.get(i));
+            }
+
+            return makeMonoContextAware(bindUserNameAndWorkspaceContext(statement))
+                    .flatMap(result -> Mono.from(result.getRowsUpdated()));
+        }));
+    }
+
     @Override
     public Mono<Comment> findById(UUID entityId, @NonNull UUID commentId) {
         return asyncTemplate.nonTransaction(connection -> {

apps/opik-backend/src/main/java/com/comet/opik/domain/CommentService.java

@@ -10,9 +10,13 @@
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import reactor.core.publisher.Mono;
+import com.comet.opik.infrastructure.db.TransactionTemplateAsync;
+import java.util.List;
+import java.util.Map;
 
 import java.util.Set;
 import java.util.UUID;
+import jakarta.ws.rs.BadRequestException;
 
 import static com.comet.opik.utils.ErrorUtils.failWithNotFound;
 
@@ -28,6 +32,10 @@ public interface CommentService {
     Mono<Void> delete(BatchDelete batchDelete);
 
     Mono<Long> deleteByEntityIds(CommentDAO.EntityType entityType, Set<UUID> entityIds);
+
+    Mono<Long> createBatchForTraces(Set<UUID> traceIds, String text);
+
+    Mono<Long> createBatchForSpans(Set<UUID> spanIds, String text);
 }
 
 @Slf4j
@@ -40,6 +48,7 @@ class CommentServiceImpl implements CommentService {
     private final @NonNull SpanDAO spanDAO;
     private final @NonNull TraceThreadDAO traceThreadDAO;
     private final @NonNull IdGenerator idGenerator;
+    private final @NonNull TransactionTemplateAsync transactionTemplateAsync;
 
     @Override
     public Mono<UUID> create(@NonNull UUID entityId, @NonNull Comment comment, CommentDAO.EntityType entityType) {
@@ -82,4 +91,43 @@ public Mono<Long> deleteByEntityIds(CommentDAO.EntityType entityType, Set<UUID>
         }
         return commentDAO.deleteByEntityIds(entityType, entityIds);
     }
+
+    @Override
+    public Mono<Long> createBatchForTraces(@NonNull Set<UUID> traceIds, @NonNull String text) {
+        return createBatch(
+                traceIds,
+                text,
+                transactionTemplateAsync.nonTransaction(connection -> traceDAO.getProjectIdsByTraceIds(traceIds, connection)),
+                CommentDAO.EntityType.TRACE);
+    }
+
+    @Override
+    public Mono<Long> createBatchForSpans(@NonNull Set<UUID> spanIds, @NonNull String text) {
+        return createBatch(
+                spanIds,
+                text,
+                spanDAO.getProjectIdsBySpanIds(spanIds),
+                CommentDAO.EntityType.SPAN);
+    }
+
+    private Mono<Long> createBatch(Set<UUID> idsSet,
+                                   String text,
+                                   Mono<Map<UUID, UUID>> idsToProjectIds,
+                                   CommentDAO.EntityType entityType) {
+        if (idsSet.isEmpty()) {
+            return Mono.just(0L);
+        }
+
+        List<UUID> ids = idsSet.stream().toList();
+
+        return idsToProjectIds.flatMap(map -> {
+            if (map.size() != ids.size()) {
+                return Mono.error(new BadRequestException("Some entities were not found in the workspace"));
+            }
+            List<UUID> projectIds = ids.stream().map(map::get).toList();
+            List<UUID> generatedIds = ids.stream().map(__ -> idGenerator.generateId()).toList();
+            return commentDAO.addCommentsBatch(entityType, ids, generatedIds, projectIds,
+                    Comment.builder().text(text).build());
+        });
+    }
 }

apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java

@@ -36,6 +36,7 @@
 import org.stringtemplate.v4.ST;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import java.util.Map;
 
 import java.math.BigDecimal;
 import java.time.Instant;
@@ -2008,6 +2009,16 @@ public Mono<List<WorkspaceAndResourceId>> getSpanWorkspace(@NonNull Set<UUID> sp
                 .collectList();
     }
 
+    private static final String SELECT_SPAN_ID_AND_PROJECT = """
+            SELECT
+                id,
+                project_id
+            FROM spans
+            WHERE id IN :spanIds
+            AND workspace_id = :workspace_id
+            ;
+            """;
+
     @WithSpan
     public Mono<UUID> getProjectIdFromSpan(@NonNull UUID spanId) {
 
@@ -2023,6 +2034,27 @@ public Mono<UUID> getProjectIdFromSpan(@NonNull UUID spanId) {
                 .singleOrEmpty();
     }
 
+    @WithSpan
+    public Mono<Map<UUID, UUID>> getProjectIdsBySpanIds(@NonNull Set<UUID> spanIds) {
+        if (spanIds.isEmpty()) {
+            return Mono.just(Map.of());
+        }
+
+        return Mono.from(connectionFactory.create())
+                .flatMapMany(connection -> {
+
+                    var statement = connection.createStatement(SELECT_SPAN_ID_AND_PROJECT)
+                            .bind("spanIds", spanIds.toArray(UUID[]::new));
+
+                    return makeFluxContextAware(bindWorkspaceIdToFlux(statement));
+                })
+                .flatMap(result -> result.map((row, rowMetadata) -> Map.entry(
+                        row.get("id", UUID.class),
+                        row.get("project_id", UUID.class)
+                )))
+                .collectMap(Map.Entry::getKey, Map.Entry::getValue);
+    }
+
     @WithSpan
     public Mono<ProjectStats> getStats(@NonNull SpanSearchCriteria searchCriteria) {