[Security GenAI] Remove assistantBedrockChat feature flag

api_docs/kbn_elastic_assistant_common.devdocs.json

@@ -844,7 +844,7 @@
           "\nInterface for features available to the elastic assistant"
         ],
         "signature": [
-          "{ readonly assistantKnowledgeBaseByDefault: boolean; readonly assistantModelEvaluation: boolean; readonly assistantBedrockChat: boolean; }"
+          "{ readonly assistantKnowledgeBaseByDefault: boolean; readonly assistantModelEvaluation: boolean; }"
         ],
         "path": "x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts",
         "deprecated": false,
@@ -2604,7 +2604,7 @@
         "label": "GetCapabilitiesResponse",
         "description": [],
         "signature": [
-          "{ assistantKnowledgeBaseByDefault: boolean; assistantModelEvaluation: boolean; assistantBedrockChat: boolean; }"
+          "{ assistantKnowledgeBaseByDefault: boolean; assistantModelEvaluation: boolean; }"
         ],
         "path": "x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts",
         "deprecated": false,
@@ -4567,7 +4567,7 @@
           "\nDefault features available to the elastic assistant"
         ],
         "signature": [
-          "{ readonly assistantKnowledgeBaseByDefault: false; readonly assistantModelEvaluation: false; readonly assistantBedrockChat: true; }"
+          "{ readonly assistantKnowledgeBaseByDefault: false; readonly assistantModelEvaluation: false; }"
         ],
         "path": "x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts",
         "deprecated": false,
@@ -5032,7 +5032,7 @@
         "label": "GetCapabilitiesResponse",
         "description": [],
         "signature": [
-          "Zod.ZodObject<{ assistantBedrockChat: Zod.ZodBoolean; assistantKnowledgeBaseByDefault: Zod.ZodBoolean; assistantModelEvaluation: Zod.ZodBoolean; }, \"strip\", Zod.ZodTypeAny, { assistantKnowledgeBaseByDefault: boolean; assistantModelEvaluation: boolean; assistantBedrockChat: boolean; }, { assistantKnowledgeBaseByDefault: boolean; assistantModelEvaluation: boolean; assistantBedrockChat: boolean; }>"
+          "Zod.ZodObject<{ assistantKnowledgeBaseByDefault: Zod.ZodBoolean; assistantModelEvaluation: Zod.ZodBoolean; }, \"strip\", Zod.ZodTypeAny, { assistantKnowledgeBaseByDefault: boolean; assistantModelEvaluation: boolean; }, { assistantKnowledgeBaseByDefault: boolean; assistantModelEvaluation: boolean; }>"
         ],
         "path": "x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts",
         "deprecated": false,
@@ -5986,4 +5986,4 @@
       }
     ]
   }
-}
+}

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

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

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts

@@ -18,7 +18,6 @@ import { z } from '@kbn/zod';
 
 export type GetCapabilitiesResponse = z.infer<typeof GetCapabilitiesResponse>;
 export const GetCapabilitiesResponse = z.object({
-  assistantBedrockChat: z.boolean(),
   assistantKnowledgeBaseByDefault: z.boolean(),
   assistantModelEvaluation: z.boolean(),
 });

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml

@@ -20,14 +20,11 @@ paths:
               schema:
                 type: object
                 properties:
-                  assistantBedrockChat:
-                    type: boolean
                   assistantKnowledgeBaseByDefault:
                     type: boolean
                   assistantModelEvaluation:
                     type: boolean
                 required:
-                  - assistantBedrockChat
                   - assistantKnowledgeBaseByDefault
                   - assistantModelEvaluation
         '400':

x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts

