[Security Solution] AI Assistant - Turn streaming on by default (#180095)

This commit is contained in:
Steph Milovic 2024-04-05 10:10:36 -06:00 committed by GitHub
parent 29fe28539b
commit 04c21e6ca9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 119 additions and 25 deletions

View file

@ -15,5 +15,4 @@ export type AssistantFeatures = { [K in keyof typeof defaultAssistantFeatures]:
*/
export const defaultAssistantFeatures = Object.freeze({
assistantModelEvaluation: false,
assistantStreamingEnabled: false,
});

View file

@ -19,5 +19,4 @@ import { z } from 'zod';
export type GetCapabilitiesResponse = z.infer<typeof GetCapabilitiesResponse>;
export const GetCapabilitiesResponse = z.object({
assistantModelEvaluation: z.boolean(),
assistantStreamingEnabled: z.boolean(),
});

View file

@ -21,11 +21,8 @@ paths:
properties:
assistantModelEvaluation:
type: boolean
assistantStreamingEnabled:
type: boolean
required:
- assistantModelEvaluation
- assistantStreamingEnabled
'400':
description: Generic Error
content:

View file

@ -5,7 +5,15 @@
* 2.0.
*/
import { EuiFormRow, EuiLink, EuiTitle, EuiText, EuiHorizontalRule, EuiSpacer } from '@elastic/eui';
import {
EuiFormRow,
EuiLink,
EuiTitle,
EuiText,
EuiHorizontalRule,
EuiSpacer,
EuiSwitch,
} from '@elastic/eui';
import React, { useCallback, useMemo } from 'react';
import { HttpSetup } from '@kbn/core-http-browser';
@ -32,9 +40,11 @@ export interface ConversationSettingsProps {
conversationSettings: Record<string, Conversation>;
conversationsSettingsBulkActions: ConversationsBulkActions;
defaultConnector?: AIConnector;
assistantStreamingEnabled: boolean;
http: HttpSetup;
onSelectedConversationChange: (conversation?: Conversation) => void;
selectedConversation?: Conversation;
setAssistantStreamingEnabled: React.Dispatch<React.SetStateAction<boolean>>;
setConversationSettings: React.Dispatch<React.SetStateAction<Record<string, Conversation>>>;
setConversationsSettingsBulkActions: React.Dispatch<
React.SetStateAction<ConversationsBulkActions>
@ -48,12 +58,14 @@ export interface ConversationSettingsProps {
export const ConversationSettings: React.FC<ConversationSettingsProps> = React.memo(
({
allSystemPrompts,
assistantStreamingEnabled,
defaultConnector,
selectedConversation,
onSelectedConversationChange,
conversationSettings,
http,
isDisabled = false,
setAssistantStreamingEnabled,
setConversationSettings,
conversationsSettingsBulkActions,
setConversationsSettingsBulkActions,
@ -342,7 +354,6 @@ export const ConversationSettings: React.FC<ConversationSettingsProps> = React.m
setConversationsSettingsBulkActions,
]
);
return (
<>
<EuiTitle size={'s'}>
@ -418,6 +429,21 @@ export const ConversationSettings: React.FC<ConversationSettingsProps> = React.m
/>
</EuiFormRow>
)}
<EuiSpacer size="l" />
<EuiTitle size={'s'}>
<h2>{i18n.SETTINGS_ALL_TITLE}</h2>
</EuiTitle>
<EuiSpacer size="xs" />
<EuiText size={'s'}>{i18n.SETTINGS_ALL_DESCRIPTION}</EuiText>
<EuiHorizontalRule margin={'s'} />
<EuiFormRow fullWidth display="rowCompressed" label={i18n.STREAMING_TITLE}>
<EuiSwitch
label={<EuiText size="xs">{i18n.STREAMING_HELP_TEXT_TITLE}</EuiText>}
checked={assistantStreamingEnabled}
onChange={(e) => setAssistantStreamingEnabled(e.target.checked)}
compressed
/>
</EuiFormRow>
</>
);
}

View file

@ -21,6 +21,20 @@ export const SETTINGS_DESCRIPTION = i18n.translate(
}
);
export const SETTINGS_ALL_TITLE = i18n.translate(
'xpack.elasticAssistant.assistant.conversations.settings.settingsAllTitle',
{
defaultMessage: 'All conversations',
}
);
export const SETTINGS_ALL_DESCRIPTION = i18n.translate(
'xpack.elasticAssistant.assistant.conversations.settings.settingsAllDescription',
{
defaultMessage: 'These settings apply to all conversations.',
}
);
export const CONNECTOR_TITLE = i18n.translate(
'xpack.elasticAssistant.assistant.conversations.settings.connectorTitle',
{
@ -41,3 +55,17 @@ export const SETTINGS_PROMPT_HELP_TEXT_TITLE = i18n.translate(
defaultMessage: 'Context provided as part of every conversation',
}
);
export const STREAMING_TITLE = i18n.translate(
'xpack.elasticAssistant.assistant.conversations.settings.streamingTitle',
{
defaultMessage: 'Streaming',
}
);
export const STREAMING_HELP_TEXT_TITLE = i18n.translate(
'xpack.elasticAssistant.assistant.conversations.settings.streamingHelpTextTitle',
{
defaultMessage: 'Controls whether streaming responses are used in conversations',
}
);

View file

@ -97,6 +97,8 @@ export const AssistantSettings: React.FC<Props> = React.memo(
knowledgeBase,
quickPromptSettings,
systemPromptSettings,
assistantStreamingEnabled,
setUpdatedAssistantStreamingEnabled,
setUpdatedDefaultAllow,
setUpdatedDefaultAllowReplacement,
setUpdatedKnowledgeBaseSettings,
@ -299,6 +301,8 @@ export const AssistantSettings: React.FC<Props> = React.memo(
allSystemPrompts={systemPromptSettings}
selectedConversation={selectedConversation}
isDisabled={selectedConversation == null}
assistantStreamingEnabled={assistantStreamingEnabled}
setAssistantStreamingEnabled={setUpdatedAssistantStreamingEnabled}
onSelectedConversationChange={onHandleSelectedConversationChange}
http={http}
/>

View file

@ -35,10 +35,13 @@ const initialDefaultAllowReplacement = ['replacement1'];
const setAllQuickPromptsMock = jest.fn();
const setAllSystemPromptsMock = jest.fn();
const setDefaultAllowMock = jest.fn();
const setAssistantStreamingEnabled = jest.fn();
const setDefaultAllowReplacementMock = jest.fn();
const setKnowledgeBaseMock = jest.fn();
const reportAssistantSettingToggled = jest.fn();
const mockValues = {
assistantStreamingEnabled: true,
setAssistantStreamingEnabled,
assistantTelemetry: { reportAssistantSettingToggled },
allSystemPrompts: mockSystemPrompts,
allQuickPrompts: mockQuickPrompts,
@ -69,6 +72,7 @@ const updatedValues = {
isEnabledKnowledgeBase: false,
latestAlerts: DEFAULT_LATEST_ALERTS,
},
assistantStreamingEnabled: false,
};
jest.mock('../../../assistant_context', () => {
@ -95,6 +99,7 @@ describe('useSettingsUpdater', () => {
setUpdatedDefaultAllow,
setUpdatedDefaultAllowReplacement,
setUpdatedKnowledgeBaseSettings,
setUpdatedAssistantStreamingEnabled,
resetSettings,
} = result.current;
@ -105,6 +110,7 @@ describe('useSettingsUpdater', () => {
setUpdatedDefaultAllow(updatedValues.defaultAllow);
setUpdatedDefaultAllowReplacement(updatedValues.defaultAllowReplacement);
setUpdatedKnowledgeBaseSettings(updatedValues.knowledgeBase);
setUpdatedAssistantStreamingEnabled(updatedValues.assistantStreamingEnabled);
expect(result.current.conversationSettings).toEqual(updatedValues.conversations);
expect(result.current.quickPromptSettings).toEqual(updatedValues.allQuickPrompts);
@ -112,6 +118,9 @@ describe('useSettingsUpdater', () => {
expect(result.current.defaultAllow).toEqual(updatedValues.defaultAllow);
expect(result.current.defaultAllowReplacement).toEqual(updatedValues.defaultAllowReplacement);
expect(result.current.knowledgeBase).toEqual(updatedValues.knowledgeBase);
expect(result.current.assistantStreamingEnabled).toEqual(
updatedValues.assistantStreamingEnabled
);
resetSettings();
@ -121,6 +130,9 @@ describe('useSettingsUpdater', () => {
expect(result.current.defaultAllow).toEqual(mockValues.defaultAllow);
expect(result.current.defaultAllowReplacement).toEqual(mockValues.defaultAllowReplacement);
expect(result.current.knowledgeBase).toEqual(mockValues.knowledgeBase);
expect(result.current.assistantStreamingEnabled).toEqual(
mockValues.assistantStreamingEnabled
);
});
});

View file

@ -15,6 +15,7 @@ import {
} from '../../api/conversations/use_bulk_actions_conversations';
interface UseSettingsUpdater {
assistantStreamingEnabled: boolean;
conversationSettings: Record<string, Conversation>;
conversationsSettingsBulkActions: ConversationsBulkActions;
defaultAllow: string[];
@ -32,6 +33,7 @@ interface UseSettingsUpdater {
setUpdatedKnowledgeBaseSettings: React.Dispatch<React.SetStateAction<KnowledgeBaseConfig>>;
setUpdatedQuickPromptSettings: React.Dispatch<React.SetStateAction<QuickPrompt[]>>;
setUpdatedSystemPromptSettings: React.Dispatch<React.SetStateAction<Prompt[]>>;
setUpdatedAssistantStreamingEnabled: React.Dispatch<React.SetStateAction<boolean>>;
saveSettings: () => Promise<boolean>;
}
@ -50,6 +52,8 @@ export const useSettingsUpdater = (
setAllSystemPrompts,
setDefaultAllow,
setDefaultAllowReplacement,
assistantStreamingEnabled,
setAssistantStreamingEnabled,
setKnowledgeBase,
http,
toasts,
@ -73,6 +77,8 @@ export const useSettingsUpdater = (
const [updatedDefaultAllow, setUpdatedDefaultAllow] = useState<string[]>(defaultAllow);
const [updatedDefaultAllowReplacement, setUpdatedDefaultAllowReplacement] =
useState<string[]>(defaultAllowReplacement);
const [updatedAssistantStreamingEnabled, setUpdatedAssistantStreamingEnabled] =
useState<boolean>(assistantStreamingEnabled);
// Knowledge Base
const [updatedKnowledgeBaseSettings, setUpdatedKnowledgeBaseSettings] =
useState<KnowledgeBaseConfig>(knowledgeBase);
@ -85,12 +91,14 @@ export const useSettingsUpdater = (
setConversationsSettingsBulkActions({});
setUpdatedQuickPromptSettings(allQuickPrompts);
setUpdatedKnowledgeBaseSettings(knowledgeBase);
setUpdatedAssistantStreamingEnabled(assistantStreamingEnabled);
setUpdatedSystemPromptSettings(allSystemPrompts);
setUpdatedDefaultAllow(defaultAllow);
setUpdatedDefaultAllowReplacement(defaultAllowReplacement);
}, [
allQuickPrompts,
allSystemPrompts,
assistantStreamingEnabled,
conversations,
defaultAllow,
defaultAllowReplacement,
@ -116,7 +124,9 @@ export const useSettingsUpdater = (
knowledgeBase.isEnabledKnowledgeBase !== updatedKnowledgeBaseSettings.isEnabledKnowledgeBase;
const didUpdateRAGAlerts =
knowledgeBase.isEnabledRAGAlerts !== updatedKnowledgeBaseSettings.isEnabledRAGAlerts;
if (didUpdateKnowledgeBase || didUpdateRAGAlerts) {
const didUpdateAssistantStreamingEnabled =
assistantStreamingEnabled !== updatedAssistantStreamingEnabled;
if (didUpdateKnowledgeBase || didUpdateRAGAlerts || didUpdateAssistantStreamingEnabled) {
assistantTelemetry?.reportAssistantSettingToggled({
...(didUpdateKnowledgeBase
? { isEnabledKnowledgeBase: updatedKnowledgeBaseSettings.isEnabledKnowledgeBase }
@ -124,8 +134,12 @@ export const useSettingsUpdater = (
...(didUpdateRAGAlerts
? { isEnabledRAGAlerts: updatedKnowledgeBaseSettings.isEnabledRAGAlerts }
: {}),
...(didUpdateAssistantStreamingEnabled
? { assistantStreamingEnabled: updatedAssistantStreamingEnabled }
: {}),
});
}
setAssistantStreamingEnabled(updatedAssistantStreamingEnabled);
setKnowledgeBase(updatedKnowledgeBaseSettings);
setDefaultAllow(updatedDefaultAllow);
setDefaultAllowReplacement(updatedDefaultAllowReplacement);
@ -141,7 +155,10 @@ export const useSettingsUpdater = (
toasts,
knowledgeBase.isEnabledKnowledgeBase,
knowledgeBase.isEnabledRAGAlerts,
updatedAssistantStreamingEnabled,
updatedKnowledgeBaseSettings,
assistantStreamingEnabled,
setAssistantStreamingEnabled,
setKnowledgeBase,
setDefaultAllow,
updatedDefaultAllow,
@ -156,6 +173,7 @@ export const useSettingsUpdater = (
defaultAllow: updatedDefaultAllow,
defaultAllowReplacement: updatedDefaultAllowReplacement,
knowledgeBase: updatedKnowledgeBaseSettings,
assistantStreamingEnabled: updatedAssistantStreamingEnabled,
quickPromptSettings: updatedQuickPromptSettings,
resetSettings,
systemPromptSettings: updatedSystemPromptSettings,
@ -163,6 +181,7 @@ export const useSettingsUpdater = (
setUpdatedDefaultAllow,
setUpdatedDefaultAllowReplacement,
setUpdatedKnowledgeBaseSettings,
setUpdatedAssistantStreamingEnabled,
setUpdatedQuickPromptSettings,
setUpdatedSystemPromptSettings,
setConversationSettings,

View file

@ -12,6 +12,7 @@ export const QUICK_PROMPT_LOCAL_STORAGE_KEY = 'quickPrompts';
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';
/** The default `n` latest alerts, ordered by risk score, sent as context to the assistant */
export const DEFAULT_LATEST_ALERTS = 20;

View file

@ -33,6 +33,7 @@ import {
KNOWLEDGE_BASE_LOCAL_STORAGE_KEY,
LAST_CONVERSATION_TITLE_LOCAL_STORAGE_KEY,
QUICK_PROMPT_LOCAL_STORAGE_KEY,
STREAMING_LOCAL_STORAGE_KEY,
SYSTEM_PROMPT_LOCAL_STORAGE_KEY,
} from './constants';
import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings';
@ -133,6 +134,7 @@ export interface UseAssistantContext {
setAllSystemPrompts: React.Dispatch<React.SetStateAction<Prompt[] | undefined>>;
setDefaultAllow: React.Dispatch<React.SetStateAction<string[]>>;
setDefaultAllowReplacement: React.Dispatch<React.SetStateAction<string[]>>;
setAssistantStreamingEnabled: React.Dispatch<React.SetStateAction<boolean | undefined>>;
setKnowledgeBase: React.Dispatch<React.SetStateAction<KnowledgeBaseConfig | undefined>>;
setLastConversationTitle: React.Dispatch<React.SetStateAction<string | undefined>>;
setSelectedSettingsTab: React.Dispatch<React.SetStateAction<SettingsTabs>>;
@ -197,6 +199,15 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
DEFAULT_KNOWLEDGE_BASE_SETTINGS
);
/**
* Local storage for streaming configuration, prefixed by assistant nameSpace
*/
// can be undefined from localStorage, if not defined, default to true
const [localStorageStreaming, setLocalStorageStreaming] = useLocalStorage<boolean>(
`${nameSpace}.${STREAMING_LOCAL_STORAGE_KEY}`,
true
);
/**
* Prompt contexts are used to provide components a way to register and make their data available to the assistant.
*/
@ -253,7 +264,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
// Fetch assistant capabilities
const { data: capabilities } = useCapabilities({ http, toasts });
const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } =
const { assistantModelEvaluation: modelEvaluatorEnabled } =
capabilities ?? defaultAssistantFeatures;
const value = useMemo(
@ -261,7 +272,6 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
actionTypeRegistry,
alertsIndexPattern,
assistantAvailability,
assistantStreamingEnabled,
assistantTelemetry,
augmentMessageCodeBlocks,
allQuickPrompts: localStorageQuickPrompts ?? [],
@ -283,6 +293,9 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
nameSpace,
registerPromptContext,
selectedSettingsTab,
// can be undefined from localStorage, if not defined, default to true
assistantStreamingEnabled: localStorageStreaming ?? true,
setAssistantStreamingEnabled: setLocalStorageStreaming,
setAllQuickPrompts: setLocalStorageQuickPrompts,
setAllSystemPrompts: setLocalStorageSystemPrompts,
setDefaultAllow,
@ -302,7 +315,6 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
actionTypeRegistry,
alertsIndexPattern,
assistantAvailability,
assistantStreamingEnabled,
assistantTelemetry,
augmentMessageCodeBlocks,
localStorageQuickPrompts,
@ -324,6 +336,8 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
nameSpace,
registerPromptContext,
selectedSettingsTab,
localStorageStreaming,
setLocalStorageStreaming,
setLocalStorageQuickPrompts,
setLocalStorageSystemPrompts,
setDefaultAllow,

View file

@ -79,6 +79,7 @@ export interface AssistantTelemetry {
reportAssistantSettingToggled: (params: {
isEnabledKnowledgeBase?: boolean;
isEnabledRAGAlerts?: boolean;
assistantStreamingEnabled?: boolean;
}) => void;
}

View file

@ -57,7 +57,6 @@ export const getCapabilitiesRoute = (router: IRouter<ElasticAssistantRequestHand
logger,
});
const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName);
return response.ok({ body: registeredFeatures });
} catch (err) {
const error = transformError(err);

View file

@ -46,7 +46,6 @@ describe('Post Evaluate Route', () => {
it('returns a 404 if evaluate feature is not registered', async () => {
context.elasticAssistant.getRegisteredFeatures.mockReturnValueOnce({
assistantModelEvaluation: false,
assistantStreamingEnabled: false,
});
const response = await server.inject(

View file

@ -54,7 +54,6 @@ describe('AppContextService', () => {
appContextService.start(mockAppContext);
appContextService.registerFeatures('super', {
assistantModelEvaluation: true,
assistantStreamingEnabled: true,
});
appContextService.stop();
@ -104,7 +103,6 @@ describe('AppContextService', () => {
const pluginName = 'pluginName';
const features: AssistantFeatures = {
assistantModelEvaluation: true,
assistantStreamingEnabled: true,
};
appContextService.start(mockAppContext);
@ -119,12 +117,10 @@ describe('AppContextService', () => {
const pluginOne = 'plugin1';
const featuresOne: AssistantFeatures = {
assistantModelEvaluation: true,
assistantStreamingEnabled: false,
};
const pluginTwo = 'plugin2';
const featuresTwo: AssistantFeatures = {
assistantModelEvaluation: false,
assistantStreamingEnabled: true,
};
appContextService.start(mockAppContext);
@ -139,11 +135,9 @@ describe('AppContextService', () => {
const pluginName = 'pluginName';
const featuresOne: AssistantFeatures = {
assistantModelEvaluation: true,
assistantStreamingEnabled: false,
};
const featuresTwo: AssistantFeatures = {
assistantModelEvaluation: false,
assistantStreamingEnabled: true,
};
appContextService.start(mockAppContext);

View file

@ -46,11 +46,6 @@ export const allowedExperimentalValues = Object.freeze({
*/
extendedRuleExecutionLoggingEnabled: false,
/**
* Enables streaming for Security AI Assistant - non-langchain only (knowledge base off)
*/
assistantStreamingEnabled: false,
/**
* Enables the SOC trends timerange and stats on D&R page
*/

View file

@ -99,5 +99,12 @@ export const assistantSettingToggledEvent: TelemetryEvent = {
optional: true,
},
},
assistantStreamingEnabled: {
type: 'boolean',
_meta: {
description: 'Is streaming enabled',
optional: true,
},
},
},
};

View file

@ -28,6 +28,7 @@ export interface ReportAssistantQuickPromptParams {
export interface ReportAssistantSettingToggledParams {
isEnabledKnowledgeBase?: boolean;
isEnabledRAGAlerts?: boolean;
assistantStreamingEnabled?: boolean;
}
export type ReportAssistantTelemetryEventParams =

View file

@ -559,7 +559,6 @@ export class Plugin implements ISecuritySolutionPlugin {
plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools());
plugins.elasticAssistant.registerFeatures(APP_UI_ID, {
assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation,
assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled,
});
if (this.lists && plugins.taskManager && plugins.fleet) {