Here is my messageRow and messageContent files respcetively: import { usePreviousDistinct } from ‘@react-hookz/web’;
import { forwardRef, useEffect, useState, useRef } from ‘react’;
import { useLongPress } from ‘react-aria’;

import { Avatar } from ‘@/components/Avatar’;
import IconButton from ‘@/components/IconButton’;
import { LongPressMenu } from ‘@/components/LongPressMenu’;
import { MessageContent } from ‘@/components/MessageContent’;
import {
Button,
CopyToClipboardButton,
CopyToClipboardIconButton,
Icon,
Markdown,
Tooltip,
} from ‘@/components/Shared’;
import { ToolEvents } from ‘@/components/ToolEvents’;
import { ReservedClasses } from ‘@/constants’;
import { Breakpoint, useBreakpoint } from ‘@/hooks/breakpoint’;
import { getMessageRowId } from ‘@/hooks/citations’;
import { useCitationsStore } from ‘@/stores’;
import {
type ChatMessage,
isAbortedMessage,
isErroredMessage,
isFulfilledMessage,
isFulfilledOrTypingMessage,
isFulfilledOrTypingMessageWithCitations,
isUserMessage,
Annotation,
} from ‘@/types/message’;
import { cn } from ‘@/utils’;

import { STYLE_LEVEL_TO_CLASSES } from ‘@/components/Shared’;

type Props = {
isLast: boolean;
message: ChatMessage;
delay?: boolean;
className?: string;
onCopy?: VoidFunction;
onRetry?: VoidFunction;
};