@@ -37,7 +37,6 @@ export interface AgentExecutorParams<T extends boolean> {
   abortSignal?: AbortSignal;
   alertsIndexPattern?: string;
   actionsClient: PublicMethodsOf<ActionsClient>;
-  bedrockChatEnabled: boolean;
   assistantTools?: AssistantTool[];
   connectorId: string;
   conversationId?: string;

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

@@ -86,10 +86,6 @@ export const getDefaultAssistantGraph = ({
         value: (x: string, y?: string) => y ?? x,
         default: () => 'unknown',
       },
-      bedrockChatEnabled: {
-        value: (x: boolean, y?: boolean) => y ?? x,
-        default: () => false,
-      },
       isStream: {
         value: (x: boolean, y?: boolean) => y ?? x,
         default: () => false,

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

@@ -49,6 +49,18 @@ describe('streamGraph', () => {
     streamEvents: mockStreamEvents,
   } as unknown as DefaultAssistantGraph;
   const mockOnLlmResponse = jest.fn().mockResolvedValue(null);
+  const requestArgs = {
+    apmTracer: mockApmTracer,
+    assistantGraph: mockAssistantGraph,
+    inputs: {
+      input: 'input',
+      responseLanguage: 'English',
+      llmType: 'openai',
+    },
+    logger: mockLogger,
+    onLlmResponse: mockOnLlmResponse,
+    request: mockRequest,
+  };
 
   beforeEach(() => {
     jest.clearAllMocks();
@@ -59,8 +71,8 @@ describe('streamGraph', () => {
       transaction: { ids: { 'transaction.id': 'transactionId' } },
     });
   });
-  describe('ActionsClientChatOpenAI', () => {
-    it('should execute the graph in streaming mode', async () => {
+  describe('OpenAI Function Agent streaming', () => {
+    it('should execute the graph in streaming mode - OpenAI + isOssModel = false', async () => {
       mockStreamEvents.mockReturnValue({
         next: jest
           .fn()
@@ -93,19 +105,7 @@ describe('streamGraph', () => {
         return: jest.fn(),
       });
 
-      const response = await streamGraph({
-        apmTracer: mockApmTracer,
-        assistantGraph: mockAssistantGraph,
-        inputs: {
-          input: 'input',
-          bedrockChatEnabled: false,
-          llmType: 'openai',
-          responseLanguage: 'English',
-        },
-        logger: mockLogger,
-        onLlmResponse: mockOnLlmResponse,
-        request: mockRequest,
-      });
+      const response = await streamGraph(requestArgs);
 
       expect(response).toBe(mockResponseWithHeaders);
       expect(mockPush).toHaveBeenCalledWith({ payload: 'content', type: 'content' });
@@ -119,78 +119,49 @@ describe('streamGraph', () => {
     });
   });
 
-  describe('ActionsClientSimpleChatModel', () => {
-    it('should execute the graph in streaming mode', async () => {
-      mockStreamEvents.mockReturnValue({
-        next: jest
-          .fn()
-          .mockResolvedValueOnce({
-            value: {
-              name: 'ActionsClientSimpleChatModel',
-              event: 'on_llm_stream',
-              data: {
-                chunk: {
-                  content:
-                    '```json\n\n  "action": "Final Answer",\n  "action_input": "Look at these',
-                },
-              },
-              tags: [AGENT_NODE_TAG],
+  describe('Tool Calling Agent and Structured Chat Agent streaming', () => {
+    const mockAsyncIterator = {
+      async *[Symbol.asyncIterator]() {
+        yield {
+          event: 'on_chat_model_stream',
+          data: {
+            chunk: {
+              content: 'Look at these',
             },
-            done: false,
-          })
-          .mockResolvedValueOnce({
-            value: {
-              name: 'ActionsClientSimpleChatModel',
-              event: 'on_llm_stream',
-              data: {
-                chunk: {
-                  content: ' rare IP',
-                },
-              },
-              tags: [AGENT_NODE_TAG],
+          },
+          tags: [AGENT_NODE_TAG],
+        };
+        yield {
+          event: 'on_chat_model_stream',
+          data: {
+            chunk: {
+              content: ' rare IP',
             },
-            done: false,
-          })
-          .mockResolvedValueOnce({
-            value: {
-              name: 'ActionsClientSimpleChatModel',
-              event: 'on_llm_stream',
-              data: {
-                chunk: {
-                  content: ' addresses." }```',
-                },
-              },
-              tags: [AGENT_NODE_TAG],
+          },
+          tags: [AGENT_NODE_TAG],
+        };
+        yield {
+          event: 'on_chat_model_stream',
+          data: {
+            chunk: {
+              content: ' addresses.',
             },
-            done: false,
-          })
-          .mockResolvedValueOnce({
-            value: {
-              name: 'ActionsClientSimpleChatModel',
-              event: 'on_llm_end',
-              tags: [AGENT_NODE_TAG],
+          },
+          tags: [AGENT_NODE_TAG],
+        };
+        yield {
+          event: 'on_chat_model_end',
+          data: {
+            output: {
+              content: 'Look at these rare IP addresses.',
             },
-          })
-          .mockResolvedValue({
-            done: true,
-          }),
-        return: jest.fn(),
-      });
-
-      const response = await streamGraph({
-        apmTracer: mockApmTracer,
-        assistantGraph: mockAssistantGraph,
-        inputs: {
-          input: 'input',
-          bedrockChatEnabled: false,
-          responseLanguage: 'English',
-          llmType: 'gemini',
-        },
-        logger: mockLogger,
-        onLlmResponse: mockOnLlmResponse,
-        request: mockRequest,
-      });
+          },
+          tags: [AGENT_NODE_TAG],
+        };
+      },
+    };
 
+    const expectConditions = async (response: unknown) => {
       expect(response).toBe(mockResponseWithHeaders);
 
       await waitFor(() => {
@@ -203,6 +174,50 @@ describe('streamGraph', () => {
           false
         );
       });
+    };
+    it('should execute the graph in streaming mode - Gemini', async () => {
+      const mockAssistantGraphAsyncIterator = {
+        streamEvents: () => mockAsyncIterator,
+      } as unknown as DefaultAssistantGraph;
+      const response = await streamGraph({
+        ...requestArgs,
+        assistantGraph: mockAssistantGraphAsyncIterator,
+        inputs: {
+          ...requestArgs.inputs,
+          llmType: 'gemini',
+        },
+      });
+
+      await expectConditions(response);
+    });
+    it('should execute the graph in streaming mode - Bedrock', async () => {
+      const mockAssistantGraphAsyncIterator = {
+        streamEvents: () => mockAsyncIterator,
+      } as unknown as DefaultAssistantGraph;
+      const response = await streamGraph({
+        ...requestArgs,
+        assistantGraph: mockAssistantGraphAsyncIterator,
+        inputs: {
+          ...requestArgs.inputs,
+          llmType: 'bedrock',
+        },
+      });
+
+      await expectConditions(response);
+    });
+    it('should execute the graph in streaming mode - OpenAI + isOssModel = false', async () => {
+      const mockAssistantGraphAsyncIterator = {
+        streamEvents: () => mockAsyncIterator,
+      } as unknown as DefaultAssistantGraph;
+      const response = await streamGraph({
+        ...requestArgs,
+        assistantGraph: mockAssistantGraphAsyncIterator,
+        inputs: {
+          ...requestArgs.inputs,
+          isOssModel: true,
+        },
+      });
+      await expectConditions(response);
     });
   });
 });