VinF Hybrid Inference #4: ChromeAdapter in stream methods (rebased)

common/api-review/util.api.md

@@ -264,6 +264,11 @@ export function isBrowser(): boolean;
 // @public (undocumented)
 export function isBrowserExtension(): boolean;
 
+// Warning: (ae-missing-release-tag) "isChrome" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
+//
+// @public (undocumented)
+export function isChrome(): boolean;
+
 // Warning: (ae-missing-release-tag) "isCloudflareWorker" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
 //
 // @public

e2e/sample-apps/modular.js

@@ -58,12 +58,7 @@ import {
   onValue,
   off
 } from 'firebase/database';
-import {
-  getGenerativeModel,
-  getVertexAI,
-  InferenceMode,
-  VertexAI
-} from 'firebase/vertexai';
+import { getGenerativeModel, getVertexAI } from 'firebase/vertexai';
 import { getDataConnect, DataConnect } from 'firebase/data-connect';
 
 /**
@@ -318,8 +313,13 @@ function callPerformance(app) {
 async function callVertexAI(app) {
   console.log('[VERTEXAI] start');
   const vertexAI = getVertexAI(app);
-  const model = getGenerativeModel(vertexAI, { model: 'gemini-1.5-flash' });
-  const result = await model.countTokens('abcdefg');
+  const model = getGenerativeModel(vertexAI, {
+    mode: 'prefer_in_cloud'
+  });
+  const result = await model.generateContentStream("What is Roko's Basalisk?");
+  for await (const chunk of result.stream) {
+    console.log(chunk.text());
+  }
   console.log(`[VERTEXAI] counted tokens: ${result.totalTokens}`);
 }
 

packages/util/src/environment.ts

@@ -173,6 +173,12 @@ export function isSafari(): boolean {
   );
 }
 
+export function isChrome(): boolean {
+  return (
+    !isNode() && !!navigator.userAgent && navigator.userAgent.includes('Chrome')
+  );
+}
+
 /**
  * This method checks if indexedDB is supported by current browser/service worker context
  * @return true if indexedDB is supported by current browser/service worker context

packages/vertexai/package.json

@@ -58,6 +58,7 @@
   "devDependencies": {
     "@firebase/app": "0.11.4",
     "@rollup/plugin-json": "6.1.0",
+    "@types/dom-chromium-ai": "0.0.6",
     "rollup": "2.79.2",
     "rollup-plugin-replace": "2.2.0",
     "rollup-plugin-typescript2": "0.36.0",

packages/vertexai/src/methods/chat-session.ts

@@ -149,6 +149,7 @@ export class ChatSession {
       this._apiSettings,
       this.model,
       generateContentRequest,
+      this.chromeAdapter,
       this.requestOptions
     );
 

packages/vertexai/src/methods/chrome-adapter.ts

@@ -103,23 +103,19 @@ export class ChromeAdapter {
     );
     const messages = ChromeAdapter.toLanguageModelMessages(request.contents);
     const text = await session.prompt(messages);
-    return {
-      json: () =>
-        Promise.resolve({
-          candidates: [
-            {
-              content: {
-                parts: [{ text }]
-              }
-            }
-          ]
-        })
-    } as Response;
+    return ChromeAdapter.toResponse(text);
+  }
+  async generateContentStreamOnDevice(
+    request: GenerateContentRequest
+  ): Promise<Response> {
+    const session = await this.createSession(
+      // TODO: normalize on-device params during construction.
+      this.onDeviceParams || {}
+    );
+    const messages = ChromeAdapter.toLanguageModelMessages(request.contents);
+    const stream = await session.promptStreaming(messages);
+    return ChromeAdapter.toStreamResponse(stream);
   }
-
-  /**
-   * Asserts inference for the given request can be performed by an on-device model.
-   */
   private static isOnDeviceRequest(request: GenerateContentRequest): boolean {
     // Returns false if the prompt is empty.
     if (request.contents.length === 0) {
@@ -224,4 +220,47 @@ export class ChromeAdapter {
     this.oldSession = newSession;
     return newSession;
   }
+
+  /**
+   * Formats string returned by Chrome as a {@link Response} returned by Vertex.
+   */
+  private static toResponse(text: string): Response {
+    return {
+      json: async () => ({
+        candidates: [
+          {
+            content: {
+              parts: [{ text }]
+            }
+          }
+        ]
+      })
+    } as Response;
+  }
+
+  /**
+   * Formats string stream returned by Chrome as SSE returned by Vertex.
+   */
+  private static toStreamResponse(stream: ReadableStream<string>): Response {
+    const encoder = new TextEncoder();
+    return {
+      body: stream.pipeThrough(
+        new TransformStream({
+          transform(chunk, controller) {
+            const json = JSON.stringify({
+              candidates: [
+                {
+                  content: {
+                    role: 'model',
+                    parts: [{ text: chunk }]
+                  }
+                }
+              ]
+            });
+            controller.enqueue(encoder.encode(`data: ${json}\n\n`));
+          }
+        })
+      )
+    } as Response;
+  }
 }

packages/vertexai/src/methods/generate-content.ts

@@ -28,20 +28,40 @@ import { processStream } from '../requests/stream-reader';
 import { ApiSettings } from '../types/internal';
 import { ChromeAdapter } from './chrome-adapter';
 
-export async function generateContentStream(
+async function generateContentStreamOnCloud(
   apiSettings: ApiSettings,
   model: string,
   params: GenerateContentRequest,
   requestOptions?: RequestOptions
-): Promise<GenerateContentStreamResult> {
-  const response = await makeRequest(
+): Promise<Response> {
+  return makeRequest(
     model,
     Task.STREAM_GENERATE_CONTENT,
     apiSettings,
     /* stream */ true,
     JSON.stringify(params),
     requestOptions
   );
+}
+
+export async function generateContentStream(
+  apiSettings: ApiSettings,
+  model: string,
+  params: GenerateContentRequest,
+  chromeAdapter: ChromeAdapter,
+  requestOptions?: RequestOptions
+): Promise<GenerateContentStreamResult> {
+  let response;
+  if (await chromeAdapter.isAvailable(params)) {
+    response = await chromeAdapter.generateContentStreamOnDevice(params);
+  } else {
+    response = await generateContentStreamOnCloud(
+      apiSettings,
+      model,
+      params,
+      requestOptions
+    );
+  }
   return processStream(response);
 }
 

packages/vertexai/src/models/generative-model.ts

@@ -123,6 +123,7 @@ export class GenerativeModel extends VertexAIModel {
         systemInstruction: this.systemInstruction,
         ...formattedParams
       },
+      this.chromeAdapter,
       this.requestOptions
     );
   }

yarn.lock

@@ -2947,6 +2947,11 @@
   resolved "https://registry.npmjs.org/@types/deep-eql/-/deep-eql-4.0.2.tgz#334311971d3a07121e7eb91b684a605e7eea9cbd"
   integrity sha512-c9h9dVVMigMPc4bwTvC5dxqtqJZwQPePsWjPlpSOnojbor6pGqdk541lfA7AqFQr5pB1BRdq0juY9db81BwyFw==
 
+"@types/[email protected]":
+  version "0.0.6"
+  resolved "https://registry.npmjs.org/@types/dom-chromium-ai/-/dom-chromium-ai-0.0.6.tgz#0c9e5712d8db3d26586cd9f175001b509cd2e514"
+  integrity sha512-/jUGe9a3BLzsjjg18Olk/Ul64PZ0P4aw8uNxrXeXVTni5PSxyCfyhHb4UohsXNVByOnwYGzlqUcb3vYKVsG4mg==
+
 "@types/eslint-scope@^3.7.7":
   version "3.7.7"
   resolved "https://registry.npmjs.org/@types/eslint-scope/-/eslint-scope-3.7.7.tgz#3108bd5f18b0cdb277c867b3dd449c9ed7079ac5"