[8.15] [Security Assistant] Fix langgraph issues (elastic#189287) (el…

…astic#189678)

# Backport

This will backport the following commits from `main` to `8.15`:
- [[Security Assistant] Fix langgraph issues
(elastic#189287)](elastic#189287)

<!--- Backport version: 9.4.3 -->

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

<!--BACKPORT [{"author":{"name":"Patryk
Kopyciński","email":"contact@patrykkopycinski.com"},"sourceCommit":{"committedDate":"2024-07-31T20:32:29Z","message":"[Security
Assistant] Fix langgraph issues (elastic#189287)\n\n## Summary\r\n\r\n- Fixed
`Regenerate` button for Gemini\r\n- Fixed auto title functionality\r\n-
Enabled `assistantBedrockChat` by
default\r\n\r\n---------\r\n\r\nCo-authored-by: kibanamachine
<42973632+kibanamachine@users.noreply.github.com>","sha":"0be8295ed0028870074772c8d82093376482c4eb","branchLabelMapping":{"^v8.16.0$":"main","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["bug","release_note:skip","ci:cloud-deploy","v8.15.0","v8.16.0"],"title":"[Security
Assistant] Fix langgraph
issues","number":189287,"url":"https://github.com/elastic/kibana/pull/189287","mergeCommit":{"message":"[Security
Assistant] Fix langgraph issues (elastic#189287)\n\n## Summary\r\n\r\n- Fixed
`Regenerate` button for Gemini\r\n- Fixed auto title functionality\r\n-
Enabled `assistantBedrockChat` by
default\r\n\r\n---------\r\n\r\nCo-authored-by: kibanamachine
<42973632+kibanamachine@users.noreply.github.com>","sha":"0be8295ed0028870074772c8d82093376482c4eb"}},"sourceBranch":"main","suggestedTargetBranches":["8.15"],"targetPullRequestStates":[{"branch":"8.15","label":"v8.15.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"main","label":"v8.16.0","branchLabelMappingKey":"^v8.16.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/189287","number":189287,"mergeCommit":{"message":"[Security
Assistant] Fix langgraph issues (elastic#189287)\n\n## Summary\r\n\r\n- Fixed
`Regenerate` button for Gemini\r\n- Fixed auto title functionality\r\n-
Enabled `assistantBedrockChat` by
default\r\n\r\n---------\r\n\r\nCo-authored-by: kibanamachine
<42973632+kibanamachine@users.noreply.github.com>","sha":"0be8295ed0028870074772c8d82093376482c4eb"}}]}]
BACKPORT-->

Co-authored-by: Patryk Kopyciński <contact@patrykkopycinski.com>

package.json

@@ -80,15 +80,15 @@
   "resolutions": {
     "**/@bazel/typescript/protobufjs": "6.11.4",
     "**/@hello-pangea/dnd": "16.6.0",
-    "**/@langchain/core": "^0.2.17",
+    "**/@langchain/core": "^0.2.18",
     "**/@types/node": "20.10.5",
     "**/@typescript-eslint/utils": "5.62.0",
     "**/chokidar": "^3.5.3",
     "**/d3-scale/**/d3-color": "npm:@elastic/kibana-d3-color@2.0.1",
     "**/globule/minimatch": "^3.1.2",
     "**/hoist-non-react-statics": "^3.3.2",
     "**/isomorphic-fetch/node-fetch": "^2.6.7",
-    "**/langchain": "^0.2.10",
+    "**/langchain": "^0.2.11",
     "**/react-intl/**/@types/react": "^17.0.45",
     "**/remark-parse/trim": "1.0.1",
     "**/sharp": "0.32.6",
@@ -938,9 +938,9 @@
     "@kbn/xstate-utils": "link:packages/kbn-xstate-utils",
     "@kbn/zod-helpers": "link:packages/kbn-zod-helpers",
     "@langchain/community": "0.2.18",
-    "@langchain/core": "^0.2.17",
+    "@langchain/core": "^0.2.18",
     "@langchain/google-genai": "^0.0.23",
-    "@langchain/langgraph": "^0.0.29",
+    "@langchain/langgraph": "^0.0.31",
     "@langchain/openai": "^0.1.3",
     "@langtrase/trace-attributes": "^3.0.8",
     "@launchdarkly/node-server-sdk": "^9.4.6",
@@ -1080,8 +1080,8 @@
     "jsonwebtoken": "^9.0.2",
     "jsts": "^1.6.2",
     "kea": "^2.6.0",
-    "langchain": "^0.2.10",
-    "langsmith": "^0.1.37",
+    "langchain": "^0.2.11",
+    "langsmith": "^0.1.39",
     "launchdarkly-js-client-sdk": "^3.4.0",
     "launchdarkly-node-server-sdk": "^7.0.3",
     "load-json-file": "^6.2.0",

x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts

@@ -21,5 +21,5 @@ export type AssistantFeatureKey = keyof AssistantFeatures;
 export const defaultAssistantFeatures = Object.freeze({
   assistantKnowledgeBaseByDefault: false,
   assistantModelEvaluation: false,
-  assistantBedrockChat: false,
+  assistantBedrockChat: true,
 });

x-pack/packages/kbn-langchain/server/language_models/bedrock_chat.ts

@@ -69,7 +69,7 @@ export class ActionsClientBedrockChatModel extends _BedrockChat {
           params: {
             subAction: 'invokeAIRaw',
             subActionParams: {
-              messages: inputBody.messages,
+              messages: prepareMessages(inputBody.messages),
               temperature: params.temperature ?? inputBody.temperature,
               stopSequences: inputBody.stop_sequences,
               system: inputBody.system,
@@ -99,3 +99,20 @@ export class ActionsClientBedrockChatModel extends _BedrockChat {
     });
   }
 }
+
+const prepareMessages = (messages: Array<{ role: string; content: string[] }>) =>
+  messages.reduce((acc, { role, content }) => {
+    const lastMessage = acc[acc.length - 1];
+
+    if (!lastMessage || lastMessage.role !== role) {
+      acc.push({ role, content });
+      return acc;
+    }
+
+    if (lastMessage.role === role) {
+      acc[acc.length - 1].content = lastMessage.content.concat(content);
+      return acc;
+    }
+
+    return acc;
+  }, [] as Array<{ role: string; content: string[] }>);

x-pack/packages/kbn-langchain/server/language_models/gemini_chat.ts

@@ -69,7 +69,7 @@ export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI {
   }
 
   async completionWithRetry(
-    request: string | GenerateContentRequest | Array<string | Part>,
+    request: GenerateContentRequest,
     options?: this['ParsedCallOptions']
   ): Promise<GenerateContentResult> {
     return this.caller.callWithOptions({ signal: options?.signal }, async () => {
@@ -80,7 +80,8 @@ export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI {
             subAction: 'invokeAIRaw',
             subActionParams: {
               model: this.#model,
-              messages: request,
+              messages: request.contents,
+              tools: request.tools,
               temperature: this.#temperature,
             },
           },

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts

@@ -90,7 +90,7 @@ export const getDefaultAssistantGraph = ({
         default: () => undefined,
       },
       messages: {
-        value: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y),
+        value: (x: BaseMessage[], y: BaseMessage[]) => y ?? x,
         default: () => [],
       },
       chatTitle: {

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts

@@ -90,6 +90,7 @@ export const streamGraph = async ({
         runName: DEFAULT_ASSISTANT_GRAPH_ID,
         tags: traceOptions?.tags ?? [],
         version: 'v2',
+        streamMode: 'values',
       },
       llmType === 'bedrock' ? { includeNames: ['Summarizer'] } : undefined
     );

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts

@@ -127,7 +127,7 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
         streamRunnable: isStream,
       })
     : llmType && ['bedrock', 'gemini'].includes(llmType) && bedrockChatEnabled
-    ? createToolCallingAgent({
+    ? await createToolCallingAgent({
         llm,
         tools,
         prompt:

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/execute_tools.ts

@@ -8,7 +8,8 @@
 import { RunnableConfig } from '@langchain/core/runnables';
 import { StructuredTool } from '@langchain/core/tools';
 import { ToolExecutor } from '@langchain/langgraph/prebuilt';
-import { isArray } from 'lodash';
+import { castArray } from 'lodash';
+import { AgentAction } from 'langchain/agents';
 import { AgentState, NodeParamsBase } from '../types';
 
 export interface ExecuteToolsParams extends NodeParamsBase {
@@ -34,23 +35,30 @@ export const executeTools = async ({ config, logger, state, tools }: ExecuteTool
   logger.debug(`Node state:\n${JSON.stringify(state, null, 2)}`);
 
   const toolExecutor = new ToolExecutor({ tools });
-  const agentAction = isArray(state.agentOutcome) ? state.agentOutcome[0] : state.agentOutcome;
+  const agentAction = state.agentOutcome;
 
   if (!agentAction || 'returnValues' in agentAction) {
     throw new Error('Agent has not been run yet');
   }
 
-  let out;
-  try {
-    out = await toolExecutor.invoke(agentAction, config);
-  } catch (err) {
-    return {
-      steps: [{ action: agentAction, observation: JSON.stringify(`Error: ${err}`, null, 2) }],
-    };
-  }
+  const steps = await Promise.all(
+    castArray(state.agentOutcome as AgentAction)?.map(async (action) => {
+      let out;
+      try {
+        out = await toolExecutor.invoke(action, config);
+      } catch (err) {
+        return {
+          action,
+          observation: JSON.stringify(`Error: ${err}`, null, 2),
+        };
+      }
+
+      return {
+        action,
+        observation: JSON.stringify(out, null, 2),
+      };
+    })
+  );
 
-  return {
-    ...state,
-    steps: [{ action: agentAction, observation: JSON.stringify(out, null, 2) }],
-  };
+  return { steps };
 };

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts

@@ -46,7 +46,6 @@ export const generateChatTitle = async ({
   logger.debug(`chatTitle: ${chatTitle}`);
 
   return {
-    ...state,
     chatTitle,
   };
 };

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/get_persisted_conversation.ts

@@ -27,7 +27,6 @@ export const getPersistedConversation = async ({
   if (!conversationId) {
     logger.debug('Cannot get conversation, because conversationId is undefined');
     return {
-      ...state,
       conversation: undefined,
       messages: [],
       chatTitle: '',
@@ -39,7 +38,6 @@ export const getPersistedConversation = async ({
   if (!conversation) {
     logger.debug('Requested conversation, because conversation is undefined');
     return {
-      ...state,
       conversation: undefined,
       messages: [],
       chatTitle: '',
@@ -50,11 +48,21 @@ export const getPersistedConversation = async ({
   logger.debug(`conversationId: ${conversationId}`);
 
   const messages = getLangChainMessages(conversation.messages ?? []);
+
+  if (!state.input) {
+    const lastMessage = messages?.splice(-1)[0];
+    return {
+      conversation,
+      messages,
+      chatTitle: conversation.title,
+      input: lastMessage?.content as string,
+    };
+  }
+
   return {
-    ...state,
     conversation,
     messages,
     chatTitle: conversation.title,
-    input: !state.input ? conversation.messages?.slice(-1)[0].content : state.input,
+    input: state.input,
   };
 };

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/persist_conversation_changes.ts

@@ -57,7 +57,6 @@ export const persistConversationChanges = async ({
     const langChainMessages = getLangChainMessages(state.conversation.messages ?? []);
     const messages = langChainMessages.slice(0, -1); // all but the last message
     return {
-      ...state,
       conversation: state.conversation,
       messages,
     };
@@ -78,15 +77,14 @@ export const persistConversationChanges = async ({
   });
   if (!updatedConversation) {
     logger.debug('Not updated conversation');
-    return { ...state, conversation: undefined, messages: [] };
+    return { conversation: undefined, messages: [] };
   }
 
   logger.debug(`conversationId: ${conversationId}`);
   const langChainMessages = getLangChainMessages(updatedConversation.messages ?? []);
   const messages = langChainMessages.slice(0, -1); // all but the last message
 
   return {
-    ...state,
     conversation: updatedConversation,
     messages,
   };

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts

@@ -42,13 +42,12 @@ export const runAgent = async ({
   const agentOutcome = await agentRunnable.withConfig({ tags: [AGENT_NODE_TAG] }).invoke(
     {
       ...state,
-      messages: state.messages.splice(-1),
       chat_history: state.messages, // TODO: Message de-dupe with ...state spread
     },
     config
   );
+
   return {
-    ...state,
     agentOutcome,
   };
 };

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/should_continue.ts

@@ -31,7 +31,7 @@ export const shouldContinue = ({ logger, state }: ShouldContinueParams) => {
 export const shouldContinueGenerateTitle = ({ logger, state }: ShouldContinueParams) => {
   logger.debug(`Node state:\n${JSON.stringify(state, null, 2)}`);
 
-  if (state.conversation?.title !== NEW_CHAT) {
+  if (state.conversation?.title?.length && state.conversation?.title !== NEW_CHAT) {
     return 'end';
   }
 

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts

@@ -17,7 +17,7 @@ export const openAIFunctionAgentPrompt = ChatPromptTemplate.fromMessages([
 export const bedrockToolCallingAgentPrompt = ChatPromptTemplate.fromMessages([
   [
     'system',
-    'You are a helpful assistant. ALWAYS use the provided tools. Use tools as often as possible, as they have access to the latest data and syntax.',
+    'You are a helpful assistant. ALWAYS use the provided tools. Use tools as often as possible, as they have access to the latest data and syntax. Always return value from ESQLKnowledgeBaseTool as is. Never return <thinking> tags in the response, but make sure to include <result> tags content in the response. Do not reflect on the quality of the returned search results in your response.',
   ],
   ['placeholder', '{chat_history}'],
   ['human', '{input}'],

x-pack/plugins/security_solution/common/experimental_features.ts

@@ -130,7 +130,7 @@ export const allowedExperimentalValues = Object.freeze({
   /**
    * Enables the Assistant BedrockChat Langchain model, introduced in `8.15.0`.
    */
-  assistantBedrockChat: false,
+  assistantBedrockChat: true,
 
   /**
    * Enables the Managed User section inside the new user details flyout.

x-pack/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts

@@ -22,7 +22,7 @@ export interface OpenAndAcknowledgedAlertsToolParams extends AssistantToolParams
 }
 
 export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION =
-  'Call this for knowledge about the latest n open and acknowledged alerts (sorted by `kibana.alert.risk_score`) in the environment, or when answering questions about open alerts. Do not call this tool for alert count or quantity. Input should be an empty object. The output is an array of the latest n open and acknowledged alerts.';
+  'Call this for knowledge about the latest n open and acknowledged alerts (sorted by `kibana.alert.risk_score`) in the environment, or when answering questions about open alerts. Do not call this tool for alert count or quantity. The output is an array of the latest n open and acknowledged alerts.';
 
 /**
  * Returns a tool for querying open and acknowledged alerts, or null if the

x-pack/plugins/stack_connectors/common/gemini/schema.ts

@@ -72,6 +72,7 @@ export const InvokeAIRawActionParamsSchema = schema.object({
   stopSequences: schema.maybe(schema.arrayOf(schema.string())),
   signal: schema.maybe(schema.any()),
   timeout: schema.maybe(schema.number()),
+  tools: schema.maybe(schema.arrayOf(schema.any())),
 });
 
 export const InvokeAIActionResponseSchema = schema.object({

x-pack/plugins/stack_connectors/server/connector_types/gemini/gemini.ts

@@ -299,9 +299,10 @@ export class GeminiConnector extends SubActionConnector<Config, Secrets> {
     temperature = 0,
     signal,
     timeout,
+    tools,
   }: InvokeAIRawActionParams): Promise<InvokeAIRawActionResponse> {
     const res = await this.runApi({
-      body: JSON.stringify(formatGeminiPayload(messages, temperature)),
+      body: JSON.stringify({ ...formatGeminiPayload(messages, temperature), tools }),
       model,
       signal,
       timeout,