align language to uppercase

Signed-off-by: Joshua Li <[email protected]>
opensearch-project · Jun 4, 2024 · 0771a0d · 0771a0d
1 parent 1a9d1d9
commit 0771a0d
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 67 deletions.
diff --git a/plugins-extra/query_enhancements/public/query_assist/create_extension.tsx b/plugins-extra/query_enhancements/public/query_assist/create_extension.tsx
@@ -5,14 +5,14 @@ import { SearchBarExtensionConfig } from '../../../../src/plugins/data/public/ui
 
 export const createQueryAssistExtension = (http: HttpSetup): SearchBarExtensionConfig => {
   return {
-    id: 'query-assist',
+    id: 'query-assist-ppl',
     order: 1000,
     isEnabled: (() => {
       let agentConfigured: boolean;
       return async () => {
         if (agentConfigured === undefined) {
           agentConfigured = await http
-            .get<{ configured: boolean }>('/api/ql/query_assist/configured/ppl')
+            .get<{ configured: boolean }>('/api/ql/query_assist/configured/PPL')
             .then((response) => response.configured)
             .catch(() => false);
         }

diff --git a/plugins-extra/query_enhancements/public/query_assist/query_assist_input.tsx b/plugins-extra/query_enhancements/public/query_assist/query_assist_input.tsx
@@ -1,5 +1,5 @@
 import { EuiFieldText, EuiIcon, EuiOutsideClickDetector, EuiPortal } from '@elastic/eui';
-import React, { useState } from 'react';
+import React, { useMemo, useState } from 'react';
 import {
   PersistedLog,
   QuerySuggestionTypes,
@@ -17,25 +17,21 @@ interface QueryAssistInputProps {
 
 export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
   const [isSuggestionsVisible, setIsSuggestionsVisible] = useState(false);
-  const [index, setIndex] = useState<number | null>(null);
+  const [suggestionIndex, setSuggestionIndex] = useState<number | null>(null);
   const [value, setValue] = useState(props.initialValue ?? '');
 
-  const getRecentSearchSuggestions = (query: string) => {
-    if (!props.persistedLog) {
-      return [];
-    }
-    const recentSearches = props.persistedLog.get();
-    const matchingRecentSearches = recentSearches.filter((recentQuery) => {
-      const recentQueryString = typeof recentQuery === 'object' ? recentQuery : recentQuery;
-      return recentQueryString.includes(query);
-    });
-    return matchingRecentSearches.map((recentSearch) => {
-      const text = recentSearch;
-      const start = 0;
-      const end = query.length;
-      return { type: QuerySuggestionTypes.RecentSearch, text, start, end };
-    });
-  };
+  const recentSearchSuggestions = useMemo(() => {
+    if (!props.persistedLog) return [];
+    return props.persistedLog
+      .get()
+      .filter((recentSearch) => recentSearch.includes(value))
+      .map((recentSearch) => ({
+        type: QuerySuggestionTypes.RecentSearch,
+        text: recentSearch,
+        start: 0,
+        end: value.length,
+      }));
+  }, [props.persistedLog, value]);
 
   return (
     <EuiOutsideClickDetector onOutsideClick={() => setIsSuggestionsVisible(false)}>
@@ -58,16 +54,16 @@ export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
         <EuiPortal>
           <SuggestionsComponent
             show={isSuggestionsVisible}
-            suggestions={getRecentSearchSuggestions(value)}
-            index={index}
+            suggestions={recentSearchSuggestions}
+            index={suggestionIndex}
             onClick={(suggestion) => {
               if (!props.inputRef.current) return;
               setValue(suggestion.text);
               setIsSuggestionsVisible(false);
-              setIndex(null);
+              setSuggestionIndex(null);
               props.inputRef.current.focus();
             }}
-            onMouseEnter={(i) => setIndex(i)}
+            onMouseEnter={(i) => setSuggestionIndex(i)}
             loadMore={() => {}}
             queryBarRect={props.inputRef.current?.getBoundingClientRect()}
             size="s"

diff --git a/plugins-extra/query_enhancements/server/routes/index.ts b/plugins-extra/query_enhancements/server/routes/index.ts
@@ -10,7 +10,7 @@ import {
   IOpenSearchDashboardsSearchRequest,
 } from '../../../../src/plugins/data/common';
 import { ISearchStrategy } from '../../../../src/plugins/data/server';
-import { registerPplQueryAssistRoute } from './query_assist';
+import { registerQueryAssistRoutes } from './query_assist';
 
 export function defineRoutes(
   logger: Logger,
@@ -20,7 +20,7 @@ export function defineRoutes(
     ISearchStrategy<IOpenSearchDashboardsSearchRequest, IDataFrameResponse>
   >
 ) {
-  registerPplQueryAssistRoute(logger, router);
+  registerQueryAssistRoutes(logger, router);
   router.post(
     {
       path: `/api/pplql/search`,

diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts b/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts
@@ -15,7 +15,7 @@ const AGENT_REQUEST_OPTIONS = {
   maxRetries: 0,
 };
 
-type AgentResponse = ApiResponse<{
+export type AgentResponse = ApiResponse<{
   inference_results: Array<{
     output: Array<{ name: string; result?: string }>;
   }>;

diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/index.ts b/plugins-extra/query_enhancements/server/routes/query_assist/index.ts
@@ -1,5 +1,5 @@
-export { registerPplQueryAssistRoute } from './routes';
+export { registerQueryAssistRoutes } from './routes';
 
 export const AGENT_CONFIG_NAME_MAP = {
-  ppl: 'os_query_assist_ppl',
+  PPL: 'os_query_assist_ppl',
 } as const;
diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/ppl/create_response.ts b/plugins-extra/query_enhancements/server/routes/query_assist/ppl/create_response.ts
@@ -0,0 +1,17 @@
+import { QueryAssistResponse } from '../../../../common/query_assist';
+import { AgentResponse } from '../agents';
+
+export const createPPLResponseBody = (agentResponse: AgentResponse): QueryAssistResponse => {
+  if (!agentResponse.body.inference_results[0].output[0].result)
+    throw new Error('Generated query not found.');
+  const result = JSON.parse(agentResponse.body.inference_results[0].output[0].result!) as {
+    ppl: string;
+  };
+  const ppl = result.ppl
+    .replace(/[\r\n]/g, ' ')
+    .trim()
+    .replace(/ISNOTNULL/g, 'isnotnull') // https://github.com/opensearch-project/sql/issues/2431
+    .replace(/`/g, '') // https://github.com/opensearch-project/dashboards-observability/issues/509, https://github.com/opensearch-project/dashboards-observability/issues/557
+    .replace(/\bSPAN\(/g, 'span('); // https://github.com/opensearch-project/dashboards-observability/issues/759
+  return { query: ppl };
+};
diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts b/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts
@@ -1,19 +1,22 @@
 import { schema, Type } from '@osd/config-schema';
 import { IRouter, Logger } from 'opensearch-dashboards/server';
 import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors';
+import { ERROR_DETAILS } from '../../../common/query_assist';
 import { getAgentIdByConfig, requestAgentByConfig } from './agents';
 import { AGENT_CONFIG_NAME_MAP } from './index';
-import { ERROR_DETAILS } from '../../../common/query_assist';
+import { createPPLResponseBody } from './ppl/create_response';
+
+export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
+  const languageSchema = schema.oneOf(
+    Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'PPL'>]
+  );
 
-export function registerPplQueryAssistRoute(logger: Logger, router: IRouter) {
   router.get(
     {
       path: '/api/ql/query_assist/configured/{language}',
       validate: {
         params: schema.object({
-          language: schema.oneOf(
-            Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'ppl'>]
-          ),
+          language: languageSchema,
         }),
       },
     },
@@ -36,53 +39,36 @@ export function registerPplQueryAssistRoute(logger: Logger, router: IRouter) {
         body: schema.object({
           index: schema.string(),
           question: schema.string(),
-          language: schema.oneOf(
-            Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'ppl'>]
-          ),
+          language: languageSchema,
         }),
       },
     },
     async (context, request, response) => {
-      return response.ok({
-        body: {
-          query: 'source=' + request.body.index + ' | head ' + Math.ceil(Math.random() * 100),
-        },
-      });
       try {
-        if (request.body.language.toLowerCase() !== 'ppl') throw new Error('Unsupported language.');
-        const pplRequest = await requestAgentByConfig({
+        if (!(request.body.language in AGENT_CONFIG_NAME_MAP))
+          throw new Error('Unsupported language.');
+        const agentResponse = await requestAgentByConfig({
           context,
-          configName: AGENT_CONFIG_NAME_MAP.ppl,
+          configName: AGENT_CONFIG_NAME_MAP[request.body.language],
           body: {
             parameters: {
               index: request.body.index,
               question: request.body.question,
             },
           },
         });
-        if (!pplRequest.body.inference_results[0].output[0].result)
-          throw new Error('Generated PPL query not found.');
-        const result = JSON.parse(pplRequest.body.inference_results[0].output[0].result!) as {
-          ppl: string;
-        };
-        const ppl = result.ppl
-          .replace(/[\r\n]/g, ' ')
-          .trim()
-          .replace(/ISNOTNULL/g, 'isnotnull') // https://github.com/opensearch-project/sql/issues/2431
-          .replace(/`/g, '') // https://github.com/opensearch-project/dashboards-observability/issues/509, https://github.com/opensearch-project/dashboards-observability/issues/557
-          .replace(/\bSPAN\(/g, 'span('); // https://github.com/opensearch-project/dashboards-observability/issues/759
-        return response.ok({
-          body: {
-            query: ppl,
-          },
-        });
+        const responseBody = createPPLResponseBody(agentResponse);
+        return response.ok({ body: responseBody });
       } catch (error) {
-        if (
-          isResponseError(error) &&
-          error.statusCode === 400 &&
-          error.body.includes(ERROR_DETAILS.GUARDRAILS_TRIGGERED)
-        ) {
-          return response.badRequest({ body: ERROR_DETAILS.GUARDRAILS_TRIGGERED });
+        if (isResponseError(error)) {
+          if (error.statusCode === 400 && error.body.includes(ERROR_DETAILS.GUARDRAILS_TRIGGERED))
+            return response.badRequest({ body: ERROR_DETAILS.GUARDRAILS_TRIGGERED });
+          return response.badRequest({
+            body:
+              typeof error.meta.body === 'string'
+                ? error.meta.body
+                : JSON.stringify(error.meta.body),
+          });
         }
         return response.custom({ statusCode: error.statusCode || 500, body: error.message });
       }