Fix duplicate status in trace explorer header, and remove filtering out session level assessments (#20131)

Signed-off-by: Daniel Lok <daniel.lok@databricks.com>

Author: daniellok-db

mlflow/server/js/src/lang/default/en.json

@@ -738,6 +738,10 @@
     "defaultMessage": "Chart view",
     "description": "Label for the table view toggle button in the logged model list page"
   },
+  "451ETc": {
+    "defaultMessage": "Session",
+    "description": "Label for the session to which an assessment belongs"
+  },
   "46xd2Z": {
     "defaultMessage": "Compare",
     "description": "Runs charts > components > config > RunsChartsConfigureDifferenceChart > Compare config section"

mlflow/server/js/src/shared/web-shared/model-trace-explorer/ModelTraceExplorer.utils.tsx

@@ -1354,6 +1354,7 @@ export const isSessionLevelAssessment = (assessment: Assessment): boolean => {
  */
 export const getTraceLevelAssessments = (assessments?: Assessment[]) =>
   assessments?.filter((assessment) => !isSessionLevelAssessment(assessment)) ?? [];
+
 export const isValidException = (
   event: ModelTraceEvent,
 ): event is ModelTraceEvent & {

mlflow/server/js/src/shared/web-shared/model-trace-explorer/ModelTraceExplorerMetricSection.tsx

@@ -9,6 +9,7 @@ export const ModelTraceHeaderMetricSection = ({
   color,
   getTruncatedLabel,
   onCopy,
+  hideLabel = false,
 }: {
   label: React.ReactNode;
   /**
@@ -23,6 +24,7 @@ export const ModelTraceHeaderMetricSection = ({
   color?: TagColors;
   getTruncatedLabel: (label: string) => string;
   onCopy: () => void;
+  hideLabel?: boolean;
 }) => {
   const { theme } = useDesignSystemTheme();
 
@@ -36,14 +38,16 @@ export const ModelTraceHeaderMetricSection = ({
       css={{
         display: 'flex',
         alignItems: 'center',
-        justifyContent: 'center',
+        justifyContent: hideLabel ? 'flex-start' : 'center',
         flexDirection: 'row',
         gap: theme.spacing.sm,
       }}
     >
-      <Typography.Text size="md" color="secondary">
-        {label}
-      </Typography.Text>
+      {!hideLabel && (
+        <Typography.Text size="md" color="secondary">
+          {label}
+        </Typography.Text>
+      )}
       <Tooltip componentId="shared.model-trace-explorer.header-details.tooltip" content={value} maxWidth={400}>
         <Tag
           componentId="shared.model-trace-explorer.header-details.tag"

mlflow/server/js/src/shared/web-shared/model-trace-explorer/ModelTraceHeaderDetails.tsx

@@ -120,7 +120,6 @@ export const ModelTraceHeaderDetails = ({ modelTraceInfo }: { modelTraceInfo: Mo
             onCopy={handleCopy}
           />
         )}
-        {statusState && <ModelTraceHeaderStatusTag statusState={statusState} getTruncatedLabel={getTruncatedLabel} />}
         {sessionId && experimentId && (
           <ModelTraceHeaderSessionIdTag
             handleCopy={handleCopy}

mlflow/server/js/src/shared/web-shared/model-trace-explorer/ModelTraceHeaderSessionIdTag.tsx

@@ -15,11 +15,13 @@ export const ModelTraceHeaderSessionIdTag = ({
   sessionId,
   traceId,
   handleCopy,
+  hideLabel = false,
 }: {
   experimentId: string;
   sessionId: string;
   traceId?: string;
   handleCopy: () => void;
+  hideLabel?: boolean;
 }) => {
   const { theme } = useDesignSystemTheme();
   const location = useLocation();
@@ -47,6 +49,7 @@ export const ModelTraceHeaderSessionIdTag = ({
         color="default"
         getTruncatedLabel={getTruncatedLabel}
         onCopy={handleCopy}
+        hideLabel={hideLabel}
       />
     );
   }
@@ -56,14 +59,16 @@ export const ModelTraceHeaderSessionIdTag = ({
       css={{
         display: 'flex',
         alignItems: 'center',
-        justifyContent: 'center',
+        justifyContent: hideLabel ? 'flex-start' : 'center',
         flexDirection: 'row',
         gap: theme.spacing.sm,
       }}
     >
-      <Typography.Text size="md" color="secondary">
-        <FormattedMessage defaultMessage="Session ID" description="Label for the session id section" />
-      </Typography.Text>
+      {!hideLabel && (
+        <Typography.Text size="md" color="secondary">
+          <FormattedMessage defaultMessage="Session ID" description="Label for the session id section" />
+        </Typography.Text>
+      )}
       <Tooltip
         componentId="mlflow.model-trace-explorer.session-id-tag"
         content={<FormattedMessage defaultMessage="View chat session" description="Tooltip for the session id tag" />}

mlflow/server/js/src/shared/web-shared/model-trace-explorer/assessments-pane/ExpectationItem.tsx

@@ -13,18 +13,26 @@ import { SpanNameDetailViewLink } from './SpanNameDetailViewLink';
 import { getSourceIcon } from './utils';
 import type { ExpectationAssessment } from '../ModelTrace.types';
 import { useModelTraceExplorerViewState } from '../ModelTraceExplorerViewStateContext';
+import { ASSESSMENT_SESSION_METADATA_KEY } from '../constants';
+import { isEmpty } from 'lodash';
+import { ModelTraceHeaderSessionIdTag } from '../ModelTraceHeaderSessionIdTag';
+import { useParams } from '../RoutingUtils';
 
 export const ExpectationItem = ({ expectation }: { expectation: ExpectationAssessment }) => {
   const { theme } = useDesignSystemTheme();
   const [isEditing, setIsEditing] = useState(false);
   const [showDeleteModal, setShowDeleteModal] = useState(false);
   const [isExpanded, setIsExpanded] = useState(false);
   const { nodeMap, activeView } = useModelTraceExplorerViewState();
+  const { experimentId } = useParams();
 
   const associatedSpan = expectation.span_id ? nodeMap[expectation.span_id] : null;
+  // indicate if the assessment is session-level
+  const sessionId = expectation.metadata?.[ASSESSMENT_SESSION_METADATA_KEY];
+  const showSessionTag = activeView === 'summary' && !isEmpty(sessionId);
   // the summary view displays all assessments regardless of span, so
   // we need some way to indicate which span an assessment is associated with.
-  const showAssociatedSpan = activeView === 'summary' && associatedSpan;
+  const showAssociatedSpan = activeView === 'summary' && associatedSpan && !showSessionTag;
 
   const parsedValue = getParsedExpectationValue(expectation.expectation);
   const SourceIcon = getSourceIcon(expectation.source);
@@ -112,6 +120,23 @@ export const ExpectationItem = ({ expectation }: { expectation: ExpectationAsses
           <SpanNameDetailViewLink node={associatedSpan} />
         </div>
       )}
+      {showSessionTag && (
+        <div css={{ display: 'flex', flexDirection: 'column', gap: theme.spacing.xs }}>
+          <Typography.Text size="sm" color="secondary">
+            <FormattedMessage
+              defaultMessage="Session"
+              description="Label for the session to which an assessment belongs"
+            />
+          </Typography.Text>
+          <ModelTraceHeaderSessionIdTag
+            experimentId={experimentId ?? ''}
+            sessionId={sessionId ?? ''}
+            traceId={expectation.trace_id}
+            handleCopy={() => {}}
+            hideLabel
+          />
+        </div>
+      )}
     </div>
   );
 };

mlflow/server/js/src/shared/web-shared/model-trace-explorer/assessments-pane/FeedbackItemContent.tsx

@@ -1,4 +1,4 @@
-import { isNil } from 'lodash';
+import { isEmpty, isNil } from 'lodash';
 import { useState } from 'react';
 
 import { Typography, useDesignSystemTheme, NewWindowIcon } from '@databricks/design-system';
@@ -12,8 +12,14 @@ import { SpanNameDetailViewLink } from './SpanNameDetailViewLink';
 import type { FeedbackAssessment } from '../ModelTrace.types';
 import { useModelTraceExplorerViewState } from '../ModelTraceExplorerViewStateContext';
 import { Link, useParams } from '../RoutingUtils';
-import { MLFLOW_ASSESSMENT_JUDGE_COST, MLFLOW_ASSESSMENT_SCORER_TRACE_ID } from '../constants';
+import {
+  ASSESSMENT_SESSION_METADATA_KEY,
+  MLFLOW_ASSESSMENT_JUDGE_COST,
+  MLFLOW_ASSESSMENT_SCORER_TRACE_ID,
+} from '../constants';
 import { getExperimentPageTracesTabRoute } from '../routes';
+import { isSessionLevelAssessment } from '../ModelTraceExplorer.utils';
+import { ModelTraceHeaderSessionIdTag } from '../ModelTraceHeaderSessionIdTag';
 
 export const FeedbackItemContent = ({ feedback }: { feedback: FeedbackAssessment }) => {
   const [isHistoryModalVisible, setIsHistoryModalVisible] = useState(false);
@@ -24,9 +30,12 @@ export const FeedbackItemContent = ({ feedback }: { feedback: FeedbackAssessment
   const value = feedback.feedback.value;
 
   const associatedSpan = feedback.span_id ? nodeMap[feedback.span_id] : null;
+  // indicate if the assessment is session-level
+  const sessionId = feedback.metadata?.[ASSESSMENT_SESSION_METADATA_KEY];
+  const showSessionTag = activeView === 'summary' && !isEmpty(sessionId);
   // the summary view displays all assessments regardless of span, so
   // we need some way to indicate which span an assessment is associated with.
-  const showAssociatedSpan = activeView === 'summary' && associatedSpan;
+  const showAssociatedSpan = activeView === 'summary' && associatedSpan && !showSessionTag;
 
   const judgeTraceId = feedback.metadata?.[MLFLOW_ASSESSMENT_SCORER_TRACE_ID];
   const judgeTraceHref = judgeTraceId && experimentId ? getJudgeTraceHref(experimentId, judgeTraceId) : undefined;
@@ -68,6 +77,29 @@ export const FeedbackItemContent = ({ feedback }: { feedback: FeedbackAssessment
           <SpanNameDetailViewLink node={associatedSpan} />
         </div>
       )}
+      {showSessionTag && (
+        <div
+          css={{
+            display: 'flex',
+            flexDirection: 'column',
+            gap: theme.spacing.xs,
+          }}
+        >
+          <Typography.Text size="sm" color="secondary">
+            <FormattedMessage
+              defaultMessage="Session"
+              description="Label for the session to which an assessment belongs"
+            />
+          </Typography.Text>
+          <ModelTraceHeaderSessionIdTag
+            experimentId={experimentId ?? ''}
+            sessionId={sessionId ?? ''}
+            traceId={feedback.trace_id}
+            handleCopy={() => {}}
+            hideLabel
+          />
+        </div>
+      )}
       {isNil(feedback.feedback.error) && (
         <div css={{ display: 'flex', flexDirection: 'column', gap: theme.spacing.xs }}>
           <Typography.Text size="sm" color="secondary">

mlflow/server/js/src/shared/web-shared/model-trace-explorer/summary-view/ModelTraceExplorerSummaryView.tsx

@@ -33,9 +33,6 @@ export const ModelTraceExplorerSummaryView = () => {
     [updatePaneSizeRatios],
   );
 
-  // Get only the trace-level assessments (exclude session-level assessments)
-  const displayedAssessments = useMemo(() => getTraceLevelAssessments(allAssessments), [allAssessments]);
-
   const intermediateNodes = useIntermediateNodes(rootNode);
 
   if (!rootNode) {
@@ -53,7 +50,7 @@ export const ModelTraceExplorerSummaryView = () => {
     );
   }
   const AssessmentsPaneComponent = (
-    <AssessmentsPane assessments={displayedAssessments} traceId={rootNode.traceId} activeSpanId={undefined} />
+    <AssessmentsPane assessments={allAssessments} traceId={rootNode.traceId} activeSpanId={undefined} />
   );
 
   return !isInComparisonView && assessmentsPaneEnabled && assessmentsPaneExpanded ? (