Merge 86d47b0 into d546aec

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.test.ts

@@ -108,42 +108,6 @@ describe('ChromeAdapter', () => {
         })
       ).to.be.false;
     });
-    it('returns false if request system instruction has function role', async () => {
-      const adapter = new ChromeAdapter({} as AI, 'prefer_on_device');
-      expect(
-        await adapter.isAvailable({
-          contents: [],
-          systemInstruction: {
-            role: 'function',
-            parts: []
-          }
-        })
-      ).to.be.false;
-    });
-    it('returns false if request system instruction has multiple parts', async () => {
-      const adapter = new ChromeAdapter({} as AI, 'prefer_on_device');
-      expect(
-        await adapter.isAvailable({
-          contents: [],
-          systemInstruction: {
-            role: 'function',
-            parts: [{ text: 'a' }, { text: 'b' }]
-          }
-        })
-      ).to.be.false;
-    });
-    it('returns false if request system instruction has non-text part', async () => {
-      const adapter = new ChromeAdapter({} as AI, 'prefer_on_device');
-      expect(
-        await adapter.isAvailable({
-          contents: [],
-          systemInstruction: {
-            role: 'function',
-            parts: [{ inlineData: { mimeType: 'a', data: 'b' } }]
-          }
-        })
-      ).to.be.false;
-    });
     it('returns true if model is readily available', async () => {
       const aiProvider = {
         languageModel: {
@@ -173,13 +137,18 @@ describe('ChromeAdapter', () => {
       const createStub = stub(aiProvider.languageModel, 'create').resolves(
         {} as AILanguageModel
       );
-      const adapter = new ChromeAdapter(aiProvider, 'prefer_on_device');
+      const onDeviceParams = {} as AILanguageModelCreateOptionsWithSystemPrompt;
+      const adapter = new ChromeAdapter(
+        aiProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
       expect(
         await adapter.isAvailable({
           contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
         })
       ).to.be.false;
-      expect(createStub).to.have.been.calledOnce;
+      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
     });
     it('avoids redundant downloads', async () => {
       const aiProvider = {
@@ -250,4 +219,83 @@ describe('ChromeAdapter', () => {
       ).to.be.false;
     });
   });
+  describe('generateContentOnDevice', () => {
+    it('Extracts and concats initial prompts', async () => {
+      const aiProvider = {
+        languageModel: {
+          create: () => Promise.resolve({})
+        }
+      } as AI;
+      const factoryStub = stub(aiProvider.languageModel, 'create').resolves({
+        prompt: s => Promise.resolve(s)
+      } as AILanguageModel);
+      const text = ['first', 'second', 'third'];
+      const onDeviceParams = {
+        initialPrompts: [{ role: 'user', content: text[0] }]
+      } as AILanguageModelCreateOptionsWithSystemPrompt;
+      const adapter = new ChromeAdapter(
+        aiProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
+      const response = await adapter.generateContentOnDevice({
+        contents: [
+          { role: 'model', parts: [{ text: text[1] }] },
+          { role: 'user', parts: [{ text: text[2] }] }
+        ]
+      });
+      expect(factoryStub).to.have.been.calledOnceWith({
+        initialPrompts: [
+          { role: 'user', content: text[0] },
+          // Asserts tail is passed as initial prompts, and
+          // role is normalized from model to assistant.
+          { role: 'assistant', content: text[1] }
+        ]
+      });
+      expect(await response.json()).to.deep.equal({
+        candidates: [
+          {
+            content: {
+              parts: [{ text: text[2] }]
+            }
+          }
+        ]
+      });
+    });
+    it('Extracts system prompt', async () => {
+      const aiProvider = {
+        languageModel: {
+          create: () => Promise.resolve({})
+        }
+      } as AI;
+      const factoryStub = stub(aiProvider.languageModel, 'create').resolves({
+        prompt: s => Promise.resolve(s)
+      } as AILanguageModel);
+      const onDeviceParams = {
+        systemPrompt: 'be yourself'
+      } as AILanguageModelCreateOptionsWithSystemPrompt;
+      const adapter = new ChromeAdapter(
+        aiProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
+      const text = 'hi';
+      const response = await adapter.generateContentOnDevice({
+        contents: [{ role: 'user', parts: [{ text }] }]
+      });
+      expect(factoryStub).to.have.been.calledOnceWith({
+        initialPrompts: [],
+        systemPrompt: onDeviceParams.systemPrompt
+      });
+      expect(await response.json()).to.deep.equal({
+        candidates: [
+          {
+            content: {
+              parts: [{ text }]
+            }
+          }
+        ]
+      });
+    });
+  });
 });