diff --git a/x-pack/plugins/inference/common/index.ts b/x-pack/plugins/inference/common/index.ts index 79433cbc71a68..6098c3fbd9327 100644 --- a/x-pack/plugins/inference/common/index.ts +++ b/x-pack/plugins/inference/common/index.ts @@ -5,13 +5,7 @@ * 2.0. */ -export { - correctCommonEsqlMistakes, - splitIntoCommands, -} from './tasks/nl_to_esql/correct_common_esql_mistakes'; - +export { correctCommonEsqlMistakes, splitIntoCommands } from './tasks/nl_to_esql'; export { generateFakeToolCallId } from './utils/generate_fake_tool_call_id'; - export { createOutputApi } from './output'; - export type { ChatCompleteRequestBody, GetConnectorsResponseBody } from './http_apis'; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/ast_tools/timespan.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/ast_tools/timespan.ts new file mode 100644 index 0000000000000..25ff8c1a00488 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/ast_tools/timespan.ts @@ -0,0 +1,76 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { ESQLTimeInterval } from '@kbn/esql-ast'; + +const units = [ + 'millisecond', + 'milliseconds', + 'ms', + // + 'second', + 'seconds', + 'sec', + 's', + // + 'minute', + 'minutes', + 'min', + // + 'hour', + 'hours', + 'h', + // + 'day', + 'days', + 'd', + // + 'week', + 'weeks', + 'w', + // + 'month', + 'months', + 'mo', + // + 'quarter', + 'quarters', + 'q', + // + 'year', + 'years', + 'yr', + 'y', +]; + +const timespanStringRegexp = new RegExp(`^["']?([0-9]+)?\\s*?(${units.join('|')})["']?$`, 'i'); + +export function createTimespanLiteral(unit: string, quantity: number): ESQLTimeInterval { + return { + type: 'timeInterval', + quantity, + unit, + text: `${unit}${quantity}`, + name: `${unit} ${quantity}`, + incomplete: false, + location: { min: 0, max: 0 }, + }; +} + +export function isTimespanString(str: string): boolean { + return Boolean(str.match(timespanStringRegexp)); +} + +export function stringToTimespanLiteral(str: string): ESQLTimeInterval { + const match = timespanStringRegexp.exec(str); + if (!match) { + throw new Error(`String "${str}" cannot be converted to timespan literal`); + } + const [_, quantity, unit] = match; + + return createTimespanLiteral(unit.toLowerCase(), quantity ? parseInt(quantity, 10) : 1); +} diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/correct_with_ast.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/correct_with_ast.ts new file mode 100644 index 0000000000000..b3a90b79e03a0 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/correct_with_ast.ts @@ -0,0 +1,35 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { BasicPrettyPrinter, parse } from '@kbn/esql-ast'; +import { correctAll, type QueryCorrection } from './corrections'; + +interface CorrectWithAstResult { + output: string; + corrections: QueryCorrection[]; +} + +export const correctQueryWithAst = (query: string): CorrectWithAstResult => { + const { root, errors } = parse(query); + // don't try modifying anything if the query is not syntactically correct + if (errors) { + return { + output: query, + corrections: [], + }; + } + + const corrections = correctAll(root); + + const multiline = /\r?\n/.test(query); + const formattedQuery = BasicPrettyPrinter.print(root, { multiline, pipeTab: '' }); + + return { + output: formattedQuery, + corrections, + }; +}; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/index.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/index.ts new file mode 100644 index 0000000000000..9d7218a6f77cb --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/index.ts @@ -0,0 +1,18 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { ESQLAstQueryExpression } from '@kbn/esql-ast'; +import type { QueryCorrection } from './types'; +import { applyTimespanLiteralsCorrections } from './timespan_literals'; + +export type { QueryCorrection } from './types'; + +export const correctAll = (query: ESQLAstQueryExpression): QueryCorrection[] => { + const corrections: QueryCorrection[] = []; + corrections.push(...applyTimespanLiteralsCorrections(query)); + return corrections; +}; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/timespan_literals.test.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/timespan_literals.test.ts new file mode 100644 index 0000000000000..616ed70802e96 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/timespan_literals.test.ts @@ -0,0 +1,144 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { parse, BasicPrettyPrinter } from '@kbn/esql-ast'; +import { applyTimespanLiteralsCorrections } from './timespan_literals'; + +describe('getTimespanLiteralsCorrections', () => { + describe('with DATE_TRUNC', () => { + it('replaces a timespan with a proper timespan literal', () => { + const query = 'FROM logs | EVAL truncated = DATE_TRUNC("1 year", date)'; + const { root } = parse(query); + + applyTimespanLiteralsCorrections(root); + + const output = BasicPrettyPrinter.print(root); + + expect(output).toMatchInlineSnapshot( + `"FROM logs | EVAL truncated = DATE_TRUNC(1 year, date)"` + ); + }); + + it('replaces a timespan without quantity', () => { + const query = 'FROM logs | EVAL truncated = DATE_TRUNC("month", date)'; + const { root } = parse(query); + + applyTimespanLiteralsCorrections(root); + + const output = BasicPrettyPrinter.print(root); + + expect(output).toMatchInlineSnapshot( + `"FROM logs | EVAL truncated = DATE_TRUNC(1 month, date)"` + ); + }); + + it('replaces uppercase literals', () => { + const query = 'FROM logs | EVAL truncated = DATE_TRUNC("1 YEAR", date)'; + const { root } = parse(query); + + applyTimespanLiteralsCorrections(root); + + const output = BasicPrettyPrinter.print(root); + + expect(output).toMatchInlineSnapshot( + `"FROM logs | EVAL truncated = DATE_TRUNC(1 year, date)"` + ); + }); + + it('returns info about the correction', () => { + const query = 'FROM logs | EVAL truncated = DATE_TRUNC("1 year", date)'; + const { root } = parse(query); + + const corrections = applyTimespanLiteralsCorrections(root); + + expect(corrections).toHaveLength(1); + expect(corrections[0]).toEqual({ + type: 'string_as_timespan_literal', + description: + 'Replaced string literal with timespan literal in DATE_TRUNC function at position 29', + node: expect.any(Object), + }); + }); + }); + + describe('with BUCKET', () => { + it('replaces a timespan with a proper timespan literal', () => { + const query = 'FROM logs | STATS hires = COUNT(*) BY week = BUCKET(hire_date, "1 week")'; + const { root } = parse(query); + + applyTimespanLiteralsCorrections(root); + + const output = BasicPrettyPrinter.print(root); + + expect(output).toMatchInlineSnapshot( + `"FROM logs | STATS hires = COUNT(*) BY week = BUCKET(hire_date, 1 week)"` + ); + }); + + it('replaces a timespan without quantity', () => { + const query = 'FROM logs | STATS hires = COUNT(*) BY hour = BUCKET(hire_date, "hour")'; + const { root } = parse(query); + + applyTimespanLiteralsCorrections(root); + + const output = BasicPrettyPrinter.print(root); + + expect(output).toMatchInlineSnapshot( + `"FROM logs | STATS hires = COUNT(*) BY hour = BUCKET(hire_date, 1 hour)"` + ); + }); + + it('replaces uppercase literals', () => { + const query = 'FROM logs | STATS hires = COUNT(*) BY week = BUCKET(hire_date, "1 WEEK")'; + const { root } = parse(query); + + applyTimespanLiteralsCorrections(root); + + const output = BasicPrettyPrinter.print(root); + + expect(output).toMatchInlineSnapshot( + `"FROM logs | STATS hires = COUNT(*) BY week = BUCKET(hire_date, 1 week)"` + ); + }); + + it('returns info about the correction', () => { + const query = 'FROM logs | STATS hires = COUNT(*) BY hour = BUCKET(hire_date, "hour")'; + const { root } = parse(query); + + const corrections = applyTimespanLiteralsCorrections(root); + + expect(corrections).toHaveLength(1); + expect(corrections[0]).toEqual({ + type: 'string_as_timespan_literal', + description: + 'Replaced string literal with timespan literal in BUCKET function at position 45', + node: expect.any(Object), + }); + }); + }); + + describe('with mixed usages', () => { + it('find all occurrences in a complex query', () => { + const query = `FROM logs + | EVAL trunc_year = DATE_TRUNC("1 year", date) + | EVAL trunc_month = DATE_TRUNC("month", date) + | STATS hires = COUNT(*) BY hour = BUCKET(hire_date, "3 hour")`; + const { root } = parse(query); + + applyTimespanLiteralsCorrections(root); + + const output = BasicPrettyPrinter.print(root, { multiline: true, pipeTab: '' }); + + expect(output).toMatchInlineSnapshot(` + "FROM logs + | EVAL trunc_year = DATE_TRUNC(1 year, date) + | EVAL trunc_month = DATE_TRUNC(1 month, date) + | STATS hires = COUNT(*) BY hour = BUCKET(hire_date, 3 hour)" + `); + }); + }); +}); diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/timespan_literals.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/timespan_literals.ts new file mode 100644 index 0000000000000..c3fbe636a2de1 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/timespan_literals.ts @@ -0,0 +1,84 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { Walker, type ESQLAstQueryExpression } from '@kbn/esql-ast'; +import { isDateTruncFunctionNode, isBucketFunctionNode, isStringLiteralNode } from '../typeguards'; +import type { ESQLDateTruncFunction, ESQLBucketFunction } from '../types'; +import { stringToTimespanLiteral, isTimespanString } from '../ast_tools/timespan'; +import { QueryCorrection } from './types'; + +/** + * Correct timespan literal grammar mistakes, and returns the list of corrections that got applied. + * + * E.g. + * `DATE_TRUNC("YEAR", @timestamp)` => `DATE_TRUNC(1 year, @timestamp)` + * `BUCKET(@timestamp, "1 week")` => `BUCKET(@timestamp, 1 week)` + * + */ +export const applyTimespanLiteralsCorrections = ( + query: ESQLAstQueryExpression +): QueryCorrection[] => { + const corrections: QueryCorrection[] = []; + + Walker.walk(query, { + visitFunction: (node) => { + if (isDateTruncFunctionNode(node)) { + corrections.push(...checkDateTrunc(node)); + } + if (isBucketFunctionNode(node)) { + corrections.push(...checkBucket(node)); + } + }, + }); + + return corrections; +}; + +function checkDateTrunc(node: ESQLDateTruncFunction): QueryCorrection[] { + if (node.args.length !== 2) { + return []; + } + + const firstArg = node.args[0]; + + if (isStringLiteralNode(firstArg) && isTimespanString(firstArg.value)) { + const replacement = stringToTimespanLiteral(firstArg.value); + node.args[0] = replacement; + + const correction: QueryCorrection = { + type: 'string_as_timespan_literal', + node, + description: `Replaced string literal with timespan literal in DATE_TRUNC function at position ${node.location.min}`, + }; + return [correction]; + } + + return []; +} + +function checkBucket(node: ESQLBucketFunction): QueryCorrection[] { + // only checking the 2 args version - e.g. BUCKET(hire_date, 1 week) + if (node.args.length !== 2) { + return []; + } + + const secondArg = node.args[1]; + + if (isStringLiteralNode(secondArg) && isTimespanString(secondArg.value)) { + const replacement = stringToTimespanLiteral(secondArg.value); + node.args[1] = replacement; + + const correction: QueryCorrection = { + type: 'string_as_timespan_literal', + node, + description: `Replaced string literal with timespan literal in BUCKET function at position ${node.location.min}`, + }; + return [correction]; + } + + return []; +} diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/types.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/types.ts new file mode 100644 index 0000000000000..dc1210c84acc2 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/corrections/types.ts @@ -0,0 +1,20 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { ESQLSingleAstItem } from '@kbn/esql-ast'; + +/** + * Represents a correction that was applied to the query + */ +export interface QueryCorrection { + /** The type of correction */ + type: string; + /** A human-friendly-ish description of the correction */ + description: string; + /** The parent node the correction was applied to */ + node: ESQLSingleAstItem; +} diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/index.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/index.ts new file mode 100644 index 0000000000000..d0e1a4808829a --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/index.ts @@ -0,0 +1,8 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export { correctQueryWithAst } from './correct_with_ast'; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/typeguards.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/typeguards.ts new file mode 100644 index 0000000000000..233625673b872 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/typeguards.ts @@ -0,0 +1,33 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { ESQLSingleAstItem, ESQLAstItem, ESQLFunction, ESQLLiteral } from '@kbn/esql-ast'; +import type { ESQLStringLiteral, ESQLDateTruncFunction, ESQLBucketFunction } from './types'; + +export function isSingleItem(item: ESQLAstItem): item is ESQLSingleAstItem { + return Object.hasOwn(item, 'type'); +} + +export function isFunctionNode(node: ESQLAstItem): node is ESQLFunction { + return isSingleItem(node) && node.type === 'function'; +} + +export function isLiteralNode(node: ESQLAstItem): node is ESQLLiteral { + return isSingleItem(node) && node.type === 'literal'; +} + +export function isStringLiteralNode(node: ESQLAstItem): node is ESQLStringLiteral { + return isLiteralNode(node) && node.literalType === 'keyword'; +} + +export function isDateTruncFunctionNode(node: ESQLAstItem): node is ESQLDateTruncFunction { + return isFunctionNode(node) && node.subtype === 'variadic-call' && node.name === 'date_trunc'; +} + +export function isBucketFunctionNode(node: ESQLAstItem): node is ESQLBucketFunction { + return isFunctionNode(node) && node.subtype === 'variadic-call' && node.name === 'bucket'; +} diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/types.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/types.ts new file mode 100644 index 0000000000000..dd2a9810e359e --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/ast/types.ts @@ -0,0 +1,23 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { ESQLFunction, ESQLLiteral } from '@kbn/esql-ast'; + +/** + * represents a DATE_TRUNC function node. + */ +export type ESQLDateTruncFunction = ESQLFunction<'variadic-call', 'date_trunc'>; + +/** + * represents a BUCKET function node. + */ +export type ESQLBucketFunction = ESQLFunction<'variadic-call', 'bucket'>; + +/** + * represents an ESQL string literal. + */ +export type ESQLStringLiteral = Extract; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_esql_query.test.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_esql_query.test.ts new file mode 100644 index 0000000000000..22ae55f772e95 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_esql_query.test.ts @@ -0,0 +1,84 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { correctCommonEsqlMistakes } from './correct_esql_query'; + +jest.mock('./ast'); +jest.mock('./non_ast'); +import { correctQueryWithAst } from './ast'; +import { correctCommonEsqlMistakes as correctQueryWithoutAst } from './non_ast'; + +const correctQueryWithAstMock = correctQueryWithAst as jest.MockedFn; +const correctQueryWithoutAstMock = correctQueryWithoutAst as jest.MockedFn< + typeof correctQueryWithoutAst +>; + +describe('correctCommonEsqlMistakes', () => { + beforeEach(() => { + correctQueryWithoutAstMock.mockImplementation((query) => { + return { + input: query, + output: query, + isCorrection: false, + }; + }); + correctQueryWithAstMock.mockImplementation((query) => { + return { + output: query, + corrections: [], + }; + }); + }); + + afterEach(() => { + correctQueryWithoutAstMock.mockReset(); + correctQueryWithAstMock.mockReset(); + }); + + it('calls correctQueryWithoutAst with the right parameters', () => { + const inputQuery = 'FROM logs | WHERE foo > bar'; + correctCommonEsqlMistakes(inputQuery); + + expect(correctQueryWithoutAstMock).toHaveBeenCalledTimes(1); + expect(correctQueryWithoutAstMock).toHaveBeenCalledWith(inputQuery); + }); + + it('calls correctQueryWithAst with the right parameters', () => { + const inputQuery = 'FROM logs | WHERE foo > bar'; + const correctQueryWithoutAstResult = 'FROM logs | WHERE foo > dolly'; + + correctQueryWithoutAstMock.mockImplementation((query) => { + return { + input: inputQuery, + output: correctQueryWithoutAstResult, + isCorrection: true, + }; + }); + + correctCommonEsqlMistakes(inputQuery); + + expect(correctQueryWithAstMock).toHaveBeenCalledTimes(1); + expect(correctQueryWithAstMock).toHaveBeenCalledWith(correctQueryWithoutAstResult); + }); + + it('returns the corrected query', () => { + const inputQuery = 'FROM logs | WHERE foo > bar'; + const correctQueryWithAstResult = 'FROM logs | WHERE foo > dolly'; + + correctQueryWithAstMock.mockImplementation((query) => { + return { + output: correctQueryWithAstResult, + corrections: [], + }; + }); + + const { input, output } = correctCommonEsqlMistakes(inputQuery); + + expect(input).toEqual(inputQuery); + expect(output).toEqual(correctQueryWithAstResult); + }); +}); diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_esql_query.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_esql_query.ts new file mode 100644 index 0000000000000..4ad96323cc7ea --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_esql_query.ts @@ -0,0 +1,27 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { correctQueryWithAst } from './ast'; +import { correctCommonEsqlMistakes as correctQueryWithoutAst } from './non_ast'; + +/** + * Correct some common ES|QL syntax and grammar mistakes that LLM can potentially do. + * + * Correcting the query is done in two steps: + * 1. we try to correct the *syntax*, without AST (given it requires a valid syntax) + * 2. we try to correct the *grammar*, using AST this time. + */ +export const correctCommonEsqlMistakes = (query: string) => { + const { output: outputWithoutAst, isCorrection: correctionWithoutAst } = + correctQueryWithoutAst(query); + const { output: corrected, corrections } = correctQueryWithAst(outputWithoutAst); + return { + input: query, + output: corrected, + isCorrection: correctionWithoutAst || corrections.length > 0, + }; +}; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.test.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.test.ts deleted file mode 100644 index 818f4854e038d..0000000000000 --- a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.test.ts +++ /dev/null @@ -1,46 +0,0 @@ -/* - * 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; you may not use this file except in compliance with the Elastic License - * 2.0. - */ -import { correctQueryWithActions } from './correct_query_with_actions'; - -describe('correctQueryWithActions', () => { - it(`fixes errors correctly for a query with one syntax error for stats`, async () => { - const fixedQuery = await correctQueryWithActions('from logstash-* | stats aveg(bytes)'); - expect(fixedQuery).toBe('from logstash-* | stats avg(bytes)'); - }); - - it(`fixes errors correctly for a query with 2 syntax errors for stats`, async () => { - const fixedQuery = await correctQueryWithActions( - 'from logstash-* | stats aveg(bytes), max2(bytes)' - ); - expect(fixedQuery).toBe('from logstash-* | stats avg(bytes), max(bytes)'); - }); - - it(`fixes errors correctly for a query with one syntax error for eval`, async () => { - const fixedQuery = await correctQueryWithActions( - 'from logstash-* | stats var0 = max(bytes) | eval ab(var0) | limit 1' - ); - expect(fixedQuery).toBe('from logstash-* | stats var0 = max(bytes) | eval abs(var0) | limit 1'); - }); - - it(`fixes errors correctly for a query with two syntax error for eval`, async () => { - const fixedQuery = await correctQueryWithActions( - 'from logstash-* | stats var0 = max2(bytes) | eval ab(var0) | limit 1' - ); - expect(fixedQuery).toBe('from logstash-* | stats var0 = max(bytes) | eval abs(var0) | limit 1'); - }); - - it(`doesnt complain for @timestamp column`, async () => { - const queryWithTimestamp = `FROM logstash-* - | WHERE @timestamp >= NOW() - 15 minutes - | EVAL bucket = DATE_TRUNC(1 minute, @timestamp) - | STATS avg_cpu = AVG(system.cpu.total.norm.pct) BY service.name, bucket - | SORT avg_cpu DESC - | LIMIT 10`; - const fixedQuery = await correctQueryWithActions(queryWithTimestamp); - expect(fixedQuery).toBe(queryWithTimestamp); - }); -}); diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.ts deleted file mode 100644 index 30e2c11adb6de..0000000000000 --- a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.ts +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -import { validateQuery, getActions } from '@kbn/esql-validation-autocomplete'; -import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; - -const fixedQueryByOneAction = async (queryString: string) => { - const { errors } = await validateQuery(queryString, getAstAndSyntaxErrors, { - ignoreOnMissingCallbacks: true, - }); - - const actions = await getActions(queryString, errors, getAstAndSyntaxErrors, { - relaxOnMissingCallbacks: true, - }); - - if (actions.length) { - const [firstAction] = actions; - const range = firstAction.edits[0].range; - const correctText = firstAction.edits[0].text; - const problematicString = queryString.substring(range.startColumn - 1, range.endColumn - 1); - const fixedQuery = queryString.replace(problematicString, correctText); - - return { - query: fixedQuery, - shouldRunAgain: Boolean(actions.length), - }; - } - return { - query: queryString, - shouldRunAgain: false, - }; -}; - -/** - * @param queryString - * @returns corrected queryString - * The cases that are handled are: - * - Query stats / eval functions have typos e.g. aveg instead of avg - * - Unquoted fields e.g. keep field-1 instead of keep `field-1` - * - Unquoted fields in stats or eval e.g. stats avg(field-1) instead of stats avg(`field-1`) - * - Combination of the above - */ - -export const correctQueryWithActions = async (queryString: string) => { - let shouldCorrectQuery = true; - let fixedQuery = queryString; - // this is an escape hatch, the loop will end automatically if the ast doesnt return more actions - // in case it goes wrong, we allow it to loop 10 times - let limit = 10; - - while (shouldCorrectQuery && limit >= 0) { - const { query, shouldRunAgain } = await fixedQueryByOneAction(fixedQuery); - shouldCorrectQuery = shouldRunAgain; - fixedQuery = query; - limit--; - } - - return fixedQuery; -}; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/get_errors_with_commands.test.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/get_errors_with_commands.test.ts deleted file mode 100644 index 04d99aede62b8..0000000000000 --- a/x-pack/plugins/inference/common/tasks/nl_to_esql/get_errors_with_commands.test.ts +++ /dev/null @@ -1,25 +0,0 @@ -/* - * 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; you may not use this file except in compliance with the Elastic License - * 2.0. - */ -import { getErrorsWithCommands } from './get_errors_with_commands'; - -describe('getErrorsWithCommands', () => { - it('returns the command associated with the error', () => { - expect( - getErrorsWithCommands(`FROM logs-* | WHERE @timestamp <= NOW() | STATS BY host.name`, [ - { - type: 'error', - text: 'Syntax error', - code: '', - location: { - min: 24, - max: 36, - }, - }, - ]) - ).toEqual(['Error in `| WHERE @timestamp <= NOW()`:\n Syntax error']); - }); -}); diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/get_errors_with_commands.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/get_errors_with_commands.ts deleted file mode 100644 index b25c79161f79b..0000000000000 --- a/x-pack/plugins/inference/common/tasks/nl_to_esql/get_errors_with_commands.ts +++ /dev/null @@ -1,26 +0,0 @@ -/* - * 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; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -import type { EditorError, ESQLMessage } from '@kbn/esql-ast'; -import { splitIntoCommands } from './correct_common_esql_mistakes'; - -export function getErrorsWithCommands(query: string, errors: Array) { - const asCommands = splitIntoCommands(query); - - const errorMessages = errors.map((error) => { - if ('location' in error) { - const commandsUntilEndOfError = splitIntoCommands(query.substring(0, error.location.max)); - const lastCompleteCommand = asCommands[commandsUntilEndOfError.length - 1]; - if (lastCompleteCommand) { - return `Error in \`| ${lastCompleteCommand.command}\`:\n ${error.text}`; - } - } - return 'text' in error ? error.text : error.message; - }); - - return errorMessages; -} diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/index.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/index.ts new file mode 100644 index 0000000000000..2efacc374e8cd --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/index.ts @@ -0,0 +1,9 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export { correctCommonEsqlMistakes } from './correct_esql_query'; +export { splitIntoCommands } from './non_ast'; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_common_esql_mistakes.test.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/correct_common_esql_mistakes.test.ts similarity index 99% rename from x-pack/plugins/inference/common/tasks/nl_to_esql/correct_common_esql_mistakes.test.ts rename to x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/correct_common_esql_mistakes.test.ts index 8c6b90bca76ab..9d4b064b35f61 100644 --- a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_common_esql_mistakes.test.ts +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/correct_common_esql_mistakes.test.ts @@ -4,6 +4,7 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes'; describe('correctCommonEsqlMistakes', () => { diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_common_esql_mistakes.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/correct_common_esql_mistakes.ts similarity index 100% rename from x-pack/plugins/inference/common/tasks/nl_to_esql/correct_common_esql_mistakes.ts rename to x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/correct_common_esql_mistakes.ts diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/index.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/index.ts new file mode 100644 index 0000000000000..2752f82ce9c43 --- /dev/null +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/non_ast/index.ts @@ -0,0 +1,8 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +export { correctCommonEsqlMistakes, splitIntoCommands } from './correct_common_esql_mistakes'; diff --git a/x-pack/plugins/inference/scripts/load_esql_docs/load_esql_docs.ts b/x-pack/plugins/inference/scripts/load_esql_docs/load_esql_docs.ts index a35491a476040..fba1d75956bf0 100644 --- a/x-pack/plugins/inference/scripts/load_esql_docs/load_esql_docs.ts +++ b/x-pack/plugins/inference/scripts/load_esql_docs/load_esql_docs.ts @@ -13,7 +13,7 @@ import Path from 'path'; import yargs, { Argv } from 'yargs'; import { REPO_ROOT } from '@kbn/repo-info'; import { INLINE_ESQL_QUERY_REGEX } from '../../common/tasks/nl_to_esql/constants'; -import { correctCommonEsqlMistakes } from '../../common/tasks/nl_to_esql/correct_common_esql_mistakes'; +import { correctCommonEsqlMistakes } from '../../common/tasks/nl_to_esql'; import { connectorIdOption, elasticsearchOption, kibanaOption } from '../util/cli_options'; import { getServiceUrls } from '../util/get_service_urls'; import { KibanaClient } from '../util/kibana_client'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts index 823344f52a891..049e627303d1b 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts @@ -11,7 +11,7 @@ import type { ElasticsearchClient } from '@kbn/core/server'; import { ESQLSearchResponse, ESQLRow } from '@kbn/es-types'; import { esFieldTypeToKibanaFieldType } from '@kbn/field-types'; import { DatatableColumn, DatatableColumnType } from '@kbn/expressions-plugin/common'; -import { splitIntoCommands } from '../../../common/tasks/nl_to_esql/correct_common_esql_mistakes'; +import { splitIntoCommands } from '../../../common'; export async function runAndValidateEsqlQuery({ query,