diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/scalar_functions.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/scalar_functions.ts index 9ac4adb40a8ed..f7adb71f6b318 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/scalar_functions.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/scalar_functions.ts @@ -16595,6 +16595,12 @@ const textEmbeddingDefinition: FunctionDefinition = { optional: false, description: 'Identifier of an existing inference endpoint the that will generate the embeddings. The inference endpoint must have the `text_embedding` task type and should use the same model that was used to embed your indexed data.', + hint: { + entityType: 'inference_endpoint', + constraints: { + task_type: 'text_embedding', + }, + }, }, ], returnType: 'dense_vector', @@ -16614,8 +16620,6 @@ const textEmbeddingDefinition: FunctionDefinition = { Location.JOIN, ], examples: [ - 'ROW input="Who is Victor Hugo?"\n| EVAL embedding = TEXT_EMBEDDING("Who is Victor Hugo?", "test_dense_inference")', - 'FROM dense_vector_text METADATA _score\n| EVAL query_embedding = TEXT_EMBEDDING("be excellent to each other", "test_dense_inference")\n| WHERE KNN(text_embedding_field, query_embedding)', 'FROM dense_vector_text METADATA _score\n| WHERE KNN(text_embedding_field, TEXT_EMBEDDING("be excellent to each other", "test_dense_inference"))', ], }; @@ -19439,6 +19443,7 @@ const topSnippetsDefinition: FunctionDefinition = { Location.JOIN, ], examples: [ + 'FROM books\n| EVAL snippets = TOP_SNIPPETS(description, "Tolkien")', 'FROM books\n| WHERE MATCH(title, "Return")\n| EVAL snippets = TOP_SNIPPETS(description, "Tolkien", { "num_snippets": 3, "num_words": 25 })', ], }; @@ -19786,7 +19791,6 @@ const vCosineDefinition: FunctionDefinition = { description: i18n.translate('kbn-esql-ast.esql.definitions.v_cosine', { defaultMessage: 'Calculates the cosine similarity between two dense_vectors.', }), - ignoreAsSuggestion: true, preview: true, alias: undefined, signatures: [ @@ -19833,7 +19837,6 @@ const vDotProductDefinition: FunctionDefinition = { description: i18n.translate('kbn-esql-ast.esql.definitions.v_dot_product', { defaultMessage: 'Calculates the dot product between two dense_vectors.', }), - ignoreAsSuggestion: true, preview: true, alias: undefined, signatures: [ @@ -19880,7 +19883,6 @@ const vHammingDefinition: FunctionDefinition = { description: i18n.translate('kbn-esql-ast.esql.definitions.v_hamming', { defaultMessage: 'Calculates the Hamming distance between two dense vectors.', }), - ignoreAsSuggestion: true, preview: true, alias: undefined, signatures: [ @@ -19927,7 +19929,6 @@ const vL1NormDefinition: FunctionDefinition = { description: i18n.translate('kbn-esql-ast.esql.definitions.v_l1_norm', { defaultMessage: 'Calculates the l1 norm between two dense_vectors.', }), - ignoreAsSuggestion: true, preview: true, alias: undefined, signatures: [ @@ -19974,7 +19975,6 @@ const vL2NormDefinition: FunctionDefinition = { description: i18n.translate('kbn-esql-ast.esql.definitions.v_l2_norm', { defaultMessage: 'Calculates the l2 norm between two dense_vectors.', }), - ignoreAsSuggestion: true, preview: true, alias: undefined, signatures: [ diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/settings.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/settings.ts index 15d6751251558..66116d0acfd55 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/settings.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/generated/settings.ts @@ -13,7 +13,7 @@ const projectRouting = { type: 'keyword', serverlessOnly: true, preview: true, - snapshotOnly: true, + snapshotOnly: false, description: 'A project routing expression, used to define which projects to route the query to. Only supported if Cross-Project Search is enabled.', ignoreAsSuggestion: false, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/types.ts index fb76f42b9c237..fd66f1579b484 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/types.ts +++ b/src/platform/packages/shared/kbn-esql-ast/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; @@ -128,6 +135,12 @@ export interface FunctionParameter { * This indicates that the parameter can accept multiple values, which will be passed as an array. */ supportsMultiValues?: boolean; + + /** + * 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-ast/src/commands/definitions/utils/autocomplete/expressions/positions/empty_expression.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/utils/autocomplete/expressions/positions/empty_expression.ts index 59d923abd79dc..465582cd0a8dc 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/utils/autocomplete/expressions/positions/empty_expression.ts +++ b/src/platform/packages/shared/kbn-esql-ast/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 { parametersFromHintsMap } from '../../parameters_from_hints'; type FunctionParamContext = NonNullable; @@ -70,7 +72,7 @@ async function handleFunctionParameterContext( } // Try exclusive suggestions first (COUNT(*), enum values) - const exclusiveSuggestions = tryExclusiveSuggestions(functionParamContext, ctx); + const exclusiveSuggestions = await tryExclusiveSuggestions(functionParamContext, ctx); if (exclusiveSuggestions.length > 0) { return exclusiveSuggestions; @@ -81,10 +83,10 @@ async function handleFunctionParameterContext( } /** Try suggestions that are exclusive (if present, return only these) */ -function tryExclusiveSuggestions( +async function tryExclusiveSuggestions( functionParamContext: FunctionParamContext, ctx: ExpressionContext -): ISuggestionItem[] { +): Promise { const { functionDefinition, paramDefinitions } = functionParamContext; const { options } = ctx; @@ -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 = await buildSuggestionsFromHints(paramDefinitions, ctx); + if (itemsFromHints.length > 0) { + return itemsFromHints; + } + return []; } @@ -374,6 +381,23 @@ function buildEnumValueSuggestions( }); } +async function buildSuggestionsFromHints( + paramDefinitions: FunctionParameter[], + ctx: ExpressionContext +): Promise { + // 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 = await Promise.all( + hints.map((hint) => parametersFromHintsMap[hint.entityType]?.(hint, ctx) ?? []) + ); + + return results.flat(); +} + /** Builds suggestions for constant-only literal parameters */ function buildConstantOnlyLiteralSuggestions( paramDefinitions: FunctionParameter[], diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/utils/autocomplete/parameters_from_hints.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/utils/autocomplete/parameters_from_hints.ts new file mode 100644 index 0000000000000..9fd1ecf39eecf --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands/definitions/utils/autocomplete/parameters_from_hints.ts @@ -0,0 +1,51 @@ +/* + * 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 { i18n } from '@kbn/i18n'; +import type { ParameterHint, ParameterHintEntityType } from '../../..'; +import type { ISuggestionItem } from '../../../registry/types'; +import type { ExpressionContext } from './expressions/types'; +import { createInferenceEndpointToCompletionItem } from './helpers'; + +/** + * For some parameters, ES gives as hints about the nature of it, that we use to provide + * custom autocompletion handlers. + */ +export const parametersFromHintsMap: Record< + ParameterHintEntityType, + (hint: ParameterHint, ctx: ExpressionContext) => Promise +> = { + ['inference_endpoint']: inferenceEndpointHandler, +}; + +async function inferenceEndpointHandler( + hint: ParameterHint, + ctx: ExpressionContext +): Promise { + if (hint.constraints?.task_type) { + const inferenceEnpoints = + ( + await ctx.callbacks?.getInferenceEndpoints?.( + hint.constraints?.task_type as InferenceTaskType + ) + )?.inferenceEndpoints || []; + + return inferenceEnpoints.map(createInferenceEndpointToCompletionItem).map((item) => { + return { + ...item, + detail: i18n.translate('kbn-esql-ast.esql.definitions.inferenceEndpointInFunction', { + defaultMessage: 'Inference endpoint used for this function', + }), + text: `"${item.text}"`, + }; + }); + } + return []; +} diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands/registry/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands/registry/types.ts index 88a8fad0dafc0..d1707585fdfca 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands/registry/types.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands/registry/types.ts @@ -13,9 +13,11 @@ import type { ESQLControlVariable, ESQLSourceResult, ESQLFieldWithMetadata, + InferenceEndpointsAutocompleteResult, } from '@kbn/esql-types'; import type { LicenseType } from '@kbn/licensing-types'; import type { PricingProduct } from '@kbn/core-pricing-common/src/types'; +import type { InferenceTaskType } from '@elastic/elasticsearch/lib/api/types'; import type { ESQLLocation } from '../../types'; import type { SupportedDataType } from '../definitions/types'; import type { EditorExtensions } from './options/recommended_queries'; @@ -138,6 +140,9 @@ export interface ICommandCallbacks { getJoinIndices?: () => Promise<{ indices: IndexAutocompleteItem[] }>; canCreateLookupIndex?: (indexName: string) => Promise; isServerless?: boolean; + getInferenceEndpoints?: ( + taskType: InferenceTaskType + ) => Promise; } export interface ICommandContext { diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/functions.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/functions.test.ts index 36614f2a158f2..2188c0a737614 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/functions.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/functions.test.ts @@ -2010,4 +2010,14 @@ describe('functions arg suggestions', () => { expect(labels).not.toContain(','); }); }); + + describe('function parameter built from hint', () => { + it('suggests inference endpoints for TEXT_EMBEDDING function', async () => { + const { suggest } = await setup(); + const suggestions = await suggest('FROM index | EVAL result = TEXT_EMBEDDING("text", /'); + const labels = suggestions.map(({ label }) => label); + + expect(labels).toEqual(['inference_1']); + }); + }); }); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts index 1d0e582eaf241..a667143844bb7 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts @@ -273,6 +273,7 @@ async function getSuggestionsWithinCommandExpression( hasMinimumLicenseRequired, canCreateLookupIndex: callbacks?.canCreateLookupIndex, isServerless: callbacks?.isServerless, + getInferenceEndpoints: callbacks?.getInferenceEndpoints, }, context, offset