[Security AI Assistant] Adds ability to specify LangSmith config and APM URL for tracing in cloud environments

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts

@@ -41,6 +41,8 @@ export const ExecuteConnectorRequestBody = z.object({
   isEnabledRAGAlerts: z.boolean().optional(),
   replacements: Replacements,
   size: z.number().optional(),
+  langSmithProject: z.string().optional(),
+  langSmithApiKey: z.string().optional(),
 });
 export type ExecuteConnectorRequestBodyInput = z.input<typeof ExecuteConnectorRequestBody>;
 

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml

@@ -61,6 +61,10 @@ paths:
                   $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacements'
                 size:
                   type: number
+                langSmithProject:
+                  type: string
+                langSmithApiKey:
+                  type: string
       responses:
         '200':
           description: Successful static response

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx

@@ -10,6 +10,7 @@ import { IHttpFetchError } from '@kbn/core-http-browser';
 import { ApiConfig, Replacements } from '@kbn/elastic-assistant-common';
 import { API_ERROR } from '../translations';
 import { getOptionalRequestParams } from '../helpers';
+import { TraceOptions } from '../types';
 export * from './conversations';
 
 export interface FetchConnectorExecuteAction {
@@ -26,6 +27,7 @@ export interface FetchConnectorExecuteAction {
   replacements: Replacements;
   signal?: AbortSignal | undefined;
   size?: number;
+  traceOptions?: TraceOptions;
 }
 
 export interface FetchConnectorExecuteResponse {
@@ -52,6 +54,7 @@ export const fetchConnectorExecuteAction = async ({
   apiConfig,
   signal,
   size,
+  traceOptions,
 }: FetchConnectorExecuteAction): Promise<FetchConnectorExecuteResponse> => {
   const isStream =
     assistantStreamingEnabled &&
@@ -77,6 +80,8 @@ export const fetchConnectorExecuteAction = async ({
     replacements,
     isEnabledKnowledgeBase,
     isEnabledRAGAlerts,
+    langSmithProject: traceOptions?.langSmithProject,
+    langSmithApiKey: traceOptions?.langSmithApiKey,
     ...optionalRequestParams,
   };
 

x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx

@@ -17,7 +17,7 @@ import { useConnectorSetup } from '../connectorland/connector_setup';
 import { UseQueryResult } from '@tanstack/react-query';
 import { WELCOME_CONVERSATION_TITLE } from './use_conversation/translations';
 
-import { useLocalStorage } from 'react-use';
+import { useLocalStorage, useSessionStorage } from 'react-use';
 import { PromptEditor } from './prompt_editor';
 import { QuickPrompts } from './quick_prompts/quick_prompts';
 import { mockAssistantAvailability, TestProviders } from '../mock/test_providers/test_providers';
@@ -108,16 +108,23 @@ describe('Assistant', () => {
   });
 
   let persistToLocalStorage: jest.Mock;
+  let persistToSessionStorage: jest.Mock;
 
   beforeEach(() => {
     jest.clearAllMocks();
     persistToLocalStorage = jest.fn();
+    persistToSessionStorage = jest.fn();
 
     jest
       .mocked(useLocalStorage)
       .mockReturnValue([undefined, persistToLocalStorage] as unknown as ReturnType<
         typeof useLocalStorage
       >);
+    jest
+      .mocked(useSessionStorage)
+      .mockReturnValue([undefined, persistToSessionStorage] as unknown as ReturnType<
+        typeof useSessionStorage
+      >);
   });
 
   describe('persistent storage', () => {

x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/evaluation_settings/evaluation_settings.tsx

@@ -52,7 +52,8 @@ interface Props {
  * Evaluation Settings -- development-only feature for evaluating models
  */
 export const EvaluationSettings: React.FC<Props> = React.memo(({ onEvaluationSettingsChange }) => {
-  const { actionTypeRegistry, basePath, http } = useAssistantContext();
+  const { actionTypeRegistry, basePath, http, setTraceOptions, traceOptions } =
+    useAssistantContext();
   const { data: connectors } = useLoadConnectors({ http });
   const {
     data: evalResponse,
@@ -92,7 +93,27 @@ export const EvaluationSettings: React.FC<Props> = React.memo(({ onEvaluationSet
     },
     [setOutputIndex]
   );
-  // Dataset
+  /** Trace Options **/
+  const [showTraceOptions, setShowTraceOptions] = useState(false);
+  const onApmUrlChange = useCallback(
+    (e) => {
+      setTraceOptions({ ...traceOptions, apmUrl: e.target.value });
+    },
+    [setTraceOptions, traceOptions]
+  );
+  const onLangSmithProjectChange = useCallback(
+    (e) => {
+      setTraceOptions({ ...traceOptions, langSmithProject: e.target.value });
+    },
+    [setTraceOptions, traceOptions]
+  );
+  const onLangSmithApiKeyChange = useCallback(
+    (e) => {
+      setTraceOptions({ ...traceOptions, langSmithApiKey: e.target.value });
+    },
+    [setTraceOptions, traceOptions]
+  );
+  /** Dataset **/
   const [useLangSmithDataset, setUseLangSmithDataset] = useState(true);
   const datasetToggleButton = useMemo(() => {
     return (
@@ -423,6 +444,59 @@ export const EvaluationSettings: React.FC<Props> = React.memo(({ onEvaluationSet
             aria-label="evaluation-output-index-textfield"
           />
         </EuiFormRow>
+        <EuiText
+          size={'xs'}
+          css={css`
+            margin-top: 16px;
+          `}
+        >
+          <EuiLink color={'primary'} onClick={() => setShowTraceOptions(!showTraceOptions)}>
+            {i18n.SHOW_TRACE_OPTIONS}
+          </EuiLink>
+        </EuiText>
+        {showTraceOptions && (
+          <>
+            <EuiFormRow
+              display="rowCompressed"
+              label={i18n.APM_URL_LABEL}
+              fullWidth
+              helpText={i18n.APM_URL_DESCRIPTION}
+              css={css`
+                margin-top: 16px;
+              `}
+            >
+              <EuiFieldText
+                value={traceOptions.apmUrl}
+                onChange={onApmUrlChange}
+                aria-label="apm-url-textfield"
+              />
+            </EuiFormRow>
+            <EuiFormRow
+              display="rowCompressed"
+              label={i18n.LANGSMITH_PROJECT_LABEL}
+              fullWidth
+              helpText={i18n.LANGSMITH_PROJECT_DESCRIPTION}
+            >
+              <EuiFieldText
+                value={traceOptions.langSmithProject}
+                onChange={onLangSmithProjectChange}
+                aria-label="langsmith-project-textfield"
+              />
+            </EuiFormRow>
+            <EuiFormRow
+              display="rowCompressed"
+              label={i18n.LANGSMITH_API_KEY_LABEL}
+              fullWidth
+              helpText={i18n.LANGSMITH_API_KEY_DESCRIPTION}
+            >
+              <EuiFieldText
+                value={traceOptions.langSmithApiKey}
+                onChange={onLangSmithApiKeyChange}
+                aria-label="langsmith-api-key-textfield"
+              />
+            </EuiFormRow>
+          </>
+        )}
       </EuiAccordion>
       <EuiHorizontalRule margin={'s'} />
       {/* Prediction Details*/}

x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/evaluation_settings/translations.ts

@@ -193,6 +193,57 @@ export const EVALUATOR_OUTPUT_INDEX_DESCRIPTION = i18n.translate(
   }
 );
 
+export const SHOW_TRACE_OPTIONS = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.showTraceOptionsLabel',
+  {
+    defaultMessage: 'Show Trace Options (for internal use only)',
+  }
+);
+
+export const APM_URL_LABEL = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.apmUrlLabel',
+  {
+    defaultMessage: 'APM URL',
+  }
+);
+
+export const APM_URL_DESCRIPTION = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.apmUrlDescription',
+  {
+    defaultMessage:
+      'URL for the Kibana APM app. Used to link to APM traces for evaluation results. Defaults to "$\\{basePath\\}/app/apm"',
+  }
+);
+
+export const LANGSMITH_PROJECT_LABEL = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.langSmithProjectLabel',
+  {
+    defaultMessage: 'LangSmith Project',
+  }
+);
+
+export const LANGSMITH_PROJECT_DESCRIPTION = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.langSmithProjectDescription',
+  {
+    defaultMessage: 'LangSmith Project to write traces to',
+  }
+);
+
+export const LANGSMITH_API_KEY_LABEL = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.langSmithApiKeyLabel',
+  {
+    defaultMessage: 'LangSmith API Key',
+  }
+);
+
+export const LANGSMITH_API_KEY_DESCRIPTION = i18n.translate(
+  'xpack.elasticAssistant.assistant.settings.evaluationSettings.langSmithApiKeyDescription',
+  {
+    defaultMessage:
+      'API Key for writing traces to LangSmith. Stored in Session Storage. Close tab to clear session.',
+  }
+);
+
 export const EVALUATOR_DATASET_LABEL = i18n.translate(
   'xpack.elasticAssistant.assistant.settings.evaluationSettings.evaluatorDatasetLabel',
   {

x-pack/packages/kbn-elastic-assistant/impl/assistant/types.ts

@@ -21,3 +21,9 @@ export interface KnowledgeBaseConfig {
   isEnabledKnowledgeBase: boolean;
   latestAlerts: number;
 }
+
+export interface TraceOptions {
+  apmUrl: string;
+  langSmithProject: string;
+  langSmithApiKey: string;
+}

x-pack/packages/kbn-elastic-assistant/impl/assistant/use_send_message/index.tsx

@@ -38,6 +38,7 @@ export const useSendMessage = (): UseSendMessage => {
     defaultAllow,
     defaultAllowReplacement,
     knowledgeBase,
+    traceOptions,
   } = useAssistantContext();
   const [isLoading, setIsLoading] = useState(false);
   const abortController = useRef(new AbortController());
@@ -60,19 +61,21 @@ export const useSendMessage = (): UseSendMessage => {
           replacements,
           signal: abortController.current.signal,
           size: knowledgeBase.latestAlerts,
+          traceOptions,
         });
       } finally {
         setIsLoading(false);
       }
     },
     [
-      alertsIndexPattern,
-      assistantStreamingEnabled,
-      defaultAllow,
-      defaultAllowReplacement,
       knowledgeBase.isEnabledRAGAlerts,
       knowledgeBase.isEnabledKnowledgeBase,
       knowledgeBase.latestAlerts,
+      alertsIndexPattern,
+      defaultAllow,
+      defaultAllowReplacement,
+      assistantStreamingEnabled,
+      traceOptions,
     ]
   );
 

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx

@@ -13,6 +13,7 @@ export const SYSTEM_PROMPT_LOCAL_STORAGE_KEY = 'systemPrompts';
 export const LAST_CONVERSATION_TITLE_LOCAL_STORAGE_KEY = 'lastConversationTitle';
 export const KNOWLEDGE_BASE_LOCAL_STORAGE_KEY = 'knowledgeBase';
 export const STREAMING_LOCAL_STORAGE_KEY = 'streaming';
+export const TRACE_OPTIONS_SESSION_STORAGE_KEY = 'traceOptions';
 
 /** The default `n` latest alerts, ordered by risk score, sent as context to the assistant */
 export const DEFAULT_LATEST_ALERTS = 20;

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx

@@ -13,6 +13,7 @@ import { TestProviders } from '../mock/test_providers/test_providers';
 
 jest.mock('react-use', () => ({
   useLocalStorage: jest.fn().mockReturnValue(['456', jest.fn()]),
+  useSessionStorage: jest.fn().mockReturnValue(['456', jest.fn()]),
 }));
 
 describe('AssistantContext', () => {

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx

@@ -11,7 +11,7 @@ import { omit, uniq } from 'lodash/fp';
 import React, { useCallback, useMemo, useState } from 'react';
 import type { IToasts } from '@kbn/core-notifications-browser';
 import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public';
-import { useLocalStorage } from 'react-use';
+import { useLocalStorage, useSessionStorage } from 'react-use';
 import type { DocLinksStart } from '@kbn/core-doc-links-browser';
 import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
 import { updatePromptContexts } from './helpers';
@@ -25,7 +25,7 @@ import { DEFAULT_ASSISTANT_TITLE } from '../assistant/translations';
 import { CodeBlockDetails } from '../assistant/use_conversation/helpers';
 import { PromptContextTemplate } from '../assistant/prompt_context/types';
 import { QuickPrompt } from '../assistant/quick_prompts/types';
-import type { KnowledgeBaseConfig, Prompt } from '../assistant/types';
+import { KnowledgeBaseConfig, Prompt, TraceOptions } from '../assistant/types';
 import { BASE_SYSTEM_PROMPTS } from '../content/prompts/system';
 import {
   DEFAULT_ASSISTANT_NAMESPACE,
@@ -35,6 +35,7 @@ import {
   QUICK_PROMPT_LOCAL_STORAGE_KEY,
   STREAMING_LOCAL_STORAGE_KEY,
   SYSTEM_PROMPT_LOCAL_STORAGE_KEY,
+  TRACE_OPTIONS_SESSION_STORAGE_KEY,
 } from './constants';
 import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings';
 import { AssistantAvailability, AssistantTelemetry } from './types';
@@ -140,8 +141,14 @@ export interface UseAssistantContext {
   setSelectedSettingsTab: React.Dispatch<React.SetStateAction<SettingsTabs>>;
   setShowAssistantOverlay: (showAssistantOverlay: ShowAssistantOverlay) => void;
   showAssistantOverlay: ShowAssistantOverlay;
+  setTraceOptions: (traceOptions: {
+    apmUrl: string;
+    langSmithProject: string;
+    langSmithApiKey: string;
+  }) => void;
   title: string;
   toasts: IToasts | undefined;
+  traceOptions: TraceOptions;
   unRegisterPromptContext: UnRegisterPromptContext;
 }
 
@@ -172,6 +179,20 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
   title = DEFAULT_ASSISTANT_TITLE,
   toasts,
 }) => {
+  /**
+   * Session storage for traceOptions, including APM URL and LangSmith Project/API Key
+   */
+  const defaultTraceOptions: TraceOptions = {
+    apmUrl: `${http.basePath.serverBasePath}/app/apm`,
+    langSmithProject: '',
+    langSmithApiKey: '',
+  };
+  const [sessionStorageTraceOptions = defaultTraceOptions, setSessionStorageTraceOptions] =
+    useSessionStorage<TraceOptions>(
+      `${nameSpace}.${TRACE_OPTIONS_SESSION_STORAGE_KEY}`,
+      defaultTraceOptions
+    );
+
   /**
    * Local storage for all quick prompts, prefixed by assistant nameSpace
    */
@@ -303,9 +324,11 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       setKnowledgeBase: setLocalStorageKnowledgeBase,
       setSelectedSettingsTab,
       setShowAssistantOverlay,
+      setTraceOptions: setSessionStorageTraceOptions,
       showAssistantOverlay,
       title,
       toasts,
+      traceOptions: sessionStorageTraceOptions,
       unRegisterPromptContext,
       getLastConversationTitle,
       setLastConversationTitle: setLocalStorageLastConversationTitle,
@@ -343,9 +366,11 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       setDefaultAllow,
       setDefaultAllowReplacement,
       setLocalStorageKnowledgeBase,
+      setSessionStorageTraceOptions,
       showAssistantOverlay,
       title,
       toasts,
+      sessionStorageTraceOptions,
       unRegisterPromptContext,
       getLastConversationTitle,
       setLocalStorageLastConversationTitle,