[8.14] [Security Solution] [AI Insights] Enable LangSmith tracing in cloud deployments (#181159)

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;

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[];
+}

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(),

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:

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<ElasticAssistantRequestH
           // get parameters from the request body
           const alertsIndexPattern = decodeURIComponent(request.body.alertsIndexPattern);
           const connectorId = decodeURIComponent(request.body.connectorId);
-          const { actionTypeId, anonymizationFields, replacements, size } = request.body;
+          const {
+            actionTypeId,
+            anonymizationFields,
+            langSmithApiKey,
+            langSmithProject,
+            replacements,
+            size,
+          } = request.body;
 
           // get an Elasticsearch client for the authenticated user:
           const esClient = (await context.core).elasticsearch.client.asCurrentUser;
@@ -91,13 +99,25 @@ export const postAlertsInsightsRoute = (router: IRouter<ElasticAssistantRequestH
             return response.notFound(); // insights tool not found
           }
 
+          const traceOptions = {
+            projectName: langSmithProject,
+            tracers: [
+              ...getLangSmithTracer({
+                apiKey: langSmithApiKey,
+                projectName: langSmithProject,
+                logger,
+              }),
+            ],
+          };
+
           const llm = new ActionsClientLlm({
             actions,
             connectorId,
             llmType: getLlmType(actionTypeId),
             logger,
             request,
             temperature: 0, // zero temperature for insights, because we want structured JSON output
+            traceOptions,
           });
 
           const assistantToolParams = getAssistantToolParams({

x-pack/plugins/security_solution/public/ai_insights/use_insights/index.tsx

@@ -16,6 +16,7 @@ import {
   AlertsInsightsPostResponse,
   ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION,
 } from '@kbn/elastic-assistant-common';
+import { isEmpty } from 'lodash/fp';
 import moment from 'moment';
 import React, { useCallback, useMemo, useState } from 'react';
 import { useLocalStorage, useSessionStorage } from 'react-use';
@@ -60,7 +61,7 @@ export const useInsights = ({
   const [isLoading, setIsLoading] = useState(false);
 
   // get alerts index pattern and allow lists from the assistant context:
-  const { alertsIndexPattern, knowledgeBase } = useAssistantContext();
+  const { alertsIndexPattern, knowledgeBase, traceOptions } = useAssistantContext();
 
   const { data: anonymizationFields } = useFetchAnonymizationFields();
 
@@ -116,6 +117,12 @@ export const useInsights = ({
       alertsIndexPattern: alertsIndexPattern ?? '',
       anonymizationFields: anonymizationFields?.data ?? [],
       connectorId: connectorId ?? '',
+      langSmithProject: isEmpty(traceOptions?.langSmithProject)
+        ? undefined
+        : traceOptions?.langSmithProject,
+      langSmithApiKey: isEmpty(traceOptions?.langSmithApiKey)
+        ? undefined
+        : traceOptions?.langSmithApiKey,
       size: knowledgeBase.latestAlerts,
       replacements: {}, // no need to re-use replacements in the current implementation
       subAction: 'invokeAI', // non-streaming
@@ -227,6 +234,8 @@ export const useInsights = ({
     setLocalStorageGenerationIntervals,
     setSessionStorageCachedInsights,
     toasts,
+    traceOptions?.langSmithApiKey,
+    traceOptions?.langSmithProject,
   ]);
 
   return {