[Security solution] AI Assistant, replace LLM with SimpleChatModel + Bedrock streaming

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

@@ -7,3 +7,4 @@
 
 export { ActionsClientChatOpenAI } from './chat_openai';
 export { ActionsClientLlm } from './llm';
+export { ActionsClientSimpleChatModel } from './simple_chat_model';

x-pack/packages/kbn-elastic-assistant-common/impl/language_models/simple_chat_model.ts

@@ -0,0 +1,136 @@
+/*
+ * 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.
+ */
+
+import {
+  SimpleChatModel,
+  type BaseChatModelParams,
+} from '@langchain/core/language_models/chat_models';
+import { type BaseMessage } from '@langchain/core/messages';
+import type { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server';
+import { Logger } from '@kbn/logging';
+import { KibanaRequest } from '@kbn/core-http-server';
+import { v4 as uuidv4 } from 'uuid';
+import { get } from 'lodash/fp';
+import { getDefaultArguments } from './constants';
+import { ExecuteConnectorRequestBody } from '../..';
+
+export const getMessageContentAndRole = (prompt: string, role = 'user') => ({
+  content: prompt,
+  role: role === 'human' ? 'user' : role,
+});
+
+export interface CustomChatModelInput extends BaseChatModelParams {
+  actions: ActionsPluginStart;
+  connectorId: string;
+  logger: Logger;
+  llmType?: string;
+  model?: string;
+  temperature?: number;
+  request: KibanaRequest<unknown, unknown, ExecuteConnectorRequestBody>;
+}
+
+export class ActionsClientSimpleChatModel extends SimpleChatModel {
+  #actions: ActionsPluginStart;
+  #connectorId: string;
+  #logger: Logger;
+  #request: KibanaRequest<unknown, unknown, ExecuteConnectorRequestBody>;
+  #traceId: string;
+  llmType: string;
+  model?: string;
+  temperature?: number;
+
+  constructor({
+    actions,
+    connectorId,
+    llmType,
+    logger,
+    model,
+    request,
+    temperature,
+  }: CustomChatModelInput) {
+    super({});
+
+    this.#actions = actions;
+    this.#connectorId = connectorId;
+    this.#traceId = uuidv4();
+    this.#logger = logger;
+    this.#request = request;
+    this.llmType = llmType ?? 'openai';
+    this.model = model;
+    this.temperature = temperature;
+  }
+
+  _llmType() {
+    return this.llmType;
+  }
+
+  // Model type needs to be `base_chat_model` to work with LangChain OpenAI Tools
+  // We may want to make this configurable (ala _llmType) if different agents end up requiring different model types
+  // See: https://github.com/langchain-ai/langchainjs/blob/fb699647a310c620140842776f4a7432c53e02fa/langchain/src/agents/openai/index.ts#L185
+  _modelType() {
+    return 'base_chat_model';
+  }
+
+  async _call(messages: BaseMessage[], options: this['ParsedCallOptions']): Promise<string> {
+    if (!messages.length) {
+      throw new Error('No messages provided.');
+    }
+    const formattedMessages = [];
+    if (messages.length === 2) {
+      messages.forEach((message, i) => {
+        if (typeof message.content !== 'string') {
+          throw new Error('Multimodal messages are not supported.');
+        }
+        formattedMessages.push(getMessageContentAndRole(message.content, message._getType()));
+      });
+    } else {
+      if (typeof messages[0].content !== 'string') {
+        throw new Error('Multimodal messages are not supported.');
+      }
+      formattedMessages.push(getMessageContentAndRole(messages[0].content));
+    }
+    this.#logger.debug(
+      `ActionsClientSimpleChatModel#_call\ntraceId: ${
+        this.#traceId
+      }\nassistantMessage:\n${JSON.stringify(formattedMessages)} `
+    );
+    // create a new connector request body with the assistant message:
+    const requestBody = {
+      actionId: this.#connectorId,
+      params: {
+        // hard code to non-streaming subaction as this class only supports non-streaming
+        subAction: 'invokeAI',
+        subActionParams: {
+          model: this.#request.body.model,
+          messages: formattedMessages,
+          ...getDefaultArguments(this.llmType, this.temperature, options.stop),
+        },
+      },
+    };
+
+    // create an actions client from the authenticated request context:
+    const actionsClient = await this.#actions.getActionsClientWithRequest(this.#request);
+
+    const actionResult = await actionsClient.execute(requestBody);
+
+    if (actionResult.status === 'error') {
+      throw new Error(
+        `ActionsClientSimpleChatModel: action result status is error: ${actionResult?.message} - ${actionResult?.serviceMessage}`
+      );
+    }
+
+    const content = get('data.message', actionResult);
+
+    if (typeof content !== 'string') {
+      throw new Error(
+        `ActionsClientSimpleChatModel: content should be a string, but it had an unexpected type: ${typeof content}`
+      );
+    }
+
+    return content; // per the contact of _call, return a string
+  }
+}

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

@@ -14,7 +14,7 @@ import { transformError } from '@kbn/securitysolution-es-utils';
 import { RetrievalQAChain } from 'langchain/chains';
 import {
   ActionsClientChatOpenAI,
-  ActionsClientLlm,
+  ActionsClientSimpleChatModel,
 } from '@kbn/elastic-assistant-common/impl/language_models';
 import { getDefaultArguments } from '@kbn/elastic-assistant-common/impl/language_models/constants';
 import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store';
