Private & Fast: Building a Local-First Mental Health Assistant with WebLLM and WebGPU 🧠💻

Privacy is no longer just a feature; it's a human right—especially when it comes to mental health. Imagine a Cognitive Behavioral Therapy (CBT) assistant that lives entirely in your browser, never sends a single byte of your conversation to a remote server, and runs at lightning speed.

Thanks to the explosion of Edge AI and the maturing WebGPU standard, this isn't science fiction anymore. In this tutorial, we will explore how to use WebLLM, TVM Unity, and React to build a high-performance, offline-capable mental health assistant. By leveraging WebLLM and Local-first AI principles, we can provide low-latency support while keeping sensitive user data exactly where it belongs: on the user's device.

Why Local AI for Mental Health?

Mental health data is incredibly sensitive. Using traditional LLM APIs (like OpenAI or Claude) means sending private thoughts to the cloud. By using WebGPU acceleration, we can run models like Llama 3 or Mistral directly on the client's GPU via the browser.

The Architecture: From Weights to WebGPU

The magic happens through TVM Unity, which compiles machine learning models into high-performance kernels that the browser can execute via the WebGPU API.

graph TD
    A[User Input] --> B[React UI State]
    B --> C[WebLLM Worker]
    subgraph Browser Environment
    C --> D[TVM Runtime]
    D --> E[WebGPU API]
    E --> F[Local GPU / VRAM]
    F --> G[Model Inference]
    G --> D
    end
    D --> H[Streaming Response]
    H --> B
    I[(IndexedDB Cache)] -.-> C

Prerequisites

To follow along, you’ll need:

• Node.js & npm/pnpm
• A browser with WebGPU support (Chrome 113+, Edge, or Firefox Nightly)
• The following stack: React, WebLLM, and Vite.

Step 1: Setting up the WebLLM Engine

First, let's install the core dependencies:

npm install @mlc-ai/web-llm react

The heart of our application is the ChatWorker. We want to run the LLM in a Web Worker to ensure the UI remains responsive during heavy inference.

// engine.ts
import { CreateWebWorkerEngine, type ChatOptions } from "@mlc-ai/web-llm";

const SELECTED_MODEL = "Llama-3-8B-Instruct-q4f16_1-MLC";

export async function initializeEngine(onProgress: (p: any) => void) {
  // This downloads the model and initializes the WebGPU pipeline
  const engine = await CreateWebWorkerEngine(
    new Worker(new URL("./worker.ts", import.meta.url), { type: "module" }),
    SELECTED_MODEL,
    { initProgressCallback: onProgress }
  );
  return engine;
}

Step 2: Defining the CBT Specialist Persona

For a mental health assistant, the system prompt is everything. We need to steer the model toward Cognitive Behavioral Therapy techniques—identifying cognitive distortions and suggesting reframing exercises.

const CBT_SYSTEM_PROMPT = `
You are a supportive Mental Health Assistant specialized in Cognitive Behavioral Therapy (CBT).
Your goal is to help users identify negative thought patterns.
Rules:
1. Be empathetic and non-judgmental.
2. If a user mentions self-harm, immediately provide emergency resources.
3. Use Socratic questioning to help users reach their own conclusions.
4. Keep data privacy top-of-mind: remind users this is a local-only session.
`;

Step 3: The React Integration

Now, let's build a custom hook to manage the chat state and the engine lifecycle. 🚀

// useWebLLM.ts
import { useState, useEffect } from 'react';
import * as webllm from "@mlc-ai/web-llm";

export function useWebLLM() {
  const [engine, setEngine] = useState<webllm.EngineInterface | null>(null);
  const [progress, setProgress] = useState("");

  useEffect(() => {
    async function init() {
      const instance = await webllm.CreateMLCEngine(
        "Llama-3-8B-Instruct-q4f16_1-MLC", 
        { initProgressCallback: (p) => setProgress(p.text) }
      );
      setEngine(instance);
    }
    init();
  }, []);

  const chat = async (messages: webllm.ChatCompletionMessageParam[]) => {
    if (!engine) return;

    const chunks = await engine.chat.completions.create({
      messages: [
        { role: "system", content: CBT_SYSTEM_PROMPT },
        ...messages
      ],
      stream: true,
    });

    return chunks;
  };

  return { chat, progress, ready: !!engine };
}

The "Official" Way: Production-Ready Edge AI 🥑

While this demo gets you up and running, deploying Edge AI at scale requires deep optimization regarding model quantization, caching strategies, and cross-browser compatibility.

If you are looking for advanced patterns on model sharding or more production-ready examples of local-first architectures, you should definitely check out the deep-dives at WellAlly Tech Blog. They cover the nuances of TVM Unity and memory management that are crucial for high-traffic applications.

Step 4: Building the UI

Finally, we connect our hook to a simple chat interface. Note how we handle the streaming response to create that "typing" effect users love.

// App.tsx
import { useState } from 'react';
import { useWebLLM } from './useWebLLM';

function App() {
  const { chat, progress, ready } = useWebLLM();
  const [input, setInput] = useState("");
  const [messages, setMessages] = useState([]);

  const handleSend = async () => {
    const userMsg = { role: "user", content: input };
    setMessages(prev => [...prev, userMsg]);
    setInput("");

    const stream = await chat([...messages, userMsg]);
    let fullReply = "";

    for await (const chunk of stream) {
      fullReply += chunk.choices[0]?.delta?.content || "";
      // Update UI with the stream
      setMessages(prev => [...prev.slice(0, -1), { role: "assistant", content: fullReply }]);
    }
  };

  if (!ready) return <div>Loading AI Engine: {progress}</div>;

  return (
    <div className="chat-container">
      <h2>Local CBT Assistant 🌿</h2>
      <div className="messages">
        {messages.map((m, i) => (
          <div key={i} className={m.role}>{m.content}</div>
        ))}
      </div>
      <input value={input} onChange={(e) => setInput(e.target.value)} />
      <button onClick={handleSend}>Talk</button>
    </div>
  );
}

Conclusion: The Future is On-Device

By combining WebLLM and WebGPU, we've built a tool that is:

1. Private: No data leaves the device.
2. Cost-Effective: Zero API costs per token.
3. Low Latency: No network round-trips for inference.

Edge AI is transforming how we think about sensitive applications like mental health, finance, and personal journaling. By moving the "brain" to the browser, we empower users to own their data without sacrificing the power of modern LLMs.

What are you planning to build with WebGPU? Let me know in the comments! And don't forget to head over to wellally.tech/blog for more technical guides on the future of decentralized AI! 🚀🔥