diff --git a/src/plugins/dashboard/public/application/components/dashboard_top_nav/__snapshots__/dashboard_top_nav.test.tsx.snap b/src/plugins/dashboard/public/application/components/dashboard_top_nav/__snapshots__/dashboard_top_nav.test.tsx.snap
index 9b5860426ffd..2ab763d87c7a 100644
--- a/src/plugins/dashboard/public/application/components/dashboard_top_nav/__snapshots__/dashboard_top_nav.test.tsx.snap
+++ b/src/plugins/dashboard/public/application/components/dashboard_top_nav/__snapshots__/dashboard_top_nav.test.tsx.snap
@@ -317,11 +317,13 @@ exports[`Dashboard top nav render in embed mode 1`] = `
               },
               "query": Object {
                 "addToQueryLog": [MockFunction],
-                "dataSet": Object {
+                "dataSetManager": Object {
+                  "fetchDefaultDataSet": [MockFunction],
                   "getDataSet": [MockFunction],
                   "getDefaultDataSet": [MockFunction],
                   "getUpdates$": [MockFunction],
                   "init": [MockFunction],
+                  "initWithIndexPattern": [MockFunction],
                   "setDataSet": [MockFunction],
                 },
                 "filterManager": Object {
@@ -1338,11 +1340,13 @@ exports[`Dashboard top nav render in embed mode, and force hide filter bar 1`] =
               },
               "query": Object {
                 "addToQueryLog": [MockFunction],
-                "dataSet": Object {
+                "dataSetManager": Object {
+                  "fetchDefaultDataSet": [MockFunction],
                   "getDataSet": [MockFunction],
                   "getDefaultDataSet": [MockFunction],
                   "getUpdates$": [MockFunction],
                   "init": [MockFunction],
+                  "initWithIndexPattern": [MockFunction],
                   "setDataSet": [MockFunction],
                 },
                 "filterManager": Object {
@@ -2359,11 +2363,13 @@ exports[`Dashboard top nav render in embed mode, components can be forced show b
               },
               "query": Object {
                 "addToQueryLog": [MockFunction],
-                "dataSet": Object {
+                "dataSetManager": Object {
+                  "fetchDefaultDataSet": [MockFunction],
                   "getDataSet": [MockFunction],
                   "getDefaultDataSet": [MockFunction],
                   "getUpdates$": [MockFunction],
                   "init": [MockFunction],
+                  "initWithIndexPattern": [MockFunction],
                   "setDataSet": [MockFunction],
                 },
                 "filterManager": Object {
@@ -3380,11 +3386,13 @@ exports[`Dashboard top nav render in full screen mode with appended URL param bu
               },
               "query": Object {
                 "addToQueryLog": [MockFunction],
-                "dataSet": Object {
+                "dataSetManager": Object {
+                  "fetchDefaultDataSet": [MockFunction],
                   "getDataSet": [MockFunction],
                   "getDefaultDataSet": [MockFunction],
                   "getUpdates$": [MockFunction],
                   "init": [MockFunction],
+                  "initWithIndexPattern": [MockFunction],
                   "setDataSet": [MockFunction],
                 },
                 "filterManager": Object {
@@ -4401,11 +4409,13 @@ exports[`Dashboard top nav render in full screen mode, no componenets should be
               },
               "query": Object {
                 "addToQueryLog": [MockFunction],
-                "dataSet": Object {
+                "dataSetManager": Object {
+                  "fetchDefaultDataSet": [MockFunction],
                   "getDataSet": [MockFunction],
                   "getDefaultDataSet": [MockFunction],
                   "getUpdates$": [MockFunction],
                   "init": [MockFunction],
+                  "initWithIndexPattern": [MockFunction],
                   "setDataSet": [MockFunction],
                 },
                 "filterManager": Object {
@@ -5422,11 +5432,13 @@ exports[`Dashboard top nav render with all components 1`] = `
               },
               "query": Object {
                 "addToQueryLog": [MockFunction],
-                "dataSet": Object {
+                "dataSetManager": Object {
+                  "fetchDefaultDataSet": [MockFunction],
                   "getDataSet": [MockFunction],
                   "getDefaultDataSet": [MockFunction],
                   "getUpdates$": [MockFunction],
                   "init": [MockFunction],
+                  "initWithIndexPattern": [MockFunction],
                   "setDataSet": [MockFunction],
                 },
                 "filterManager": Object {
diff --git a/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts b/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts
index 2f1f5144274c..93dcdfb22b82 100644
--- a/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts
+++ b/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts
@@ -12,6 +12,8 @@ const createSetupContractMock = () => {
     setDataSet: jest.fn(),
     getUpdates$: jest.fn(),
     getDefaultDataSet: jest.fn(),
+    fetchDefaultDataSet: jest.fn(),
+    initWithIndexPattern: jest.fn(),
   };
   return dataSetManagerMock;
 };
diff --git a/src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts b/src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts
index d989eb057686..8f87f063d9be 100644
--- a/src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts
+++ b/src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts
@@ -15,11 +15,11 @@ describe('DataSetManager', () => {
     service = new DataSetManager(uiSettingsMock);
   });
 
-  test('getUpdates$ emits initially and after data set changes', () => {
+  test('getUpdates$ is a cold emits only after dataset changes', () => {
     const obs$ = service.getUpdates$();
-    const emittedValues: Array<SimpleDataSet | undefined> = [];
+    const emittedValues: SimpleDataSet[] = [];
     obs$.subscribe((v) => {
-      emittedValues.push(v);
+      emittedValues.push(v!);
     });
     expect(emittedValues).toHaveLength(0);
     expect(emittedValues[0]).toEqual(undefined);
diff --git a/src/plugins/data/public/query/mocks.ts b/src/plugins/data/public/query/mocks.ts
index 200ef46a5175..87bd2edb49bc 100644
--- a/src/plugins/data/public/query/mocks.ts
+++ b/src/plugins/data/public/query/mocks.ts
@@ -42,7 +42,7 @@ const createSetupContractMock = () => {
     filterManager: createFilterManagerMock(),
     timefilter: timefilterServiceMock.createSetupContract(),
     queryString: queryStringManagerMock.createSetupContract(),
-    dataSet: dataSetManagerMock.createSetupContract(),
+    dataSetManager: dataSetManagerMock.createSetupContract(),
     state$: new Observable(),
   };
 
@@ -57,7 +57,7 @@ const createStartContractMock = () => {
     savedQueries: jest.fn() as any,
     state$: new Observable(),
     timefilter: timefilterServiceMock.createStartContract(),
-    dataSet: dataSetManagerMock.createStartContract(),
+    dataSetManager: dataSetManagerMock.createStartContract(),
     getOpenSearchQuery: jest.fn(),
   };
 
diff --git a/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx b/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx
index 94674ab2bcff..b6a6274e13cf 100644
--- a/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx
+++ b/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx
@@ -15,7 +15,7 @@ export function createDataSetNavigator(
   dataSetManager: DataSetContract
 ) {
   // Return a function that takes props, omitting the dependencies from the props type
-  return (props: Omit<DataSetNavigatorProps, 'savedObjectsClient' | 'http' | 'dataSet'>) => (
+  return (props: Omit<DataSetNavigatorProps, 'savedObjectsClient' | 'http' | 'dataSetManager'>) => (
     <DataSetNavigator
       {...props}
       savedObjectsClient={savedObjectsClient}
diff --git a/src/plugins/data/public/ui/query_editor/query_editor.tsx b/src/plugins/data/public/ui/query_editor/query_editor.tsx
index 476a05fe4965..3abcfac2d340 100644
--- a/src/plugins/data/public/ui/query_editor/query_editor.tsx
+++ b/src/plugins/data/public/ui/query_editor/query_editor.tsx
@@ -406,6 +406,10 @@ export default class QueryEditorUI extends Component<Props, State> {
           {languageSelector}
           {this.props.queryActions}
         </div>
+        <div
+          ref={this.headerRef}
+          className={classNames('osdQueryEditor__header', this.props.headerClassName)}
+        />
         {!this.state.isCollapsed && (
           <div className="osdQueryEditor__body">{languageEditor.Body()}</div>
         )}
@@ -505,7 +509,6 @@ export default class QueryEditorUI extends Component<Props, State> {
            </EuiFlexItem>
 
            <EuiFlexItem onClick={this.onClickInput} grow={true}>
-             <div ref={this.headerRef} className={headerClassName} />
              {!this.state.isCollapsed && useQueryEditor && (
                <CodeEditor
                  height={70}
diff --git a/src/plugins/query_enhancements/public/index.scss b/src/plugins/query_enhancements/public/index.scss
index ff7112406eac..f2424381b9f1 100644
--- a/src/plugins/query_enhancements/public/index.scss
+++ b/src/plugins/query_enhancements/public/index.scss
@@ -1 +1,6 @@
-/* stylelint-disable no-empty-source */
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+@import "./query_assist";
diff --git a/src/plugins/query_enhancements/public/query_assist/_index.scss b/src/plugins/query_enhancements/public/query_assist/_index.scss
new file mode 100644
index 000000000000..f80ad45850c9
--- /dev/null
+++ b/src/plugins/query_enhancements/public/query_assist/_index.scss
@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+.queryAssist {
+  &.queryAssist__form {
+    padding: 0 $euiSizeXS $euiSizeXS;
+  }
+
+  &.queryAssist__callout {
+    margin-top: $euiSizeXS;
+  }
+}
diff --git a/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap b/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap
index cfe754d3ae0c..f929a6d1fbf9 100644
--- a/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap
+++ b/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/call_outs.test.tsx.snap
@@ -3,7 +3,7 @@
 exports[`CallOuts spec should display empty_index call out 1`] = `
 <div>
   <div
-    class="euiCallOut euiCallOut--warning euiCallOut--small"
+    class="euiCallOut euiCallOut--warning euiCallOut--small queryAssist queryAssist__callout"
     data-test-subj="query-assist-empty-index-callout"
   >
     <div
@@ -30,7 +30,7 @@ exports[`CallOuts spec should display empty_index call out 1`] = `
 exports[`CallOuts spec should display empty_query call out 1`] = `
 <div>
   <div
-    class="euiCallOut euiCallOut--warning euiCallOut--small"
+    class="euiCallOut euiCallOut--warning euiCallOut--small queryAssist queryAssist__callout"
     data-test-subj="query-assist-empty-query-callout"
   >
     <div
@@ -57,7 +57,7 @@ exports[`CallOuts spec should display empty_query call out 1`] = `
 exports[`CallOuts spec should display invalid_query call out 1`] = `
 <div>
   <div
-    class="euiCallOut euiCallOut--danger euiCallOut--small"
+    class="euiCallOut euiCallOut--danger euiCallOut--small queryAssist queryAssist__callout"
     data-test-subj="query-assist-guard-callout"
   >
     <div
@@ -84,7 +84,7 @@ exports[`CallOuts spec should display invalid_query call out 1`] = `
 exports[`CallOuts spec should display query_generated call out 1`] = `
 <div>
   <div
-    class="euiCallOut euiCallOut--success euiCallOut--small"
+    class="euiCallOut euiCallOut--success euiCallOut--small queryAssist queryAssist__callout"
     data-test-subj="query-assist-query-generated-callout"
   >
     <div
diff --git a/src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx b/src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx
index 16d218bed67e..57f70390ede1 100644
--- a/src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx
+++ b/src/plugins/query_enhancements/public/query_assist/components/call_outs.tsx
@@ -22,6 +22,7 @@ export type QueryAssistCallOutType =
 
 const EmptyIndexCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-empty-index-callout"
     title={
       <FormattedMessage
@@ -39,6 +40,7 @@ const EmptyIndexCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
 
 const ProhibitedQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-guard-callout"
     title={
       <FormattedMessage
@@ -56,6 +58,7 @@ const ProhibitedQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
 
 const EmptyQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-empty-query-callout"
     title={
       <FormattedMessage
@@ -73,6 +76,7 @@ const EmptyQueryCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
 
 const QueryGeneratedCallOut: React.FC<QueryAssistCallOutProps> = (props) => (
   <EuiCallOut
+    className="queryAssist queryAssist__callout"
     data-test-subj="query-assist-query-generated-callout"
     title={
       <FormattedMessage
diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx
index 53b20266191f..6a1240e93b00 100644
--- a/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx
+++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx
@@ -40,6 +40,7 @@ export const QueryAssistBanner: React.FC<QueryAssistBannerProps> = (props) => {
 
   return (
     <EuiCallOut
+      className="queryAssist"
       size="s"
       title={
         <EuiFlexGroup responsive={false} gutterSize="s" alignItems="center">
diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx
index 9433047833b8..4270177b7969 100644
--- a/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx
+++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx
@@ -35,7 +35,9 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
   const { generateQuery, loading } = useGenerateQuery();
   const [callOutType, setCallOutType] = useState<QueryAssistCallOutType>();
   const dismissCallout = () => setCallOutType(undefined);
-  const [selectedDataSet, setSelectedDataSet] = useState<SimpleDataSet>();
+  const [selectedDataSet, setSelectedDataSet] = useState<SimpleDataSet | undefined>(
+    services.data.query.dataSetManager.getDataSet()
+  );
   const selectedIndex = selectedDataSet?.title;
   const previousQuestionRef = useRef<string>();
 
@@ -85,7 +87,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
   if (props.dependencies.isCollapsed) return null;
 
   return (
-    <EuiForm component="form" onSubmit={onSubmit}>
+    <EuiForm component="form" onSubmit={onSubmit} className="queryAssist queryAssist__form">
       <EuiFormRow fullWidth>
         <EuiFlexGroup gutterSize="s" responsive={false} alignItems="center">
           <EuiFlexItem>
diff --git a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx
index c3fb47aef887..cf99c79798b0 100644
--- a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx
+++ b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx
@@ -15,24 +15,6 @@ import { DataSetContract } from '../../../../data/public/query';
 import { ConfigSchema } from '../../../common/config';
 import { createQueryAssistExtension } from './create_extension';
 
-const mockSimpleDataSet = {
-  id: 'mock-data-set-id',
-  title: 'mock-title',
-  dataSourceRef: {
-    id: 'mock-data-source-id',
-  },
-} as SimpleDataSet;
-
-const mockDataSetManager: jest.Mocked<DataSetContract> = {
-  getUpdates$: jest.fn().mockReturnValue(of(mockSimpleDataSet)),
-  getDataSet: jest.fn().mockReturnValue(mockSimpleDataSet),
-  setDataSet: jest.fn(),
-  getDefaultDataSet: jest.fn().mockReturnValue(mockSimpleDataSet),
-  fetchDefaultDataSet: jest.fn().mockResolvedValue(mockSimpleDataSet),
-  init: jest.fn(),
-  initWithIndexPattern: jest.fn(),
-};
-
 const coreSetupMock = coreMock.createSetup({
   pluginStartDeps: {
     data: {
@@ -41,13 +23,19 @@ const coreSetupMock = coreMock.createSetup({
   },
 });
 const httpMock = coreSetupMock.http;
-const dataMock = {
-  ...dataPluginMock.createSetupContract(),
-  query: {
-    ...dataPluginMock.createSetupContract().query,
-    dataSetManager: mockDataSetManager,
+const dataMock = dataPluginMock.createSetupContract();
+const dataSetMock = (dataMock.query.dataSetManager as unknown) as jest.Mocked<DataSetContract>;
+
+const mockSimpleDataSet = {
+  id: 'mock-data-set-id',
+  title: 'mock-title',
+  dataSourceRef: {
+    id: 'mock-data-source-id',
   },
-};
+} as SimpleDataSet;
+
+dataSetMock.getDataSet.mockReturnValue(mockSimpleDataSet);
+dataSetMock.getUpdates$.mockReturnValue(of(mockSimpleDataSet));
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
@@ -71,17 +59,7 @@ describe('CreateExtension', () => {
 
   it('should be enabled if at least one language is configured', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(
-      httpMock,
-      {
-        ...dataMock,
-        query: {
-          ...dataMock.query,
-          dataSetManager: mockDataSetManager,
-        },
-      },
-      config
-    );
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const isEnabled = await firstValueFrom(extension.isEnabled$(dependencies));
     expect(isEnabled).toBeTruthy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
diff --git a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx
index eaafaf88f7ce..b6a55fa3ebb6 100644
--- a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx
+++ b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx
@@ -5,7 +5,7 @@
 
 import { HttpSetup } from 'opensearch-dashboards/public';
 import React, { useEffect, useState } from 'react';
-import { distinctUntilChanged, map, switchMap } from 'rxjs/operators';
+import { distinctUntilChanged, map, startWith, switchMap } from 'rxjs/operators';
 import { SIMPLE_DATA_SOURCE_TYPES } from '../../../../data/common';
 import {
   DataPublicPluginSetup,
@@ -26,6 +26,7 @@ const getAvailableLanguages$ = (
   data: DataPublicPluginSetup
 ) =>
   data.query.dataSetManager.getUpdates$().pipe(
+    startWith(data.query.dataSetManager.getDataSet()),
     distinctUntilChanged(),
     switchMap(async (simpleDataSet) => {
       // currently query assist tool relies on opensearch API to get index