OPIK-968: Add trace count to project (#1258)

comet-ml · Feb 12, 2025 · c9ae181 · c9ae181
1 parent 3cb7c15
commit c9ae181
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 3 deletions.
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/Project.java b/apps/opik-backend/src/main/java/com/comet/opik/api/Project.java
@@ -40,7 +40,9 @@ public record Project(
         @JsonView({
                 Project.View.Detailed.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) @Nullable Double totalEstimatedCost,
         @JsonView({
-                Project.View.Detailed.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) @Nullable Map<String, Double> usage){
+                Project.View.Detailed.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) @Nullable Map<String, Double> usage,
+        @JsonView({
+                Project.View.Detailed.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) @Nullable Long traceCount){
 
     public static class View {
         public static class Write {

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/ProjectStatsSummary.java b/apps/opik-backend/src/main/java/com/comet/opik/api/ProjectStatsSummary.java
@@ -22,7 +22,8 @@ public record ProjectStatsSummaryItem(
             List<FeedbackScoreAverage> feedbackScores,
             ProjectStats.PercentageValues duration,
             Double totalEstimatedCost,
-            Map<String, Double> usage) {
+            Map<String, Double> usage,
+            Long traceCount) {
     }
 
 }
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/ProjectService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/ProjectService.java
@@ -242,6 +242,7 @@ private ProjectStatsSummaryItem getStats(UUID projectId, Map<String, Object> pro
                 .duration(StatsMapper.getStatsDuration(projectStats))
                 .totalEstimatedCost(StatsMapper.getStatsTotalEstimatedCost(projectStats))
                 .usage(StatsMapper.getStatsUsage(projectStats))
+                .traceCount(StatsMapper.getStatsTraceCount(projectStats))
                 .build();
     }
 
@@ -490,6 +491,7 @@ public Project retrieveByName(@NonNull String projectName) {
                                 .duration(StatsMapper.getStatsDuration(projectStats.get(project.id())))
                                 .totalEstimatedCost(
                                         StatsMapper.getStatsTotalEstimatedCost(projectStats.get(project.id())))
+                                .traceCount(StatsMapper.getStatsTraceCount(projectStats.get(project.id())))
                                 .build();
                     })
                     .orElseThrow(this::createNotFoundError);

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/stats/StatsMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/stats/StatsMapper.java
@@ -26,6 +26,7 @@ public class StatsMapper {
     public static final String OUTPUT = "output";
     public static final String METADATA = "metadata";
     public static final String TAGS = "tags";
+    public static final String TRACE_COUNT = "trace_count";
 
     public static ProjectStats mapProjectStats(Row row, String entityCountLabel) {
 
@@ -111,4 +112,10 @@ public static PercentageValues getStatsDuration(Map<String, ?> stats) {
                 .map(map -> (PercentageValues) map.get(DURATION))
                 .orElse(null);
     }
+
+    public static Long getStatsTraceCount(Map<String, Object> projectStats) {
+        return Optional.ofNullable(projectStats)
+                .map(map -> (Long) map.get(TRACE_COUNT))
+                .orElse(null);
+    }
 }
diff --git a/...opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java b/...opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java
@@ -117,7 +117,7 @@ class ProjectsResourceTest {
     public static final String URL_TEMPLATE = "%s/v1/private/projects";
     public static final String URL_TEMPLATE_TRACE = "%s/v1/private/traces";
     public static final String[] IGNORED_FIELDS = {"createdBy", "lastUpdatedBy", "createdAt", "lastUpdatedAt",
-            "lastUpdatedTraceAt", "feedbackScores", "duration", "totalEstimatedCost", "usage"};
+            "lastUpdatedTraceAt", "feedbackScores", "duration", "totalEstimatedCost", "usage", "traceCount"};
     public static final String[] IGNORED_FIELD_MIN = {"createdBy", "lastUpdatedBy", "createdAt", "lastUpdatedAt",
             "lastUpdatedTraceAt"};
 
@@ -1254,6 +1254,7 @@ private List<ProjectStatsSummaryItem> getProjectStatsSummaryItems(String apiKey,
                             .usage(project.usage())
                             .feedbackScores(project.feedbackScores())
                             .projectId(project.id())
+                            .traceCount(project.traceCount())
                             .build())
                     .toList();
         }
@@ -1495,6 +1496,7 @@ private Project buildProjectStats(Project project, String apiKey, String workspa
                         .collect(groupingBy(Map.Entry::getKey, averagingDouble(Map.Entry::getValue))))
                 .feedbackScores(getScoreAverages(traces))
                 .lastUpdatedTraceAt(traces.stream().map(Trace::lastUpdatedAt).max(Instant::compareTo).orElse(null))
+                .traceCount((long) traces.size())
                 .build();
     }