From 3fa8ebf23d6f87ad2b45b85fcc71b36f7b61ad17 Mon Sep 17 00:00:00 2001
From: Griffin Tarpenning <griffin.tarpenning@wandb.com>
Date: Tue, 19 Nov 2024 14:53:30 -0800
Subject: [PATCH] chore(weave): human annotation labeling review comments
 (#3020)

---
 .../StructuredFeedback/FeedbackSidebar.tsx    | 41 +++++++++++++++++-
 .../StructuredFeedback/HumanAnnotation.tsx    | 39 +++++++++++------
 .../humanAnnotationTypes.ts                   |  8 +++-
 .../StructuredFeedback/tsHumanFeedback.ts     | 35 ++++++++-------
 .../Home/Browse3/pages/CallPage/CallPage.tsx  | 21 +++++----
 .../pages/CallPage/PaginationControls.tsx     | 41 ++++++++++++++----
 .../pages/CallsPage/callsTableColumns.tsx     |  3 ++
 .../ScorersPage/AnnotationScorerForm.tsx      | 21 ++++++---
 .../pages/ScorersPage/NewScorerDrawer.tsx     |  1 +
 .../src/integrations/analytics/viewEvents.ts  | 43 +++++++++++++++++++
 10 files changed, 194 insertions(+), 59 deletions(-)

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/FeedbackSidebar.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/FeedbackSidebar.tsx
index 249a2335704..a593f79b9ac 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/FeedbackSidebar.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/FeedbackSidebar.tsx
@@ -1,9 +1,13 @@
 import {toast} from '@wandb/weave/common/components/elements/Toast';
+import {useOrgName} from '@wandb/weave/common/hooks/useOrganization';
 import {useViewerInfo} from '@wandb/weave/common/hooks/useViewerInfo';
+import {useViewerUserInfo2} from '@wandb/weave/common/hooks/useViewerUserInfo';
 import {Button} from '@wandb/weave/components/Button';
 import {Icon} from '@wandb/weave/components/Icon';
+import {Loading} from '@wandb/weave/components/Loading';
+import {annotationsViewed} from '@wandb/weave/integrations/analytics/viewEvents';
 import {makeRefCall} from '@wandb/weave/util/refs';
-import React, {useState} from 'react';
+import React, {useEffect, useState} from 'react';
 import {useHistory} from 'react-router-dom';
 
 import {useWeaveflowRouteContext} from '../../context';
@@ -14,6 +18,7 @@ import {tsHumanAnnotationSpec} from './humanAnnotationTypes';
 
 type FeedbackSidebarProps = {
   humanAnnotationSpecs: tsHumanAnnotationSpec[];
+  specsLoading: boolean;
   callID: string;
   entity: string;
   project: string;
@@ -21,6 +26,7 @@ type FeedbackSidebarProps = {
 
 export const FeedbackSidebar = ({
   humanAnnotationSpecs,
+  specsLoading,
   callID,
   entity,
   project,
@@ -32,6 +38,11 @@ export const FeedbackSidebar = ({
     Record<string, () => Promise<boolean>>
   >({});
 
+  const {loading: viewerLoading, userInfo} = useViewerUserInfo2();
+  const {loading: orgNameLoading, orgName} = useOrgName({
+    entityName: entity,
+  });
+
   const save = async () => {
     setIsSaving(true);
     try {
@@ -58,6 +69,28 @@ export const FeedbackSidebar = ({
     }
   };
 
+  useEffect(() => {
+    if (!viewerLoading && !orgNameLoading) {
+      annotationsViewed({
+        traceId: callID,
+        userId: userInfo.id,
+        organizationName: orgName,
+        entityName: entity,
+        projectName: project,
+        numAnnotationSpecs: humanAnnotationSpecs.length,
+      });
+    }
+  }, [
+    viewerLoading,
+    orgNameLoading,
+    userInfo,
+    orgName,
+    entity,
+    project,
+    humanAnnotationSpecs.length,
+    callID,
+  ]);
+
   return (
     <div className="flex h-full flex-col bg-white">
       <div className="justify-left flex w-full border-b border-moon-300 p-12">
@@ -75,7 +108,7 @@ export const FeedbackSidebar = ({
               setUnsavedFeedbackChanges={setUnsavedFeedbackChanges}
             />
           </div>
-          <div className="flex w-full border-t border-moon-300 p-6 pr-10">
+          <div className="flex w-full border-t border-moon-300 p-12 pr-18">
             <Button
               onClick={save}
               variant="primary"
@@ -88,6 +121,10 @@ export const FeedbackSidebar = ({
             </Button>
           </div>
         </>
+      ) : specsLoading ? (
+        <div className="mt-12 w-full items-center justify-center">
+          <Loading centered />
+        </div>
       ) : (
         <div className="mt-12 w-full items-center justify-center">
           <Empty {...EMPTY_PROPS_ANNOTATIONS} />
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/HumanAnnotation.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/HumanAnnotation.tsx
index 9a1d1438786..66838942c03 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/HumanAnnotation.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/HumanAnnotation.tsx
@@ -167,7 +167,7 @@ const FeedbackComponentSelector: React.FC<{
   jsonSchema: Record<string, any>;
   focused: boolean;
   onAddFeedback: (value: any) => Promise<boolean>;
-  foundValue: string | number | null;
+  foundValue: string | number | boolean | null;
   feedbackSpecRef: string;
   setUnsavedFeedbackChanges: React.Dispatch<
     React.SetStateAction<Record<string, () => Promise<boolean>>>
@@ -194,12 +194,23 @@ const FeedbackComponentSelector: React.FC<{
     );
 
     switch (type) {
+      case 'integer':
+        return (
+          <NumericalFeedbackColumn
+            min={jsonSchema.minimum}
+            max={jsonSchema.maximum}
+            isInteger={true}
+            onAddFeedback={wrappedOnAddFeedback}
+            defaultValue={foundValue as number | null}
+            focused={focused}
+          />
+        );
       case 'number':
         return (
           <NumericalFeedbackColumn
-            min={jsonSchema.min}
-            max={jsonSchema.max}
-            isInteger={jsonSchema.is_integer}
+            min={jsonSchema.minimum}
+            max={jsonSchema.maximum}
+            isInteger={false}
             onAddFeedback={wrappedOnAddFeedback}
             defaultValue={foundValue as number | null}
             focused={focused}
@@ -226,7 +237,7 @@ const FeedbackComponentSelector: React.FC<{
         return (
           <BinaryFeedbackColumn
             onAddFeedback={wrappedOnAddFeedback}
-            defaultValue={foundValue as string | null}
+            defaultValue={foundValue as boolean | null}
             focused={focused}
           />
         );
@@ -513,35 +524,35 @@ export const BinaryFeedbackColumn = ({
   defaultValue,
   focused,
 }: {
-  onAddFeedback?: (value: string) => Promise<boolean>;
-  defaultValue: string | null;
+  onAddFeedback?: (value: any) => Promise<boolean>;
+  defaultValue: boolean | null;
   focused?: boolean;
 }) => {
-  const [value, setValue] = useState<string | null>(defaultValue);
+  const [value, setValue] = useState<boolean | null>(defaultValue);
 
   useEffect(() => {
     setValue(defaultValue);
   }, [defaultValue]);
 
-  const handleClick = (newValue: string) => {
+  const handleClick = (newValue: boolean) => {
     // If clicking the same value, deselect it
     const valueToSet = value === newValue ? null : newValue;
     setValue(valueToSet);
-    onAddFeedback?.(valueToSet ?? '');
+    onAddFeedback?.(valueToSet);
   };
 
   return (
     <Tailwind>
       <div className="flex w-full justify-center gap-10">
         <Button
-          variant={value === 'true' ? 'primary' : 'outline'}
-          onClick={() => handleClick('true')}
+          variant={value === true ? 'primary' : 'outline'}
+          onClick={() => handleClick(true)}
           autoFocus={focused}>
           True
         </Button>
         <Button
-          variant={value === 'false' ? 'primary' : 'outline'}
-          onClick={() => handleClick('false')}>
+          variant={value === false ? 'primary' : 'outline'}
+          onClick={() => handleClick(false)}>
           False
         </Button>
       </div>
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/humanAnnotationTypes.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/humanAnnotationTypes.ts
index 43b58c759d6..8d3a89e9d24 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/humanAnnotationTypes.ts
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/humanAnnotationTypes.ts
@@ -2,7 +2,13 @@ import {AnnotationSpec} from '../../pages/wfReactInterface/generatedBaseObjectCl
 import {Feedback} from '../../pages/wfReactInterface/traceServerClientTypes';
 
 export const HUMAN_ANNOTATION_BASE_TYPE = 'wandb.annotation';
-export const FEEDBACK_TYPE_OPTIONS = ['string', 'number', 'boolean', 'enum'];
+export const FEEDBACK_TYPE_OPTIONS = [
+  'string',
+  'number',
+  'boolean',
+  'enum',
+  'integer',
+];
 export type FeedbackSchemaType = (typeof FEEDBACK_TYPE_OPTIONS)[number];
 
 export const makeAnnotationFeedbackType = (type: string) =>
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/tsHumanFeedback.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/tsHumanFeedback.ts
index e6ebabd1b73..6460d474b15 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/tsHumanFeedback.ts
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/feedback/StructuredFeedback/tsHumanFeedback.ts
@@ -12,7 +12,7 @@ import {tsHumanAnnotationSpec} from './humanAnnotationTypes';
 export const useHumanAnnotationSpecs = (
   entity: string,
   project: string
-): tsHumanAnnotationSpec[] => {
+): {humanAnnotationSpecs: tsHumanAnnotationSpec[]; specsLoading: boolean} => {
   const req: TraceObjQueryReq = {
     project_id: projectIdFromParts({entity, project}),
     filter: {
@@ -31,20 +31,23 @@ export const useHumanAnnotationSpecs = (
 
   return useMemo(() => {
     if (cols == null) {
-      return [];
+      return {humanAnnotationSpecs: [], specsLoading: res.loading};
     }
-    return cols.map(col => ({
-      ...col.val,
-      // attach object refs to the columns
-      ref: objectVersionKeyToRefUri({
-        scheme: 'weave',
-        weaveKind: 'object',
-        entity,
-        project,
-        objectId: col.object_id,
-        versionHash: col.digest,
-        path: '',
-      }),
-    }));
-  }, [cols, entity, project]);
+    return {
+      humanAnnotationSpecs: cols.map(col => ({
+        ...col.val,
+        // attach object refs to the columns
+        ref: objectVersionKeyToRefUri({
+          scheme: 'weave',
+          weaveKind: 'object',
+          entity,
+          project,
+          objectId: col.object_id,
+          versionHash: col.digest,
+          path: '',
+        }),
+      })),
+      specsLoading: res.loading,
+    };
+  }, [cols, entity, project, res.loading]);
 };
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx
index d674ae48975..7bf7b17f4bf 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallPage.tsx
@@ -223,7 +223,7 @@ const CallPageInnerVertical: FC<{
       )
     );
   }, [currentRouter, history, path, showTraceTree, call, showFeedbackExpand]);
-  const humanAnnotationSpecs = useHumanAnnotationSpecs(
+  const {humanAnnotationSpecs, specsLoading} = useHumanAnnotationSpecs(
     call.entity,
     call.project
   );
@@ -275,16 +275,6 @@ const CallPageInnerVertical: FC<{
             <PaginationControls call={call} path={path} />
           )}
           <Box sx={{marginLeft: showPaginationContols ? 0 : 'auto'}}>
-            <Button
-              icon="marker"
-              tooltip={`${
-                showFeedbackExpand ? 'Hide' : 'Show'
-              } feedback sidebar`}
-              variant="ghost"
-              active={showFeedbackExpand ?? false}
-              onClick={onToggleFeedbackExpand}
-              className="mr-4"
-            />
             <Button
               icon="layout-tabs"
               tooltip={`${showTraceTree ? 'Hide' : 'Show'} trace tree`}
@@ -292,6 +282,14 @@ const CallPageInnerVertical: FC<{
               active={showTraceTree ?? false}
               onClick={onToggleTraceTree}
             />
+            <Button
+              icon="marker"
+              tooltip={`${showFeedbackExpand ? 'Hide' : 'Show'} feedback`}
+              variant="ghost"
+              active={showFeedbackExpand ?? false}
+              onClick={onToggleFeedbackExpand}
+              className="ml-4"
+            />
           </Box>
         </Box>
       }
@@ -301,6 +299,7 @@ const CallPageInnerVertical: FC<{
           <div className="flex h-full flex-col">
             <FeedbackSidebar
               humanAnnotationSpecs={humanAnnotationSpecs}
+              specsLoading={specsLoading}
               callID={currentCall.callId}
               entity={currentCall.entity}
               project={currentCall.project}
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/PaginationControls.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/PaginationControls.tsx
index e0f00f80799..03af4838d77 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/PaginationControls.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/PaginationControls.tsx
@@ -4,8 +4,11 @@ import React, {FC, useCallback, useContext, useEffect} from 'react';
 import {useHistory} from 'react-router-dom';
 
 import {TableRowSelectionContext} from '../../../Browse3';
-import {TRACETREE_PARAM, useWeaveflowCurrentRouteContext} from '../../context';
-import {isEvaluateOp} from '../common/heuristics';
+import {
+  FEEDBACK_EXPAND_PARAM,
+  TRACETREE_PARAM,
+  useWeaveflowCurrentRouteContext,
+} from '../../context';
 import {useURLSearchParamsDict} from '../util';
 import {CallSchema} from '../wfReactInterface/wfDataModelHooksInterface';
 
@@ -21,9 +24,11 @@ export const PaginationControls: FC<{
   const currentRouter = useWeaveflowCurrentRouteContext();
   const query = useURLSearchParamsDict();
   const showTraceTree =
-    TRACETREE_PARAM in query
-      ? query[TRACETREE_PARAM] === '1'
-      : !isEvaluateOp(call.spanName);
+    TRACETREE_PARAM in query ? query[TRACETREE_PARAM] === '1' : false;
+  const showFeedbackExpand =
+    FEEDBACK_EXPAND_PARAM in query
+      ? query[FEEDBACK_EXPAND_PARAM] === '1'
+      : false;
 
   const onNextCall = useCallback(() => {
     const nextCallId = getNextRowId?.(call.callId);
@@ -35,11 +40,20 @@ export const PaginationControls: FC<{
           call.traceId,
           nextCallId,
           path,
-          showTraceTree
+          showTraceTree,
+          showFeedbackExpand
         )
       );
     }
-  }, [call, currentRouter, history, path, showTraceTree, getNextRowId]);
+  }, [
+    call,
+    currentRouter,
+    history,
+    path,
+    showTraceTree,
+    getNextRowId,
+    showFeedbackExpand,
+  ]);
   const onPreviousCall = useCallback(() => {
     const previousRowId = getPreviousRowId?.(call.callId);
     if (previousRowId) {
@@ -50,11 +64,20 @@ export const PaginationControls: FC<{
           call.traceId,
           previousRowId,
           path,
-          showTraceTree
+          showTraceTree,
+          showFeedbackExpand
         )
       );
     }
-  }, [call, currentRouter, history, path, showTraceTree, getPreviousRowId]);
+  }, [
+    call,
+    currentRouter,
+    history,
+    path,
+    showTraceTree,
+    getPreviousRowId,
+    showFeedbackExpand,
+  ]);
   const handleKeyDown = useCallback(
     (event: KeyboardEvent) => {
       if (event.key === 'ArrowDown' && event.shiftKey) {
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx
index 29f72a4eef9..0fc4e8298a5 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx
@@ -366,6 +366,9 @@ function buildCallsTableColumns(
             return row[c];
           },
           renderCell: (params: GridRenderCellParams<TraceCallSchema>) => {
+            if (typeof params.value === 'boolean') {
+              return <div>{params.value ? 'true' : 'false'}</div>;
+            }
             return <div>{params.value}</div>;
           },
         };
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/AnnotationScorerForm.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/AnnotationScorerForm.tsx
index a931c51885a..3b14910b12b 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/AnnotationScorerForm.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/AnnotationScorerForm.tsx
@@ -11,15 +11,20 @@ import {ZSForm} from './ZodSchemaForm';
 
 const AnnotationScorerFormSchema = z.object({
   Name: z.string().min(1),
-  Description: z.string().min(1),
+  Description: z.string().optional(),
   Type: z.discriminatedUnion('type', [
     z.object({
       type: z.literal('boolean'),
     }),
+    z.object({
+      type: z.literal('integer'),
+      minimum: z.number().optional().describe('Optional minimum value'),
+      maximum: z.number().optional().describe('Optional maximum value'),
+    }),
     z.object({
       type: z.literal('number'),
-      min: z.number().optional().describe('Optional minimum value'),
-      max: z.number().optional().describe('Optional maximum value'),
+      minimum: z.number().optional().describe('Optional minimum value'),
+      maximum: z.number().optional().describe('Optional maximum value'),
     }),
     z.object({
       type: z.literal('string'),
@@ -29,16 +34,20 @@ const AnnotationScorerFormSchema = z.object({
         .describe('Optional maximum length of the string'),
     }),
     z.object({
-      type: z.literal('enum'),
+      type: z.literal('options'),
       enum: z.array(z.string()).describe('List of options to choose from'),
     }),
   ]),
 });
 
+const DEFAULT_STATE = {
+  Type: {type: 'boolean'},
+} as z.infer<typeof AnnotationScorerFormSchema>;
+
 export const AnnotationScorerForm: FC<
   ScorerFormProps<z.infer<typeof AnnotationScorerFormSchema>>
 > = ({data, onDataChange}) => {
-  const [config, setConfig] = useState(data);
+  const [config, setConfig] = useState(data ?? DEFAULT_STATE);
   const [isValid, setIsValid] = useState(false);
 
   const handleConfigChange = useCallback(
@@ -76,7 +85,7 @@ export const onAnnotationScorerSave = async (
   client: TraceServerClient
 ) => {
   let type = data.Type.type;
-  if (type === 'enum') {
+  if (type === 'options') {
     type = 'string';
   }
   return createBaseObjectInstance(client, 'AnnotationSpec', {
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/NewScorerDrawer.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/NewScorerDrawer.tsx
index b483dd19229..f562749f05b 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/NewScorerDrawer.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ScorersPage/NewScorerDrawer.tsx
@@ -119,6 +119,7 @@ export const NewScorerDrawer: FC<NewScorerDrawerProps> = ({
         getClient()
       );
       onClose();
+      setFormData(null);
     } catch (error) {
       console.error('Failed to create scorer:', error);
       // Handle error appropriately
diff --git a/weave-js/src/integrations/analytics/viewEvents.ts b/weave-js/src/integrations/analytics/viewEvents.ts
index 0881db1ad74..8d93c7897a2 100644
--- a/weave-js/src/integrations/analytics/viewEvents.ts
+++ b/weave-js/src/integrations/analytics/viewEvents.ts
@@ -103,3 +103,46 @@ export const evaluationViewed = makeTrackEvent<
     };
   }
 >('Weave evaluation viewed');
+
+export const annotationsViewed = makeTrackEvent<
+  {
+    traceId: string;
+    userId: string;
+    organizationName: string;
+    entityName: string;
+    projectName: string;
+    numAnnotationSpecs: number;
+  },
+  {
+    _description: `User viewed the annotations of a trace`;
+    _location: '/calls/:itemName';
+    _motivation: 'Used for annotations views';
+    traceId: {
+      description: 'ID of trace viewed';
+      exampleValues: [
+        'bb5621fd-91bc-42af-b017-1a34e3250330',
+        'b2136295-edc5-4778-9304-0fa36c9541a4'
+      ];
+    };
+    userId: {
+      description: 'ID of user viewing annotations';
+      exampleValues: ['VXNlcjo0NTM4MTM='];
+    };
+    organizationName: {
+      description: 'Name of organization';
+      exampleValues: ['my-org'];
+    };
+    entityName: {
+      description: 'Name of entity';
+      exampleValues: ['my-entity'];
+    };
+    projectName: {
+      description: 'Name of project';
+      exampleValues: ['my-project'];
+    };
+    numAnnotationSpecs: {
+      description: 'Number of human annotation specs';
+      exampleValues: [1, 2, 3];
+    };
+  }
+>('Weave annotations viewed');