[Obs ai assistant] Hide the grid in case of errors (#184923)

## Summary

Small enhancements on the data grid support:

- Add the error message from ES in the errorMessages
- Hide the grid when there are error messages and display them to the
users

---------

Co-authored-by: Dario Gieselaar <dario.gieselaar@elastic.co>

x-pack/plugins/observability_solution/observability_ai_assistant_app/common/functions/visualize_esql.ts

@@ -7,6 +7,8 @@
 import type { FromSchema } from 'json-schema-to-ts';
 import { FunctionVisibility } from '@kbn/observability-ai-assistant-plugin/common';
 import { VISUALIZE_ESQL_USER_INTENTIONS } from '@kbn/observability-ai-assistant-plugin/common/functions/visualize_esql';
+import type { ESQLRow } from '@kbn/es-types';
+import type { DatatableColumn } from '@kbn/expressions-plugin/common';
 
 export const visualizeESQLFunction = {
   name: 'visualize_query',
@@ -29,4 +31,22 @@ export const visualizeESQLFunction = {
   contexts: ['core'],
 };
 
+export interface VisualizeQueryResponsev0 {
+  content: DatatableColumn[];
+}
+
+export interface VisualizeQueryResponsev1 {
+  data: {
+    columns: DatatableColumn[];
+    rows: ESQLRow[];
+    userOverrides?: unknown;
+  };
+  content: {
+    message: string;
+    errorMessages: string[];
+  };
+}
+
+export type VisualizeQueryResponse = VisualizeQueryResponsev0 | VisualizeQueryResponsev1;
+
 export type VisualizeESQLFunctionArguments = FromSchema<typeof visualizeESQLFunction['parameters']>;

x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx

@@ -42,7 +42,10 @@ import React, { useCallback, useContext, useEffect, useMemo, useState } from 're
 import ReactDOM from 'react-dom';
 import useAsync from 'react-use/lib/useAsync';
 import { v4 as uuidv4 } from 'uuid';
-import { VisualizeESQLFunctionArguments } from '../../common/functions/visualize_esql';
+import type {
+  VisualizeESQLFunctionArguments,
+  VisualizeQueryResponse,
+} from '../../common/functions/visualize_esql';
 import { ObservabilityAIAssistantAppPluginStartDependencies } from '../types';
 
 enum ChartType {
@@ -58,24 +61,6 @@ enum ChartType {
   Table = 'Table',
 }
 
-interface VisualizeQueryResponsev0 {
-  content: DatatableColumn[];
-}
-
-interface VisualizeQueryResponsev1 {
-  data: {
-    columns: DatatableColumn[];
-    rows: ESQLRow[];
-    userOverrides?: unknown;
-  };
-  content: {
-    message: string;
-    errorMessages: string[];
-  };
-}
-
-type VisualizeQueryResponse = VisualizeQueryResponsev0 | VisualizeQueryResponsev1;
-
 interface VisualizeESQLProps {
   /** Lens start contract, get the ES|QL charts suggestions api */
   lens: LensPublicStart;
@@ -105,6 +90,19 @@ interface VisualizeESQLProps {
 function generateId() {
   return uuidv4();
 }
+const saveVisualizationLabel = i18n.translate(
+  'xpack.observabilityAiAssistant.lensESQLFunction.save',
+  {
+    defaultMessage: 'Save visualization',
+  }
+);
+
+const editVisualizationLabel = i18n.translate(
+  'xpack.observabilityAiAssistant.lensESQLFunction.edit',
+  {
+    defaultMessage: 'Edit visualization',
+  }
+);
 
 export function VisualizeESQL({
   lens,
@@ -259,150 +257,135 @@ export function VisualizeESQL({
 
   return (
     <>
-      {!isLensInputTable && (
-        <EuiFlexGroup direction="column">
-          {Boolean(errorMessages?.length) && (
-            <>
-              <EuiText size="s">
-                {i18n.translate('xpack.observabilityAiAssistant.lensESQLFunction.errorMessage', {
-                  defaultMessage: 'There were some errors in the generated query',
-                })}
-              </EuiText>
-              <EuiDescriptionList data-test-subj="observabilityAiAssistantErrorsList">
-                {errorMessages?.map((error, index) => {
-                  return (
-                    <EuiDescriptionListDescription key={index}>
-                      <EuiFlexGroup gutterSize="s" alignItems="center">
-                        <EuiFlexItem grow={false}>
-                          <EuiIcon type="error" color="danger" size="s" />
-                        </EuiFlexItem>
-                        <EuiFlexItem grow={false}>{error}</EuiFlexItem>
-                      </EuiFlexGroup>
-                    </EuiDescriptionListDescription>
-                  );
-                })}
-              </EuiDescriptionList>
-            </>
-          )}
-
-          <EuiFlexItem grow={false}>
-            <EuiFlexGroup direction="row" gutterSize="s" justifyContent="flexEnd">
-              <EuiFlexItem grow={false}>
-                <EuiToolTip
-                  content={
-                    isTableVisible
-                      ? i18n.translate(
-                          'xpack.observabilityAiAssistant.lensESQLFunction.visualization',
-                          {
-                            defaultMessage: 'Visualization',
-                          }
-                        )
-                      : i18n.translate('xpack.observabilityAiAssistant.lensESQLFunction.table', {
-                          defaultMessage: 'Table of results',
-                        })
-                  }
-                >
-                  <EuiButtonIcon
-                    size="xs"
-                    iconType={isTableVisible ? 'visBarVerticalStacked' : 'tableDensityExpanded'}
-                    onClick={() => setIsTableVisible(!isTableVisible)}
-                    data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton"
-                    aria-label={
+      <EuiFlexGroup direction="column">
+        {!!errorMessages?.length && (
+          <>
+            <EuiText size="s">
+              {i18n.translate('xpack.observabilityAiAssistant.lensESQLFunction.errorMessage', {
+                defaultMessage: 'There were some errors in the generated query',
+              })}
+            </EuiText>
+            <EuiDescriptionList data-test-subj="observabilityAiAssistantErrorsList">
+              {errorMessages.map((error, index) => {
+                return (
+                  <EuiDescriptionListDescription key={index}>
+                    <EuiFlexGroup gutterSize="s" alignItems="center">
+                      <EuiFlexItem grow={false}>
+                        <EuiIcon type="error" color="danger" size="s" />
+                      </EuiFlexItem>
+                      <EuiFlexItem grow={false}>{error}</EuiFlexItem>
+                    </EuiFlexGroup>
+                  </EuiDescriptionListDescription>
+                );
+              })}
+            </EuiDescriptionList>
+          </>
+        )}
+        {!isLensInputTable && (
+          <>
+            <EuiFlexItem grow={false}>
+              <EuiFlexGroup direction="row" gutterSize="s" justifyContent="flexEnd">
+                <EuiFlexItem grow={false}>
+                  <EuiToolTip
+                    content={
                       isTableVisible
                         ? i18n.translate(
-                            'xpack.observabilityAiAssistant.lensESQLFunction.displayChart',
+                            'xpack.observabilityAiAssistant.lensESQLFunction.visualization',
                             {
-                              defaultMessage: 'Display chart',
-                            }
-                          )
-                        : i18n.translate(
-                            'xpack.observabilityAiAssistant.lensESQLFunction.displayTable',
-                            {
-                              defaultMessage: 'Display results',
+                              defaultMessage: 'Visualization',
                             }
                           )
+                        : i18n.translate('xpack.observabilityAiAssistant.lensESQLFunction.table', {
+                            defaultMessage: 'Table of results',
+                          })
                     }
-                  />
-                </EuiToolTip>
-              </EuiFlexItem>
-              <EuiToolTip
-                content={i18n.translate('xpack.observabilityAiAssistant.lensESQLFunction.edit', {
-                  defaultMessage: 'Edit visualization',
-                })}
-              >
-                <EuiButtonIcon
-                  size="xs"
-                  iconType="pencil"
-                  onClick={() => {
-                    chatFlyoutSecondSlotHandler?.setVisibility?.(true);
-                    if (triggerOptions) {
-                      uiActions.getTrigger('IN_APP_EMBEDDABLE_EDIT_TRIGGER').exec(triggerOptions);
-                    }
-                  }}
-                  data-test-subj="observabilityAiAssistantLensESQLEditButton"
-                  aria-label={i18n.translate(
-                    'xpack.observabilityAiAssistant.lensESQLFunction.edit',
-                    {
-                      defaultMessage: 'Edit visualization',
-                    }
-                  )}
-                />
-              </EuiToolTip>
-              <EuiFlexItem grow={false}>
-                <EuiToolTip
-                  content={i18n.translate('xpack.observabilityAiAssistant.lensESQLFunction.save', {
-                    defaultMessage: 'Save visualization',
-                  })}
-                >
+                  >
+                    <EuiButtonIcon
+                      size="xs"
+                      iconType={isTableVisible ? 'visBarVerticalStacked' : 'tableDensityExpanded'}
+                      onClick={() => setIsTableVisible(!isTableVisible)}
+                      data-test-subj="observabilityAiAssistantLensESQLDisplayTableButton"
+                      aria-label={
+                        isTableVisible
+                          ? i18n.translate(
+                              'xpack.observabilityAiAssistant.lensESQLFunction.displayChart',
+                              {
+                                defaultMessage: 'Display chart',
+                              }
+                            )
+                          : i18n.translate(
+                              'xpack.observabilityAiAssistant.lensESQLFunction.displayTable',
+                              {
+                                defaultMessage: 'Display table',
+                              }
+                            )
+                      }
+                    />
+                  </EuiToolTip>
+                </EuiFlexItem>
+                <EuiToolTip content={editVisualizationLabel}>
                   <EuiButtonIcon
                     size="xs"
-                    iconType="save"
-                    onClick={() => setIsSaveModalOpen(true)}
-                    data-test-subj="observabilityAiAssistantLensESQLSaveButton"
-                    aria-label={i18n.translate(
-                      'xpack.observabilityAiAssistant.lensESQLFunction.save',
-                      {
-                        defaultMessage: 'Save visualization',
+                    iconType="pencil"
+                    onClick={() => {
+                      chatFlyoutSecondSlotHandler?.setVisibility?.(true);
+                      if (triggerOptions) {
+                        uiActions.getTrigger('IN_APP_EMBEDDABLE_EDIT_TRIGGER').exec(triggerOptions);
                       }
-                    )}
+                    }}
+                    data-test-subj="observabilityAiAssistantLensESQLEditButton"
+                    aria-label={editVisualizationLabel}
                   />
                 </EuiToolTip>
-              </EuiFlexItem>
-            </EuiFlexGroup>
+                <EuiFlexItem grow={false}>
+                  <EuiToolTip content={saveVisualizationLabel}>
+                    <EuiButtonIcon
+                      size="xs"
+                      iconType="save"
+                      onClick={() => setIsSaveModalOpen(true)}
+                      data-test-subj="observabilityAiAssistantLensESQLSaveButton"
+                      aria-label={saveVisualizationLabel}
+                    />
+                  </EuiToolTip>
+                </EuiFlexItem>
+              </EuiFlexGroup>
+            </EuiFlexItem>
+            <EuiFlexItem data-test-subj={visualizationComponentDataTestSubj}>
+              {isTableVisible ? (
+                <ESQLDataGrid
+                  rows={rows}
+                  columns={columns}
+                  dataView={dataViewAsync.value}
+                  query={{ esql: query }}
+                  flyoutType="overlay"
+                  isTableView
+                />
+              ) : (
+                <lens.EmbeddableComponent
+                  {...lensInput}
+                  style={{
+                    height: 240,
+                  }}
+                  onLoad={onLoad}
+                />
+              )}
+            </EuiFlexItem>
+          </>
+        )}
+        {/* hide the grid in case of errors (as the user can't fix them) */}
+        {isLensInputTable && !errorMessages?.length && (
+          <EuiFlexItem data-test-subj="observabilityAiAssistantESQLDataGrid">
+            <ESQLDataGrid
+              rows={rows}
+              columns={columns}
+              dataView={dataViewAsync.value}
+              query={{ esql: query }}
+              flyoutType="overlay"
+            />
           </EuiFlexItem>
-          <EuiFlexItem data-test-subj={visualizationComponentDataTestSubj}>
-            {isTableVisible ? (
-              <ESQLDataGrid
-                rows={rows}
-                columns={columns}
-                dataView={dataViewAsync.value}
-                query={{ esql: query }}
-                flyoutType="overlay"
-                isTableView
-              />
-            ) : (
-              <lens.EmbeddableComponent
-                {...lensInput}
-                style={{
-                  height: 240,
-                }}
-                onLoad={onLoad}
-              />
-            )}
-          </EuiFlexItem>
-        </EuiFlexGroup>
-      )}
-      {isLensInputTable && (
-        <div data-test-subj="observabilityAiAssistantESQLDataGrid">
-          <ESQLDataGrid
-            rows={rows}
-            columns={columns}
-            dataView={dataViewAsync.value}
-            query={{ esql: query }}
-            flyoutType="overlay"
-          />
-        </div>
-      )}
+        )}
+      </EuiFlexGroup>
+
       {isSaveModalOpen ? (
         <lens.SaveModalComponent
           initialInput={lensInput}

x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts

@@ -9,7 +9,8 @@ import { validateQuery } from '@kbn/esql-validation-autocomplete';
 import { getAstAndSyntaxErrors } from '@kbn/esql-ast';
 import type { ElasticsearchClient } from '@kbn/core/server';
 import { ESQLSearchResponse, ESQLRow } from '@kbn/es-types';
-import { esFieldTypeToKibanaFieldType, type KBN_FIELD_TYPES } from '@kbn/field-types';
+import { esFieldTypeToKibanaFieldType } from '@kbn/field-types';
+import { DatatableColumn, DatatableColumnType } from '@kbn/expressions-plugin/common';
 import { splitIntoCommands } from './correct_common_esql_mistakes';
 
 export async function runAndValidateEsqlQuery({
@@ -19,13 +20,7 @@ export async function runAndValidateEsqlQuery({
   query: string;
   client: ElasticsearchClient;
 }): Promise<{
-  columns?: Array<{
-    id: string;
-    name: string;
-    meta: {
-      type: KBN_FIELD_TYPES;
-    };
-  }>;
+  columns?: DatatableColumn[];
   rows?: ESQLRow[];
   error?: Error;
   errorMessages?: string[];
@@ -63,7 +58,7 @@ export async function runAndValidateEsqlQuery({
         esqlResponse.columns?.map(({ name, type }) => ({
           id: name,
           name,
-          meta: { type: esFieldTypeToKibanaFieldType(type) },
+          meta: { type: esFieldTypeToKibanaFieldType(type) as DatatableColumnType },
         })) ?? [];
 
       return { columns, rows: esqlResponse.values };

x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/visualize_esql.ts

@@ -5,8 +5,11 @@
  * 2.0.
  */
 import { VisualizeESQLUserIntention } from '@kbn/observability-ai-assistant-plugin/common/functions/visualize_esql';
-import { visualizeESQLFunction } from '../../common/functions/visualize_esql';
-import { FunctionRegistrationParameters } from '.';
+import {
+  visualizeESQLFunction,
+  type VisualizeQueryResponsev1,
+} from '../../common/functions/visualize_esql';
+import type { FunctionRegistrationParameters } from '.';
 import { runAndValidateEsqlQuery } from './query/validate_esql_query';
 
 const getMessageForLLM = (
@@ -27,23 +30,32 @@ export function registerVisualizeESQLFunction({
   functions,
   resources,
 }: FunctionRegistrationParameters) {
-  functions.registerFunction(visualizeESQLFunction, async ({ arguments: { query, intention } }) => {
-    const { columns, errorMessages, rows } = await runAndValidateEsqlQuery({
-      query,
-      client: (await resources.context.core).elasticsearch.client.asCurrentUser,
-    });
+  functions.registerFunction(
+    visualizeESQLFunction,
+    async ({ arguments: { query, intention } }): Promise<VisualizeQueryResponsev1> => {
+      // errorMessages contains the syntax errors from the client side valdation
+      // error contains the error from the server side validation, it is always one error
+      // and help us identify errors like index not found, field not found etc.
+      const { columns, errorMessages, rows, error } = await runAndValidateEsqlQuery({
+        query,
+        client: (await resources.context.core).elasticsearch.client.asCurrentUser,
+      });
 
-    const message = getMessageForLLM(intention, query, Boolean(errorMessages?.length));
+      const message = getMessageForLLM(intention, query, Boolean(errorMessages?.length));
 
-    return {
-      data: {
-        columns,
-        rows,
-      },
-      content: {
-        message,
-        errorMessages,
-      },
-    };
-  });
+      return {
+        data: {
+          columns: columns ?? [],
+          rows: rows ?? [],
+        },
+        content: {
+          message,
+          errorMessages: [
+            ...(errorMessages ? errorMessages : []),
+            ...(error ? [error.message] : []),
+          ],
+        },
+      };
+    }
+  );
 }