-
Notifications
You must be signed in to change notification settings - Fork 8.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Security Solution] [GenAi] refactor security ai assistant tools to u…
…se tool helper method (#212865) ## Summary Clean up some security ai assistant code. - Replace the usage of `new DynamicStructuredTool()` with the `tool()` helper method. This is the recommended approach today and has the correct types to work with [`Command`](https://langchain-ai.github.io/langgraphjs/concepts/low_level/#command). - Extract code such as the default assistant graph state and agentRunnableFactory to reduce cognitive overload. - Update AssistantTool type definition ### Checklist Check the PR satisfies following conditions. Reviewers should verify this PR satisfies this list as well. - [X] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md) - [X] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [X] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [X] If a plugin configuration key changed, check if it needs to be allowlisted in the cloud and added to the [docker list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker) - [X] This was checked for breaking HTTP API changes, and any breaking changes have been approved by the breaking-change committee. The `release_note:breaking` label should be applied in these situations. - [X] [Flaky Test Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was used on any tests changed - [X] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) ### Identify risks Does this PR introduce any risks? For example, consider risks like hard to test bugs, performance regression, potential of data loss. Describe the risk, its severity, and mitigation for each identified risk. Invite stakeholders and evaluate how to proceed before merging. - [ ] [See some risk examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx) - [ ] ... --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com> (cherry picked from commit d37fcb6)
- Loading branch information
Showing
14 changed files
with
318 additions
and
248 deletions.
There are no files selected for viewing
56 changes: 56 additions & 0 deletions
56
...ns/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/agentRunnable.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| /* | ||
| * 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 { ToolDefinition } from '@langchain/core/language_models/base'; | ||
| import { | ||
| ActionsClientChatBedrockConverse, | ||
| ActionsClientChatVertexAI, | ||
| ActionsClientChatOpenAI, | ||
| } from '@kbn/langchain/server'; | ||
| import type { StructuredToolInterface } from '@langchain/core/tools'; | ||
| import { | ||
| AgentRunnableSequence, | ||
| createOpenAIToolsAgent, | ||
| createStructuredChatAgent, | ||
| createToolCallingAgent, | ||
| } from 'langchain/agents'; | ||
| import { ChatPromptTemplate } from '@langchain/core/prompts'; | ||
|
|
||
| export const TOOL_CALLING_LLM_TYPES = new Set(['bedrock', 'gemini']); | ||
|
|
||
| export const agentRunableFactory = async ({ | ||
| llm, | ||
| isOpenAI, | ||
| llmType, | ||
| tools, | ||
| isStream, | ||
| prompt, | ||
| }: { | ||
| llm: ActionsClientChatBedrockConverse | ActionsClientChatVertexAI | ActionsClientChatOpenAI; | ||
| isOpenAI: boolean; | ||
| llmType: string | undefined; | ||
| tools: StructuredToolInterface[] | ToolDefinition[]; | ||
| isStream: boolean; | ||
| prompt: ChatPromptTemplate; | ||
| }): Promise<AgentRunnableSequence> => { | ||
| const params = { | ||
| llm, | ||
| tools, | ||
| streamRunnable: isStream, | ||
| prompt, | ||
| } as const; | ||
|
|
||
| if (isOpenAI || llmType === 'inference') { | ||
| return createOpenAIToolsAgent(params); | ||
| } | ||
|
|
||
| if (llmType && TOOL_CALLING_LLM_TYPES.has(llmType)) { | ||
| return createToolCallingAgent(params); | ||
| } | ||
|
|
||
| return createStructuredChatAgent(params); | ||
| }; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
86 changes: 86 additions & 0 deletions
86
...ty/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/state.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| /* | ||
| * 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 { ConversationResponse } from '@kbn/elastic-assistant-common'; | ||
| import { BaseMessage } from '@langchain/core/messages'; | ||
| import { Annotation } from '@langchain/langgraph'; | ||
| import { AgentStep, AgentAction, AgentFinish } from 'langchain/agents'; | ||
|
|
||
| export const getStateAnnotation = ({ getFormattedTime }: { getFormattedTime?: () => string }) => { | ||
| const graphAnnotation = Annotation.Root({ | ||
| input: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => '', | ||
| }), | ||
| lastNode: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => 'start', | ||
| }), | ||
| steps: Annotation<AgentStep[]>({ | ||
| reducer: (x: AgentStep[], y: AgentStep[]) => x.concat(y), | ||
| default: () => [], | ||
| }), | ||
| hasRespondStep: Annotation<boolean>({ | ||
| reducer: (x: boolean, y?: boolean) => y ?? x, | ||
| default: () => false, | ||
| }), | ||
| agentOutcome: Annotation<AgentAction | AgentFinish | undefined>({ | ||
| reducer: ( | ||
| x: AgentAction | AgentFinish | undefined, | ||
| y?: AgentAction | AgentFinish | undefined | ||
| ) => y ?? x, | ||
| default: () => undefined, | ||
| }), | ||
| messages: Annotation<BaseMessage[]>({ | ||
| reducer: (x: BaseMessage[], y: BaseMessage[]) => y ?? x, | ||
| default: () => [], | ||
| }), | ||
| chatTitle: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => '', | ||
| }), | ||
| llmType: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => 'unknown', | ||
| }), | ||
| isStream: Annotation<boolean>({ | ||
| reducer: (x: boolean, y?: boolean) => y ?? x, | ||
| default: () => false, | ||
| }), | ||
| isOssModel: Annotation<boolean>({ | ||
| reducer: (x: boolean, y?: boolean) => y ?? x, | ||
| default: () => false, | ||
| }), | ||
| connectorId: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => '', | ||
| }), | ||
| conversation: Annotation<ConversationResponse | undefined>({ | ||
| reducer: (x: ConversationResponse | undefined, y?: ConversationResponse | undefined) => | ||
| y ?? x, | ||
| default: () => undefined, | ||
| }), | ||
| conversationId: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => '', | ||
| }), | ||
| responseLanguage: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => 'English', | ||
| }), | ||
| provider: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: () => '', | ||
| }), | ||
| formattedTime: Annotation<string>({ | ||
| reducer: (x: string, y?: string) => y ?? x, | ||
| default: getFormattedTime ?? (() => ''), | ||
| }), | ||
| }); | ||
|
|
||
| return graphAnnotation; | ||
| }; |
Oops, something went wrong.