Skip to content

Commit

Permalink
Blocks Renderers: extraction, cleanups, more maintainable and optimized
Browse files Browse the repository at this point in the history
  • Loading branch information
@enricoros
enricoros committed Feb 7, 2024
1 parent 7ba315c commit d62757d
Show file tree
Hide file tree
Showing 18 changed files with 300 additions and 213 deletions.
172 changes: 42 additions & 130 deletions src/apps/chat/components/message/ChatMessage.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,8 @@
import * as React from 'react';
import TimeAgo from 'react-timeago';
import { shallow } from 'zustand/shallow';
import { cleanupEfficiency, Diff as TextDiff, makeDiff } from '@sanity/diff-match-patch';

import { Avatar, Box, Button, CircularProgress, IconButton, ListDivider, ListItem, ListItemDecorator, MenuItem, Switch, Tooltip, Typography } from '@mui/joy';
import { SxProps } from '@mui/joy/styles/types';
import { Avatar, Box, CircularProgress, IconButton, ListDivider, ListItem, ListItemDecorator, MenuItem, Switch, Tooltip, Typography } from '@mui/joy';
import AccountTreeIcon from '@mui/icons-material/AccountTree';
import ClearIcon from '@mui/icons-material/Clear';
import ContentCopyIcon from '@mui/icons-material/ContentCopy';
Expand All @@ -23,30 +21,18 @@ import VerticalAlignBottomIcon from '@mui/icons-material/VerticalAlignBottom';

import { CloseableMenu } from '~/common/components/CloseableMenu';
import { DMessage } from '~/common/state/store-chats';
import { InlineError } from '~/common/components/InlineError';
import { InlineTextarea } from '~/common/components/InlineTextarea';
import { KeyStroke } from '~/common/components/KeyStroke';
import { Link } from '~/common/components/Link';
import { SystemPurposeId, SystemPurposes } from '../../../../data';
import { copyToClipboard } from '~/common/util/clipboardUtils';
import { cssRainbowColorKeyframes, lineHeightChatText } from '~/common/app.theme';
import { cssRainbowColorKeyframes } from '~/common/app.theme';
import { prettyBaseModel } from '~/common/util/modelUtils';
import { useUIPreferencesStore } from '~/common/state/store-ui';

import { BlocksRenderer, editBlocksSx } from './blocks/BlocksRenderer';
import { useChatShowTextDiff } from '../../store-app-chat';

import { RenderCode } from './RenderCode';
import { RenderHtml } from './RenderHtml';
import { RenderImage } from './RenderImage';
import { RenderLatex } from './RenderLatex';
import { RenderMarkdown } from './RenderMarkdown';
import { RenderText } from './RenderText';
import { RenderTextDiff } from './RenderTextDiff';
import { parseBlocks } from './blocks';


// How long is the user collapsed message
const USER_COLLAPSED_LINES: number = 8;

