From 5ec4f1cc8df0ff386819c320f41f7644bcd7efe9 Mon Sep 17 00:00:00 2001
From: Sihan He <hesihan@amazon.com>
Date: Fri, 20 Dec 2024 11:55:07 +0800
Subject: [PATCH] Add flag for trace view

Signed-off-by: Sihan He <hesihan@amazon.com>
---
 common/types/config.ts                        |  1 +
 .../generate_popover_body.tsx                 |  3 ++
 public/tabs/chat/messages/message_action.tsx  |  4 +-
 .../chat/messages/message_bubble.test.tsx     | 47 +++++++++++++++++++
 public/tabs/chat/messages/message_bubble.tsx  |  5 +-
 5 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/common/types/config.ts b/common/types/config.ts
index 2b92644f..698f00bf 100644
--- a/common/types/config.ts
+++ b/common/types/config.ts
@@ -9,6 +9,7 @@ export const configSchema = schema.object({
   enabled: schema.boolean({ defaultValue: true }),
   chat: schema.object({
     enabled: schema.boolean({ defaultValue: false }),
+    trace: schema.boolean({ defaultValue: true }),
   }),
   incontextInsight: schema.object({
     enabled: schema.boolean({ defaultValue: true }),
diff --git a/public/components/incontext_insight/generate_popover_body.tsx b/public/components/incontext_insight/generate_popover_body.tsx
index ede3a119..7233fba3 100644
--- a/public/components/incontext_insight/generate_popover_body.tsx
+++ b/public/components/incontext_insight/generate_popover_body.tsx
@@ -321,6 +321,7 @@ export const GeneratePopoverBody: React.FC<{
   };
 
   const renderInnerFooter = () => {
+    const traceTip = 'Insight With RAG';
     return (
       <EuiPopoverFooter className="incontextInsightGeneratePopoverFooter" paddingSize="none">
         {
@@ -332,6 +333,7 @@ export const GeneratePopoverBody: React.FC<{
               onViewTrace={() => {
                 setShowInsight(true);
               }}
+              traceTip={traceTip}
               usageCollection={usageCollection}
               isOnTrace={showInsight}
               metricAppName={metricAppName}
@@ -345,6 +347,7 @@ export const GeneratePopoverBody: React.FC<{
               showFeedback
               showTraceIcon={insightAvailable}
               onViewTrace={() => {}}
+              traceTip={traceTip}
               usageCollection={usageCollection}
               isOnTrace={showInsight}
               metricAppName={metricAppName}
diff --git a/public/tabs/chat/messages/message_action.tsx b/public/tabs/chat/messages/message_action.tsx
index e2b84b24..d1cbdca9 100644
--- a/public/tabs/chat/messages/message_action.tsx
+++ b/public/tabs/chat/messages/message_action.tsx
@@ -22,6 +22,7 @@ interface MessageActionsProps {
   showTraceIcon?: boolean;
   isOnTrace?: boolean;
   traceInteractionId?: string;
+  traceTip?: string;
   onViewTrace?: () => void;
   shouldActionBarVisibleOnHover?: boolean;
   isFullWidth?: boolean;
@@ -44,6 +45,7 @@ export const MessageActions: React.FC<MessageActionsProps> = ({
   showTraceIcon = false,
   isOnTrace = false,
   traceInteractionId = null,
+  traceTip = 'info',
   onViewTrace,
   shouldActionBarVisibleOnHover = false,
   isFullWidth = false,
@@ -146,7 +148,7 @@ export const MessageActions: React.FC<MessageActionsProps> = ({
     trace: {
       show: showTraceIcon && onViewTrace,
       component: renderButtonWithTooltip(
-        'Insight with RAG',
+        traceTip,
         <EuiSmallButtonIcon
           aria-label="How was this generated?"
           {...(traceInteractionId && {
diff --git a/public/tabs/chat/messages/message_bubble.test.tsx b/public/tabs/chat/messages/message_bubble.test.tsx
index a41274da..8fca94ef 100644
--- a/public/tabs/chat/messages/message_bubble.test.tsx
+++ b/public/tabs/chat/messages/message_bubble.test.tsx
@@ -8,6 +8,7 @@ import { render, screen, fireEvent } from '@testing-library/react';
 
 import { MessageBubble } from './message_bubble';
 import { IOutput } from '../../../../common/types/chat_saved_object_attributes';
+import * as services from '../../../services';
 import * as useFeedbackHookExports from '../../../hooks/use_feed_back';
 import * as useChatActionsExports from '../../../hooks/use_chat_actions';
 import * as coreHookExports from '../../../contexts/core_context';
@@ -18,6 +19,7 @@ describe('<MessageBubble />', () => {
   const reportUiStatsMock = jest.fn();
 
   beforeEach(() => {
+    jest.spyOn(services, 'getConfigSchema').mockReturnValue({ chat: { trace: true } });
     jest
       .spyOn(useFeedbackHookExports, 'useFeedback')
       .mockReturnValue({ feedbackResult: undefined, sendFeedback: sendFeedbackMock });
@@ -289,4 +291,49 @@ describe('<MessageBubble />', () => {
     );
     expect(screen.queryByTestId('trace-icon-bar')).toBeNull();
   });
+
+  it('should control view trace through config flag', () => {
+    render(
+      <MessageBubble
+        showActionBar={true}
+        showRegenerate={true}
+        message={{
+          type: 'output',
+          contentType: 'markdown',
+          content: 'here are the indices in your cluster: .alert',
+          interactionId: 'bar1',
+        }}
+        interaction={{
+          input: 'foo',
+          response: 'bar1',
+          conversation_id: 'foo',
+          interaction_id: 'bar1',
+          create_time: new Date().toLocaleString(),
+        }}
+      />
+    );
+    expect(screen.queryByTestId('trace-icon-bar1')).toBeVisible();
+
+    jest.spyOn(services, 'getConfigSchema').mockReturnValue({ chat: { trace: false } });
+    render(
+      <MessageBubble
+        showActionBar={true}
+        showRegenerate={true}
+        message={{
+          type: 'output',
+          contentType: 'markdown',
+          content: 'here are the indices in your cluster: .alert',
+          interactionId: 'bar2',
+        }}
+        interaction={{
+          input: 'foo',
+          response: 'bar2',
+          conversation_id: 'foo',
+          interaction_id: 'bar2',
+          create_time: new Date().toLocaleString(),
+        }}
+      />
+    );
+    expect(screen.queryByTestId('trace-icon-bar2')).toBeNull();
+  });
 });
diff --git a/public/tabs/chat/messages/message_bubble.tsx b/public/tabs/chat/messages/message_bubble.tsx
index d86997b1..7764325e 100644
--- a/public/tabs/chat/messages/message_bubble.tsx
+++ b/public/tabs/chat/messages/message_bubble.tsx
@@ -17,6 +17,7 @@ import React from 'react';
 import { IconType } from '@elastic/eui/src/components/icon/icon';
 import { MessageActions } from './message_action';
 import { useCore } from '../../../contexts';
+import { getConfigSchema } from '../../../services';
 
 // TODO: Replace with getChrome().logos.Chat.url
 import { useChatActions } from '../../../hooks';
@@ -123,6 +124,7 @@ export const MessageBubble: React.FC<MessageBubbleProps> = React.memo((props) =>
     }
 
     const fullWidth = props.message.fullWidth;
+    const configSchema = getConfigSchema();
 
     return (
       <EuiFlexGroup
@@ -159,7 +161,8 @@ export const MessageBubble: React.FC<MessageBubbleProps> = React.memo((props) =>
                   interaction={props.interaction}
                   message={props.message as IOutput}
                   showFeedback={showFeedback}
-                  showTraceIcon={!!props.message.interactionId}
+                  showTraceIcon={configSchema.chat.trace && !!props.message.interactionId}
+                  traceTip="How was this generated?"
                   traceInteractionId={props.message.interactionId || ''}
                   onViewTrace={() => {
                     const message = props.message as IOutput;