Add app

adamelliotfields · Feb 14, 2024 · 2940ba6 · 2940ba6
1 parent e821edc
commit 2940ba6
Show file tree

Hide file tree

Showing 15 changed files with 744 additions and 0 deletions.
diff --git a/index.html b/index.html
@@ -0,0 +1,17 @@
+<!doctype html>
+<html lang="en" class="h-full antialiased [font-synthesis:none] [text-rendering:optimizeLegibility]">
+  <head>
+    <title>Chat</title>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <meta name="description" content="Chat with LLMs running in your browser." />
+    <meta name="theme-color" content="#000000" />
+    <meta name="robots" content="index, follow" />
+    <!-- <link rel="icon" type="image/svg+xml" href="/favicon.svg" /> -->
+    <link rel="icon" href="https://aef.me/favicon.ico" />
+    <link rel="canonical" href="https://aef.me/chat/" />
+  <body class="h-full bg-neutral-50">
+    <div id="root" class="h-full"></div>
+    <script type="module" src="./src/index.tsx"></script>
+  </body>
+</html>
diff --git a/src/atoms.ts b/src/atoms.ts
@@ -0,0 +1,24 @@
+import { type AppConfig } from '@mlc-ai/web-llm'
+import { atom } from 'jotai'
+import { atomWithReset } from 'jotai/utils'
+import config from './config'
+import { type Conversation } from './types'
+
+// the app config
+export const configAtom = atom<AppConfig>(config)
+
+// the loaded model
+export const activeModelIdAtom = atomWithReset<string | null>(null)
+
+// the conversation (resettable)
+export const conversationAtom = atomWithReset<Conversation>({
+  messages: [],
+  stream: null
+})
+
+// runtime stats text
+export const runtimeStatsTextAtom = atomWithReset<string | null>(null)
+
+// loading states
+export const loadingAtom = atom<boolean>(false)
+export const generatingAtom = atom<boolean>(false)
diff --git a/src/components/App.tsx b/src/components/App.tsx
@@ -0,0 +1,230 @@
+import { ChatModule } from '@mlc-ai/web-llm'
+import { useAtom, useAtomValue, useSetAtom } from 'jotai'
+import { RESET } from 'jotai/utils'
+import { Square, Trash } from 'lucide-react'
+import { type MouseEvent } from 'react'
+
+import {
+  activeModelIdAtom,
+  configAtom,
+  conversationAtom,
+  generatingAtom,
+  loadingAtom,
+  runtimeStatsTextAtom
+} from '../atoms'
+import { type Message } from '../types'
+import { Button } from './Button'
+import { Header } from './Header'
+import { MessageList } from './MessageList'
+import { ModelSelect } from './ModelSelect'
+import { PromptInput } from './PromptInput'
+import { RuntimeStats } from './RuntimeStats'
+
+type ProgressReport = { progress: number; text: string; timeElapsed: number }
+type InitProgressCallback = (report: ProgressReport) => void
+type GenerateCallback = (step: number, content: string) => void
+
+interface AppProps {
+  chat: ChatModule
+}
+
+export default function App({ chat }: AppProps) {
+  const [generating, setGenerating] = useAtom(generatingAtom)
+  const [loading, setLoading] = useAtom(loadingAtom)
+  const [conversation, setConversation] = useAtom(conversationAtom)
+  const setActiveModelId = useSetAtom(activeModelIdAtom)
+  const setStatsText = useSetAtom(runtimeStatsTextAtom)
+  const config = useAtomValue(configAtom)
+
+  const trashDisabled = conversation.messages.length < 1
+  const stopDisabled = loading || !generating
+
+  const onGenerate: GenerateCallback = (_, content) => {
+    setConversation(({ messages }) => ({
+      messages,
+      stream: { messageRole: 'assistant', content }
+    }))
+  }
+
+  const handleResetClick = async () => {
+    if (trashDisabled) return
+    try {
+      await chat.resetChat()
+      setStatsText(RESET)
+      setConversation(RESET)
+    } catch (err) {
+      setConversation(({ messages }) => ({
+        messages: [
+          ...messages,
+          { messageRole: 'status', content: `Error: ${(err as Error).message}` }
+        ],
+        stream: null
+      }))
+      console.error(err)
+    }
+  }
+
+  const handleReloadClick = async (
+    _: MouseEvent<HTMLButtonElement>,
+    modelId: string | null
+  ) => {
+    // load the selected model
+    if (typeof modelId === 'string') {
+      setLoading(true)
+
+      try {
+        await chat.resetChat()
+        setConversation(RESET)
+
+        await chat.reload(modelId, undefined, config)
+        setConversation(({ stream }) => ({
+          messages: [{ messageRole: 'status', content: (stream as Message).content }],
+          stream: null
+        }))
+        setActiveModelId(modelId)
+      } catch (err) {
+        setConversation({
+          messages: [
+            { messageRole: 'status', content: `Error: ${(err as Error).message}` }
+          ],
+          stream: null
+        })
+        console.error(err)
+      }
+
+      setLoading(false)
+      return
+    }
+
+    // unload the active model and reset the conversation
+    if (modelId === null) {
+      await chat.unload()
+      setActiveModelId(RESET)
+      setConversation(RESET)
+    }
+  }
+
+  const handleInputButtonClick = async (
+    _: MouseEvent<HTMLButtonElement>,
+    prompt: string
+  ) => {
+    if (prompt) {
+      setGenerating(true)
+
+      // user message
+      setConversation(({ messages }) => ({
+        messages: [...messages, { messageRole: 'user', content: prompt }],
+        stream: null
+      }))
+
+      // assistant message
+      try {
+        // the stop button causes the Promise to resolve
+        // this is why `setConversation` gets called
+        const message = await chat.generate(prompt, onGenerate)
+        setConversation(({ messages }) => ({
+          messages: [...messages, { messageRole: 'assistant', content: message }],
+          stream: null
+        }))
+      } catch (err) {
+        setConversation(({ messages }) => ({
+          messages: [
+            ...messages,
+            { messageRole: 'status', content: `Error: ${(err as Error).message}` }
+          ],
+          stream: null
+        }))
+        console.error(err)
+        return
+      }
+
+      setGenerating(false)
+
+      // runtime stats
+      try {
+        const text = await chat.runtimeStatsText()
+        setStatsText(text)
+      } catch (err) {
+        setConversation(({ messages }) => ({
+          messages: [
+            ...messages,
+            { messageRole: 'status', content: `Error: ${(err as Error).message}` }
+          ],
+          stream: null
+        }))
+        console.error(err)
+      }
+    }
+  }
+
+  const handleStopClick = async () => {
+    if (stopDisabled) return
+    try {
+      await chat.interruptGenerate()
+    } catch (err) {
+      console.error(err)
+    }
+  }
+
+  const initProgress: InitProgressCallback = ({ text: content }) =>
+    setConversation(({ messages }) => ({
+      messages: messages,
+      stream: { messageRole: 'status', content }
+    }))
+  chat.setInitProgressCallback(initProgress)
+
+  // the height math is one way of ensuring we don't get outer scrollbars
+  // message list is full viewport height minus header and footer
+  // header and footer get fixed heights
+  return (
+    <div className="flex flex-col">
+      <div className="grow">
+        <Header className="h-[56px]" />
+
+        <main className="flex flex-col grow">
+          <div className="w-full max-w-screen-lg mx-auto">
+            <MessageList
+              // change `scroll` to `auto` if you don't like the scrollbars being always visible
+              className="h-[calc(100vh_-_56px_-_192px)] overflow-y-scroll md:h-[calc(100vh_-_56px_-_160px)]"
+            />
+          </div>
+        </main>
+      </div>
+
+      <footer className="h-[192px] bg-neutral-50 sticky bottom-0 border-t z-20 md:h-[160px]">
+        {/* footer top row */}
+        <div className="max-w-screen-lg mx-auto">
+          <div className="p-4 flex flex-col justify-between md:border-x md:flex-row md:items-center">
+            <RuntimeStats />
+            <ModelSelect handleClick={handleReloadClick} />
+          </div>
+        </div>
+
+        {/* footer bottom row */}
+        <div className="max-w-screen-lg mx-auto">
+          <div className="py-4 px-2 border-t md:border-x">
+            <div className="flex items-center">
+              <Button
+                disabled={trashDisabled}
+                label="Reset"
+                icon={Trash}
+                className="mr-2 p-2 text-lg"
+                onClick={handleResetClick}
+              />
+              <div className="grow">
+                <PromptInput handleClick={handleInputButtonClick} />
+              </div>
+              <Button
+                disabled={stopDisabled}
+                label="Stop"
+                icon={Square}
+                className="ml-2 p-2 text-lg"
+                onClick={handleStopClick}
+              />
+            </div>
+          </div>
+        </div>
+      </footer>
+    </div>
+  )
+}
diff --git a/src/components/Button.tsx b/src/components/Button.tsx
@@ -0,0 +1,70 @@
+import clsx from 'clsx'
+import type { AnchorHTMLAttributes, ButtonHTMLAttributes, ElementType } from 'react'
+
+export type ButtonProps = Omit<
+  ButtonHTMLAttributes<HTMLButtonElement> & AnchorHTMLAttributes<HTMLAnchorElement>,
+  'onClick'
+> & {
+  active?: boolean
+  disabled?: boolean
+  icon: ElementType
+  label: string
+  onClick?: ButtonHTMLAttributes<HTMLButtonElement>['onClick']
+}
+
+// TODO: take "size" prop which controls the padding
+export function Button({
+  active = false,
+  className,
+  disabled = false,
+  href,
+  icon,
+  label,
+  onClick,
+  ...rest
+}: ButtonProps) {
+  const Icon = icon
+
+  const enabledClassNames =
+    'active:shadow-none active:scale-95 focus:ring-2 focus:ring-cyan-500 hover:transition-all hover:shadow-sm'
+  const enabledInactiveClassNames =
+    'text-neutral-400 focus:text-neutral-500 hover:text-neutral-500 hover:bg-neutral-200/50'
+  const enabledActiveClassNames = 'text-cyan-500 hover:bg-cyan-200/50'
+  const disabledClassNames = 'text-neutral-400/50 cursor-not-allowed'
+
+  const classNames = clsx(
+    'leading-none rounded-full focus-visible:outline-none',
+    disabled && disabledClassNames,
+    !disabled && enabledClassNames,
+    !disabled && active && enabledActiveClassNames,
+    !disabled && !active && enabledInactiveClassNames
+  )
+
+  if (typeof href === 'string' && href.length > 0) {
+    return (
+      <a
+        className={clsx(classNames, className)}
+        href={href}
+        target="_blank"
+        rel="noopener noreferrer"
+        {...rest}
+      >
+        <span className="sr-only">{label}</span>
+        <Icon size="1em" aria-hidden="true" />
+      </a>
+    )
+  }
+
+  return (
+    <button
+      className={clsx(classNames, className)}
+      aria-label={label}
+      type="button"
+      onClick={onClick}
+      disabled={disabled}
+      {...rest}
+    >
+      <Icon size="1em" />
+    </button>
+  )
+}
diff --git a/src/components/Header.tsx b/src/components/Header.tsx
@@ -0,0 +1,39 @@
+import clsx from 'clsx'
+import { Github } from 'lucide-react'
+import { type HTMLAttributes, memo } from 'react'
+
+import { HREF, TITLE } from '../consts'
+import { Button } from './Button'
+
+export interface HeaderProps extends HTMLAttributes<HTMLDivElement> {}
+
+export const Header = memo(function Header({ className, ...rest }: HeaderProps) {
+  return (
+    <div
+      className={clsx('bg-neutral-50 border-b sticky top-0 z-20', className)}
+      {...rest}
+    >
+      <div className="h-full max-w-screen-lg mx-auto">
+        <div className="h-full flex items-center justify-between p-4 md:border-x">
+          <a
+            href={HREF}
+            className={clsx(
+              'px-1 py-0.5 font-mono font-black text-lg text-neutral-600 tracking-wide',
+              'focus:ring-2 focus:ring-cyan-500 focus-visible:outline-none'
+            )}
+            target="_blank"
+            rel="noopener noreferrer"
+          >
+            {TITLE}
+          </a>
+          <Button
+            className="p-2 text-lg"
+            href="https://github.com/adamelliotfields/chat"
+            icon={Github}
+            label="GitHub"
+          />
+        </div>
+      </div>
+    </div>
+  )
+})