From 5f49cc50178d4abc67be420534f8cfeafb60c954 Mon Sep 17 00:00:00 2001
From: Saarika Bhasi <55930906+saarikabhasi@users.noreply.github.com>
Date: Wed, 7 Aug 2024 16:03:53 -0400
Subject: [PATCH] [Search] [Playground]Improve follow up question flow
 (#189848)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

Adds the query used for searching. This improves visibility to user on
the query using for follow up questions.

<img width="1510" alt="Screenshot 2024-08-03 at 3 24 10 PM"
src="https://github.com/user-attachments/assets/e1da0fb4-ca58-4f1e-866a-5acf070b02f3">



### Checklist

Delete any items that are not applicable to this PR.

- [x] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md)
- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

---------

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
---
 .../message_list/assistant_message.test.tsx   |   2 +-
 .../message_list/assistant_message.tsx        | 104 +++++++++++-------
 .../public/hooks/use_ai_assist_chat.ts        |   2 +-
 .../plugins/search_playground/public/types.ts |   4 +-
 .../public/utils/transform_to_messages.ts     |   3 +
 .../server/lib/conversational_chain.ts        |  11 ++
 6 files changed, 84 insertions(+), 42 deletions(-)

diff --git a/x-pack/plugins/search_playground/public/components/message_list/assistant_message.test.tsx b/x-pack/plugins/search_playground/public/components/message_list/assistant_message.test.tsx
index f3a903331cf3b..4608546616a51 100644
--- a/x-pack/plugins/search_playground/public/components/message_list/assistant_message.test.tsx
+++ b/x-pack/plugins/search_playground/public/components/message_list/assistant_message.test.tsx
@@ -38,7 +38,7 @@ describe('AssistantMessage component', () => {
     createdAt: new Date(),
     citations: [],
     retrievalDocs: [{ content: '', metadata: { _id: '1', _index: 'index', _score: 1 } }],
-    inputTokens: { context: 20, total: 10 },
+    inputTokens: { context: 20, total: 10, searchQuery: 'Test question' },
   };
 
   it('renders message content correctly', () => {
diff --git a/x-pack/plugins/search_playground/public/components/message_list/assistant_message.tsx b/x-pack/plugins/search_playground/public/components/message_list/assistant_message.tsx
index 7e94059bdc272..e0b14c2a31934 100644
--- a/x-pack/plugins/search_playground/public/components/message_list/assistant_message.tsx
+++ b/x-pack/plugins/search_playground/public/components/message_list/assistant_message.tsx
@@ -53,53 +53,79 @@ export const AssistantMessage: React.FC<AssistantMessageProps> = ({ message }) =
   return (
     <>
       {!!retrievalDocs?.length && (
-        <EuiComment
-          username={username}
-          timelineAvatar="dot"
-          data-test-subj="retrieval-docs-comment"
-          eventColor="subdued"
-          css={{
-            '.euiAvatar': { backgroundColor: euiTheme.colors.ghost },
-            '.euiCommentEvent': {
-              border: euiTheme.border.thin,
-              borderRadius: euiTheme.border.radius.medium,
-            },
-          }}
-          event={
-            <>
+        <>
+          <EuiComment
+            username={username}
+            timelineAvatar="dot"
+            data-test-subj="assistant-message-searching"
+            eventColor="subdued"
+            css={{
+              '.euiAvatar': { backgroundColor: euiTheme.colors.ghost },
+              '.euiCommentEvent': {
+                border: euiTheme.border.thin,
+                borderRadius: euiTheme.border.radius.medium,
+              },
+            }}
+            event={
               <EuiText size="s">
                 <p>
                   <FormattedMessage
-                    id="xpack.searchPlayground.chat.message.assistant.retrievalDocs"
-                    defaultMessage="Grounding answer based on"
+                    id="xpack.searchPlayground.chat.message.assistant.searchingQuestion"
+                    defaultMessage='Searching for "{question}"'
+                    values={{ question: inputTokens.searchQuery }}
                   />
-                  {` `}
                 </p>
               </EuiText>
+            }
+          />
+          <EuiComment
+            username={username}
+            timelineAvatar="dot"
+            data-test-subj="retrieval-docs-comment"
+            eventColor="subdued"
+            css={{
+              '.euiAvatar': { backgroundColor: euiTheme.colors.ghost },
+              '.euiCommentEvent': {
+                border: euiTheme.border.thin,
+                borderRadius: euiTheme.border.radius.medium,
+              },
+            }}
+            event={
+              <>
+                <EuiText size="s">
+                  <p>
+                    <FormattedMessage
+                      id="xpack.searchPlayground.chat.message.assistant.retrievalDocs"
+                      defaultMessage="Grounding answer based on"
+                    />
+                    {` `}
+                  </p>
+                </EuiText>
 
-              <EuiButtonEmpty
-                css={{ blockSize: 'auto' }}
-                size="s"
-                flush="left"
-                data-test-subj="retrieval-docs-button"
-                onClick={() => setIsDocsFlyoutOpen(true)}
-              >
-                <FormattedMessage
-                  id="xpack.searchPlayground.chat.message.assistant.retrievalDocButton"
-                  defaultMessage="{count} document sources"
-                  values={{ count: retrievalDocs.length }}
-                />
-              </EuiButtonEmpty>
+                <EuiButtonEmpty
+                  css={{ blockSize: 'auto' }}
+                  size="s"
+                  flush="left"
+                  data-test-subj="retrieval-docs-button"
+                  onClick={() => setIsDocsFlyoutOpen(true)}
+                >
+                  <FormattedMessage
+                    id="xpack.searchPlayground.chat.message.assistant.retrievalDocButton"
+                    defaultMessage="{count} document sources"
+                    values={{ count: retrievalDocs.length }}
+                  />
+                </EuiButtonEmpty>
 
-              {isDocsFlyoutOpen && (
-                <RetrievalDocsFlyout
-                  onClose={() => setIsDocsFlyoutOpen(false)}
-                  retrievalDocs={retrievalDocs}
-                />
-              )}
-            </>
-          }
-        />
+                {isDocsFlyoutOpen && (
+                  <RetrievalDocsFlyout
+                    onClose={() => setIsDocsFlyoutOpen(false)}
+                    retrievalDocs={retrievalDocs}
+                  />
+                )}
+              </>
+            }
+          />
+        </>
       )}
       {retrievalDocs?.length === 0 && (
         <EuiComment
diff --git a/x-pack/plugins/search_playground/public/hooks/use_ai_assist_chat.ts b/x-pack/plugins/search_playground/public/hooks/use_ai_assist_chat.ts
index f29bb7de11f88..d046b9425c60a 100644
--- a/x-pack/plugins/search_playground/public/hooks/use_ai_assist_chat.ts
+++ b/x-pack/plugins/search_playground/public/hooks/use_ai_assist_chat.ts
@@ -184,7 +184,7 @@ export function useAIAssistChat({
       if (messagesRef.current.length === 0) return null;
 
       const chatRequest: ChatRequest = {
-        messages: messagesRef.current,
+        messages: messagesRef.current.slice(0, messagesRef.current.length - 1),
         options,
         data,
       };
diff --git a/x-pack/plugins/search_playground/public/types.ts b/x-pack/plugins/search_playground/public/types.ts
index 2bbd45ff16230..90f6cdc72e04f 100644
--- a/x-pack/plugins/search_playground/public/types.ts
+++ b/x-pack/plugins/search_playground/public/types.ts
@@ -100,8 +100,9 @@ export interface AnnotationDoc {
 }
 
 export interface AnnotationTokens {
-  type: 'prompt_token_count' | 'context_token_count' | 'context_clipped';
+  type: 'prompt_token_count' | 'context_token_count' | 'context_clipped' | 'search_query';
   count: number;
+  question?: string;
 }
 
 export interface Doc {
@@ -117,6 +118,7 @@ export interface AIMessage extends Message {
     context: number;
     total: number;
     contextClipped?: number;
+    searchQuery: string;
   };
 }
 
diff --git a/x-pack/plugins/search_playground/public/utils/transform_to_messages.ts b/x-pack/plugins/search_playground/public/utils/transform_to_messages.ts
index 4500ac6f10890..af8d75a8ee7f1 100644
--- a/x-pack/plugins/search_playground/public/utils/transform_to_messages.ts
+++ b/x-pack/plugins/search_playground/public/utils/transform_to_messages.ts
@@ -44,6 +44,9 @@ export const transformFromChatMessages = (messages: UseChatHelpers['messages']):
           contextClipped: annotations?.find(
             (annotation): annotation is AnnotationTokens => annotation.type === 'context_clipped'
           )?.count,
+          searchQuery: annotations?.find(
+            (annotation): annotation is AnnotationTokens => annotation.type === 'search_query'
+          )?.question,
         },
       } as AIMessage;
     }
diff --git a/x-pack/plugins/search_playground/server/lib/conversational_chain.ts b/x-pack/plugins/search_playground/server/lib/conversational_chain.ts
index c63481e93c98f..922f672bda5c6 100644
--- a/x-pack/plugins/search_playground/server/lib/conversational_chain.ts
+++ b/x-pack/plugins/search_playground/server/lib/conversational_chain.ts
@@ -198,6 +198,13 @@ class ConversationalChainFn {
         context: RunnableSequence.from([(input) => input.question, retrievalChain]),
         question: (input) => input.question,
       },
+      RunnableLambda.from((inputs) => {
+        data.appendMessageAnnotation({
+          type: 'search_query',
+          question: inputs.question,
+        });
+        return inputs;
+      }),
       RunnableLambda.from(clipContext(this.options?.rag?.inputTokensLimit, prompt, data)),
       RunnableLambda.from(registerContextTokenCounts(data)),
       prompt,
@@ -236,6 +243,10 @@ class ConversationalChainFn {
                   type: 'prompt_token_count',
                   count: getTokenEstimateFromMessages(msg),
                 });
+                data.appendMessageAnnotation({
+                  type: 'search_query',
+                  question,
+                });
               }
             },
             // callback for prompt based models (Bedrock uses ActionsClientLlm)