diff --git a/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/types.ts b/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/types.ts index 757189a21bf0fd..fd66f1579b4845 100644 --- a/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/types.ts +++ b/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; +} + 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 { diff --git a/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/expressions/positions/empty_expression.ts b/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/expressions/positions/empty_expression.ts index 59d923abd79dc1..fafcc9733a3b08 100644 --- a/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/expressions/positions/empty_expression.ts +++ b/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; @@ -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[], diff --git a/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/parameters_from_hints.test.ts b/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/parameters_from_hints.test.ts new file mode 100644 index 00000000000000..55416680d9cceb --- /dev/null +++ b/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; +} diff --git a/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/parameters_from_hints.ts b/src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/parameters_from_hints.ts new file mode 100644 index 00000000000000..b5169ed8fcac35 --- /dev/null +++ b/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, + callbacks: ESQLCallbacks +) => Promise>; + +/** + * 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, + callbacks: ESQLCallbacks +): Promise> { + 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 {}; +} diff --git a/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/__tests__/autocomplete.parameters_with_hints.test.ts b/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/__tests__/autocomplete.parameters_with_hints.test.ts new file mode 100644 index 00000000000000..a03bbc21c17225 --- /dev/null +++ b/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); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/autocomplete.ts b/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/autocomplete.ts index a3b97a10136db4..e23e7bcff8f7d9 100644 --- a/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/autocomplete.ts +++ b/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, diff --git a/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/get_command_context.ts b/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/get_command_context.ts index 289f4300b13990..9ba2d9abb94e26 100644 --- a/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/get_command_context.ts +++ b/src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/get_command_context.ts @@ -8,62 +8,80 @@ */ import type { ESQLCallbacks } from '@kbn/esql-types'; +import { isEqual, uniqWith } from 'lodash'; +import type { ParameterHint } from '../../..'; +import { walk } from '../../..'; +import type { ESQLAstAllCommands } from '../../types'; +import { getFunctionDefinition } from '../../commands/definitions/utils'; +import { parametersFromHintsResolvers } from '../../commands/definitions/utils/autocomplete/parameters_from_hints'; +import type { ICommandContext } from '../../commands/registry/types'; import { getPolicyHelper, getSourcesHelper } from '../shared/resources_helpers'; export const getCommandContext = async ( - commandName: string, + command: ESQLAstAllCommands, queryString: string, callbacks?: ESQLCallbacks -) => { +): Promise> => { const getSources = getSourcesHelper(callbacks); const helpers = getPolicyHelper(callbacks); - switch (commandName) { + + let context: Partial = {}; + + switch (command.name) { case 'completion': - return { + context = { inferenceEndpoints: (await callbacks?.getInferenceEndpoints?.('completion'))?.inferenceEndpoints || [], }; + break; case 'rerank': - return { + context = { inferenceEndpoints: (await callbacks?.getInferenceEndpoints?.('rerank'))?.inferenceEndpoints || [], }; + break; case 'enrich': const policies = await helpers.getPolicies(); const policiesMap = new Map(policies.map((policy) => [policy.name, policy])); - return { + context = { policies: policiesMap, }; + break; case 'from': const editorExtensions = (await callbacks?.getEditorExtensions?.(queryString)) ?? { recommendedQueries: [], + recommendedFields: [], }; - return { + context = { sources: await getSources(), editorExtensions, }; + break; case 'join': const joinSources = await callbacks?.getJoinIndices?.(); - return { + context = { joinSources: joinSources?.indices || [], supportsControls: callbacks?.canSuggestVariables?.() ?? false, }; + break; case 'stats': const histogramBarTarget = (await callbacks?.getPreferences?.())?.histogramBarTarget || 50; - return { + context = { histogramBarTarget, supportsControls: callbacks?.canSuggestVariables?.() ?? false, variables: callbacks?.getVariables?.(), }; + break; case 'inline stats': - return { + context = { histogramBarTarget: (await callbacks?.getPreferences?.())?.histogramBarTarget || 50, supportsControls: callbacks?.canSuggestVariables?.() ?? false, variables: callbacks?.getVariables?.(), }; + break; case 'fork': const enrichPolicies = await helpers.getPolicies(); - return { + context = { histogramBarTarget: (await callbacks?.getPreferences?.())?.histogramBarTarget || 50, joinSources: (await callbacks?.getJoinIndices?.())?.indices || [], supportsControls: callbacks?.canSuggestVariables?.() ?? false, @@ -71,16 +89,74 @@ export const getCommandContext = async ( inferenceEndpoints: (await callbacks?.getInferenceEndpoints?.('completion'))?.inferenceEndpoints || [], }; + break; case 'ts': const timeseriesSources = await callbacks?.getTimeseriesIndices?.(); - return { + context = { timeSeriesSources: timeseriesSources?.indices || [], sources: await getSources(), editorExtensions: (await callbacks?.getEditorExtensions?.(queryString)) ?? { recommendedQueries: [], + recommendedFields: [], }, }; + break; default: - return {}; + break; + } + + // Check if the functions used within the command needs additional context + context = await enhanceWithFunctionsContext(command, context, callbacks); + + return context; +}; + +/** + * Returns the context needed by the functions used within a command. + */ +export const enhanceWithFunctionsContext = async ( + command: ESQLAstAllCommands, + context: Partial, + callbacks?: ESQLCallbacks +): Promise> => { + const hints: ParameterHint[] = []; + const newContext: Partial = Object.assign({}, context); + + // Gathers all hints from all functions used within the command + walk(command, { + visitFunction: (funcNode) => { + const functionDefinition = getFunctionDefinition(funcNode.name); + + if (functionDefinition) { + for (const signature of functionDefinition.signatures) { + for (const param of signature.params) { + if (param.hint) { + hints.push(param.hint); + } + } + } + } + }, + }); + + // Remove duplicate hints + const uniqueHints = uniqWith( + hints, + (a, b) => a.entityType === b.entityType && isEqual(a.constraints, b.constraints) + ); + + // If the hint needs new data to build the suggestions, we add that data to the context + for (const hint of uniqueHints) { + const parameterHandler = parametersFromHintsResolvers[hint.entityType]; + if (parameterHandler?.contextResolver) { + const resolvedContext = await parameterHandler.contextResolver( + hint, + context, + callbacks ?? {} + ); + Object.assign(newContext, resolvedContext); + } } + + return newContext; };