[ES|QL] Build function arguments suggestions from hints

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/types.ts

@@ -99,6 +99,13 @@ export const isParameterType = (str: string | undefined): str is FunctionParamet
 export const isReturnType = (str: string | FunctionParameterType): str is FunctionReturnType =>
   str !== 'unsupported' && (str === 'unknown' || str === 'any' || dataTypes.includes(str));
 
+export const parameterHintEntityTypes = ['inference_endpoint'] as const;
+export type ParameterHintEntityType = (typeof parameterHintEntityTypes)[number];
+export interface ParameterHint {
+  entityType: ParameterHintEntityType;
+  constraints?: Record<string, string>;
+}
+
 export interface FunctionParameter {
   name: string;
   type: FunctionParameterType;
@@ -129,8 +136,11 @@ export interface FunctionParameter {
    */
   supportsMultiValues?: boolean;
 
-  /** Additional hint information for the parameter */
-  hint?: unknown;
+  /**
+   * Provides information that is useful for getting parameter values from external sources.
+   * For example, an inference endpoint
+   */
+  hint?: ParameterHint;
 }
 
 export interface ElasticsearchCommandDefinition {

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/expressions/positions/empty_expression.ts

@@ -8,7 +8,7 @@
  */
 
 import { ControlTriggerSource, ESQLVariableType } from '@kbn/esql-types';
-import { uniq } from 'lodash';
+import { isEqual, uniq, uniqWith } from 'lodash';
 import { matchesSpecialFunction } from '../utils';
 import { shouldSuggestComma, type CommaContext } from '../comma_decision_engine';
 import type { ExpressionContext } from '../types';
@@ -21,6 +21,7 @@ import type {
   FunctionDefinition,
   FunctionParameter,
   FunctionParameterType,
+  ParameterHint,
 } from '../../../../types';
 import { type ISuggestionItem } from '../../../../../registry/types';
 import { FULL_TEXT_SEARCH_FUNCTIONS } from '../../../../constants';
@@ -29,6 +30,7 @@ import {
   valuePlaceholderConstant,
   defaultValuePlaceholderConstant,
 } from '../../../../../registry/complete_items';
+import { parametersFromHintsResolvers } from '../../parameters_from_hints';
 
 type FunctionParamContext = NonNullable<ExpressionContext['options']['functionParameterContext']>;
 
@@ -95,11 +97,16 @@ function tryExclusiveSuggestions(
     Boolean(functionParamContext.hasMoreMandatoryArgs),
     options.isCursorFollowedByComma ?? false
   );
-
   if (enumItems.length > 0) {
     return enumItems;
   }
 
+  // Some parameters suggests special values that are deduced from the hints object provided by ES.
+  const itemsFromHints = buildSuggestionsFromHints(paramDefinitions, ctx);
+  if (itemsFromHints.length > 0) {
+    return itemsFromHints;
+  }
+
   return [];
 }
 
@@ -374,6 +381,24 @@ function buildEnumValueSuggestions(
   });
 }
 
