From 2b29f63e969ef031038bbafc3de0115e30a9fbf8 Mon Sep 17 00:00:00 2001 From: Andrew Macri Date: Thu, 18 Apr 2024 15:01:27 -0400 Subject: [PATCH] [Security Solution] [AI Insights] Enable LangSmith tracing in cloud deployments (#181159) ## [Security Solution] [AI Insights] Enable LangSmith tracing in cloud deployments ### Summary This PR enables LangSmith tracing for the [AI Insights](https://github.com/elastic/kibana/pull/180611) feature in cloud deployments. LangSmith project names and API keys are specified using the same UI and patterns introduced by @spong in ### Details To enable LangSmith tracing in cloud deployments, configure the following `xpack.securitySolution.enableExperimental` feature flags: ``` xpack.securitySolution.enableExperimental: ['assistantModelEvaluation', 'assistantAlertsInsights'] ``` - The `assistantModelEvaluation` feature flag enables the `Evaluation` settings category in the assistant. The LangSmith project name and API key are configured here - The `assistantAlertsInsights` feature flag enables the AI Insights feature, which is off by default at the time of this writing After enabling the feature flags above: 1) Click the `AI Assistant` button anywhere in the Security Solution to launch the assistant 2) Click the settings gear in the assistant 3) Click the `Evaluation` settings category 4) Click `Show Trace Options (for internal use only)` 5) Specify a `LangSmith Project` and `LangSmith API Key` per the screenshot below: ![langsmith_settings](https://github.com/elastic/kibana/assets/4459398/896c8155-3a06-4381-becf-b846b5aba426) (cherry picked from commit 89609fe596d79b7d2eb4f209c5388824f9b279c1) --- .../impl/llm/actions_client_llm.ts | 11 +++++++--- .../impl/llm/types.ts | 10 +++++++++ .../alerts/post_alerts_insights_route.gen.ts | 2 ++ .../post_alerts_insights_route.schema.yaml | 4 ++++ .../insights/alerts/post_alerts_insights.ts | 22 ++++++++++++++++++- .../public/ai_insights/use_insights/index.tsx | 11 +++++++++- 6 files changed, 55 insertions(+), 5 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/llm/actions_client_llm.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/llm/actions_client_llm.ts index ab77e90c1c3f4..8df4c60817c6c 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/llm/actions_client_llm.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/llm/actions_client_llm.ts @@ -5,13 +5,14 @@ * 2.0. */ -import { v4 as uuidv4 } from 'uuid'; -import { KibanaRequest, Logger } from '@kbn/core/server'; import type { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server'; +import { KibanaRequest, Logger } from '@kbn/core/server'; import { LLM } from '@langchain/core/language_models/llms'; import { get } from 'lodash/fp'; +import { v4 as uuidv4 } from 'uuid'; import { getMessageContentAndRole } from './helpers'; +import { TraceOptions } from './types'; const LLM_TYPE = 'ActionsClientLlm'; @@ -27,6 +28,7 @@ interface ActionsClientLlmParams { model?: string; temperature?: number; traceId?: string; + traceOptions?: TraceOptions; } export class ActionsClientLlm extends LLM { @@ -52,8 +54,11 @@ export class ActionsClientLlm extends LLM { model, request, temperature, + traceOptions, }: ActionsClientLlmParams) { - super({}); + super({ + callbacks: [...(traceOptions?.tracers ?? [])], + }); this.#actions = actions; this.#connectorId = connectorId; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/llm/types.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/llm/types.ts index 85c970b1b948a..5fd08e6ee4e2d 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/llm/types.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/llm/types.ts @@ -5,6 +5,7 @@ * 2.0. */ +import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; import { ChatCompletionContentPart, ChatCompletionCreateParamsNonStreaming, @@ -38,3 +39,12 @@ export interface InvokeAIActionParamsSchema { functions?: ChatCompletionCreateParamsNonStreaming['functions']; signal?: AbortSignal; } + +export interface TraceOptions { + evaluationId?: string; + exampleId?: string; + projectName?: string; + runName?: string; + tags?: string[]; + tracers?: LangChainTracer[]; +} diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.gen.ts index 2f64f3c6e521d..cfae1a004a54d 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.gen.ts @@ -56,6 +56,8 @@ export const AlertsInsightsPostRequestBody = z.object({ anonymizationFields: z.array(AnonymizationFieldResponse), connectorId: z.string(), actionTypeId: z.string(), + langSmithProject: z.string().optional(), + langSmithApiKey: z.string().optional(), model: z.string().optional(), replacements: Replacements.optional(), size: z.number(), diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.schema.yaml index 5ab58c6023ee5..a4a647784ee29 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/insights/alerts/post_alerts_insights_route.schema.yaml @@ -73,6 +73,10 @@ paths: type: string actionTypeId: type: string + langSmithProject: + type: string + langSmithApiKey: + type: string model: type: string replacements: diff --git a/x-pack/plugins/elastic_assistant/server/routes/insights/alerts/post_alerts_insights.ts b/x-pack/plugins/elastic_assistant/server/routes/insights/alerts/post_alerts_insights.ts index 455f50703e836..e58e4cd0f8e20 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/insights/alerts/post_alerts_insights.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/insights/alerts/post_alerts_insights.ts @@ -19,6 +19,7 @@ import { transformError } from '@kbn/securitysolution-es-utils'; import { INSIGHTS_ALERTS } from '../../../../common/constants'; import { getAssistantToolParams, isInsightsFeatureEnabled } from './helpers'; import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../../helpers'; +import { getLangSmithTracer } from '../../evaluate/utils'; import { buildResponse } from '../../../lib/build_response'; import { ElasticAssistantRequestHandlerContext } from '../../../types'; import { getLlmType } from '../../utils'; @@ -73,7 +74,14 @@ export const postAlertsInsightsRoute = (router: IRouter