github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-04-24 09:48:58 -04:00
Code Issues Projects Releases Packages Wiki Activity

[Obs AI Assistant] Refactor hooks, recall on every message (#171965)

Browse source 
Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
This commit is contained in:
Dario Gieselaar 2023-11-29 17:27:06 +01:00 • committed by GitHub
parent 3437e6d878
commit 450b2c5a2d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
34 changed files with 2827 additions and 1726 deletions
Download patch file Download diff file Expand all files Collapse all files

2
x-pack/plugins/observability_ai_assistant/jest.config.js
View file

@ -10,4 +10,6 @@ module.exports = {
rootDir: '../../..',
roots: ['<rootDir>/x-pack/plugins/observability_ai_assistant'],
setupFiles: ['<rootDir>/x-pack/plugins/observability_ai_assistant/.storybook/jest_setup.js'],
collectCoverage: true,
coverageReporters: ['html'],
};

35
x-pack/plugins/observability_ai_assistant/public/components/action_menu_item/action_menu_item.tsx
View file

@ -6,19 +6,15 @@
*/
import { EuiFlexGroup, EuiFlexItem, EuiHeaderLink, EuiLoadingSpinner } from '@elastic/eui';
import { i18n } from '@kbn/i18n';
import React, { useState } from 'react';
import React, { useMemo, useState } from 'react';
import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
import { useAbortableAsync } from '../../hooks/use_abortable_async';
import { useConversation } from '../../hooks/use_conversation';
import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
import { EMPTY_CONVERSATION_TITLE } from '../../i18n';
import { AssistantAvatar } from '../assistant_avatar';
import { ChatFlyout } from '../chat/chat_flyout';
export function ObservabilityAIAssistantActionMenuItem() {
const service = useObservabilityAIAssistant();
const connectors = useGenAIConnectors();
const [isOpen, setIsOpen] = useState(false);
@ -32,14 +28,7 @@ export function ObservabilityAIAssistantActionMenuItem() {
[service, isOpen]
);
const [conversationId, setConversationId] = useState<string>();
const { conversation, displayedMessages, setDisplayedMessages, save, saveTitle } =
useConversation({
conversationId,
connectorId: connectors.selectedConnector,
chatService: chatService.value,
});
const initialMessages = useMemo(() => [], []);
if (!service.isEnabled()) {
return null;
@ -72,26 +61,12 @@ export function ObservabilityAIAssistantActionMenuItem() {
{chatService.value ? (
<ObservabilityAIAssistantChatServiceProvider value={chatService.value}>
<ChatFlyout
initialTitle=""
initialMessages={initialMessages}
isOpen={isOpen}
title={conversation.value?.conversation.title ?? EMPTY_CONVERSATION_TITLE}
messages={displayedMessages}
conversationId={conversationId}
startedFrom="appTopNavbar"
onClose={() => {
setIsOpen(() => false);
}}
onChatComplete={(messages) => {
save(messages)
.then((nextConversation) => {
setConversationId(nextConversation.conversation.id);
})
.catch(() => {});
}}
onChatUpdate={(nextMessages) => {
setDisplayedMessages(nextMessages);
}}
onChatTitleSave={(newTitle) => {
saveTitle(newTitle);
setIsOpen(false);
}}
/>
</ObservabilityAIAssistantChatServiceProvider>

6
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.stories.tsx
View file

@ -21,8 +21,8 @@ const meta: ComponentMeta<typeof Component> = {
export default meta;
const defaultProps: ComponentStoryObj<typeof Component> = {
args: {
title: 'My Conversation',
messages: [
initialTitle: 'My Conversation',
initialMessages: [
getAssistantSetupMessage({ contexts: [] }),
{
'@timestamp': new Date().toISOString(),
@ -64,8 +64,6 @@ const defaultProps: ComponentStoryObj<typeof Component> = {
currentUser: {
username: 'elastic',
},
onChatUpdate: () => {},
onChatComplete: () => {},
},
render: (props) => {
return (

196
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx
View file

@ -5,9 +5,8 @@
* 2.0.
*/
import React, { useEffect, useRef, useState } from 'react';
import { flatten, last } from 'lodash';
import {
EuiCallOut,
EuiFlexGroup,
EuiFlexItem,
EuiHorizontalRule,
@ -16,21 +15,26 @@ import {
EuiSpacer,
} from '@elastic/eui';
import { css } from '@emotion/css';
import { euiThemeVars } from '@kbn/ui-theme';
import type { AuthenticatedUser } from '@kbn/security-plugin/common';
import type { Message } from '../../../common/types';
import { euiThemeVars } from '@kbn/ui-theme';
import React, { useEffect, useRef, useState } from 'react';
import { i18n } from '@kbn/i18n';
import { Conversation, Message, MessageRole } from '../../../common/types';
import { ChatState } from '../../hooks/use_chat';
import { useConversation } from '../../hooks/use_conversation';
import type { UseGenAIConnectorsResult } from '../../hooks/use_genai_connectors';
import type { UseKnowledgeBaseResult } from '../../hooks/use_knowledge_base';
import { useTimeline } from '../../hooks/use_timeline';
import { useLicense } from '../../hooks/use_license';
import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
import { ExperimentalFeatureBanner } from './experimental_feature_banner';
import { InitialSetupPanel } from './initial_setup_panel';
import { IncorrectLicensePanel } from './incorrect_license_panel';
import { StartedFrom } from '../../utils/get_timeline_items_from_conversation';
import { ChatHeader } from './chat_header';
import { ChatPromptEditor } from './chat_prompt_editor';
import { ChatTimeline } from './chat_timeline';
import { StartedFrom } from '../../utils/get_timeline_items_from_conversation';
import { ExperimentalFeatureBanner } from './experimental_feature_banner';
import { IncorrectLicensePanel } from './incorrect_license_panel';
import { InitialSetupPanel } from './initial_setup_panel';
import { ChatActionClickType } from './types';
import { EMPTY_CONVERSATION_TITLE } from '../../i18n';
const timelineClassName = css`
overflow-y: auto;
@ -45,48 +49,45 @@ const incorrectLicenseContainer = css`
padding: ${euiThemeVars.euiPanelPaddingModifiers.paddingMedium};
`;
const chatBodyContainerClassNameWithError = css`
align-self: center;
`;
export function ChatBody({
title,
loading,
messages,
initialTitle,
initialMessages,
initialConversationId,
connectors,
knowledgeBase,
connectorsManagementHref,
modelsManagementHref,
conversationId,
currentUser,
startedFrom,
onChatUpdate,
onChatComplete,
onSaveTitle,
onConversationUpdate,
}: {
title: string;
loading: boolean;
messages: Message[];
initialTitle?: string;
initialMessages?: Message[];
initialConversationId?: string;
connectors: UseGenAIConnectorsResult;
knowledgeBase: UseKnowledgeBaseResult;
connectorsManagementHref: string;
modelsManagementHref: string;
conversationId?: string;
currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
startedFrom?: StartedFrom;
onChatUpdate: (messages: Message[]) => void;
onChatComplete: (messages: Message[]) => void;
onSaveTitle: (title: string) => void;
onConversationUpdate: (conversation: Conversation) => void;
}) {
const license = useLicense();
const hasCorrectLicense = license?.hasAtLeast('enterprise');
const chatService = useObservabilityAIAssistantChatService();
const timeline = useTimeline({
const { conversation, messages, next, state, stop, saveTitle } = useConversation({
initialConversationId,
initialMessages,
initialTitle,
chatService,
connectors,
currentUser,
messages,
startedFrom,
onChatUpdate,
onChatComplete,
connectorId: connectors.selectedConnector,
onConversationUpdate,
});
const timelineContainerRef = useRef<HTMLDivElement | null>(null);
@ -94,7 +95,10 @@ export function ChatBody({
let footer: React.ReactNode;
const isLoading = Boolean(
connectors.loading || knowledgeBase.status.loading || last(flatten(timeline.items))?.loading
connectors.loading ||
knowledgeBase.status.loading ||
state === ChatState.Loading ||
conversation.loading
);
const containerClassName = css`
@ -139,12 +143,12 @@ export function ChatBody({
});
const handleCopyConversation = () => {
const content = JSON.stringify({ title, messages });
const content = JSON.stringify({ title: initialTitle, messages });
navigator.clipboard?.writeText(content || '');
};
if (!hasCorrectLicense && !conversationId) {
if (!hasCorrectLicense && !initialConversationId) {
footer = (
<>
<EuiFlexItem grow className={incorrectLicenseContainer}>
@ -155,19 +159,29 @@ export function ChatBody({
</EuiFlexItem>
<EuiFlexItem grow={false}>
<EuiPanel hasBorder={false} hasShadow={false} paddingSize="m">
<ChatPromptEditor loading={isLoading} disabled onSubmit={timeline.onSubmit} />
<ChatPromptEditor
loading={isLoading}
disabled
onSubmit={(message) => {
next(messages.concat(message));
}}
/>
<EuiSpacer size="s" />
</EuiPanel>
</EuiFlexItem>
</>
);
} else if (connectors.loading || knowledgeBase.status.loading) {
} else if (
connectors.loading ||
knowledgeBase.status.loading ||
(!conversation.value && conversation.loading)
) {
footer = (
<EuiFlexItem className={loadingSpinnerContainerClassName}>
<EuiLoadingSpinner />
</EuiFlexItem>
);
} else if (connectors.connectors?.length === 0 && !conversationId) {
} else if (connectors.connectors?.length === 0 && !initialConversationId) {
footer = (
<InitialSetupPanel
connectors={connectors}
@ -183,15 +197,47 @@ export function ChatBody({
<div ref={timelineContainerRef}>
<EuiPanel hasBorder={false} hasShadow={false} paddingSize="m">
<ChatTimeline
items={timeline.items}
startedFrom={startedFrom}
messages={messages}
knowledgeBase={knowledgeBase}
onEdit={timeline.onEdit}
onFeedback={timeline.onFeedback}
onRegenerate={timeline.onRegenerate}
onStopGenerating={timeline.onStopGenerating}
chatService={chatService}
currentUser={currentUser}
chatState={state}
hasConnector={!!connectors.connectors?.length}
onEdit={(editedMessage, newMessage) => {
const indexOf = messages.indexOf(editedMessage);
next(messages.slice(0, indexOf).concat(newMessage));
}}
onFeedback={(message, feedback) => {}}
onRegenerate={(message) => {
const indexOf = messages.indexOf(message);
next(messages.slice(0, indexOf));
}}
onStopGenerating={() => {
stop();
}}
onActionClick={(payload) => {
setStickToBottom(true);
return timeline.onActionClick(payload);
switch (payload.type) {
case ChatActionClickType.executeEsqlQuery:
next(
messages.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'execute_query',
arguments: JSON.stringify({
query: payload.query,
}),
trigger: MessageRole.User,
},
},
})
);
break;
}
}}
/>
</EuiPanel>
@ -207,7 +253,7 @@ export function ChatBody({
disabled={!connectors.selectedConnector || !hasCorrectLicense}
onSubmit={(message) => {
setStickToBottom(true);
return timeline.onSubmit(message);
return next(messages.concat(message));
}}
/>
<EuiSpacer size="s" />
@ -217,6 +263,33 @@ export function ChatBody({
);
}
if (conversation.error) {
return (
<EuiFlexGroup
direction="column"
gutterSize="none"
className={containerClassName}
justifyContent="center"
>
<EuiFlexItem grow={false} className={chatBodyContainerClassNameWithError}>
<EuiCallOut
color="danger"
title={i18n.translate('xpack.observabilityAiAssistant.couldNotFindConversationTitle', {
defaultMessage: 'Conversation not found',
})}
iconType="warning"
>
{i18n.translate('xpack.observabilityAiAssistant.couldNotFindConversationContent', {
defaultMessage:
'Could not find a conversation with id {conversationId}. Make sure the conversation exists and you have access to it.',
values: { conversationId: initialConversationId },
})}
</EuiCallOut>
</EuiFlexItem>
</EuiFlexGroup>
);
}
return (
<EuiFlexGroup direction="column" gutterSize="none" className={containerClassName}>
{connectors.selectedConnector ? (
@ -224,20 +297,45 @@ export function ChatBody({
<ExperimentalFeatureBanner />
</EuiFlexItem>
) : null}
<EuiFlexItem
grow={false}
className={conversation.error ? chatBodyContainerClassNameWithError : undefined}
>
{conversation.error ? (
<EuiCallOut
color="danger"
title={i18n.translate('xpack.observabilityAiAssistant.couldNotFindConversationTitle', {
defaultMessage: 'Conversation not found',
})}
iconType="warning"
>
{i18n.translate('xpack.observabilityAiAssistant.couldNotFindConversationContent', {
defaultMessage:
'Could not find a conversation with id {conversationId}. Make sure the conversation exists and you have access to it.',
values: { conversationId: initialConversationId },
})}
</EuiCallOut>
) : null}
</EuiFlexItem>
<EuiFlexItem grow={false}>
<ChatHeader
connectors={connectors}
conversationId={conversationId}
conversationId={
conversation.value?.conversation && 'id' in conversation.value.conversation
? conversation.value.conversation.id
: undefined
}
connectorsManagementHref={connectorsManagementHref}
modelsManagementHref={modelsManagementHref}
knowledgeBase={knowledgeBase}
licenseInvalid={!hasCorrectLicense && !conversationId}
loading={loading}
licenseInvalid={!hasCorrectLicense && !initialConversationId}
loading={isLoading}
startedFrom={startedFrom}
title={title}
title={conversation.value?.conversation.title || initialTitle || EMPTY_CONVERSATION_TITLE}
onCopyConversation={handleCopyConversation}
onSaveTitle={onSaveTitle}
onSaveTitle={(newTitle) => {
saveTitle(newTitle);
}}
/>
</EuiFlexItem>
<EuiFlexItem grow={false}>

6
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_consolidated_items.tsx
View file

@ -120,12 +120,12 @@ export function ChatConsolidatedItems({
key={index}
{...item}
onFeedbackClick={(feedback) => {
onFeedback(item, feedback);
onFeedback(item.message, feedback);
}}
onRegenerateClick={() => {
onRegenerate(item);
onRegenerate(item.message);
}}
onEditSubmit={(message) => onEditSubmit(item, message)}
onEditSubmit={(message) => onEditSubmit(item.message, message)}
onStopGeneratingClick={onStopGenerating}
onActionClick={onActionClick}
/>

7
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx
View file

@ -29,13 +29,10 @@ const Template: ComponentStory<typeof Component> = (props: ChatFlyoutProps) => {
const defaultProps: ChatFlyoutProps = {
isOpen: true,
title: 'How is this working',
messages: [getAssistantSetupMessage({ contexts: [] })],
initialTitle: 'How is this working',
initialMessages: [getAssistantSetupMessage({ contexts: [] })],
startedFrom: 'appTopNavbar',
onClose: () => {},
onChatComplete: () => {},
onChatTitleSave: () => {},
onChatUpdate: () => {},
};
export const ChatFlyout = Template.bind({});

42
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.tsx
View file

@ -7,7 +7,7 @@
import { EuiFlexGroup, EuiFlexItem, EuiFlyout, EuiLink, EuiPanel, useEuiTheme } from '@elastic/eui';
import { css } from '@emotion/css';
import { i18n } from '@kbn/i18n';
import React from 'react';
import React, { useState } from 'react';
import type { Message } from '../../../common/types';
import { useCurrentUser } from '../../hooks/use_current_user';
import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
@ -28,25 +28,17 @@ const bodyClassName = css`
`;
export function ChatFlyout({
title,
messages,
conversationId,
initialTitle,
initialMessages,
onClose,
isOpen,
startedFrom,
onClose,
onChatUpdate,
onChatComplete,
onChatTitleSave,
}: {
title: string;
messages: Message[];
conversationId?: string;
initialTitle: string;
initialMessages: Message[];
isOpen: boolean;
startedFrom: StartedFrom;
onClose: () => void;
onChatUpdate: (messages: Message[]) => void;
onChatComplete: (messages: Message[]) => void;
onChatTitleSave: (title: string) => void;
}) {
const { euiTheme } = useEuiTheme();
const {
@ -61,6 +53,8 @@ export function ChatFlyout({
const knowledgeBase = useKnowledgeBase();
const [conversationId, setConversationId] = useState<string | undefined>(undefined);
return isOpen ? (
<EuiFlyout onClose={onClose}>
<EuiFlexGroup
@ -101,28 +95,16 @@ export function ChatFlyout({
</EuiFlexItem>
<EuiFlexItem grow className={bodyClassName}>
<ChatBody
loading={false}
connectors={connectors}
title={title}
messages={messages}
initialTitle={initialTitle}
initialMessages={initialMessages}
currentUser={currentUser}
connectorsManagementHref={getConnectorsManagementHref(http)}
modelsManagementHref={getModelsManagementHref(http)}
conversationId={conversationId}
knowledgeBase={knowledgeBase}
startedFrom={startedFrom}
onChatUpdate={(nextMessages) => {
if (onChatUpdate) {
onChatUpdate(nextMessages);
}
}}
onChatComplete={(nextMessages) => {
if (onChatComplete) {
onChatComplete(nextMessages);
}
}}
onSaveTitle={(newTitle) => {
onChatTitleSave(newTitle);
onConversationUpdate={(conversation) => {
setConversationId(conversation.conversation.id);
}}
/>
</EuiFlexItem>

7
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_header.tsx
View file

@ -90,7 +90,12 @@ export function ChatHeader({
{ defaultMessage: 'Edit conversation' }
)}
editModeProps={{ inputProps: { inputRef } }}
isReadOnly={!connectors.selectedConnector || licenseInvalid || !Boolean(onSaveTitle)}
isReadOnly={
!conversationId ||
!connectors.selectedConnector ||
licenseInvalid ||
!Boolean(onSaveTitle)
}
onSave={onSaveTitle}
/>
) : null}

7
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item.tsx
View file

@ -27,7 +27,7 @@ import { FailedToLoadResponse } from '../message_panel/failed_to_load_response';
import { ChatActionClickHandler } from './types';
export interface ChatItemProps extends ChatTimelineItem {
onEditSubmit: (message: Message) => Promise<void>;
onEditSubmit: (message: Message) => void;
onFeedbackClick: (feedback: Feedback) => void;
onRegenerateClick: () => void;
onStopGeneratingClick: () => void;
@ -66,13 +66,14 @@ const noPanelMessageClassName = css`
export function ChatItem({
actions: { canCopy, canEdit, canGiveFeedback, canRegenerate },
display: { collapsed },
message: {
message: { function_call: functionCall, role },
},
content,
currentUser,
element,
error,
function_call: functionCall,
loading,
role,
title,
onEditSubmit,
onFeedbackClick,

2
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx
View file

@ -22,7 +22,7 @@ interface Props {
| undefined;
loading: boolean;
editing: boolean;
onSubmit: (message: Message) => Promise<void>;
onSubmit: (message: Message) => void;
onActionClick: ChatActionClickHandler;
}
export function ChatItemContentInlinePromptEditor({

16
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx
View file

@ -5,20 +5,20 @@
* 2.0.
*/
import React, { useCallback, useEffect, useRef, useState } from 'react';
import {
EuiButtonEmpty,
EuiButtonIcon,
EuiFlexGroup,
EuiFlexItem,
EuiFocusTrap,
EuiPanel,
EuiSpacer,
EuiTextArea,
keys,
EuiFocusTrap,
} from '@elastic/eui';
import { i18n } from '@kbn/i18n';
import { CodeEditor } from '@kbn/kibana-react-plugin/public';
import React, { useCallback, useEffect, useRef, useState } from 'react';
import { MessageRole, type Message } from '../../../common';
import { useJsonEditorModel } from '../../hooks/use_json_editor_model';
import { FunctionListPopover } from './function_list_popover';
@ -30,7 +30,7 @@ export interface ChatPromptEditorProps {
initialSelectedFunctionName?: string;
initialFunctionPayload?: string;
trigger?: MessageRole;
onSubmit: (message: Message) => Promise<void>;
onSubmit: (message: Message) => void;
}
export function ChatPromptEditor({
@ -216,7 +216,10 @@ export function ChatPromptEditor({
{selectedFunctionName ? (
<EuiPanel borderRadius="none" color="subdued" hasShadow={false} paddingSize="xs">
<CodeEditor
aria-label="payloadEditor"
aria-label={i18n.translate(
'xpack.observabilityAiAssistant.chatPromptEditor.codeEditor.payloadEditorLabel',
{ defaultMessage: 'payloadEditor' }
)}
data-test-subj="observabilityAiAssistantChatPromptEditorCodeEditor"
fullWidth
height={functionEditorLineCount > 8 ? '200px' : '120px'}
@ -284,7 +287,10 @@ export function ChatPromptEditor({
<EuiSpacer size="xl" />
<EuiButtonIcon
data-test-subj="observabilityAiAssistantChatPromptEditorButtonIcon"
aria-label="Submit"
aria-label={i18n.translate(
'xpack.observabilityAiAssistant.chatPromptEditor.euiButtonIcon.submitLabel',
{ defaultMessage: 'Submit' }
)}
disabled={selectedFunctionName ? false : !prompt?.trim() || loading || disabled}
display={
selectedFunctionName ? (functionPayload ? 'fill' : 'base') : prompt ? 'fill' : 'base'

105
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.stories.tsx
View file

@ -9,11 +9,13 @@ import { EuiButton, EuiSpacer } from '@elastic/eui';
import { ComponentStory } from '@storybook/react';
import React, { ComponentProps, useState } from 'react';
import { MessageRole } from '../../../common';
import { ChatState } from '../../hooks/use_chat';
import { ObservabilityAIAssistantChatService } from '../../types';
import {
buildAssistantChatItem,
buildChatInitItem,
buildFunctionChatItem,
buildUserChatItem,
buildAssistantMessage,
buildFunctionResponseMessage,
buildSystemMessage,
buildUserMessage,
} from '../../utils/builders';
import { ChatTimeline as Component, ChatTimelineProps } from './chat_timeline';
@ -30,17 +32,17 @@ export default {
};
const Template: ComponentStory<typeof Component> = (props: ChatTimelineProps) => {
const [count, setCount] = useState(props.items.length - 1);
const [count, setCount] = useState(props.messages.length - 1);
return (
<>
<Component {...props} items={props.items.filter((_, index) => index <= count)} />
<Component {...props} messages={props.messages.filter((_, index) => index <= count)} />
<EuiSpacer />
<EuiButton
data-test-subj="observabilityAiAssistantTemplateAddMessageButton"
onClick={() => setCount(count >= 0 && count < props.items.length - 1 ? count + 1 : 0)}
onClick={() => setCount(count >= 0 && count < props.messages.length - 1 ? count + 1 : 0)}
>
Add message
</EuiButton>
@ -61,13 +63,23 @@ const defaultProps: ComponentProps<typeof Component> = {
installError: undefined,
install: async () => {},
},
items: [
buildChatInitItem(),
buildUserChatItem(),
buildAssistantChatItem(),
buildUserChatItem({ content: 'How does it work?' }),
buildAssistantChatItem({
content: `The way functions work depends on whether we are talking about mathematical functions or programming functions. Let's explore both:
chatService: {
hasRenderFunction: () => false,
} as unknown as ObservabilityAIAssistantChatService,
chatState: ChatState.Ready,
hasConnector: true,
currentUser: {
full_name: 'John Doe',
username: 'johndoe',
},
messages: [
buildSystemMessage(),
buildUserMessage(),
buildAssistantMessage(),
buildUserMessage({ message: { content: 'How does it work?' } }),
buildAssistantMessage({
message: {
content: `The way functions work depends on whether we are talking about mathematical functions or programming functions. Let's explore both:
Mathematical Functions:
In mathematics, a function maps input values to corresponding output values based on a specific rule or expression. The general process of how a mathematical function works can be summarized as follows:
@ -78,54 +90,33 @@ const defaultProps: ComponentProps<typeof Component> = {
Step 3: Output - After processing the input, the function produces an output value, denoted as 'f(x)' or 'y'. This output represents the dependent variable and is the result of applying the function's rule to the input.
Step 4: Uniqueness - A well-defined mathematical function ensures that each input value corresponds to exactly one output value. In other words, the function should yield the same output for the same input whenever it is called.`,
}),
buildUserChatItem({
content: 'Can you execute a function?',
}),
buildAssistantChatItem({
content: 'Sure, I can do that.',
title: 'suggested a function',
function_call: {
name: 'a_function',
arguments: '{ "foo": "bar" }',
trigger: MessageRole.Assistant,
},
actions: {
canEdit: false,
canCopy: true,
canGiveFeedback: true,
canRegenerate: true,
},
}),
buildFunctionChatItem({
content: '{ "message": "The arguments are wrong" }',
error: new Error(),
actions: {
canRegenerate: false,
canEdit: true,
canGiveFeedback: false,
canCopy: true,
buildUserMessage({
message: { content: 'Can you execute a function?' },
}),
buildAssistantMessage({
message: {
content: 'Sure, I can do that.',
function_call: {
name: 'a_function',
arguments: '{ "foo": "bar" }',
trigger: MessageRole.Assistant,
},
},
}),
buildAssistantChatItem({
content: '',
title: 'suggested a function',
function_call: {
name: 'a_function',
arguments: '{ "bar": "foo" }',
trigger: MessageRole.Assistant,
},
actions: {
canEdit: true,
canCopy: true,
canGiveFeedback: true,
canRegenerate: true,
},
buildFunctionResponseMessage({
message: { content: '{ "message": "The arguments are wrong" }' },
}),
buildFunctionChatItem({
content: '',
title: 'are executing a function',
loading: true,
buildAssistantMessage({
message: {
content: '',
function_call: {
name: 'a_function',
arguments: '{ "bar": "foo" }',
trigger: MessageRole.Assistant,
},
},
}),
],
onEdit: async () => {},

77
x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.tsx
View file

@ -5,7 +5,7 @@
* 2.0.
*/
import React, { ReactNode } from 'react';
import React, { ReactNode, useMemo } from 'react';
import { css } from '@emotion/css';
import { EuiCommentList } from '@elastic/eui';
import type { AuthenticatedUser } from '@kbn/security-plugin/common';
@ -16,6 +16,12 @@ import type { Feedback } from '../feedback_buttons';
import { type Message } from '../../../common';
import type { UseKnowledgeBaseResult } from '../../hooks/use_knowledge_base';
import type { ChatActionClickHandler } from './types';
import {
getTimelineItemsfromConversation,
StartedFrom,
} from '../../utils/get_timeline_items_from_conversation';
import { ObservabilityAIAssistantChatService } from '../../types';
import { ChatState } from '../../hooks/use_chat';
export interface ChatTimelineItem
extends Pick<Message['message'], 'role' | 'content' | 'function_call'> {
@ -35,27 +41,70 @@ export interface ChatTimelineItem
element?: React.ReactNode;
currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
error?: any;
message: Message;
}
export interface ChatTimelineProps {
items: Array<ChatTimelineItem | ChatTimelineItem[]>;
messages: Message[];
knowledgeBase: UseKnowledgeBaseResult;
onEdit: (item: ChatTimelineItem, message: Message) => Promise<void>;
onFeedback: (item: ChatTimelineItem, feedback: Feedback) => void;
onRegenerate: (item: ChatTimelineItem) => void;
chatService: ObservabilityAIAssistantChatService;
hasConnector: boolean;
chatState: ChatState;
currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
startedFrom?: StartedFrom;
onEdit: (message: Message, messageAfterEdit: Message) => void;
onFeedback: (message: Message, feedback: Feedback) => void;
onRegenerate: (message: Message) => void;
onStopGenerating: () => void;
onActionClick: ChatActionClickHandler;
}
export function ChatTimeline({
items = [],
messages,
knowledgeBase,
chatService,
hasConnector,
currentUser,
startedFrom,
onEdit,
onFeedback,
onRegenerate,
onStopGenerating,
onActionClick,
chatState,
}: ChatTimelineProps) {
const items = useMemo(() => {
const timelineItems = getTimelineItemsfromConversation({
chatService,
hasConnector,
messages,
currentUser,
startedFrom,
chatState,
});
const consolidatedChatItems: Array<ChatTimelineItem | ChatTimelineItem[]> = [];
let currentGroup: ChatTimelineItem[] | null = null;
for (const item of timelineItems) {
if (item.display.hide || !item) continue;
if (item.display.collapsed) {
if (currentGroup) {
currentGroup.push(item);
} else {
currentGroup = [item];
consolidatedChatItems.push(currentGroup);
}
} else {
consolidatedChatItems.push(item);
currentGroup = null;
}
}
return consolidatedChatItems;
}, [chatService, hasConnector, messages, currentUser, startedFrom, chatState]);
return (
<EuiCommentList
css={css`
@ -65,8 +114,8 @@ export function ChatTimeline({
{items.length <= 1 ? (
<ChatWelcomePanel knowledgeBase={knowledgeBase} />
) : (
items.map((item, index) =>
Array.isArray(item) ? (
items.map((item, index) => {
return Array.isArray(item) ? (
<ChatConsolidatedItems
key={index}
consolidatedItem={item}
@ -82,17 +131,19 @@ export function ChatTimeline({
key={index}
{...item}
onFeedbackClick={(feedback) => {
onFeedback(item, feedback);
onFeedback(item.message, feedback);
}}
onRegenerateClick={() => {
onRegenerate(item);
onRegenerate(item.message);
}}
onEditSubmit={(message) => {
onEdit(item.message, message);
}}
onEditSubmit={(message) => onEdit(item, message)}
onStopGeneratingClick={onStopGenerating}
onActionClick={onActionClick}
/>
)
)
);
})
)}
</EuiCommentList>
);

2
x-pack/plugins/observability_ai_assistant/public/components/chat/types.ts
View file

@ -20,4 +20,4 @@ export enum ChatActionClickType {
executeEsqlQuery = 'executeEsqlQuery',
}
export type ChatActionClickHandler = (payload: ChatActionClickPayload) => Promise<unknown>;
export type ChatActionClickHandler = (payload: ChatActionClickPayload) => void;

214
x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx
View file

@ -4,20 +4,17 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { first } from 'lodash';
import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
import { AbortError } from '@kbn/kibana-utils-plugin/common';
import { isObservable, Subscription } from 'rxjs';
import { last } from 'lodash';
import React, { useEffect, useRef, useState } from 'react';
import { MessageRole, type Message } from '../../../common/types';
import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
import { useKibana } from '../../hooks/use_kibana';
import { useAbortableAsync } from '../../hooks/use_abortable_async';
import { useConversation } from '../../hooks/use_conversation';
import { ChatState, useChat } from '../../hooks/use_chat';
import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
import { useKibana } from '../../hooks/use_kibana';
import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
import type { PendingMessage } from '../../types';
import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
import { StartChatButton } from '../buttons/start_chat_button';
@ -40,197 +37,40 @@ function ChatContent({
}) {
const chatService = useObservabilityAIAssistantChatService();
const [pendingMessage, setPendingMessage] = useState<PendingMessage | undefined>();
const initialMessagesRef = useRef(initialMessages);
const [loading, setLoading] = useState(false);
const [subscription, setSubscription] = useState<Subscription | undefined>();
const [conversationId, setConversationId] = useState<string>();
const {
conversation,
displayedMessages,
setDisplayedMessages,
getSystemMessage,
save,
saveTitle,
} = useConversation({
conversationId,
connectorId,
const { messages, next, state, stop } = useChat({
chatService,
connectorId,
initialMessages,
});
const conversationTitle = conversationId
? conversation.value?.conversation.title || ''
: defaultTitle;
const controllerRef = useRef(new AbortController());
const reloadRecalledMessages = useCallback(
async (messages: Message[]) => {
controllerRef.current.abort();
const controller = (controllerRef.current = new AbortController());
const isStartOfConversation =
messages.some((message) => message.message.role === MessageRole.Assistant) === false;
if (isStartOfConversation && chatService.hasFunction('recall')) {
// manually execute recall function and append to list of
// messages
const functionCall = {
name: 'recall',
args: JSON.stringify({ queries: [], contexts: [] }),
};
const response = await chatService.executeFunction({
...functionCall,
messages,
signal: controller.signal,
connectorId,
});
if (isObservable(response)) {
throw new Error('Recall function unexpectedly returned an Observable');
}
return [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: functionCall.name,
arguments: functionCall.args,
trigger: MessageRole.User as const,
},
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name: functionCall.name,
content: JSON.stringify(response.content),
},
},
];
}
return [];
},
[chatService, connectorId]
const lastAssistantResponse = last(
messages.filter((message) => message.message.role === MessageRole.Assistant)
);
const reloadConversation = useCallback(async () => {
setLoading(true);
setDisplayedMessages(initialMessages);
setPendingMessage(undefined);
const messages = [getSystemMessage(), ...initialMessages];
const recalledMessages = await reloadRecalledMessages(messages);
const next = messages.concat(recalledMessages);
setDisplayedMessages(next);
let lastPendingMessage: PendingMessage | undefined;
const nextSubscription = chatService
.chat({ messages: next, connectorId, function: 'none' })
.subscribe({
next: (msg) => {
lastPendingMessage = msg;
setPendingMessage(() => msg);
},
complete: () => {
setDisplayedMessages((prev) =>
prev.concat({
'@timestamp': new Date().toISOString(),
...lastPendingMessage!,
})
);
setPendingMessage(lastPendingMessage);
setLoading(false);
},
});
setSubscription(nextSubscription);
}, [
reloadRecalledMessages,
chatService,
connectorId,
initialMessages,
getSystemMessage,
setDisplayedMessages,
]);
useEffect(() => {
reloadConversation();
}, [reloadConversation]);
useEffect(() => {
setDisplayedMessages(initialMessages);
}, [initialMessages, setDisplayedMessages]);
next(initialMessagesRef.current);
}, [next]);
const [isOpen, setIsOpen] = useState(false);
const messagesWithPending = useMemo(() => {
return pendingMessage
? displayedMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
...pendingMessage.message,
},
})
: displayedMessages;
}, [pendingMessage, displayedMessages]);
const firstAssistantMessage = first(
messagesWithPending.filter(
(message) =>
message.message.role === MessageRole.Assistant &&
(!message.message.function_call?.trigger ||
message.message.function_call.trigger === MessageRole.Assistant)
)
);
return (
<>
<MessagePanel
body={
<MessageText
content={firstAssistantMessage?.message.content ?? ''}
loading={loading}
content={lastAssistantResponse?.message.content ?? ''}
loading={state === ChatState.Loading}
onActionClick={async () => {}}
/>
}
error={pendingMessage?.error}
error={state === ChatState.Error}
controls={
loading ? (
state === ChatState.Loading ? (
<StopGeneratingButton
onClick={() => {
subscription?.unsubscribe();
setLoading(false);
setDisplayedMessages((prevMessages) =>
prevMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
...pendingMessage!.message,
},
})
);
setPendingMessage((prev) => ({
message: {
role: MessageRole.Assistant,
...prev?.message,
},
aborted: true,
error: new AbortError(),
}));
stop();
}}
/>
) : (
@ -238,7 +78,7 @@ function ChatContent({
<EuiFlexItem grow={false}>
<RegenerateResponseButton
onClick={() => {
reloadConversation();
next(initialMessages);
}}
/>
</EuiFlexItem>
@ -254,27 +94,13 @@ function ChatContent({
}
/>
<ChatFlyout
title={conversationTitle}
isOpen={isOpen}
onClose={() => {
setIsOpen(() => false);
setIsOpen(false);
}}
messages={displayedMessages}
conversationId={conversationId}
initialMessages={messages}
initialTitle={defaultTitle}
startedFrom="contextualInsight"
onChatComplete={(nextMessages) => {
save(nextMessages)
.then((nextConversation) => {
setConversationId(nextConversation.conversation.id);
})
.catch(() => {});
}}
onChatUpdate={(nextMessages) => {
setDisplayedMessages(nextMessages);
}}
onChatTitleSave={(newTitle) => {
saveTitle(newTitle);
}}
/>
</>
);

4
x-pack/plugins/observability_ai_assistant/public/components/message_panel/message_panel.stories.tsx
View file

@ -71,7 +71,7 @@ export const ContentFailed: ComponentStoryObj<typeof Component> = {
onActionClick={async () => {}}
/>
),
error: new Error(),
error: true,
},
};
@ -111,7 +111,7 @@ export const Controls: ComponentStoryObj<typeof Component> = {
onActionClick={async () => {}}
/>
),
error: new Error(),
error: true,
controls: <FeedbackButtons onClickFeedback={() => {}} />,
},
};

2
x-pack/plugins/observability_ai_assistant/public/components/message_panel/message_panel.tsx
View file

@ -9,7 +9,7 @@ import React from 'react';
import { FailedToLoadResponse } from './failed_to_load_response';
interface Props {
error?: Error;
error?: boolean;
body?: React.ReactNode;
controls?: React.ReactNode;
}

6
x-pack/plugins/observability_ai_assistant/public/hooks/use_abortable_async.ts
View file

@ -18,9 +18,9 @@ export type AbortableAsyncState<T> = (T extends Promise<infer TReturn>
: State<T>) & { refresh: () => void };
export function useAbortableAsync<T>(
fn: ({}: { signal: AbortSignal }) => T,
fn: ({}: { signal: AbortSignal }) => T | Promise<T>,
deps: any[],
options?: { clearValueOnNext?: boolean }
options?: { clearValueOnNext?: boolean; defaultValue?: () => T }
): AbortableAsyncState<T> {
const clearValueOnNext = options?.clearValueOnNext;
@ -30,7 +30,7 @@ export function useAbortableAsync<T>(
const [error, setError] = useState<Error>();
const [loading, setLoading] = useState(false);
const [value, setValue] = useState<T>();
const [value, setValue] = useState<T | undefined>(options?.defaultValue);
useEffect(() => {
controllerRef.current.abort();

464
x-pack/plugins/observability_ai_assistant/public/hooks/use_chat.test.ts Normal file
View file

@ -0,0 +1,464 @@
/*
* 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 type { DeeplyMockedKeys } from '@kbn/utility-types-jest';
import { type RenderHookResult, renderHook, act } from '@testing-library/react-hooks';
import { Subject } from 'rxjs';
import { MessageRole } from '../../common';
import type { ObservabilityAIAssistantChatService, PendingMessage } from '../types';
import { type UseChatResult, useChat, type UseChatProps, ChatState } from './use_chat';
import * as useKibanaModule from './use_kibana';
type MockedChatService = DeeplyMockedKeys<ObservabilityAIAssistantChatService>;
const mockChatService: MockedChatService = {
chat: jest.fn(),
executeFunction: jest.fn(),
getContexts: jest.fn().mockReturnValue([{ name: 'core', description: '' }]),
getFunctions: jest.fn().mockReturnValue([]),
hasFunction: jest.fn().mockReturnValue(false),
hasRenderFunction: jest.fn().mockReturnValue(true),
renderFunction: jest.fn(),
};
const addErrorMock = jest.fn();
jest.spyOn(useKibanaModule, 'useKibana').mockReturnValue({
services: {
notifications: {
toasts: {
addError: addErrorMock,
},
},
},
} as any);
let hookResult: RenderHookResult<UseChatProps, UseChatResult>;
describe('useChat', () => {
beforeEach(() => {
jest.clearAllMocks();
});
describe('initially', () => {
beforeEach(() => {
hookResult = renderHook(useChat, {
initialProps: {
connectorId: 'my-connector',
chatService: mockChatService,
initialMessages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'hello',
},
},
],
} as UseChatProps,
});
});
it('returns the initial messages including the system message', () => {
const { messages } = hookResult.result.current;
expect(messages.length).toBe(2);
expect(messages[0].message.role).toBe('system');
expect(messages[1].message.content).toBe('hello');
});
it('sets chatState to ready', () => {
expect(hookResult.result.current.state).toBe(ChatState.Ready);
});
});
describe('when calling next()', () => {
let subject: Subject<PendingMessage>;
beforeEach(() => {
hookResult = renderHook(useChat, {
initialProps: {
connectorId: 'my-connector',
chatService: mockChatService,
initialMessages: [],
} as UseChatProps,
});
subject = new Subject();
mockChatService.chat.mockReturnValueOnce(subject);
act(() => {
hookResult.result.current.next([
...hookResult.result.current.messages,
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'hello',
},
},
]);
});
});
it('sets the chatState to loading', () => {
expect(hookResult.result.current.state).toBe(ChatState.Loading);
});
describe('after asking for another response', () => {
beforeEach(() => {
act(() => {
hookResult.result.current.next([]);
subject.next({
message: {
role: MessageRole.User,
content: 'goodbye',
},
});
subject.complete();
});
});
it('shows an empty list of messages', () => {
expect(hookResult.result.current.messages.length).toBe(1);
expect(hookResult.result.current.messages[0].message.role).toBe(MessageRole.System);
});
it('aborts the running request', () => {
expect(subject.observed).toBe(false);
});
});
describe('after a partial response', () => {
it('updates the returned messages', () => {
act(() => {
subject.next({
message: {
content: 'good',
role: MessageRole.Assistant,
},
});
});
expect(hookResult.result.current.messages[2].message.content).toBe('good');
});
});
describe('after a completed response', () => {
it('updates the returned messages and the loading state', () => {
act(() => {
subject.next({
message: {
content: 'good',
role: MessageRole.Assistant,
},
});
subject.next({
message: {
content: 'goodbye',
role: MessageRole.Assistant,
},
});
subject.complete();
});
expect(hookResult.result.current.messages[2].message.content).toBe('goodbye');
expect(hookResult.result.current.state).toBe(ChatState.Ready);
});
});
describe('after aborting a response', () => {
beforeEach(() => {
act(() => {
subject.next({
message: {
content: 'good',
role: MessageRole.Assistant,
},
aborted: true,
});
subject.complete();
});
});
it('shows the partial message and sets chatState to aborted', () => {
expect(hookResult.result.current.messages[2].message.content).toBe('good');
expect(hookResult.result.current.state).toBe(ChatState.Aborted);
});
it('does not show an error toast', () => {
expect(addErrorMock).not.toHaveBeenCalled();
});
});
describe('after a response errors out', () => {
beforeEach(() => {
act(() => {
subject.next({
message: {
content: 'good',
role: MessageRole.Assistant,
},
error: new Error('foo'),
});
subject.complete();
});
});
it('shows the partial message and sets chatState to error', () => {
expect(hookResult.result.current.messages[2].message.content).toBe('good');
expect(hookResult.result.current.state).toBe(ChatState.Error);
});
it('shows an error toast', () => {
expect(addErrorMock).toHaveBeenCalled();
});
});
describe('after the LLM responds with a function call', () => {
let resolve: (data: any) => void;
let reject: (error: Error) => void;
beforeEach(() => {
mockChatService.executeFunction.mockResolvedValueOnce(
new Promise((...args) => {
resolve = args[0];
reject = args[1];
})
);
act(() => {
subject.next({
message: {
content: '',
role: MessageRole.Assistant,
function_call: {
name: 'my_function',
arguments: JSON.stringify({ foo: 'bar' }),
trigger: MessageRole.Assistant,
},
},
});
subject.complete();
});
});
it('the chat state stays loading', () => {
expect(hookResult.result.current.state).toBe(ChatState.Loading);
});
it('adds a message', () => {
const { messages } = hookResult.result.current;
expect(messages.length).toBe(3);
expect(messages[2]).toEqual({
'@timestamp': expect.any(String),
message: {
content: '',
function_call: {
arguments: JSON.stringify({ foo: 'bar' }),
name: 'my_function',
trigger: MessageRole.Assistant,
},
role: MessageRole.Assistant,
},
});
});
describe('the function call succeeds', () => {
beforeEach(async () => {
subject = new Subject();
mockChatService.chat.mockReturnValueOnce(subject);
await act(async () => {
resolve({ content: { foo: 'bar' }, data: { bar: 'foo' } });
});
});
it('adds a message', () => {
const { messages } = hookResult.result.current;
expect(messages.length).toBe(4);
expect(messages[3]).toEqual({
'@timestamp': expect.any(String),
message: {
content: JSON.stringify({ foo: 'bar' }),
data: JSON.stringify({ bar: 'foo' }),
name: 'my_function',
role: MessageRole.User,
},
});
});
it('keeps the chat state in loading', () => {
expect(hookResult.result.current.state).toBe(ChatState.Loading);
});
it('sends the function call back to the LLM for a response', () => {
expect(mockChatService.chat).toHaveBeenCalledTimes(2);
expect(mockChatService.chat).toHaveBeenLastCalledWith({
connectorId: 'my-connector',
messages: hookResult.result.current.messages,
});
});
});
describe('the function call fails', () => {
beforeEach(async () => {
subject = new Subject();
mockChatService.chat.mockReturnValue(subject);
await act(async () => {
reject(new Error('connection error'));
});
});
it('keeps the chat state in loading', () => {
expect(hookResult.result.current.state).toBe(ChatState.Loading);
});
it('adds a message', () => {
const { messages } = hookResult.result.current;
expect(messages.length).toBe(4);
expect(messages[3]).toEqual({
'@timestamp': expect.any(String),
message: {
content: JSON.stringify({
message: 'Error: connection error',
error: {},
}),
name: 'my_function',
role: MessageRole.User,
},
});
});
it('does not show an error toast', () => {
expect(addErrorMock).not.toHaveBeenCalled();
});
it('sends the function call back to the LLM for a response', () => {
expect(mockChatService.chat).toHaveBeenCalledTimes(2);
expect(mockChatService.chat).toHaveBeenLastCalledWith({
connectorId: 'my-connector',
messages: hookResult.result.current.messages,
});
});
});
describe('stop() is called', () => {
beforeEach(() => {
act(() => {
hookResult.result.current.stop();
});
});
it('sets the chatState to aborted', () => {
expect(hookResult.result.current.state).toBe(ChatState.Aborted);
});
it('has called the abort controller', () => {
const signal = mockChatService.executeFunction.mock.calls[0][0].signal;
expect(signal.aborted).toBe(true);
});
it('is not updated after the promise is rejected', () => {
const numRenders = hookResult.result.all.length;
act(() => {
reject(new Error('Request aborted'));
});
expect(numRenders).toBe(hookResult.result.all.length);
});
it('removes all subscribers', () => {
expect(subject.observed).toBe(false);
});
});
describe('setMessages() is called', () => {});
});
});
describe('when calling next() with the recall function available', () => {
let subject: Subject<PendingMessage>;
beforeEach(async () => {
hookResult = renderHook(useChat, {
initialProps: {
connectorId: 'my-connector',
chatService: mockChatService,
initialMessages: [],
} as UseChatProps,
});
subject = new Subject();
mockChatService.hasFunction.mockReturnValue(true);
mockChatService.executeFunction.mockResolvedValueOnce({
content: [
{
id: 'my_document',
text: 'My text',
},
],
});
mockChatService.chat.mockReturnValueOnce(subject);
await act(async () => {
hookResult.result.current.next([
...hookResult.result.current.messages,
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'hello',
},
},
]);
});
});
it('adds a user message and a recall function request', () => {
expect(hookResult.result.current.messages[1].message.content).toBe('hello');
expect(hookResult.result.current.messages[2].message.function_call?.name).toBe('recall');
expect(hookResult.result.current.messages[2].message.content).toBe('');
expect(hookResult.result.current.messages[2].message.function_call?.arguments).toBe(
JSON.stringify({ queries: [], contexts: [] })
);
expect(hookResult.result.current.messages[3].message.name).toBe('recall');
expect(hookResult.result.current.messages[3].message.content).toBe(
JSON.stringify([
{
id: 'my_document',
text: 'My text',
},
])
);
});
it('executes the recall function', () => {
expect(mockChatService.executeFunction).toHaveBeenCalled();
expect(mockChatService.executeFunction).toHaveBeenCalledWith({
signal: expect.any(AbortSignal),
connectorId: 'my-connector',
args: JSON.stringify({ queries: [], contexts: [] }),
name: 'recall',
messages: [...hookResult.result.current.messages.slice(0, -1)],
});
});
it('sends the user message, function request and recall response to the LLM', () => {
expect(mockChatService.chat).toHaveBeenCalledWith({
connectorId: 'my-connector',
messages: [...hookResult.result.current.messages],
});
});
});
});

260
x-pack/plugins/observability_ai_assistant/public/hooks/use_chat.ts Normal file
View file

@ -0,0 +1,260 @@
/*
* 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 { i18n } from '@kbn/i18n';
import { last } from 'lodash';
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { isObservable } from 'rxjs';
import { type Message, MessageRole } from '../../common';
import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
import type { ObservabilityAIAssistantChatService, PendingMessage } from '../types';
import { useKibana } from './use_kibana';
import { useOnce } from './use_once';
export enum ChatState {
Ready = 'ready',
Loading = 'loading',
Error = 'error',
Aborted = 'aborted',
}
export interface UseChatResult {
messages: Message[];
setMessages: (messages: Message[]) => void;
state: ChatState;
next: (messages: Message[]) => void;
stop: () => void;
}
export interface UseChatProps {
initialMessages: Message[];
chatService: ObservabilityAIAssistantChatService;
connectorId?: string;
onChatComplete?: (messages: Message[]) => void;
}
export function useChat({
initialMessages,
chatService,
connectorId,
onChatComplete,
}: UseChatProps): UseChatResult {
const [chatState, setChatState] = useState(ChatState.Ready);
const systemMessage = useMemo(() => {
return getAssistantSetupMessage({ contexts: chatService.getContexts() });
}, [chatService]);
useOnce(initialMessages);
const [messages, setMessages] = useState<Message[]>(initialMessages);
const [pendingMessage, setPendingMessage] = useState<PendingMessage>();
const abortControllerRef = useRef(new AbortController());
const {
services: { notifications },
} = useKibana();
const onChatCompleteRef = useRef(onChatComplete);
onChatCompleteRef.current = onChatComplete;
const handleSignalAbort = useCallback(() => {
setChatState(ChatState.Aborted);
}, []);
const next = useCallback(
async (nextMessages: Message[]) => {
// make sure we ignore any aborts for the previous signal
abortControllerRef.current.signal.removeEventListener('abort', handleSignalAbort);
// cancel running requests
abortControllerRef.current.abort();
const lastMessage = last(nextMessages);
const allMessages = [
systemMessage,
...nextMessages.filter((message) => message.message.role !== MessageRole.System),
];
setMessages(allMessages);
if (!lastMessage || !connectorId) {
setChatState(ChatState.Ready);
onChatCompleteRef.current?.(nextMessages);
return;
}
const isUserMessage = lastMessage.message.role === MessageRole.User;
const functionCall = lastMessage.message.function_call;
const isAssistantMessageWithFunctionRequest =
lastMessage.message.role === MessageRole.Assistant && functionCall && !!functionCall.name;
const isFunctionResult = isUserMessage && !!lastMessage.message.name;
const isRecallFunctionAvailable = chatService.hasFunction('recall');
if (!isUserMessage && !isAssistantMessageWithFunctionRequest) {
setChatState(ChatState.Ready);
onChatCompleteRef.current?.(nextMessages);
return;
}
const abortController = (abortControllerRef.current = new AbortController());
abortController.signal.addEventListener('abort', handleSignalAbort);
setChatState(ChatState.Loading);
if (isUserMessage && !isFunctionResult && isRecallFunctionAvailable) {
const allMessagesWithRecall = allMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'recall',
arguments: JSON.stringify({ queries: [], contexts: [] }),
trigger: MessageRole.Assistant,
},
},
});
next(allMessagesWithRecall);
return;
}
function handleError(error: Error) {
setChatState(ChatState.Error);
notifications.toasts.addError(error, {
title: i18n.translate('xpack.observabilityAiAssistant.failedToLoadResponse', {
defaultMessage: 'Failed to load response from the AI Assistant',
}),
});
}
const response = isAssistantMessageWithFunctionRequest
? await chatService
.executeFunction({
name: functionCall.name,
signal: abortController.signal,
args: functionCall.arguments,
connectorId,
messages: allMessages,
})
.catch((error) => {
return {
content: {
message: error.toString(),
error,
},
data: undefined,
};
})
: chatService.chat({
messages: allMessages,
connectorId,
});
if (abortController.signal.aborted) {
return;
}
if (isObservable(response)) {
let localPendingMessage: PendingMessage = {
message: {
content: '',
role: MessageRole.User,
},
};
const subscription = response.subscribe({
next: (nextPendingMessage) => {
localPendingMessage = nextPendingMessage;
setPendingMessage(nextPendingMessage);
},
complete: () => {
setPendingMessage(undefined);
const allMessagesWithResolved = allMessages.concat({
message: {
...localPendingMessage.message,
},
'@timestamp': new Date().toISOString(),
});
if (localPendingMessage.aborted) {
setChatState(ChatState.Aborted);
setMessages(allMessagesWithResolved);
} else if (localPendingMessage.error) {
handleError(localPendingMessage.error);
setMessages(allMessagesWithResolved);
} else {
next(allMessagesWithResolved);
}
},
error: (error) => {
handleError(error);
},
});
abortController.signal.addEventListener('abort', () => {
subscription.unsubscribe();
});
} else {
const allMessagesWithFunctionReply = allMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
name: functionCall!.name,
role: MessageRole.User,
content: JSON.stringify(response.content),
data: JSON.stringify(response.data),
},
});
next(allMessagesWithFunctionReply);
}
},
[connectorId, chatService, handleSignalAbort, notifications.toasts, systemMessage]
);
useEffect(() => {
return () => {
abortControllerRef.current.abort();
};
}, []);
const memoizedMessages = useMemo(() => {
const includingSystemMessage = [
systemMessage,
...messages.filter((message) => message.message.role !== MessageRole.System),
];
return pendingMessage
? includingSystemMessage.concat({
...pendingMessage,
'@timestamp': new Date().toISOString(),
})
: includingSystemMessage;
}, [systemMessage, messages, pendingMessage]);
const setMessagesWithAbort = useCallback((nextMessages: Message[]) => {
abortControllerRef.current.abort();
setPendingMessage(undefined);
setChatState(ChatState.Ready);
setMessages(nextMessages);
}, []);
return {
messages: memoizedMessages,
setMessages: setMessagesWithAbort,
state: chatState,
next,
stop: () => {
abortControllerRef.current.abort();
},
};
}

703
x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.test.tsx Normal file
View file

@ -0,0 +1,703 @@
/*
* 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 React from 'react';
import {
useConversation,
type UseConversationProps,
type UseConversationResult,
} from './use_conversation';
import {
act,
renderHook,
type RenderHookResult,
type WrapperComponent,
} from '@testing-library/react-hooks';
import type { ObservabilityAIAssistantService, PendingMessage } from '../types';
import type { DeeplyMockedKeys } from '@kbn/utility-types-jest';
import { ObservabilityAIAssistantProvider } from '../context/observability_ai_assistant_provider';
import * as useKibanaModule from './use_kibana';
import { Message, MessageRole } from '../../common';
import { ChatState } from './use_chat';
import { createMockChatService } from '../service/create_mock_chat_service';
import { Subject } from 'rxjs';
import { EMPTY_CONVERSATION_TITLE } from '../i18n';
import { merge, omit } from 'lodash';
let hookResult: RenderHookResult<UseConversationProps, UseConversationResult>;
type MockedService = DeeplyMockedKeys<ObservabilityAIAssistantService>;
const mockService: MockedService = {
callApi: jest.fn(),
getCurrentUser: jest.fn(),
getLicense: jest.fn(),
getLicenseManagementLocator: jest.fn(),
isEnabled: jest.fn(),
start: jest.fn(),
};
const mockChatService = createMockChatService();
const addErrorMock = jest.fn();
jest.spyOn(useKibanaModule, 'useKibana').mockReturnValue({
services: {
notifications: {
toasts: {
addError: addErrorMock,
},
},
},
} as any);
describe('useConversation', () => {
let wrapper: WrapperComponent<UseConversationProps>;
beforeEach(() => {
jest.clearAllMocks();
wrapper = ({ children }) => (
<ObservabilityAIAssistantProvider value={mockService}>
{children}
</ObservabilityAIAssistantProvider>
);
});
describe('with initial messages and a conversation id', () => {
beforeEach(() => {
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialMessages: [
{
'@timestamp': new Date().toISOString(),
message: { content: '', role: MessageRole.User },
},
],
initialConversationId: 'foo',
},
wrapper,
});
});
it('throws an error', () => {
expect(hookResult.result.error).toBeTruthy();
});
});
describe('without initial messages and a conversation id', () => {
beforeEach(() => {
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
},
wrapper,
});
});
it('returns only the system message', () => {
expect(hookResult.result.current.messages).toEqual([
{
'@timestamp': expect.any(String),
message: {
content: '',
role: MessageRole.System,
},
},
]);
});
it('returns a ready state', () => {
expect(hookResult.result.current.state).toBe(ChatState.Ready);
});
it('does not call the fetch api', () => {
expect(mockService.callApi).not.toHaveBeenCalled();
});
});
describe('with initial messages', () => {
beforeEach(() => {
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialMessages: [
{
'@timestamp': new Date().toISOString(),
message: {
content: 'Test',
role: MessageRole.User,
},
},
],
},
wrapper,
});
});
it('returns the system message and the initial messages', () => {
expect(hookResult.result.current.messages).toEqual([
{
'@timestamp': expect.any(String),
message: {
content: '',
role: MessageRole.System,
},
},
{
'@timestamp': expect.any(String),
message: {
content: 'Test',
role: MessageRole.User,
},
},
]);
});
});
describe('with a conversation id that successfully loads', () => {
beforeEach(async () => {
mockService.callApi.mockResolvedValueOnce({
conversation: {
id: 'my-conversation-id',
},
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'User',
},
},
],
});
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialConversationId: 'my-conversation-id',
},
wrapper,
});
await act(async () => {});
});
it('returns the loaded conversation', () => {
expect(hookResult.result.current.conversation.value).toEqual({
conversation: {
id: 'my-conversation-id',
},
messages: [
{
'@timestamp': expect.any(String),
message: {
content: 'System',
role: MessageRole.System,
},
},
{
'@timestamp': expect.any(String),
message: {
content: 'User',
role: MessageRole.User,
},
},
],
});
});
it('sets messages to the messages of the conversation', () => {
expect(hookResult.result.current.messages).toEqual([
{
'@timestamp': expect.any(String),
message: {
content: expect.any(String),
role: MessageRole.System,
},
},
{
'@timestamp': expect.any(String),
message: {
content: 'User',
role: MessageRole.User,
},
},
]);
});
it('overrides the system message', () => {
expect(hookResult.result.current.messages[0].message.content).toBe('');
});
});
describe('with a conversation id that fails to load', () => {
beforeEach(async () => {
mockService.callApi.mockRejectedValueOnce(new Error('failed to load'));
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialConversationId: 'my-conversation-id',
},
wrapper,
});
await act(async () => {});
});
it('returns an error', () => {
expect(hookResult.result.current.conversation.error).toBeTruthy();
});
it('resets the messages', () => {
expect(hookResult.result.current.messages.length).toBe(1);
});
});
describe('when chat completes without an initial conversation id', () => {
let subject: Subject<PendingMessage>;
const expectedMessages = [
{
'@timestamp': expect.any(String),
message: {
role: MessageRole.System,
content: '',
},
},
{
'@timestamp': expect.any(String),
message: {
role: MessageRole.User,
content: 'Hello',
},
},
{
'@timestamp': expect.any(String),
message: {
role: MessageRole.Assistant,
content: 'Goodbye',
},
},
{
'@timestamp': expect.any(String),
message: {
role: MessageRole.User,
content: 'Hello again',
},
},
{
'@timestamp': expect.any(String),
message: {
role: MessageRole.Assistant,
content: 'Goodbye again',
},
},
];
beforeEach(() => {
subject = new Subject();
mockService.callApi.mockImplementation(async (endpoint, request) =>
merge(
{
conversation: {
id: 'my-conversation-id',
},
messages: expectedMessages,
},
(request as any).params.body
)
);
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialMessages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: 'Goodbye',
},
},
],
},
wrapper,
});
mockChatService.chat.mockImplementationOnce(() => {
return subject;
});
act(() => {
hookResult.result.current.next(
hookResult.result.current.messages.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello again',
},
})
);
});
});
describe('when chat completes with an error', () => {
beforeEach(async () => {
mockService.callApi.mockClear();
act(() => {
subject.next({
message: {
role: MessageRole.Assistant,
content: 'Goodbye',
},
error: new Error(),
});
subject.complete();
});
await act(async () => {});
});
it('does not store the conversation', () => {
expect(mockService.callApi).not.toHaveBeenCalled();
});
});
describe('when chat completes without an error', () => {
beforeEach(async () => {
act(() => {
subject.next({
message: {
role: MessageRole.Assistant,
content: 'Goodbye again',
},
});
subject.complete();
});
await act(async () => {});
});
it('the conversation is created including the initial messages', async () => {
expect(mockService.callApi.mock.calls[0]).toEqual([
'POST /internal/observability_ai_assistant/conversation',
{
params: {
body: {
conversation: {
'@timestamp': expect.any(String),
conversation: {
title: EMPTY_CONVERSATION_TITLE,
},
messages: expectedMessages,
labels: {},
numeric_labels: {},
public: false,
},
},
},
signal: null,
},
]);
expect(hookResult.result.current.conversation.error).toBeUndefined();
expect(hookResult.result.current.messages).toEqual(expectedMessages);
});
});
});
describe('when chat completes with an initial conversation id', () => {
let subject: Subject<PendingMessage>;
const initialMessages: Message[] = [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: '',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'user',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: 'assistant',
},
},
];
beforeEach(async () => {
mockService.callApi.mockImplementation(async (endpoint, request) => ({
'@timestamp': new Date().toISOString(),
conversation: {
id: 'my-conversation-id',
title: EMPTY_CONVERSATION_TITLE,
},
labels: {},
numeric_labels: {},
public: false,
messages: initialMessages,
}));
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialConversationId: 'my-conversation-id',
},
wrapper,
});
await act(async () => {});
});
it('the conversation is loadeded', async () => {
expect(mockService.callApi.mock.calls[0]).toEqual([
'GET /internal/observability_ai_assistant/conversation/{conversationId}',
{
signal: expect.anything(),
params: {
path: {
conversationId: 'my-conversation-id',
},
},
},
]);
expect(hookResult.result.current.messages).toEqual(
initialMessages.map((msg) => ({ ...msg, '@timestamp': expect.any(String) }))
);
});
describe('after chat completes', () => {
const nextUserMessage: Message = {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello again',
},
};
const nextAssistantMessage: Message = {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: 'Goodbye again',
},
};
beforeEach(async () => {
mockService.callApi.mockClear();
subject = new Subject();
mockChatService.chat.mockImplementationOnce(() => {
return subject;
});
act(() => {
hookResult.result.current.next(
hookResult.result.current.messages.concat(nextUserMessage)
);
subject.next(omit(nextAssistantMessage, '@timestamp'));
subject.complete();
});
await act(async () => {});
});
it('saves the updated message', () => {
expect(mockService.callApi.mock.calls[0]).toEqual([
'PUT /internal/observability_ai_assistant/conversation/{conversationId}',
{
params: {
path: {
conversationId: 'my-conversation-id',
},
body: {
conversation: {
'@timestamp': expect.any(String),
conversation: {
title: EMPTY_CONVERSATION_TITLE,
id: 'my-conversation-id',
},
messages: initialMessages
.concat([nextUserMessage, nextAssistantMessage])
.map((msg) => ({ ...msg, '@timestamp': expect.any(String) })),
labels: {},
numeric_labels: {},
public: false,
},
},
},
signal: null,
},
]);
});
});
});
describe('when the title is updated', () => {
describe('without a stored conversation', () => {
beforeEach(() => {
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialMessages: [
{
'@timestamp': new Date().toISOString(),
message: { content: '', role: MessageRole.User },
},
],
initialConversationId: 'foo',
},
wrapper,
});
});
it('throws an error', () => {
expect(() => hookResult.result.current.saveTitle('my-new-title')).toThrow();
});
});
describe('with a stored conversation', () => {
let resolve: (value: unknown) => void;
beforeEach(async () => {
mockService.callApi.mockImplementation(async (endpoint, request) => {
if (
endpoint === 'PUT /internal/observability_ai_assistant/conversation/{conversationId}'
) {
return new Promise((_resolve) => {
resolve = _resolve;
});
}
return {
'@timestamp': new Date().toISOString(),
conversation: {
id: 'my-conversation-id',
title: EMPTY_CONVERSATION_TITLE,
},
labels: {},
numeric_labels: {},
public: false,
messages: [],
};
});
await act(async () => {
hookResult = renderHook(useConversation, {
initialProps: {
chatService: mockChatService,
connectorId: 'my-connector',
initialConversationId: 'my-conversation-id',
},
wrapper,
});
});
});
it('does not throw an error', () => {
expect(() => hookResult.result.current.saveTitle('my-new-title')).not.toThrow();
});
it('calls the update API', async () => {
act(() => {
hookResult.result.current.saveTitle('my-new-title');
});
expect(resolve).not.toBeUndefined();
expect(mockService.callApi.mock.calls[1]).toEqual([
'PUT /internal/observability_ai_assistant/conversation/{conversationId}',
{
signal: null,
params: {
path: {
conversationId: 'my-conversation-id',
},
body: {
conversation: {
'@timestamp': expect.any(String),
conversation: {
title: 'my-new-title',
id: 'my-conversation-id',
},
labels: expect.anything(),
messages: expect.anything(),
numeric_labels: expect.anything(),
public: expect.anything(),
},
},
},
},
]);
mockService.callApi.mockImplementation(async (endpoint, request) => {
return {
'@timestamp': new Date().toISOString(),
conversation: {
id: 'my-conversation-id',
title: 'my-new-title',
},
labels: {},
numeric_labels: {},
public: false,
messages: [],
};
});
await act(async () => {
resolve({
conversation: {
title: 'my-new-title',
},
});
});
expect(mockService.callApi.mock.calls[2]).toEqual([
'GET /internal/observability_ai_assistant/conversation/{conversationId}',
{
signal: expect.anything(),
params: {
path: {
conversationId: 'my-conversation-id',
},
},
},
]);
expect(hookResult.result.current.conversation.value?.conversation.title).toBe(
'my-new-title'
);
});
});
});
});

293
x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts
View file

@ -6,145 +6,127 @@
*/
import { i18n } from '@kbn/i18n';
import { merge, omit } from 'lodash';
import { Dispatch, SetStateAction, useCallback, useMemo, useState } from 'react';
import { type Conversation, type Message } from '../../common';
import { ConversationCreateRequest, MessageRole } from '../../common/types';
import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
import { ObservabilityAIAssistantChatService } from '../types';
import { useState } from 'react';
import type { Conversation, Message } from '../../common';
import type { ConversationCreateRequest } from '../../common/types';
import { EMPTY_CONVERSATION_TITLE } from '../i18n';
import type { ObservabilityAIAssistantChatService } from '../types';
import { useAbortableAsync, type AbortableAsyncState } from './use_abortable_async';
import { useChat, UseChatResult } from './use_chat';
import { useKibana } from './use_kibana';
import { useObservabilityAIAssistant } from './use_observability_ai_assistant';
import { createNewConversation } from './use_timeline';
import { useOnce } from './use_once';
function createNewConversation({
title = EMPTY_CONVERSATION_TITLE,
}: { title?: string } = {}): ConversationCreateRequest {
return {
'@timestamp': new Date().toISOString(),
messages: [],
conversation: {
title,
},
labels: {},
numeric_labels: {},
public: false,
};
}
export interface UseConversationProps {
initialConversationId?: string;
initialMessages?: Message[];
initialTitle?: string;
chatService: ObservabilityAIAssistantChatService;
connectorId: string | undefined;
onConversationUpdate?: (conversation: Conversation) => void;
}
export type UseConversationResult = {
conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined>;
saveTitle: (newTitle: string) => void;
} & Omit<UseChatResult, 'setMessages'>;
const DEFAULT_INITIAL_MESSAGES: Message[] = [];
export function useConversation({
conversationId,
initialConversationId: initialConversationIdFromProps,
initialMessages: initialMessagesFromProps = DEFAULT_INITIAL_MESSAGES,
initialTitle: initialTitleFromProps,
chatService,
connectorId,
initialMessages = [],
}: {
conversationId?: string;
chatService?: ObservabilityAIAssistantChatService; // will eventually resolve to a non-nullish value
connectorId: string | undefined;
initialMessages?: Message[];
}): {
conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined>;
displayedMessages: Message[];
setDisplayedMessages: Dispatch<SetStateAction<Message[]>>;
getSystemMessage: () => Message;
save: (messages: Message[], handleRefreshConversations?: () => void) => Promise<Conversation>;
saveTitle: (
title: string,
handleRefreshConversations?: () => void
) => Promise<Conversation | void>;
} {
onConversationUpdate,
}: UseConversationProps): UseConversationResult {
const service = useObservabilityAIAssistant();
const {
services: { notifications },
} = useKibana();
const [displayedMessages, setDisplayedMessages] = useState<Message[]>(initialMessages);
const initialConversationId = useOnce(initialConversationIdFromProps);
const initialMessages = useOnce(initialMessagesFromProps);
const initialTitle = useOnce(initialTitleFromProps);
const getSystemMessage = useCallback(() => {
return getAssistantSetupMessage({ contexts: chatService?.getContexts() || [] });
}, [chatService]);
if (initialMessages.length && initialConversationId) {
throw new Error('Cannot set initialMessages if initialConversationId is set');
}
const displayedMessagesWithHardcodedSystemMessage = useMemo(() => {
if (!chatService) {
return displayedMessages;
}
const systemMessage = getSystemMessage();
if (displayedMessages[0]?.message.role === MessageRole.User) {
return [systemMessage, ...displayedMessages];
}
return [systemMessage, ...displayedMessages.slice(1)];
}, [displayedMessages, chatService, getSystemMessage]);
const conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined> =
useAbortableAsync(
({ signal }) => {
if (!conversationId) {
const nextConversation = createNewConversation({
contexts: chatService?.getContexts() || [],
});
setDisplayedMessages(nextConversation.messages);
return nextConversation;
}
return service
.callApi('GET /internal/observability_ai_assistant/conversation/{conversationId}', {
signal,
params: { path: { conversationId } },
})
.then((nextConversation) => {
setDisplayedMessages(nextConversation.messages);
return nextConversation;
})
.catch((error) => {
setDisplayedMessages([]);
throw error;
});
},
[conversationId, chatService]
);
return {
conversation,
displayedMessages: displayedMessagesWithHardcodedSystemMessage,
setDisplayedMessages,
getSystemMessage,
save: (messages: Message[], handleRefreshConversations?: () => void) => {
const conversationObject = conversation.value!;
return conversationId
? service
.callApi(`PUT /internal/observability_ai_assistant/conversation/{conversationId}`, {
signal: null,
params: {
path: {
conversationId,
},
body: {
conversation: merge(
{
'@timestamp': conversationObject['@timestamp'],
conversation: {
id: conversationId,
},
},
omit(
conversationObject,
'conversation.last_updated',
'namespace',
'user',
'messages'
),
{ messages }
),
const update = (nextConversationObject: Conversation) => {
return service
.callApi(`PUT /internal/observability_ai_assistant/conversation/{conversationId}`, {
signal: null,
params: {
path: {
conversationId: nextConversationObject.conversation.id,
},
body: {
conversation: merge(
{
'@timestamp': nextConversationObject['@timestamp'],
conversation: {
id: nextConversationObject.conversation.id,
},
},
})
.catch((err) => {
notifications.toasts.addError(err, {
title: i18n.translate('xpack.observabilityAiAssistant.errorUpdatingConversation', {
defaultMessage: 'Could not update conversation',
}),
});
throw err;
})
omit(nextConversationObject, 'conversation.last_updated', 'namespace', 'user')
),
},
},
})
.catch((err) => {
notifications.toasts.addError(err, {
title: i18n.translate('xpack.observabilityAiAssistant.errorUpdatingConversation', {
defaultMessage: 'Could not update conversation',
}),
});
throw err;
});
};
const save = (nextMessages: Message[]) => {
const conversationObject = conversation.value!;
const nextConversationObject = merge({}, omit(conversationObject, 'messages'), {
messages: nextMessages,
});
return (
displayedConversationId
? update(
merge(
{ conversation: { id: displayedConversationId } },
nextConversationObject
) as Conversation
)
: service
.callApi(`POST /internal/observability_ai_assistant/conversation`, {
signal: null,
params: {
body: {
conversation: merge({}, conversationObject, { messages }),
conversation: nextConversationObject,
},
},
})
.then((nextConversation) => {
setDisplayedConversationId(nextConversation.conversation.id);
if (connectorId) {
service
.callApi(
@ -162,7 +144,7 @@ export function useConversation({
}
)
.then(() => {
handleRefreshConversations?.();
onConversationUpdate?.(nextConversation);
return conversation.refresh();
});
}
@ -175,27 +157,78 @@ export function useConversation({
}),
});
throw err;
});
})
).then((nextConversation) => {
onConversationUpdate?.(nextConversation);
return nextConversation;
});
};
const { next, messages, setMessages, state, stop } = useChat({
initialMessages,
chatService,
connectorId,
onChatComplete: (nextMessages) => {
save(nextMessages);
},
saveTitle: (title: string, handleRefreshConversations?: () => void) => {
if (conversationId) {
});
const [displayedConversationId, setDisplayedConversationId] = useState(initialConversationId);
const conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined> =
useAbortableAsync(
({ signal }) => {
if (!displayedConversationId) {
const nextConversation = createNewConversation({ title: initialTitle });
return nextConversation;
}
return service
.callApi('PUT /internal/observability_ai_assistant/conversation/{conversationId}/title', {
signal: null,
params: {
path: {
conversationId,
},
body: {
title,
},
},
.callApi('GET /internal/observability_ai_assistant/conversation/{conversationId}', {
signal,
params: { path: { conversationId: displayedConversationId } },
})
.then(() => {
handleRefreshConversations?.();
.then((nextConversation) => {
setMessages(nextConversation.messages);
return nextConversation;
})
.catch((error) => {
setMessages([]);
throw error;
});
},
[displayedConversationId, initialTitle],
{
defaultValue: () => {
if (!displayedConversationId) {
const nextConversation = createNewConversation({ title: initialTitle });
return nextConversation;
}
return undefined;
},
}
return Promise.resolve();
);
return {
conversation,
state,
next,
stop,
messages,
saveTitle: (title: string) => {
if (!displayedConversationId || !conversation.value) {
throw new Error('Cannot save title if conversation is not stored');
}
const nextConversation = merge({}, conversation.value as Conversation, {
conversation: { title },
});
return update(nextConversation)
.then(() => {
return conversation.refresh();
})
.then(() => {
onConversationUpdate?.(nextConversation);
});
},
};
}

16
x-pack/plugins/observability_ai_assistant/public/hooks/use_force_update.ts Normal file
View file

@ -0,0 +1,16 @@
/*
* 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 { useState } from 'react';
export function useForceUpdate() {
const [_, setCounter] = useState(0);
return () => {
setCounter((prev) => prev + 1);
};
}

1
x-pack/plugins/observability_ai_assistant/public/hooks/use_knowledge_base.tsx
View file

@ -57,6 +57,7 @@ export function useKnowledgeBase(): UseKnowledgeBaseResult {
text: i18n.translate('xpack.observabilityAiAssistant.knowledgeBaseReadyContentReload', {
defaultMessage: 'A page reload is needed to be able to use it.',
}),
toastLifeTimeMs: Number.MAX_VALUE,
});
})
.catch((error) => {

21
x-pack/plugins/observability_ai_assistant/public/hooks/use_once.ts Normal file
View file

@ -0,0 +1,21 @@
/*
* 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 { useRef } from 'react';
export function useOnce<T>(variable: T): T {
const ref = useRef(variable);
if (ref.current !== variable) {
// eslint-disable-next-line no-console
console.trace(
`Variable changed from ${ref.current} to ${variable}, but only the initial value will be taken into account`
);
}
return ref.current;
}

611
x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.test.ts
View file

@ -1,611 +0,0 @@
/*
* 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 type { FindActionResult } from '@kbn/actions-plugin/server';
import { AbortError } from '@kbn/kibana-utils-plugin/common';
import {
act,
renderHook,
type Renderer,
type RenderHookResult,
} from '@testing-library/react-hooks';
import { BehaviorSubject, Subject } from 'rxjs';
import { MessageRole } from '../../common';
import { ChatTimelineItem } from '../components/chat/chat_timeline';
import type { PendingMessage } from '../types';
import { useTimeline, UseTimelineResult } from './use_timeline';
type HookProps = Parameters<typeof useTimeline>[0];
const WAIT_OPTIONS = { timeout: 1500 };
jest.mock('./use_kibana', () => ({
useKibana: () => ({
services: {
notifications: {
toasts: {
addError: jest.fn(),
},
},
},
}),
}));
describe('useTimeline', () => {
let hookResult: RenderHookResult<HookProps, UseTimelineResult, Renderer<HookProps>>;
describe('with an empty conversation', () => {
beforeAll(() => {
hookResult = renderHook((props) => useTimeline(props), {
initialProps: {
connectors: {
loading: false,
selectedConnector: 'OpenAI',
selectConnector: () => {},
connectors: [{ id: 'OpenAI' }] as FindActionResult[],
},
chatService: {},
messages: [],
onChatComplete: jest.fn(),
onChatUpdate: jest.fn(),
} as unknown as HookProps,
});
});
it('renders the correct timeline items', () => {
expect(hookResult.result.current.items.length).toEqual(1);
expect(hookResult.result.current.items[0]).toEqual({
display: {
collapsed: false,
hide: false,
},
actions: {
canCopy: false,
canEdit: false,
canRegenerate: false,
canGiveFeedback: false,
},
role: MessageRole.User,
title: 'started a conversation',
loading: false,
id: expect.any(String),
});
});
});
describe('with an existing conversation', () => {
beforeAll(() => {
hookResult = renderHook((props) => useTimeline(props), {
initialProps: {
messages: [
{
message: {
role: MessageRole.System,
content: 'You are a helpful assistant for Elastic Observability',
},
},
{
message: {
role: MessageRole.User,
content: 'hello',
},
},
{
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'recall',
trigger: MessageRole.User,
},
},
},
{
message: {
name: 'recall',
role: MessageRole.User,
content: '',
},
},
{
message: {
content: 'goodbye',
function_call: {
name: '',
arguments: '',
trigger: MessageRole.Assistant,
},
role: MessageRole.Assistant,
},
},
],
connectors: {
selectedConnector: 'foo',
},
chatService: {
chat: () => {},
hasRenderFunction: () => {},
hasFunction: () => {},
},
} as unknown as HookProps,
});
});
it('renders the correct timeline items', () => {
expect(hookResult.result.current.items.length).toEqual(4);
expect(hookResult.result.current.items[1]).toEqual({
actions: { canCopy: true, canEdit: true, canGiveFeedback: false, canRegenerate: false },
content: 'hello',
currentUser: undefined,
display: { collapsed: false, hide: false },
element: undefined,
function_call: undefined,
id: expect.any(String),
loading: false,
role: MessageRole.User,
title: '',
});
expect(hookResult.result.current.items[3]).toEqual({
actions: { canCopy: true, canEdit: false, canGiveFeedback: false, canRegenerate: true },
content: 'goodbye',
currentUser: undefined,
display: { collapsed: false, hide: false },
element: undefined,
function_call: {
arguments: '',
name: '',
trigger: MessageRole.Assistant,
},
id: expect.any(String),
loading: false,
role: MessageRole.Assistant,
title: '',
});
// Items that are function calls are collapsed into an array.
// 'title' is a <FormattedMessage /> component. This throws Jest for a loop.
const collapsedItemsWithoutTitle = (
hookResult.result.current.items[2] as ChatTimelineItem[]
).map(({ title, ...rest }) => rest);
expect(collapsedItemsWithoutTitle).toEqual([
{
display: {
collapsed: true,
hide: false,
},
actions: {
canCopy: true,
canEdit: true,
canRegenerate: false,
canGiveFeedback: false,
},
currentUser: undefined,
function_call: {
name: 'recall',
trigger: MessageRole.User,
},
role: MessageRole.User,
content: `\`\`\`
{
\"name\": \"recall\"
}
\`\`\``,
loading: false,
id: expect.any(String),
},
{
display: {
collapsed: true,
hide: false,
},
actions: {
canCopy: true,
canEdit: false,
canRegenerate: false,
canGiveFeedback: false,
},
currentUser: undefined,
function_call: undefined,
role: MessageRole.User,
content: `\`\`\`
{}
\`\`\``,
loading: false,
id: expect.any(String),
},
]);
});
});
describe('when submitting a new prompt', () => {
let subject: Subject<PendingMessage>;
let props: Omit<HookProps, 'onChatUpdate' | 'onChatComplete' | 'chatService'> & {
onChatUpdate: jest.MockedFn<HookProps['onChatUpdate']>;
onChatComplete: jest.MockedFn<HookProps['onChatComplete']>;
chatService: Omit<HookProps['chatService'], 'executeFunction'> & {
executeFunction: jest.MockedFn<HookProps['chatService']['executeFunction']>;
};
};
beforeEach(() => {
props = {
messages: [],
connectors: {
selectedConnector: 'foo',
},
chatService: {
chat: jest.fn().mockImplementation(() => {
subject = new BehaviorSubject<PendingMessage>({
message: {
role: MessageRole.Assistant,
content: '',
},
});
return subject;
}),
executeFunction: jest.fn(),
hasFunction: jest.fn(),
hasRenderFunction: jest.fn(),
},
onChatUpdate: jest.fn().mockImplementation((messages) => {
props = { ...props, messages };
hookResult.rerender(props as unknown as HookProps);
}),
onChatComplete: jest.fn(),
} as any;
hookResult = renderHook((nextProps) => useTimeline(nextProps), {
initialProps: props as unknown as HookProps,
});
});
describe("and it's loading", () => {
beforeEach(() => {
act(() => {
hookResult.result.current.onSubmit({
'@timestamp': new Date().toISOString(),
message: { role: MessageRole.User, content: 'Hello' },
});
});
});
it('adds two items of which the last one is loading', async () => {
expect((hookResult.result.current.items[0] as ChatTimelineItem).role).toEqual(
MessageRole.User
);
expect((hookResult.result.current.items[1] as ChatTimelineItem).role).toEqual(
MessageRole.User
);
expect((hookResult.result.current.items[2] as ChatTimelineItem).role).toEqual(
MessageRole.Assistant
);
expect(hookResult.result.current.items[1]).toMatchObject({
role: MessageRole.User,
content: 'Hello',
loading: false,
});
expect(hookResult.result.current.items[2]).toMatchObject({
role: MessageRole.Assistant,
content: '',
loading: true,
actions: {
canRegenerate: false,
canGiveFeedback: false,
},
});
expect(hookResult.result.current.items.length).toBe(3);
expect(hookResult.result.current.items[2]).toMatchObject({
role: MessageRole.Assistant,
content: '',
loading: true,
actions: {
canRegenerate: false,
canGiveFeedback: false,
},
});
});
describe('and it pushes the next part', () => {
beforeEach(() => {
act(() => {
subject.next({ message: { role: MessageRole.Assistant, content: 'Goodbye' } });
});
});
it('adds the partial response', () => {
expect(hookResult.result.current.items[2]).toMatchObject({
role: MessageRole.Assistant,
content: 'Goodbye',
loading: true,
actions: {
canRegenerate: false,
canGiveFeedback: false,
},
});
});
describe('and it completes', () => {
beforeEach(async () => {
act(() => {
subject.complete();
});
await hookResult.waitForNextUpdate(WAIT_OPTIONS);
});
it('adds the completed message', () => {
expect(hookResult.result.current.items[2]).toMatchObject({
role: MessageRole.Assistant,
content: 'Goodbye',
loading: false,
actions: {
canRegenerate: true,
canGiveFeedback: false,
},
});
});
describe('and the user edits a message', () => {
beforeEach(() => {
act(() => {
hookResult.result.current.onEdit(
hookResult.result.current.items[1] as ChatTimelineItem,
{
'@timestamp': new Date().toISOString(),
message: { content: 'Edited message', role: MessageRole.User },
}
);
subject.next({ message: { role: MessageRole.Assistant, content: '' } });
subject.complete();
});
});
it('calls onChatUpdate with the edited message', () => {
expect(hookResult.result.current.items.length).toEqual(4);
expect((hookResult.result.current.items[2] as ChatTimelineItem).content).toEqual(
'Edited message'
);
expect((hookResult.result.current.items[3] as ChatTimelineItem).content).toEqual('');
});
});
});
});
describe('and it is being aborted', () => {
beforeEach(() => {
act(() => {
subject.next({ message: { role: MessageRole.Assistant, content: 'My partial' } });
subject.next({
message: {
role: MessageRole.Assistant,
content: 'My partial',
},
aborted: true,
error: new AbortError(),
});
subject.complete();
});
});
it('adds the partial response', async () => {
expect(hookResult.result.current.items.length).toBe(3);
expect(hookResult.result.current.items[2]).toEqual({
actions: {
canEdit: false,
canRegenerate: true,
canGiveFeedback: false,
canCopy: true,
},
display: {
collapsed: false,
hide: false,
},
content: 'My partial',
id: expect.any(String),
loading: false,
title: '',
role: MessageRole.Assistant,
error: expect.any(AbortError),
});
});
describe('and it is being regenerated', () => {
beforeEach(() => {
act(() => {
hookResult.result.current.onRegenerate(
hookResult.result.current.items[2] as ChatTimelineItem
);
subject.next({ message: { role: MessageRole.Assistant, content: '' } });
});
});
it('updates the last item in the array to be loading', () => {
expect(hookResult.result.current.items.length).toEqual(3);
expect(hookResult.result.current.items[2]).toEqual({
display: {
hide: false,
collapsed: false,
},
actions: {
canCopy: true,
canEdit: false,
canRegenerate: false,
canGiveFeedback: false,
},
content: '',
id: expect.any(String),
loading: true,
title: '',
role: MessageRole.Assistant,
});
});
describe('and it is regenerated again', () => {
beforeEach(async () => {
act(() => {
hookResult.result.current.onStopGenerating();
});
act(() => {
hookResult.result.current.onRegenerate(
hookResult.result.current.items[2] as ChatTimelineItem
);
});
});
it('updates the last item to be not loading again', async () => {
expect(hookResult.result.current.items.length).toBe(3);
expect(hookResult.result.current.items[2]).toEqual({
actions: {
canCopy: true,
canEdit: false,
canRegenerate: false,
canGiveFeedback: false,
},
display: {
collapsed: false,
hide: false,
},
content: '',
id: expect.any(String),
loading: true,
title: '',
role: MessageRole.Assistant,
});
act(() => {
subject.next({ message: { role: MessageRole.Assistant, content: 'Regenerated' } });
subject.complete();
});
await hookResult.waitForNextUpdate(WAIT_OPTIONS);
expect(hookResult.result.current.items.length).toBe(3);
expect(hookResult.result.current.items[2]).toEqual({
display: {
collapsed: false,
hide: false,
},
actions: {
canCopy: true,
canEdit: false,
canRegenerate: true,
canGiveFeedback: false,
},
content: 'Regenerated',
currentUser: undefined,
function_call: undefined,
id: expect.any(String),
element: undefined,
loading: false,
title: '',
role: MessageRole.Assistant,
});
});
});
});
});
describe('and a function call is returned', () => {
it('the function call is executed and its response is sent as a user reply', async () => {
jest.clearAllMocks();
act(() => {
subject.next({
message: {
role: MessageRole.Assistant,
function_call: {
trigger: MessageRole.Assistant,
name: 'my_function',
arguments: '{}',
},
},
});
subject.complete();
});
props.chatService.executeFunction.mockResolvedValueOnce({
content: {
message: 'my-response',
},
});
await hookResult.waitForNextUpdate(WAIT_OPTIONS);
expect(props.onChatUpdate).toHaveBeenCalledTimes(2);
expect(
props.onChatUpdate.mock.calls[0][0].map(
(msg) => msg.message.content || msg.message.function_call?.name
)
).toEqual(['Hello', 'my_function']);
expect(
props.onChatUpdate.mock.calls[1][0].map(
(msg) => msg.message.content || msg.message.function_call?.name
)
).toEqual(['Hello', 'my_function', JSON.stringify({ message: 'my-response' })]);
expect(props.onChatComplete).not.toHaveBeenCalled();
expect(props.chatService.executeFunction).toHaveBeenCalledWith({
name: 'my_function',
args: '{}',
connectorId: 'foo',
messages: [
{
'@timestamp': expect.any(String),
message: {
content: 'Hello',
role: MessageRole.User,
},
},
],
signal: expect.any(Object),
});
act(() => {
subject.next({
message: {
role: MessageRole.Assistant,
content: 'looks like my-function returned my-response',
},
});
subject.complete();
});
await hookResult.waitForNextUpdate(WAIT_OPTIONS);
expect(props.onChatComplete).toHaveBeenCalledTimes(1);
expect(
props.onChatComplete.mock.calls[0][0].map(
(msg) => msg.message.content || msg.message.function_call?.name
)
).toEqual([
'Hello',
'my_function',
JSON.stringify({ message: 'my-response' }),
'looks like my-function returned my-response',
]);
});
});
});
});
});

387
x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts
View file

@ -1,387 +0,0 @@
/*
* 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 { i18n } from '@kbn/i18n';
import { AbortError } from '@kbn/kibana-utils-plugin/common';
import type { AuthenticatedUser } from '@kbn/security-plugin/common';
import { flatten, last } from 'lodash';
import { useEffect, useMemo, useRef, useState } from 'react';
import usePrevious from 'react-use/lib/usePrevious';
import { isObservable, Observable, Subscription } from 'rxjs';
import {
ContextDefinition,
MessageRole,
type ConversationCreateRequest,
type Message,
} from '../../common/types';
import type { ChatPromptEditorProps } from '../components/chat/chat_prompt_editor';
import type { ChatTimelineItem, ChatTimelineProps } from '../components/chat/chat_timeline';
import { ChatActionClickType } from '../components/chat/types';
import { EMPTY_CONVERSATION_TITLE } from '../i18n';
import type { ObservabilityAIAssistantChatService, PendingMessage } from '../types';
import {
getTimelineItemsfromConversation,
StartedFrom,
} from '../utils/get_timeline_items_from_conversation';
import type { UseGenAIConnectorsResult } from './use_genai_connectors';
import { useKibana } from './use_kibana';
export function createNewConversation({
contexts,
}: {
contexts: ContextDefinition[];
}): ConversationCreateRequest {
return {
'@timestamp': new Date().toISOString(),
messages: [],
conversation: {
title: EMPTY_CONVERSATION_TITLE,
},
labels: {},
numeric_labels: {},
public: false,
};
}
export type UseTimelineResult = Pick<
ChatTimelineProps,
'onEdit' | 'onFeedback' | 'onRegenerate' | 'onStopGenerating' | 'onActionClick' | 'items'
> &
Pick<ChatPromptEditorProps, 'onSubmit'>;
export function useTimeline({
messages,
connectors,
conversationId,
currentUser,
chatService,
startedFrom,
onChatUpdate,
onChatComplete,
}: {
messages: Message[];
conversationId?: string;
connectors: UseGenAIConnectorsResult;
currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
chatService: ObservabilityAIAssistantChatService;
startedFrom?: StartedFrom;
onChatUpdate: (messages: Message[]) => void;
onChatComplete: (messages: Message[]) => void;
}): UseTimelineResult {
const connectorId = connectors.selectedConnector;
const hasConnector = !!connectorId;
const {
services: { notifications },
} = useKibana();
const conversationItems = useMemo(() => {
const items = getTimelineItemsfromConversation({
currentUser,
chatService,
hasConnector,
messages,
startedFrom,
});
return items;
}, [currentUser, chatService, hasConnector, messages, startedFrom]);
const [subscription, setSubscription] = useState<Subscription | undefined>();
const controllerRef = useRef(new AbortController());
const [pendingMessage, setPendingMessage] = useState<PendingMessage>();
const [isFunctionLoading, setIsFunctionLoading] = useState(false);
const prevConversationId = usePrevious(conversationId);
useEffect(() => {
if (prevConversationId !== conversationId && pendingMessage?.error) {
setPendingMessage(undefined);
}
}, [conversationId, pendingMessage?.error, prevConversationId]);
function chat(
nextMessages: Message[],
response$: Observable<PendingMessage> | undefined = undefined
): Promise<Message[]> {
const controller = new AbortController();
return new Promise<PendingMessage | undefined>(async (resolve, reject) => {
try {
if (!connectorId) {
reject(new Error('Can not add a message without a connector'));
return;
}
const isStartOfConversation =
nextMessages.some((message) => message.message.role === MessageRole.Assistant) === false;
if (isStartOfConversation && chatService.hasFunction('recall')) {
nextMessages = nextMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'recall',
arguments: JSON.stringify({ queries: [], contexts: [] }),
trigger: MessageRole.User,
},
},
});
}
onChatUpdate(nextMessages);
const lastMessage = last(nextMessages);
if (lastMessage?.message.function_call?.name) {
// the user has edited a function suggestion, no need to talk to the LLM
resolve(undefined);
return;
}
response$ =
response$ ||
chatService!.chat({
messages: nextMessages,
connectorId,
});
let pendingMessageLocal = pendingMessage;
const nextSubscription = response$.subscribe({
next: (nextPendingMessage) => {
pendingMessageLocal = nextPendingMessage;
setPendingMessage(() => nextPendingMessage);
},
error: reject,
complete: () => {
const error = pendingMessageLocal?.error;
if (error) {
notifications.toasts.addError(error, {
title: i18n.translate('xpack.observabilityAiAssistant.failedToLoadResponse', {
defaultMessage: 'Failed to load response from the AI Assistant',
}),
});
}
resolve(pendingMessageLocal!);
},
});
setSubscription(() => {
controllerRef.current = controller;
return nextSubscription;
});
} catch (error) {
reject(error);
}
}).then(async (reply) => {
if (reply?.error) {
return nextMessages;
}
if (reply?.aborted) {
return nextMessages;
}
setPendingMessage(undefined);
const messagesAfterChat = reply
? nextMessages.concat({
'@timestamp': new Date().toISOString(),
message: {
...reply.message,
},
})
: nextMessages;
onChatUpdate(messagesAfterChat);
const lastMessage = last(messagesAfterChat);
if (lastMessage?.message.function_call?.name) {
const name = lastMessage.message.function_call.name;
setIsFunctionLoading(true);
try {
let message = await chatService!.executeFunction({
name,
args: lastMessage.message.function_call.arguments,
messages: messagesAfterChat.slice(0, -1),
signal: controller.signal,
connectorId: connectorId!,
});
let nextResponse$: Observable<PendingMessage> | undefined;
if (isObservable(message)) {
nextResponse$ = message;
message = { content: '', data: '' };
}
return await chat(
messagesAfterChat.concat({
'@timestamp': new Date().toISOString(),
message: {
name,
role: MessageRole.User,
content: JSON.stringify(message.content),
data: JSON.stringify(message.data),
},
}),
nextResponse$
);
} catch (error) {
return await chat(
messagesAfterChat.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name,
content: JSON.stringify({
message: error.toString(),
error,
}),
},
})
);
} finally {
setIsFunctionLoading(false);
}
}
return messagesAfterChat;
});
}
const itemsWithAddedLoadingStates = useMemo(() => {
// While we're loading we add an empty loading chat item:
if (pendingMessage || isFunctionLoading) {
const nextItems = conversationItems.concat({
id: '',
actions: {
canCopy: true,
canEdit: false,
canGiveFeedback: false,
canRegenerate: pendingMessage?.aborted || !!pendingMessage?.error,
},
display: {
collapsed: false,
hide: pendingMessage?.message.role === MessageRole.System,
},
content: pendingMessage?.message.content,
currentUser,
error: pendingMessage?.error,
function_call: pendingMessage?.message.function_call,
loading: !pendingMessage?.aborted && !pendingMessage?.error,
role: pendingMessage?.message.role || MessageRole.Assistant,
title: '',
});
return nextItems;
}
if (!isFunctionLoading) {
return conversationItems;
}
return conversationItems.map((item, index) => {
// When we're done loading we remove the placeholder item again
if (index < conversationItems.length - 1) {
return item;
}
return {
...item,
loading: true,
};
});
}, [conversationItems, pendingMessage, currentUser, isFunctionLoading]);
const items = useMemo(() => {
const consolidatedChatItems: Array<ChatTimelineItem | ChatTimelineItem[]> = [];
let currentGroup: ChatTimelineItem[] | null = null;
for (const item of itemsWithAddedLoadingStates) {
if (item.display.hide || !item) continue;
if (item.display.collapsed) {
if (currentGroup) {
currentGroup.push(item);
} else {
currentGroup = [item];
consolidatedChatItems.push(currentGroup);
}
} else {
consolidatedChatItems.push(item);
currentGroup = null;
}
}
return consolidatedChatItems;
}, [itemsWithAddedLoadingStates]);
useEffect(() => {
return () => {
subscription?.unsubscribe();
};
}, [subscription]);
return {
items,
onEdit: async (item, newMessage) => {
const indexOf = flatten(items).indexOf(item);
const sliced = messages.slice(0, indexOf);
const nextMessages = await chat(sliced.concat(newMessage));
onChatComplete(nextMessages);
},
onFeedback: (item, feedback) => {},
onRegenerate: (item) => {
const indexOf = flatten(items).indexOf(item);
chat(messages.slice(0, indexOf)).then((nextMessages) => onChatComplete(nextMessages));
},
onStopGenerating: () => {
subscription?.unsubscribe();
setPendingMessage((prevPendingMessage) => ({
message: {
role: MessageRole.Assistant,
...prevPendingMessage?.message,
},
aborted: true,
error: new AbortError(),
}));
setSubscription(undefined);
},
onSubmit: async (message) => {
const nextMessages = await chat(messages.concat(message));
onChatComplete(nextMessages);
},
onActionClick: async (payload) => {
switch (payload.type) {
case ChatActionClickType.executeEsqlQuery:
const nextMessages = await chat(
messages.concat({
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'execute_query',
arguments: JSON.stringify({
query: payload.query,
}),
trigger: MessageRole.User,
},
},
})
);
onChatComplete(nextMessages);
break;
}
},
};
}

15
x-pack/plugins/observability_ai_assistant/public/routes/config.tsx
View file

@ -31,11 +31,18 @@ const observabilityAIAssistantRoutes = {
element: <ConversationView />,
},
'/conversations/{conversationId}': {
params: t.type({
path: t.type({
conversationId: t.string,
params: t.intersection([
t.type({
path: t.type({
conversationId: t.string,
}),
}),
}),
t.partial({
state: t.partial({
prevConversationKey: t.string,
}),
}),
]),
element: <ConversationView />,
},
'/conversations': {

160
x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx
View file

@ -4,36 +4,34 @@
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/
import React, { useMemo, useState } from 'react';
import { EuiCallOut, EuiFlexGroup, EuiFlexItem, EuiLoadingSpinner, EuiSpacer } from '@elastic/eui';
import { EuiFlexGroup, EuiFlexItem, EuiLoadingSpinner, EuiSpacer } from '@elastic/eui';
import { css } from '@emotion/css';
import { i18n } from '@kbn/i18n';
import { euiThemeVars } from '@kbn/ui-theme';
import React, { useMemo, useRef, useState } from 'react';
import usePrevious from 'react-use/lib/usePrevious';
import { v4 } from 'uuid';
import { ChatBody } from '../../components/chat/chat_body';
import { ConversationList } from '../../components/chat/conversation_list';
import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
import { useAbortableAsync } from '../../hooks/use_abortable_async';
import { useConfirmModal } from '../../hooks/use_confirm_modal';
import { useConversation } from '../../hooks/use_conversation';
import { useCurrentUser } from '../../hooks/use_current_user';
import { useForceUpdate } from '../../hooks/use_force_update';
import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
import { useKibana } from '../../hooks/use_kibana';
import { useKnowledgeBase } from '../../hooks/use_knowledge_base';
import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
import { useObservabilityAIAssistantParams } from '../../hooks/use_observability_ai_assistant_params';
import { useObservabilityAIAssistantRouter } from '../../hooks/use_observability_ai_assistant_router';
import { EMPTY_CONVERSATION_TITLE } from '../../i18n';
import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
import { getModelsManagementHref } from '../../utils/get_models_management_href';
import { EMPTY_CONVERSATION_TITLE } from '../../i18n';
const containerClassName = css`
max-width: 100%;
`;
const chatBodyContainerClassNameWithError = css`
align-self: center;
`;
const conversationListContainerName = css`
min-width: 250px;
width: 250px;
@ -80,12 +78,24 @@ export function ConversationView() {
const conversationId = 'conversationId' in path ? path.conversationId : undefined;
const { conversation, displayedMessages, setDisplayedMessages, save, saveTitle } =
useConversation({
conversationId,
chatService: chatService.value,
connectorId: connectors.selectedConnector,
});
// Regenerate the key only when the id changes, except after
// creating the conversation. Ideally this happens by adding
// state to the current route, but I'm not keen on adding
// the concept of state to the router, due to a mismatch
// between router.link() and router.push(). So, this is a
// pretty gross workaround for persisting a key under some
// conditions.
const chatBodyKeyRef = useRef(v4());
const keepPreviousKeyRef = useRef(false);
const prevConversationId = usePrevious(conversationId);
if (conversationId !== prevConversationId && keepPreviousKeyRef.current === false) {
chatBodyKeyRef.current = v4();
}
keepPreviousKeyRef.current = false;
const forceUpdate = useForceUpdate();
const conversations = useAbortableAsync(
({ signal }) => {
@ -111,14 +121,17 @@ export function ConversationView() {
];
}, [conversations.value?.conversations, conversationId, observabilityAIAssistantRouter]);
function navigateToConversation(nextConversationId?: string) {
observabilityAIAssistantRouter.push(
nextConversationId ? '/conversations/{conversationId}' : '/conversations/new',
{
path: { conversationId: nextConversationId },
function navigateToConversation(nextConversationId?: string, usePrevConversationKey?: boolean) {
if (nextConversationId) {
observabilityAIAssistantRouter.push('/conversations/{conversationId}', {
path: {
conversationId: nextConversationId,
},
query: {},
}
);
});
} else {
observabilityAIAssistantRouter.push('/conversations/new', { path: {}, query: {} });
}
}
function handleRefreshConversations() {
@ -136,10 +149,16 @@ export function ConversationView() {
error={conversations.error}
conversations={displayedConversations}
onClickNewChat={() => {
observabilityAIAssistantRouter.push('/conversations/new', {
path: {},
query: {},
});
if (conversationId) {
observabilityAIAssistantRouter.push('/conversations/new', {
path: {},
query: {},
});
} else {
// clear the chat
chatBodyKeyRef.current = v4();
forceUpdate();
}
}}
onClickDeleteConversation={(id) => {
confirmDeleteFunction()
@ -194,69 +213,36 @@ export function ConversationView() {
/>
<EuiSpacer size="s" />
</EuiFlexItem>
<EuiFlexItem
grow
className={conversation.error ? chatBodyContainerClassNameWithError : undefined}
>
{conversation.error ? (
<EuiCallOut
color="danger"
title={i18n.translate(
'xpack.observabilityAiAssistant.couldNotFindConversationTitle',
{
defaultMessage: 'Conversation not found',
{!chatService.value ? (
<EuiFlexGroup direction="column" alignItems="center" gutterSize="l">
<EuiFlexItem grow={false}>
<EuiSpacer size="xl" />
<EuiLoadingSpinner size="l" />
</EuiFlexItem>
</EuiFlexGroup>
) : null}
{chatService.value && (
<ObservabilityAIAssistantChatServiceProvider value={chatService.value}>
<ChatBody
key={chatBodyKeyRef.current}
currentUser={currentUser}
connectors={connectors}
connectorsManagementHref={getConnectorsManagementHref(http)}
modelsManagementHref={getModelsManagementHref(http)}
initialConversationId={conversationId}
knowledgeBase={knowledgeBase}
startedFrom="conversationView"
onConversationUpdate={(conversation) => {
if (!conversationId) {
keepPreviousKeyRef.current = true;
navigateToConversation(conversation.conversation.id);
}
)}
iconType="warning"
>
{i18n.translate('xpack.observabilityAiAssistant.couldNotFindConversationContent', {
defaultMessage:
'Could not find a conversation with id {conversationId}. Make sure the conversation exists and you have access to it.',
values: { conversationId },
})}
</EuiCallOut>
) : null}
{!chatService.value ? (
<EuiFlexGroup direction="column" alignItems="center" gutterSize="l">
<EuiFlexItem grow={false}>
<EuiSpacer size="xl" />
<EuiLoadingSpinner size="l" />
</EuiFlexItem>
</EuiFlexGroup>
) : null}
{conversation.value && chatService.value && !conversation.error ? (
<ObservabilityAIAssistantChatServiceProvider value={chatService.value}>
<ChatBody
loading={conversation.loading}
currentUser={currentUser}
connectors={connectors}
connectorsManagementHref={getConnectorsManagementHref(http)}
modelsManagementHref={getModelsManagementHref(http)}
conversationId={conversationId}
knowledgeBase={knowledgeBase}
messages={displayedMessages}
title={conversation.value.conversation.title}
startedFrom="conversationView"
onChatUpdate={(messages) => {
setDisplayedMessages(messages);
}}
onChatComplete={(messages) => {
save(messages, handleRefreshConversations)
.then((nextConversation) => {
conversations.refresh();
if (!conversationId && nextConversation?.conversation?.id) {
navigateToConversation(nextConversation.conversation.id);
}
})
.catch((e) => {});
}}
onSaveTitle={(title) => {
saveTitle(title, handleRefreshConversations);
}}
/>
</ObservabilityAIAssistantChatServiceProvider>
) : null}
</EuiFlexItem>
handleRefreshConversations();
}}
/>
</ObservabilityAIAssistantChatServiceProvider>
)}
</EuiFlexGroup>
</>
);

24
x-pack/plugins/observability_ai_assistant/public/service/create_mock_chat_service.ts Normal file
View file

@ -0,0 +1,24 @@
/*
* 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 type { DeeplyMockedKeys } from '@kbn/utility-types-jest';
import type { ObservabilityAIAssistantChatService } from '../types';
type MockedChatService = DeeplyMockedKeys<ObservabilityAIAssistantChatService>;
export const createMockChatService = (): MockedChatService => {
const mockChatService: MockedChatService = {
chat: jest.fn(),
executeFunction: jest.fn(),
getContexts: jest.fn().mockReturnValue([{ name: 'core', description: '' }]),
getFunctions: jest.fn().mockReturnValue([]),
hasFunction: jest.fn().mockReturnValue(false),
hasRenderFunction: jest.fn().mockReturnValue(true),
renderFunction: jest.fn(),
};
return mockChatService;
};

160
x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
View file

@ -6,106 +6,98 @@
*/
import { merge, uniqueId } from 'lodash';
import { MessageRole, Conversation, FunctionDefinition } from '../../common/types';
import { ChatTimelineItem } from '../components/chat/chat_timeline';
import { DeepPartial } from 'utility-types';
import { MessageRole, Conversation, FunctionDefinition, Message } from '../../common/types';
import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
type ChatItemBuildProps = Omit<Partial<ChatTimelineItem>, 'actions' | 'display' | 'currentUser'> & {
actions?: Partial<ChatTimelineItem['actions']>;
display?: Partial<ChatTimelineItem['display']>;
currentUser?: Partial<ChatTimelineItem['currentUser']>;
} & Pick<ChatTimelineItem, 'role'>;
type BuildMessageProps = DeepPartial<Message> & {
message: {
role: MessageRole;
function_call?: {
name: string;
trigger: MessageRole.Assistant | MessageRole.User | MessageRole.Elastic;
};
};
};
export function buildChatItem(params: ChatItemBuildProps): ChatTimelineItem {
export function buildMessage(params: BuildMessageProps): Message {
return merge(
{
id: uniqueId(),
title: '',
actions: {
canCopy: true,
canEdit: false,
canGiveFeedback: false,
canRegenerate: params.role === MessageRole.Assistant,
},
display: {
collapsed: false,
hide: false,
},
currentUser: {
username: 'elastic',
},
loading: false,
'@timestamp': new Date().toISOString(),
},
params
);
}
export function buildSystemChatItem(params?: Omit<ChatItemBuildProps, 'role'>) {
return buildChatItem({
role: MessageRole.System,
...params,
});
export function buildSystemMessage(
params?: Omit<BuildMessageProps, 'message'> & {
message: DeepPartial<Omit<Message['message'], 'role'>>;
}
) {
return buildMessage(
merge({}, params, {
message: { role: MessageRole.System },
})
);
}
export function buildChatInitItem() {
return buildChatItem({
role: MessageRole.User,
title: 'started a conversation',
actions: {
canEdit: false,
canCopy: true,
canGiveFeedback: false,
canRegenerate: false,
},
});
export function buildUserMessage(
params?: Omit<BuildMessageProps, 'message'> & {
message?: DeepPartial<Omit<Message['message'], 'role'>>;
}
) {
return buildMessage(
merge(
{
message: {
content: "What's a function?",
},
},
params,
{
message: { role: MessageRole.User },
}
)
);
}
export function buildUserChatItem(params?: Omit<ChatItemBuildProps, 'role'>) {
return buildChatItem({
role: MessageRole.User,
content: "What's a function?",
actions: {
canCopy: true,
canEdit: true,
canGiveFeedback: false,
canRegenerate: true,
},
...params,
});
export function buildAssistantMessage(
params?: Omit<BuildMessageProps, 'message'> & {
message: DeepPartial<Omit<Message['message'], 'role'>>;
}
) {
return buildMessage(
merge(
{
message: {
content: `In computer programming and mathematics, a function is a fundamental concept that represents a relationship between input values and output values. It takes one or more input values (also known as arguments or parameters) and processes them to produce a result, which is the output of the function. The input values are passed to the function, and the function performs a specific set of operations or calculations on those inputs to produce the desired output.
A function is often defined with a name, which serves as an identifier to call and use the function in the code. It can be thought of as a reusable block of code that can be executed whenever needed, and it helps in organizing code and making it more modular and maintainable.`,
},
},
params,
{
message: { role: MessageRole.Assistant },
}
)
);
}
export function buildAssistantChatItem(params?: Omit<ChatItemBuildProps, 'role'>) {
return buildChatItem({
role: MessageRole.Assistant,
content: `In computer programming and mathematics, a function is a fundamental concept that represents a relationship between input values and output values. It takes one or more input values (also known as arguments or parameters) and processes them to produce a result, which is the output of the function. The input values are passed to the function, and the function performs a specific set of operations or calculations on those inputs to produce the desired output.
A function is often defined with a name, which serves as an identifier to call and use the function in the code. It can be thought of as a reusable block of code that can be executed whenever needed, and it helps in organizing code and making it more modular and maintainable.`,
actions: {
canCopy: true,
canEdit: false,
canRegenerate: true,
canGiveFeedback: true,
},
...params,
});
}
export function buildFunctionChatItem(params: Omit<ChatItemBuildProps, 'role'>) {
return buildChatItem({
role: MessageRole.User,
title: 'executed a function',
function_call: {
name: 'leftpad',
arguments: '{ foo: "bar" }',
trigger: MessageRole.Assistant,
},
...params,
});
}
export function buildTimelineItems() {
return {
items: [buildSystemChatItem(), buildUserChatItem(), buildAssistantChatItem()],
};
export function buildFunctionResponseMessage(
params?: Omit<BuildMessageProps, 'message'> & {
message: DeepPartial<Omit<Message['message'], 'role'>>;
}
) {
return buildUserMessage(
merge(
{},
{
message: {
name: 'leftpad',
},
...params,
}
)
);
}
export function buildConversation(params?: Partial<Conversation>) {

597
x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.test.tsx Normal file
View file

@ -0,0 +1,597 @@
/*
* 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 React from 'react';
import { last, pick } from 'lodash';
import { render } from '@testing-library/react';
import { Message, MessageRole } from '../../common';
import { createMockChatService } from '../service/create_mock_chat_service';
import { getTimelineItemsfromConversation } from './get_timeline_items_from_conversation';
import { __IntlProvider as IntlProvider } from '@kbn/i18n-react';
import { ObservabilityAIAssistantChatServiceProvider } from '../context/observability_ai_assistant_chat_service_provider';
import { ChatState } from '../hooks/use_chat';
const mockChatService = createMockChatService();
let items: ReturnType<typeof getTimelineItemsfromConversation>;
describe('getTimelineItemsFromConversation', () => {
describe('returns an opening message only', () => {
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
messages: [],
chatState: ChatState.Ready,
});
expect(items.length).toBe(1);
expect(items[0].title).toBe('started a conversation');
});
describe('with a start of a conversation', () => {
beforeEach(() => {
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
currentUser: {
username: 'johndoe',
full_name: 'John Doe',
},
chatState: ChatState.Ready,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'User',
},
},
],
});
});
it('excludes the system message', () => {
expect(items.length).toBe(2);
expect(items[0].title).toBe('started a conversation');
});
it('includes the rest of the conversation', () => {
expect(items[1].currentUser?.full_name).toEqual('John Doe');
expect(items[1].content).toEqual('User');
});
it('formats the user message', () => {
expect(pick(items[1], 'title', 'actions', 'display', 'loading')).toEqual({
title: '',
actions: {
canCopy: true,
canEdit: true,
canGiveFeedback: false,
canRegenerate: false,
},
display: {
collapsed: false,
hide: false,
},
loading: false,
});
});
});
describe('with function calling', () => {
beforeEach(() => {
mockChatService.hasRenderFunction.mockImplementation(() => false);
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
chatState: ChatState.Ready,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
function_call: {
name: 'recall',
arguments: JSON.stringify({ queries: [], contexts: [] }),
trigger: MessageRole.Assistant,
},
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name: 'recall',
content: JSON.stringify([]),
},
},
],
});
});
it('formats the function request', () => {
expect(pick(items[2], 'actions', 'display', 'loading')).toEqual({
actions: {
canCopy: true,
canEdit: true,
canGiveFeedback: false,
canRegenerate: true,
},
display: {
collapsed: true,
hide: false,
},
loading: false,
});
const { container } = render(items[2].title as React.ReactElement, {
wrapper: ({ children }) => (
<IntlProvider locale="en" messages={{}}>
{children}
</IntlProvider>
),
});
expect(container.textContent).toBe('requested the function recall');
});
it('formats the function response', () => {
expect(pick(items[3], 'actions', 'display', 'loading')).toEqual({
actions: {
canCopy: true,
canEdit: false,
canGiveFeedback: false,
canRegenerate: false,
},
display: {
collapsed: true,
hide: false,
},
loading: false,
});
const { container } = render(items[3].title as React.ReactElement, {
wrapper: ({ children }) => (
<IntlProvider locale="en" messages={{}}>
{children}
</IntlProvider>
),
});
expect(container.textContent).toBe('executed the function recall');
});
});
describe('with a render function', () => {
beforeEach(() => {
mockChatService.hasRenderFunction.mockImplementation(() => true);
mockChatService.renderFunction.mockImplementation(() => 'Rendered');
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
chatState: ChatState.Ready,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
function_call: {
name: 'my_render_function',
arguments: JSON.stringify({ foo: 'bar' }),
trigger: MessageRole.Assistant,
},
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name: 'my_render_function',
content: JSON.stringify([]),
},
},
],
});
});
it('renders a display element', () => {
expect(mockChatService.hasRenderFunction).toHaveBeenCalledWith('my_render_function');
expect(pick(items[3], 'actions', 'display')).toEqual({
actions: {
canCopy: true,
canEdit: false,
canGiveFeedback: false,
canRegenerate: false,
},
display: {
collapsed: false,
hide: false,
},
});
expect(items[3].element).toBeTruthy();
const { container } = render(items[3].element as React.ReactElement, {
wrapper: ({ children }) => (
<IntlProvider locale="en" messages={{}}>
<ObservabilityAIAssistantChatServiceProvider value={mockChatService}>
{children}
</ObservabilityAIAssistantChatServiceProvider>
</IntlProvider>
),
});
expect(mockChatService.renderFunction).toHaveBeenCalledWith(
'my_render_function',
JSON.stringify({ foo: 'bar' }),
{ content: '[]', name: 'my_render_function', role: 'user' }
);
expect(container.textContent).toEqual('Rendered');
});
});
describe('with a function that errors out', () => {
beforeEach(() => {
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
chatState: ChatState.Ready,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Hello',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
function_call: {
name: 'my_render_function',
arguments: JSON.stringify({ foo: 'bar' }),
trigger: MessageRole.Assistant,
},
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name: 'my_render_function',
content: JSON.stringify({
error: {
message: 'An error occurred',
},
}),
},
},
],
});
});
it('returns a title that reflects a failure to execute the function', () => {
const { container } = render(items[3].title as React.ReactElement, {
wrapper: ({ children }) => (
<IntlProvider locale="en" messages={{}}>
{children}
</IntlProvider>
),
});
expect(container.textContent).toBe('failed to execute the function my_render_function');
});
it('formats the messages correctly', () => {
expect(pick(items[3], 'actions', 'display', 'loading')).toEqual({
actions: {
canCopy: true,
canEdit: false,
canGiveFeedback: false,
canRegenerate: false,
},
display: {
collapsed: true,
hide: false,
},
loading: false,
});
});
});
describe('with an invalid JSON response', () => {
beforeEach(() => {
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
currentUser: {
username: 'johndoe',
full_name: 'John Doe',
},
chatState: ChatState.Ready,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
function_call: {
name: 'my_function',
arguments: JSON.stringify({}),
trigger: MessageRole.User,
},
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'invalid-json',
name: 'my_function',
},
},
],
});
});
it('sets the invalid json as content', () => {
expect(items[2].content).toBe(
`\`\`\`
{
"content": "invalid-json"
}
\`\`\``
);
});
});
describe('when starting from a contextual insight', () => {
beforeEach(() => {
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
currentUser: {
username: 'johndoe',
full_name: 'John Doe',
},
chatState: ChatState.Ready,
startedFrom: 'contextualInsight',
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Test',
},
},
],
});
});
it('hides the first user message', () => {
expect(items[1].display.collapsed).toBe(true);
});
});
describe('with function calling suggested by the user', () => {
beforeEach(() => {
mockChatService.hasRenderFunction.mockImplementation(() => false);
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
chatState: ChatState.Ready,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
function_call: {
name: 'recall',
arguments: JSON.stringify({ queries: [], contexts: [] }),
trigger: MessageRole.User,
},
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name: 'recall',
content: JSON.stringify([]),
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: 'Reply from assistant',
},
},
],
});
});
it('formats the function request', () => {
expect(pick(items[1], 'actions', 'display')).toEqual({
actions: {
canCopy: true,
canRegenerate: false,
canEdit: true,
canGiveFeedback: false,
},
display: {
collapsed: true,
hide: false,
},
});
});
it('formats the assistant response', () => {
expect(pick(items[3], 'actions', 'display')).toEqual({
actions: {
canCopy: true,
canRegenerate: true,
canEdit: false,
canGiveFeedback: false,
},
display: {
collapsed: false,
hide: false,
},
});
});
});
describe('while the chat is loading', () => {
const renderWithLoading = (extraMessages: Message[]) => {
items = getTimelineItemsfromConversation({
chatService: mockChatService,
hasConnector: true,
chatState: ChatState.Loading,
messages: [
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.System,
content: 'System',
},
},
{
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
content: 'Test',
},
},
...extraMessages,
],
});
};
describe('with a user message last', () => {
beforeEach(() => {
renderWithLoading([]);
});
it('adds an assistant message which is loading', () => {
expect(pick(last(items), 'display', 'actions', 'loading', 'role', 'content')).toEqual({
loading: true,
role: MessageRole.Assistant,
actions: {
canCopy: false,
canRegenerate: false,
canEdit: false,
canGiveFeedback: false,
},
display: {
collapsed: false,
hide: false,
},
content: '',
});
});
});
describe('with a function request as the last message', () => {
beforeEach(() => {
renderWithLoading([
{
'@timestamp': new Date().toISOString(),
message: {
function_call: {
name: 'my_function_call',
trigger: MessageRole.Assistant,
},
role: MessageRole.Assistant,
},
},
]);
});
it('adds an assistant message which is loading', () => {
expect(pick(last(items), 'display', 'actions', 'loading', 'role', 'content')).toEqual({
loading: true,
role: MessageRole.Assistant,
actions: {
canCopy: false,
canRegenerate: false,
canEdit: false,
canGiveFeedback: false,
},
display: {
collapsed: false,
hide: false,
},
content: '',
});
});
});
});
});

102
x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.tsx
View file

@ -6,7 +6,7 @@
*/
import React from 'react';
import { v4 } from 'uuid';
import { isEmpty, omitBy } from 'lodash';
import { isEmpty, last, omitBy } from 'lodash';
import { useEuiTheme } from '@elastic/eui';
import { i18n } from '@kbn/i18n';
import { FormattedMessage } from '@kbn/i18n-react';
@ -15,6 +15,15 @@ import { Message, MessageRole } from '../../common';
import type { ChatTimelineItem } from '../components/chat/chat_timeline';
import { RenderFunction } from '../components/render_function';
import type { ObservabilityAIAssistantChatService } from '../types';
import { ChatState } from '../hooks/use_chat';
function safeParse(jsonStr: string) {
try {
return JSON.parse(jsonStr);
} catch (err) {
return jsonStr;
}
}
function convertMessageToMarkdownCodeBlock(message: Message['message']) {
let value: object;
@ -22,7 +31,7 @@ function convertMessageToMarkdownCodeBlock(message: Message['message']) {
if (!message.name) {
const name = message.function_call?.name;
const args = message.function_call?.arguments
? JSON.parse(message.function_call.arguments)
? safeParse(message.function_call.arguments)
: undefined;
value = {
@ -32,9 +41,9 @@ function convertMessageToMarkdownCodeBlock(message: Message['message']) {
} else {
const content =
message.role !== MessageRole.Assistant && message.content
? JSON.parse(message.content)
? safeParse(message.content)
: message.content;
const data = message.data ? JSON.parse(message.data) : undefined;
const data = message.data ? safeParse(message.data) : undefined;
value = omitBy(
{
content,
@ -61,26 +70,36 @@ export function getTimelineItemsfromConversation({
hasConnector,
messages,
startedFrom,
chatState,
}: {
chatService: ObservabilityAIAssistantChatService;
currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
hasConnector: boolean;
messages: Message[];
startedFrom?: StartedFrom;
chatState: ChatState;
}): ChatTimelineItem[] {
return [
const messagesWithoutSystem = messages.filter(
(message) => message.message.role !== MessageRole.System
);
const items: ChatTimelineItem[] = [
{
id: v4(),
actions: { canCopy: false, canEdit: false, canGiveFeedback: false, canRegenerate: false },
display: { collapsed: false, hide: false },
currentUser,
loading: false,
role: MessageRole.User,
message: {
'@timestamp': new Date().toISOString(),
message: { role: MessageRole.User },
},
title: i18n.translate('xpack.observabilityAiAssistant.conversationStartTitle', {
defaultMessage: 'started a conversation',
}),
role: MessageRole.User,
},
...messages.map((message, index) => {
...messagesWithoutSystem.map((message, index) => {
const id = v4();
let title: React.ReactNode = '';
@ -88,8 +107,10 @@ export function getTimelineItemsfromConversation({
let element: React.ReactNode | undefined;
const prevFunctionCall =
message.message.name && messages[index - 1] && messages[index - 1].message.function_call
? messages[index - 1].message.function_call
message.message.name &&
messagesWithoutSystem[index - 1] &&
messagesWithoutSystem[index - 1].message.function_call
? messagesWithoutSystem[index - 1].message.function_call
: undefined;
let role = message.message.function_call?.trigger || message.message.role;
@ -107,10 +128,6 @@ export function getTimelineItemsfromConversation({
};
switch (role) {
case MessageRole.System:
display.hide = true;
break;
case MessageRole.User:
actions.canCopy = true;
actions.canGiveFeedback = false;
@ -120,16 +137,15 @@ export function getTimelineItemsfromConversation({
// User executed a function:
if (message.message.name && prevFunctionCall) {
let parsedContent;
let isError = false;
try {
parsedContent = JSON.parse(message.message.content ?? 'null');
const parsedContent = JSON.parse(message.message.content ?? 'null');
isError =
parsedContent && typeof parsedContent === 'object' && 'error' in parsedContent;
} catch (error) {
parsedContent = message.message.content;
isError = true;
}
const isError =
parsedContent && typeof parsedContent === 'object' && 'error' in parsedContent;
title = !isError ? (
<FormattedMessage
id="xpack.observabilityAiAssistant.userExecutedFunctionEvent"
@ -190,7 +206,7 @@ export function getTimelineItemsfromConversation({
display.collapsed = false;
if (startedFrom === 'contextualInsight') {
const firstUserMessageIndex = messages.findIndex(
const firstUserMessageIndex = messagesWithoutSystem.findIndex(
(el) => el.message.role === MessageRole.User
);
@ -252,7 +268,53 @@ export function getTimelineItemsfromConversation({
currentUser,
function_call: message.message.function_call,
loading: false,
message,
};
}),
];
const isLoading = chatState === ChatState.Loading;
let lastMessage = last(items);
const isNaturalLanguageOnlyAnswerFromAssistant =
lastMessage?.message.message.role === MessageRole.Assistant &&
!lastMessage.message.message.function_call?.name;
const addLoadingPlaceholder = isLoading && !isNaturalLanguageOnlyAnswerFromAssistant;
if (addLoadingPlaceholder) {
items.push({
id: v4(),
actions: {
canCopy: false,
canEdit: false,
canGiveFeedback: false,
canRegenerate: false,
},
display: {
collapsed: false,
hide: false,
},
content: '',
currentUser,
loading: chatState === ChatState.Loading,
role: MessageRole.Assistant,
title: '',
message: {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.Assistant,
content: '',
},
},
});
lastMessage = last(items);
}
if (isLoading && lastMessage) {
lastMessage.loading = isLoading;
}
return items;
}

3
x-pack/plugins/observability_ai_assistant/tsconfig.json
View file

@ -46,7 +46,8 @@
"@kbn/es-query",
"@kbn/rule-registry-plugin",
"@kbn/licensing-plugin",
"@kbn/share-plugin"
"@kbn/share-plugin",
"@kbn/utility-types-jest"
],
"exclude": ["target/**/*"]
}