src/index.js

import * as webllm from "@mlc-ai/web-llm";
import rehypeStringify from "rehype-stringify";
import remarkFrontmatter from "remark-frontmatter";
import remarkGfm from "remark-gfm";
import RemarkBreaks from "remark-breaks";
import remarkParse from "remark-parse";
import remarkRehype from "remark-rehype";
import RehypeKatex from "rehype-katex";
import { unified } from "unified";
import remarkMath from "remark-math";
import rehypeHighlight from "rehype-highlight";
import { visit } from 'unist-util-visit';

// Add custom plugin to escape <think> tags
function escapeThinkTags() {
  return (tree) => {
    visit(tree, 'text', (node) => {
      node.value = node.value.replace(/<think>/g, '&lt;think&gt;')
                            .replace(/<\/think>/g, '&lt;/think&gt;');
    });
  };
}

/*************** WebLLM logic ***************/
const messageFormatter = unified()
  .use(remarkParse)
  .use(remarkFrontmatter)
  .use(remarkMath)
  .use(remarkGfm)
  .use(RemarkBreaks)
  .use(remarkRehype, { allowDangerousHtml: true })
  .use(escapeThinkTags)
  .use(rehypeStringify)
  .use(RehypeKatex)
  .use(rehypeHighlight, {
    detect: true,
    ignoreMissing: true,
  });

const messages = [
  {
    content: "You are a helpful AI agent helping users.",
    role: "system",
  },
];

// Callback function for initializing progress
function updateEngineInitProgressCallback(report) {
  console.log("initialize", report.progress);
  document.getElementById("download-status").textContent = report.text;
}

// Create engine instance
let modelLoaded = false;
const engine = new webllm.MLCEngine();
engine.setLogLevel("INFO");
engine.setInitProgressCallback(updateEngineInitProgressCallback);

async function initializeWebLLMEngine() {
  const model = document.getElementById("model").value;
  const quantization = document.getElementById("quantization").value;
  const context_window_size = parseInt(document.getElementById("context").value);
  const temperature = parseFloat(document.getElementById("temperature").value);
  const top_p = parseFloat(document.getElementById("top_p").value);
  const presence_penalty = parseFloat(document.getElementById("presence_penalty").value);
  const frequency_penalty = parseFloat(document.getElementById("frequency_penalty").value);

  document.getElementById("download-status").classList.remove("hidden");

  const selectedModel = `DeepSeek-R1-Distill-${model}-${quantization}-MLC`;
  const config = {
    temperature,
    top_p,
    frequency_penalty,
    presence_penalty,
    context_window_size,
  };
  console.log(`Loading Model: ${selectedModel}`);
  console.log(`Config: ${JSON.stringify(config)}`);
  await engine.reload(selectedModel, config);
  modelLoaded = true;
}

async function streamingGenerating(messages, onUpdate, onFinish, onError) {
  try {
    let curMessage = "";
    let usage;
    const completion = await engine.chat.completions.create({
      stream: true,
      messages,
      stream_options: { include_usage: true },
    });
    for await (const chunk of completion) {
      const curDelta = chunk.choices[0]?.delta.content;
      if (curDelta) {
        curMessage += curDelta;
      }
      if (chunk.usage) {
        usage = chunk.usage;
      }
      onUpdate(curMessage);
    }
    const finalMessage = await engine.getMessage();
    onFinish(finalMessage, usage);
  } catch (err) {
    onError(err);
  }
}

/*************** UI logic ***************/
function onMessageSend() {
  if (!modelLoaded) {
    return;
  }
  const input = document.getElementById("user-input").value.trim();
  const message = {
    content: input,
    role: "user",
  };
  if (input.length === 0) {
    return;
  }
  document.getElementById("send").disabled = true;

  messages.push(message);
  appendMessage(message);

  document.getElementById("user-input").value = "";
  document
    .getElementById("user-input")
    .setAttribute("placeholder", "Generating...");

  const aiMessage = {
    content: "typing...",
    role: "assistant",
  };
  appendMessage(aiMessage);

  const onFinishGenerating = async (finalMessage, usage) => {
    updateLastMessage(finalMessage);
    document.getElementById("send").disabled = false;
    const usageText =
      `prompt_tokens: ${usage.prompt_tokens}, ` +
      `completion_tokens: ${usage.completion_tokens}, ` +
      `prefill: ${usage.extra.prefill_tokens_per_s.toFixed(4)} tokens/sec, ` +
      `decoding: ${usage.extra.decode_tokens_per_s.toFixed(4)} tokens/sec`;
    document.getElementById("chat-stats").classList.remove("hidden");
    document.getElementById("chat-stats").textContent = usageText;

    document
      .getElementById("user-input")
      .setAttribute("placeholder", "Type a message...");
  };

  streamingGenerating(
    messages,
    updateLastMessage,
    onFinishGenerating,
    console.error
  );
}

function appendMessage(message) {
  const chatBox = document.getElementById("chat-box");
  const container = document.createElement("div");
  container.classList.add("message-container");
  const newMessage = document.createElement("div");
  newMessage.classList.add("message");
  newMessage.textContent = message.content;

  if (message.role === "user") {
    container.classList.add("user");
  } else {
    container.classList.add("assistant");
  }

  container.appendChild(newMessage);
  chatBox.appendChild(container);
  chatBox.scrollTop = chatBox.scrollHeight; // Scroll to the latest message
}

async function updateLastMessage(content) {
  const formattedMessage = await messageFormatter.process(content);
  const messageDoms = document
    .getElementById("chat-box")
    .querySelectorAll(".message");
  const lastMessageDom = messageDoms[messageDoms.length - 1];
  lastMessageDom.innerHTML = formattedMessage;
}

/*************** UI binding ***************/
document.addEventListener('DOMContentLoaded', function() {
  document.getElementById("download").addEventListener("click", function () {
    document.getElementById("send").disabled = true;
    initializeWebLLMEngine().then(() => {
      document.getElementById("send").disabled = false;
    });
  });
  document.getElementById("send").addEventListener("click", function () {
    onMessageSend();
  });
  document.getElementById("user-input").addEventListener("keydown", (event) => {
    if (event.key === "Enter") {
      onMessageSend();
    }
  });

  document.getElementById('model_size').addEventListener('change', function() {
    const quantizationSelect = document.getElementById('quantization');
    const selectedSize = document.getElementById('model_size').value;
    const q0Options = Array.from(quantizationSelect.options).filter(option => 
      option.value === 'q0f32' || option.value === 'q0f16'
    );

    if (selectedSize === '3B') {
      q0Options.forEach(option => option.style.display = 'none');
    } else {
      q0Options.forEach(option => option.style.display = '');
    }
    if (quantizationSelect.selectedOptions[0].style.display === 'none') {
      quantizationSelect.value = quantizationSelect.options[0].value;
    }
  });
});

window.onload = function () {
  document.getElementById("download").textContent = "Download";
  document.getElementById("download").disabled = false;
}