src/chat/Chat.tsx

import { Button, Card, Slider } from "@theme";
import { useChatSettings } from "@utils/context/chatSettings/useChatSettings.ts";
import { formatBytes } from "@utils/format.ts";
import { CONVERSATION_STARTERS, MODELS } from "@utils/models.ts";
import { TOOLS } from "@utils/tools.ts";
import { RotateCcw, Settings } from "lucide-react";
import {
  useEffect,
  useMemo,
  useRef,
  useState,
  useSyncExternalStore,
} from "react";

import TextGeneration from "../textGeneration/TextGeneration.ts";
import type {
  ChatMessageAssistant,
  ChatMessageUser,
} from "../textGeneration/types.ts";
import cn from "../utils/classnames.ts";
import ChatForm from "./ChatForm.tsx";
import Message from "./Message.tsx";

enum State {
  IDLE,
  INITIALIZING,
  READY,
  GENERATING,
}

export default function Chat({ className = "" }: { className?: string }) {
  const { openSettingsModal, settings, downloadedModels } = useChatSettings();
  const initializedModelKey = useRef<string>(null);
  const [downloadProgress, setDownloadProgress] = useState<number>(0);
  const [state, setState] = useState<State>(State.IDLE);
  const settingsKey = useRef<string>(null);
  const [starterCategory, setStarterCategory] = useState<number>(0);
  const messagesEndRef = useRef<HTMLDivElement>(null);

  const generator = useMemo(() => new TextGeneration(), []);
  const getSnapshot = () => generator.chatMessages;
  const messages = useSyncExternalStore(
    generator.onChatMessageUpdate,
    getSnapshot
  );

  useEffect(() => {
    const newSettingsKey = JSON.stringify(settings);
    if (!settings || settingsKey.current === newSettingsKey) return;
    settingsKey.current = newSettingsKey;
    initializeConversation();
  }, [settings]);

  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  }, [messages]);

  if (!settings) return;

  const modelDownloaded = downloadedModels.includes(settings.modelKey);

  const initializeConversation = () =>
    generator.initializeConversation(
      TOOLS.filter((tool) => settings.tools.includes(tool.name)),
      settings.temperature,
      settings.enableThinking,
      settings.systemPrompt
    );

  const initializeModel = async () => {
    setState(State.INITIALIZING);
    setDownloadProgress(0);
    await generator.initializeModel(settings.modelKey, (percentage) =>
      setDownloadProgress(percentage)
    );
    initializedModelKey.current = settings.modelKey;
    setState(State.READY);
  };

  const generate = async (prompt: string) => {
    if (initializedModelKey.current !== settings.modelKey) {
      await initializeModel();
    }
    setState(State.GENERATING);
    await generator.runAgent(prompt);

    setState(State.READY);
  };

  const model = MODELS[settings.modelKey];
  const ready: boolean = state === State.READY || modelDownloaded;

  const conversationMessages = messages.filter(({ role }) => role !== "system");

  return (
    <>
      <div
        className={cn(
          className,
          "mx-auto min-h-screen w-full max-w-4xl px-4 pt-24 pb-36"
        )}
      >
        {state === State.IDLE && !modelDownloaded ? (
          <Card className="mt-26 flex max-w-[500px] flex-col gap-4 self-center justify-self-center">
            <p>
              You are about to load <b>{model.title}</b>.<br />
              Once downloaded, the model ({formatBytes(model.size)}) will be
              cached and reused when you revisit the page.
            </p>
            <p>
              Everything runs directly in your browser using 🤗 Transformers.js
              and ONNX Runtime Web, meaning your conversations aren't sent to a
              server. You can even disconnect from the internet after the model
              has loaded!
            </p>
            <Button onClick={initializeModel}>
              Download Model ({formatBytes(model.size)})
            </Button>
          </Card>
        ) : state === State.INITIALIZING ? (
          <div className="mt-46 flex h-full items-center justify-center">
            <div className="flex h-full w-full max-w-[400px] flex-col items-center justify-center gap-2">
              <p className="flex w-full items-center justify-between">
                <span>
                  {modelDownloaded
                    ? "initializing the model..."
                    : "downloading the model..."}
                </span>
                <span>{downloadProgress.toFixed(2)}%</span>
              </p>
              <Slider width={Math.round(downloadProgress)} />
            </div>
          </div>
        ) : conversationMessages.length === 0 ? (
          <div className="mt-26 flex h-full items-center justify-center">
            <div className="flex flex-col gap-8">
              <h2 className="text-2xl font-bold">How can I help you?</h2>
              <div className="flex gap-4">
                {CONVERSATION_STARTERS.map((starter, index) => (
                  <Button
                    iconLeft={starter.icon}
                    key={index}
                    onClick={() => setStarterCategory(index)}
                    variant={index === starterCategory ? "solid" : "outline"}
                  >
                    {starter.category}
                  </Button>
                ))}
              </div>
              <div className="gap- flex flex-col border-t border-gray-300 dark:border-gray-700">
                {CONVERSATION_STARTERS[starterCategory].prompts.map(
                  (prompt, index) => (
                    <div
                      key={prompt}
                      className="border-b border-gray-300 py-1 dark:border-gray-700"
                    >
                      <Button
                        key={index}
                        onClick={() => generate(prompt)}
                        variant="ghost"
                        color="mono"
                        className="justify-start"
                      >
                        {prompt}
                      </Button>
                    </div>
                  )
                )}
              </div>
            </div>
          </div>
        ) : (
          <div className="flex flex-col gap-4">
            {conversationMessages.map((message, index) => (
              <div
                className={cn("w-3/4", {
                  "self-end": message.role === "user",
                })}
                key={index}
              >
                {message.role === "user" ? (
                  <Card className="mt-2 mb-12">
                    <p>{(message as ChatMessageUser).content}</p>
                  </Card>
                ) : message.role === "assistant" ? (
                  <Message message={message as ChatMessageAssistant} />
                ) : null}
              </div>
            ))}
            <div ref={messagesEndRef} />
          </div>
        )}
      </div>
      <div className="fixed right-0 bottom-6 left-0 mx-auto w-full max-w-4xl space-y-2">
        <Card className="mt-auto">
          <ChatForm
            disabled={!ready}
            onSubmit={(prompt) => generate(prompt)}
            isGenerating={state === State.GENERATING}
            onAbort={() => generator.abort()}
          />
        </Card>
        <div className="flex justify-between dark:border-gray-700">
          <Button
            iconLeft={<Settings />}
            size="xs"
            variant="ghost"
            color="mono"
            onClick={openSettingsModal}
          >
            {settings?.modelKey ? MODELS[settings.modelKey].title : ""}
            <span className="mx-2 text-gray-400 dark:text-gray-500">•</span>
            {settings?.tools.length}/{TOOLS.length} tool
            {settings?.tools.length !== 1 ? "s" : ""} active
          </Button>
          <Button
            iconLeft={<RotateCcw />}
            size="xs"
            variant="ghost"
            color="mono"
            onClick={initializeConversation}
          >
            new conversation
          </Button>
        </div>
      </div>
    </>
  );
}