Merge branch 'main' into thiagohora/OPIK-147_add_filter_to_dataset_it…

…em_experiment_compare_endpoint
comet-ml · Oct 9, 2024 · 97898c5 · 97898c5
2 parents e72e12e + 76144c5
commit 97898c5
Show file tree

Hide file tree

Showing 68 changed files with 1,378 additions and 385 deletions.
diff --git a/.github/workflows/sanity.yml b/.github/workflows/sanity.yml
@@ -0,0 +1,63 @@
+name: Install Local Version of Opik
+
+on:
+    workflow_dispatch:
+
+jobs:
+    test_installation:
+        runs-on: ubuntu-20.04
+
+        steps:
+            - name: Checkout repo
+              uses: actions/checkout@v3
+              with:
+                ref: ${{ github.ref }}
+
+            - name: Setup Python
+              uses: actions/setup-python@v5
+              with:
+                python-version: 3.12
+
+            - name: Install Opik
+              run: pip install ${{ github.workspace }}/sdks/python
+
+            - name: Install Test Dependencies
+              run: |
+                pip install -r ${{ github.workspace }}/tests_end_to_end/test_requirements.txt
+                playwright install
+
+            - name: Install Opik
+              env:
+                OPIK_USAGE_REPORT_ENABLED: false
+              run: |
+                cd ${{ github.workspace }}/deployment/docker-compose
+                docker compose up -d --build
+                
+            - name: Check Docker pods are up
+              run: |
+                chmod +x ./tests_end_to_end/installer/check_docker_compose_pods.sh
+                ./tests_end_to_end/installer/check_docker_compose_pods.sh
+              shell: bash
+
+            - name: Check backend health
+              run: |
+                chmod +x ./tests_end_to_end/installer/check_backend.sh
+                ./tests_end_to_end/installer/check_backend.sh
+              shell: bash
+
+            - name: Check app is up via the UI
+              run: |
+                pytest -v -s ${{ github.workspace }}/tests_end_to_end/installer/test_app_status.py
+            
+            - name: Run sanity suite
+              run: |
+                cd ${{ github.workspace }}/tests_end_to_end
+                export PYTHONPATH='.'
+                pytest -s application_sanity/test_sanity.py --browser chromium --base-url http://localhost:5173 --setup-show
+
+            - name: Stop Opik server
+              if: always()
+              run: |
+                cd ${{ github.workspace }}/deployment/docker-compose
+                docker compose down
+                cd -
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
@@ -5,6 +5,7 @@
 import com.fasterxml.jackson.databind.PropertyNamingStrategies;
 import com.fasterxml.jackson.databind.annotation.JsonNaming;
 import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.annotation.Nullable;
 import jakarta.validation.constraints.NotBlank;
 import jakarta.validation.constraints.Pattern;
 import lombok.Builder;
