support disabling summarization

Signed-off-by: Joshua Li <[email protected]>
joshuali925 · Jan 20, 2024 · 4cfd171 · 4cfd171
1 parent fa8e856
commit 4cfd171
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 24 deletions.
diff --git a/public/components/event_analytics/explorer/query_assist/__tests__/input.test.tsx b/public/components/event_analytics/explorer/query_assist/__tests__/input.test.tsx
@@ -85,7 +85,23 @@ describe('<QueryAssistInput /> spec', () => {
     );
   });
 
+  it('should not call summarize if disabled', async () => {
+    coreRefs.summarizeEnabled = false;
+    httpMock.post.mockRejectedValueOnce({ body: { statusCode: 429 } });
+
+    const { component } = renderQueryAssistInput();
+    await waitFor(() => {
+      fireEvent.click(component.getByTestId('query-assist-generate-and-run-button'));
+    });
+
+    expect(httpMock.post).toBeCalledWith(QUERY_ASSIST_API.GENERATE_PPL, {
+      body: '{"question":"test-input","index":"selected-test-index"}',
+    });
+    expect(httpMock.post).not.toBeCalledWith(QUERY_ASSIST_API.SUMMARIZE, expect.anything());
+  });
+
   it('should call summarize for generate and run errors', async () => {
+    coreRefs.summarizeEnabled = true;
     httpMock.post.mockRejectedValueOnce({ body: { statusCode: 429 } }).mockResolvedValueOnce({
       summary: 'too many requests',
       suggestedQuestions: ['1', '2'],

diff --git a/public/components/event_analytics/explorer/query_assist/input.tsx b/public/components/event_analytics/explorer/query_assist/input.tsx
@@ -23,6 +23,7 @@ import { ResponseError } from '@opensearch-project/opensearch/lib/errors';
 import React, { useEffect, useState } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 import { RAW_QUERY } from '../../../../../common/constants/explorer';
+import { QUERY_ASSIST_API } from '../../../../../common/constants/query_assist';
 import { getOSDHttp } from '../../../../../common/utils';
 import { coreRefs } from '../../../../framework/core_refs';
 import chatLogo from '../../../datasources/icons/query-assistant-logo.svg';
@@ -34,7 +35,6 @@ import {
 } from '../../redux/slices/query_assistant_summarization_slice';
 import { reset, selectQueryResult } from '../../redux/slices/query_result_slice';
 import { changeQuery, selectQueries } from '../../redux/slices/query_slice';
-import { QUERY_ASSIST_API } from '../../../../../common/constants/query_assist';
 
 interface SummarizationContext {
   question: string;
@@ -87,22 +87,23 @@ export const QueryAssistInput: React.FC<Props> = (props) => {
 
   useEffect(() => {
     if (
-      summaryData.responseForSummaryStatus === 'success' ||
-      summaryData.responseForSummaryStatus === 'failure'
-    ) {
-      void (async () => {
-        await dispatch(
-          changeSummary({
-            tabId: props.tabId,
-            data: {
-              summaryLoading: false,
-            },
-          })
-        );
-        if (explorerData.total > 0 || summaryData.responseForSummaryStatus === 'failure')
-          generateSummary();
-      })();
-    }
+      props.nlqInput.trim().length === 0 ||
+      (summaryData.responseForSummaryStatus !== 'success' &&
+        summaryData.responseForSummaryStatus !== 'failure')
+    )
+      return;
+    void (async () => {
+      await dispatch(
+        changeSummary({
+          tabId: props.tabId,
+          data: {
+            summaryLoading: false,
+          },
+        })
+      );
+      if (explorerData.total > 0 || summaryData.responseForSummaryStatus === 'failure')
+        generateSummary();
+    })();
   }, [summaryData.responseForSummaryStatus]);
 
   const [barSelected, setBarSelected] = useState(false);
@@ -174,6 +175,11 @@ export const QueryAssistInput: React.FC<Props> = (props) => {
     }
   };
   const generateSummary = async (context?: Partial<SummarizationContext>) => {
+    if (!coreRefs.summarizeEnabled) {
+      if (summaryData.responseForSummaryStatus === 'failure')
+        coreRefs.toasts?.addError(explorerData.error.body, { title: 'Invalid query' });
+      return;
+    }
     try {
       const isError = summaryData.responseForSummaryStatus === 'failure';
       const summarizationContext: SummarizationContext = {
@@ -334,7 +340,9 @@ export const QueryAssistInput: React.FC<Props> = (props) => {
                   <EuiButton
                     isLoading={generating}
                     onClick={generatePPL}
-                    isDisabled={generating || generatingOrRunning}
+                    isDisabled={
+                      generating || generatingOrRunning || props.nlqInput.trim().length === 0
+                    }
                     iconSide="right"
                     fill={false}
                     data-test-subj="query-assist-generate-button"
@@ -347,7 +355,9 @@ export const QueryAssistInput: React.FC<Props> = (props) => {
                   <EuiButton
                     isLoading={generatingOrRunning}
                     onClick={runAndSummarize}
-                    isDisabled={generating || generatingOrRunning}
+                    isDisabled={
+                      generating || generatingOrRunning || props.nlqInput.trim().length === 0
+                    }
                     iconType="returnKey"
                     iconSide="right"
                     type="submit"

diff --git a/public/framework/core_refs.ts b/public/framework/core_refs.ts
@@ -25,6 +25,7 @@ class CoreRefs {
   public chrome?: ChromeStart;
   public application?: ApplicationStart;
   public queryAssistEnabled?: boolean;
+  public summarizeEnabled?: boolean;
   public dashboard?: DashboardStart;
   public dashboardProviders?: unknown;
   private constructor() {

diff --git a/public/plugin.ts b/public/plugin.ts
@@ -84,6 +84,9 @@ interface PublicConfig {
   query_assist: {
     enabled: boolean;
   };
+  summarize: {
+    enabled: boolean;
+  };
 }
 
 export class ObservabilityPlugin
@@ -336,6 +339,7 @@ export class ObservabilityPlugin
     coreRefs.application = core.application;
     coreRefs.dashboard = startDeps.dashboard;
     coreRefs.queryAssistEnabled = this.config.query_assist.enabled;
+    coreRefs.summarizeEnabled = this.config.summarize.enabled;
 
     const { dataSourceService, dataSourceFactory } = startDeps.data.dataSources;
 

diff --git a/server/index.ts b/server/index.ts
@@ -17,9 +17,27 @@ const observabilityConfig = {
   schema: schema.object({
     query_assist: schema.object({
       enabled: schema.boolean({ defaultValue: false }),
-      ppl_agent_name: schema.maybe(schema.string()),
-      response_summary_agent_name: schema.maybe(schema.string()),
-      error_summary_agent_name: schema.maybe(schema.string()),
+      ppl_agent_name: schema.conditional(
+        schema.siblingRef('enabled'),
+        true,
+        schema.string(),
+        schema.maybe(schema.string())
+      ),
+    }),
+    summarize: schema.object({
+      enabled: schema.boolean({ defaultValue: false }),
+      response_summary_agent_name: schema.conditional(
+        schema.siblingRef('enabled'),
+        true,
+        schema.string(),
+        schema.maybe(schema.string())
+      ),
+      error_summary_agent_name: schema.conditional(
+        schema.siblingRef('enabled'),
+        true,
+        schema.string(),
+        schema.maybe(schema.string())
+      ),
     }),
   }),
 };
@@ -30,5 +48,6 @@ export const config: PluginConfigDescriptor<ObservabilityConfig> = {
   schema: observabilityConfig.schema,
   exposeToBrowser: {
     query_assist: true,
+    summarize: true,
   },
 };
diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts
@@ -16,11 +16,11 @@ import { generateFieldContext } from '../../common/helpers/query_assist/generate
 import { requestWithRetryAgentSearch } from './utils/agents';
 
 export function registerQueryAssistRoutes(router: IRouter, config: ObservabilityConfig) {
+  const { ppl_agent_name: pplAgentName } = config.query_assist;
   const {
-    ppl_agent_name: pplAgentName,
     response_summary_agent_name: responseSummaryAgentName,
     error_summary_agent_name: errorSummaryAgentName,
-  } = config.query_assist;
+  } = config.summarize;
 
   router.post(
     {