[Security GenAI] Use AI setting to set langsmith tracing to the Integration Assistant (#187466)

## Summary

Enables tracing Langchain invocations in the integrations assistant
using the Langsmith settings stored by the Security AI Settings.
The evaluation settings tab is still under an experimental flag, to see
it:

```
xpack.securitySolution.enableExperimental: ['assistantModelEvaluation']
```

### Screenshots

<img width="1317" alt="Settings"
src="6aed1ef6-3750-4259-9fe2-b8bf1aed5504">

After one execution of the integration assistant:

<img width="1240" alt="langsmith"
src="dd3dd99c-7c83-4a35-95b2-789e7a341031">

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

x-pack/packages/kbn-langchain/server/tracers/apm/index.ts

@@ -0,0 +1,8 @@
+/*
+ * 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.
+ */
+
+export { APMTracer } from './apm_tracer';

x-pack/packages/kbn-langchain/server/tracers/langsmith/index.ts

@@ -0,0 +1,7 @@
+/*
+ * 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.
+ */
+export { getLangSmithTracer, isLangSmithEnabled } from './langsmith_tracer';

x-pack/packages/kbn-langchain/server/tracers/langsmith/langsmith_tracer.ts

@@ -0,0 +1,64 @@
+/*
+ * 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 { Client } from 'langsmith';
+import type { Logger } from '@kbn/core/server';
+import { ToolingLog } from '@kbn/tooling-log';
+import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
+
+/**
+ * Returns a custom LangChainTracer which adds the `exampleId` so Dataset 'Test' runs are written to LangSmith
+ * If `exampleId` is present (and a corresponding example exists in LangSmith) trace is written to the Dataset's `Tests`
+ * section, otherwise it is written to the `Project` provided
+ *
+ * @param apiKey API Key for LangSmith (will fetch from env vars if not provided)
+ * @param projectName Name of project to trace results to
+ * @param exampleId Dataset exampleId to associate trace with
+ * @param logger
+ */
+export const getLangSmithTracer = ({
+  apiKey,
+  projectName,
+  exampleId,
+  logger,
+}: {
+  apiKey?: string;
+  projectName?: string;
+  exampleId?: string;
+  logger: Logger | ToolingLog;
+}): LangChainTracer[] => {
+  try {
+    if (!isLangSmithEnabled() || apiKey == null) {
+      return [];
+    }
+    const lcTracer = new LangChainTracer({
+      projectName, // Shows as the 'test' run's 'name' in langsmith ui
+      exampleId,
+      client: new Client({ apiKey }),
+    });
+
+    return [lcTracer];
+  } catch (e) {
+    // Note: creating a tracer can fail if the LangSmith env vars are not set correctly
+    logger.error(`Error creating LangSmith tracer: ${e.message}`);
+  }
+
+  return [];
+};
+
+/**
+ * Returns true if LangSmith/tracing is enabled
+ */
+export const isLangSmithEnabled = (): boolean => {
+  try {
+    // Just checking if apiKey is available, if better way to check for enabled that is not env var please update
+    const config = Client.getDefaultClientConfig();
+    return config.apiKey != null;
+  } catch (e) {
+    return false;
+  }
+};

x-pack/packages/kbn-langchain/tsconfig.json

@@ -18,6 +18,7 @@
     "@kbn/logging",
     "@kbn/actions-plugin",
     "@kbn/logging-mocks",
-    "@kbn/utility-types"
+    "@kbn/utility-types",
+    "@kbn/tooling-log"
   ]
 }

x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts

