Skip to content

Commit

Permalink
OPIK-52 Add trace usage search (#352)
Browse files Browse the repository at this point in the history
  • Loading branch information
@andrescrz
andrescrz authored Oct 7, 2024
1 parent a422d75 commit 92f7857
Show file tree
Hide file tree
Showing 9 changed files with 455 additions and 207 deletions.
3 changes: 3 additions & 0 deletions apps/opik-backend/src/main/java/com/comet/opik/api/filter/Field.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ public interface Field {
String OUTPUT_QUERY_PARAM = "output";
String METADATA_QUERY_PARAM = "metadata";
String TAGS_QUERY_PARAM = "tags";
String USAGE_COMPLETION_TOKENS_QUERY_PARAM = "usage.completion_tokens";
String USAGE_PROMPT_TOKENS_QUERY_PARAM = "usage.prompt_tokens";
String USAGE_TOTAL_TOKEN_QUERY_PARAMS = "usage.total_tokens";
String FEEDBACK_SCORES_QUERY_PARAM = "feedback_scores";

@JsonValue
Expand Down
6 changes: 3 additions & 3 deletions apps/opik-backend/src/main/java/com/comet/opik/api/filter/SpanField.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,9 @@ public enum SpanField implements Field {
OUTPUT(OUTPUT_QUERY_PARAM, FieldType.STRING),
METADATA(METADATA_QUERY_PARAM, FieldType.DICTIONARY),
TAGS(TAGS_QUERY_PARAM, FieldType.LIST),
USAGE_COMPLETION_TOKENS("usage.completion_tokens", FieldType.NUMBER),
USAGE_PROMPT_TOKENS("usage.prompt_tokens", FieldType.NUMBER),
USAGE_TOTAL_TOKENS("usage.total_tokens", FieldType.NUMBER),
USAGE_COMPLETION_TOKENS(USAGE_COMPLETION_TOKENS_QUERY_PARAM, FieldType.NUMBER),
USAGE_PROMPT_TOKENS(USAGE_PROMPT_TOKENS_QUERY_PARAM, FieldType.NUMBER),
USAGE_TOTAL_TOKENS(USAGE_TOTAL_TOKEN_QUERY_PARAMS, FieldType.NUMBER),
FEEDBACK_SCORES(FEEDBACK_SCORES_QUERY_PARAM, FieldType.FEEDBACK_SCORES_NUMBER),
;

Expand Down
3 changes: 3 additions & 0 deletions apps/opik-backend/src/main/java/com/comet/opik/api/filter/TraceField.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@ public enum TraceField implements Field {
OUTPUT(OUTPUT_QUERY_PARAM, FieldType.STRING),
METADATA(METADATA_QUERY_PARAM, FieldType.DICTIONARY),
TAGS(TAGS_QUERY_PARAM, FieldType.LIST),
USAGE_COMPLETION_TOKENS(USAGE_COMPLETION_TOKENS_QUERY_PARAM, FieldType.NUMBER),
USAGE_PROMPT_TOKENS(USAGE_PROMPT_TOKENS_QUERY_PARAM, FieldType.NUMBER),
USAGE_TOTAL_TOKENS(USAGE_TOTAL_TOKEN_QUERY_PARAMS, FieldType.NUMBER),
FEEDBACK_SCORES(FEEDBACK_SCORES_QUERY_PARAM, FieldType.FEEDBACK_SCORES_NUMBER),
;

Expand Down
75 changes: 51 additions & 24 deletions apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,9 @@ LEFT JOIN (
) AS s ON t.id = s.trace_id
GROUP BY
t.*
<if(trace_aggregation_filters)>
HAVING <trace_aggregation_filters>
<endif>
ORDER BY t.id DESC
;
""";
Expand All @@ -327,33 +330,53 @@ WHERE created_at BETWEEN toStartOfDay(yesterday()) AND toStartOfDay(today())
private static final String COUNT_BY_PROJECT_ID = """
SELECT
count(id) as count
FROM
(
SELECT
id
FROM traces
WHERE project_id = :project_id
AND workspace_id = :workspace_id
<if(filters)> AND <filters> <endif>
<if(feedback_scores_filters)>
AND id in (
FROM (
SELECT
t.id,
sumMap(s.usage) as usage
FROM (
SELECT
entity_id
FROM (
SELECT *
FROM feedback_scores
WHERE entity_type = 'trace'
AND workspace_id = :workspace_id
AND project_id = :project_id
ORDER BY entity_id DESC, last_updated_at DESC
LIMIT 1 BY entity_id, name
)
id
FROM traces
WHERE project_id = :project_id
AND workspace_id = :workspace_id
<if(filters)> AND <filters> <endif>
<if(feedback_scores_filters)>
AND id in (
SELECT
entity_id
FROM (
SELECT *
FROM feedback_scores
WHERE entity_type = 'trace'
AND workspace_id = :workspace_id
AND project_id = :project_id
ORDER BY entity_id DESC, last_updated_at DESC
LIMIT 1 BY entity_id, name
)
GROUP BY entity_id
HAVING <feedback_scores_filters>
)
<endif>
ORDER BY id DESC, last_updated_at DESC
LIMIT 1 BY id
)
<endif>
ORDER BY id DESC, last_updated_at DESC
LIMIT 1 BY id
) AS t
LEFT JOIN (
SELECT
trace_id,
usage
FROM spans
WHERE workspace_id = :workspace_id
AND project_id = :project_id
ORDER BY id DESC, last_updated_at DESC
LIMIT 1 BY id
) AS s ON t.id = s.trace_id
GROUP BY
t.id
<if(trace_aggregation_filters)>
HAVING <trace_aggregation_filters>
<endif>
ORDER BY t.id DESC
) AS latest_rows
;
""";
Expand Down Expand Up @@ -756,6 +779,9 @@ private ST newFindTemplate(String query, TraceSearchCriteria traceSearchCriteria
.ifPresent(filters -> {
filterQueryBuilder.toAnalyticsDbFilters(filters, FilterStrategy.TRACE)
.ifPresent(traceFilters -> template.add("filters", traceFilters));
filterQueryBuilder.toAnalyticsDbFilters(filters, FilterStrategy.TRACE_AGGREGATION)
.ifPresent(traceAggregationFilters -> template.add("trace_aggregation_filters",
traceAggregationFilters));
filterQueryBuilder.toAnalyticsDbFilters(filters, FilterStrategy.FEEDBACK_SCORES)
.ifPresent(scoresFilters -> template.add("feedback_scores_filters", scoresFilters));
});
Expand All @@ -766,6 +792,7 @@ private void bindSearchCriteria(TraceSearchCriteria traceSearchCriteria, Stateme
Optional.ofNullable(traceSearchCriteria.filters())
.ifPresent(filters -> {
filterQueryBuilder.bind(statement, filters, FilterStrategy.TRACE);
filterQueryBuilder.bind(statement, filters, FilterStrategy.TRACE_AGGREGATION);
filterQueryBuilder.bind(statement, filters, FilterStrategy.FEEDBACK_SCORES);
});
}
Expand Down
17 changes: 14 additions & 3 deletions apps/opik-backend/src/main/java/com/comet/opik/domain/filter/FilterQueryBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@ public class FilterQueryBuilder {
private static final String OUTPUT_ANALYTICS_DB = "output";
private static final String METADATA_ANALYTICS_DB = "metadata";
private static final String TAGS_ANALYTICS_DB = "tags";
private static final String USAGE_COMPLETION_TOKENS_ANALYTICS_DB = "usage['completion_tokens']";
private static final String USAGE_PROMPT_TOKENS_ANALYTICS_DB = "usage['prompt_tokens']";
private static final String USAGE_TOTAL_TOKENS_ANALYTICS_DB = "usage['total_tokens']";
private static final String VALUE_ANALYTICS_DB = "value";

private static final Map<Operator, Map<FieldType, String>> ANALYTICS_DB_OPERATOR_MAP = new EnumMap<>(Map.of(
Expand Down Expand Up @@ -92,6 +95,9 @@ public class FilterQueryBuilder {
.put(TraceField.OUTPUT, OUTPUT_ANALYTICS_DB)
.put(TraceField.METADATA, METADATA_ANALYTICS_DB)
.put(TraceField.TAGS, TAGS_ANALYTICS_DB)
.put(TraceField.USAGE_COMPLETION_TOKENS, USAGE_COMPLETION_TOKENS_ANALYTICS_DB)
.put(TraceField.USAGE_PROMPT_TOKENS, USAGE_PROMPT_TOKENS_ANALYTICS_DB)
.put(TraceField.USAGE_TOTAL_TOKENS, USAGE_TOTAL_TOKENS_ANALYTICS_DB)
.put(TraceField.FEEDBACK_SCORES, VALUE_ANALYTICS_DB)
.build());

Expand All @@ -105,9 +111,9 @@ public class FilterQueryBuilder {
.put(SpanField.OUTPUT, OUTPUT_ANALYTICS_DB)
.put(SpanField.METADATA, METADATA_ANALYTICS_DB)
.put(SpanField.TAGS, TAGS_ANALYTICS_DB)
.put(SpanField.USAGE_COMPLETION_TOKENS, "usage['completion_tokens']")
.put(SpanField.USAGE_PROMPT_TOKENS, "usage['prompt_tokens']")
.put(SpanField.USAGE_TOTAL_TOKENS, "usage['total_tokens']")
.put(SpanField.USAGE_COMPLETION_TOKENS, USAGE_COMPLETION_TOKENS_ANALYTICS_DB)
.put(SpanField.USAGE_PROMPT_TOKENS, USAGE_PROMPT_TOKENS_ANALYTICS_DB)
.put(SpanField.USAGE_TOTAL_TOKENS, USAGE_TOTAL_TOKENS_ANALYTICS_DB)
.put(SpanField.FEEDBACK_SCORES, VALUE_ANALYTICS_DB)
.build());

Expand All @@ -122,6 +128,11 @@ public class FilterQueryBuilder {
.add(TraceField.METADATA)
.add(TraceField.TAGS)
.build()),
FilterStrategy.TRACE_AGGREGATION, EnumSet.copyOf(ImmutableSet.<TraceField>builder()
.add(TraceField.USAGE_COMPLETION_TOKENS)
.add(TraceField.USAGE_PROMPT_TOKENS)
.add(TraceField.USAGE_TOTAL_TOKENS)
.build()),
FilterStrategy.SPAN, EnumSet.copyOf(ImmutableSet.<SpanField>builder()
.add(SpanField.ID)
.add(SpanField.NAME)
Expand Down
1 change: 1 addition & 0 deletions apps/opik-backend/src/main/java/com/comet/opik/domain/filter/FilterStrategy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

public enum FilterStrategy {
TRACE,
TRACE_AGGREGATION,
SPAN,
FEEDBACK_SCORES
}
7 changes: 3 additions & 4 deletions ...opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1204,10 +1204,9 @@ void create__whenDescriptionIsMultiline__thenAcceptTheRequest() {
var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
.id(null)
.description("""
Test
Description
"""
)
Test
Description
""")
.build();

createAndAssert(dataset);
Expand Down
Loading

0 comments on commit 92f7857

Please sign in to comment.