Merge pull request #12 from jeasonstudio/feat-llm-polyfill

feat: add chrome-ai polyfill when browser not support

.changeset/strange-adults-give.md

@@ -0,0 +1,5 @@
+---
+"chrome-ai": minor
+---
+
+feat: add chrome-ai polyfill when chrome version not ok

README.md

@@ -43,7 +43,7 @@
 > ⚠️ Note:
 > * This module is under development and may contain errors and frequent incompatible changes.
 > * Chrome's implementation of [built-in AI with Gemini Nano](https://developer.chrome.com/docs/ai/built-in) is an experiment and will change as they test and address feedback.
-> * If you've never heard of it before, [follow these steps](#enable-ai-in-chrome) to turn on Chrome's built-in AI.
+> * If you've never heard of it before, [follow these steps](#enabling-ai-in-chrome) to turn on Chrome's built-in AI.
 
 ## 📦 Installation
 
@@ -189,7 +189,7 @@ for await (const partialObject of result.partialObjectStream) {
 
 > Due to model reasons, `toolCall/functionCall` are not supported. We are making an effort to implement these functions by prompt engineering.
 
-## Enable AI in Chrome
+## Enabling AI in Chrome
 
 Chrome built-in AI is a preview feature, you need to use chrome version 127 or greater, now in [dev](https://www.google.com/chrome/dev/?extra=devchannel) or [canary](https://www.google.com/chrome/canary/) channel, [may release on stable chanel at Jul 17, 2024](https://chromestatus.com/roadmap).
 
@@ -198,6 +198,14 @@ After then, you should turn on these flags:
 * [chrome://flags/#optimization-guide-on-device-model](chrome://flags/#optimization-guide-on-device-model): `Enabled BypassPrefRequirement`
 * [chrome://components/](chrome://components/): Click `Optimization Guide On Device Model` to download the model.
 
+Or you can try using the experimental feature: `chrome-ai/polyfill`, to use `chrome-ai` in any browser that supports WebGPU and WebAssembly.
+
+```ts
+import 'chrome-ai/polyfill';
+// or
+requre('chrome-ai/polyfill');
+```
+
 ## License
 
 [MIT](LICENSE) License © 2024 [Jeason](https://github.com/jeasonstudio)

app/components/outputs.tsx

@@ -14,7 +14,6 @@ export const Outputs = React.forwardRef<
 >(({ children, ...props }, ref) => {
   let infoElement = null;
 
-  const [isBrowserSupport, setIsBrowserSupport] = React.useState(true);
   const [isEnabledFlags, setIsEnabledFlags] = React.useState(true);
 
   React.useEffect(() => {
@@ -23,7 +22,7 @@ export const Outputs = React.forwardRef<
       return raw ? parseInt(raw[2], 10) : 0;
     }
     const version = getChromeVersion();
-    setIsBrowserSupport(version >= 127);
+    // setIsBrowserSupport(version >= 127);
 
     setIsEnabledFlags(!!('ai' in globalThis));
 
@@ -32,7 +31,7 @@ export const Outputs = React.forwardRef<
     });
   }, []);
 
-  if (!isBrowserSupport || !isEnabledFlags) {
+  if (!isEnabledFlags) {
     infoElement = (
       <div className="w-[500px] m-auto flex flex-col gap-4">
         <h2 className="scroll-m-20 border-b pb-2 text-3xl font-semibold tracking-tight first:mt-0">
@@ -43,15 +42,13 @@ export const Outputs = React.forwardRef<
           experimental and will change as they test and address feedback.
         </p>
 
-        {isBrowserSupport ? null : (
-          <Alert variant="destructive">
-            <AlertCircle className="h-4 w-4" />
-            <AlertTitle>Your browser is not supported.</AlertTitle>
-            <AlertDescription>
-              Please update Chrome to version 127 or higher.
-            </AlertDescription>
-          </Alert>
-        )}
+        <Alert variant="destructive">
+          <AlertCircle className="h-4 w-4" />
+          <AlertTitle>Your browser is not supported.</AlertTitle>
+          <AlertDescription>
+            Please update Chrome to version 127 or higher.
+          </AlertDescription>
+        </Alert>
 
         {isEnabledFlags ? null : (
           <div className="flex flex-col items-start justify-start">

app/orders/page.tsx

@@ -1,5 +1,6 @@
 'use client';
 
+import 'chrome-ai/polyfill';
 import { LoaderCircle, SquareKanban } from 'lucide-react';
 import { Badge } from '../components/ui/badge';
 import { Button } from '../components/ui/button';

app/page.tsx

@@ -1,5 +1,6 @@
 'use client';
 
+require('chrome-ai/polyfill');
 import {
   CornerDownLeft,
   Command,

package.json

@@ -7,13 +7,23 @@
   "exports": {
     ".": {
       "import": "./dist/index.mjs",
-      "require": "./dist/index.js"
+      "require": "./dist/index.js",
+      "types": "./dist/index.d.ts"
+    },
+    "./polyfill": {
+      "import": "./dist/polyfill.mjs",
+      "require": "./dist/polyfill.js",
+      "types": "./dist/polyfill.d.ts"
     }
   },
-  "sideEffects": false,
+  "sideEffects": [
+    "./dist/polyfill.global.js",
+    "./dist/polyfill.js",
+    "./dist/polyfill.mjs"
+  ],
   "scripts": {
     "dev": "tsup --tsconfig=tsconfig.tsup.json --watch",
-    "build": "tsup --tsconfig=tsconfig.tsup.json",
+    "build": "tsup --tsconfig=tsconfig.tsup.json --clean",
     "test": "vitest",
     "test:coverage": "vitest --coverage",
     "test:ci": "vitest --run",
@@ -36,6 +46,7 @@
   },
   "dependencies": {
     "@ai-sdk/provider": "^0.0.11",
+    "@mediapipe/tasks-genai": "^0.10.14",
     "@mediapipe/tasks-text": "^0.10.14",
     "debug": "^4.3.5"
   },

src/global.d.ts

@@ -1,4 +1,4 @@
-export type ChromeAISessionAvailable = 'no' | 'readily';
+export type ChromeAISessionAvailable = 'no' | 'after-download' | 'readily';
 
 export interface ChromeAISessionOptions {
   temperature?: number;
@@ -18,11 +18,21 @@ export interface ChromePromptAPI {
   canCreateTextSession: () => Promise<ChromeAISessionAvailable>;
   defaultGenericSessionOptions: () => Promise<ChromeAISessionOptions>;
   defaultTextSessionOptions: () => Promise<ChromeAISessionOptions>;
-  createGenericSession: (options?: ChromeAISessionOptions) => Promise<ChromeAISession>;
-  createTextSession: (options?: ChromeAISessionOptions) => Promise<ChromeAISession>;
+  createGenericSession: (
+    options?: ChromeAISessionOptions
+  ) => Promise<ChromeAISession>;
+  createTextSession: (
+    options?: ChromeAISessionOptions
+  ) => Promise<ChromeAISession>;
+}
+
+export interface PolyfillChromeAIOptions {
+  llmModelAssetPath: string;
+  filesetBasePath: string;
 }
 
 declare global {
   var ai: ChromePromptAPI;
   var model = ai;
+  var __polyfill_ai_options__: Partial<PolyfillChromeAIOptions> | undefined;
 }

src/index.ts

@@ -1,3 +1,4 @@
 export * from './language-model';
 export * from './embedding-model';
 export * from './chromeai';
+export * from './polyfill/session';

src/polyfill/index.ts

@@ -0,0 +1,3 @@
+import { polyfillChromeAI } from './session';
+
+polyfillChromeAI(globalThis.__polyfill_ai_options__);

src/polyfill/session.ts

@@ -0,0 +1,124 @@
+import {
+  FilesetResolver,
+  LlmInference,
+  ProgressListener,
+} from '@mediapipe/tasks-genai';
+import {
+  ChromeAISession,
+  ChromeAISessionAvailable,
+  ChromeAISessionOptions,
+  ChromePromptAPI,
+  PolyfillChromeAIOptions,
+} from '../global';
+import createDebug from 'debug';
+
+const debug = createDebug('chromeai:polyfill');
+
+class PolyfillChromeAISession implements ChromeAISession {
+  public constructor(private llm: LlmInference) {
+    debug('PolyfillChromeAISession created', llm);
+  }
+
+  private generateText = async (prompt: string): Promise<string> => {
+    const response = await this.llm.generateResponse(prompt);
+    debug('generateText', prompt, response);
+    return response;
+  };
+
+  private streamText = (prompt: string): ReadableStream<string> => {
+    debug('streamText', prompt);
+    const stream = new ReadableStream<string>({
+      start: (controller) => {
+        const listener: ProgressListener = (
+          partialResult: string,
+          done: boolean
+        ) => {
+          controller.enqueue(partialResult);
+          if (done) {
+            controller.close();
+          }
+        };
+        this.llm.generateResponse(prompt, listener);
+      },
+      cancel: (reason) => {
+        console.warn('stream text canceled', reason);
+      },
+    });
+    debug('streamText', prompt);
+    return stream;
+  };
+
+  public destroy = async () => this.llm.close();
+  public prompt = this.generateText;
+  public execute = this.generateText;
+  public promptStreaming = this.streamText;
+  public executeStreaming = this.streamText;
+}
+
+/**
+ * Model: https://huggingface.co/oongaboongahacker/Gemini-Nano
+ */
+export class PolyfillChromeAI implements ChromePromptAPI {
+  private aiOptions: PolyfillChromeAIOptions = {
+    // About 1.78GB, should cache by browser
+    llmModelAssetPath:
+      'https://huggingface.co/oongaboongahacker/Gemini-Nano/resolve/main/weights.bin',
+    filesetBasePath: 'https://unpkg.com/@mediapipe/tasks-genai/wasm/',
+  };
+
+  public constructor(aiOptions: Partial<PolyfillChromeAIOptions> = {}) {
+    this.aiOptions = Object.assign(this.aiOptions, aiOptions);
+    debug('PolyfillChromeAI created', this.aiOptions);
+    this.modelAssetBuffer = fetch(this.aiOptions.llmModelAssetPath).then(
+      (response) => response.body!.getReader()
+    )!;
+  }
+
+  private modelAssetBuffer: Promise<ReadableStreamDefaultReader>;
+
+  private canCreateSession = async (): Promise<ChromeAISessionAvailable> => {
+    // TODO@jeasonstudio:
+    // * if browser do not support WebAssembly/WebGPU, return 'no';
+    // * check if modelAssetBuffer is downloaded, if not, return 'after-download';
+    return 'readily';
+  };
+  private defaultSessionOptions =
+    async (): Promise<ChromeAISessionOptions> => ({
+      temperature: 0.800000011920929,
+      topK: 3,
+    });
+
+  private createSession = async (
+    options?: ChromeAISessionOptions
+  ): Promise<ChromeAISession> => {
+    const argv = options ?? (await this.defaultSessionOptions());
+    const fileset = await FilesetResolver.forGenAiTasks(
+      this.aiOptions.filesetBasePath
+    );
+    const llm = await LlmInference.createFromOptions(fileset, {
+      baseOptions: {
+        modelAssetBuffer: await this.modelAssetBuffer,
+      },
+      temperature: argv.temperature,
+      topK: argv.topK,
+    });
+    const session = new PolyfillChromeAISession(llm);
+    debug('createSession', options, session);
+    return session;
+  };
+
+  public canCreateGenericSession = this.canCreateSession;
+  public canCreateTextSession = this.canCreateSession;
+  public defaultGenericSessionOptions = this.defaultSessionOptions;
+  public defaultTextSessionOptions = this.defaultSessionOptions;
+  public createGenericSession = this.createSession;
+  public createTextSession = this.createSession;
+}
+
+export const polyfillChromeAI = (
+  options?: Partial<PolyfillChromeAIOptions>
+) => {
+  const ai = new PolyfillChromeAI(options);
+  globalThis.ai = globalThis.ai || ai;
+  globalThis.model = globalThis.model || ai;
+};

tsconfig.json

@@ -2,8 +2,10 @@
   "extends": "./tsconfig.base.json",
   "compilerOptions": {
     "plugins": [{ "name": "next" }],
+    "baseUrl": "./",
     "paths": {
       "@/*": ["./*"],
+      "chrome-ai/*": ["./src/*"],
       "chrome-ai": ["./src/index.ts"]
     },
     "incremental": true

tsup.config.ts

@@ -7,4 +7,12 @@ export default defineConfig([
     format: ['cjs', 'esm'],
     sourcemap: true,
   },
+  {
+    dts: true,
+    entry: {
+      polyfill: 'src/polyfill/index.ts',
+    },
+    format: ['cjs', 'esm', 'iife'],
+    sourcemap: true,
+  },
 ]);

vitest.config.ts

@@ -8,6 +8,7 @@ export default defineConfig({
       // If you want a coverage reports even if your tests are failing, include the reportOnFailure option
       reportOnFailure: true,
       include: ['src/**/*.ts'],
+      exclude: ['src/polyfill/*'], // TODO@jeasonstudio: finish the polyfill tests
     },
   },
 });