Use @langchain/community WebLLM model

Author: jacoblee93

.env.example .env.local.example

@@ -26,6 +26,7 @@ yarn-error.log*
 
 # local env files
 .env*.local
+.env
 
 # vercel
 .vercel
@@ -34,4 +35,4 @@ yarn-error.log*
 *.tsbuildinfo
 next-env.d.ts
 
-.yarn
+.yarn

.gitignore

@@ -29,7 +29,7 @@ import { LangChainTracer } from "@langchain/core/tracers/tracer_langchain";
 import { Client } from "langsmith";
 
 import { ChatOllama } from "@langchain/community/chat_models/ollama";
-import { ChatWebLLM } from "./lib/chat_models/webllm";
+import { ChatWebLLM } from "@langchain/community/chat_models/webllm";
 
 const embeddings = new HuggingFaceTransformersEmbeddings({
   modelName: "Xenova/all-MiniLM-L6-v2",
@@ -229,15 +229,15 @@ self.addEventListener("message", async (event: { data: any }) => {
   } else {
     const modelProvider = event.data.modelProvider;
     const modelConfig = event.data.modelConfig;
-    let chatModel: BaseChatModel | LanguageModelLike =
-      modelProvider === "ollama"
-        ? new ChatOllama(modelConfig)
-        : new ChatWebLLM(modelConfig);
+    let chatModel: BaseChatModel | LanguageModelLike;
     if (modelProvider === "webllm") {
-      await (chatModel as ChatWebLLM).initialize((event) =>
+      const webllmModel = new ChatWebLLM(modelConfig);
+      await webllmModel.initialize((event) =>
         self.postMessage({ type: "init_progress", data: event }),
       );
-      chatModel = chatModel.bind({ stop: ["\nInstruct:", "Instruct:"] });
+      chatModel = webllmModel.bind({ stop: ["\nInstruct:", "Instruct:"] });
+    } else {
+      chatModel = new ChatOllama(modelConfig);
     }
     try {
       await queryVectorStore(event.data.messages, {

app/lib/chat_models/webllm.ts

@@ -6,7 +6,7 @@ import 'react-toastify/dist/ReactToastify.css';
 import { useRef, useState, useEffect } from "react";
 import type { FormEvent } from "react";
 
-import { ChatMessageBubble } from "@/components/ChatMessageBubble";
+import { ChatMessageBubble } from '@/components/ChatMessageBubble';
 import { ChatWindowMessage } from '@/schema/ChatWindowMessage';
 
 export function ChatWindow(props: {
@@ -35,25 +35,12 @@ export function ChatWindow(props: {
           controller.close();
           return;
         }
-        // Config copied from:
-        // https://github.com/mlc-ai/web-llm/blob/eaaff6a7730b6403810bb4fd2bbc4af113c36050/examples/simple-chat/src/gh-config.js
+        // See https://github.com/mlc-ai/web-llm/blob/main/src/config.ts for a list of available models
         const webLLMConfig = {
-          temperature: 0.1,
-          modelRecord: {
-            "model_url": "https://huggingface.co/mlc-ai/phi-2-q4f32_1-MLC/resolve/main/",
-            "local_id": "Phi2-q4f32_1",
-            "model_lib_url": "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/main/phi-2/phi-2-q4f32_1-ctx2k-webgpu.wasm",
-            "vram_required_MB": 4032.48,
-            "low_resource_required": false,
+          model: "Phi2-q4f32_1",
+          chatOptions: {
+            temperature: 0.1,
           },
-          // {
-          //   "model_url": "https://huggingface.co/mlc-ai/phi-2-q0f16-MLC/resolve/main/",
-          //   "local_id": "Phi2-q0f16",
-          //   "model_lib_url": "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/main/phi-2/phi-2-q0f16-ctx2k-webgpu.wasm",
-          //   "vram_required_MB": 11079.47,
-          //   "low_resource_required": false,
-          //   "required_features": ["shader-f16"],
-          // },
         };
         const ollamaConfig = {
           baseUrl: "http://localhost:11435",
@@ -89,15 +76,15 @@ export function ChatWindow(props: {
                 initProgressToastId.current = toast(
                   "Loading model weights... This may take a while",
                   {
-                    progress: e.data.data.progress,
+                    progress: e.data.data.progress || 0.01,
                     theme: "dark"
                   }
                 );
               } else {
                 if (e.data.data.progress === 1) {
                   await new Promise((resolve) => setTimeout(resolve, 2000));
                 }
-                toast.update(initProgressToastId.current, { progress: e.data.data.progress });
+                toast.update(initProgressToastId.current, { progress: e.data.data.progress || 0.01 });
               }
               break
             case "chunk":

app/worker.ts

@@ -13,8 +13,8 @@
     "node": ">=18"
   },
   "dependencies": {
-    "@langchain/community": "^0.0.40",
-    "@mlc-ai/web-llm": "^0.2.28",
+    "@langchain/community": "^0.0.57",
+    "@mlc-ai/web-llm": "^0.2.35",
     "@types/node": "20.4.5",
     "@types/react": "18.2.17",
     "@types/react-dom": "18.2.7",
@@ -37,5 +37,8 @@
   },
   "devDependencies": {
     "prettier": "3.0.0"
+  },
+  "resolutions": {
+    "@langchain/core": "0.1.63"
   }
 }