[Search][Playground] Support Multiple Context Fields (#210703)

## Summary

This PR updates the search playground to allow selecting > 1 context
fields to be included in context documents for the LLM.

### Screenshots
<img width="1399" alt="image"
src="https://github.com/user-attachments/assets/76c6bd84-1dc6-4862-b822-a7fc3595cd69"
/>


Context Fields Updated to ComboBox:
<img width="384" alt="image"
src="https://github.com/user-attachments/assets/e246628b-4952-4832-9ac3-f2203700a667"
/>

### Checklist

- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [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
- [ ] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed

---------

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

x-pack/solutions/search/plugins/search_playground/public/components/edit_context/context_fields_select.tsx

@@ -0,0 +1,85 @@
+/*
+ * 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, { useMemo, useCallback } from 'react';
+
+import { EuiCallOut, EuiComboBox, EuiComboBoxOptionOption } from '@elastic/eui';
+import { i18n } from '@kbn/i18n';
+
+import { QuerySourceFields } from '../../types';
+
+export interface ContextFieldsSelectProps {
+  indexName: string;
+  indexFields: QuerySourceFields;
+  selectedContextFields?: string[];
+  updateSelectedContextFields: (index: string, value: string[]) => void;
+}
+
+export const ContextFieldsSelect = ({
+  indexName,
+  indexFields,
+  selectedContextFields,
+  updateSelectedContextFields,
+}: ContextFieldsSelectProps) => {
+  const { options: selectOptions, selectedOptions } = useMemo(() => {
+    if (!indexFields.source_fields?.length) return { options: [], selectedOptions: [] };
+
+    const options: Array<EuiComboBoxOptionOption<unknown>> = indexFields.source_fields.map(
+      (field) => ({
+        label: field,
+        'data-test-subj': `contextField-${field}`,
+      })
+    );
+    const selected: Array<EuiComboBoxOptionOption<unknown>> =
+      selectedContextFields
+        ?.map((field) => options.find((opt) => opt.label === field))
+        ?.filter(
+          (
+            val: EuiComboBoxOptionOption<unknown> | undefined
+          ): val is EuiComboBoxOptionOption<unknown> => val !== undefined
+        ) ?? [];
+    return {
+      options,
+      selectedOptions: selected,
+    };
+  }, [indexFields.source_fields, selectedContextFields]);
+  const onSelectFields = useCallback(
+    (updatedSelectedOptions: Array<EuiComboBoxOptionOption<unknown>>) => {
+      // always require at least 1 selected field
+      if (updatedSelectedOptions.length === 0) return;
+      updateSelectedContextFields(
+        indexName,
+        updatedSelectedOptions.map((opt) => opt.label)
+      );
+    },
+    [indexName, updateSelectedContextFields]
+  );
+
+  if (selectOptions.length === 0) {
+    return (
+      <EuiCallOut
+        title={i18n.translate('xpack.searchPlayground.editContext.noSourceFieldWarning', {
+          defaultMessage: 'No source fields found',
+        })}
+        color="warning"
+        iconType="warning"
+        size="s"
+      />
+    );
+  }
+
+  return (
+    <EuiComboBox
+      data-test-subj={`contextFieldsSelectable-${indexName}`}
+      options={selectOptions}
+      selectedOptions={selectedOptions}
+      onChange={onSelectFields}
+      isClearable={false}
+      fullWidth
+    />
+  );
+};

x-pack/solutions/search/plugins/search_playground/public/components/edit_context/edit_context_panel.test.tsx

@@ -19,14 +19,14 @@ jest.mock('../../hooks/use_source_indices_field', () => ({
         elser_query_fields: [],
         dense_vector_query_fields: [],
         bm25_query_fields: ['field1', 'field2'],
-        source_fields: ['context_field1', 'context_field2'],
+        source_fields: ['title', 'description'],
         semantic_fields: [],
       },
       index2: {
         elser_query_fields: [],
         dense_vector_query_fields: [],
-        bm25_query_fields: ['field1', 'field2'],
-        source_fields: ['context_field1', 'context_field2'],
+        bm25_query_fields: ['foo', 'bar'],
+        source_fields: ['body'],
         semantic_fields: [],
       },
     },
@@ -47,8 +47,8 @@ const MockFormProvider = ({ children }: { children: React.ReactElement }) => {
       [ChatFormFields.indices]: ['index1'],
       [ChatFormFields.docSize]: 1,
       [ChatFormFields.sourceFields]: {
-        index1: ['context_field1'],
-        index2: ['context_field2'],
+        index1: ['title'],
+        index2: ['body'],
       },
     },
   });
@@ -67,9 +67,15 @@ describe('EditContextFlyout component tests', () => {
   });
 
   it('should see the context fields', async () => {
-    expect(screen.getByTestId('contextFieldsSelectable-0')).toBeInTheDocument();
-    fireEvent.click(screen.getByTestId('contextFieldsSelectable-0'));
-    const fields = await screen.findAllByTestId('contextField');
-    expect(fields.length).toBe(2);
+    expect(screen.getByTestId('contextFieldsSelectable-index1')).toBeInTheDocument();
+    const listButton = screen
+      .getByTestId('contextFieldsSelectable-index1')
+      .querySelector('[data-test-subj="comboBoxToggleListButton"]');
+    expect(listButton).not.toBeNull();
+    fireEvent.click(listButton!);
+
+    for (const field of ['title', 'description']) {
+      expect(screen.getByTestId(`contextField-${field}`)).toBeInTheDocument();
+    }
   });
 });

x-pack/solutions/search/plugins/search_playground/public/components/edit_context/edit_context_panel.tsx

@@ -5,24 +5,16 @@
  * 2.0.
  */
 
-import {
-  EuiFlexGroup,
-  EuiFlexItem,
-  EuiFormRow,
-  EuiPanel,
-  EuiSelect,
-  EuiSuperSelect,
-  EuiText,
-  EuiCallOut,
-} from '@elastic/eui';
+import { EuiFlexGroup, EuiFlexItem, EuiFormRow, EuiPanel, EuiSelect, EuiText } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 import { FormattedMessage } from '@kbn/i18n-react';
-import React from 'react';
+import React, { useCallback } from 'react';
 import { useController } from 'react-hook-form';
 import { useSourceIndicesFields } from '../../hooks/use_source_indices_field';
 import { useUsageTracker } from '../../hooks/use_usage_tracker';
 import { ChatForm, ChatFormFields } from '../../types';
 import { AnalyticsEvents } from '../../analytics/constants';
+import { ContextFieldsSelect } from './context_fields_select';
 
 export const EditContextPanel: React.FC = () => {
   const usageTracker = useUsageTracker();
@@ -40,13 +32,16 @@ export const EditContextPanel: React.FC = () => {
     name: ChatFormFields.sourceFields,
   });
 
-  const updateSourceField = (index: string, field: string) => {
-    onChangeSourceFields({
-      ...sourceFields,
-      [index]: [field],
-    });
-    usageTracker?.click(AnalyticsEvents.editContextFieldToggled);
-  };
+  const updateSourceField = useCallback(
+    (index: string, contextFields: string[]) => {
+      onChangeSourceFields({
+        ...sourceFields,
+        [index]: contextFields,
+      });
+      usageTracker?.click(AnalyticsEvents.editContextFieldToggled);
+    },
+    [onChangeSourceFields, sourceFields, usageTracker]
+  );
 
   const handleDocSizeChange = (e: React.ChangeEvent<HTMLSelectElement>) => {
     usageTracker?.click(AnalyticsEvents.editContextDocSizeChanged);
@@ -64,6 +59,7 @@ export const EditContextPanel: React.FC = () => {
             fullWidth
           >
             <EuiSelect
+              data-test-subj="contextPanelDocumentNumberSelect"
               options={[
                 {
                   value: 1,
@@ -100,32 +96,15 @@ export const EditContextPanel: React.FC = () => {
                 </h5>
               </EuiText>
             </EuiFlexItem>
-            {Object.entries(fields).map(([index, group], indexNum) => (
+            {Object.entries(fields).map(([index, group]) => (
               <EuiFlexItem grow={false} key={index}>
                 <EuiFormRow label={index} fullWidth>
-                  {!!group.source_fields?.length ? (
-                    <EuiSuperSelect
-                      data-test-subj={`contextFieldsSelectable-${indexNum}`}
-                      options={group.source_fields.map((field) => ({
-                        value: field,
-                        inputDisplay: field,
-                        'data-test-subj': 'contextField',
-                      }))}
-                      valueOfSelected={sourceFields[index]?.[0]}
-                      onChange={(value) => updateSourceField(index, value)}
-                      fullWidth
-                    />
-                  ) : (
-                    <EuiCallOut
-                      title={i18n.translate(
-                        'xpack.searchPlayground.editContext.noSourceFieldWarning',
-                        { defaultMessage: 'No source fields found' }
-                      )}
-                      color="warning"
-                      iconType="warning"
-                      size="s"
-                    />
-                  )}
+                  <ContextFieldsSelect
+                    indexName={index}
+                    indexFields={group}
+                    selectedContextFields={sourceFields[index] ?? []}
+                    updateSelectedContextFields={updateSourceField}
+                  />
                 </EuiFormRow>
               </EuiFlexItem>
             ))}