[ES|QL] Add context-aware suggestion ordering with categorization

src/platform/packages/shared/kbn-esql-ast/index.ts

@@ -74,3 +74,7 @@ export { TIME_SYSTEM_PARAMS } from './src/definitions/utils/literals';
 export { withAutoSuggest } from './src/definitions/utils/autocomplete/helpers';
 
 export { getNoValidCallSignatureError } from './src/definitions/utils/validation/utils';
+
+export { SuggestionOrderingEngine } from './src/sorting';
+export { SuggestionCategory } from './src/sorting';
+export type { SortingContext } from './src/sorting';

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts

@@ -45,6 +45,7 @@ import {
 import { ESQL_VARIABLES_PREFIX } from '../../constants';
 import { getExpressionType, isExpressionComplete } from '../../../definitions/utils/expressions';
 import { getFunctionDefinition } from '../../../definitions/utils/functions';
+import { SuggestionCategory } from '../../../sorting/types';
 
 export enum CompletionPosition {
   AFTER_COMPLETION = 'after_completion',
@@ -207,7 +208,14 @@ export async function autocomplete(
 
       if (!lastWord) {
         suggestions.push({
-          ...buildConstantsDefinitions([promptSnippetText], '', '1')[0],
+          ...buildConstantsDefinitions(
+            [promptSnippetText],
+            '',
+            '1',
+            undefined,
+            undefined,
+            SuggestionCategory.CONSTANT_VALUE
+          )[0],
           label: promptText,
           asSnippet: true,
         });

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/util.ts

@@ -11,6 +11,7 @@ import { withAutoSuggest } from '../../../definitions/utils/autocomplete/helpers
 import type { ESQLAstAllCommands } from '../../../types';
 import type { ESQLPolicy, ISuggestionItem } from '../../types';
 import { getSafeInsertText } from '../../../definitions/utils/autocomplete/helpers';
+import { SuggestionCategory } from '../../../sorting/types';
 
 export const ENRICH_MODES = [
   {
@@ -123,6 +124,7 @@ export const noPoliciesAvailableSuggestion: ISuggestionItem = {
     defaultMessage: 'Click to create',
   }),
   sortText: 'D',
+  category: SuggestionCategory.CUSTOM_ACTION,
   command: {
     id: 'esql.policies.create',
     title: i18n.translate('kbn-esql-ast.esql.autocomplete.createNewPolicy', {
@@ -187,5 +189,6 @@ export const buildMatchingFieldsDefinition = (
         },
       }),
       sortText: 'D',
+      category: SuggestionCategory.FIELD,
     })
   );

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/autocomplete.ts

@@ -9,6 +9,7 @@
 import { withAutoSuggest } from '../../../definitions/utils/autocomplete/helpers';
 import type { ESQLAstRerankCommand, ESQLSingleAstItem, ESQLAstAllCommands } from '../../../types';
 import type { ICommandCallbacks, ISuggestionItem, ICommandContext } from '../../types';
+import { SuggestionCategory } from '../../../sorting/types';
 import { Location } from '../../types';
 import { getPosition, CaretPosition } from './utils';
 import {
@@ -60,7 +61,14 @@ export async function autocomplete(
       return [
         targetField,
         {
-          ...buildConstantsDefinitions([QUERY_TEXT_SNIPPET], '', '1')[0],
+          ...buildConstantsDefinitions(
+            [QUERY_TEXT_SNIPPET],
+            '',
+            '1',
+            undefined,
+            undefined,
+            SuggestionCategory.CONSTANT_VALUE
+          )[0],
           label: QUERY_TEXT,
           asSnippet: true,
         },
@@ -74,7 +82,14 @@ export async function autocomplete(
     case CaretPosition.RERANK_AFTER_TARGET_ASSIGNMENT: {
       return [
         {
-          ...buildConstantsDefinitions([QUERY_TEXT_SNIPPET], '', '1')[0],
+          ...buildConstantsDefinitions(
+            [QUERY_TEXT_SNIPPET],
+            '',
+            '1',
+            undefined,
+            undefined,
+            SuggestionCategory.CONSTANT_VALUE
+          )[0],
           label: QUERY_TEXT,
           asSnippet: true,
         },

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/complete_items.ts

@@ -12,6 +12,7 @@ import { esqlCommandRegistry } from '.';
 import { buildDocumentation } from '../definitions/utils/documentation';
 import { TIME_SYSTEM_PARAMS } from '../definitions/utils/literals';
 import { withAutoSuggest } from '../definitions/utils/autocomplete/helpers';
+import { SuggestionCategory } from '../sorting/types';
 
 const techPreviewLabel = i18n.translate('kbn-esql-ast.esql.autocomplete.techPreviewLabel', {
   defaultMessage: `Technical Preview`,
@@ -24,7 +25,13 @@ function buildCharCompleteItem(
     sortText,
     quoted,
     advanceCursorAndOpenSuggestions,
-  }: { sortText?: string; quoted: boolean; advanceCursorAndOpenSuggestions?: boolean } = {
+    category,
+  }: {
+    sortText?: string;
+    quoted: boolean;
+    advanceCursorAndOpenSuggestions?: boolean;
+    category?: SuggestionCategory;
+  } = {
     quoted: false,
   }
 ): ISuggestionItem {
@@ -34,6 +41,7 @@ function buildCharCompleteItem(
     kind: 'Keyword',
     detail,
     sortText,
+    ...(category && { category }),
   };
   return advanceCursorAndOpenSuggestions ? withAutoSuggest(suggestion) : suggestion;
 }
@@ -46,6 +54,7 @@ export const pipeCompleteItem: ISuggestionItem = withAutoSuggest({
     defaultMessage: 'Pipe (|)',
   }),
   sortText: 'C',
+  category: SuggestionCategory.PIPE,
 });
 
 export const allStarConstant: ISuggestionItem = {
@@ -65,7 +74,7 @@ export const commaCompleteItem = buildCharCompleteItem(
   i18n.translate('kbn-esql-ast.esql.autocomplete.commaDoc', {
     defaultMessage: 'Comma (,)',
   }),
-  { sortText: 'B', quoted: false }
+  { sortText: 'B', quoted: false, category: SuggestionCategory.COMMA }
 );
 
 export const byCompleteItem: ISuggestionItem = withAutoSuggest({
@@ -74,6 +83,7 @@ export const byCompleteItem: ISuggestionItem = withAutoSuggest({
   kind: 'Reference',
   detail: 'By',
   sortText: '1',
+  category: SuggestionCategory.LANGUAGE_KEYWORD,
 });
 
 export const whereCompleteItem: ISuggestionItem = withAutoSuggest({
@@ -82,6 +92,7 @@ export const whereCompleteItem: ISuggestionItem = withAutoSuggest({
   kind: 'Reference',
   detail: 'Where',
   sortText: '1',
+  category: SuggestionCategory.LANGUAGE_KEYWORD,
 });
 
 export const onCompleteItem: ISuggestionItem = withAutoSuggest({
@@ -90,6 +101,7 @@ export const onCompleteItem: ISuggestionItem = withAutoSuggest({
   kind: 'Reference',
   detail: 'On',
   sortText: '1',
+  category: SuggestionCategory.LANGUAGE_KEYWORD,
 });
 
 export const withCompleteItem: ISuggestionItem = withAutoSuggest({
@@ -99,6 +111,7 @@ export const withCompleteItem: ISuggestionItem = withAutoSuggest({
   kind: 'Reference',
   detail: 'With',
   sortText: '1',
+  category: SuggestionCategory.LANGUAGE_KEYWORD,
 });
 
 export const withMapCompleteItem: ISuggestionItem = withAutoSuggest({
@@ -119,6 +132,7 @@ export const subqueryCompleteItem: ISuggestionItem = withAutoSuggest({
     defaultMessage: 'Adds a nested ES|QL query to your current query',
   }),
   sortText: '1',
+  category: SuggestionCategory.CUSTOM_ACTION,
 });
 
 export const minMaxValueCompleteItem: ISuggestionItem = {
@@ -127,6 +141,7 @@ export const minMaxValueCompleteItem: ISuggestionItem = {
   kind: 'Value',
   detail: 'minmax',
   sortText: '1',
+  category: SuggestionCategory.VALUE,
 };
 
 export const noneValueCompleteItem: ISuggestionItem = {
@@ -135,6 +150,7 @@ export const noneValueCompleteItem: ISuggestionItem = {
   kind: 'Value',
   detail: 'none',
   sortText: '1',
+  category: SuggestionCategory.VALUE,
 };
 
 export const getNewUserDefinedColumnSuggestion = (label: string): ISuggestionItem => {
@@ -146,6 +162,7 @@ export const getNewUserDefinedColumnSuggestion = (label: string): ISuggestionIte
       defaultMessage: 'Define a new column',
     }),
     sortText: '1',
+    category: SuggestionCategory.USER_DEFINED_COLUMN,
   });
 };
 
@@ -157,6 +174,7 @@ export const assignCompletionItem: ISuggestionItem = withAutoSuggest({
   kind: 'Variable',
   sortText: '1',
   text: '= ',
+  category: SuggestionCategory.USER_DEFINED_COLUMN,
 });
 
 export const asCompletionItem: ISuggestionItem = {
@@ -167,6 +185,7 @@ export const asCompletionItem: ISuggestionItem = {
   label: 'AS',
   sortText: '1',
   text: 'AS ',
+  category: SuggestionCategory.LANGUAGE_KEYWORD,
 };
 
 export const colonCompleteItem = buildCharCompleteItem(
@@ -254,4 +273,5 @@ export const getDateHistogramCompletionItem: (histogramBarTarget?: number) => IS
       defaultMessage: 'Add date histogram using bucket()',
     }),
     sortText: '1',
+    category: SuggestionCategory.CUSTOM_ACTION,
   });

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/options/metadata/index.ts

@@ -14,6 +14,7 @@ import { buildFieldsDefinitions } from '../../../definitions/utils/functions';
 import { handleFragment } from '../../../definitions/utils/autocomplete/helpers';
 import { commaCompleteItem, pipeCompleteItem } from '../../complete_items';
 import { isColumn, isOptionNode } from '../../../ast/is';
+import { SuggestionCategory } from '../../../sorting/types';
 
 export const METADATA_FIELDS = [
   '_version',
@@ -33,6 +34,7 @@ export const metadataSuggestion: ISuggestionItem = withAutoSuggest({
     defaultMessage: 'Metadata',
   }),
   sortText: 'C',
+  category: SuggestionCategory.VALUE,
 });
 
 export const getMetadataSuggestions = (command: ESQLAstAllCommands, queryText: string) => {

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/options/recommended_queries/index.ts

@@ -11,12 +11,21 @@ import type { RecommendedQuery, RecommendedField } from '@kbn/esql-types';
 import type { GetColumnsByTypeFn, ISuggestionItem } from '../../types';
 import { METADATA_FIELDS } from '../metadata';
 import { prettifyQueryTemplate, prettifyQuery } from './utils';
+import { SuggestionCategory } from '../../../sorting/types';
 
 export interface EditorExtensions {
   recommendedQueries: RecommendedQuery[];
   recommendedFields: RecommendedField[];
 }
 
+interface QueryTemplate {
+  label: string;
+  description: string;
+  queryString: string;
+  sortText?: string;
+  category?: SuggestionCategory;
+}
+
 // Order starts with the simple ones and goes to more complex ones
 
 export const getRecommendedQueriesTemplates = ({
@@ -27,8 +36,8 @@ export const getRecommendedQueriesTemplates = ({
   fromCommand: string;
   timeField?: string;
   categorizationField?: string;
-}) => {
-  const queries = [
+}): QueryTemplate[] => {
+  const queries: QueryTemplate[] = [
     {
       label: i18n.translate('kbn-esql-ast.recommendedQueries.searchExample.label', {
         defaultMessage: 'Search all fields',
@@ -38,6 +47,7 @@ export const getRecommendedQueriesTemplates = ({
       }),
       queryString: `${fromCommand}\n  | WHERE KQL("term") /* Search all fields using KQL – e.g. WHERE KQL("debug") */`,
       sortText: 'D',
+      category: SuggestionCategory.RECOMMENDED_QUERY_WITH_PRIORITY,
     },
     {
       label: i18n.translate('kbn-esql-ast.recommendedQueries.aggregateExample.label', {
@@ -201,6 +211,7 @@ export const getRecommendedQueriesSuggestionsFromStaticTemplates = async (
       kind: 'Issue',
       detail: query.description,
       sortText: query?.sortText ?? 'E',
+      category: query.category ?? SuggestionCategory.RECOMMENDED_QUERY,
       command: {
         id: 'esql.recommendedQuery.accept',
         title: 'Accept recommended query',
@@ -239,6 +250,7 @@ export const getRecommendedQueriesTemplatesFromExtensions = (
           : {}),
         kind: 'Issue',
         sortText: 'D',
+        category: SuggestionCategory.RECOMMENDED_QUERY_WITH_PRIORITY,
       };
     }
   );

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts

@@ -18,6 +18,7 @@ import type { PricingProduct } from '@kbn/core-pricing-common/src/types';
 import type { ESQLLocation } from '../types';
 import type { FieldType, SupportedDataType } from '../definitions/types';
 import type { EditorExtensions } from './options/recommended_queries';
+import type { SuggestionCategory } from '../sorting/types';
 
 // This is a subset of the Monaco's editor CompletitionItemKind type
 export type ItemKind =
@@ -65,6 +66,10 @@ export interface ISuggestionItem {
    * A string to use for sorting the suggestion within the suggestions list
    */
   sortText?: string;
+  /**
+   * The category of the suggestion, used for sorting and prioritization
+   */
+  category?: SuggestionCategory;
   /**
    * Suggestions can trigger a command by id. This is useful to trigger specific actions in some contexts
    */

src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts

@@ -28,6 +28,7 @@ import type { FunctionDefinition } from '../../types';
 import type { SupportedDataType } from '../../types';
 import { argMatchesParamType, getExpressionType, getParamAtPosition } from '../expressions';
 import { filterFunctionDefinitions, getAllFunctions, getFunctionSuggestion } from '../functions';
+import { SuggestionCategory } from '../../../sorting/types';
 import { buildConstantsDefinitions, getCompatibleLiterals, getDateLiterals } from '../literals';
 import { getColumnByName } from '../shared';
 
@@ -53,6 +54,7 @@ export const buildUserDefinedColumnsDefinitions = (
       defaultMessage: `Column specified by the user within the ES|QL query`,
     }),
     sortText: 'D',
+    category: SuggestionCategory.USER_DEFINED_COLUMN,
   }));
 
 export function pushItUpInTheList(suggestions: ISuggestionItem[], shouldPromote: boolean) {
@@ -327,6 +329,7 @@ export function getControlSuggestion(
         defaultMessage: 'Click to create',
       }),
       sortText: '1',
+      category: SuggestionCategory.CUSTOM_ACTION,
       command: {
         id: `esql.control.${type}.create`,
         title: i18n.translate('kbn-esql-ast.esql.autocomplete.createControlDetailLabel', {
@@ -341,7 +344,10 @@ export function getControlSuggestion(
           i18n.translate('kbn-esql-ast.esql.autocomplete.namedParamDefinition', {
             defaultMessage: 'Named parameter',
           }),
-          '1A'
+          '1A',
+          undefined,
+          undefined,
+          SuggestionCategory.USER_DEFINED_COLUMN
         )
       : []),
   ];
@@ -519,6 +525,7 @@ export function createInferenceEndpointToCompletionItem(
     label: inferenceEndpoint.inference_id,
     sortText: '1',
     text: inferenceEndpoint.inference_id,
+    category: SuggestionCategory.VALUE,
   };
 }