components/chat/ChatMessage.tsx

// Inspired by Chatbot-UI and modified to fit the needs of this project
// @see https://github.com/mckaywrigley/chatbot-ui/blob/main/components/Chat/ChatMessage.tsx

import remarkGfm from 'remark-gfm';
import remarkMath from 'remark-math';

import { useMemo } from 'react';
import { cn } from '@/lib/utils';
import { CodeBlock } from '@/components/ui/CodeBlock';
import { MemoizedReactMarkdown } from '@/components/chat/MemoizedReactMarkdown';
import { IconOpenAI, IconUser } from '@/components/ui/Icons';
import { MessageBase } from '../../lib/types';
import {
  Tooltip,
  TooltipContent,
  TooltipTrigger,
} from '@/components/ui/Tooltip';
import Img from '../ui/Img';
import { getCleanedUpMessages } from '@/lib/messageUtils';
import Loading from '../ui/Loading';

export interface ChatMessageProps {
  message: MessageBase;
  isLoading: boolean;
}

export function ChatMessage({
  message,
  isLoading,
  ...props
}: ChatMessageProps) {
  const { logs, content } = useMemo(() => {
    return getCleanedUpMessages({
      content: message.content,
      role: message.role,
    });
  }, [message.content, message.role]);
  return (
    <div className={cn('group relative mb-4 flex items-start')} {...props}>
      <div
        className={cn(
          'flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow',
          message.role === 'user'
            ? 'bg-background'
            : 'bg-primary text-primary-foreground',
        )}
      >
        {message.role === 'user' ? <IconUser /> : <IconOpenAI />}
      </div>
      <div className="flex-1 px-1 ml-4 space-y-2 overflow-hidden">
        {logs && message.role !== 'user' && (
          <div className="bg-slate-100 dark:bg-slate-900 mb-4 p-4 max-h-[400px] overflow-auto">
            <div className="text-xl font-bold">Thinking Process</div>
            <MemoizedReactMarkdown
              className="break-words text-sm"
              remarkPlugins={[remarkGfm, remarkMath]}
              components={{
                p({ children }) {
                  return (
                    <p className="my-2 last:mb-0 whitespace-pre-line">
                      {children}
                    </p>
                  );
                },
                code({ children, ...props }) {
                  return (
                    <code className="whitespace-pre-line">{children}</code>
                  );
                },
              }}
            >
              {logs}
            </MemoizedReactMarkdown>
          </div>
        )}
        <MemoizedReactMarkdown
          className="break-words"
          remarkPlugins={[remarkGfm, remarkMath]}
          components={{
            p({ children, ...props }) {
              if (
                props.node.children.some(
                  child => child.type === 'element' && child.tagName === 'img',
                )
              ) {
                return (
                  <p className="flex flex-wrap gap-2 items-start">{children}</p>
                );
              }
              return (
                <p className="my-2 last:mb-0 whitespace-pre-line">{children}</p>
              );
            },
            img(props) {
              if (props.src?.endsWith('.mp4')) {
                return (
                  <video src={props.src} controls width={500} height={500} />
                );
              }
              return (
                <Tooltip>
                  <TooltipTrigger asChild>
                    <Img
                      src={props.src ?? '/landing.png'}
                      alt={props.alt ?? 'answer-image'}
                      quality={100}
                      className="cursor-zoom-in"
                      sizes="(min-width: 66em) 25vw,
                      (min-width: 44em) 40vw,
                      100vw"
                    />
                  </TooltipTrigger>
                  <TooltipContent>
                    <Img
                      className="m-2"
                      src={props.src ?? '/landing.png'}
                      alt={props.alt ?? 'answer-image'}
                      quality={100}
                      width={500}
                    />
                  </TooltipContent>
                </Tooltip>
              );
            },
            code({ node, inline, className, children, ...props }) {
              if (children.length) {
                if (children[0] == '▍') {
                  return (
                    <span className="mt-1 cursor-default animate-pulse">▍</span>
                  );
                }

                children[0] = (children[0] as string).replace('`▍`', '▍');
              }

              const match = /language-(\w+)/.exec(className || '');
              if (inline) {
                return (
                  <code className={className} {...props}>
                    {children}
                  </code>
                );
              }

              return (
                <CodeBlock
                  key={Math.random()}
                  language={(match && match[1]) || ''}
                  value={String(children).replace(/\n$/, '')}
                  {...props}
                />
              );
            },
          }}
        >
          {content}
        </MemoizedReactMarkdown>
        {/* <ChatMessageActions message={message} /> */}
        {isLoading && <Loading />}
      </div>
    </div>
  );
}