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

# Backport

This will backport the following commits from `main` to `8.15`:
- [[Security GenAI] Use AI setting to set langsmith tracing to the
Integration Assistant
(#187466)](https://github.com/elastic/kibana/pull/187466)

<!--- Backport version: 9.4.3 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Sergi
Massaneda","email":"sergi.massaneda@elastic.co"},"sourceCommit":{"committedDate":"2024-07-08T18:14:57Z","message":"[Security
GenAI] Use AI setting to set langsmith tracing to the Integration
Assistant (#187466)\n\n## Summary\r\n\r\nEnables tracing Langchain
invocations in the integrations assistant\r\nusing the Langsmith
settings stored by the Security AI Settings.\r\nThe evaluation settings
tab is still under an experimental flag, to
see\r\nit:\r\n\r\n```\r\nxpack.securitySolution.enableExperimental:
['assistantModelEvaluation']\r\n```\r\n\r\n### Screenshots\r\n\r\n<img
width=\"1317\"
alt=\"Settings\"\r\nsrc=\"6aed1ef6-3750-4259-9fe2-b8bf1aed5504\">\r\n\r\nAfter
one execution of the integration assistant:\r\n\r\n<img width=\"1240\"
alt=\"langsmith\"\r\nsrc=\"dd3dd99c-7c83-4a35-95b2-789e7a341031\">\r\n\r\n---------\r\n\r\nCo-authored-by:
kibanamachine
<42973632+kibanamachine@users.noreply.github.com>\r\nCo-authored-by:
Elastic Machine
<elasticmachine@users.noreply.github.com>","sha":"92099b277dc5b1448d14994d280674611ca9e261","branchLabelMapping":{"^v8.16.0$":"main","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","Team:Security
Generative AI","v8.15.0","v8.16.0"],"title":"[Security GenAI] Use AI
setting to set langsmith tracing to the Integration
Assistant","number":187466,"url":"https://github.com/elastic/kibana/pull/187466","mergeCommit":{"message":"[Security
GenAI] Use AI setting to set langsmith tracing to the Integration
Assistant (#187466)\n\n## Summary\r\n\r\nEnables tracing Langchain
invocations in the integrations assistant\r\nusing the Langsmith
settings stored by the Security AI Settings.\r\nThe evaluation settings
tab is still under an experimental flag, to
see\r\nit:\r\n\r\n```\r\nxpack.securitySolution.enableExperimental:
['assistantModelEvaluation']\r\n```\r\n\r\n### Screenshots\r\n\r\n<img
width=\"1317\"
alt=\"Settings\"\r\nsrc=\"6aed1ef6-3750-4259-9fe2-b8bf1aed5504\">\r\n\r\nAfter
one execution of the integration assistant:\r\n\r\n<img width=\"1240\"
alt=\"langsmith\"\r\nsrc=\"dd3dd99c-7c83-4a35-95b2-789e7a341031\">\r\n\r\n---------\r\n\r\nCo-authored-by:
kibanamachine
<42973632+kibanamachine@users.noreply.github.com>\r\nCo-authored-by:
Elastic Machine
<elasticmachine@users.noreply.github.com>","sha":"92099b277dc5b1448d14994d280674611ca9e261"}},"sourceBranch":"main","suggestedTargetBranches":["8.15"],"targetPullRequestStates":[{"branch":"8.15","label":"v8.15.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"main","label":"v8.16.0","branchLabelMappingKey":"^v8.16.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/187466","number":187466,"mergeCommit":{"message":"[Security
GenAI] Use AI setting to set langsmith tracing to the Integration
Assistant (#187466)\n\n## Summary\r\n\r\nEnables tracing Langchain
invocations in the integrations assistant\r\nusing the Langsmith
settings stored by the Security AI Settings.\r\nThe evaluation settings
tab is still under an experimental flag, to
see\r\nit:\r\n\r\n```\r\nxpack.securitySolution.enableExperimental:
['assistantModelEvaluation']\r\n```\r\n\r\n### Screenshots\r\n\r\n<img
width=\"1317\"
alt=\"Settings\"\r\nsrc=\"6aed1ef6-3750-4259-9fe2-b8bf1aed5504\">\r\n\r\nAfter
one execution of the integration assistant:\r\n\r\n<img width=\"1240\"
alt=\"langsmith\"\r\nsrc=\"dd3dd99c-7c83-4a35-95b2-789e7a341031\">\r\n\r\n---------\r\n\r\nCo-authored-by:
kibanamachine
<42973632+kibanamachine@users.noreply.github.com>\r\nCo-authored-by:
Elastic Machine
<elasticmachine@users.noreply.github.com>","sha":"92099b277dc5b1448d14994d280674611ca9e261"}}]}]
BACKPORT-->

Co-authored-by: Sergi Massaneda <sergi.massaneda@elastic.co>

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"
   ]
 }