@@ -18,12 +18,12 @@ import {
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server';
 import { MessagesPlaceholder } from '@langchain/core/prompts';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { withAssistantSpan } from '../tracers/apm/with_assistant_span';
 import { EsAnonymizationFieldsSchema } from '../../../ai_assistant_data_clients/anonymization_fields/types';
 import { transformESSearchToAnonymizationFields } from '../../../ai_assistant_data_clients/anonymization_fields/helpers';
 import { AgentExecutor } from '../executors/types';
-import { APMTracer } from '../tracers/apm_tracer';
 import { AssistantToolParams } from '../../../types';
-import { withAssistantSpan } from '../tracers/with_assistant_span';
 export const DEFAULT_AGENT_EXECUTOR_ID = 'Elastic AI Assistant Agent Executor';
 
 /**

x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts

@@ -11,9 +11,9 @@ import { BufferMemory, ChatMessageHistory } from 'langchain/memory';
 import { ChainTool } from 'langchain/tools/chain';
 
 import { ActionsClientLlm } from '@kbn/langchain/server';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { withAssistantSpan } from '../tracers/apm/with_assistant_span';
 import { AgentExecutor } from './types';
-import { withAssistantSpan } from '../tracers/with_assistant_span';
-import { APMTracer } from '../tracers/apm_tracer';
 
 export const OPEN_AI_FUNCTIONS_AGENT_EXECUTOR_ID =
   'Elastic AI Assistant Agent Executor (OpenAI Functions)';

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts

@@ -11,11 +11,11 @@ import { streamFactory, StreamResponseWithHeaders } from '@kbn/ml-response-strea
 import { transformError } from '@kbn/securitysolution-es-utils';
 import type { KibanaRequest } from '@kbn/core-http-server';
 import type { ExecuteConnectorRequestBody, TraceData } from '@kbn/elastic-assistant-common';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { withAssistantSpan } from '../../tracers/apm/with_assistant_span';
 import { AGENT_NODE_TAG } from './nodes/run_agent';
 import { DEFAULT_ASSISTANT_GRAPH_ID, DefaultAssistantGraph } from './graph';
 import type { OnLlmResponse, TraceOptions } from '../../executors/types';
-import type { APMTracer } from '../../tracers/apm_tracer';
-import { withAssistantSpan } from '../../tracers/with_assistant_span';
 
 interface StreamGraphParams {
   apmTracer: APMTracer;

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts

@@ -13,11 +13,11 @@ import {
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server';
 import { createOpenAIFunctionsAgent, createStructuredChatAgent } from 'langchain/agents';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
 import { EsAnonymizationFieldsSchema } from '../../../../ai_assistant_data_clients/anonymization_fields/types';
 import { AssistantToolParams } from '../../../../types';
 import { AgentExecutor } from '../../executors/types';
 import { openAIFunctionAgentPrompt, structuredChatAgentPrompt } from './prompts';
-import { APMTracer } from '../../tracers/apm_tracer';
 import { getDefaultAssistantGraph } from './graph';
 import { invokeGraph, streamGraph } from './helpers';
 import { transformESSearchToAnonymizationFields } from '../../../../ai_assistant_data_clients/anonymization_fields/helpers';

x-pack/plugins/elastic_assistant/server/lib/model_evaluator/evaluation.ts

@@ -15,9 +15,10 @@ import { ToolingLog } from '@kbn/tooling-log';
 import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
 import { RunCollectorCallbackHandler } from '@langchain/core/tracers/run_collector';
 import { Dataset } from '@kbn/elastic-assistant-common';
+import { isLangSmithEnabled } from '@kbn/langchain/server/tracers/langsmith';
 import { AgentExecutorEvaluatorWithMetadata } from '../langchain/executors/types';
 import { callAgentWithRetry, getMessageFromLangChainResponse } from './utils';
-import { isLangSmithEnabled, writeLangSmithFeedback } from '../../routes/evaluate/utils';
+import { writeLangSmithFeedback } from '../../routes/evaluate/utils';
 import { ResponseBody } from '../langchain/types';
 
 export interface PerformEvaluationParams {

x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts

@@ -28,7 +28,7 @@ import type { ActionsClient } from '@kbn/actions-plugin/server';
 import moment from 'moment/moment';
 import { uniq } from 'lodash/fp';
 import { PublicMethodsOf } from '@kbn/utility-types';
-import { getLangSmithTracer } from '../evaluate/utils';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { getLlmType } from '../utils';
 import type { GetRegisteredTools } from '../../services/app_context';
 import {

x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts

@@ -18,6 +18,7 @@ import {
   ExecuteConnectorRequestBody,
 } from '@kbn/elastic-assistant-common';
 import { ActionsClientLlm } from '@kbn/langchain/server';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common';
 import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from '../knowledge_base/constants';
 import { buildResponse } from '../../lib/build_response';
@@ -29,7 +30,7 @@ import {
   indexEvaluations,
   setupEvaluationIndex,
 } from '../../lib/model_evaluator/output_index/utils';
-import { fetchLangSmithDataset, getConnectorName, getLangSmithTracer } from './utils';
+import { fetchLangSmithDataset, getConnectorName } from './utils';
 import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
 
 /**

x-pack/plugins/elastic_assistant/server/routes/evaluate/utils.ts

@@ -10,8 +10,8 @@ import type { ActionResult } from '@kbn/actions-plugin/server';
 import type { Logger } from '@kbn/core/server';
 import type { Run } from 'langsmith/schemas';
 import { ToolingLog } from '@kbn/tooling-log';
-import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
 import { Dataset } from '@kbn/elastic-assistant-common';
+import { isLangSmithEnabled } from '@kbn/langchain/server/tracers/langsmith';
 
 /**
  * Return connector name for the given connectorId/connectors
@@ -97,56 +97,3 @@ export const writeLangSmithFeedback = async (
     return '';
   }
 };
-
-/**
- * Returns a custom LangChainTracer which adds the `exampleId` so Dataset 'Test' runs are written to LangSmith
- * If `exampleId` is present (and a corresponding example exists in LangSmith) trace is written to the Dataset's `Tests`
- * section, otherwise it is written to the `Project` provided
- *
- * @param apiKey API Key for LangSmith (will fetch from env vars if not provided)
- * @param projectName Name of project to trace results to
- * @param exampleId Dataset exampleId to associate trace with
- * @param logger
- */
-export const getLangSmithTracer = ({
-  apiKey,
-  projectName,
-  exampleId,
-  logger,
-}: {
-  apiKey?: string;
-  projectName?: string;
-  exampleId?: string;
-  logger: Logger | ToolingLog;
-}): LangChainTracer[] => {
-  try {
-    if (!isLangSmithEnabled() && apiKey == null) {
-      return [];
-    }
-    const lcTracer = new LangChainTracer({
-      projectName, // Shows as the 'test' run's 'name' in langsmith ui
-      exampleId,
-      client: new Client({ apiKey }),
-    });
-
-    return [lcTracer];
-  } catch (e) {
-    // Note: creating a tracer can fail if the LangSmith env vars are not set correctly
-    logger.error(`Error creating LangSmith tracer: ${e.message}`);
-  }
-
-  return [];
-};
-
-/**
- * Returns true if LangSmith/tracing is enabled
- */
-export const isLangSmithEnabled = (): boolean => {
-  try {
-    // Just checking if apiKey is available, if better way to check for enabled that is not env var please update
-    const config = Client.getDefaultClientConfig();
-    return config.apiKey != null;
-  } catch (e) {
-    return false;
-  }
-};

x-pack/plugins/elastic_assistant/server/routes/helpers.ts

@@ -27,6 +27,7 @@ import { i18n } from '@kbn/i18n';
 import { AwaitedProperties, PublicMethodsOf } from '@kbn/utility-types';
 import { ActionsClient } from '@kbn/actions-plugin/server';
 import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabilities';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { MINIMUM_AI_ASSISTANT_LICENSE } from '../../common/constants';
 import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './knowledge_base/constants';
 import { callAgentExecutor } from '../lib/langchain/execute_custom_llm_chain';
@@ -39,7 +40,6 @@ import {
 import { executeAction, StaticResponse } from '../lib/executor';
 import { getLangChainMessages } from '../lib/langchain/helpers';
 
-import { getLangSmithTracer } from './evaluate/utils';
 import { ElasticsearchStore } from '../lib/langchain/elasticsearch_store/elasticsearch_store';
 import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations';
 import { INVOKE_ASSISTANT_SUCCESS_EVENT } from '../lib/telemetry/event_based_telemetry';

x-pack/plugins/elastic_assistant/tsconfig.json

@@ -27,7 +27,6 @@
     "@kbn/core-elasticsearch-server",
     "@kbn/logging",
     "@kbn/ml-plugin",
-    "@kbn/apm-utils",
     "@kbn/elastic-assistant-common",
     "@kbn/core-http-router-server-mocks",
     "@kbn/data-stream-adapter",
@@ -46,6 +45,7 @@
     "@kbn/langchain",
     "@kbn/stack-connectors-plugin",
     "@kbn/security-plugin",
+    "@kbn/apm-utils",
   ],
   "exclude": [
     "target/**/*",

x-pack/plugins/integration_assistant/common/api/categorization/categorization_route.schema.yaml

@@ -34,6 +34,8 @@ paths:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Pipeline"
                 connectorId:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Connector"
+                langSmithOptions:
+                  $ref: "../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
       responses:
         200:
           description: Indicates a successful call.

x-pack/plugins/integration_assistant/common/api/categorization/categorization_route.ts

@@ -10,6 +10,7 @@ import { z } from 'zod';
 import {
   Connector,
   DataStreamName,
+  LangSmithOptions,
   PackageName,
   Pipeline,
   RawSamples,
@@ -23,6 +24,7 @@ export const CategorizationRequestBody = z.object({
   rawSamples: RawSamples,
   currentPipeline: Pipeline,
   connectorId: Connector,
+  langSmithOptions: LangSmithOptions.optional(),
 });
 export type CategorizationRequestBodyInput = z.input<typeof CategorizationRequestBody>;
 

x-pack/plugins/integration_assistant/common/api/ecs/ecs_route.schema.yaml

@@ -33,6 +33,8 @@ paths:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Mapping"
                 connectorId:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Connector"
+                langSmithOptions:
+                  $ref: "../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
       responses:
         200:
           description: Indicates a successful call.

x-pack/plugins/integration_assistant/common/api/ecs/ecs_route.ts

@@ -10,6 +10,7 @@ import { z } from 'zod';
 import {
   Connector,
   DataStreamName,
+  LangSmithOptions,
   Mapping,
   PackageName,
   RawSamples,
@@ -23,6 +24,7 @@ export const EcsMappingRequestBody = z.object({
   rawSamples: RawSamples,
   mapping: Mapping.optional(),
   connectorId: Connector,
+  langSmithOptions: LangSmithOptions.optional(),
 });
 export type EcsMappingRequestBodyInput = z.input<typeof EcsMappingRequestBody>;
 

x-pack/plugins/integration_assistant/common/api/model/common_attributes.schema.yaml

@@ -143,3 +143,17 @@ components:
         logo:
           type: string
           description: The logo of the integration.
+
+    LangSmithOptions:
+      type: object
+      description: The LangSmith options object.
+      required:
+        - projectName
+        - apiKey
+      properties:
+        projectName:
+          type: string
+          description: The project name.
+        apiKey:
+          type: string
+          description: The apiKey to use for tracing.

x-pack/plugins/integration_assistant/common/api/model/common_attributes.ts

@@ -156,3 +156,18 @@ export const Integration = z.object({
    */
   logo: z.string().optional(),
 });
+
+/**
+ * The LangSmith options object.
+ */
+export type LangSmithOptions = z.infer<typeof LangSmithOptions>;
+export const LangSmithOptions = z.object({
+  /**
+   * The project name to use with tracing.
+   */
+  projectName: z.string(),
+  /**
+   * The api key for the project
+   */
+  apiKey: z.string(),
+});

x-pack/plugins/integration_assistant/common/api/related/related_route.schema.yaml

@@ -34,6 +34,8 @@ paths:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Pipeline"
                 connectorId:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Connector"
+                langSmithOptions:
+                  $ref: "../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
       responses:
         200:
           description: Indicates a successful call.

x-pack/plugins/integration_assistant/common/api/related/related_route.ts

@@ -10,6 +10,7 @@ import { z } from 'zod';
 import {
   Connector,
   DataStreamName,
+  LangSmithOptions,
   PackageName,
   Pipeline,
   RawSamples,
@@ -23,6 +24,7 @@ export const RelatedRequestBody = z.object({
   rawSamples: RawSamples,
   currentPipeline: Pipeline,
   connectorId: Connector,
+  langSmithOptions: LangSmithOptions.optional(),
 });
 export type RelatedRequestBodyInput = z.input<typeof RelatedRequestBody>;
 

x-pack/plugins/integration_assistant/kibana.jsonc

@@ -18,5 +18,8 @@
       "actions",
       "stackConnectors",
     ],
+    "requiredBundles": [
+      "kibanaUtils",
+    ]
   }
 }

x-pack/plugins/integration_assistant/public/common/lib/lang_smith.ts

@@ -0,0 +1,37 @@
+/*
+ * 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 {
+  DEFAULT_ASSISTANT_NAMESPACE,
+  TRACE_OPTIONS_SESSION_STORAGE_KEY,
+} from '@kbn/elastic-assistant/impl/assistant_context/constants';
+import { Storage } from '@kbn/kibana-utils-plugin/public';
+import type { TraceOptions } from '@kbn/elastic-assistant/impl/assistant/types';
+import type { LangSmithOptions } from '../../../common/api/model/common_attributes';
+
+const sessionStorage = new Storage(window.sessionStorage);
+
+/**
+ * Retrieves the LangSmith options from the AI Settings.
+ */
+export const getLangSmithOptions = (
+  nameSpace: string = DEFAULT_ASSISTANT_NAMESPACE
+): LangSmithOptions | undefined => {
+  // Get the LangSmith options stored by the AI Settings using the assistant context
+  // TODO: Encapsulate all AI Settings logic in a generic place.
+  const sessionStorageTraceOptions: TraceOptions = sessionStorage.get(
+    `${nameSpace}.${TRACE_OPTIONS_SESSION_STORAGE_KEY}`
+  );
+
+  if (!sessionStorageTraceOptions) {
+    return;
+  }
+  return {
+    projectName: sessionStorageTraceOptions.langSmithProject,
+    apiKey: sessionStorageTraceOptions.langSmithApiKey,
+  };
+};

x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/steps/data_stream_step/generation_modal.tsx

@@ -24,6 +24,7 @@ import {
 import { isEmpty } from 'lodash/fp';
 import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import { css } from '@emotion/react';
+import { getLangSmithOptions } from '../../../../../common/lib/lang_smith';
 import type {
   CategorizationRequestBody,
   EcsMappingRequestBody,
@@ -87,6 +88,7 @@ export const useGeneration = ({
           dataStreamName: integrationSettings.dataStreamName ?? '',
           rawSamples: integrationSettings.logsSampleParsed ?? [],
           connectorId: connector.id,
+          langSmithOptions: getLangSmithOptions(),
         };
 
         setProgress('ecs');

x-pack/plugins/integration_assistant/server/routes/categorization_routes.ts

@@ -11,6 +11,8 @@ import {
   ActionsClientChatOpenAI,
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server/language_models';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import {
   CATEGORIZATION_GRAPH_PATH,
   CategorizationRequestBody,
@@ -46,7 +48,8 @@ export function registerCategorizationRoutes(
       },
       withAvailability(
         async (context, req, res): Promise<IKibanaResponse<CategorizationResponse>> => {
-          const { packageName, dataStreamName, rawSamples, currentPipeline } = req.body;
+          const { packageName, dataStreamName, rawSamples, currentPipeline, langSmithOptions } =
+            req.body;
           const services = await context.resolve(['core']);
           const { client } = services.core.elasticsearch;
           const { getStartServices, logger } = await context.integrationAssistant;
@@ -76,13 +79,22 @@ export function registerCategorizationRoutes(
               streaming: false,
             });
 
-            const graph = await getCategorizationGraph(client, model);
-            const results = await graph.invoke({
+            const parameters = {
               packageName,
               dataStreamName,
               rawSamples,
               currentPipeline,
-            });
+            };
+            const options = {
+              callbacks: [
+                new APMTracer({ projectName: langSmithOptions?.projectName ?? 'default' }, logger),
+                ...getLangSmithTracer({ ...langSmithOptions, logger }),
+              ],
+            };
+
+            const graph = await getCategorizationGraph(client, model);
+            const results = await graph.invoke(parameters, options);
+
             return res.ok({ body: CategorizationResponse.parse(results) });
           } catch (e) {
             return res.badRequest({ body: e });

x-pack/plugins/integration_assistant/server/routes/ecs_routes.ts

@@ -11,6 +11,8 @@ import {
   ActionsClientChatOpenAI,
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server/language_models';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { ECS_GRAPH_PATH, EcsMappingRequestBody, EcsMappingResponse } from '../../common';
 import { ROUTE_HANDLER_TIMEOUT } from '../constants';
 import { getEcsGraph } from '../graphs/ecs';
@@ -39,7 +41,7 @@ export function registerEcsRoutes(router: IRouter<IntegrationAssistantRouteHandl
         },
       },
       withAvailability(async (context, req, res): Promise<IKibanaResponse<EcsMappingResponse>> => {
-        const { packageName, dataStreamName, rawSamples, mapping } = req.body;
+        const { packageName, dataStreamName, rawSamples, mapping, langSmithOptions } = req.body;
         const { getStartServices, logger } = await context.integrationAssistant;
 
         const [, { actions: actionsPlugin }] = await getStartServices();
@@ -53,6 +55,7 @@ export function registerEcsRoutes(router: IRouter<IntegrationAssistantRouteHandl
 
           const abortSignal = getRequestAbortedSignal(req.events.aborted$);
           const isOpenAI = connector.actionTypeId === '.gen-ai';
+
           const llmClass = isOpenAI ? ActionsClientChatOpenAI : ActionsClientSimpleChatModel;
 
           const model = new llmClass({
@@ -67,22 +70,23 @@ export function registerEcsRoutes(router: IRouter<IntegrationAssistantRouteHandl
             streaming: false,
           });
 
-          const graph = await getEcsGraph(model);
+          const parameters = {
+            packageName,
+            dataStreamName,
+            rawSamples,
+            ...(mapping && { mapping }),
+          };
+
+          const options = {
+            callbacks: [
+              new APMTracer({ projectName: langSmithOptions?.projectName ?? 'default' }, logger),
+              ...getLangSmithTracer({ ...langSmithOptions, logger }),
+            ],
+          };
+
+          const graph = await getEcsGraph(model);
+          const results = await graph.invoke(parameters, options);
 
-          let results;
-          if (req.body?.mapping) {
-            results = await graph.invoke({
-              packageName,
-              dataStreamName,
-              rawSamples,
-              mapping,
-            });
-          } else
-            results = await graph.invoke({
-              packageName,
-              dataStreamName,
-              rawSamples,
-            });
           return res.ok({ body: EcsMappingResponse.parse(results) });
         } catch (e) {
           return res.badRequest({ body: e });

x-pack/plugins/integration_assistant/server/routes/related_routes.ts

@@ -11,6 +11,8 @@ import {
   ActionsClientChatOpenAI,
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server/language_models';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { RELATED_GRAPH_PATH, RelatedRequestBody, RelatedResponse } from '../../common';
 import { ROUTE_HANDLER_TIMEOUT } from '../constants';
 import { getRelatedGraph } from '../graphs/related';
@@ -39,7 +41,8 @@ export function registerRelatedRoutes(router: IRouter<IntegrationAssistantRouteH
         },
       },
       withAvailability(async (context, req, res): Promise<IKibanaResponse<RelatedResponse>> => {
-        const { packageName, dataStreamName, rawSamples, currentPipeline } = req.body;
+        const { packageName, dataStreamName, rawSamples, currentPipeline, langSmithOptions } =
+          req.body;
         const services = await context.resolve(['core']);
         const { client } = services.core.elasticsearch;
         const { getStartServices, logger } = await context.integrationAssistant;
@@ -68,13 +71,21 @@ export function registerRelatedRoutes(router: IRouter<IntegrationAssistantRouteH
             streaming: false,
           });
 
-          const graph = await getRelatedGraph(client, model);
-          const results = await graph.invoke({
+          const parameters = {
             packageName,
             dataStreamName,
             rawSamples,
             currentPipeline,
-          });
+          };
+          const options = {
+            callbacks: [
+              new APMTracer({ projectName: langSmithOptions?.projectName ?? 'default' }, logger),
+              ...getLangSmithTracer({ ...langSmithOptions, logger }),
+            ],
+          };
+
+          const graph = await getRelatedGraph(client, model);
+          const results = await graph.invoke(parameters, options);
           return res.ok({ body: RelatedResponse.parse(results) });
         } catch (e) {
           return res.badRequest({ body: e });

x-pack/plugins/integration_assistant/tsconfig.json

@@ -36,6 +36,7 @@
     "@kbn/licensing-plugin",
     "@kbn/core-http-request-handler-context-server",
     "@kbn/core-http-router-server-mocks",
-    "@kbn/core-http-server"
+    "@kbn/core-http-server",
+    "@kbn/kibana-utils-plugin"
   ]
 }