From f51584cc3a159d1522525e1ac5fecb852dd709c7 Mon Sep 17 00:00:00 2001
From: Dario Gieselaar <dario.gieselaar@elastic.co>
Date: Fri, 26 Jan 2024 12:49:45 +0100
Subject: [PATCH] [8.12] [Obs AI Assistant] Increase no of max function calls
 to 5 (#175588) (#175668)

# Backport

This will backport the following commits from `main` to `8.12`:
- [[Obs AI Assistant] Increase no of max function calls to 5
(#175588)](https://github.com/elastic/kibana/pull/175588)

<!--- Backport version: 8.9.8 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Dario
Gieselaar","email":"dario.gieselaar@elastic.co"},"sourceCommit":{"committedDate":"2024-01-26T09:17:36Z","message":"[Obs
AI Assistant] Increase no of max function calls to 5 (#175588)\n\nThe
function limit of 3 is too low, it runs into issues w/ ES|QL
query\r\ngeneration if the `get_dataset_info` and `execute_query`
functions
are\r\nused.","sha":"beff74e19e94c7d3c609a2be0cb1d4b2f28da043","branchLabelMapping":{"^v8.13.0$":"main","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:fix","v8.12.1","v8.13.0"],"number":175588,"url":"https://github.com/elastic/kibana/pull/175588","mergeCommit":{"message":"[Obs
AI Assistant] Increase no of max function calls to 5 (#175588)\n\nThe
function limit of 3 is too low, it runs into issues w/ ES|QL
query\r\ngeneration if the `get_dataset_info` and `execute_query`
functions
are\r\nused.","sha":"beff74e19e94c7d3c609a2be0cb1d4b2f28da043"}},"sourceBranch":"main","suggestedTargetBranches":["8.12"],"targetPullRequestStates":[{"branch":"8.12","label":"v8.12.1","labelRegex":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"main","label":"v8.13.0","labelRegex":"^v8.13.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/175588","number":175588,"mergeCommit":{"message":"[Obs
AI Assistant] Increase no of max function calls to 5 (#175588)\n\nThe
function limit of 3 is too low, it runs into issues w/ ES|QL
query\r\ngeneration if the `get_dataset_info` and `execute_query`
functions
are\r\nused.","sha":"beff74e19e94c7d3c609a2be0cb1d4b2f28da043"}}]}]
BACKPORT-->
---
 .../server/service/client/index.test.ts                     | 6 +++++-
 .../server/service/client/index.ts                          | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.test.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.test.ts
index 3e7660c427be7..41beecee82245 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.test.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.test.ts
@@ -1250,11 +1250,15 @@ describe('Observability AI Assistant client', () => {
 
       await requestAlertsFunctionCall();
 
+      await requestAlertsFunctionCall();
+
+      await requestAlertsFunctionCall();
+
       await finished(stream);
     });
 
     it('executed the function no more than three times', () => {
-      expect(functionClientMock.executeFunction).toHaveBeenCalledTimes(3);
+      expect(functionClientMock.executeFunction).toHaveBeenCalledTimes(5);
     });
 
     it('does not give the LLM the choice to call a function anymore', () => {
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index 0e3927f82370d..c100cd05cd7ae 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -175,7 +175,7 @@ export class ObservabilityAIAssistantClient {
 
     let numFunctionsCalled: number = 0;
 
-    const MAX_FUNCTION_CALLS = 3;
+    const MAX_FUNCTION_CALLS = 5;
     const MAX_FUNCTION_RESPONSE_TOKEN_COUNT = 4000;
 
     const next = async (nextMessages: Message[]): Promise<void> => {