Skip to content

Commit

Permalink
[8.x] [Security GenAI] Remove assistantBedrockChat feature flag (#1…
Browse files Browse the repository at this point in the history 
…95461) (#195516)

# Backport

This will backport the following commits from `main` to `8.x`:
- [[Security GenAI] Remove `assistantBedrockChat` feature flag
(#195461)](#195461)

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

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

<!--BACKPORT [{"author":{"name":"Steph
Milovic","email":"[email protected]"},"sourceCommit":{"committedDate":"2024-10-08T22:31:05Z","message":"[Security
GenAI] Remove `assistantBedrockChat` feature flag
(#195461)","sha":"38d0bdd3deeb72ca4e502361fb0df1f68fa7f443","branchLabelMapping":{"^v9.0.0$":"main","^v8.16.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:fix","v9.0.0","Team:
SecuritySolution","backport:prev-minor","Team:Security Generative
AI","v8.16.0"],"number":195461,"url":"https://github.com/elastic/kibana/pull/195461","mergeCommit":{"message":"[Security
GenAI] Remove `assistantBedrockChat` feature flag
(#195461)","sha":"38d0bdd3deeb72ca4e502361fb0df1f68fa7f443"}},"sourceBranch":"main","suggestedTargetBranches":["8.x"],"targetPullRequestStates":[{"branch":"main","label":"v9.0.0","labelRegex":"^v9.0.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/195461","number":195461,"mergeCommit":{"message":"[Security
GenAI] Remove `assistantBedrockChat` feature flag
(#195461)","sha":"38d0bdd3deeb72ca4e502361fb0df1f68fa7f443"}},{"branch":"8.x","label":"v8.16.0","labelRegex":"^v8.16.0$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

Co-authored-by: Elastic Machine <[email protected]>
  • Loading branch information
@stephmilovic @elasticmachine
stephmilovic and elasticmachine authored Oct 9, 2024
1 parent abd6766 commit 6b9c758
Show file tree
Hide file tree
Showing 18 changed files with 131 additions and 187 deletions.
10 changes: 5 additions & 5 deletions api_docs/kbn_elastic_assistant_common.devdocs.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -2574,7 +2574,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,
Expand Down Expand Up @@ -4522,7 +4522,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,
Expand Down Expand Up @@ -4987,7 +4987,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,
Expand Down Expand Up @@ -5941,4 +5941,4 @@
}
]
}
}
}
14 changes: 7 additions & 7 deletions api_docs/security_solution.devdocs.json
View file

Large diffs are not rendered by default.

1 change: 0 additions & 1 deletion x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,5 +21,4 @@ export type AssistantFeatureKey = keyof AssistantFeatures;
export const defaultAssistantFeatures = Object.freeze({
assistantKnowledgeBaseByDefault: false,
assistantModelEvaluation: false,
assistantBedrockChat: true,
});
1 change: 0 additions & 1 deletion ...ages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
});
3 changes: 0 additions & 3 deletions ...kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,11 @@ paths:
schema:
type: object
properties:
assistantBedrockChat:
type: boolean
assistantKnowledgeBaseByDefault:
type: boolean
assistantModelEvaluation:
type: boolean
required:
- assistantBedrockChat
- assistantKnowledgeBaseByDefault
- assistantModelEvaluation
'400':
Expand Down
1 change: 0 additions & 1 deletion x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
4 changes: 0 additions & 4 deletions ...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
179 changes: 97 additions & 82 deletions ...ins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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()
Expand Down Expand Up @@ -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' });
Expand All @@ -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(() => {
Expand All @@ -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);
});
});
});
Loading

0 comments on commit 6b9c758

Please sign in to comment.