From 1c96c31aee5dc3b9b4f21639cd21e81c200bf76d Mon Sep 17 00:00:00 2001 From: Andrew Macri Date: Wed, 15 May 2024 20:26:20 -0400 Subject: [PATCH] [Security Solution] [Attack discovery] Overrides default Attack discovery timeouts (#183575) ## [Security Solution] [Attack discovery] Overrides default Attack discovery timeouts ### Summary This PR fixes an issue where Attack discovery requests may be retried when responses from the LLM take longer than two minutes. In LangSmith, the retry looks like the following _before_ screenshot: #### Before ![langsmith_before](https://github.com/elastic/kibana/assets/4459398/b02f016c-c260-43f3-a6cc-1260ca8d99c2) _Above: Before the fix, a retry, shown in LangSmith, for an LLM call > 2 minutes_ After the fix, a single pair for runs > 2 minutes are observed in LangSmith: #### After ![langsmith_after](https://github.com/elastic/kibana/assets/4459398/864ef2d4-f845-4d62-ab30-686211aadf30) _Above: After the fix, a single pair in LangSmith, for an LLM call > 2 minutes_ ### Details This PR overrides the following default timeouts: 1) The attack discovery route's `idleSocket` socket timeout in `x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts` 2) The connector timeout (also in `x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts`) 3) The chain timeout in `x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.ts` with the following defaults: ```typescript const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes const LANG_CHAIN_TIMEOUT = ROUTE_HANDLER_TIMEOUT - 10_000; // 9 minutes 50 seconds const CONNECTOR_TIMEOUT = LANG_CHAIN_TIMEOUT - 10_000; // 9 minutes 40 seconds ``` ### Desk testing 1) Verify there are ~ 100 open alerts in the last 24 hours in your testing environment 2) Navigate to Security > Attack discovery 3) Select an Azure / OpenAI connector 4) Click Generate **Expected results** - LangSmith displays a single pair of `LLMChain` and `AttackDiscovery` runs when the LLM responds (with the final answer) in less than 2 minutes - LangSmith displays a single pair of `LLMChain` and `AttackDiscovery` runs when the LLM takes longer than two minutes to respond (with the final answer), as illustrated by the `before` / `after` screenshots in the description above --- .../server/routes/attack_discovery/helpers.ts | 3 +++ .../routes/attack_discovery/post_attack_discovery.ts | 9 +++++++++ x-pack/plugins/elastic_assistant/server/types.ts | 1 + .../tools/attack_discovery/attack_discovery_tool.ts | 2 ++ 4 files changed, 15 insertions(+) diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts index f11802c206693..c473c21da6e84 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts @@ -37,6 +37,7 @@ export const getAssistantToolParams = ({ alertsIndexPattern, anonymizationFields, esClient, + langChainTimeout, latestReplacements, llm, onNewReplacements, @@ -46,6 +47,7 @@ export const getAssistantToolParams = ({ alertsIndexPattern: string; anonymizationFields?: AnonymizationFieldResponse[]; esClient: ElasticsearchClient; + langChainTimeout: number; latestReplacements: Replacements; llm: ActionsClientLlm; onNewReplacements: (newReplacements: Replacements) => void; @@ -61,6 +63,7 @@ export const getAssistantToolParams = ({ isEnabledKnowledgeBase: false, // not required for attack discovery chain: undefined, // not required for attack discovery esClient, + langChainTimeout, llm, modelExists: false, // not required for attack discovery onNewReplacements, diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts index 12546ddbb25b8..923f33dc44f3d 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/post_attack_discovery.ts @@ -24,6 +24,10 @@ import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantRequestHandlerContext } from '../../types'; import { getLlmType } from '../utils'; +const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes +const LANG_CHAIN_TIMEOUT = ROUTE_HANDLER_TIMEOUT - 10_000; // 9 minutes 50 seconds +const CONNECTOR_TIMEOUT = LANG_CHAIN_TIMEOUT - 10_000; // 9 minutes 40 seconds + export const postAttackDiscoveryRoute = ( router: IRouter ) => { @@ -33,6 +37,9 @@ export const postAttackDiscoveryRoute = ( path: ATTACK_DISCOVERY, options: { tags: ['access:elasticAssistant'], + timeout: { + idleSocket: ROUTE_HANDLER_TIMEOUT, + }, }, }) .addVersion( @@ -109,6 +116,7 @@ export const postAttackDiscoveryRoute = ( logger, request, temperature: 0, // zero temperature for attack discovery, because we want structured JSON output + timeout: CONNECTOR_TIMEOUT, traceOptions, }); @@ -117,6 +125,7 @@ export const postAttackDiscoveryRoute = ( anonymizationFields, esClient, latestReplacements, + langChainTimeout: LANG_CHAIN_TIMEOUT, llm, onNewReplacements, request, diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index 063573be37c3e..63b62e8943c20 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -211,6 +211,7 @@ export interface AssistantToolParams { isEnabledKnowledgeBase: boolean; chain?: RetrievalQAChain; esClient: ElasticsearchClient; + langChainTimeout?: number; llm?: ActionsClientLlm | ActionsClientChatOpenAI; modelExists: boolean; onNewReplacements?: (newReplacements: Replacements) => void; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.ts index 0d321e7402400..264862d76b8f5 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.ts @@ -53,6 +53,7 @@ export const ATTACK_DISCOVERY_TOOL: AssistantTool = { alertsIndexPattern, anonymizationFields, esClient, + langChainTimeout, llm, onNewReplacements, replacements, @@ -102,6 +103,7 @@ export const ATTACK_DISCOVERY_TOOL: AssistantTool = { const result = await answerFormattingChain.call({ query: getAttackDiscoveryPrompt({ anonymizedAlerts }), + timeout: langChainTimeout, }); const attackDiscoveries = result.records;