@@ -29,7 +30,9 @@ public record Project(
         @JsonView({Project.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant createdAt,
         @JsonView({Project.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String createdBy,
         @JsonView({Project.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant lastUpdatedAt,
-        @JsonView({Project.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String lastUpdatedBy){
+        @JsonView({Project.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String lastUpdatedBy,
+        @JsonView({
+                Project.View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) @Nullable Instant lastUpdatedTraceAt){
 
     public static class View {
         public static class Write {
@@ -47,5 +50,9 @@ public record ProjectPage(
             @JsonView({Project.View.Public.class}) List<Project> content)
             implements
                 com.comet.opik.api.Page<Project>{
+
+        public static ProjectPage empty(int page) {
+            return new ProjectPage(page, 0, 0, List.of());
+        }
     }
 }
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/filter/Field.java b/apps/opik-backend/src/main/java/com/comet/opik/api/filter/Field.java
@@ -13,6 +13,9 @@ public interface Field {
     String METADATA_QUERY_PARAM = "metadata";
     String EXPECTED_OUTPUT_QUERY_PARAM = "expected_output";
     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

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/filter/SpanField.java b/apps/opik-backend/src/main/java/com/comet/opik/api/filter/SpanField.java
@@ -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),
     ;
 

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/filter/TraceField.java b/apps/opik-backend/src/main/java/com/comet/opik/api/filter/TraceField.java
@@ -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),
     ;
 

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemDAO.java
@@ -8,7 +8,7 @@
 import com.comet.opik.api.ScoreSource;
 import com.comet.opik.domain.filter.FilterQueryBuilder;
 import com.comet.opik.domain.filter.FilterStrategy;
-import com.comet.opik.infrastructure.db.TransactionTemplate;
+import com.comet.opik.infrastructure.db.TransactionTemplateAsync;
 import com.comet.opik.utils.JsonUtils;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.inject.ImplementedBy;
@@ -404,7 +404,7 @@ LEFT JOIN (
             ;
             """;
 
-    private final @NonNull TransactionTemplate asyncTemplate;
+    private final @NonNull TransactionTemplateAsync asyncTemplate;
     private final @NonNull FilterQueryBuilder filterQueryBuilder;
 
     @Override

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java
@@ -29,8 +29,8 @@
 
 import static com.comet.opik.api.Dataset.DatasetPage;
 import static com.comet.opik.domain.ExperimentItemDAO.ExperimentSummary;
-import static com.comet.opik.infrastructure.db.TransactionTemplate.READ_ONLY;
-import static com.comet.opik.infrastructure.db.TransactionTemplate.WRITE;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.WRITE;
 import static java.util.stream.Collectors.toMap;
 
 @ImplementedBy(DatasetServiceImpl.class)

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackDefinitionService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackDefinitionService.java
@@ -22,8 +22,8 @@
 import java.util.List;
 import java.util.UUID;
 
-import static com.comet.opik.infrastructure.db.TransactionTemplate.READ_ONLY;
-import static com.comet.opik.infrastructure.db.TransactionTemplate.WRITE;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.WRITE;
 
 @ImplementedBy(FeedbackDefinitionServiceImpl.class)
 public interface FeedbackDefinitionService {

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackScoreService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackScoreService.java
@@ -7,7 +7,7 @@
 import com.comet.opik.api.error.ErrorMessage;
 import com.comet.opik.api.error.IdentifierMismatchException;
 import com.comet.opik.infrastructure.auth.RequestContext;
-import com.comet.opik.infrastructure.db.TransactionTemplate;
+import com.comet.opik.infrastructure.db.TransactionTemplateAsync;
 import com.comet.opik.infrastructure.lock.LockService;
 import com.comet.opik.utils.WorkspaceUtils;
 import com.google.inject.ImplementedBy;
@@ -30,8 +30,8 @@
 import java.util.function.Function;
 
 import static com.comet.opik.domain.FeedbackScoreDAO.EntityType;
-import static com.comet.opik.infrastructure.db.TransactionTemplate.READ_ONLY;
-import static com.comet.opik.infrastructure.db.TransactionTemplate.WRITE;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.WRITE;
 import static java.util.stream.Collectors.groupingBy;
 import static java.util.stream.Collectors.toMap;
 
@@ -58,7 +58,7 @@ class FeedbackScoreServiceImpl implements FeedbackScoreService {
 
     private final @NonNull FeedbackScoreDAO dao;
     private final @NonNull ru.vyarus.guicey.jdbi3.tx.TransactionTemplate syncTemplate;
-    private final @NonNull TransactionTemplate asyncTemplate;
+    private final @NonNull TransactionTemplateAsync asyncTemplate;
     private final @NonNull IdGenerator idGenerator;
     private final @NonNull LockService lockService;
 

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
@@ -2,12 +2,14 @@
 
 import com.comet.opik.api.Page;
 import com.comet.opik.api.Project;
+import com.comet.opik.api.Project.ProjectPage;
 import com.comet.opik.api.ProjectCriteria;
 import com.comet.opik.api.ProjectUpdate;
 import com.comet.opik.api.error.CannotDeleteProjectException;
 import com.comet.opik.api.error.EntityAlreadyExistsException;
 import com.comet.opik.api.error.ErrorMessage;
 import com.comet.opik.infrastructure.auth.RequestContext;
+import com.comet.opik.infrastructure.db.TransactionTemplateAsync;
 import com.google.inject.ImplementedBy;
 import jakarta.inject.Inject;
 import jakarta.inject.Provider;
@@ -21,12 +23,16 @@
 import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate;
 
 import java.sql.SQLIntegrityConstraintViolationException;
+import java.time.Instant;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 
-import static com.comet.opik.infrastructure.db.TransactionTemplate.READ_ONLY;
-import static com.comet.opik.infrastructure.db.TransactionTemplate.WRITE;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY;
+import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.WRITE;
+import static java.util.stream.Collectors.toSet;
 
 @ImplementedBy(ProjectServiceImpl.class)
 public interface ProjectService {
@@ -58,10 +64,15 @@ public interface ProjectService {
 @RequiredArgsConstructor(onConstructor_ = @Inject)
 class ProjectServiceImpl implements ProjectService {
 
+    record ProjectRecordSet(List<Project> content, long total) {
+    }
+
     private static final String PROJECT_ALREADY_EXISTS = "Project already exists";
     private final @NonNull TransactionTemplate template;
     private final @NonNull IdGenerator idGenerator;
     private final @NonNull Provider<RequestContext> requestContext;
+    private final @NonNull TraceDAO traceDAO;
+    private final @NonNull TransactionTemplateAsync transactionTemplateAsync;
 
     private NotFoundException createNotFoundError() {
         String message = "Project not found";
@@ -89,14 +100,16 @@ private Project createProject(Project project, UUID projectId, String userName,
                 .build();
 
         try {
-            return template.inTransaction(WRITE, handle -> {
+            template.inTransaction(WRITE, handle -> {
 
                 var repository = handle.attach(ProjectDAO.class);
 
                 repository.save(workspaceId, newProject);
 
-                return repository.findById(projectId, workspaceId);
+                return newProject;
             });
+
+            return get(newProject.id(), workspaceId);
         } catch (UnableToExecuteStatementException e) {
             if (e.getCause() instanceof SQLIntegrityConstraintViolationException) {
                 throw newConflict();
@@ -117,7 +130,7 @@ public Project update(@NonNull UUID id, @NonNull ProjectUpdate projectUpdate) {
         String workspaceId = requestContext.get().getWorkspaceId();
 
         try {
-            return template.inTransaction(WRITE, handle -> {
+            template.inTransaction(WRITE, handle -> {
 
                 var repository = handle.attach(ProjectDAO.class);
 
@@ -130,9 +143,10 @@ public Project update(@NonNull UUID id, @NonNull ProjectUpdate projectUpdate) {
                         projectUpdate.description(),
                         userName);
 
-                return repository.findById(id, workspaceId);
+                return null;
             });
 
+            return get(id, workspaceId);
         } catch (UnableToExecuteStatementException e) {
             if (e.getCause() instanceof SQLIntegrityConstraintViolationException) {
                 throw newConflict();
@@ -151,12 +165,20 @@ public Project get(@NonNull UUID id) {
 
     @Override
     public Project get(@NonNull UUID id, @NonNull String workspaceId) {
-        return template.inTransaction(READ_ONLY, handle -> {
+        Project project = template.inTransaction(READ_ONLY, handle -> {
 
             var repository = handle.attach(ProjectDAO.class);
 
             return repository.fetch(id, workspaceId).orElseThrow(this::createNotFoundError);
         });
+
+        Map<UUID, Instant> lastUpdatedTraceAt = transactionTemplateAsync
+                .nonTransaction(connection -> traceDAO.getLastUpdatedTraceAt(Set.of(id), workspaceId, connection))
+                .block();
+
+        return project.toBuilder()
+                .lastUpdatedTraceAt(lastUpdatedTraceAt.get(id))
+                .build();
     }
 
     @Override
@@ -188,19 +210,37 @@ public void delete(@NonNull UUID id) {
 
     @Override
     public Page<Project> find(int page, int size, @NonNull ProjectCriteria criteria) {
+
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        return template.inTransaction(READ_ONLY, handle -> {
+        ProjectRecordSet projectRecordSet = template.inTransaction(READ_ONLY, handle -> {
 
             ProjectDAO repository = handle.attach(ProjectDAO.class);
 
             int offset = (page - 1) * size;
 
-            List<Project> projects = repository.find(size, offset, workspaceId, criteria.projectName());
-
-            return new Project.ProjectPage(page, projects.size(),
-                    repository.findCount(workspaceId, criteria.projectName()), projects);
+            return new ProjectRecordSet(
+                    repository.find(size, offset, workspaceId, criteria.projectName()),
+                    repository.findCount(workspaceId, criteria.projectName()));
         });
+
+        if (projectRecordSet.content().isEmpty()) {
+            return ProjectPage.empty(page);
+        }
+
+        Map<UUID, Instant> projectLastUpdatedTraceAtMap = transactionTemplateAsync.nonTransaction(connection -> {
+            Set<UUID> projectIds = projectRecordSet.content().stream().map(Project::id).collect(toSet());
+            return traceDAO.getLastUpdatedTraceAt(projectIds, workspaceId, connection);
+        }).block();
+
+        List<Project> projects = projectRecordSet.content()
+                .stream()
+                .map(project -> project.toBuilder()
+                        .lastUpdatedTraceAt(projectLastUpdatedTraceAtMap.get(project.id()))
+                        .build())
+                .toList();
+
+        return new ProjectPage(page, projects.size(), projectRecordSet.total(), projects);
     }
 
     @Override

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanService.java
@@ -262,11 +262,12 @@ public Mono<Long> create(@NonNull SpanBatch batch) {
         List<String> projectNames = batch.spans()
                 .stream()
                 .map(Span::projectName)
+                .map(WorkspaceUtils::getProjectName)
                 .distinct()
                 .toList();
 
         Mono<List<Span>> resolveProjects = Flux.fromIterable(projectNames)
-                .flatMap(this::resolveProject)
+                .flatMap(this::getOrCreateProject)
                 .collectList()
                 .map(projects -> bindSpanToProjectAndId(batch, projects))
                 .subscribeOn(Schedulers.boundedElastic());
@@ -293,7 +294,4 @@ private List<Span> bindSpanToProjectAndId(SpanBatch batch, List<Project> project
                 .toList();
     }
 
-    private Mono<Project> resolveProject(String projectName) {
-        return getOrCreateProject(WorkspaceUtils.getProjectName(projectName));
-    }
 }