// Enable the menu on text selection
const ENABLE_SELECTION_RIGHT_CLICK_MENU: boolean = true;
Expand Down Expand Up @@ -206,13 +192,13 @@ export const ChatMessageMemo = React.memo(ChatMessage);
* or collapsing long user messages.
*
*/
export function ChatMessage(props: {
function ChatMessage(props: {
message: DMessage,
showDate?: boolean, diffPreviousText?: string,
hideAvatars?: boolean, codeBackground?: string,
noMarkdown?: boolean, diagramMode?: boolean,
isBottom?: boolean, noBottomBorder?: boolean,
isImagining?: boolean, isSpeaking?: boolean,
diffPreviousText?: string,
isBottom?: boolean,
isImagining?: boolean,
isSpeaking?: boolean,
blocksShowDate?: boolean,
onConversationBranch?: (messageId: string) => void,
onConversationRestartFrom?: (messageId: string, offset: number) => Promise<void>,
onConversationTruncate?: (messageId: string) => void,
Expand All @@ -221,11 +207,9 @@ export function ChatMessage(props: {
onTextDiagram?: (messageId: string, text: string) => Promise<void>
onTextImagine?: (text: string) => Promise<void>
onTextSpeak?: (text: string) => Promise<void>
sx?: SxProps,
}) {

// state
const [forceUserExpanded, setForceUserExpanded] = React.useState(false);
const [isHovering, setIsHovering] = React.useState(false);
const [opsMenuAnchor, setOpsMenuAnchor] = React.useState<HTMLElement | null>(null);
const [selMenuAnchor, setSelMenuAnchor] = React.useState<HTMLElement | null>(null);
Expand Down Expand Up @@ -257,10 +241,9 @@ export function ChatMessage(props: {

const fromAssistant = messageRole === 'assistant';
const fromSystem = messageRole === 'system';
const fromUser = messageRole === 'user';
const wasEdited = !!messageUpdated;

const showAvatars = props.hideAvatars !== true && !cleanerLooks;
const showAvatars = !cleanerLooks;

const textSel = selMenuText ? selMenuText : messageText;
const isSpecialT2I = textSel.startsWith('https://images.prodia.xyz/') || textSel.startsWith('/draw ') || textSel.startsWith('/imagine ') || textSel.startsWith('/img ');
Expand All @@ -275,8 +258,6 @@ export function ChatMessage(props: {
props.onMessageEdit(messageId, editedText);
};

const handleUncollapse = () => setForceUserExpanded(true);


// Operations Menu

Expand All @@ -289,12 +270,12 @@ export function ChatMessage(props: {
closeSelectionMenu();
};

const handleOpsEdit = (e: React.MouseEvent) => {
const handleOpsEdit = React.useCallback((e: React.MouseEvent) => {
if (messageTyping && !isEditing) return; // don't allow editing while typing
setIsEditing(!isEditing);
e.preventDefault();
closeOpsMenu();
};
}, [isEditing, messageTyping]);

const handleOpsConversationBranch = (e: React.MouseEvent) => {
e.preventDefault();
Expand Down Expand Up @@ -396,6 +377,17 @@ export function ChatMessage(props: {
}, [openSelectionMenu]);


// Blocks renderer

const handleBlocksContextMenu = React.useCallback((event: React.MouseEvent) => {
handleMouseUp(event.nativeEvent);
}, [handleMouseUp]);

const handleBlocksDoubleClick = React.useCallback((event: React.MouseEvent) => {
doubleClickToEdit && props.onMessageEdit && handleOpsEdit(event);
}, [doubleClickToEdit, handleOpsEdit, props.onMessageEdit]);


// prettier upstream errors
const { isAssistantError, errorMessage } = React.useMemo(
() => explainErrorInMessage(messageText, fromAssistant, messageOriginLLM),
Expand All @@ -411,50 +403,19 @@ export function ChatMessage(props: {
[messageAvatar, messageOriginLLM, messagePurposeId, messageRole, messageSender, messageTyping, showAvatars],
);

// per-blocks css
const blockSx: SxProps = {
my: 'auto',
lineHeight: lineHeightChatText,
};
const typographySx: SxProps = {
lineHeight: lineHeightChatText,
};
const codeSx: SxProps = {
// backgroundColor: fromAssistant ? 'background.level1' : 'background.level1',
backgroundColor: props.codeBackground ? props.codeBackground : fromAssistant ? 'neutral.plainHoverBg' : 'primary.plainActiveBg',
boxShadow: 'xs',
fontFamily: 'code',
fontSize: '0.875rem',
fontVariantLigatures: 'none',
lineHeight: lineHeightChatText,
borderRadius: 'var(--joy-radius-sm)',
};

// user message truncation
let collapsedText = messageText;
let isCollapsed = false;
if (fromUser && !forceUserExpanded) {
const lines = messageText.split('\n');
if (lines.length > USER_COLLAPSED_LINES) {
collapsedText = lines.slice(0, USER_COLLAPSED_LINES).join('\n');
isCollapsed = true;
}
}


return (
<ListItem
sx={{
display: 'flex', flexDirection: !fromAssistant ? 'row-reverse' : 'row', alignItems: 'flex-start',
gap: { xs: 0, md: 1 }, px: { xs: 1, md: 2 }, py: 2,
gap: { xs: 0, md: 1 },
px: { xs: 1, md: 2 },
py: 2,
backgroundColor,
...(props.noBottomBorder !== true && {
borderBottom: '1px solid',
borderBottomColor: 'divider',
}),
borderBottom: '1px solid',
borderBottomColor: 'divider',
...(ENABLE_COPY_MESSAGE_OVERLAY && { position: 'relative' }),
'&:hover > button': { opacity: 1 },
...props.sx,
}}
>

Expand Down Expand Up @@ -500,73 +461,24 @@ export function ChatMessage(props: {

<InlineTextarea
initialText={messageText} onEdit={handleTextEdited}
sx={{
...blockSx,
flexGrow: 1,
}} />
sx={editBlocksSx}
/>

) : (

<Box
onContextMenu={(ENABLE_SELECTION_RIGHT_CLICK_MENU && props.onMessageEdit) ? event => handleMouseUp(event.nativeEvent) : undefined}
onDoubleClick={event => (doubleClickToEdit && props.onMessageEdit) ? handleOpsEdit(event) : null}
sx={{
...blockSx,
flexGrow: 0,
overflowX: 'auto',
...(!!props.diagramMode && {
// width: '100%',
boxShadow: 'md',
}),
}}>

{props.showDate === true && (
<Typography level='body-sm' sx={{ mx: 1.5, textAlign: fromAssistant ? 'left' : 'right' }}>
<TimeAgo date={messageUpdated || messageCreated} />
</Typography>
)}

{/* Warn about user-edited system message */}
{fromSystem && wasEdited && (
<Typography level='body-sm' color='warning' sx={{ mt: 1, mx: 1.5 }}>modified by user - auto-update disabled</Typography>
)}

{errorMessage && (
<Tooltip title={<Typography sx={{ maxWidth: 800 }}>{collapsedText}</Typography>} variant='soft'>
<InlineError error={errorMessage} />
</Tooltip>
)}
<BlocksRenderer
text={messageText}
fromRole={messageRole}
renderTextAsMarkdown={renderMarkdown}
errorMessage={errorMessage}
isBottom={props.isBottom}
showDate={props.blocksShowDate === true ? messageUpdated || messageCreated || undefined : undefined}
textDiffs={textDiffs || undefined}
wasUserEdited={wasEdited}
onContextMenu={(props.onMessageEdit && ENABLE_SELECTION_RIGHT_CLICK_MENU) ? handleBlocksContextMenu : undefined}
onDoubleClick={(props.onMessageEdit && doubleClickToEdit) ? handleBlocksDoubleClick : undefined}
/>

{/* sequence of render components, for each Block */}
{!errorMessage && parseBlocks(collapsedText, fromSystem, textDiffs)
.filter((block, _, blocks) => !props.diagramMode || block.type === 'code' || blocks.length === 1)
.map(
(block, index) =>
block.type === 'html'
? <RenderHtml key={'html-' + index} htmlBlock={block} sx={codeSx} />
: block.type === 'code'
? <RenderCode key={'code-' + index} codeBlock={block} sx={codeSx} noCopyButton={props.diagramMode} />
: block.type === 'image'
? <RenderImage key={'image-' + index} imageBlock={block} isFirst={!index} allowRunAgain={props.isBottom === true} onRunAgain={handleOpsConversationRestartFrom} />
: block.type === 'latex'
? <RenderLatex key={'latex-' + index} latexBlock={block} sx={typographySx} />
: block.type === 'diff'
? <RenderTextDiff key={'latex-' + index} diffBlock={block} sx={typographySx} />
: (renderMarkdown && props.noMarkdown !== true && !fromSystem && !(fromUser && block.content.startsWith('/')))
? <RenderMarkdown key={'text-md-' + index} textBlock={block} />
: <RenderText key={'text-' + index} textBlock={block} sx={typographySx} />)}

{isCollapsed && (
<Button variant='plain' color='neutral' onClick={handleUncollapse}>... expand ...</Button>
)}

{/* import VisibilityIcon from '@mui/icons-material/Visibility'; */}
{/*<br />*/}
{/*<Chip variant='outlined' color='warning' sx={{ mt: 1, fontSize: '0.75em' }} startDecorator={<VisibilityIcon />}>*/}
{/* BlockAction*/}
{/*</Chip>*/}

</Box>
)}


Expand Down
Loading

0 comments on commit d62757d

Please sign in to comment.