From d38d3c3df6b036f05fa694ed808ef114dc615baa Mon Sep 17 00:00:00 2001
From: "github-actions[bot]" <github-actions[bot]@users.noreply.github.com>
Date: Thu, 1 Feb 2024 22:31:41 +0000
Subject: [PATCH] Hide query assist UI if PPL agent is not created (#1386)

Signed-off-by: Joshua Li <joshuali925@gmail.com>
(cherry picked from commit c1a89321e35bed51f0eb2c25ad9fb672f33d33ac)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 common/constants/query_assist.ts              |  1 +
 .../event_analytics/explorer/explorer.tsx     | 22 ++++++++++++++--
 server/routes/query_assist/routes.ts          | 26 ++++++++++++++++++-
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/common/constants/query_assist.ts b/common/constants/query_assist.ts
index 04b38c1299..e0b3bb29e6 100644
--- a/common/constants/query_assist.ts
+++ b/common/constants/query_assist.ts
@@ -5,6 +5,7 @@
 
 const QUERY_ASSIST_API_PREFIX = '/api/observability/query_assist';
 export const QUERY_ASSIST_API = {
+  CONFIGURED: `${QUERY_ASSIST_API_PREFIX}/configured`,
   GENERATE_PPL: `${QUERY_ASSIST_API_PREFIX}/generate_ppl`,
   SUMMARIZE: `${QUERY_ASSIST_API_PREFIX}/summarize`,
 };
diff --git a/public/components/event_analytics/explorer/explorer.tsx b/public/components/event_analytics/explorer/explorer.tsx
index 0d7b0efec3..279f234326 100644
--- a/public/components/event_analytics/explorer/explorer.tsx
+++ b/public/components/event_analytics/explorer/explorer.tsx
@@ -58,6 +58,7 @@ import {
   TAB_EVENT_TITLE,
   TIME_INTERVAL_OPTIONS,
 } from '../../../../common/constants/explorer';
+import { QUERY_ASSIST_API } from '../../../../common/constants/query_assist';
 import {
   LIVE_END_TIME,
   LIVE_OPTIONS,
@@ -101,8 +102,8 @@ import { selectSearchMetaData } from '../../event_analytics/redux/slices/search_
 import { getVizContainerProps } from '../../visualizations/charts/helpers';
 import { TabContext, useFetchEvents, useFetchPatterns, useFetchVisualizations } from '../hooks';
 import {
-  selectCountDistribution,
   render as updateCountDistribution,
+  selectCountDistribution,
 } from '../redux/slices/count_distribution_slice';
 import { selectFields, updateFields } from '../redux/slices/field_slice';
 import { selectQueryResult } from '../redux/slices/query_result_slice';
@@ -112,8 +113,8 @@ import { selectExplorerVisualization } from '../redux/slices/visualization_slice
 import {
   change as changeVisualizationConfig,
   change as changeVizConfig,
-  selectVisualizationConfig,
   change as updateVizConfig,
+  selectVisualizationConfig,
 } from '../redux/slices/viualization_config_slice';
 import { getDefaultVisConfig } from '../utils';
 import { formatError, getContentTabTitle, getDateRange } from '../utils/utils';
@@ -217,6 +218,7 @@ export const Explorer = ({
   const [subType, setSubType] = useState('visualization');
   const [_metricMeasure, setMetricMeasure] = useState('');
   const [_metricChecked, setMetricChecked] = useState(false);
+  const [_refresh, setRefresh] = useState({});
   const queryRef = useRef();
   const appBasedRef = useRef('');
   appBasedRef.current = appBaseQuery;
@@ -272,6 +274,22 @@ export const Explorer = ({
     };
   }, []);
 
+  useEffect(() => {
+    // query assist UI should only be enabled when the feature is enabled and configured.
+    if (coreRefs.queryAssistEnabled) {
+      http
+        .get<{ configured: boolean; error?: string }>(QUERY_ASSIST_API.CONFIGURED)
+        .catch(() => {
+          console.warn('Failed to check if query assist is configured');
+          return { configured: false };
+        })
+        .then((response) => {
+          coreRefs.queryAssistEnabled = response.configured;
+          setRefresh({});
+        });
+    }
+  }, []);
+
   const getErrorHandler = (title: string) => {
     return (error: any) => {
       if (coreRefs.summarizeEnabled) return;
diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts
index e99f5f9729..aef881f8c8 100644
--- a/server/routes/query_assist/routes.ts
+++ b/server/routes/query_assist/routes.ts
@@ -13,7 +13,7 @@ import {
 import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors';
 import { QUERY_ASSIST_API } from '../../../common/constants/query_assist';
 import { generateFieldContext } from '../../common/helpers/query_assist/generate_field_context';
-import { requestWithRetryAgentSearch } from './utils/agents';
+import { requestWithRetryAgentSearch, searchAgentIdByName } from './utils/agents';
 
 export function registerQueryAssistRoutes(router: IRouter, config: ObservabilityConfig) {
   const { ppl_agent_name: pplAgentName } = config.query_assist;
@@ -22,6 +22,30 @@ export function registerQueryAssistRoutes(router: IRouter, config: Observability
     error_summary_agent_name: errorSummaryAgentName,
   } = config.summarize;
 
+  /**
+   * Returns whether the PPL agent is configured.
+   */
+  router.get(
+    {
+      path: QUERY_ASSIST_API.CONFIGURED,
+      validate: {},
+    },
+    async (
+      context,
+      request,
+      response
+    ): Promise<IOpenSearchDashboardsResponse<any | ResponseError>> => {
+      const client = context.core.opensearch.client.asCurrentUser;
+      try {
+        // if the call does not throw any error, then the agent is properly configured
+        await searchAgentIdByName(client, pplAgentName!);
+        return response.ok({ body: { configured: true } });
+      } catch (error) {
+        return response.ok({ body: { configured: false, error: error.message } });
+      }
+    }
+  );
+
   router.post(
     {
       path: QUERY_ASSIST_API.GENERATE_PPL,