feat(annotations): refetch annotations on annotation changes

app/schema.graphql

@@ -1194,6 +1194,7 @@ type SpanAnnotation implements Node & Annotation {
 
 type SpanAnnotationMutationPayload {
   spanAnnotations: [SpanAnnotation!]!
+  query: Query!
 }
 
 type SpanAsExampleRevision implements ExampleRevision {
@@ -1339,6 +1340,7 @@ type TraceAnnotation implements Node {
 
 type TraceAnnotationMutationPayload {
   traceAnnotations: [TraceAnnotation!]!
+  query: Query!
 }
 
 type TraceEvaluation implements Annotation {

app/src/pages/project/DatasetSelectorPopoverContent.tsx

@@ -16,6 +16,7 @@ import {
 } from "@arizeai/components";
 
 import { DatasetSelectorPopoverContent_datasets$key } from "./__generated__/DatasetSelectorPopoverContent_datasets.graphql";
+import { DatasetSelectorPopoverContentDatasetsQuery } from "./__generated__/DatasetSelectorPopoverContentDatasetsQuery.graphql";
 import { DatasetSelectorPopoverContentQuery } from "./__generated__/DatasetSelectorPopoverContentQuery.graphql";
 
 export type DatasetSelectorPopoverContentProps = {
@@ -83,7 +84,10 @@ function DatasetsListBox(props: {
   onDatasetSelected: (datasetId: string) => void;
 }) {
   const { search, onDatasetSelected } = props;
-  const [data] = useRefetchableFragment(
+  const [data] = useRefetchableFragment<
+    DatasetSelectorPopoverContentDatasetsQuery,
+    DatasetSelectorPopoverContent_datasets$key
+  >(
     graphql`
       fragment DatasetSelectorPopoverContent_datasets on Query
       @refetchable(queryName: "DatasetSelectorPopoverContentDatasetsQuery") {

app/src/pages/trace/EditSpanAnnotationsDialog.tsx

@@ -1,5 +1,5 @@
 import React, { Suspense, useMemo, useState } from "react";
-import { graphql, useLazyLoadQuery } from "react-relay";
+import { graphql, useLazyLoadQuery, useRefetchableFragment } from "react-relay";
 import { css } from "@emotion/react";
 
 import {
@@ -22,12 +22,14 @@ import { Empty } from "@phoenix/components/Empty";
 import { useNotifySuccess } from "@phoenix/contexts";
 import { formatFloat } from "@phoenix/utils/numberFormatUtils";
 
-import { EditSpanAnnotationsDialogNewAnnotationQuery } from "./__generated__/EditSpanAnnotationsDialogNewAnnotationQuery.graphql";
 import {
   AnnotatorKind,
-  EditSpanAnnotationsDialogQuery,
-  EditSpanAnnotationsDialogQuery$data,
-} from "./__generated__/EditSpanAnnotationsDialogQuery.graphql";
+  EditSpanAnnotationsDialog_spanAnnotations$data,
+  EditSpanAnnotationsDialog_spanAnnotations$key,
+} from "./__generated__/EditSpanAnnotationsDialog_spanAnnotations.graphql";
+import { EditSpanAnnotationsDialogNewAnnotationQuery } from "./__generated__/EditSpanAnnotationsDialogNewAnnotationQuery.graphql";
+import { EditSpanAnnotationsDialogQuery } from "./__generated__/EditSpanAnnotationsDialogQuery.graphql";
+import { EditSpanAnnotationsDialogSpanAnnotationsQuery } from "./__generated__/EditSpanAnnotationsDialogSpanAnnotationsQuery.graphql";
 import { NewSpanAnnotationForm } from "./NewSpanAnnotationForm";
 import { SpanAnnotationActionMenu } from "./SpanAnnotationActionMenu";
 import { SpanAnnotationForm } from "./SpanAnnotationForm";
@@ -43,7 +45,6 @@ export function EditSpanAnnotationsDialog(
   const [newAnnotationName, setNewAnnotationName] = useState<string | null>(
     null
   );
-  const [fetchKey, setFetchKey] = useState(0);
   const notifySuccess = useNotifySuccess();
   return (
     <Dialog
@@ -76,7 +77,6 @@ export function EditSpanAnnotationsDialog(
               }}
               onCreated={() => {
                 setNewAnnotationName(null);
-                setFetchKey((key) => key + 1);
                 notifySuccess({
                   title: `New Span Annotation`,
                   message: `Annotation ${newAnnotationName} has been created.`,
@@ -86,7 +86,7 @@ export function EditSpanAnnotationsDialog(
           </View>
         )}
         <Suspense>
-          <EditSpanAnnotations {...props} fetchKey={fetchKey} />
+          <EditSpanAnnotations {...props} />
         </Suspense>
       </div>
     </Dialog>
@@ -166,36 +166,52 @@ function NewAnnotationPopover(props: NewAnnotationPopoverProps) {
     </Card>
   );
 }
-type EditSpanAnnotationsProps = EditSpanAnnotationsDialogProps & {
-  /**
-   * Key used to force a refetch of the annotations
-   */
-  fetchKey: number;
-};
+type EditSpanAnnotationsProps = EditSpanAnnotationsDialogProps;
+
 function EditSpanAnnotations(props: EditSpanAnnotationsProps) {
-  const { fetchKey } = props;
   const data = useLazyLoadQuery<EditSpanAnnotationsDialogQuery>(
     graphql`
       query EditSpanAnnotationsDialogQuery($spanId: GlobalID!) {
         span: node(id: $spanId) {
           id
           ... on Span {
-            spanAnnotations {
-              id
-              name
-              annotatorKind
-              score
-              label
-              explanation
-            }
+            ...EditSpanAnnotationsDialog_spanAnnotations
           }
         }
       }
     `,
     { spanId: props.spanNodeId },
-    { fetchKey, fetchPolicy: "store-and-network" }
+    { fetchPolicy: "store-and-network" }
+  );
+  return <SpanAnnotationsList span={data.span} />;
+}
+
+function SpanAnnotationsList(props: {
+  span: EditSpanAnnotationsDialog_spanAnnotations$key;
+}) {
+  const { span } = props;
+  const [data] = useRefetchableFragment<
+    EditSpanAnnotationsDialogSpanAnnotationsQuery,
+    EditSpanAnnotationsDialog_spanAnnotations$key
+  >(
+    graphql`
+      fragment EditSpanAnnotationsDialog_spanAnnotations on Span
+      @refetchable(queryName: "EditSpanAnnotationsDialogSpanAnnotationsQuery") {
+        id
+        spanAnnotations {
+          id
+          name
+          annotatorKind
+          score
+          label
+          explanation
+        }
+      }
+    `,
+    span
   );
-  const annotations = data.span.spanAnnotations || [];
+
+  const annotations = data.spanAnnotations || [];
   const hasAnnotations = annotations.length > 0;
   return (
     <div>
@@ -211,7 +227,7 @@ function EditSpanAnnotations(props: EditSpanAnnotationsProps) {
       >
         {annotations.map((annotation, idx) => (
           <li key={idx}>
-            <SpanAnnotationCard annotation={annotation} />
+            <SpanAnnotationCard annotation={annotation} spanNodeId={data.id} />
           </li>
         ))}
       </ul>
@@ -263,11 +279,14 @@ function NewSpanAnnotationCard(props: {
 }
 
 type Annotation = NonNullable<
-  EditSpanAnnotationsDialogQuery$data["span"]["spanAnnotations"]
+  EditSpanAnnotationsDialog_spanAnnotations$data["spanAnnotations"]
 >[number];
 
-function SpanAnnotationCard(props: { annotation: Annotation }) {
-  const { annotation } = props;
+function SpanAnnotationCard(props: {
+  annotation: Annotation;
+  spanNodeId: string;
+}) {
+  const { annotation, spanNodeId } = props;
   const [error, setError] = useState<Error | null>(null);
   const notifySuccess = useNotifySuccess();
   return (
@@ -287,6 +306,7 @@ function SpanAnnotationCard(props: { annotation: Annotation }) {
           <AnnotatorKindLabel kind={annotation.annotatorKind} />
           <SpanAnnotationActionMenu
             annotationId={annotation.id}
+            spanNodeId={spanNodeId}
             annotationName={annotation.name}
             onSpanAnnotationDelete={() => {
               notifySuccess({
@@ -412,8 +432,9 @@ function NewAnnotationPopoverContent(props: {
               if (keys === "all" || keys.size === 0) {
                 return;
               }
-              const name = keys.entries().next().value[0];
-              setNewName(name);
+              const nameKey = keys.values().next().value;
+              const name = nameKey as string;
+              setNewName(name || "");
             }}
             disabledKeys={existingAnnotationNames}
           >

app/src/pages/trace/NewSpanAnnotationForm.tsx

@@ -19,11 +19,15 @@ export function NewSpanAnnotationForm(props: NewSpanAnnotationFormProps) {
     graphql`
       mutation NewSpanAnnotationFormMutation(
         $input: CreateSpanAnnotationInput!
+        $spanId: GlobalID!
       ) {
         createSpanAnnotations(input: [$input]) {
-          spanAnnotations {
-            id
-            name
+          query {
+            node(id: $spanId) {
+              ... on Span {
+                ...EditSpanAnnotationsDialog_spanAnnotations
+              }
+            }
           }
         }
       }
@@ -37,6 +41,7 @@ export function NewSpanAnnotationForm(props: NewSpanAnnotationFormProps) {
           annotatorKind: "HUMAN",
           ...data,
         },
+        spanId: spanNodeId,
       },
       onCompleted: () => {
         onCreated();

app/src/pages/trace/SpanAnnotationActionMenu.tsx

@@ -20,6 +20,7 @@ import { AnnotationActionMenu } from "./AnnotationActionMenu";
 
 type SpanAnnotationActionMenuProps = {
   annotationId: string;
+  spanNodeId: string;
   annotationName: string;
   onSpanAnnotationDelete: () => void;
   onSpanAnnotationDeleteError: (error: Error) => void;
@@ -28,6 +29,7 @@ type SpanAnnotationActionMenuProps = {
 export function SpanAnnotationActionMenu(props: SpanAnnotationActionMenuProps) {
   const {
     annotationId,
+    spanNodeId,
     annotationName,
     onSpanAnnotationDelete,
     onSpanAnnotationDeleteError,
@@ -37,9 +39,16 @@ export function SpanAnnotationActionMenu(props: SpanAnnotationActionMenuProps) {
     useMutation<SpanAnnotationActionMenuDeleteMutation>(graphql`
       mutation SpanAnnotationActionMenuDeleteMutation(
         $annotationId: GlobalID!
+        $spanId: GlobalID!
       ) {
         deleteSpanAnnotations(input: { annotationIds: [$annotationId] }) {
-          __typename
+          query {
+            node(id: $spanId) {
+              ... on Span {
+                ...EditSpanAnnotationsDialog_spanAnnotations
+              }
+            }
+          }
         }
       }
     `);
@@ -49,6 +58,7 @@ export function SpanAnnotationActionMenu(props: SpanAnnotationActionMenuProps) {
       commitDelete({
         variables: {
           annotationId,
+          spanId: spanNodeId,
         },
         onCompleted: () => {
           onSpanAnnotationDelete();
@@ -61,6 +71,7 @@ export function SpanAnnotationActionMenu(props: SpanAnnotationActionMenuProps) {
   }, [
     commitDelete,
     annotationId,
+    spanNodeId,
     onSpanAnnotationDelete,
     onSpanAnnotationDeleteError,
   ]);