From 6bc1dac6a314f5001a49b842b475a4c438438464 Mon Sep 17 00:00:00 2001
From: Yan Savitski <yan.savitski@elastic.co>
Date: Tue, 20 Feb 2024 20:57:46 +0100
Subject: [PATCH 1/3] [AI Playground] Add chat sidebar with accordions

---
 .../kbn-ai-playground/components/chat.tsx     |  55 +-------
 .../components/chat_sidebar.tsx               | 123 ++++++++++++++++++
 .../sources_panel/sources_panel_sidebar.tsx   |  60 +++------
 3 files changed, 146 insertions(+), 92 deletions(-)
 create mode 100644 packages/kbn-ai-playground/components/chat_sidebar.tsx

diff --git a/packages/kbn-ai-playground/components/chat.tsx b/packages/kbn-ai-playground/components/chat.tsx
index 4c74d6ce6c803..0e132629a6274 100644
--- a/packages/kbn-ai-playground/components/chat.tsx
+++ b/packages/kbn-ai-playground/components/chat.tsx
@@ -27,15 +27,11 @@ import { ChatForm, ChatFormFields, MessageRole } from '../types';
 
 import { MessageList } from './message_list/message_list';
 import { QuestionInput } from './question_input';
-import { OpenAIKeyField } from './open_ai_key_field';
-import { InstructionsField } from './instructions_field';
-import { IncludeCitationsField } from './include_citations_field';
-import { SourcesPanelSidebar } from './sources_panel/sources_panel_sidebar';
 
 import { TelegramIcon } from './telegram_icon';
 
 import { transformFromChatMessages } from '../utils/transformToMessages';
