[Automatic Import V2] Add langsmith tracing

x-pack/platform/plugins/shared/automatic_import_v2/common/index.ts

@@ -39,4 +39,5 @@ export type {
   AllIntegrationsResponseIntegration,
   OriginalSource,
   OriginalSourceType,
+  LangSmithOptions,
 } from './model/common_attributes.gen';

x-pack/platform/plugins/shared/automatic_import_v2/common/model/api/data_streams/data_stream.gen.ts

@@ -17,7 +17,7 @@
 import { z } from '@kbn/zod';
 
 import { NonEmptyString } from '../../primitive.gen';
-import { OriginalSource } from '../../common_attributes.gen';
+import { OriginalSource, LangSmithOptions } from '../../common_attributes.gen';
 
 export type DeleteDataStreamRequestParams = z.infer<typeof DeleteDataStreamRequestParams>;
 export const DeleteDataStreamRequestParams = z.object({
@@ -78,6 +78,10 @@ export const UploadSamplesToDataStreamRequestBody = z.object({
    * The original source of the samples
    */
   originalSource: OriginalSource,
+  /**
+   * The LangSmith tracing options
+   */
+  langSmithOptions: LangSmithOptions.optional(),
 });
 export type UploadSamplesToDataStreamRequestBodyInput = z.input<
   typeof UploadSamplesToDataStreamRequestBody

x-pack/platform/plugins/shared/automatic_import_v2/common/model/api/data_streams/data_stream.schema.yaml

@@ -76,6 +76,9 @@ paths:
                 originalSource:
                   description: The original source of the samples
                   $ref: "../../../model/common_attributes.schema.yaml#/components/schemas/OriginalSource"
+                langSmithOptions:
+                  description: The LangSmith tracing options
+                  $ref: "../../../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
       responses:
         200:
           description: The samples are uploaded successfully.

x-pack/platform/plugins/shared/automatic_import_v2/common/model/api/integrations/integration.gen.ts

@@ -18,6 +18,7 @@ import { z } from '@kbn/zod';
 
 import { NonEmptyString } from '../../primitive.gen';
 import {
+  LangSmithOptions,
   DataStream,
   AllIntegrationsResponseIntegration,
   InputType,
@@ -31,6 +32,10 @@ export const ApproveIntegrationRequest = z
      * The version of the integration
      */
     version: NonEmptyString,
+    /**
+     * The LangSmith tracing options
+     */
+    langSmithOptions: LangSmithOptions.optional(),
   })
   .strict();
 
@@ -76,6 +81,10 @@ export const CreateAutoImportIntegrationRequestBody = z
      * The description of the integration
      */
     description: NonEmptyString,
+    /**
+     * The LangSmith tracing options
+     */
+    langSmithOptions: LangSmithOptions.optional(),
     /**
      * The logo of the integration
      */
@@ -166,6 +175,10 @@ export const UpdateAutoImportIntegrationRequestBody = z
      * Integration logo image blob
      */
     logo: NonEmptyString.optional(),
+    /**
+     * The LangSmith tracing options
+     */
+    langSmithOptions: LangSmithOptions.optional(),
     /**
      * The data streams of the integration
      */

x-pack/platform/plugins/shared/automatic_import_v2/common/model/api/integrations/integration.schema.yaml

@@ -38,6 +38,9 @@ paths:
                 description:
                   description: The description of the integration
                   $ref: "../../../model/primitive.schema.yaml#/components/schemas/NonEmptyString"
+                langSmithOptions:
+                  description: The LangSmith tracing options
+                  $ref: "../../../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
                 logo:
                   description: The logo of the integration
                   $ref: "../../../model/primitive.schema.yaml#/components/schemas/NonEmptyString"
@@ -123,6 +126,9 @@ paths:
                   description: Integration logo image blob
                   $ref: "../../../model/primitive.schema.yaml#/components/schemas/NonEmptyString"
                   format: binary
+                langSmithOptions:
+                  description: The LangSmith tracing options
+                  $ref: "../../../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
                 dataStreams:
                   description: The data streams of the integration
                   type: array
@@ -273,3 +279,6 @@ components:
         version:
           description: The version of the integration
           $ref: "../../../model/primitive.schema.yaml#/components/schemas/NonEmptyString"
+        langSmithOptions:
+          description: The LangSmith tracing options
+          $ref: "../../../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"

x-pack/platform/plugins/shared/automatic_import_v2/common/model/common_attributes.gen.ts

@@ -217,3 +217,20 @@ export const AllIntegrationsResponseIntegration = z.object({
    */
   status: TaskStatus,
 });
+
+/**
+ * The LangSmith options object.
+ */
+export type LangSmithOptions = z.infer<typeof LangSmithOptions>;
+export const LangSmithOptions = z
+  .object({
+    /**
+     * The project name.
+     */
+    projectName: z.string(),
+    /**
+     * The apiKey to use for tracing.
+     */
+    apiKey: z.string(),
+  })
+  .strict();

x-pack/platform/plugins/shared/automatic_import_v2/common/model/common_attributes.schema.yaml

@@ -198,3 +198,18 @@ components:
             status:
               description: The status of the integration
               $ref: "#/components/schemas/TaskStatus"
+
+    LangSmithOptions:
+      type: object
+      description: The LangSmith options object.
+      required:
+        - projectName
+        - apiKey
+      additionalProperties: false
+      properties:
+        projectName:
+          type: string
+          description: The project name.
+        apiKey:
+          type: string
+          description: The apiKey to use for tracing.

x-pack/platform/plugins/shared/automatic_import_v2/kibana.jsonc

@@ -22,7 +22,8 @@
       "inference",
       "security",
       "features",
-      "taskManager"
+      "taskManager",
+      "kibanaUtils"
     ],
     "requiredBundles": [
       "esUiShared"

x-pack/platform/plugins/shared/automatic_import_v2/moon.yml

@@ -52,6 +52,8 @@ dependsOn:
   - '@kbn/core-http-server-mocks'
   - '@kbn/shared-ux-router'
   - '@kbn/react-query'
+  - '@kbn/elastic-assistant'
+  - '@kbn/kibana-utils-plugin'
 tags:
   - plugin
   - prod

x-pack/platform/plugins/shared/automatic_import_v2/public/common/index.ts

@@ -30,4 +30,7 @@ export {
 
 export type { CreateUpdateIntegrationRequest } from './lib/api';
 
+export { getLangSmithOptions } from './lib/lang_smith';
+export type { LangSmithOptions } from './lib/lang_smith';
+
 export { generateId } from './lib/helper_functions';

x-pack/platform/plugins/shared/automatic_import_v2/public/common/lib/api.ts

@@ -11,6 +11,8 @@ import type {
   GetAutoImportIntegrationResponse,
 } from '../../../common/model/api/integrations/integration.gen';
 import type { DataStream } from '../../../common/model/common_attributes.gen';
+import { getLangSmithOptions } from './lang_smith';
+import type { LangSmithOptions } from './lang_smith';
 
 export const FLEET_PACKAGES_PATH = `/api/fleet/epm/packages`;
 export const AUTOMATIC_IMPORT_INTEGRATIONS_PATH = `/api/automatic_import_v2/integrations`;
@@ -49,6 +51,7 @@ export interface CreateUpdateIntegrationRequest {
   description: string;
   logo?: string;
   dataStreams?: DataStream[];
+  langSmithOptions?: LangSmithOptions;
 }
 
 export const createIntegration = async ({
@@ -58,7 +61,10 @@ export const createIntegration = async ({
 }: RequestDeps & CreateUpdateIntegrationRequest): Promise<CreateAutoImportIntegrationResponse> =>
   http.put<CreateAutoImportIntegrationResponse>(AUTOMATIC_IMPORT_INTEGRATIONS_PATH, {
     version: '1',
-    body: JSON.stringify(body),
+    body: JSON.stringify({
+      ...body,
+      langSmithOptions: body.langSmithOptions ?? getLangSmithOptions(),
+    }),
     signal: abortSignal,
   });
 

x-pack/platform/plugins/shared/automatic_import_v2/public/common/lib/lang_smith.ts

@@ -0,0 +1,47 @@
+/*
+ * 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';
+
+export interface LangSmithOptions {
+  projectName: string;
+  apiKey: string;
+}
+
+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;
+  }
+
+  const { langSmithProject, langSmithApiKey } = sessionStorageTraceOptions;
+  if (!langSmithProject || !langSmithApiKey) {
+    return;
+  }
+
+  return {
+    projectName: langSmithProject,
+    apiKey: langSmithApiKey,
+  };
+};

x-pack/platform/plugins/shared/automatic_import_v2/server/routes/integrations_route.ts

@@ -153,7 +153,15 @@ const createIntegrationRoute = (
       async (context, request, response) => {
         const { automaticImportService, getCurrentUser, esClient } =
           await context.automaticImportv2;
-        const { integrationId, title, logo, description, connectorId, dataStreams } = request.body;
+        const {
+          integrationId,
+          title,
+          logo,
+          description,
+          connectorId,
+          dataStreams,
+          langSmithOptions,
+        } = request.body;
         try {
           const authenticatedUser = await getCurrentUser();
 
@@ -183,6 +191,7 @@ const createIntegrationRoute = (
                     dataStreamParams,
                     esClient,
                     connectorId,
+                    langSmithOptions,
                   },
                   request
                 )

x-pack/platform/plugins/shared/automatic_import_v2/server/routes/types.ts

@@ -8,6 +8,11 @@
 import type { AuthenticatedUser, ElasticsearchClient } from '@kbn/core/server';
 import type { InputType } from '../../common';
 
+export interface LangSmithOptions {
+  projectName: string;
+  apiKey: string;
+}
+
 export interface CreateUpdateIntegrationParams {
   integrationParams: IntegrationParams;
   authenticatedUser: AuthenticatedUser;
@@ -31,6 +36,10 @@ export interface CreateDataStreamParams {
    * Inference connector to use when the background task runs.
    */
   connectorId: string;
+  /**
+   * Optional LangSmith tracing options to propagate to background tasks.
+   */
+  langSmithOptions?: LangSmithOptions;
   /**
    * Minimal set of auth headers required to reconstruct a scoped client
    * as the original user inside the Task Manager runner.

x-pack/platform/plugins/shared/automatic_import_v2/server/services/agents/agent_service.ts

@@ -21,6 +21,7 @@ import {
 } from '../../agents/tools';
 import type { AutomaticImportSamplesIndexService } from '../samples_index/index_service';
 import { INGEST_PIPELINE_GENERATOR_PROMPT } from '../../agents/prompts';
+import type { LangSmithOptions } from '../../routes/types';
 
 export class AgentService {
   private logger: Logger;
@@ -49,7 +50,8 @@ export class AgentService {
     integrationId: string,
     dataStreamId: string,
     esClient: ElasticsearchClient,
-    model: InferenceChatModel
+    model: InferenceChatModel,
+    langSmithOptions?: LangSmithOptions
   ) {
     this.logger.debug(
       `invokeAutomaticImportAgent: Invoking automatic import agent for integration ${integrationId} and data stream ${dataStreamId}`
@@ -100,16 +102,14 @@ export class AgentService {
       subagents: [logsAnalyzerSubAgent, pipelineGeneratorSubAgent, textToEcsSubAgent],
     });
 
-    const traceOptions = {
-      tracers: [
-        ...getLangSmithTracer({
-          // TODO: Get apiKey from config
-          apiKey: 'apiKey',
-          projectName: 'projectName',
-          logger: this.logger,
-        }),
-      ],
-    };
+    const langSmithTracers =
+      langSmithOptions?.apiKey && langSmithOptions?.projectName
+        ? getLangSmithTracer({
+            apiKey: langSmithOptions.apiKey,
+            projectName: langSmithOptions.projectName,
+            logger: this.logger,
+          })
+        : [];
 
     const result = await automaticImportAgent.invoke(
       {
@@ -121,7 +121,7 @@ export class AgentService {
         ],
       },
       {
-        callbacks: [...(traceOptions.tracers ?? [])],
+        callbacks: [...langSmithTracers],
         runName: 'automatic_import_agent',
         tags: ['automatic_import_agent'],
       }

x-pack/platform/plugins/shared/automatic_import_v2/server/services/automatic_import_service.ts

@@ -233,13 +233,14 @@ export class AutomaticImportService {
     request: KibanaRequest
   ): Promise<void> {
     assert(this.savedObjectService, 'Saved Objects service not initialized.');
-    const { authenticatedUser, dataStreamParams, connectorId } = params;
+    const { authenticatedUser, dataStreamParams, connectorId, langSmithOptions } = params;
 
     // Schedule the data stream creation background task
     const dataStreamTaskParams: DataStreamTaskParams = {
       integrationId: dataStreamParams.integrationId,
       dataStreamId: dataStreamParams.dataStreamId,
       connectorId,
+      ...(langSmithOptions ? { langSmithOptions } : {}),
     };
     const { taskId } = await this.taskManagerService.scheduleDataStreamCreationTask(
       dataStreamTaskParams,