Implement visualization summary in text2viz page

CHANGELOG.md

@@ -25,4 +25,5 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Add experimental feature to support text to visualization ([#218](https://github.com/opensearch-project/dashboards-assistant/pull/218))
 - Be compatible with ML configuration index mapping change ([#239](https://github.com/opensearch-project/dashboards-assistant/pull/239))
 - Support context aware alert analysis by reusing incontext insight component([#215](https://github.com/opensearch-project/dashboards-assistant/pull/215))
-Use smaller and compressed variants of buttons and form components ([#250](https://github.com/opensearch-project/dashboards-assistant/pull/250))
+Use smaller and compressed variants of buttons and form components ([#250](https://github.com/opensearch-project/dashboards-assistant/pull/250))
+- Support visualization summary in text to visualization ([#257](https://github.com/opensearch-project/dashboards-assistant/pull/257))

common/constants/llm.ts

@@ -24,6 +24,10 @@ export const TEXT2VIZ_API = {
   TEXT2VEGA: `${API_BASE}/text2vega`,
 };
 
+export const SUMMARY_ASSISTANT_API = {
+  SUMMARIZE_VIZ: `${API_BASE}/summarize_viz`,
+};
+
 export const NOTEBOOK_API = {
   CREATE_NOTEBOOK: `${NOTEBOOK_PREFIX}/note`,
   SET_PARAGRAPH: `${NOTEBOOK_PREFIX}/set_paragraphs/`,

public/components/visualization/text2vega.ts

@@ -11,7 +11,7 @@ import { DataPublicPluginStart } from '../../../../../src/plugins/data/public';
 
 const DATA_SOURCE_DELIMITER = '::';
 
-const topN = (ppl: string, n: number) => `${ppl} | head ${n}`;
+export const topN = (ppl: string, n: number) => `${ppl} | head ${n}`;
 
 const getDataSourceAndIndexFromLabel = (label: string) => {
   if (label.includes(DATA_SOURCE_DELIMITER)) {

public/components/visualization/text2viz.tsx

@@ -23,6 +23,7 @@ import { i18n } from '@osd/i18n';
 import { useCallback } from 'react';
 import { useObservable } from 'react-use';
 import { useMemo } from 'react';
+import { BehaviorSubject } from 'rxjs';
 import { SourceSelector } from './source_selector';
 import type { DataSourceOption } from '../../../../../src/plugins/data/public';
 import chatIcon from '../../assets/chat.svg';
@@ -41,11 +42,14 @@ import {
 import './text2viz.scss';
 import { Text2VizEmpty } from './text2viz_empty';
 import { Text2VizLoading } from './text2viz_loading';
-import { Text2Vega } from './text2vega';
+import { Text2Vega, topN } from './text2vega';
 import {
   OnSaveProps,
   SavedObjectSaveModalOrigin,
 } from '../../../../../src/plugins/saved_objects/public';
+import { VizSummary } from './viz_summary';
+import sparkleSvg from '../../assets/sparkle.svg';
+import shinySparkleSvg from '../../assets/shiny_sparkle.svg';
 
 export const Text2Viz = () => {
   const [selectedSource, setSelectedSource] = useState<DataSourceOption>();
@@ -67,6 +71,10 @@ export const Text2Viz = () => {
   const [vegaSpec, setVegaSpec] = useState<Record<string, any>>();
   const text2vegaRef = useRef(new Text2Vega(http, data.search));
   const status = useObservable(text2vegaRef.current.status$);
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const sampleData$ = useRef(new BehaviorSubject<any>(null));
+  const [dataSourceId, setDataSourceId] = useState<string | undefined>(undefined);
+  const [showVizSummary, setShowVizSummary] = useState<boolean>(true); // By default, we enable to show the visualization summary.
 
   useEffect(() => {
     const text2vega = text2vegaRef.current;
@@ -89,17 +97,45 @@ export const Text2Viz = () => {
     };
   }, [http, notifications]);
 
+  useEffect(() => {
+    const fetchData = async () => {
+      if (vegaSpec) {
+        // Today's ppl query could return at most 10000 rows, make a safe limit to avoid data explosion of llm input for now
+        const ppl = topN(vegaSpec.data.url.body.query, 50);
+
+        try {
+          const sampleData = await data.search
+            .search({ params: { body: { query: ppl } }, dataSourceId }, { strategy: 'pplraw' })
+            .toPromise();
+
+          sampleData$.current.next(sampleData.rawResponse);
+        } catch (error) {
+          notifications.toasts.addError(error, {
+            title: i18n.translate('dashboardAssistant.feature.vizSummary.fetchData.error', {
+              defaultMessage: 'Error while fetching data to summarize visualization',
+            }),
+          });
+        }
+      }
+    };
+
+    fetchData();
+  }, [data.search, vegaSpec]);
+
   const onInputChange = useCallback((e: React.ChangeEvent<HTMLInputElement>) => {
     setInput(e.target.value);
   }, []);
 
   const onSubmit = useCallback(async () => {
     setVegaSpec(undefined);
+    setDataSourceId(undefined);
+    sampleData$.current.next(undefined);
     const text2vega = text2vegaRef.current;
     if (selectedSource?.label) {
       const dataSource = (await selectedSource.ds.getDataSet()).dataSets.find(
         (ds) => ds.title === selectedSource.label
       );
+      setDataSourceId(dataSource?.dataSourceId);
       text2vega.invoke({
         index: selectedSource.label,
         prompt: input,
@@ -247,7 +283,25 @@ export const Text2Viz = () => {
                   iconType="returnKey"
                 />
               </EuiFlexItem>
+              <EuiFlexItem grow={false}>
+                <EuiButtonIcon
+                  aria-label="show visualization summary"
+                  iconSize={'l'}
+                  iconType={showVizSummary ? shinySparkleSvg : sparkleSvg}
+                  onClick={() => {
+                    setShowVizSummary((isOn) => !isOn);
+                  }}
+                />
+              </EuiFlexItem>
             </EuiFlexGroup>
+            {showVizSummary && (
+              <VizSummary
+                http={http}
+                sampleData$={sampleData$.current}
+                vizParams={vegaSpec}
+                dataSourceId={dataSourceId}
+              />
+            )}
             {status === 'STOPPED' && !vegaSpec && (
               <EuiFlexGroup>
                 <EuiFlexItem>

public/components/visualization/viz_summary.test.tsx

@@ -0,0 +1,108 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { render, screen, waitFor } from '@testing-library/react';
+import '@testing-library/jest-dom/extend-expect';
+import { BehaviorSubject } from 'rxjs';
+import { HttpSetup } from '../../../../../src/core/public';
+import { SUMMARY_ASSISTANT_API } from '../../../common/constants/llm';
+import { VizSummary } from './viz_summary';
+import { getAssistantRole } from '../../utils/constants';
+import { getNotifications } from '../../services';
+
+// Mock necessary functions and modules
+jest.mock('../../services', () => ({
+  getNotifications: jest.fn().mockReturnValue({
+    toasts: {
+      addDanger: jest.fn(),
+    },
+  }),
+}));
+
+const mockHttpPost = jest.fn();
+const mockEscape = jest.fn((text) => text);
+jest.mock('lodash', () => ({
+  escape: (text) => mockEscape(text),
+}));
+
+describe('VizSummary', () => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  let sampleData$: BehaviorSubject<any>;
+  let http: HttpSetup;
+
+  beforeEach(() => {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    sampleData$ = new BehaviorSubject<any>(null);
+    http = ({
+      post: mockHttpPost,
+    } as unknown) as HttpSetup;
+    mockEscape.mockClear();
+    mockHttpPost.mockClear();
+  });
+
+  it('renders correctly when there is no summary and not generating', () => {
+    render(<VizSummary http={http} sampleData$={sampleData$} />);
+    expect(screen.getByText('Ask a question to generate a summary.')).toBeInTheDocument();
+  });
+
+  it('renders correctly and firstly shows the "Generating response..." message and then show summary', async () => {
+    sampleData$.next({ size: 100 });
+    mockHttpPost.mockResolvedValue('Summary text');
+    render(<VizSummary http={http} sampleData$={sampleData$} vizParams={{}} />);
+
+    expect(screen.queryByText('Generating response...')).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.queryByText('Generating response...')).not.toBeInTheDocument();
+      expect(screen.getByText('Summary text')).toBeInTheDocument();
+    });
+  });
+
+  it('correctly sends data to the server', () => {
+    const testData = { some: 'data' };
+    const testParams = { test: 'params' };
+    const expectedPrompt = getAssistantRole('vizSummary');
+
+    render(
+      <VizSummary http={http} sampleData$={new BehaviorSubject(testData)} vizParams={testParams} />
+    );
+
+    expect(mockHttpPost).toHaveBeenCalledWith(SUMMARY_ASSISTANT_API.SUMMARIZE_VIZ, {
+      body: JSON.stringify({
+        vizData: JSON.stringify(testData),
+        vizParams: JSON.stringify(testParams),
+        prompt: expectedPrompt,
+      }),
+      query: { dataSourceId: undefined },
+    });
+  });
+
+  it('handles API errors gracefully', async () => {
+    sampleData$.next({ size: 100 });
+    mockHttpPost.mockRejectedValue(new Error('API Error'));
+
+    render(<VizSummary http={http} sampleData$={sampleData$} vizParams={{}} />);
+
+    // Wait for the error handling logic to be executed
+    await waitFor(() => {
+      expect(getNotifications().toasts.addDanger).toHaveBeenCalled();
+      expect(screen.getByText('Ask a question to generate a summary.')).toBeInTheDocument();
+    });
+  });
+
+  it('does not make API call if sampleData is undefined', () => {
+    sampleData$.next(undefined);
+    render(<VizSummary http={http} sampleData$={sampleData$} vizParams={{}} />);
+
+    expect(mockHttpPost).not.toHaveBeenCalled();
+  });
+
+  it('does not make API call if vizParams is undefined', () => {
+    sampleData$.next({ size: 100 });
+    render(<VizSummary http={http} sampleData$={sampleData$} />);
+
+    expect(mockHttpPost).not.toHaveBeenCalled();
+  });
+});