chat working

Author: mikeesto

package-lock.json

@@ -11,7 +11,9 @@
   },
   "dependencies": {
     "@huggingface/transformers": "^3.0.2",
+    "@radix-ui/react-avatar": "^1.1.1",
     "@radix-ui/react-progress": "^1.1.0",
+    "@radix-ui/react-scroll-area": "^1.2.1",
     "@radix-ui/react-slot": "^1.1.0",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.1",

package.json

@@ -1,10 +1,11 @@
 import { useState } from "react";
 import { ModelLoader } from "@/components/ModelLoader";
 import { useWorker } from "./context/workerContext";
+import { Chat } from "@/components/Chat";
 
 export default function App() {
   const [progressItems, setProgressItems] = useState<ProgressItem[]>([]);
   const { status } = useWorker();
 
-  return <ModelLoader />;
+  return <>{status === "ready" ? <Chat /> : <ModelLoader />}</>;
 }

src/App.tsx

@@ -0,0 +1,79 @@
+import { useState } from "react";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { ScrollArea } from "@/components/ui/scroll-area";
+import { useWorker } from "@/context/workerContext";
+
+interface Message {
+  role: "user" | "assistant";
+  content: string;
+}
+
+export function Chat() {
+  const [messages, setMessages] = useState<Message[]>([]);
+  const [input, setInput] = useState("");
+  const [isLoading, setIsLoading] = useState(false);
+  const { generate } = useWorker();
+
+  const sendMessage = async (e: React.FormEvent) => {
+    e.preventDefault();
+    if (!input.trim()) return;
+
+    const newMessage: Message = { role: "user", content: input };
+    setInput("");
+    setIsLoading(true);
+
+    console.log("chat.tsx", [...messages, newMessage]);
+
+    const response = await generate([...messages, newMessage]);
+
+    setMessages((prev) => [
+      ...prev,
+      newMessage,
+      { role: "assistant", content: response },
+    ]);
+    setIsLoading(false);
+  };
+
+  return (
+    <div className="flex h-screen flex-col">
+      <h1>Talk with SmollLM</h1>
+      <ScrollArea className="flex-1 p-4">
+        <div className="space-y-4">
+          {messages.map((message, i) => (
+            <div
+              key={i}
+              className={`flex ${
+                message.role === "user" ? "justify-end" : "justify-start"
+              }`}
+            >
+              <div
+                className={`rounded-lg p-4 max-w-[80%] ${
+                  message.role === "user"
+                    ? "bg-primary text-primary-foreground"
+                    : "bg-muted"
+                }`}
+              >
+                {message.content}
+              </div>
+            </div>
+          ))}
+        </div>
+      </ScrollArea>
+
+      <form onSubmit={sendMessage} className="p-4 border-t">
+        <div className="flex gap-2">
+          <Input
+            value={input}
+            onChange={(e) => setInput(e.target.value)}
+            placeholder="Type a message..."
+            disabled={isLoading}
+          />
+          <Button type="submit" disabled={isLoading}>
+            Send
+          </Button>
+        </div>
+      </form>
+    </div>
+  );
+}

src/components/Chat.tsx

@@ -0,0 +1,22 @@
+import * as React from "react"
+
+import { cn } from "@/lib/utils"
+
+const Input = React.forwardRef<HTMLInputElement, React.ComponentProps<"input">>(
+  ({ className, type, ...props }, ref) => {
+    return (
+      <input
+        type={type}
+        className={cn(
+          "flex h-9 w-full rounded-md border border-input bg-transparent px-3 py-1 text-base shadow-sm transition-colors file:border-0 file:bg-transparent file:text-sm file:font-medium file:text-foreground placeholder:text-muted-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:cursor-not-allowed disabled:opacity-50 md:text-sm",
+          className
+        )}
+        ref={ref}
+        {...props}
+      />
+    )
+  }
+)
+Input.displayName = "Input"
+
+export { Input }

src/components/ui/input.tsx

@@ -0,0 +1,46 @@
+import * as React from "react"
+import * as ScrollAreaPrimitive from "@radix-ui/react-scroll-area"
+
+import { cn } from "@/lib/utils"
+
+const ScrollArea = React.forwardRef<
+  React.ElementRef<typeof ScrollAreaPrimitive.Root>,
+  React.ComponentPropsWithoutRef<typeof ScrollAreaPrimitive.Root>
+>(({ className, children, ...props }, ref) => (
+  <ScrollAreaPrimitive.Root
+    ref={ref}
+    className={cn("relative overflow-hidden", className)}
+    {...props}
+  >
+    <ScrollAreaPrimitive.Viewport className="h-full w-full rounded-[inherit]">
+      {children}
+    </ScrollAreaPrimitive.Viewport>
+    <ScrollBar />
+    <ScrollAreaPrimitive.Corner />
+  </ScrollAreaPrimitive.Root>
+))
+ScrollArea.displayName = ScrollAreaPrimitive.Root.displayName
+
+const ScrollBar = React.forwardRef<
+  React.ElementRef<typeof ScrollAreaPrimitive.ScrollAreaScrollbar>,
+  React.ComponentPropsWithoutRef<typeof ScrollAreaPrimitive.ScrollAreaScrollbar>
+>(({ className, orientation = "vertical", ...props }, ref) => (
+  <ScrollAreaPrimitive.ScrollAreaScrollbar
+    ref={ref}
+    orientation={orientation}
+    className={cn(
+      "flex touch-none select-none transition-colors",
+      orientation === "vertical" &&
+        "h-full w-2.5 border-l border-l-transparent p-[1px]",
+      orientation === "horizontal" &&
+        "h-2.5 flex-col border-t border-t-transparent p-[1px]",
+      className
+    )}
+    {...props}
+  >
+    <ScrollAreaPrimitive.ScrollAreaThumb className="relative flex-1 rounded-full bg-border" />
+  </ScrollAreaPrimitive.ScrollAreaScrollbar>
+))
+ScrollBar.displayName = ScrollAreaPrimitive.ScrollAreaScrollbar.displayName
+
+export { ScrollArea, ScrollBar }

src/components/ui/scroll-area.tsx

@@ -1,5 +1,10 @@
 import { createContext, useContext, useRef, useState, useEffect } from "react";
 
+interface Message {
+  role: "user" | "assistant";
+  content: string;
+}
+
 const WorkerContext = createContext(null);
 
 export function useWorker() {
@@ -60,7 +65,7 @@ export function WorkerProvider({ children }) {
       worker.current?.removeEventListener("message", onMessageReceived);
   }, []);
 
-  const generate = async (prompt: string): Promise<string> => {
+  const generate = async (messages: Message[]): Promise<string> => {
     return new Promise((resolve, reject) => {
       let modelResponse = "";
 
@@ -88,7 +93,7 @@ export function WorkerProvider({ children }) {
       setIsRunning(true);
       worker.current.postMessage({
         type: "generate",
-        data: [{ role: "user", content: prompt }],
+        data: messages,
       });
     });
   };