Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Security solution] Elastic Assistant, OpenAI stream #169729

Merged
merged 70 commits into from
Nov 7, 2023
Merged
Changes from 1 commit
Commits
Show all changes
70 commits
Select commit Hold shift + click to select a range
e88a785
wip
stephmilovic Oct 13, 2023
1e8c7f9
Merge branch 'main' into ai_stream
stephmilovic Oct 13, 2023
556e4bd
more wip
stephmilovic Oct 13, 2023
bed104b
wip
stephmilovic Oct 16, 2023
343bef6
Merge branch 'main' into ai_stream
stephmilovic Oct 16, 2023
939afd7
rm logs
stephmilovic Oct 16, 2023
693fcae
wip
stephmilovic Oct 16, 2023
f58a5e1
its streaming baby!
stephmilovic Oct 16, 2023
a36253c
Merge branch 'main' into ai_stream
stephmilovic Oct 18, 2023
ccfb94e
merge
stephmilovic Oct 18, 2023
1ebf9b6
update w main
stephmilovic Oct 18, 2023
ed49c2d
Merge branch 'main' into ai_stream
stephmilovic Oct 19, 2023
c37a1d7
an attempt
stephmilovic Oct 19, 2023
fa00236
rm extra files
stephmilovic Oct 19, 2023
f2b354c
rm more
stephmilovic Oct 19, 2023
eb856ac
console logs
stephmilovic Oct 19, 2023
8a71d4e
saved to memory and stop button works
stephmilovic Oct 19, 2023
6d14e5e
rm logs
stephmilovic Oct 19, 2023
b821018
wip
stephmilovic Oct 23, 2023
682d69a
use stream for all comments
stephmilovic Oct 23, 2023
4fe7b81
regenerate button done
stephmilovic Oct 23, 2023
e8fc9f0
fixing
stephmilovic Oct 24, 2023
a6c749e
test
stephmilovic Oct 24, 2023
b6c976a
align regenerate button right
stephmilovic Oct 24, 2023
80b9ea1
Merge branch 'main' into ai_stream
stephmilovic Oct 24, 2023
ebaab59
fix
stephmilovic Oct 24, 2023
d52e9cf
[CI] Auto-commit changed files from 'node scripts/eslint --no-cache -…
kibanamachine Oct 24, 2023
c4f8a8e
fixing
stephmilovic Oct 25, 2023
d806f99
Merge branch 'ai_stream' of github.com:stephmilovic/kibana into ai_st…
stephmilovic Oct 25, 2023
87b07a6
Merge branch 'main' into ai_stream
stephmilovic Oct 25, 2023
cd8a3b5
trying to make it better but made it worse
stephmilovic Oct 25, 2023
2dcb57d
stash dumb changes
stephmilovic Oct 25, 2023
358881a
add comment
stephmilovic Oct 26, 2023
523ecb5
Merge branch 'main' into ai_stream
stephmilovic Oct 26, 2023
5154e32
fix stream to be pretty
stephmilovic Oct 26, 2023
db30dd8
add comments
stephmilovic Oct 26, 2023
6b79554
cleanup
stephmilovic Oct 26, 2023
b701323
better UX
stephmilovic Oct 26, 2023
8ff7c3b
cleanup
stephmilovic Oct 26, 2023
e41cd26
Merge branch 'ai_stream' of github.com:stephmilovic/kibana into ai_st…
stephmilovic Oct 30, 2023
f3258a4
update amendMessage
stephmilovic Oct 30, 2023
30f2624
tests wip
stephmilovic Oct 30, 2023
8781484
stream observable tests
stephmilovic Oct 30, 2023
b7ae9e9
Merge branch 'main' into ai_stream
stephmilovic Oct 31, 2023
a34a980
fix tests
stephmilovic Oct 31, 2023
14d7ce1
more tests
stephmilovic Oct 31, 2023
cc389fe
fix types
stephmilovic Oct 31, 2023
99c526d
more tests
stephmilovic Oct 31, 2023
db40b5a
fix types in test
stephmilovic Oct 31, 2023
2e92ed8
type and lint fixes
stephmilovic Oct 31, 2023
5cf4f06
more fixing
stephmilovic Oct 31, 2023
487a166
Merge branch 'main' into ai_stream
stephmilovic Oct 31, 2023
fcc4dc7
transform message added to stream
stephmilovic Oct 31, 2023
8a6c547
cleanup
stephmilovic Oct 31, 2023
480ef87
fix
stephmilovic Oct 31, 2023
d92c804
fix type
stephmilovic Oct 31, 2023
1c6d412
PR fixing
stephmilovic Nov 3, 2023
471e5a0
Merge branch 'main' into ai_stream
stephmilovic Nov 3, 2023
cae753a
solves weird streaming issue
stephmilovic Nov 3, 2023
d72bd00
lint fix
stephmilovic Nov 3, 2023
61c6ed4
comment code
stephmilovic Nov 3, 2023
301aa51
fix type
stephmilovic Nov 3, 2023
826048f
Merge branch 'main' into ai_stream
stephmilovic Nov 3, 2023
a920258
update w main
stephmilovic Nov 3, 2023
c5ce1aa
Merge branch 'main' into ai_stream
stephmilovic Nov 6, 2023
40e2b79
fixed
stephmilovic Nov 6, 2023
be53e17
fix tests
stephmilovic Nov 6, 2023
9a2330b
fix scrolling
stephmilovic Nov 6, 2023
7b0d156
Merge branch 'main' into ai_stream
kibanamachine Nov 6, 2023
78bf5e2
Merge branch 'main' into ai_stream
kibanamachine Nov 7, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
regenerate button done
  • Loading branch information
