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 rehypeRaw from 'rehype-raw';

import { useMemo, useState } from 'react';
import { cn } from '@/lib/utils';
import { CodeBlock } from '@/components/ui/CodeBlock';
import { MemoizedReactMarkdown } from '@/components/chat/MemoizedReactMarkdown';
import {
  IconCheckCircle,
  IconChevronDoubleRight,
  IconCodeWrap,
  IconCrossCircle,
  IconLandingAI,
  IconListUnordered,
  IconTerminalWindow,
  IconUser,
} from '@/components/ui/Icons';
import { MessageBase } from '../../lib/types';
import Img from '../ui/Img';
import { ChunkBody, CodeResult, formatStreamLogs } from '@/lib/messageUtils';
import {
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableHeader,
  TableRow,
} from '../ui/Table';
import { Button } from '../ui/Button';
import { Separator } from '../ui/Separator';
import {
  Tooltip,
  TooltipContent,
  TooltipTrigger,
} from '@/components/ui/Tooltip';

import { Dialog, DialogContent, DialogTrigger } from '../ui/Dialog';

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

const Markdown: React.FC<{
  content: string;
  setDetails?: (val: string) => void;
}> = ({ content, setDetails }) => {
  return (
    <>
      <MemoizedReactMarkdown
        className="break-words overflow-auto"
        remarkPlugins={[remarkGfm, remarkMath]}
        rehypePlugins={[rehypeRaw] as any}
        components={{
          table({ children, ...props }) {
            return <Table {...props}>{children}</Table>;
          },
          thead({ children, ...props }) {
            return <TableHeader {...props}>{children}</TableHeader>;
          },
          th({ children, ...props }) {
            return <TableHead {...props}>{children}</TableHead>;
          },
          tr({ children, ...props }) {
            return <TableRow {...props}>{children}</TableRow>;
          },
          td({ children, ...props }) {
            return <TableCell {...props}>{children}</TableCell>;
          },
          button({ children, ...props }) {
            if ('data-details' in props && setDetails) {
              return (
                <Button
                  {...props}
                  onClick={() =>
                    setDetails(decodeURI(props['data-details'] as string))
                  }
                >
                  {children}
                </Button>
              );
            }
            return <Button {...props}>{children}</Button>;
          },
          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="mb-2 whitespace-pre-line">{children}</p>;
          },
          img(props) {
            if (props.src?.endsWith('.mp4')) {
              return (
                <video src={props.src} controls width={500} height={500} />
              );
            }
            return (
              <Img
                src={props.src ?? '/landing.png'}
                alt={props.alt ?? 'answer-image'}
                quality={100}
                sizes="(min-width: 66em) 15vw,
                        (min-width: 44em) 20vw,
                        100vw"
              />
            );
          },
          code({ node, inline, className, children, ...props }) {
            const match = /language-(\w+)/.exec(className || '');

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

export function ChatMessage({ message, isLoading }: ChatMessageProps) {
  const { role, content } = message;

  return role === 'user' ? (
    <UserChatMessage content={content} />
  ) : (
    <AssistantChatMessage content={content} />
  );
}

const UserChatMessage: React.FC<{
  content: string;
}> = ({ content }) => {
  return (
    <div className="group relative mb-6 flex rounded-md bg-muted p-4 ml-auto mr-0 w-3/5">
      <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-background">
        <IconUser />
      </div>
      <div className="flex-1 px-1 ml-4 space-y-2 overflow-hidden">
        {content && <Markdown content={content} />}
      </div>
    </div>
  );
};

const ChunkStatusToIconDict: Record<ChunkBody['status'], React.ReactElement> = {
  started: <IconChevronDoubleRight className="text-cyan-500" />,
  completed: <IconCheckCircle className="text-green-500" />,
  running: <IconTerminalWindow className="text-teal-500" />,
  failed: <IconCrossCircle className="text-red-500" />,
};
const ChunkTypeToTextDict: Record<ChunkBody['type'], string> = {
  plans: 'Creating instructions',
  tools: 'Retrieving tools',
  code: 'Generating code',
  final_code: 'Final result',
};
const ChunkPayloadAction: React.FC<{
  payload: ChunkBody['payload'];
}> = ({ payload }) => {
  if (Array.isArray(payload)) {
    // [{title: 123, content, 345}, {title: ..., content: ...}] => ['title', 'content']
    const keyArray = Array.from(
      payload.reduce((acc, curr) => {
        Object.keys(curr).forEach(key => acc.add(key));
        return acc;
      }, new Set<string>()),
    );

    return (
      <Dialog>
        <DialogTrigger asChild>
          <Button variant="ghost" size="icon">
            <IconListUnordered />
          </Button>
        </DialogTrigger>
        <DialogContent className="max-w-5xl">
          <Table className="border rounded-lg bg-zinc-700 overflow-hidden">
            <TableHeader>
              <TableRow className="border-primary/50">
                {keyArray.map(header => (
                  <TableHead key={header}>{header}</TableHead>
                ))}
              </TableRow>
            </TableHeader>
            <TableBody>
              {payload.map((line, index) => (
                <TableRow className="border-primary/50" key={index}>
                  {keyArray.map(header => (
                    <TableCell key={header}>{line[header]}</TableCell>
                  ))}
                </TableRow>
              ))}
            </TableBody>
          </Table>
        </DialogContent>
      </Dialog>
    );
  } else {
    return (
      <Dialog>
        <DialogTrigger asChild>
          <Button variant="ghost" size="icon">
            <IconCodeWrap />
          </Button>
        </DialogTrigger>
        <DialogContent className="max-w-5xl">
          <CodeResultDisplay codeResult={payload as CodeResult} />
        </DialogContent>
      </Dialog>
    );
  }
};

const CodeResultDisplay: React.FC<{
  codeResult: CodeResult;
}> = ({ codeResult }) => {
  const { code, test, result } = codeResult;
  return (
    <div className="rounded-lg overflow-hidden relative max-w-5xl">
      <CodeBlock language="python" value={code} />
      <div className="rounded-lg relative">
        <Separator />
        <Tooltip>
          <TooltipTrigger asChild>
            <Button
              variant="ghost"
              size="icon"
              className="size-8 absolute left-1/2 -translate-x-1/2 -top-4 z-10"
            >
              <IconTerminalWindow className="text-teal-500 size-4" />
            </Button>
          </TooltipTrigger>
          <TooltipContent>
            <CodeBlock language="python" value={test} />
          </TooltipContent>
        </Tooltip>
      </div>
      <CodeBlock language="output" value={result} />
    </div>
  );
};

const AssistantChatMessage: React.FC<{
  content: string;
}> = ({ content }) => {
  const [formattedSections, codeResult] = useMemo(
    () => formatStreamLogs(content),
    [content],
  );

  return (
    <div className="group relative mb-6 flex rounded-md bg-muted p-4 w-full">
      <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-primary text-primary-foreground">
        <IconLandingAI />
      </div>
      <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
        <Table className="border rounded-lg bg-zinc-700 overflow-hidden w-[400px]">
          <TableBody>
            {formattedSections.map(section => (
              <TableRow className="border-primary/50" key={section.type}>
                <TableCell>{ChunkStatusToIconDict[section.status]}</TableCell>
                <TableCell className="font-medium">
                  {ChunkTypeToTextDict[section.type]}
                </TableCell>
                <TableCell className="text-right">
                  <ChunkPayloadAction payload={section.payload} />
                </TableCell>
              </TableRow>
            ))}
          </TableBody>
        </Table>
        {codeResult && <CodeResultDisplay codeResult={codeResult} />}
      </div>
    </div>
  );
};

export default UserChatMessage;