Merge branch 'thiagohora/OPIK-218_remove_limitations_in_dataset_items…

…' of https://github.com/comet-ml/opik into thiagohora/OPIK-218_remove_limitations_in_dataset_items
comet-ml · Oct 11, 2024 · c8edc9e · c8edc9e
2 parents 3f09312 + f416602
commit c8edc9e
Show file tree

Hide file tree

Showing 8 changed files with 119 additions and 43 deletions.
diff --git a/.github/workflows/sdk-e2e-tests.yaml b/.github/workflows/sdk-e2e-tests.yaml
@@ -38,7 +38,7 @@ jobs:
             cd deployment/docker-compose
             docker compose up -d --build
         
-        - name: Check Opik server avialability
+        - name: Check Opik server availability
           shell: bash
           run: |
             chmod +x ${{ github.workspace }}/tests_end_to_end/installer/*.sh

diff --git a/...opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java b/...opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java
@@ -76,6 +76,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -2834,39 +2835,29 @@ private void getItemAndAssert(DatasetItem expectedDatasetItem, String workspaceN
 
     private static DatasetItem mergeInputMap(DatasetItem expectedDatasetItem,
             Map<String, DatasetItemInputValue<?>> inputData) {
-        if (expectedDatasetItem.expectedOutput() != null) {
 
-            Map<String, DatasetItemInputValue.JsonValue> oldColumns = Map.of(
-                    "input", new DatasetItemInputValue.JsonValue(expectedDatasetItem.input()),
-                    "expected_output", new DatasetItemInputValue.JsonValue(expectedDatasetItem.expectedOutput()));
+        Map<String, DatasetItemInputValue<?>> newMap = new HashMap<>();
 
-            Map<String, DatasetItemInputValue<?>> mergedMap = Stream
-                    .concat(inputData.entrySet().stream(), oldColumns.entrySet().stream())
-                    .collect(toMap(
-                            Map.Entry::getKey,
-                            Map.Entry::getValue,
-                            (v1, v2) -> v2 // In case of conflict, use the value from map2
-                    ));
+        if (expectedDatasetItem.expectedOutput() != null) {
+            newMap.put("expected_output", new DatasetItemInputValue.JsonValue(expectedDatasetItem.expectedOutput()));
+        }
 
-            expectedDatasetItem = expectedDatasetItem.toBuilder()
-                    .inputData(mergedMap)
-                    .build();
-        } else {
-            Map<String, DatasetItemInputValue.JsonValue> oldColumns = Map.of("input",
-                    new DatasetItemInputValue.JsonValue(expectedDatasetItem.input()));
+        if (expectedDatasetItem.input() != null) {
+            newMap.put("input", new DatasetItemInputValue.JsonValue(expectedDatasetItem.input()));
+        }
 
-            Map<String, DatasetItemInputValue<?>> mergedMap = Stream
-                    .concat(inputData.entrySet().stream(), oldColumns.entrySet().stream())
-                    .collect(toMap(
-                            Map.Entry::getKey,
-                            Map.Entry::getValue,
-                            (v1, v2) -> v2 // In case of conflict, use the value from map2
-                    ));
+        Map<String, DatasetItemInputValue<?>> mergedMap = Stream
+                .concat(inputData.entrySet().stream(), newMap.entrySet().stream())
+                .collect(toMap(
+                        Map.Entry::getKey,
+                        Map.Entry::getValue,
+                        (v1, v2) -> v2 // In case of conflict, use the value from map2
+                ));
+
+        expectedDatasetItem = expectedDatasetItem.toBuilder()
+                .inputData(mergedMap)
+                .build();
 
-            expectedDatasetItem = expectedDatasetItem.toBuilder()
-                    .inputData(mergedMap)
-                    .build();
-        }
         return expectedDatasetItem;
     }
 

diff --git a/apps/opik-frontend/src/api/datasets/useCompareExperimentsList.ts b/apps/opik-frontend/src/api/datasets/useCompareExperimentsList.ts
@@ -1,12 +1,15 @@
 import { QueryFunctionContext, useQuery } from "@tanstack/react-query";
 import api, { DATASETS_REST_ENDPOINT, QueryConfig } from "@/api/api";
 import { ExperimentsCompare } from "@/types/datasets";
+import { Filters } from "@/types/filters";
+import { processFilters } from "@/lib/filters";
 
 type UseCompareExperimentsListParams = {
   workspaceName: string;
   datasetId: string;
   experimentsIds: string[];
   search?: string;
+  filters?: Filters;
   page: number;
   size: number;
 };
@@ -23,6 +26,7 @@ const getCompareExperimentsList = async (
     datasetId,
     experimentsIds,
     search,
+    filters,
     size,
     page,
   }: UseCompareExperimentsListParams,
@@ -34,6 +38,7 @@ const getCompareExperimentsList = async (
       params: {
         workspace_name: workspaceName,
         experiment_ids: JSON.stringify(experimentsIds),
+        ...processFilters(filters),
         ...(search && { name: search }),
         size,
         page,

diff --git a/...end/src/components/pages/CompareExperimentsPage/ExperimentItemsTab/ExperimentItemsTab.tsx b/...end/src/components/pages/CompareExperimentsPage/ExperimentItemsTab/ExperimentItemsTab.tsx
@@ -2,7 +2,12 @@ import React, { useCallback, useMemo } from "react";
 import isObject from "lodash/isObject";
 import findIndex from "lodash/findIndex";
 import find from "lodash/find";
-import { NumberParam, StringParam, useQueryParam } from "use-query-params";
+import {
+  JsonParam,
+  NumberParam,
+  StringParam,
+  useQueryParam,
+} from "use-query-params";
 import { keepPreviousData } from "@tanstack/react-query";
 import useLocalStorageState from "use-local-storage-state";
 
@@ -24,6 +29,7 @@ import CompareExperimentAddHeader from "@/components/pages/CompareExperimentsPag
 import TraceDetailsPanel from "@/components/shared/TraceDetailsPanel/TraceDetailsPanel";
 import CompareExperimentsPanel from "@/components/pages/CompareExperimentsPage/CompareExperimentsPanel/CompareExperimentsPanel";
 import ColumnsButton from "@/components/shared/ColumnsButton/ColumnsButton";
+import FiltersButton from "@/components/shared/FiltersButton/FiltersButton";
 import Loader from "@/components/shared/Loader/Loader";
 import useCompareExperimentsList from "@/api/datasets/useCompareExperimentsList";
 import useAppStore from "@/store/AppStore";
@@ -49,13 +55,7 @@ const SELECTED_COLUMNS_KEY = "compare-experiments-selected-columns";
 const COLUMNS_WIDTH_KEY = "compare-experiments-columns-width";
 const COLUMNS_ORDER_KEY = "compare-experiments-columns-order";
 
-export const DEFAULT_COLUMNS: ColumnData<ExperimentsCompare>[] = [
-  {
-    id: "id",
-    label: "Item ID",
-    type: COLUMN_TYPE.string,
-    cell: IdCell as never,
-  },
+export const SHARED_COLUMNS: ColumnData<ExperimentsCompare>[] = [
   {
     id: "input",
     label: "Input",
@@ -90,6 +90,16 @@ export const DEFAULT_COLUMNS: ColumnData<ExperimentsCompare>[] = [
         : row.metadata || "",
     cell: CodeCell as never,
   },
+];
+
+export const DEFAULT_COLUMNS: ColumnData<ExperimentsCompare>[] = [
+  {
+    id: "id",
+    label: "Item ID",
+    type: COLUMN_TYPE.string,
+    cell: IdCell as never,
+  },
+  ...SHARED_COLUMNS,
   {
     id: "created_at",
     label: "Created",
@@ -98,6 +108,20 @@ export const DEFAULT_COLUMNS: ColumnData<ExperimentsCompare>[] = [
   },
 ];
 
+export const FILTER_COLUMNS: ColumnData<ExperimentsCompare>[] = [
+  ...SHARED_COLUMNS,
+  {
+    id: "feedback_scores",
+    label: "Feedback scores",
+    type: COLUMN_TYPE.numberDictionary,
+  },
+  {
+    id: "output",
+    label: "Output",
+    type: COLUMN_TYPE.string,
+  },
+];
+
 export const DEFAULT_SELECTED_COLUMNS: string[] = ["id", "input"];
 
 export type ExperimentItemsTabProps = {
@@ -140,6 +164,10 @@ const ExperimentItemsTab: React.FunctionComponent<ExperimentItemsTabProps> = ({
     },
   );
 
+  const [filters = [], setFilters] = useQueryParam("filters", JsonParam, {
+    updateType: "replaceIn",
+  });
+
   const [columnsWidth, setColumnsWidth] = useLocalStorageState<
     Record<string, number>
   >(COLUMNS_WIDTH_KEY, {
@@ -210,6 +238,7 @@ const ExperimentItemsTab: React.FunctionComponent<ExperimentItemsTabProps> = ({
       workspaceName,
       datasetId,
       experimentsIds,
+      filters,
       page: page as number,
       size: size as number,
     },
@@ -264,7 +293,13 @@ const ExperimentItemsTab: React.FunctionComponent<ExperimentItemsTabProps> = ({
   return (
     <div>
       <div className="mb-6 flex items-center justify-between gap-8">
-        <div className="flex items-center gap-2"></div>
+        <div className="flex items-center gap-2">
+          <FiltersButton
+            columns={FILTER_COLUMNS}
+            filters={filters}
+            onChange={setFilters}
+          />
+        </div>
         <div className="flex items-center gap-2">
           <DataTableRowHeightSelector
             type={height as ROW_HEIGHT}

diff --git a/sdks/python/src/opik/integrations/openai/opik_tracker.py b/sdks/python/src/opik/integrations/openai/opik_tracker.py
@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Optional, Union
 
 import openai
 
@@ -7,6 +7,7 @@
 
 def track_openai(
     openai_client: Union[openai.OpenAI, openai.AsyncOpenAI],
+    project_name: Optional[str] = None,
 ) -> Union[openai.OpenAI, openai.AsyncOpenAI]:
     """Adds Opik tracking to an OpenAI client.
 
@@ -15,6 +16,7 @@ def track_openai(
 
     Args:
         openai_client: An instance of OpenAI or AsyncOpenAI client.
+        project_name: The name of the project to log data.
 
     Returns:
         The modified OpenAI client with Opik tracking enabled.
@@ -24,6 +26,7 @@ def track_openai(
         type="llm",
         name="chat_completion_create",
         generations_aggregator=chunks_aggregator.aggregate,
+        project_name=project_name,
     )
     openai_client.chat.completions.create = wrapper(
         openai_client.chat.completions.create