[ES|QL] Suggest adding curly braces after the WITH keyword for Rerank and Completion

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.test.ts

@@ -113,6 +113,12 @@ describe('COMPLETION Autocomplete', () => {
     await completionExpectSuggestions(`FROM a | COMPLETION "prompt" WIT/`, ['WITH { $0 }']);
   });
 
+  it('suggests opening braces when WITH is already typed', async () => {
+    await completionExpectSuggestions(`FROM a | COMPLETION "prompt" WITH `, [
+      '{ "inference_id": "$0" }',
+    ]);
+  });
+
   it('suggests inference_id parameter within the named parameters map', async () => {
     await completionExpectSuggestions(`FROM a | COMPLETION "prompt" WITH {`, [
       '"inference_id": "$0"',

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts

@@ -21,6 +21,7 @@ import {
   assignCompletionItem,
   getNewUserDefinedColumnSuggestion,
   withCompleteItem,
+  withMapCompleteItem,
 } from '../../complete_items';
 import {
   getFieldsSuggestions,
@@ -50,6 +51,7 @@ export enum CompletionPosition {
   AFTER_TARGET_ID = 'after_target_id',
   AFTER_PROMPT_OR_TARGET = 'after_prompt_or_target',
   AFTER_PROMPT = 'after_prompt',
+  AFTER_WITH_KEYWORD = 'after_with_keyword',
   WITHIN_MAP_EXPRESSION = 'within_map_expression',
   AFTER_COMMAND = 'after_command',
 }
@@ -62,7 +64,16 @@ function getPosition(
 ): CompletionPosition | undefined {
   const { prompt, targetField } = command as ESQLAstCompletionCommand;
 
-  const paramsMap = command.args[1] as ast.ESQLMap | undefined;
+  const arg1 = command.args[1];
+  let paramsMap: ast.ESQLMap | undefined;
+
+  if (arg1 && 'type' in arg1 && arg1.type === 'option') {
+    paramsMap = (arg1 as ast.ESQLCommandOption).args[0] as ast.ESQLMap;
+
+    if (paramsMap && paramsMap.incomplete && !paramsMap.text) {
+      return CompletionPosition.AFTER_WITH_KEYWORD;
+    }
+  }
 
   if (paramsMap?.text && paramsMap.incomplete) {
     return CompletionPosition.WITHIN_MAP_EXPRESSION;
@@ -242,6 +253,9 @@ export async function autocomplete(
       ];
     }
 
+    case CompletionPosition.AFTER_WITH_KEYWORD:
+      return [withMapCompleteItem];
+
     case CompletionPosition.WITHIN_MAP_EXPRESSION:
       const availableParameters: MapParameters = {
         inference_id: {

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/autocomplete.test.ts

@@ -307,6 +307,16 @@ describe('RERANK Autocomplete', () => {
   // ============================================================================
 
   describe('WITH clause functionality', () => {
+    test('suggests opening braces with inference_id when WITH is already typed', async () => {
+      const query =
+        buildRerankQuery({
+          query: '"search query"',
+          onClause: 'textField',
+        }) + ' WITH ';
+
+      await expectRerankSuggestions(query, ['{ "inference_id": "$0" }']);
+    });
+
     test('suggests inference_id key in WITH map', async () => {
       const query = buildRerankQuery({
         query: '"search query"',

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/autocomplete.ts

@@ -15,6 +15,7 @@ import {
   getNewUserDefinedColumnSuggestion,
   onCompleteItem,
   assignCompletionItem,
+  withMapCompleteItem,
 } from '../../complete_items';
 import {
   withinQuotes,
@@ -116,6 +117,9 @@ export async function autocomplete(
       });
     }
 
+    case CaretPosition.AFTER_WITH_KEYWORD:
+      return [withMapCompleteItem];
+
     case CaretPosition.WITHIN_MAP_EXPRESSION: {
       const endpoints = context?.inferenceEndpoints;
       const availableParameters: MapParameters = {

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/utils.ts

@@ -11,6 +11,7 @@ import type {
   ESQLMap,
   ESQLSingleAstItem,
   ESQLAstAllCommands,
+  ESQLCommandOption,
 } from '../../../types';
 import { isAssignment, isFunctionExpression } from '../../../ast/is';
 import { within } from '../../../ast/location';
@@ -35,6 +36,7 @@ export enum CaretPosition {
   ON_WITHIN_FIELD_LIST, // After "ON": suggest field names
   ON_KEEP_SUGGESTIONS_AFTER_TRAILING_SPACE, // Special case: After a complete field with space, suggest next actions or assignment
   ON_EXPRESSION, // After "ON": handle all field list expressions like EVAL
+  AFTER_WITH_KEYWORD, // After "WITH " but before opening brace: suggest opening braces with params
   WITHIN_MAP_EXPRESSION, // After "WITH": suggest a json of params
   AFTER_COMMAND, // Command is complete, suggest pipe
 }
@@ -50,9 +52,15 @@ export function getPosition(
   const rerankCommand = command as ESQLAstRerankCommand;
   const innerText = query.substring(rerankCommand.location.min);
   const onMap = rerankCommand.args[1];
-  const withMap = rerankCommand.args[2] as ESQLMap | undefined;
+  const withArg = rerankCommand.args[2];
+
+  if (withArg && 'type' in withArg && withArg.type === 'option') {
+    const withMap = (withArg as ESQLCommandOption).args[0] as ESQLMap | undefined;
+
+    if (!withMap || (withMap.incomplete && !withMap.text)) {
+      return { position: CaretPosition.AFTER_WITH_KEYWORD };
+    }
 
-  if (withMap) {
     if (withMap.text && withMap.incomplete) {
       return { position: CaretPosition.WITHIN_MAP_EXPRESSION };
     }

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

@@ -101,6 +101,15 @@ export const withCompleteItem: ISuggestionItem = withAutoSuggest({
   sortText: '1',
 });
 
+export const withMapCompleteItem: ISuggestionItem = withAutoSuggest({
+  label: 'inference_id',
+  text: '{ "inference_id": "$0" }',
+  asSnippet: true,
+  kind: 'Reference',
+  detail: 'Inference endpoint',
+  sortText: '1',
+});
+
 export const subqueryCompleteItem: ISuggestionItem = withAutoSuggest({
   label: '(FROM ...)',
   text: '(FROM $0)',