From dd5f08db894a5d144ffd59e5d529b7c8fec45360 Mon Sep 17 00:00:00 2001 From: "Quynh Nguyen (Quinn)" <43350163+qn895@users.noreply.github.com> Date: Tue, 26 Aug 2025 14:37:18 -0500 Subject: [PATCH 1/2] [NL2ESQL] Add max function calls allowed to prevent tool being called recursively (#231719) Follow up of https://github.com/elastic/kibana/pull/231217. This PR fixes https://github.com/elastic/kibana/issues/229979 and implements a mechanism to stop the nl2esql tool being called indefinitely, and to stop after 5 tries maximum. image https://35-187-109-62.sslip.io/projects/UHJvamVjdDoxMTAy/traces/e035f0e5f983d2e61c1153f6f41abd80?selected=&selectedSpanNodeId=U3Bhbjo2NjczNjE%3D Check the PR satisfies following conditions. Reviewers should verify this PR satisfies this list as well. - [ ] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md) - [ ] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [ ] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [ ] If a plugin configuration key changed, check if it needs to be allowlisted in the cloud and added to the [docker list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker) - [ ] This was checked for breaking HTTP API changes, and any breaking changes have been approved by the breaking-change committee. The `release_note:breaking` label should be applied in these situations. - [ ] [Flaky Test Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was used on any tests changed - [ ] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) - [ ] Review the [backport guidelines](https://docs.google.com/document/d/1VyN5k91e5OVumlc0Gb9RPa3h1ewuPE705nRtioPiTvY/edit?usp=sharing) and apply applicable `backport:*` labels. Does this PR introduce any risks? For example, consider risks like hard to test bugs, performance regression, potential of data loss. Describe the risk, its severity, and mitigation for each identified risk. Invite stakeholders and evaluate how to proceed before merging. - [ ] [See some risk examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx) - [ ] ... (cherry picked from commit 58869ae814518b61c9d2203d735077466efde722) --- .../tasks/nl_to_esql/actions/generate_esql.ts | 71 +++++++++++++------ 1 file changed, 49 insertions(+), 22 deletions(-) diff --git a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index 54fa27b2bdafa..3e5acc3de9863 100644 --- a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -5,7 +5,8 @@ * 2.0. */ -import { Observable, map, merge, of, switchMap } from 'rxjs'; +import type { Observable } from 'rxjs'; +import { map, merge, of, switchMap } from 'rxjs'; import type { Logger } from '@kbn/logging'; import { ToolCall, @@ -18,14 +19,17 @@ import { OutputEventType, FunctionCallingMode, ChatCompleteMetadata, + ToolChoiceType, } from '@kbn/inference-common'; import { correctCommonEsqlMistakes, generateFakeToolCallId } from '../../../../common'; import { InferenceClient } from '../../..'; import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; -import { EsqlDocumentBase } from '../doc_base'; +import type { EsqlDocumentBase } from '../doc_base'; import { requestDocumentationSchema } from './shared'; import type { NlToEsqlTaskEvent } from '../types'; +const MAX_CALLS = 5; + export const generateEsqlTask = ({ chatCompleteApi, connectorId, @@ -34,9 +38,12 @@ export const generateEsqlTask = ({ toolOptions: { tools, toolChoice }, docBase, functionCalling, + maxRetries, + retryConfiguration, logger, system, metadata, + maxCallsAllowed = MAX_CALLS, }: { connectorId: string; systemMessage: string; @@ -48,12 +55,16 @@ export const generateEsqlTask = ({ logger: Pick; metadata?: ChatCompleteMetadata; system?: string; + maxCallsAllowed?: number; }) => { return function askLlmToRespond({ documentationRequest: { commands, functions }, + callCount = 0, }: { documentationRequest: { commands?: string[]; functions?: string[] }; + callCount?: number; }): Observable> { + const functionLimitReached = callCount >= maxCallsAllowed; const keywords = [...(commands ?? []), ...(functions ?? [])]; const requestedDocumentation = docBase.getDocumentation(keywords); const fakeRequestDocsToolCall = createFakeTooCall(commands, functions); @@ -120,14 +131,16 @@ export const generateEsqlTask = ({ toolCallId: fakeRequestDocsToolCall.toolCallId, }, ], - toolChoice, - tools: { - ...tools, - request_documentation: { - description: 'Request additional ES|QL documentation if needed', - schema: requestDocumentationSchema, - }, - }, + toolChoice: !functionLimitReached ? toolChoice : ToolChoiceType.none, + tools: functionLimitReached + ? {} + : { + ...tools, + request_documentation: { + description: 'Request additional ES|QL documentation if needed', + schema: requestDocumentationSchema, + }, + }, }).pipe( withoutTokenCountEvents(), map((generateEvent) => { @@ -144,18 +157,32 @@ export const generateEsqlTask = ({ }), switchMap((generateEvent) => { if (isChatCompletionMessageEvent(generateEvent)) { - const onlyToolCall = - generateEvent.toolCalls.length === 1 ? generateEvent.toolCalls[0] : undefined; - - if (onlyToolCall?.function.name === 'request_documentation') { - const args = onlyToolCall.function.arguments; - - return askLlmToRespond({ - documentationRequest: { - commands: args.commands, - functions: args.functions, - }, - }); + const toolCalls = generateEvent.toolCalls as ToolCall[]; + const onlyToolCall = toolCalls.length === 1 ? toolCalls[0] : undefined; + + if (onlyToolCall && onlyToolCall.function.name === 'request_documentation') { + if (functionLimitReached) { + return of({ + ...generateEvent, + content: `You have reached the maximum number of documentation requests. Do not try to request documentation again for commands ${commands?.join( + ', ' + )} and functions ${functions?.join( + ', ' + )}. Try to answer the user's question using currently available information.`, + }); + } + + const args = + 'arguments' in onlyToolCall.function ? onlyToolCall.function.arguments : undefined; + if (args && (args.commands?.length || args.functions?.length)) { + return askLlmToRespond({ + documentationRequest: { + commands: args.commands ?? [], + functions: args.functions ?? [], + }, + callCount: callCount + 1, + }); + } } } From c9e39c62672e423a5e84e0bcf5bf4388cf5d8757 Mon Sep 17 00:00:00 2001 From: "Quynh Nguyen (Quinn)" <43350163+qn895@users.noreply.github.com> Date: Wed, 27 Aug 2025 12:34:19 -0500 Subject: [PATCH 2/2] Fix arg --- .../inference/server/tasks/nl_to_esql/actions/generate_esql.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index 3e5acc3de9863..2610c099a9814 100644 --- a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -38,8 +38,6 @@ export const generateEsqlTask = ({ toolOptions: { tools, toolChoice }, docBase, functionCalling, - maxRetries, - retryConfiguration, logger, system, metadata,