Integrate chrome adapter into stream methods

Author: erikeldridge

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, VertexAI } 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}`);
 }
 
@@ -337,17 +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: InferenceMode.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.
@@ -357,19 +346,18 @@ async function main() {
   const app = initializeApp(config);
   setLogLevel('warn');
 
-  callAppCheck(app);
-  await authLogin(app);
-  await callStorage(app);
-  await callFirestore(app);
-  await callDatabase(app);
-  await callMessaging(app);
-  callAnalytics(app);
-  callPerformance(app);
-  await callFunctions(app);
+  // callAppCheck(app);
+  // await authLogin(app);
+  // await callStorage(app);
+  // await callFirestore(app);
+  // await callDatabase(app);
+  // await callMessaging(app);
+  // callAnalytics(app);
+  // callPerformance(app);
+  // await callFunctions(app);
   await callVertexAI(app);
-  callDataConnect(app);
-  await authLogout(app);
-  await callVertex(app);
+  // callDataConnect(app);
+  // await authLogout(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

@@ -78,7 +78,9 @@ export class ChromeAdapter {
   ): Promise<EnhancedGenerateContentResponse> {
     const createOptions = this.onDeviceParams || {};
     createOptions.initialPrompts ??= [];
-    const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(request.contents);
+    const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(
+      request.contents
+    );
     // Assumes validation asserted there is at least one initial prompt.
     const prompt = extractedInitialPrompts.pop()!;
     createOptions.initialPrompts.push(...extractedInitialPrompts);
@@ -89,6 +91,46 @@ export class ChromeAdapter {
       functionCalls: () => undefined
     };
   }
+  async generateContentStreamOnDevice(
+    request: GenerateContentRequest
+  ): Promise<Response> {
+    const createOptions = this.onDeviceParams || {};
+    createOptions.initialPrompts ??= [];
+    const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(
+      request.contents
+    );
+    // Assumes validation asserted there is at least one initial prompt.
+    const prompt = extractedInitialPrompts.pop()!;
+    createOptions.initialPrompts.push(...extractedInitialPrompts);
+    const session = await this.session(createOptions);
+    const stream = await session.promptStreaming(prompt.content);
+    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;
+  }
   private static isOnDeviceRequest(request: GenerateContentRequest): boolean {
     // Returns false if the prompt is empty.
     if (request.contents.length === 0) {

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

@@ -29,20 +29,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
     );
   }