[Search] [Playground]Improve follow up question flow (#189848)

## 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 <[email protected]>
elastic · Aug 7, 2024 · 5f49cc5 · 5f49cc5
1 parent 412473f
commit 5f49cc5
Show file tree

Hide file tree

Showing 6 changed files with 84 additions and 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)