[Agent Builder] add row limit and custom instructions to index_search tool type (#243490)

## Summary

Add `Row Limit` and `Custom instructions` to index search tool. 

- Row limit: configurable row limit to prevent overflowing agent context
with large response
- We achieve this by forking inference plugin instructions for
generating esql (no syntax tho) and making it templated to have query
limits configurable
- Custom instructions: additional guidance used to construct esql query,
e.g. select only certain fields to reduce the response size
- Add missing FTR tests

### Demo


https://github.com/user-attachments/assets/36755c7d-6cb7-4545-93cc-1150d23a3e21

---------

Co-authored-by: Elastic Machine <[email protected]>

Author: jedrazb

x-pack/platform/packages/shared/onechat/onechat-common/tools/types/index_search.ts

@@ -11,6 +11,8 @@ import { ToolType, type ToolDefinition, type ToolDefinitionWithSchema } from '..
 // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
 export type IndexSearchToolConfig = {
   pattern: string;
+  row_limit?: number;
+  custom_instructions?: string;
 };
 
 export type IndexSearchToolDefinition = ToolDefinition<

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql/graph.ts

@@ -34,6 +34,7 @@ const StateAnnotation = Annotation.Root({
   maxRetries: Annotation<number>(),
   additionalInstructions: Annotation<string | undefined>(),
   additionalContext: Annotation<string | undefined>(),
+  rowLimit: Annotation<number | undefined>(),
   // internal
   resource: Annotation<ResolvedResourceWithSampling>(),
   currentTry: Annotation<number>({ reducer: (a, b) => b, default: () => 0 }),
@@ -126,6 +127,7 @@ export const createNlToEsqlGraph = ({
         previousActions: state.actions,
         additionalInstructions: state.additionalInstructions,
         additionalContext: state.additionalContext,
+        rowLimit: state.rowLimit,
       })
     );
 

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql/nl_to_esql.ts

@@ -70,6 +70,10 @@ export interface GenerateEsqlOptions {
    * Defaults to `3`
    * */
   maxRetries?: number;
+  /**
+   * Maximum row limit to use in generated ES|QL queries.
+   */
+  rowLimit?: number;
 }
 
 export type GenerateEsqlParams = GenerateEsqlOptions & GenerateEsqlDeps;
@@ -81,13 +85,21 @@ export const generateEsql = async ({
   additionalInstructions,
   additionalContext,
   maxRetries = 3,
+  rowLimit,
   model,
   esClient,
   logger,
   events,
 }: GenerateEsqlParams): Promise<GenerateEsqlResponse> => {
   const docBase = await EsqlDocumentBase.load();
-  const graph = createNlToEsqlGraph({ model, esClient, logger, docBase, events });
+
+  const graph = createNlToEsqlGraph({
+    model,
+    esClient,
+    logger,
+    docBase,
+    events,
+  });
 
   return withActiveInferenceSpan(
     'GenerateEsqlGraph',
@@ -128,6 +140,7 @@ export const generateEsql = async ({
             maxRetries,
             additionalInstructions,
             additionalContext,
+            rowLimit,
           },
           {
             recursionLimit: 25,

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql/prompts.ts

@@ -11,6 +11,18 @@ import type { ResolvedResourceWithSampling } from '../utils/resources';
 import { formatResourceWithSampledValues } from '../utils/resources';
 import type { Action } from './actions';
 import { formatAction } from './actions';
+import { getEsqlInstructions } from './prompts/instructions_template';
+
+const getInstructionsWithRowLimit = (rowLimit?: number): string => {
+  if (!rowLimit) {
+    return getEsqlInstructions();
+  }
+
+  const defaultLimit = rowLimit;
+  const maxAllLimit = rowLimit;
+
+  return getEsqlInstructions({ defaultLimit, maxAllLimit });
+};
 
 export const createRequestDocumentationPrompt = ({
   nlQuery,
@@ -59,13 +71,15 @@ export const createGenerateEsqlPrompt = ({
   prompts,
   additionalInstructions,
   additionalContext,
+  rowLimit,
 }: {
   nlQuery: string;
   resource: ResolvedResourceWithSampling;
   prompts: EsqlPrompts;
   previousActions: Action[];
   additionalInstructions?: string;
   additionalContext?: string;
+  rowLimit?: number;
 }): BaseMessageLike[] => {
   return [
     [
@@ -85,11 +99,11 @@ ${prompts.syntax}
 
 ${prompts.examples}
 
-${prompts.instructions}
+${getInstructionsWithRowLimit(rowLimit)}
 
 ${
   additionalInstructions
-    ? `<additional_instructions>\n${additionalInstructions}\n</<additional_instructions>`
+    ? `<additional_instructions>\n${additionalInstructions}\n</additional_instructions>`
     : ''
 }
 

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql/prompts/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 { getEsqlInstructions, type InstructionsTemplateParams } from './instructions_template';

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql/prompts/instructions_template.ts

@@ -0,0 +1,93 @@
+/*
+ * 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 interface InstructionsTemplateParams {
+  /**
+   * The default LIMIT to use when no specific limit is requested by the user.
+   */
+  defaultLimit?: number;
+  /**
+   * The maximum LIMIT to use when the user asks for "all" results.
+   */
+  maxAllLimit?: number;
+}
+
+const DEFAULT_LIMIT = 100;
+const MAX_ALL_LIMIT = 250;
+
+/**
+ * Generates ES|QL query generation instructions with configurable limit values.
+ * This is a copy of the instructions from the inference plugin, modified to support
+ * custom row limits for Agent Builder's index search tool.
+ */
+export const getEsqlInstructions = (params: InstructionsTemplateParams = {}): string => {
+  const { defaultLimit = DEFAULT_LIMIT, maxAllLimit = MAX_ALL_LIMIT } = params;
+
+  return `<instructions>
+
+    ## Follow the syntax
+
+    It is CRUCIAL and MANDATORY to only use commands and functions which are present in the syntax definition,
+    and to follow the syntax as described in the documentation and its examples. Do not try to guess
+    new functions or commands based on other query languages. Assume that ONLY the set of capabilities described
+    in the provided ES|QL documentation is valid, and do not try to guess parameters or syntax based
+    on other query languages.
+
+    ## Respect the mappings or field definitions
+
+    If the user, or a tool, provides in the discussion the mappings or a list of fields present in the index, you should **ONLY** use
+    the provided fields to create your query. Do not assume other fields may exist. Only use the set of fields
+    which were provided by the user.
+
+    ## Use a safety LIMIT
+
+    1. **LIMIT is Mandatory:** All multi-row queries **must** end with a \`LIMIT\`. The only exception is for single-row aggregations (e.g., \`STATS\` without a \`GROUP BY\`).
+
+    2. **Applying Limits:**
+        * **User-Specified:** If the user provides a number ("top 10", "get 50"), use it for the \`LIMIT\`.
+        * **Default:** If no number is given, default to \`LIMIT ${defaultLimit}\` for both raw events and \`GROUP BY\` results. Notify the user when you apply this default (e.g., "I've added a \`LIMIT ${defaultLimit}\` for safety.").
+
+    3. **Handling "All Data" Requests:** If a user asks for "all" results, apply a safety \`LIMIT ${maxAllLimit}\` and state that this limit was added to protect the system.
+
+    ## Don't use tech preview features unless specified otherwise
+
+    Using tech preview commands, functions or other features should be avoided unless specifically asked by the user.
+
+    ## Use MATCH for full text search
+
+    Unless specified otherwise, full text searches should always be done using MATCH in favor of other search functions.
+
+    ## ES|QL query formatting
+
+    - All generated ES|QL queries must be wrapped with \`\`\`esql and \`\`\`
+    - Queries must be properly formatted, with a carriage return after each function
+
+    Example:
+    \`\`\`
+    FROM logs-*
+    | WHERE @timestamp <= NOW() - 24 hours
+    | STATS count = COUNT(*) BY log.level
+    | SORT count DESC
+    \`\`\`
+
+    ## Do not invent things to please the user
+
+    If what the user is asking for is not technically achievable with ES|QL's capabilities, just inform
+    the user. DO NOT invent capabilities not described in the documentation just to provide
+    a positive answer to the user.
+
+    When converting queries from one language to ES|QL, make sure that the functions are available
+    and documented in ES|QL. E.g., for SPL's LEN, use LENGTH. For IF, use CASE.
+
+    ## Tool Usage Restrictions
+
+    **CRITICAL**: Only use the tools that are explicitly defined in your available tool set. Do not call
+    tools from other contexts or systems.
+
+</instructions>
+`;
+};

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/nl_search.ts

@@ -33,13 +33,17 @@ export const naturalLanguageSearch = async ({
   esClient,
   logger,
   events,
+  rowLimit,
+  customInstructions,
 }: {
   nlQuery: string;
   target: string;
   model: ScopedModel;
   esClient: ElasticsearchClient;
   logger: Logger;
   events: ToolEventEmitter;
+  rowLimit?: number;
+  customInstructions?: string;
 }): Promise<NaturalLanguageSearchResponse> => {
   const queryGenResponse = await generateEsql({
     nlQuery,
@@ -49,6 +53,8 @@ export const naturalLanguageSearch = async ({
     esClient,
     logger,
     events,
+    rowLimit,
+    additionalInstructions: customInstructions,
   });
 
   return {

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/search/graph.ts

@@ -24,6 +24,8 @@ const StateAnnotation = Annotation.Root({
   // inputs
   nlQuery: Annotation<string>(),
   targetPattern: Annotation<string | undefined>(),
+  rowLimit: Annotation<number | undefined>(),
+  customInstructions: Annotation<string | undefined>(),
   // inner
   indexIsValid: Annotation<boolean>(),
   searchTarget: Annotation<SearchTarget>(),
@@ -52,12 +54,7 @@ export const createSearchToolGraph = ({
   logger: Logger;
   events: ToolEventEmitter;
 }) => {
-  const tools = [
-    createRelevanceSearchTool({ model, esClient, events }),
-    createNaturalLanguageSearchTool({ model, esClient, events, logger }),
-  ];
-
-  const toolNode = new ToolNode<typeof StateAnnotation.State.messages>(tools);
+  const relevanceTool = createRelevanceSearchTool({ model, esClient, events });
 
   const selectAndValidateIndex = async (state: StateType) => {
     events?.reportProgress(progressMessages.selectingTarget());
@@ -89,18 +86,33 @@ export const createSearchToolGraph = ({
     return state.indexIsValid ? 'agent' : '__end__';
   };
 
-  const searchModel = model.chatModel.bindTools(tools).withConfig({
-    tags: ['onechat-search-tool'],
-  });
-
   const callSearchAgent = async (state: StateType) => {
     events?.reportProgress(
       progressMessages.resolvingSearchStrategy({
         target: state.searchTarget.name,
       })
     );
+
+    const nlSearchTool = createNaturalLanguageSearchTool({
+      model,
+      esClient,
+      events,
+      logger,
+      rowLimit: state.rowLimit,
+      customInstructions: state.customInstructions,
+    });
+
+    const tools = [relevanceTool, nlSearchTool];
+    const searchModel = model.chatModel.bindTools(tools).withConfig({
+      tags: ['onechat-search-tool'],
+    });
+
     const response = await searchModel.invoke(
-      getSearchPrompt({ nlQuery: state.nlQuery, searchTarget: state.searchTarget })
+      getSearchPrompt({
+        nlQuery: state.nlQuery,
+        searchTarget: state.searchTarget,
+        customInstructions: state.customInstructions,
+      })
     );
     return {
       messages: [response],
@@ -113,6 +125,18 @@ export const createSearchToolGraph = ({
   };
 
   const executeTool = async (state: StateType) => {
+    const nlSearchTool = createNaturalLanguageSearchTool({
+      model,
+      esClient,
+      events,
+      logger,
+      rowLimit: state.rowLimit,
+      customInstructions: state.customInstructions,
+    });
+
+    const tools = [relevanceTool, nlSearchTool];
+    const toolNode = new ToolNode<typeof StateAnnotation.State.messages>(tools);
+
     const toolNodeResult = await toolNode.invoke(state.messages);
     const toolResults = extractToolResults(toolNodeResult[toolNodeResult.length - 1]);
 

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/search/inner_tools.ts

@@ -95,11 +95,15 @@ export const createNaturalLanguageSearchTool = ({
   esClient,
   events,
   logger,
+  rowLimit,
+  customInstructions,
 }: {
   model: ScopedModel;
   esClient: ElasticsearchClient;
   events: ToolEventEmitter;
   logger: Logger;
+  rowLimit?: number;
+  customInstructions?: string;
 }) => {
   return toTool(
     async ({ query, index }) => {
@@ -115,6 +119,8 @@ export const createNaturalLanguageSearchTool = ({
             esClient,
             events,
             logger,
+            rowLimit,
+            customInstructions,
           });
 
           const results: ToolResult[] = response.esqlData

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/search/prompts.ts

@@ -16,9 +16,11 @@ import {
 export const getSearchPrompt = ({
   nlQuery,
   searchTarget,
+  customInstructions,
 }: {
   nlQuery: string;
   searchTarget: SearchTarget;
+  customInstructions?: string;
 }): BaseMessageLike[] => {
   const systemPrompt = `You are an expert search dispatcher. Your sole task is to analyze a user's request and call the single most appropriate tool to answer it.
 You **must** call **one** of the available tools. Do not answer the user directly or ask clarifying questions.
@@ -42,7 +44,13 @@ You **must** call **one** of the available tools. Do not answer the user directl
 
 ## Additional instructions
 
-- The search will be performed against the \`${searchTarget.name}\` ${searchTarget.type}, so you should use that value for the \`index\` parameters of the tool you will call.`;
+- The search will be performed against the \`${searchTarget.name}\` ${
+    searchTarget.type
+  }, so you should use that value for the \`index\` parameters of the tool you will call.${
+    customInstructions ? `\n- User provided additional instructions: ${customInstructions}` : ''
+  }
+
+`;
 
   const userPrompt = `Execute the following user query: "${nlQuery}"`;