Test content generation

Author: erikeldridge

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

@@ -19,9 +19,14 @@ import { expect, use } from 'chai';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import { ChromeAdapter } from './chrome-adapter';
-import { Availability, LanguageModel } from '../types/language-model';
+import {
+  Availability,
+  LanguageModel,
+  LanguageModelCreateOptions
+} from '../types/language-model';
 import { stub } from 'sinon';
 import * as util from '@firebase/util';
+import { GenerateContentRequest } from '../types';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -191,16 +196,18 @@ describe('ChromeAdapter', () => {
       const createStub = stub(languageModelProvider, 'create').resolves(
         {} as LanguageModel
       );
+      const onDeviceParams = {} as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
-        'prefer_on_device'
+        '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 languageModelProvider = {
@@ -260,4 +267,55 @@ describe('ChromeAdapter', () => {
       ).to.be.false;
     });
   });
+  describe('generateContentOnDevice', () => {
+    it('generates content', async () => {
+      const languageModelProvider = {
+        create: () => Promise.resolve({})
+      } as LanguageModel;
+      const languageModel = {
+        prompt: i => Promise.resolve(i)
+      } as LanguageModel;
+      const createStub = stub(languageModelProvider, 'create').resolves(
+        languageModel
+      );
+      const promptOutput = 'hi';
+      const promptStub = stub(languageModel, 'prompt').resolves(promptOutput);
+      const onDeviceParams = {
+        systemPrompt: 'be yourself'
+      } as LanguageModelCreateOptions;
+      const adapter = new ChromeAdapter(
+        languageModelProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
+      const request = {
+        contents: [{ role: 'user', parts: [{ text: 'anything' }] }]
+      } as GenerateContentRequest;
+      const response = await adapter.generateContentOnDevice(request);
+      // Asserts initialization params are proxied.
+      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
+      // Asserts Vertex input type is mapped to Chrome type.
+      expect(promptStub).to.have.been.calledOnceWith([
+        {
+          role: request.contents[0].role,
+          content: [
+            {
+              type: 'text',
+              content: request.contents[0].parts[0].text
+            }
+          ]
+        }
+      ]);
+      // Asserts expected output.
+      expect(await response.json()).to.deep.equal({
+        candidates: [
+          {
+            content: {
+              parts: [{ text: promptOutput }]
+            }
+          }
+        ]
+      });
+    });
+  });
 });

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

@@ -27,8 +27,9 @@ import {
   Availability,
   LanguageModel,
   LanguageModelCreateOptions,
+  LanguageModelMessage,
   LanguageModelMessageRole,
-  LanguageModelMessageShorthand
+  LanguageModelMessageContent
 } from '../types/language-model';
 import { isChrome } from '@firebase/util';
 
@@ -85,17 +86,13 @@ export class ChromeAdapter {
   async generateContentOnDevice(
     request: GenerateContentRequest
   ): Promise<Response> {
-    const initialPrompts = ChromeAdapter.toInitialPrompts(request.contents);
-    // Assumes validation asserted there is at least one initial prompt.
-    const prompt = initialPrompts.pop()!;
-    const systemPrompt = ChromeAdapter.toSystemPrompt(
-      request.systemInstruction
+    const session = await this.session(
+      // TODO: normalize on-device params during construction.
+      this.onDeviceParams || {}
     );
-    const session = await this.session({
-      initialPrompts,
-      systemPrompt
-    });
-    const text = await session.prompt(prompt.content);
+    const messages = ChromeAdapter.toLanguageModelMessages(request.contents);
+    const text = await session.prompt(messages);
+    console.log(text);
     return {
       json: () =>
         Promise.resolve({
@@ -162,45 +159,30 @@ export class ChromeAdapter {
         this.isDownloading = false;
       });
   }
-  private static toSystemPrompt(
-    prompt: string | Content | Part | undefined
-  ): string | undefined {
-    if (!prompt) {
-      return undefined;
-    }
-
-    if (typeof prompt === 'string') {
-      return prompt;
-    }
-
-    const systemContent = prompt as Content;
-    if (
-      systemContent.parts &&
-      systemContent.parts[0] &&
-      systemContent.parts[0].text
-    ) {
-      return systemContent.parts[0].text;
-    }
-
-    const systemPart = prompt as Part;
-    if (systemPart.text) {
-      return systemPart.text;
-    }
-
-    return undefined;
-  }
   private static toOnDeviceRole(role: Role): LanguageModelMessageRole {
     return role === 'model' ? 'assistant' : 'user';
   }
-  private static toInitialPrompts(
+  private static toLanguageModelMessages(
     contents: Content[]
-  ): LanguageModelMessageShorthand[] {
+  ): LanguageModelMessage[] {
     return contents.map(c => ({
       role: ChromeAdapter.toOnDeviceRole(c.role),
-      // Assumes contents have been verified to contain only a single TextPart.
-      content: c.parts[0].text!
+      content: c.parts.map(ChromeAdapter.toLanguageModelMessageContent)
     }));
   }
+  private static toLanguageModelMessageContent(
+    part: Part
+  ): LanguageModelMessageContent {
+    if (part.text) {
+      return {
+        type: 'text',
+        content: part.text
+      };
+    }
+    // Assumes contents have been verified to contain only a single TextPart.
+    // TODO: support other input types
+    throw new Error('Not yet implemented');
+  }
   private async session(
     opts: LanguageModelCreateOptions
   ): Promise<LanguageModel> {

packages/vertexai/src/types/language-model.ts

@@ -56,22 +56,22 @@ interface LanguageModelExpectedInput {
   type: LanguageModelMessageType;
   languages?: string[];
 }
-type LanguageModelPrompt =
+export type LanguageModelPrompt =
   | LanguageModelMessage[]
   | LanguageModelMessageShorthand[]
   | string;
 type LanguageModelInitialPrompts =
   | LanguageModelMessage[]
   | LanguageModelMessageShorthand[];
-interface LanguageModelMessage {
+export interface LanguageModelMessage {
   role: LanguageModelMessageRole;
   content: LanguageModelMessageContent[];
 }
-export interface LanguageModelMessageShorthand {
+interface LanguageModelMessageShorthand {
   role: LanguageModelMessageRole;
   content: string;
 }
-interface LanguageModelMessageContent {
+export interface LanguageModelMessageContent {
   type: LanguageModelMessageType;
   content: LanguageModelMessageContentValue;
 }