/**
* Renders a single message row from the user or from our models.
*/
const MessageRow = forwardRef<HTMLDivElement, Props>(function MessageRowInternal(
{ message, delay = false, isLast, className = ‘’, onCopy, onRetry },
ref
) {
const breakpoint = useBreakpoint();

const [isShowing, setIsShowing] = useState(false);
const [isLongPressMenuOpen, setIsLongPressMenuOpen] = useState(false);
const [isStepsExpanded, setIsStepsExpanded] = useState<boolean>(isLast);
const {
citations: { selectedCitation, hoveredGenerationId },
hoverCitation,
} = useCitationsStore();
const hasSteps =
(isFulfilledOrTypingMessage(message) ||
isErroredMessage(message) ||
isAbortedMessage(message)) &&
!!message.toolEvents &&
message.toolEvents.length > 0;

const getMessageText = () => {
if (isFulfilledMessage(message)) {
return message.originalText;
}

return message.text;
};

const enableLongPress =
(isFulfilledMessage(message) || isUserMessage(message)) && breakpoint === Breakpoint.sm;
const { longPressProps } = useLongPress({
onLongPress: () => setIsLongPressMenuOpen(true),
});

// Delay the appearance of the message to make it feel more natural.
useEffect(() => {
if (delay) {
setTimeout(() => setIsShowing(true), 300);
}
}, []);

useEffect(() => {
if (isLast) {
setIsStepsExpanded(true);
}
}, [isLast]);

const [highlightMessage, setHighlightMessage] = useState(false);
const prevSelectedCitationGenId = usePreviousDistinct(selectedCitation?.generationId);

useEffect(() => {
if (isFulfilledOrTypingMessage(message) && message.citations && message.generationId) {
if (
selectedCitation?.generationId === message.generationId &&
prevSelectedCitationGenId !== message.generationId
) {
setHighlightMessage(true);

setTimeout(() => {
setHighlightMessage(false);
}, 1000);
}
}
}, [selectedCitation?.generationId, prevSelectedCitationGenId]);

if (delay && !isShowing) return null;

const handleOnMouseEnter = () => {
if (isFulfilledOrTypingMessageWithCitations(message)) {
hoverCitation(message.generationId);
}
};

const handleOnMouseLeave = () => {
if (isFulfilledOrTypingMessageWithCitations(message)) {
hoverCitation(null);
}
};

//

// Annotations code for error highlighing in prompts!!!

//


// State to manage the selected text and annotations
const [selectedText, setSelectedText] = useState<string>(“”);
const [annotations, setAnnotations] = useState<Annotation[]>([]);
const inputRef = useRef<HTMLInputElement>(null);

// Function to add an annotation
const addAnnotation = (annotationText: string, index: number) => {
const newAnnotations = annotations.map((annotation, i) =>
i === index ? { …annotation, annotation: annotationText } : annotation
);

setAnnotations(newAnnotations);
setSelectedText(‘’);
const inputBox = inputRef.current;
if (inputBox) inputBox.style.display = “none”;
};

// Function to handle text selection
// Function to handle text selection
const handleMouseUp = () => {
const selection = window.getSelection();
if (selection && selection.rangeCount > 0) {
const range = selection.getRangeAt(0);
if (range && selection.toString().length > 0) {
const start = range.startOffset;
const end = range.endOffset;

setAnnotations([…annotations, {
text: selection.toString(),
annotation: ‘’,
start,
end
}]);

setSelectedText(selection.toString());

const boundingRect = range.getBoundingClientRect();
const inputBox = inputRef.current;
if (inputBox) {
const containerRect = range.commonAncestorContainer.parentElement!.getBoundingClientRect();
inputBox.style.position = ‘absolute’;
inputBox.style.top = ${boundingRect.bottom - containerRect.top}px;
inputBox.style.left = ${boundingRect.left - containerRect.left}px;
inputBox.style.display = ‘block’;
inputBox.focus();
}
}
}
};


return (
<div
id={
isFulfilledOrTypingMessage(message) && message.generationId
? getMessageRowId(message.generationId)
: undefined
}
className={cn(ReservedClasses.MESSAGE, ‘flex’, className)}
onMouseEnter={handleOnMouseEnter}
onMouseLeave={handleOnMouseLeave}
ref={ref}
>
<LongPressMenu
isOpen={isLongPressMenuOpen}
close={() => setIsLongPressMenuOpen(false)}
className=“md:hidden”
>
<div className={cn(‘flex flex-col divide-y’, ‘divide-marble-300’)}>
<div className=“flex flex-col gap-y-4 pt-4”>
<CopyToClipboardButton
value={getMessageText()}
label=“Copy text”
kind=“secondary”
iconAtStart
onClick={onCopy}
/>
{hasSteps && (
<Button
label={${isStepsExpanded ? 'Hide' : 'Show'} steps}
startIcon={<Icon name=“list” />}
kind=“secondary”
size=“md”
aria-label={${isStepsExpanded ? 'Hide' : 'Show'} steps}
animate={false}
onClick={() => setIsStepsExpanded((prevIsExpanded) => !prevIsExpanded)}
/>
)}
</div>
</div>
</LongPressMenu>
<div
className={cn(
‘group flex h-fit w-full flex-col gap-2 rounded-md p-2 text-left md:flex-row’,
‘transition-colors ease-in-out’,
‘hover:bg-secondary-100’,

{
‘bg-secondary-50’:
isFulfilledOrTypingMessage(message) &&
message.generationId &&
hoveredGenerationId === message.generationId,
‘bg-primary-50 hover:bg-primary-50’: highlightMessage,
}
)}
{…(enableLongPress && longPressProps)}
>
<div className=“flex w-full gap-x-2”>
<Avatar message={message} />
<div className=“flex w-full min-w-0 max-w-message flex-1 flex-col items-center gap-x-3 md:flex-row”>
<div className=“w-full”>
{hasSteps && <ToolEvents show={isStepsExpanded} events={message.toolEvents} />}

{//This is where we render the annotated box
}

<MessageContent
isLast={isLast}
message={message}
annotations={annotations}
onAnnotate={handleMouseUp}
onRetry={onRetry}
/>

{inputRef.current && (
<input
ref={inputRef}
type=‘text’
style={{
fontSize: ‘14px’, // Adjust the font size as needed
height: ‘2rem’, // Adjust the height as needed
width: ‘200px’, // Adjust the width as needed
padding: ‘0.5rem’, // Adjust the padding as needed
lineHeight: ‘1.5’, // Adjust the line height as needed
fontFamily: ‘Arial, sans-serif’ // Adjust the font family as needed
}}
className={cn(
‘min-h-[1rem] md:min-h-[2rem]’,
“w-auto”,
‘self-center’,
‘rounded’,
‘px-1 px-2’,
‘border’,
‘bg-danger-50’,
‘text-lg’,

‘border-secondary-400’,
‘transition ease-in-out’,
‘focus:border-secondary-700’,
‘focus:outline-none’,

‘placeholder-base’

)}

placeholder=“Add annotation . . .”
onKeyDown={(e) => {
if (e.key === “Enter”) {
const target = e.target as HTMLInputElement;
const index = annotations.findIndex(a => a.text === selectedText);
if (index !== -1) {
addAnnotation(target.value, index);
}
target.value = “”;
}
}}
/>
)}

{annotations.map((annotation, index) => (
<div key={index} className=‘bg-secondary-50 p-1 my-1’>
<span className=‘font-bold’>{annotation.text}</span>: {annotation.annotation}
</div>
))}
</div>



<div
className={cn(‘flex h-full items-end justify-end self-end’, {
‘hidden md:invisible md:flex’:
!isFulfilledMessage(message) && !isUserMessage(message),
‘hidden md:invisible md:flex md:group-hover:visible’: !isLast,
})}
>
{hasSteps && (
<Tooltip label={${isStepsExpanded ? 'Hide' : 'Show'} steps} hover>
<IconButton
iconName=“list”
className=“rounded hover:bg-secondary-100”
iconClassName={cn(
‘text-volcanic-800 group-hover/icon-button:text-secondary-800’,
{
‘hidden md:invisible md:flex’: !isFulfilledMessage(message),
}
)}
onClick={() => setIsStepsExpanded((prevIsExpanded) => !prevIsExpanded)}
/>
</Tooltip>
)}
<CopyToClipboardIconButton value={getMessageText()} onClick={onCopy} />
</div>
</div>
</div>
</div>
</div>
);
});
export default MessageRow;
import { Transition } from ‘@headlessui/react’;
import { PropsWithChildren } from ‘react’;

