feat(annotations): add feedback column to spans / traces tables with …

…all annotations (#4100) * feat(annotations): add feedback column to spans / traces tables with all annotations * fix column selector and spans table * add annotator to AnnotationTooltip kind section
Arize-ai · Aug 1, 2024 · 193b309 · 193b309
1 parent 2b313a0
commit 193b309
Show file tree

Hide file tree

Showing 13 changed files with 197 additions and 151 deletions.
diff --git a/app/src/components/annotation/AnnotationTooltip.tsx b/app/src/components/annotation/AnnotationTooltip.tsx
@@ -45,12 +45,14 @@ export function AnnotationTooltip({
             </Text>
             <Text color="inherit">{annotation.label || "--"}</Text>
           </Flex>
-          <Flex direction="row" justifyContent="space-between">
-            <Text weight="heavy" color="inherit">
-              kind
-            </Text>
-            <Text color="inherit">{annotation.annotatorKind}</Text>
-          </Flex>
+          {annotation.annotatorKind ? (
+            <Flex direction="row" justifyContent="space-between">
+              <Text weight="heavy" color="inherit">
+                annotator kind
+              </Text>
+              <Text color="inherit">{annotation.annotatorKind}</Text>
+            </Flex>
+          ) : null}
         </View>
         {annotation.explanation ? (
           <View paddingTop="size-50">

diff --git a/app/src/components/annotation/types.ts b/app/src/components/annotation/types.ts
@@ -3,5 +3,5 @@ export interface Annotation {
   label?: string | null;
   score?: number | null;
   explanation?: string | null;
-  annotatorKind: string;
+  annotatorKind?: string;
 }
diff --git a/app/src/pages/project/SpanColumnSelector.tsx b/app/src/pages/project/SpanColumnSelector.tsx
@@ -7,7 +7,7 @@ import { Dropdown, Flex, Icon, Icons, View } from "@arizeai/components";
 
 import { useTracingContext } from "@phoenix/contexts/TracingContext";
 
-import { SpanColumnSelector_evaluations$key } from "./__generated__/SpanColumnSelector_evaluations.graphql";
+import { SpanColumnSelector_annotations$key } from "./__generated__/SpanColumnSelector_annotations.graphql";
 
 const UN_HIDABLE_COLUMN_IDS = ["spanKind", "name"];
 
@@ -19,7 +19,7 @@ type SpanColumnSelectorProps = {
    */
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   columns: Column<any>[];
-  query: SpanColumnSelector_evaluations$key;
+  query: SpanColumnSelector_annotations$key;
 };
 
 export function SpanColumnSelector(props: SpanColumnSelectorProps) {
@@ -145,33 +145,33 @@ function ColumnSelectorMenu(props: SpanColumnSelectorProps) {
 function EvaluationColumnSelector({
   query,
 }: Pick<SpanColumnSelectorProps, "query">) {
-  const data = useFragment<SpanColumnSelector_evaluations$key>(
+  const data = useFragment<SpanColumnSelector_annotations$key>(
     graphql`
-      fragment SpanColumnSelector_evaluations on Project {
-        spanEvaluationNames
+      fragment SpanColumnSelector_annotations on Project {
+        spanAnnotationNames
       }
     `,
     query
   );
-  const evaluationVisibility = useTracingContext(
-    (state) => state.evaluationVisibility
+  const annotationColumnVisibility = useTracingContext(
+    (state) => state.annotationColumnVisibility
   );
-  const setEvaluationVisibility = useTracingContext(
-    (state) => state.setEvaluationVisibility
+  const setAnnotationColumnVisibility = useTracingContext(
+    (state) => state.setAnnotationColumnVisibility
   );
   const allVisible = useMemo(() => {
-    return data.spanEvaluationNames.every((name) => {
-      const stateValue = evaluationVisibility[name];
+    return data.spanAnnotationNames.every((name) => {
+      const stateValue = annotationColumnVisibility[name];
       return stateValue || false;
     });
-  }, [evaluationVisibility, data.spanEvaluationNames]);
+  }, [data.spanAnnotationNames, annotationColumnVisibility]);
 
   const onToggleEvaluations = useCallback(() => {
-    const newVisibilityState = data.spanEvaluationNames.reduce((acc, name) => {
+    const newVisibilityState = data.spanAnnotationNames.reduce((acc, name) => {
       return { ...acc, [name]: !allVisible };
     }, {});
-    setEvaluationVisibility(newVisibilityState);
-  }, [setEvaluationVisibility, allVisible, data.spanEvaluationNames]);
+    setAnnotationColumnVisibility(newVisibilityState);
+  }, [data.spanAnnotationNames, setAnnotationColumnVisibility, allVisible]);
   return (
     <section>
       <View
@@ -189,13 +189,13 @@ function EvaluationColumnSelector({
               checked={allVisible}
               onChange={onToggleEvaluations}
             />
-            evaluations
+            feedback
           </label>
         </div>
       </View>
       <ul>
-        {data.spanEvaluationNames.map((name) => {
-          const isVisible = evaluationVisibility[name];
+        {data.spanAnnotationNames.map((name) => {
+          const isVisible = annotationColumnVisibility[name] ?? false;
           return (
             <li key={name} css={columCheckboxItemCSS}>
               <label>
@@ -204,8 +204,8 @@ function EvaluationColumnSelector({
                   name={name}
                   checked={isVisible}
                   onChange={() => {
-                    setEvaluationVisibility({
-                      ...evaluationVisibility,
+                    setAnnotationColumnVisibility({
+                      ...annotationColumnVisibility,
                       [name]: !isVisible,
                     });
                   }}

diff --git a/app/src/pages/project/SpansTable.tsx b/app/src/pages/project/SpansTable.tsx
@@ -48,9 +48,9 @@ import { SpanFilterConditionField } from "./SpanFilterConditionField";
 import { SpanSelectionToolbar } from "./SpanSelectionToolbar";
 import { spansTableCSS } from "./styles";
 import {
+  ANNOTATIONS_COLUMN_PREFIX,
+  ANNOTATIONS_KEY_SEPARATOR,
   DEFAULT_SORT,
-  EVALS_COLUMN_PREFIX,
-  EVALS_KEY_SEPARATOR,
   getGqlSort,
 } from "./tableUtils";
 type SpansTableProps = {
@@ -82,7 +82,7 @@ export function SpansTable(props: SpansTableProps) {
           }
           filterCondition: { type: "String", defaultValue: null }
         ) {
-          ...SpanColumnSelector_evaluations
+          ...SpanColumnSelector_annotations
           spans(
             first: $first
             after: $after
@@ -112,10 +112,11 @@ export function SpansTable(props: SpansTableProps) {
                 output {
                   value: truncatedValue
                 }
-                spanEvaluations {
+                spanAnnotations {
                   name
                   label
                   score
+                  annotatorKind
                 }
                 documentRetrievalMetrics {
                   evaluationName
@@ -131,14 +132,14 @@ export function SpansTable(props: SpansTableProps) {
       props.project
     );
 
-  const evaluationVisibility = useTracingContext(
-    (state) => state.evaluationVisibility
+  const annotationColumnVisibility = useTracingContext(
+    (state) => state.annotationColumnVisibility
   );
-  const visibleEvaluationColumnNames = useMemo(() => {
-    return Object.keys(evaluationVisibility).filter(
-      (name) => evaluationVisibility[name]
+  const visibleAnnotationColumnNames = useMemo(() => {
+    return Object.keys(annotationColumnVisibility).filter(
+      (name) => annotationColumnVisibility[name]
     );
-  }, [evaluationVisibility]);
+  }, [annotationColumnVisibility]);
 
   const tableData = useMemo(() => {
     const tableData = data.spans.edges.map(({ span }) => span);
@@ -147,58 +148,54 @@ export function SpansTable(props: SpansTableProps) {
   }, [data]);
   type TableRow = (typeof tableData)[number];
 
-  const dynamicEvaluationColumns: ColumnDef<TableRow>[] =
-    visibleEvaluationColumnNames.map((name) => {
+  const dynamicAnnotationColumns: ColumnDef<TableRow>[] =
+    visibleAnnotationColumnNames.map((name) => {
       return {
         header: name,
         columns: [
           {
             header: `label`,
-            accessorKey: `${EVALS_COLUMN_PREFIX}${EVALS_KEY_SEPARATOR}label${EVALS_KEY_SEPARATOR}${name}`,
+            accessorKey: `${ANNOTATIONS_COLUMN_PREFIX}${ANNOTATIONS_KEY_SEPARATOR}label${ANNOTATIONS_KEY_SEPARATOR}${name}`,
             cell: ({ row }) => {
-              const evaluation = row.original.spanEvaluations.find(
-                (evaluation) => evaluation.name === name
+              const annotation = row.original.spanAnnotations.find(
+                (annotation) => annotation.name === name
               );
-              if (!evaluation) {
+              if (!annotation) {
                 return null;
               }
-              return evaluation.label;
+              return annotation.label;
             },
           } as ColumnDef<TableRow>,
           {
             header: `score`,
-            accessorKey: `${EVALS_COLUMN_PREFIX}${EVALS_KEY_SEPARATOR}score${EVALS_KEY_SEPARATOR}${name}`,
+            accessorKey: `${ANNOTATIONS_COLUMN_PREFIX}${ANNOTATIONS_KEY_SEPARATOR}score${ANNOTATIONS_KEY_SEPARATOR}${name}`,
             cell: ({ row }) => {
-              const evaluation = row.original.spanEvaluations.find(
-                (evaluation) => evaluation.name === name
+              const annotation = row.original.spanAnnotations.find(
+                (annotation) => annotation.name === name
               );
-              if (!evaluation) {
+              if (!annotation) {
                 return null;
               }
-              return evaluation.score;
+              return annotation.score;
             },
           } as ColumnDef<TableRow>,
         ],
       };
     });
 
-  const evaluationColumns: ColumnDef<TableRow>[] = [
+  const annotationColumns: ColumnDef<TableRow>[] = [
     {
-      header: "evaluations",
-      accessorKey: "spanEvaluations",
+      header: "feedback",
+      accessorKey: "spanAnnotations",
       enableSorting: false,
 
       cell: ({ row }) => {
         return (
           <Flex direction="row" gap="size-50" wrap="wrap">
-            {row.original.spanEvaluations.map((evaluation) => {
-              const annotation = {
-                ...evaluation,
-                annotatorKind: "LLM",
-              };
+            {row.original.spanAnnotations.map((annotation) => {
               return (
                 <AnnotationTooltip
-                  key={evaluation.name}
+                  key={annotation.name}
                   annotation={annotation}
                 >
                   <AnnotationLabel
@@ -236,7 +233,7 @@ export function SpansTable(props: SpansTableProps) {
         );
       },
     },
-    ...dynamicEvaluationColumns,
+    ...dynamicAnnotationColumns,
   ];
   const columns: ColumnDef<TableRow>[] = [
     {
@@ -303,7 +300,7 @@ export function SpansTable(props: SpansTableProps) {
       cell: TextCell,
       enableSorting: false,
     },
-    ...evaluationColumns, // TODO: consider hiding this column if there are no evals. For now we want people to know that there are evals
+    ...annotationColumns, // TODO: consider hiding this column if there are no evals. For now we want people to know that there are evals
     {
       header: "start time",
       accessorKey: "startTime",