Integrate chrome adapter into stream methods

Author: erikeldridge

e2e/sample-apps/modular.js

@@ -58,16 +58,7 @@ import {
   onValue,
   off
 } from 'firebase/database';
-<<<<<<< HEAD
 import { getGenerativeModel, getVertexAI } from 'firebase/vertexai';
-=======
-import {
-  getGenerativeModel,
-  getVertexAI,
-  InferenceMode,
-  VertexAI
-} from 'firebase/vertexai';
->>>>>>> 3f02db006 (Run yarn format)
 import { getDataConnect, DataConnect } from 'firebase/data-connect';
 
 /**
@@ -322,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}`);
 }
 
@@ -341,15 +337,6 @@ function callDataConnect(app) {
   console.log('[DATACONNECT] initialized');
 }
 
-async function callVertex(app) {
-  console.log('[VERTEX] start');
-  const vertex = getVertexAI(app);
-  const model = getGenerativeModel(vertex, { mode: 'prefer_on_device' });
-  const result = await model.generateContent("What is Roko's Basalisk?");
-  console.log(result.response.text());
-  console.log('[VERTEX] initialized');
-}
-
 /**
  * Run smoke tests for all products.
  * Comment out any products you want to ignore.
@@ -371,7 +358,6 @@ async function main() {
   await callVertexAI(app);
   callDataConnect(app);
   await authLogout(app);
-  await callVertex(app);
   console.log('DONE');
 }
 

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

@@ -117,6 +117,46 @@ export class ChromeAdapter {
     } as Response;
   }
 
+  /**
+   * Generates a stream of content.
+   */
+  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);
+  }
+  // Formats string stream returned by Chrome as SSE returned by Vertex.
+  private static async toStreamResponse(
+    stream: ReadableStream<string>
+  ): Promise<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;
+  }
+
   /**
    * Asserts inference for the given request can be performed by an on-device model.
    */

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
     );
   }