Skip to content

Commit

Permalink
[OPIK-397] Add trace stats endpoint (#692)
Browse files Browse the repository at this point in the history
* [OPIK-397] Add trace stats endpoint

* Add to string

* Test

* Fix test quantiles calc
  • Loading branch information
thiagohora authored Nov 26, 2024
1 parent 1d78f1c commit 6376b74
Show file tree
Hide file tree
Showing 6 changed files with 2,903 additions and 2 deletions.
105 changes: 105 additions & 0 deletions apps/opik-backend/src/main/java/com/comet/opik/api/ProjectStats.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
package com.comet.opik.api;

import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
import com.fasterxml.jackson.databind.annotation.JsonNaming;
import io.swagger.v3.oas.annotations.media.DiscriminatorMapping;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.Builder;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.ToString;
import lombok.experimental.SuperBuilder;

import java.beans.ConstructorProperties;
import java.math.BigDecimal;
import java.util.List;

@Builder(toBuilder = true)
@JsonIgnoreProperties(ignoreUnknown = true)
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
public record ProjectStats(List<ProjectStatItem<?>> stats) {

public static ProjectStats empty() {
return new ProjectStats(List.of());
}

@JsonIgnoreProperties(ignoreUnknown = true)
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "type", visible = true)
@JsonSubTypes({
@JsonSubTypes.Type(value = PercentageValueStat.class, name = "PERCENTAGE"),
@JsonSubTypes.Type(value = CountValueStat.class, name = "COUNT"),
@JsonSubTypes.Type(value = AvgValueStat.class, name = "AVG")
})
@Schema(discriminatorProperty = "type", discriminatorMapping = {
@DiscriminatorMapping(value = "PERCENTAGE", schema = PercentageValueStat.class),
@DiscriminatorMapping(value = "COUNT", schema = CountValueStat.class),
@DiscriminatorMapping(value = "AVG", schema = AvgValueStat.class),
})
@RequiredArgsConstructor
@Getter
@SuperBuilder(toBuilder = true)
@EqualsAndHashCode
@ToString(callSuper = true)
public abstract static sealed class ProjectStatItem<T> {
private final String name;
private final T value;
private final StatsType type;
}

@ToString(callSuper = true)
@EqualsAndHashCode(callSuper = true)
@SuperBuilder(toBuilder = true)
@Getter
public abstract static sealed class SingleValueStat<T extends Number> extends ProjectStatItem<T> {
}

@ToString(callSuper = true)
@EqualsAndHashCode(callSuper = true)
@SuperBuilder(toBuilder = true)
@Getter
public static final class CountValueStat extends SingleValueStat<Long> {

@ConstructorProperties({"name", "value"})
public CountValueStat(String name, Long value) {
super(CountValueStat.builder().value(value).name(name).type(StatsType.COUNT));
}
}

@ToString(callSuper = true)
@EqualsAndHashCode(callSuper = true)
@SuperBuilder(toBuilder = true)
@Getter
public static final class AvgValueStat extends SingleValueStat<Number> {

@ConstructorProperties({"name", "value"})
public AvgValueStat(String name, Number value) {
super(AvgValueStat.builder().value(value).name(name).type(StatsType.AVG));
}
}

public record PercentageValues(double p50, double p90, double p99) {
}

@ToString(callSuper = true)
@EqualsAndHashCode(callSuper = true)
@SuperBuilder(toBuilder = true)
@Getter
public static final class PercentageValueStat extends ProjectStatItem<PercentageValues> {

@ConstructorProperties({"name", "value"})
public PercentageValueStat(String name, PercentageValues value) {
super(PercentageValueStat.builder().value(value).name(name).type(StatsType.PERCENTAGE));
}
}

public enum StatsType {
COUNT,
PERCENTAGE,
AVG
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import com.comet.opik.api.FeedbackScore;
import com.comet.opik.api.FeedbackScoreBatch;
import com.comet.opik.api.FeedbackScoreNames;
import com.comet.opik.api.ProjectStats;
import com.comet.opik.api.Trace;
import com.comet.opik.api.Trace.TracePage;
import com.comet.opik.api.TraceBatch;
Expand Down Expand Up @@ -294,6 +295,37 @@ public Response scoreBatchOfTraces(
return Response.noContent().build();
}

@GET
@Path("/stats")
@Operation(operationId = "getTraceStats", summary = "Get trace stats", description = "Get trace stats", responses = {
@ApiResponse(responseCode = "200", description = "Trace stats resource", content = @Content(schema = @Schema(implementation = ProjectStats.class)))
})
public Response getStats(@QueryParam("project_id") UUID projectId,
@QueryParam("project_name") String projectName,
@QueryParam("filters") String filters) {

validateProjectNameAndProjectId(projectName, projectId);
var traceFilters = filtersFactory.newFilters(filters, TraceFilter.LIST_TYPE_REFERENCE);
var searchCriteria = TraceSearchCriteria.builder()
.projectName(projectName)
.projectId(projectId)
.filters(traceFilters)
.build();

String workspaceId = requestContext.get().getWorkspaceId();

log.info("Get trace stats by '{}' on workspaceId '{}'", searchCriteria, workspaceId);

ProjectStats projectStats = service.getStats(searchCriteria)
.contextWrite(ctx -> setRequestContext(ctx, requestContext))
.block();

log.info("Found trace stats by '{}', count '{}' on workspaceId '{}'", searchCriteria,
projectStats.stats().size(), workspaceId);

return Response.ok(projectStats).build();
}

@GET
@Path("/feedback-scores/names")
@Operation(operationId = "findFeedbackScoreNames", summary = "Find Feedback Score names", description = "Find Feedback Score names", responses = {
Expand Down
Loading

0 comments on commit 6376b74

Please sign in to comment.