[Security AI Assistant] Adds ability to specify LangSmith config and …

…APM URL for tracing in cloud environments (#180227) ## Summary While we wait for #178304, this is a PR for allowing users to specify their LangSmith config for tracing in cloud environments by only storing them in session storage. This is also behind an experimental feature flag and must be enabled with the `assistantModelEvaluation` flag ala: ``` xpack.securitySolution.enableExperimental: [ 'assistantModelEvaluation'] ``` ~Note I: `xpack.securitySolution.enableExperimental` should be allowlisted in cloud, but I have manually enabled via source for initial testing.~ Note II: I have verified the above is configurable on cloud deployments 👍 The new `traceOptions` are stored with the `elasticAssistantDefault.traceOptions` key, and the following keys: ``` { apmUrl : "${basepath}/app/apm" langSmithApiKey: "🫣" langSmithProject: "Cloud Testing" } ``` The `langSmithApiKey` and `langSmithProject` are then sent along with the request to `/actions/connector/{connectorId}/_execute`, and a new `LangChainTracer` is created using the values. The tracing infrastructue was already in place for evaluation, so no other changes were necessary. The `apmUrl` value is now used for the `View APM trace for message` action, so if you have set up a remote APM server, you can now link directly to that instance from the message. A basic UI was added for these fields under the `Run` step of the Evaluation Settings. No need to save or run an evaluation once entering. Fields are immediately stored in session storage upon entry. <p align="center"> <img width="500" src="https://github.com/elastic/kibana/assets/2946766/02445b24-9d4b-40a9-bbad-f261ec098faa" /> </p> ### Test Instructions Click on the [latest Kibana Buildkite build](https://buildkite.com/elastic/kibana-pull-request/builds/201924#annotation-cloud), go to the `ci:cloud-deploy` cluster (grabbing creds from vault), then set a LangChain Project/API key in the above UI, then make a request to the LLM and verify the trace is collected in the LangSmith UI: > [!NOTE] > Only LangChain codepaths can be traced to LangSmith, so you must ensure LangChain is enabled by either turning on the Knowledge Base or enabling the Alert tools. The former can't be done in default `ci:cloud-deploy` deployments as they only have a 1GB ML nodes, so it is easiest to just turn on the Alert tools. <p align="center"> <img width="500" src="https://github.com/elastic/kibana/assets/2946766/b7c6747c-3314-44e2-8d58-f9d2bfdda687" /> </p> ### Checklist Delete any items that are not applicable to this PR. - [X] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md)
elastic · Apr 8, 2024 · 16c0ab8 · 16c0ab8
1 parent 2af3219
commit 16c0ab8
Show file tree

Hide file tree

Showing 16 changed files with 228 additions and 23 deletions.
diff --git a/...sistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts b/...sistant-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>;
 

diff --git a/...nt-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml b/...nt-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

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx b/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,
   };
 

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx b/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', () => {

diff --git a/...kbn-elastic-assistant/impl/assistant/settings/evaluation_settings/evaluation_settings.tsx b/...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*/}

diff --git a/...ackages/kbn-elastic-assistant/impl/assistant/settings/evaluation_settings/translations.ts b/...ackages/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',
   {

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/types.ts b/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;
+}
diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_send_message/index.tsx b/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,
     ]
   );
 

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx b/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;

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx b/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', () => {

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/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,