+function buildSuggestionsFromHints(
+  paramDefinitions: FunctionParameter[],
+  ctx: ExpressionContext
+): ISuggestionItem[] {
+  // Keep the hints that are unique by entityType + constraints
+  const hints: ParameterHint[] = uniqWith(
+    paramDefinitions.flatMap(({ hint }) => hint ?? []),
+    (a, b) => a.entityType === b.entityType && isEqual(a.constraints, b.constraints)
+  );
+
+  const results = hints.map(
+    (hint) =>
+      parametersFromHintsResolvers[hint.entityType]?.suggestionResolver?.(hint, ctx.context) ?? []
+  );
+
+  return results.flat();
+}
+
 /** Builds suggestions for constant-only literal parameters */
 function buildConstantOnlyLiteralSuggestions(
   paramDefinitions: FunctionParameter[],

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/parameters_from_hints.test.ts

@@ -0,0 +1,113 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { ICommandContext } from '../../../registry/types';
+import type { ParameterHint } from '../../types';
+import { parametersFromHintsResolvers } from './parameters_from_hints';
+import type { ESQLCallbacks, InferenceEndpointAutocompleteItem } from '@kbn/esql-types';
+
+describe('Parameters from hints handlers', () => {
+  describe('inference_endpoint hint', () => {
+    const inferenceEndpoints: InferenceEndpointAutocompleteItem[] = [
+      {
+        inference_id: 'text_embedding_endpoint',
+        task_type: 'text_embedding',
+      },
+    ];
+
+    const mockCallbacks: ESQLCallbacks = {
+      getInferenceEndpoints: jest.fn(async () => ({ inferenceEndpoints })),
+    };
+
+    const hint: ParameterHint = {
+      entityType: 'inference_endpoint' as const,
+      constraints: {
+        task_type: 'text_embedding',
+      },
+    };
+
+    beforeEach(() => {
+      jest.clearAllMocks();
+    });
+
+    it('should return inference endpoint suggestions filtered by task_type constraint', async () => {
+      const suggestions = await getSuggestionsForHint(hint, undefined, mockCallbacks);
+      expect(suggestions).toEqual(['text_embedding_endpoint']);
+    });
+
+    it('should not refetch inference endpoints if the context already has endpoints for the task type', async () => {
+      const suggestions = await getSuggestionsForHint(
+        hint,
+        {
+          columns: new Map(),
+          inferenceEndpoints,
+        },
+        mockCallbacks
+      );
+
+      expect(suggestions).toEqual(['text_embedding_endpoint']);
+      expect(mockCallbacks.getInferenceEndpoints).not.toHaveBeenCalled();
+    });
+
+    it('should fetch inference endpoints if the context already has endpoints, but not of the requested task type, also, it should preserve both', async () => {
+      const otherInferenceEndpoints: InferenceEndpointAutocompleteItem[] = [
+        {
+          inference_id: 'completion_endpoint',
+          task_type: 'completion',
+        },
+      ];
+
+      const suggestions = await getSuggestionsForHint(
+        hint,
+        {
+          columns: new Map(),
+          inferenceEndpoints: otherInferenceEndpoints,
+        },
+        mockCallbacks
+      );
+
+      expect(suggestions).toEqual(['text_embedding_endpoint']);
+      expect(mockCallbacks.getInferenceEndpoints).toHaveBeenCalledWith('text_embedding');
+    });
+  });
+});
+
+/**
+ * Calculates which would be the suggestions for a given parameter hint
+ * given certain callbacks and former context.
+ */
+export async function getSuggestionsForHint(
+  hint: ParameterHint,
+  formerContext?: ICommandContext,
+  callbacks: ESQLCallbacks = {}
+) {
+  const resolversEntry = parametersFromHintsResolvers[hint.entityType];
+
+  if (!resolversEntry) {
+    throw new Error(`No resolvers found for hint type: ${hint.entityType}`);
+  }
+
+  const { suggestionResolver, contextResolver } = resolversEntry;
+  if (!suggestionResolver) {
+    throw new Error(`No suggestionResolver found for hint type: ${hint.entityType}`);
+  }
+
+  // Build the context using the context resolver if available
+  let context: ICommandContext = formerContext ?? { columns: new Map() };
+  if (contextResolver) {
+    context = {
+      ...context,
+      ...((await contextResolver?.(hint, context, callbacks)) ?? {}),
+    };
+  }
+
+  const suggestions = suggestionResolver(hint, context).map((s) => s.label);
+
+  return suggestions;
+}

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/parameters_from_hints.ts

@@ -0,0 +1,104 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { InferenceTaskType } from '@elastic/elasticsearch/lib/api/types';
+import type { ESQLCallbacks } from '@kbn/esql-types';
+import { uniqBy } from 'lodash';
+import type { ParameterHint, ParameterHintEntityType } from '../../..';
+import type { ICommandContext, ISuggestionItem } from '../../../registry/types';
+import { createInferenceEndpointToCompletionItem } from './helpers';
+
+type SuggestionResolver = (hint: ParameterHint, ctx?: ICommandContext) => ISuggestionItem[];
+
+type ContextResolver = (
+  hint: ParameterHint,
+  ctx: Partial<ICommandContext>,
+  callbacks: ESQLCallbacks
+) => Promise<Record<string, unknown>>;
+
+/**
+ * For some parameters, ES gives us hints about the nature of it, that we use to provide
+ * custom autocompletion handlers.
+ *
+ * For each hint we need to provide:
+ * - a suggestionResolver to generate the autocompletion items for this param.
+ * - optionally, a contextResolver that populates the context with the data needed by the suggestionResolver.
+ *
+ * Important!
+ * Be mindful while implementing context resolvers, context is shared by the command and all functions used within it.
+ * If the data you need is already present, don't overwrite it, prefer merging it.
+ */
+export const parametersFromHintsResolvers: Partial<
+  Record<
+    ParameterHintEntityType,
+    {
+      suggestionResolver: SuggestionResolver;
+      contextResolver?: ContextResolver;
+    }
+  >
+> = {
+  ['inference_endpoint']: {
+    suggestionResolver: inferenceEndpointSuggestionResolver,
+    contextResolver: inferenceEndpointContextResolver,
+  },
+};
+
+// -------- INFERENCE ENDPOINT HINT -------- //
+function inferenceEndpointSuggestionResolver(
+  hint: ParameterHint,
+  ctx?: ICommandContext
+): ISuggestionItem[] {
+  if (hint.constraints?.task_type) {
+    const inferenceEnpoints =
+      ctx?.inferenceEndpoints?.filter((endpoint) => {
+        return endpoint.task_type === hint.constraints?.task_type;
+      }) ?? [];
+
+    return inferenceEnpoints.map((inferenceEndpoint) => {
+      const item = createInferenceEndpointToCompletionItem(inferenceEndpoint);
+      return {
+        ...item,
+        detail: '',
+        text: `"${item.text}"`,
+      };
+    });
+  }
+  return [];
+}
+
+async function inferenceEndpointContextResolver(
+  hint: ParameterHint,
+  ctx: Partial<ICommandContext>,
+  callbacks: ESQLCallbacks
+): Promise<Record<string, unknown>> {
+  if (hint.constraints?.task_type) {
+    const inferenceEndpointsFromContext = ctx.inferenceEndpoints ?? [];
+
+    // If the context already has an endpoint for the task type, we don't need to fetch them again
+    if (
+      inferenceEndpointsFromContext.find(
+        (endpoint) => endpoint.task_type === hint.constraints?.task_type
+      )
+    ) {
+      return {};
+    }
+
+    const inferenceEnpoints =
+      (await callbacks?.getInferenceEndpoints?.(hint.constraints?.task_type as InferenceTaskType))
+        ?.inferenceEndpoints || [];
+
+    return {
+      inferenceEndpoints: uniqBy(
+        [...inferenceEndpointsFromContext, ...inferenceEnpoints],
+        'inference_id'
+      ),
+    };
+  }
+  return {};
+}

src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/__tests__/autocomplete.parameters_with_hints.test.ts

@@ -0,0 +1,67 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { getSuggestionsForHint } from '../../../commands/definitions/utils/autocomplete/parameters_from_hints.test';
+import type { ESQLCallbacks } from '@kbn/esql-types';
+import { setup } from './helpers';
+import { getAllFunctions } from '../../../commands/definitions/utils/functions';
+import { uniqBy } from 'lodash';
+import { setTestFunctions } from '../../../commands/definitions/utils/test_functions';
+import { FunctionDefinitionTypes } from '../../../commands';
+import { Location } from '../../../commands/registry/types';
+
+const allUniqueParameterHints = uniqBy(
+  getAllFunctions()
+    .flatMap((fn) => fn.signatures)
+    .flatMap((signature) => signature.params)
+    .flatMap((param) => (param.hint ? [param.hint] : [])),
+  'entityType'
+);
+
+describe('function parameters autocomplete from hints', () => {
+  const callbacks: ESQLCallbacks = {
+    getInferenceEndpoints: async () => {
+      return {
+        inferenceEndpoints: [{ inference_id: 'inference_endpoint_1', task_type: 'text_embedding' }],
+      };
+    },
+  };
+
+  it.each(allUniqueParameterHints)('should resolve suggestions for $entityType', async (hint) => {
+    const functionName = `test_hint_${hint.entityType}`;
+
+    // Define a fake function to test the parameter hint
+    // (real functions can have the hinted param in different positions, making it difficult to generalize)
+    setTestFunctions([
+      {
+        type: FunctionDefinitionTypes.SCALAR,
+        name: functionName,
+        description: '',
+        signatures: [
+          {
+            params: [{ name: 'field', type: 'keyword', hint }],
+            returnType: 'double',
+          },
+        ],
+        locationsAvailable: [Location.EVAL],
+      },
+    ]);
+
+    const { suggest } = await setup();
+    const suggestions = (
+      await suggest(`FROM index | EVAL result = ${functionName}(/`, {
+        callbacks,
+      })
+    ).map((s) => s.label);
+
+    const suggestionsForHint = await getSuggestionsForHint(hint, undefined, callbacks);
+
+    expect(suggestions).toEqual(suggestionsForHint);
+  });
+});

src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/autocomplete.ts

@@ -245,11 +245,13 @@ async function getSuggestionsWithinCommandExpression(
     return findNewUserDefinedColumn(allUserDefinedColumns);
   };
 
+  // Get the context that might be needed by the command itself
   const additionalCommandContext = await getCommandContext(
-    astContext.command.name,
+    astContext.command,
     innerText,
     callbacks
   );
+
   const context = {
     ...references,
     ...additionalCommandContext,