stephmilovic committed Oct 23, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
commit 4fe7b815192386a4209cbeb8529ec3941c2dbb9a
Original file line number Diff line number Diff line change
@@ -34,6 +34,7 @@ interface UseChatSend {
handleOnChatCleared: () => void;
handlePromptChange: (prompt: string) => void;
handleSendMessage: (promptText: string) => void;
handleRegenerateResponse: () => void;
isLoading: boolean;
}

@@ -54,7 +55,8 @@ export const useChatSend = ({
setUserPrompt,
}: UseChatSendProps): UseChatSend => {
const { isLoading, sendMessages } = useSendMessages();
const { appendMessage, appendReplacements, clearConversation } = useConversation();
const { appendMessage, appendReplacements, clearConversation, removeLastMessage } =
useConversation();

const handlePromptChange = (prompt: string) => {
setPromptTextPreview(prompt);
@@ -112,6 +114,24 @@ export const useChatSend = ({
]
);

const handleRegenerateResponse = useCallback(async () => {
const updatedMessages = removeLastMessage(currentConversation.id);
const rawResponse = await sendMessages({
Comment on lines +118 to +120
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it's related to the regenerate response action, but I am sometimes seeing the scroll to bottom behavior not working. Can't repro 100% yet, but will add steps once I can (nearly positive I didn't touch the mouse/keyboard, so shouldn't be an interrupt here).

http,
apiConfig: currentConversation.apiConfig,
messages: updatedMessages,
});
const responseMessage: Message = getMessageFromRawResponse(rawResponse);
appendMessage({ conversationId: currentConversation.id, message: responseMessage });
}, [
appendMessage,
currentConversation.apiConfig,
currentConversation.id,
http,
removeLastMessage,
sendMessages,
]);

const handleButtonSendMessage = useCallback(
(message: string) => {
handleSendMessage(message);
@@ -146,6 +166,7 @@ export const useChatSend = ({
handleOnChatCleared,
handlePromptChange,
handleSendMessage,
handleRegenerateResponse,
isLoading,
};
};
15 changes: 15 additions & 0 deletions x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
Original file line number Diff line number Diff line change
@@ -31,6 +31,7 @@ import { css } from '@emotion/react';

import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/common/openai/constants';
import { ActionConnectorProps } from '@kbn/triggers-actions-ui-plugin/public/types';
import { useChatSend } from './chat_send/use_chat_send';
import { ChatSend } from './chat_send';
import { BlockBotCallToAction } from './block_bot/cta';
import { AssistantHeader } from './assistant_header';
@@ -339,6 +340,18 @@ const AssistantComponent: React.FC<Props> = ({
[messageCodeBlocks]
);

const { handleRegenerateResponse } = useChatSend({
allSystemPrompts,
currentConversation,
setPromptTextPreview,
setUserPrompt,
editingSystemPromptId,
http,
setEditingSystemPromptId,
selectedPromptContexts,
setSelectedPromptContexts,
});

const chatbotComments = useMemo(
() => (
<>
@@ -348,6 +361,7 @@ const AssistantComponent: React.FC<Props> = ({
lastCommentRef,
showAnonymizedValues,
amendMessage,
regenerateMessage: handleRegenerateResponse,
})}
css={css`
margin-right: 20px;
@@ -382,6 +396,7 @@ const AssistantComponent: React.FC<Props> = ({
editingSystemPromptId,
getComments,
handleOnSystemPromptSelectionChange,
handleRegenerateResponse,
isSettingsModalVisible,
promptContexts,
promptTextPreview,
Original file line number Diff line number Diff line change
@@ -60,7 +60,6 @@ interface SetConversationProps {
}

interface UseConversation {
appendStreamMessage: ({ conversationId, message }: AppendMessageProps) => void;
appendMessage: ({ conversationId, message }: AppendMessageProps) => Message[];
amendMessage: ({ conversationId, content }: AmendMessageProps) => Message[];
appendReplacements: ({
@@ -70,6 +69,7 @@ interface UseConversation {
clearConversation: (conversationId: string) => void;
createConversation: ({ conversationId, messages }: CreateConversationProps) => Conversation;
deleteConversation: (conversationId: string) => void;
removeLastMessage: (conversationId: string) => Message[];
setApiConfig: ({ conversationId, apiConfig }: SetApiConfigProps) => void;
setConversation: ({ conversation }: SetConversationProps) => void;
}
@@ -78,21 +78,52 @@ export const useConversation = (): UseConversation => {
const { allSystemPrompts, assistantTelemetry, setConversations } = useAssistantContext();

/**
* Replaces the last message of conversation[] for a given conversationId
* Removes the last message of conversation[] for a given conversationId
*/
const removeLastMessage = useCallback(
(conversationId: string) => {
let messages: Message[] = [];
setConversations((prev: Record<string, Conversation>) => {
const prevConversation: Conversation | undefined = prev[conversationId];

if (prevConversation != null) {
prevConversation.messages.pop();
messages = prevConversation.messages;
const newConversation = {
...prevConversation,
messages,
};
console.log('removeLastMessage newConversation', newConversation);
return {
...prev,
[conversationId]: newConversation,
};
} else {
return prev;
}
});
return messages;
},
[setConversations]
);

/**
* Updates the last message of conversation[] for a given conversationId with provided content
*/
const amendMessage = useCallback(
({ conversationId, content }) => {
({ conversationId, content }: AmendMessageProps) => {
let messages: Message[] = [];
setConversations((prev: Record<string, Conversation>) => {
const prevConversation: Conversation | undefined = prev[conversationId];

if (prevConversation != null) {
const message = prevConversation.messages.pop();
const message = prevConversation.messages.pop() as unknown as Message;
stephmilovic marked this conversation as resolved.
Show resolved Hide resolved
messages = [...prevConversation.messages, { ...message, content }];
const newConversation = {
...prevConversation,
messages,
};
console.log('amendMessage newConversation', newConversation);
return {
...prev,
[conversationId]: newConversation,
@@ -303,6 +334,7 @@ export const useConversation = (): UseConversation => {
clearConversation,
createConversation,
deleteConversation,
removeLastMessage,
setApiConfig,
setConversation,
};
Original file line number Diff line number Diff line change
@@ -13,7 +13,6 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react';
import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public';
import { useLocalStorage } from 'react-use';
import type { DocLinksStart } from '@kbn/core-doc-links-browser';
import { AppendMessageProps } from '../assistant/use_conversation';
import { updatePromptContexts } from './helpers';
import type {
PromptContext,
@@ -109,7 +108,14 @@ export interface UseAssistantContext {
}: {
currentConversation: Conversation;
lastCommentRef: React.MutableRefObject<HTMLDivElement | null>;
amendMessage: ({ conversationId: string, message: Message }: AppendMessageProps) => Message[];
amendMessage: ({
conversationId,
content,
}: {
conversationId: string;
content: string;
}) => Message[];
regenerateMessage: () => void;
showAnonymizedValues: boolean;
}) => EuiCommentProps[];
http: HttpSetup;
Original file line number Diff line number Diff line change
@@ -22,6 +22,7 @@ export const getComments = ({
amendMessage,
currentConversation,
lastCommentRef,
regenerateMessage,
showAnonymizedValues,
}: {
amendMessage: ({
@@ -33,6 +34,7 @@ export const getComments = ({
}) => Message[];
currentConversation: Conversation;
lastCommentRef: React.MutableRefObject<HTMLDivElement | null>;
regenerateMessage: (conversationId: string) => void;
showAnonymizedValues: boolean;
}): EuiCommentProps[] => {
const amendMessageOfConversation = (content: string) => {
@@ -42,6 +44,10 @@ export const getComments = ({
});
};

const regenerateMessageOfConversation = () => {
regenerateMessage(currentConversation.id);
};

return currentConversation.messages.map((message, index) => {
const isUser = message.role === 'user';
const replacements = currentConversation.replacements;
@@ -70,18 +76,22 @@ export const getComments = ({
username: isUser ? i18n.YOU : i18n.ASSISTANT,
...(message.isError ? errorStyles : {}),
};
const isLastComment = index === currentConversation.messages.length - 1;

// message still needs to stream, no response manipulation
if (!(message.content && message.content.length)) {
return {
...messageProps,
children: (
<StreamComment
amendMessage={amendMessageOfConversation}
reader={message.reader}
index={index}
lastCommentRef={lastCommentRef}
isLastComment={index === currentConversation.messages.length - 1}
/>
<>
<StreamComment
amendMessage={amendMessageOfConversation}
reader={message.reader}
regenerateMessage={regenerateMessageOfConversation}
isLastComment={isLastComment}
/>
{isLastComment ? <span ref={lastCommentRef} /> : null}
</>
),
};
}
@@ -107,18 +117,12 @@ export const getComments = ({
amendMessage={amendMessageOfConversation}
content={showAnonymizedValues ? message.content : transformedMessage.content}
reader={message.reader}
lastCommentRef={lastCommentRef}
isLastComment={index === currentConversation.messages.length - 1}
regenerateMessage={regenerateMessageOfConversation}
isLastComment={isLastComment}
/>
{index !== currentConversation.messages.length - 1 ? null : <span ref={lastCommentRef} />}
{isLastComment ? <span ref={lastCommentRef} /> : null}
</>
),
};

// return {
// ...messageProps,
// children: <EuiText>{'oh no an error happened'}</EuiText>,
// ...errorStyles,
// };
});
};
Original file line number Diff line number Diff line change
@@ -17,80 +17,55 @@ interface Props {
amendMessage: (message: string) => void;
content?: string;
isLastComment: boolean;
lastCommentRef: React.MutableRefObject<HTMLDivElement | null>;
regenerateMessage: () => void;
reader?: ReadableStreamDefaultReader<Uint8Array>;
}

export const StreamComment = ({
amendMessage,
content,
isLastComment,
lastCommentRef,
reader,
regenerateMessage,
}: Props) => {
const { error, isLoading, isStreaming, pendingMessage, setComplete, subscription } = useStream({
const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({
amendMessage,
content,
reader,
});
const message = content ?? pendingMessage;
const controls = useMemo(
() =>
reader != null ? (
isLoading || isStreaming ? (
<StopGeneratingButton
onClick={() => {
subscription?.unsubscribe();
setComplete(true);
console.log('stop generating');
// 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(),
// }));
}}
/>
) : isLastComment ? (
<EuiFlexGroup direction="row">
<EuiFlexItem grow={false}>
<RegenerateResponseButton
onClick={() => {
console.log('RegenerateResponseButton');
// reloadRecalledMessages();
}}
/>
</EuiFlexItem>
</EuiFlexGroup>
) : null
) : null,
[isLastComment, isLoading, isStreaming, reader, setComplete, subscription]
);
const controls = useMemo(() => {
if (reader == null || !isLastComment) {
return;
}
if (isLoading || isStreaming) {
return (
<StopGeneratingButton
onClick={() => {
setComplete(true);
}}
/>
);
}
return (
<EuiFlexGroup direction="row">
<EuiFlexItem grow={false}>
<RegenerateResponseButton onClick={regenerateMessage} />
</EuiFlexItem>
</EuiFlexGroup>
);
}, [isLastComment, isLoading, isStreaming, reader, regenerateMessage, setComplete]);
return (
<>
<MessagePanel
body={
<MessageText
content={message}
loading={isLoading || isStreaming}
onActionClick={async () => {}}
/>
}
error={error ? new Error(error) : undefined}
controls={controls}
/>
{isLastComment && <span ref={lastCommentRef} />}
</>
<MessagePanel
body={
<MessageText
content={message}
loading={isLoading || isStreaming}
onActionClick={async () => {}}
/>
}
error={error ? new Error(error) : undefined}
controls={controls}
/>
);
};
Loading