Skip to content

Commit 6f7bef3

Browse files
authored
[ES|QL] Implement RERANK Ast Parser + basic template (elastic#232885)
## Summary Part of: elastic#217285 Grammar: ` rerankCommand: RERANK (targetField=qualifiedName ASSIGN)? queryText=constant ON rerankFields commandNamedParameters ` Grammar -> commandParameters expanded: `commandNamedParameters: (WITH mapExpression)?` `mapExpression: LEFT_BRACES (entryExpression (COMMA entryExpression)*)? RIGHT_BRACES` Syntax: `RERANK [column =] query ON field [, field, ...] [WITH { "inference_id" : "my_inference_endpoint" }]`
1 parent e9eed98 commit 6f7bef3

File tree

8 files changed

+609
-1
lines changed

8 files changed

+609
-1
lines changed
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
import type { ESQLCommand } from '../../../types';
11+
import { pipeCompleteItem } from '../../complete_items';
12+
import type { ICommandCallbacks } from '../../types';
13+
import { type ISuggestionItem, type ICommandContext } from '../../types';
14+
15+
export async function autocomplete(
16+
query: string,
17+
command: ESQLCommand,
18+
callbacks?: ICommandCallbacks,
19+
context?: ICommandContext,
20+
cursorPosition?: number
21+
): Promise<ISuggestionItem[]> {
22+
if (!callbacks?.getByType) {
23+
return [];
24+
}
25+
26+
const suggestions: ISuggestionItem[] = [];
27+
28+
// TODO: Add more specific autocomplete logic for RERANK command
29+
const fieldSuggestions = await callbacks.getByType('any');
30+
suggestions.push(...fieldSuggestions);
31+
32+
// Add the pipe completion item when appropriate
33+
suggestions.push(pipeCompleteItem);
34+
35+
return suggestions;
36+
}
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
import { i18n } from '@kbn/i18n';
11+
import type { ICommandMethods } from '../../registry';
12+
import { autocomplete } from './autocomplete';
13+
import { validate } from './validate';
14+
import type { ICommandContext } from '../../types';
15+
16+
const rerankCommandMethods: ICommandMethods<ICommandContext> = {
17+
autocomplete,
18+
validate,
19+
};
20+
21+
export const rerankCommand = {
22+
name: 'rerank',
23+
methods: rerankCommandMethods,
24+
hidden: true,
25+
metadata: {
26+
description: i18n.translate('kbn-esql-ast.esql.definitions.rerankDoc', {
27+
defaultMessage:
28+
'Uses an inference model to compute new relevance scores for documents, directly within your ES|QL queries.',
29+
}),
30+
declaration:
31+
'RERANK [target_field =] query_text ON field1 [, field2, ...] WITH { "inference_id": "model_id" }',
32+
examples: [
33+
'FROM movies | RERANK "star wars" ON title WITH { "inference_id": "reranker" }',
34+
'FROM books | RERANK rerank_score = "hobbit" ON title, description WITH { "inference_id": "my_reranker" }',
35+
],
36+
preview: true,
37+
},
38+
};
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
import { i18n } from '@kbn/i18n';
11+
import type { ESQLCommand, ESQLMessage, ESQLAst, ESQLAstRerankCommand } from '../../../types';
12+
import type { ICommandContext, ICommandCallbacks } from '../../types';
13+
import { validateCommandArguments } from '../../../definitions/utils/validation';
14+
15+
export const validate = (
16+
command: ESQLCommand,
17+
ast: ESQLAst,
18+
context?: ICommandContext,
19+
callbacks?: ICommandCallbacks
20+
): ESQLMessage[] => {
21+
const messages: ESQLMessage[] = [];
22+
23+
// Run standard argument validation
24+
messages.push(...validateCommandArguments(command, ast, context, callbacks));
25+
26+
// Cast to RERANK command for type-specific validation
27+
const rerankCommand = command as ESQLAstRerankCommand;
28+
29+
// Validate that query text is provided
30+
if (!rerankCommand.query) {
31+
messages.push({
32+
location: command.location,
33+
text: i18n.translate('kbn-esql-ast.esql.validation.rerankMissingQuery', {
34+
defaultMessage: '[RERANK] Query text is required.',
35+
}),
36+
type: 'error',
37+
code: 'rerankMissingQuery',
38+
});
39+
}
40+
41+
// Validate that at least one field is specified in ON clause
42+
if (!rerankCommand.fields || rerankCommand.fields.length === 0) {
43+
messages.push({
44+
location: command.location,
45+
text: i18n.translate('kbn-esql-ast.esql.validation.rerankMissingFields', {
46+
defaultMessage: '[RERANK] At least one field must be specified in the ON clause.',
47+
}),
48+
type: 'error',
49+
code: 'rerankMissingFields',
50+
});
51+
}
52+
53+
return messages;
54+
};

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import { showCommand } from './commands/show';
3131
import { timeseriesCommand } from './commands/timeseries';
3232
import { whereCommand } from './commands/where';
3333
import { fuseCommand } from './commands/fuse';
34+
import { rerankCommand } from './commands/rerank';
3435
import { mergeCommandWithGeneratedCommandData } from './elastisearch_command_data_loader';
3536

3637
const esqlCommandRegistry = new CommandRegistry();
@@ -59,6 +60,7 @@ const baseCommands = [
5960
timeseriesCommand,
6061
whereCommand,
6162
fuseCommand,
63+
rerankCommand,
6264
];
6365

6466
baseCommands.forEach((command) => {

0 commit comments

Comments
 (0)