-import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import { ChatSidebar } from '@kbn/ai-playground/components/chat_sidebar';
 
 export const Chat = () => {
   const { euiTheme } = useEuiTheme();
@@ -153,53 +149,8 @@ export const Chat = () => {
           </EuiFlexGroup>
         </EuiFlexItem>
 
-        <EuiFlexItem
-          grow={1}
-          css={{
-            padding: euiTheme.size.l,
-          }}
-        >
-          <Controller
-            name={ChatFormFields.openAIKey}
-            control={control}
-            defaultValue=""
-            render={({ field }) => <OpenAIKeyField apiKey={field.value} onSave={field.onChange} />}
-          />
-
-          <Controller
-            name={ChatFormFields.prompt}
-            control={control}
-            defaultValue=""
-            render={({ field }) => (
-              <InstructionsField value={field.value} onChange={field.onChange} />
-            )}
-          />
-
-          <Controller
-            name={ChatFormFields.citations}
-            control={control}
-            defaultValue={true}
-            render={({ field }) => (
-              <IncludeCitationsField checked={field.value} onChange={field.onChange} />
-            )}
-          />
-
-          <Controller
-            name={ChatFormFields.indices}
-            control={control}
-            defaultValue={[]}
-            render={({ field }) => (
-              <SourcesPanelSidebar
-                selectedIndices={field.value}
-                addIndex={(newIndex: IndexName) => {
-                  field.onChange([...field.value, newIndex]);
-                }}
-                removeIndex={(index: IndexName) => {
-                  field.onChange(field.value.filter((indexName) => indexName !== index));
-                }}
-              />
-            )}
-          />
+        <EuiFlexItem grow={1}>
+          <ChatSidebar control={control} />
         </EuiFlexItem>
       </EuiFlexGroup>
     </EuiForm>
diff --git a/packages/kbn-ai-playground/components/chat_sidebar.tsx b/packages/kbn-ai-playground/components/chat_sidebar.tsx
new file mode 100644
index 0000000000000..0405f9373b006
--- /dev/null
+++ b/packages/kbn-ai-playground/components/chat_sidebar.tsx
@@ -0,0 +1,123 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import {
+  EuiAccordion,
+  EuiFlexGroup,
+  EuiFlexItem,
+  EuiSpacer,
+  EuiTitle,
+  useEuiTheme,
+  useGeneratedHtmlId,
+} from '@elastic/eui';
+import React, { useState } from 'react';
+import { i18n } from '@kbn/i18n';
+import { Control, Controller } from 'react-hook-form';
+import { ChatForm, ChatFormFields } from '../types';
+import { OpenAIKeyField } from './open_ai_key_field';
+import { InstructionsField } from './instructions_field';
+import { IncludeCitationsField } from './include_citations_field';
+import { SourcesPanelSidebar } from './sources_panel/sources_panel_sidebar';
+import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+
+interface ChatSidebarProps {
+  control: Control<ChatForm>;
+}
+
+export const ChatSidebar: React.FC<ChatSidebarProps> = ({ control }) => {
+  const { euiTheme } = useEuiTheme();
+  const accordions = [
+    {
+      id: useGeneratedHtmlId({ prefix: 'summarizationAccordion' }),
+      title: i18n.translate('aiPlayground.sidebar.summarizationTitle', {
+        defaultMessage: 'Summarization',
+      }),
+      children: (
+        <>
+          <Controller
+            name={ChatFormFields.openAIKey}
+            control={control}
+            defaultValue=""
+            render={({ field }) => <OpenAIKeyField apiKey={field.value} onSave={field.onChange} />}
+          />
+
+          <Controller
+            name={ChatFormFields.prompt}
+            control={control}
+            defaultValue=""
+            render={({ field }) => (
+              <InstructionsField value={field.value} onChange={field.onChange} />
+            )}
+          />
+
+          <Controller
+            name={ChatFormFields.citations}
+            control={control}
+            defaultValue={true}
+            render={({ field }) => (
+              <IncludeCitationsField checked={field.value} onChange={field.onChange} />
+            )}
+          />
+        </>
+      ),
+    },
+    {
+      id: useGeneratedHtmlId({ prefix: 'sourcesAccordion' }),
+      title: i18n.translate('aiPlayground.sidebar.sourceTitle', { defaultMessage: 'Sources' }),
+      children: (
+        <Controller
+          name={ChatFormFields.indices}
+          control={control}
+          defaultValue={[]}
+          render={({ field }) => (
+            <SourcesPanelSidebar
+              selectedIndices={field.value}
+              addIndex={(newIndex: IndexName) => {
+                field.onChange([...field.value, newIndex]);
+              }}
+              removeIndex={(index: IndexName) => {
+                field.onChange(field.value.filter((indexName: string) => indexName !== index));
+              }}
+            />
+          )}
+        />
+      ),
+    },
+  ];
+  const [openAccordionId, setOpenAccordionId] = useState(accordions[0].id);
+
+  return (
+    <EuiFlexGroup direction="column" gutterSize="none">
+      {accordions.map(({ id, title, children }, index) => (
+        <EuiFlexItem
+          key={id}
+          css={{
+            borderBottom: index === accordions.length - 1 ? 'none' : euiTheme.border.thin,
+            padding: `0 ${euiTheme.size.l}`,
+            flexGrow: openAccordionId === id ? 1 : 0,
+            transition: `${euiTheme.animation.normal} ease-in-out`,
+          }}
+        >
+          <EuiAccordion
+            id={id}
+            buttonContent={
+              <EuiTitle size="xs">
+                <h5>{title}</h5>
+              </EuiTitle>
+            }
+            buttonProps={{ paddingSize: 'l' }}
+            forceState={openAccordionId === id ? 'open' : 'closed'}
+            onToggle={() => setOpenAccordionId(openAccordionId === id ? '' : id)}
+          >
+            {children}
+            <EuiSpacer size="l" />
+          </EuiAccordion>
+        </EuiFlexItem>
+      ))}
+    </EuiFlexGroup>
+  );
+};
diff --git a/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx b/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx
index 2ebdb1fbe0294..00252761a5b82 100644
--- a/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx
+++ b/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx
@@ -7,14 +7,7 @@
 
 import React from 'react';
 import { i18n } from '@kbn/i18n';
-import {
-  EuiAccordion,
-  EuiCallOut,
-  EuiFlexGroup,
-  EuiFlexItem,
-  EuiTitle,
-  useGeneratedHtmlId,
-} from '@elastic/eui';
+import { EuiCallOut, EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 import { IndicesList } from './indices_list';
 import { AddIndicesField } from './add_indices_field';
 import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
@@ -29,37 +22,24 @@ export const SourcesPanelSidebar: React.FC<SourcesPanelSidebarProps> = ({
   selectedIndices,
   addIndex,
   removeIndex,
-}) => {
-  const accordionId = useGeneratedHtmlId({ prefix: 'sourceAccordion' });
+}) => (
+  <EuiFlexGroup direction="column">
+    <EuiFlexItem>
+      <EuiCallOut
+        title={i18n.translate('aiPlayground.sources.callout', {
+          defaultMessage: 'Changes here will reset your custom query',
+        })}
+        iconType="warning"
+        size="s"
+      />
+    </EuiFlexItem>
 
-  return (
-    <EuiAccordion
-      id={accordionId}
-      buttonContent={
-        <EuiTitle>
-          <h3>{i18n.translate('aiPlayground.sources.title', { defaultMessage: 'Sources' })}</h3>
-        </EuiTitle>
-      }
-    >
-      <EuiFlexGroup direction="column">
-        <EuiFlexItem>
-          <EuiCallOut
-            title={i18n.translate('aiPlayground.sources.callout', {
-              defaultMessage: 'Changes here will reset your custom query',
-            })}
-            iconType="warning"
-            size="s"
-          />
-        </EuiFlexItem>
+    <EuiFlexItem>
+      <IndicesList indices={selectedIndices} onRemoveClick={removeIndex} hasBorder />
+    </EuiFlexItem>
 
-        <EuiFlexItem>
-          <IndicesList indices={selectedIndices} onRemoveClick={removeIndex} hasBorder />
-        </EuiFlexItem>
-
-        <EuiFlexItem>
-          <AddIndicesField selectedIndices={selectedIndices} onIndexSelect={addIndex} />
-        </EuiFlexItem>
-      </EuiFlexGroup>
-    </EuiAccordion>
-  );
-};
+    <EuiFlexItem>
+      <AddIndicesField selectedIndices={selectedIndices} onIndexSelect={addIndex} />
+    </EuiFlexItem>
+  </EuiFlexGroup>
+);

From 932bff44fa64e575600b7c036d07ec9a5df1755b Mon Sep 17 00:00:00 2001
From: Yan Savitski <yan.savitski@elastic.co>
Date: Tue, 20 Feb 2024 22:16:24 +0100
Subject: [PATCH 2/3] [AI Playground] Add form provider to use chat and do
 refactor

---
 .../kbn-ai-playground/components/chat.tsx     | 165 +++++++++---------
 .../components/chat_sidebar.tsx               |  61 +------
 .../sources_panel/sources_panel_sidebar.tsx   |  60 ++++---
 .../include_citations_field.tsx               |   0
 .../instructions_field.tsx                    |   0
 .../open_ai_key_callout.tsx                   |   0
 .../open_ai_key_field.tsx                     |   0
 .../open_ai_key_flyout.tsx                    |   0
 .../summarization_panel.tsx                   |  44 +++++
 9 files changed, 164 insertions(+), 166 deletions(-)
 rename packages/kbn-ai-playground/components/{ => summarization_panel}/include_citations_field.tsx (100%)
 rename packages/kbn-ai-playground/components/{ => summarization_panel}/instructions_field.tsx (100%)
 rename packages/kbn-ai-playground/components/{ => summarization_panel}/open_ai_key_callout.tsx (100%)
 rename packages/kbn-ai-playground/components/{ => summarization_panel}/open_ai_key_field.tsx (100%)
 rename packages/kbn-ai-playground/components/{ => summarization_panel}/open_ai_key_flyout.tsx (100%)
 create mode 100644 packages/kbn-ai-playground/components/summarization_panel/summarization_panel.tsx

diff --git a/packages/kbn-ai-playground/components/chat.tsx b/packages/kbn-ai-playground/components/chat.tsx
index 0e132629a6274..8e76ace0d9f6f 100644
--- a/packages/kbn-ai-playground/components/chat.tsx
+++ b/packages/kbn-ai-playground/components/chat.tsx
@@ -8,7 +8,7 @@
 
 import React, { useMemo } from 'react';
 
-import { Controller, useForm } from 'react-hook-form';
+import { Controller, FormProvider, useForm } from 'react-hook-form';
 import {
   EuiButtonIcon,
   EuiFlexGroup,
@@ -35,12 +35,13 @@ import { ChatSidebar } from '@kbn/ai-playground/components/chat_sidebar';
 
 export const Chat = () => {
   const { euiTheme } = useEuiTheme();
+  const form = useForm<ChatForm>();
   const {
     control,
     formState: { isValid, isSubmitting },
     resetField,
     handleSubmit,
-  } = useForm<ChatForm>();
+  } = form;
   const { messages, append, stop: stopRequest } = useChat();
 
   const onSubmit = async (data: ChatForm) => {
@@ -71,88 +72,90 @@ export const Chat = () => {
   );
 
   return (
-    <EuiForm
-      component="form"
-      css={{ display: 'flex', flexGrow: 1 }}
-      onSubmit={handleSubmit(onSubmit)}
-    >
-      <EuiFlexGroup gutterSize="none">
-        <EuiFlexItem
-          grow={2}
-          css={{
-            borderRight: euiTheme.border.thin,
-            paddingTop: euiTheme.size.l,
-            paddingBottom: euiTheme.size.l,
-          }}
-        >
-          <EuiFlexGroup direction="column" className="eui-fullHeight">
-            {/* // Set scroll at the border of parent element*/}
-            <EuiFlexItem
-              grow={1}
-              className="eui-yScroll"
-              css={{ paddingLeft: euiTheme.size.l, paddingRight: euiTheme.size.l }}
-              tabIndex={0}
-            >
-              <MessageList messages={chatMessages} />
-            </EuiFlexItem>
+    <FormProvider {...form}>
+      <EuiForm
+        component="form"
+        css={{ display: 'flex', flexGrow: 1 }}
+        onSubmit={handleSubmit(onSubmit)}
+      >
+        <EuiFlexGroup gutterSize="none">
+          <EuiFlexItem
+            grow={2}
+            css={{
+              borderRight: euiTheme.border.thin,
+              paddingTop: euiTheme.size.l,
+              paddingBottom: euiTheme.size.l,
+            }}
+          >
+            <EuiFlexGroup direction="column" className="eui-fullHeight">
+              {/* // Set scroll at the border of parent element*/}
+              <EuiFlexItem
+                grow={1}
+                className="eui-yScroll"
+                css={{ paddingLeft: euiTheme.size.l, paddingRight: euiTheme.size.l }}
+                tabIndex={0}
+              >
+                <MessageList messages={chatMessages} />
+              </EuiFlexItem>
 
-            <EuiFlexItem
-              grow={false}
-              css={{ paddingLeft: euiTheme.size.l, paddingRight: euiTheme.size.l }}
-            >
-              <EuiHorizontalRule margin="none" />
+              <EuiFlexItem
+                grow={false}
+                css={{ paddingLeft: euiTheme.size.l, paddingRight: euiTheme.size.l }}
+              >
+                <EuiHorizontalRule margin="none" />
 
-              <EuiSpacer size="m" />
+                <EuiSpacer size="m" />
 
-              <Controller
-                name={ChatFormFields.question}
-                control={control}
-                defaultValue=""
-                rules={{
-                  required: true,
-                  validate: (rule) => !!rule?.trim(),
-                }}
-                render={({ field }) => (
-                  <QuestionInput
-                    value={field.value}
-                    onChange={field.onChange}
-                    isDisabled={isSubmitting}
-                    button={
-                      isSubmitting ? (
-                        <EuiButtonIcon
-                          aria-label={i18n.translate('aiPlayground.chat.stopButtonAriaLabel', {
-                            defaultMessage: 'Stop request',
-                          })}
-                          display="base"
-                          size="s"
-                          iconType="stop"
-                          onClick={stopRequest}
-                        />
-                      ) : (
-                        <EuiButtonIcon
-                          aria-label={i18n.translate('aiPlayground.chat.sendButtonAriaLabel', {
-                            defaultMessage: 'Send a question',
-                          })}
-                          display={isValid ? 'base' : 'empty'}
-                          size="s"
-                          type="submit"
-                          isLoading={isSubmitting}
-                          isDisabled={!isValid}
-                          iconType={TelegramIcon}
-                        />
-                      )
-                    }
-                  />
-                )}
-              />
-            </EuiFlexItem>
-          </EuiFlexGroup>
-        </EuiFlexItem>
+                <Controller
+                  name={ChatFormFields.question}
+                  control={control}
+                  defaultValue=""
+                  rules={{
+                    required: true,
+                    validate: (rule) => !!rule?.trim(),
+                  }}
+                  render={({ field }) => (
+                    <QuestionInput
+                      value={field.value}
+                      onChange={field.onChange}
+                      isDisabled={isSubmitting}
+                      button={
+                        isSubmitting ? (
+                          <EuiButtonIcon
+                            aria-label={i18n.translate('aiPlayground.chat.stopButtonAriaLabel', {
+                              defaultMessage: 'Stop request',
+                            })}
+                            display="base"
+                            size="s"
+                            iconType="stop"
+                            onClick={stopRequest}
+                          />
+                        ) : (
+                          <EuiButtonIcon
+                            aria-label={i18n.translate('aiPlayground.chat.sendButtonAriaLabel', {
+                              defaultMessage: 'Send a question',
+                            })}
+                            display={isValid ? 'base' : 'empty'}
+                            size="s"
+                            type="submit"
+                            isLoading={isSubmitting}
+                            isDisabled={!isValid}
+                            iconType={TelegramIcon}
+                          />
+                        )
+                      }
+                    />
+                  )}
+                />
+              </EuiFlexItem>
+            </EuiFlexGroup>
+          </EuiFlexItem>
 
-        <EuiFlexItem grow={1}>
-          <ChatSidebar control={control} />
-        </EuiFlexItem>
-      </EuiFlexGroup>
-    </EuiForm>
+          <EuiFlexItem grow={1}>
+            <ChatSidebar />
+          </EuiFlexItem>
+        </EuiFlexGroup>
+      </EuiForm>
+    </FormProvider>
   );
 };
diff --git a/packages/kbn-ai-playground/components/chat_sidebar.tsx b/packages/kbn-ai-playground/components/chat_sidebar.tsx
index 0405f9373b006..e9a9dfed6630b 100644
--- a/packages/kbn-ai-playground/components/chat_sidebar.tsx
+++ b/packages/kbn-ai-playground/components/chat_sidebar.tsx
@@ -16,19 +16,10 @@ import {
 } from '@elastic/eui';
 import React, { useState } from 'react';
 import { i18n } from '@kbn/i18n';
-import { Control, Controller } from 'react-hook-form';
-import { ChatForm, ChatFormFields } from '../types';
-import { OpenAIKeyField } from './open_ai_key_field';
-import { InstructionsField } from './instructions_field';
-import { IncludeCitationsField } from './include_citations_field';
 import { SourcesPanelSidebar } from './sources_panel/sources_panel_sidebar';
-import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import { SummarizationPanel } from './summarization_panel/summarization_panel';
 
-interface ChatSidebarProps {
-  control: Control<ChatForm>;
-}
-
-export const ChatSidebar: React.FC<ChatSidebarProps> = ({ control }) => {
+export const ChatSidebar: React.FC = () => {
   const { euiTheme } = useEuiTheme();
   const accordions = [
     {
@@ -36,56 +27,12 @@ export const ChatSidebar: React.FC<ChatSidebarProps> = ({ control }) => {
       title: i18n.translate('aiPlayground.sidebar.summarizationTitle', {
         defaultMessage: 'Summarization',
       }),
-      children: (
-        <>
-          <Controller
-            name={ChatFormFields.openAIKey}
-            control={control}
-            defaultValue=""
-            render={({ field }) => <OpenAIKeyField apiKey={field.value} onSave={field.onChange} />}
-          />
-
-          <Controller
-            name={ChatFormFields.prompt}
-            control={control}
-            defaultValue=""
-            render={({ field }) => (
-              <InstructionsField value={field.value} onChange={field.onChange} />
-            )}
-          />
-
-          <Controller
-            name={ChatFormFields.citations}
-            control={control}
-            defaultValue={true}
-            render={({ field }) => (
-              <IncludeCitationsField checked={field.value} onChange={field.onChange} />
-            )}
-          />
-        </>
-      ),
+      children: <SummarizationPanel />,
     },
     {
       id: useGeneratedHtmlId({ prefix: 'sourcesAccordion' }),
       title: i18n.translate('aiPlayground.sidebar.sourceTitle', { defaultMessage: 'Sources' }),
-      children: (
-        <Controller
-          name={ChatFormFields.indices}
-          control={control}
-          defaultValue={[]}
-          render={({ field }) => (
-            <SourcesPanelSidebar
-              selectedIndices={field.value}
-              addIndex={(newIndex: IndexName) => {
-                field.onChange([...field.value, newIndex]);
-              }}
-              removeIndex={(index: IndexName) => {
-                field.onChange(field.value.filter((indexName: string) => indexName !== index));
-              }}
-            />
-          )}
-        />
-      ),
+      children: <SourcesPanelSidebar />,
     },
   ];
   const [openAccordionId, setOpenAccordionId] = useState(accordions[0].id);
diff --git a/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx b/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx
index 00252761a5b82..279233a55c703 100644
--- a/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx
+++ b/packages/kbn-ai-playground/components/sources_panel/sources_panel_sidebar.tsx
@@ -11,35 +11,39 @@ import { EuiCallOut, EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 import { IndicesList } from './indices_list';
 import { AddIndicesField } from './add_indices_field';
 import { IndexName } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import { useController } from 'react-hook-form';
+import { ChatFormFields } from '../../types';
 
-interface SourcesPanelSidebarProps {
-  selectedIndices: IndexName[];
-  addIndex: (newIndex: IndexName) => void;
-  removeIndex: (index: IndexName) => void;
-}
+export const SourcesPanelSidebar: React.FC = () => {
+  const {
+    field: { value: selectedIndices, onChange },
+  } = useController({ name: ChatFormFields.indices, defaultValue: [] });
+  const addIndex = (newIndex: IndexName) => {
+    onChange([...selectedIndices, newIndex]);
+  };
+  const removeIndex = (index: IndexName) => {
+    onChange(selectedIndices.filter((indexName: string) => indexName !== index));
+  };
 
-export const SourcesPanelSidebar: React.FC<SourcesPanelSidebarProps> = ({
-  selectedIndices,
-  addIndex,
-  removeIndex,
-}) => (
-  <EuiFlexGroup direction="column">
-    <EuiFlexItem>
-      <EuiCallOut
-        title={i18n.translate('aiPlayground.sources.callout', {
-          defaultMessage: 'Changes here will reset your custom query',
-        })}
-        iconType="warning"
-        size="s"
-      />
-    </EuiFlexItem>
+  return (
+    <EuiFlexGroup direction="column">
+      <EuiFlexItem>
+        <EuiCallOut
+          title={i18n.translate('aiPlayground.sources.callout', {
+            defaultMessage: 'Changes here will reset your custom query',
+          })}
+          iconType="warning"
+          size="s"
+        />
+      </EuiFlexItem>
 
-    <EuiFlexItem>
-      <IndicesList indices={selectedIndices} onRemoveClick={removeIndex} hasBorder />
-    </EuiFlexItem>
+      <EuiFlexItem>
+        <IndicesList indices={selectedIndices} onRemoveClick={removeIndex} hasBorder />
+      </EuiFlexItem>
 
-    <EuiFlexItem>
-      <AddIndicesField selectedIndices={selectedIndices} onIndexSelect={addIndex} />
-    </EuiFlexItem>
-  </EuiFlexGroup>
-);
+      <EuiFlexItem>
+        <AddIndicesField selectedIndices={selectedIndices} onIndexSelect={addIndex} />
+      </EuiFlexItem>
+    </EuiFlexGroup>
+  );
+};
diff --git a/packages/kbn-ai-playground/components/include_citations_field.tsx b/packages/kbn-ai-playground/components/summarization_panel/include_citations_field.tsx
similarity index 100%
rename from packages/kbn-ai-playground/components/include_citations_field.tsx
rename to packages/kbn-ai-playground/components/summarization_panel/include_citations_field.tsx
diff --git a/packages/kbn-ai-playground/components/instructions_field.tsx b/packages/kbn-ai-playground/components/summarization_panel/instructions_field.tsx
similarity index 100%
rename from packages/kbn-ai-playground/components/instructions_field.tsx
rename to packages/kbn-ai-playground/components/summarization_panel/instructions_field.tsx
diff --git a/packages/kbn-ai-playground/components/open_ai_key_callout.tsx b/packages/kbn-ai-playground/components/summarization_panel/open_ai_key_callout.tsx
similarity index 100%
rename from packages/kbn-ai-playground/components/open_ai_key_callout.tsx
rename to packages/kbn-ai-playground/components/summarization_panel/open_ai_key_callout.tsx
diff --git a/packages/kbn-ai-playground/components/open_ai_key_field.tsx b/packages/kbn-ai-playground/components/summarization_panel/open_ai_key_field.tsx
similarity index 100%
rename from packages/kbn-ai-playground/components/open_ai_key_field.tsx
rename to packages/kbn-ai-playground/components/summarization_panel/open_ai_key_field.tsx
diff --git a/packages/kbn-ai-playground/components/open_ai_key_flyout.tsx b/packages/kbn-ai-playground/components/summarization_panel/open_ai_key_flyout.tsx
similarity index 100%
rename from packages/kbn-ai-playground/components/open_ai_key_flyout.tsx
rename to packages/kbn-ai-playground/components/summarization_panel/open_ai_key_flyout.tsx
diff --git a/packages/kbn-ai-playground/components/summarization_panel/summarization_panel.tsx b/packages/kbn-ai-playground/components/summarization_panel/summarization_panel.tsx
new file mode 100644
index 0000000000000..3673c66112a20
--- /dev/null
+++ b/packages/kbn-ai-playground/components/summarization_panel/summarization_panel.tsx
@@ -0,0 +1,44 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React from 'react';
+import { Controller, useFormContext } from 'react-hook-form';
+import { IncludeCitationsField } from './include_citations_field';
+import { InstructionsField } from './instructions_field';
+import { OpenAIKeyField } from './open_ai_key_field';
+import { ChatFormFields } from '../../types';
+
+export const SummarizationPanel: React.FC = () => {
+  const { control } = useFormContext();
+
+  return (
+    <>
+      <Controller
+        name={ChatFormFields.openAIKey}
+        control={control}
+        defaultValue=""
+        render={({ field }) => <OpenAIKeyField apiKey={field.value} onSave={field.onChange} />}
+      />
+
+      <Controller
+        name={ChatFormFields.prompt}
+        control={control}
+        defaultValue=""
+        render={({ field }) => <InstructionsField value={field.value} onChange={field.onChange} />}
+      />
+
+      <Controller
+        name={ChatFormFields.citations}
+        control={control}
+        defaultValue={true}
+        render={({ field }) => (
+          <IncludeCitationsField checked={field.value} onChange={field.onChange} />
+        )}
+      />
+    </>
+  );
+};

From 13342e128dbe05b9d0078a15813536bbd1e776b4 Mon Sep 17 00:00:00 2001
From: Yan Savitski <yan.savitski@elastic.co>
Date: Tue, 20 Feb 2024 22:40:41 +0100
Subject: [PATCH 3/3] [AI Playground] Add extra action for sources sidebar

---
 .../kbn-ai-playground/components/chat.tsx     |  4 +++-
 .../components/chat_sidebar.tsx               | 20 +++++++++++++++++--
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/packages/kbn-ai-playground/components/chat.tsx b/packages/kbn-ai-playground/components/chat.tsx
index 8e76ace0d9f6f..1dbebe4a34fd2 100644
--- a/packages/kbn-ai-playground/components/chat.tsx
+++ b/packages/kbn-ai-playground/components/chat.tsx
@@ -38,11 +38,13 @@ export const Chat = () => {
   const form = useForm<ChatForm>();
   const {
     control,
+    watch,
     formState: { isValid, isSubmitting },
     resetField,
     handleSubmit,
   } = form;
   const { messages, append, stop: stopRequest } = useChat();
+  const selectedIndicesCount = watch(ChatFormFields.indices, []).length;
 
   const onSubmit = async (data: ChatForm) => {
     await append(
@@ -152,7 +154,7 @@ export const Chat = () => {
           </EuiFlexItem>
 
           <EuiFlexItem grow={1}>
-            <ChatSidebar />
+            <ChatSidebar selectedIndicesCount={selectedIndicesCount} />
           </EuiFlexItem>
         </EuiFlexGroup>
       </EuiForm>
diff --git a/packages/kbn-ai-playground/components/chat_sidebar.tsx b/packages/kbn-ai-playground/components/chat_sidebar.tsx
index e9a9dfed6630b..0fba4df001805 100644
--- a/packages/kbn-ai-playground/components/chat_sidebar.tsx
+++ b/packages/kbn-ai-playground/components/chat_sidebar.tsx
@@ -10,6 +10,7 @@ import {
   EuiFlexGroup,
   EuiFlexItem,
   EuiSpacer,
+  EuiText,
   EuiTitle,
   useEuiTheme,
   useGeneratedHtmlId,
@@ -19,7 +20,11 @@ import { i18n } from '@kbn/i18n';
 import { SourcesPanelSidebar } from './sources_panel/sources_panel_sidebar';
 import { SummarizationPanel } from './summarization_panel/summarization_panel';
 
-export const ChatSidebar: React.FC = () => {
+interface ChatSidebarProps {
+  selectedIndicesCount: number;
+}
+
+export const ChatSidebar: React.FC<ChatSidebarProps> = ({ selectedIndicesCount }) => {
   const { euiTheme } = useEuiTheme();
   const accordions = [
     {
@@ -32,6 +37,16 @@ export const ChatSidebar: React.FC = () => {
     {
       id: useGeneratedHtmlId({ prefix: 'sourcesAccordion' }),
       title: i18n.translate('aiPlayground.sidebar.sourceTitle', { defaultMessage: 'Sources' }),
+      extraAction: !!selectedIndicesCount && (
+        <EuiText size="xs">
+          <p>
+            {i18n.translate('aiPlayground.sidebar.sourceIndicesCount', {
+              defaultMessage: '{count, number} {count, plural, one {Index} other {Indices}}',
+              values: { count: Number(selectedIndicesCount) },
+            })}
+          </p>
+        </EuiText>
+      ),
       children: <SourcesPanelSidebar />,
     },
   ];
@@ -39,7 +54,7 @@ export const ChatSidebar: React.FC = () => {
 
   return (
     <EuiFlexGroup direction="column" gutterSize="none">
-      {accordions.map(({ id, title, children }, index) => (
+      {accordions.map(({ id, title, extraAction, children }, index) => (
         <EuiFlexItem
           key={id}
           css={{
@@ -56,6 +71,7 @@ export const ChatSidebar: React.FC = () => {
                 <h5>{title}</h5>
               </EuiTitle>
             }
+            extraAction={extraAction}
             buttonProps={{ paddingSize: 'l' }}
             forceState={openAccordionId === id ? 'open' : 'closed'}
             onToggle={() => setOpenAccordionId(openAccordionId === id ? '' : id)}