feat(experiments): add experiment run filter to compare experiments p…

…age (#5738)
Arize-ai · Jan 9, 2025 · 0bf194d · 0bf194d
1 parent 4aa4ec2
commit 0bf194d
Show file tree

Hide file tree

Showing 16 changed files with 3,402 additions and 117 deletions.
diff --git a/app/schema.graphql b/app/schema.graphql
@@ -1388,7 +1388,8 @@ type Query {
   projectsLastUpdatedAt: DateTime
   datasets(first: Int = 50, last: Int, after: String, before: String, sort: DatasetSort): DatasetConnection!
   datasetsLastUpdatedAt: DateTime
-  compareExperiments(experimentIds: [GlobalID!]!): [ExperimentComparison!]!
+  compareExperiments(experimentIds: [GlobalID!]!, filterCondition: String): [ExperimentComparison!]!
+  validateExperimentRunFilterCondition(condition: String!, experimentIds: [GlobalID!]!): ValidationResult!
   functionality: Functionality!
   model: Model!
   node(id: GlobalID!): Node!

diff --git a/app/src/pages/experiment/ExperimentComparePage.tsx b/app/src/pages/experiment/ExperimentComparePage.tsx
@@ -9,6 +9,7 @@ import { Flex, Heading, Loading, View } from "@phoenix/components";
 import { experimentCompareLoaderQuery$data } from "./__generated__/experimentCompareLoaderQuery.graphql";
 import { ExperimentCompareTable } from "./ExperimentCompareTable";
 import { ExperimentMultiSelector } from "./ExperimentMultiSelector";
+import { ExperimentRunFilterConditionProvider } from "./ExperimentRunFilterConditionContext";
 
 export function ExperimentComparePage() {
   const data = useLoaderData() as experimentCompareLoaderQuery$data;
@@ -62,13 +63,15 @@ export function ExperimentComparePage() {
         </Flex>
       </View>
       {experimentIdsSelected ? (
-        <Suspense fallback={<Loading />}>
-          <ExperimentCompareTable
-            datasetId={data.dataset.id}
-            experimentIds={experimentIds}
-            displayFullText={displayFullText}
-          />
-        </Suspense>
+        <ExperimentRunFilterConditionProvider>
+          <Suspense fallback={<Loading />}>
+            <ExperimentCompareTable
+              datasetId={data.dataset.id}
+              experimentIds={experimentIds}
+              displayFullText={displayFullText}
+            />
+          </Suspense>
+        </ExperimentRunFilterConditionProvider>
       ) : (
         <View padding="size-200">
           <Alert variant="info" title="No Experiment Selected">

diff --git a/app/src/pages/experiment/ExperimentCompareTable.tsx b/app/src/pages/experiment/ExperimentCompareTable.tsx
@@ -65,6 +65,7 @@ import {
   ExperimentCompareTableQuery,
   ExperimentCompareTableQuery$data,
 } from "./__generated__/ExperimentCompareTableQuery.graphql";
+import { ExperimentRunFilterConditionField } from "./ExperimentRunFilterConditionField";
 
 type ExampleCompareTableProps = {
   datasetId: string;
@@ -131,13 +132,19 @@ const annotationTooltipExtraCSS = css`
 
 export function ExperimentCompareTable(props: ExampleCompareTableProps) {
   const { datasetId, experimentIds, displayFullText } = props;
+  const [filterCondition, setFilterCondition] = useState("");
+
   const data = useLazyLoadQuery<ExperimentCompareTableQuery>(
     graphql`
       query ExperimentCompareTableQuery(
         $experimentIds: [GlobalID!]!
         $datasetId: GlobalID!
+        $filterCondition: String
       ) {
-        comparisons: compareExperiments(experimentIds: $experimentIds) {
+        comparisons: compareExperiments(
+          experimentIds: $experimentIds
+          filterCondition: $filterCondition
+        ) {
           example {
             id
             revision {
@@ -198,6 +205,7 @@ export function ExperimentCompareTable(props: ExampleCompareTableProps) {
     {
       experimentIds,
       datasetId,
+      filterCondition,
     }
   );
   const experimentInfoById = useMemo(() => {
@@ -440,64 +448,79 @@ export function ExperimentCompareTable(props: ExampleCompareTableProps) {
   // Make sure the table is at least 1280px wide
 
   return (
-    <div css={tableWrapCSS}>
-      <table
-        css={css(tableCSS, borderedTableCSS)}
-        style={{
-          ...columnSizeVars,
-          width: table.getTotalSize(),
-          minWidth: "100%",
-        }}
+    <Flex direction="column" height="100%">
+      <View
+        paddingTop="size-100"
+        paddingBottom="size-100"
+        paddingStart="size-200"
+        paddingEnd="size-200"
+        borderBottomColor="grey-300"
+        borderBottomWidth="thin"
+        flex="none"
       >
-        <thead>
-          {table.getHeaderGroups().map((headerGroup) => (
-            <tr key={headerGroup.id}>
-              {headerGroup.headers.map((header) => (
-                <th
-                  key={header.id}
-                  style={{
-                    width: `calc(var(--header-${header?.id}-size) * 1px)`,
-                  }}
-                >
-                  <div>
-                    {flexRender(
-                      header.column.columnDef.header,
-                      header.getContext()
-                    )}
-                  </div>
-                  <div
-                    {...{
-                      onMouseDown: header.getResizeHandler(),
-                      onTouchStart: header.getResizeHandler(),
-                      className: `resizer ${
-                        header.column.getIsResizing() ? "isResizing" : ""
-                      }`,
+        <ExperimentRunFilterConditionField
+          onValidCondition={setFilterCondition}
+        />
+      </View>
+      <div css={tableWrapCSS}>
+        <table
+          css={css(tableCSS, borderedTableCSS)}
+          style={{
+            ...columnSizeVars,
+            width: table.getTotalSize(),
+            minWidth: "100%",
+          }}
+        >
+          <thead>
+            {table.getHeaderGroups().map((headerGroup) => (
+              <tr key={headerGroup.id}>
+                {headerGroup.headers.map((header) => (
+                  <th
+                    key={header.id}
+                    style={{
+                      width: `calc(var(--header-${header?.id}-size) * 1px)`,
                     }}
-                  />
-                </th>
-              ))}
-            </tr>
-          ))}
-        </thead>
-        {isEmpty ? (
-          <TableEmpty />
-        ) : /* When resizing any column we will render this special memoized version of our table body */
-        table.getState().columnSizingInfo.isResizingColumn ? (
-          <MemoizedTableBody table={table} />
-        ) : (
-          <TableBody table={table} />
-        )}
-      </table>
-      <DialogContainer
-        isDismissable
-        type="slideOver"
-        onDismiss={() => {
-          setDialog(null);
-        }}
-      >
-        {dialog}
-      </DialogContainer>
-    </div>
+                  >
+                    <div>
+                      {flexRender(
+                        header.column.columnDef.header,
+                        header.getContext()
+                      )}
+                    </div>
+                    <div
+                      {...{
+                        onMouseDown: header.getResizeHandler(),
+                        onTouchStart: header.getResizeHandler(),
+                        className: `resizer ${
+                          header.column.getIsResizing() ? "isResizing" : ""
+                        }`,
+                      }}
+                    />
+                  </th>
+                ))}
+              </tr>
+            ))}
+          </thead>
+          {isEmpty ? (
+            <TableEmpty />
+          ) : /* When resizing any column we will render this special memoized version of our table body */
+          table.getState().columnSizingInfo.isResizingColumn ? (
+            <MemoizedTableBody table={table} />
+          ) : (
+            <TableBody table={table} />
+          )}
+        </table>
+        <DialogContainer
+          isDismissable
+          type="slideOver"
+          onDismiss={() => {
+            setDialog(null);
+          }}
+        >
+          {dialog}
+        </DialogContainer>
+      </div>
+    </Flex>
   );
 }
 

diff --git a/app/src/pages/experiment/ExperimentRunFilterConditionContext.tsx b/app/src/pages/experiment/ExperimentRunFilterConditionContext.tsx
@@ -0,0 +1,56 @@
+import React, {
+  createContext,
+  PropsWithChildren,
+  startTransition,
+  useCallback,
+  useContext,
+  useState,
+} from "react";
+
+export type ExperimentRunFilterConditionContextType = {
+  filterCondition: string;
+  setFilterCondition: (condition: string) => void;
+  appendFilterCondition: (condition: string) => void;
+};
+
+export const ExperimentRunFilterConditionContext =
+  createContext<ExperimentRunFilterConditionContextType | null>(null);
+
+export function useExperimentRunFilterCondition() {
+  const context = useContext(ExperimentRunFilterConditionContext);
+  if (context === null) {
+    throw new Error(
+      "useExperimentRunFilterCondition must be used within a ExperimentRunFilterConditionProvider"
+    );
+  }
+  return context;
+}
+
+export function ExperimentRunFilterConditionProvider(props: PropsWithChildren) {
+  const [filterCondition, _setFilterCondition] = useState<string>("");
+  const setFilterCondition = useCallback((condition: string) => {
+    startTransition(() => {
+      _setFilterCondition(condition);
+    });
+  }, []);
+  const appendFilterCondition = useCallback(
+    (condition: string) => {
+      startTransition(() => {
+        if (filterCondition.length > 0) {
+          _setFilterCondition(filterCondition + " and " + condition);
+        } else {
+          _setFilterCondition(condition);
+        }
+      });
+    },
+    [filterCondition]
+  );
+
+  return (
+    <ExperimentRunFilterConditionContext.Provider
+      value={{ filterCondition, setFilterCondition, appendFilterCondition }}
+    >
+      {props.children}
+    </ExperimentRunFilterConditionContext.Provider>
+  );
+}