import { CitationTextHighlighter } from ‘@/components/Citations/CitationTextHighlighter’;
import { DataTable } from ‘@/components/DataTable’;
import { MarkdownImage } from ‘@/components/MarkdownImage’;
import { Icon } from ‘@/components/Shared’;
import { Markdown, Text } from ‘@/components/Shared’;
import { UploadedFile } from ‘@/components/UploadedFile’;
import {
type ChatMessage,
MessageType,
isAbortedMessage,
isErroredMessage,
isFulfilledOrTypingMessage,
isLoadingMessage,
Annotation,
} from ‘@/types/message’;
import { cn } from ‘@/utils’;

import { useState, useRef } from ‘react’;
import { STYLE_LEVEL_TO_CLASSES } from ‘@/components/Shared’;
import internal from ‘stream’;

type Props = {
isLast: boolean;
message: ChatMessage;
annotations?: Annotation[];
onRetry?: VoidFunction;
onAnnotate?: (start: number, end: number) => void;
};

const BOT_ERROR_MESSAGE = 'Unable to generate a response since an error was encountered. ';

export const MessageContent: React.FC<Props> = ({
isLast,
message,
onRetry,
annotations = [],
onAnnotate
}) => {
const isUser = message.type === MessageType.USER;
const isLoading = isLoadingMessage(message);
const isBotError = isErroredMessage(message);
const isUserError = isUser && message.error;
const isAborted = isAbortedMessage(message);
const isTypingOrFulfilledMessage = isFulfilledOrTypingMessage(message);

const contentRef = useRef<HTMLDivElement>(null);

const handleMouseUp = () => {
const selection = window.getSelection();
if (selection && selection.toString()) {
const range = selection.getRangeAt(0);
const start = range.startOffset;
const end = range.endOffset;

// Call the annotation handler passed as props
if (onAnnotate) {
onAnnotate(start, end);
}
}
};

const renderAnnotatedText = (text: string) => {
if (annotations.length === 0) {
return <Markdown
text={text}
className={cn(STYLE_LEVEL_TO_CLASSES.p)}
/>
}

const parts: React.ReactNode[] = [];
let lastIndex = 0;

annotations.forEach(({ start, end, text: annotatedText }, index) => {
if (start !== -1) {
parts.push(
<span key={part-${index}-${lastIndex}}>{text.substring(lastIndex, start)}</span>,
<span key={index} className={cn(‘bg-yellow-100”, “annotated-text’)}>
{annotatedText}
</span>
);
lastIndex = end;
}
});

if (lastIndex < text.length) {
parts.push(<span key={part-${lastIndex}}>{text.substring(lastIndex)}</span>);
}

return parts;
};

let content: React.ReactNode = null;

if (isUserError) {
content = (
<>
<Text>{message.text}</Text>
<MessageInfo type=“error”>
{message.error}
{isLast && (
<button className=“underline underline-offset-1” type=“button” onClick={onRetry}>
Retry?
</button>
)}
</MessageInfo>
</>
);
// } else if (isUser) {
// content = (
// <div ref={contentRef} onMouseUp={handleMouseUp}>
// {renderAnnotatedText(message.text)}
// </div>
// );
} else if (isLoading) {
const hasLoadingMessage = message.text.length > 0;
content = (
<Text className={cn(‘flex min-w-0 text-volcanic-700’)} as=“span”>
{hasLoadingMessage && (
<Transition
appear={true}
show={true}
enterFrom=“opacity-0”
enterTo=“opacity-full”
enter=“transition-opacity ease-in-out duration-500”
>
{message.text}
</Transition>
)}
{!hasLoadingMessage && (
<span className=“w-max”>
<div className=“animate-typing-ellipsis overflow-hidden whitespace-nowrap pr-1”>
…
</div>
</span>
)}
</Text>
);
} else if (isBotError) {
content = (
<>
{message.text.length > 0 ? (
<Markdown text={message.text} />
) : (
<Text className={cn(‘text-volcanic-700’)}>{BOT_ERROR_MESSAGE}</Text>
)}
<MessageInfo type=“error”>{message.error}</MessageInfo>
</>
);
} else {
const hasCitations =
isTypingOrFulfilledMessage && message.citations && message.citations.length > 0;
content = (
<>
<div ref={contentRef} onMouseUp={handleMouseUp}>
{renderAnnotatedText(message.text)}
</div>

{isAborted && (
<MessageInfo>
This generation was stopped.{" "}
{isLast && isAborted && (
<button className=‘underline underline-offset-1’ type=‘button’ onClick={onRetry}>
Retry?
</button>
)}
</MessageInfo>
)}
</>
);
}

return (
<div className=“flex w-full flex-col justify-center gap-y-1 py-1”>
<Text
as=“div”
className=“flex flex-col gap-y-1 whitespace-pre-wrap [overflow-wrap:anywhere] md:max-w-4xl”
>
{content}
</Text>
</div>
);
};

const MessageInfo = ({
type = ‘default’,
children,
}: PropsWithChildren & { type?: ‘default’ | ‘error’ }) => (
<div
className={cn(‘flex items-start gap-1’, {
‘text-volcanic-700’: type === ‘default’,
‘text-danger-500’: type === ‘error’,
})}
>
<Icon name=“warning” size=“md” className=“flex items-center text-p” />
<Text as=“span”>{children}</Text>
</div>
);