@@ -26,7 +26,7 @@ import { AssistantToolParams } from '../../../types';
 export const DEFAULT_AGENT_EXECUTOR_ID = 'Elastic AI Assistant Agent Executor';
 
 /**
- * The default agent executor used by the Elastic AI Assistant. Main agent/chain that wraps the ActionsClientLlm,
+ * The default agent executor used by the Elastic AI Assistant. Main agent/chain that wraps the ActionsClientSimpleChatModel,
  * sets up a conversation BufferMemory from chat history, and registers tools like the ESQLKnowledgeBaseTool.
  *
  */
@@ -55,7 +55,7 @@ export const callAgentExecutor: AgentExecutor<true | false> = async ({
 }) => {
   // TODO implement llmClass for bedrock streaming
   // tracked here: https://github.com/elastic/security-team/issues/7363
-  const llmClass = isStream ? ActionsClientChatOpenAI : ActionsClientLlm;
+  const llmClass = isStream ? ActionsClientChatOpenAI : ActionsClientSimpleChatModel;
 
   const llm = new llmClass({
     actions,
@@ -130,11 +130,16 @@ export const callAgentExecutor: AgentExecutor<true | false> = async ({
         verbose: false,
       })
     : await initializeAgentExecutorWithOptions(tools, llm, {
-        agentType: 'chat-conversational-react-description',
+        agentType: 'structured-chat-zero-shot-react-description',
         memory,
         verbose: false,
+        returnIntermediateSteps: false,
+        handleParsingErrors: 'Try again, paying close attention to the allowed tool input',
+        agentArgs: {
+          // this is important to help LangChain correctly format tool input
+          humanMessageTemplate: `Question: {input}\n\n{agent_scratchpad}`,
+        },
       });
-
   // Sets up tracer for tracing executions to APM. See x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/README.mdx
   // If LangSmith env vars are set, executions will be traced there as well. See https://docs.smith.langchain.com/tracing
   const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger);

x-pack/plugins/elastic_assistant/server/types.ts

@@ -37,8 +37,10 @@ import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'
 import {
   ActionsClientChatOpenAI,
   ActionsClientLlm,
+  ActionsClientSimpleChatModel,
 } from '@kbn/elastic-assistant-common/impl/language_models';
 
+import { DynamicStructuredTool } from '@langchain/core/dist/tools';
 import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/conversations';
 import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context';
 import { AIAssistantDataClient } from './ai_assistant_data_clients';
@@ -202,7 +204,7 @@ export interface AssistantTool {
   description: string;
   sourceRegister: string;
   isSupported: (params: AssistantToolParams) => boolean;
-  getTool: (params: AssistantToolParams) => Tool | null;
+  getTool: (params: AssistantToolParams) => Tool | DynamicStructuredTool | null;
 }
 
 export interface AssistantToolParams {
@@ -211,7 +213,7 @@ export interface AssistantToolParams {
   isEnabledKnowledgeBase: boolean;
   chain?: RetrievalQAChain;
   esClient: ElasticsearchClient;
-  llm?: ActionsClientLlm | ActionsClientChatOpenAI;
+  llm?: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel;
   modelExists: boolean;
   onNewReplacements?: (newReplacements: Replacements) => void;
   replacements?: Replacements;

x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts

@@ -5,15 +5,16 @@
  * 2.0.
  */
 
-import { DynamicTool } from '@langchain/core/tools';
+import { DynamicStructuredTool } from '@langchain/core/tools';
+import { z } from 'zod';
 import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server';
 import { APP_UI_ID } from '../../../../common';
 
 export type EsqlKnowledgeBaseToolParams = AssistantToolParams;
 
 const toolDetails = {
   description:
-    'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language. Input must always be the query on a single line, with no other text. Only output valid ES|QL queries as described above. Do not add any additional text to describe your output.',
+    'Call this for knowledge on how to build an ESQL query, or answer questions about the ES|QL query language. Input must always be the query on a single line, with no other text. Only output valid ES|QL queries as described above. Format ESQL correctly by surrounding it with back ticks. Do not add any additional text to describe your output.',
   id: 'esql-knowledge-base-tool',
   name: 'ESQLKnowledgeBaseTool',
 };
@@ -30,13 +31,16 @@ export const ESQL_KNOWLEDGE_BASE_TOOL: AssistantTool = {
     const { chain } = params as EsqlKnowledgeBaseToolParams;
     if (chain == null) return null;
 
-    return new DynamicTool({
+    return new DynamicStructuredTool({
       name: toolDetails.name,
       description: toolDetails.description,
+      schema: z.object({
+        question: z.string().describe(`The user's exact question about ESQL`),
+      }),
       func: async (input, _, cbManager) => {
         const result = await chain.invoke(
           {
-            query: input,
+            query: input.question,
           },
           cbManager
         );

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

@@ -68,12 +68,15 @@ export const InvokeAIActionParamsSchema = schema.object({
         {
           name: schema.string(),
           description: schema.string(),
-          parameters: schema.object({
-            type: schema.string(),
-            properties: schema.object({}, { unknowns: 'allow' }),
-            additionalProperties: schema.boolean(),
-            $schema: schema.string(),
-          }),
+          parameters: schema.object(
+            {
+              type: schema.string(),
+              properties: schema.object({}, { unknowns: 'allow' }),
+              additionalProperties: schema.boolean(),
+              $schema: schema.string(),
+            },
+            { unknowns: 'allow' }
+          ),
         },
         // Not sure if this will include other properties, we should pass them if it does
         { unknowns: 'allow' }