Narrow Chrome prompt input type to seq<part>

Author: erikeldridge

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

@@ -22,7 +22,8 @@ import { ChromeAdapter } from './chrome-adapter';
 import {
   Availability,
   LanguageModel,
-  LanguageModelCreateOptions
+  LanguageModelCreateOptions,
+  LanguageModelMessageContent
 } from '../types/language-model';
 import { stub } from 'sinon';
 import { GenerateContentRequest } from '../types';
@@ -281,7 +282,8 @@ describe('ChromeAdapter', () => {
         create: () => Promise.resolve({})
       } as LanguageModel;
       const languageModel = {
-        prompt: i => Promise.resolve(i)
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        prompt: (p: LanguageModelMessageContent[]) => Promise.resolve('')
       } as LanguageModel;
       const createStub = stub(languageModelProvider, 'create').resolves(
         languageModel
@@ -305,13 +307,8 @@ describe('ChromeAdapter', () => {
       // 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
-            }
-          ]
+          type: 'text',
+          content: request.contents[0].parts[0].text
         }
       ]);
       // Asserts expected output.
@@ -366,21 +363,16 @@ describe('ChromeAdapter', () => {
       // Asserts Vertex input type is mapped to Chrome type.
       expect(measureInputUsageStub).to.have.been.calledOnceWith([
         {
-          role: 'user',
-          content: [
-            {
-              type: 'text',
-              content: inputText
-            }
-          ]
+          type: 'text',
+          content: inputText
         }
       ]);
       expect(await response.json()).to.deep.equal({
         totalTokens: expectedCount
       });
     });
   });
-  describe('generateContentStreamOnDevice', () => {
+  describe('generateContentStream', () => {
     it('generates content stream', async () => {
       const languageModelProvider = {
         create: () => Promise.resolve({})
@@ -413,13 +405,8 @@ describe('ChromeAdapter', () => {
       expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
       expect(promptStub).to.have.been.calledOnceWith([
         {
-          role: request.contents[0].role,
-          content: [
-            {
-              type: 'text',
-              content: request.contents[0].parts[0].text
-            }
-          ]
+          type: 'text',
+          content: request.contents[0].parts[0].text
         }
       ]);
       const actual = await toStringArray(response.body!);

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

@@ -16,19 +16,15 @@
  */
 
 import {
-  Content,
   CountTokensRequest,
   GenerateContentRequest,
   InferenceMode,
-  Part,
-  Role
+  Part
 } from '../types';
 import {
   Availability,
   LanguageModel,
   LanguageModelCreateOptions,
-  LanguageModelMessage,
-  LanguageModelMessageRole,
   LanguageModelMessageContent
 } from '../types/language-model';
 
@@ -100,8 +96,12 @@ export class ChromeAdapter {
       // TODO: normalize on-device params during construction.
       this.onDeviceParams || {}
     );
-    const messages = ChromeAdapter.toLanguageModelMessages(request.contents);
-    const text = await session.prompt(messages);
+    // TODO: support multiple content objects when Chrome supports
+    // sequence<LanguageModelMessage>
+    const contents = request.contents[0].parts.map(
+      ChromeAdapter.toLanguageModelMessageContent
+    );
+    const text = await session.prompt(contents);
     return ChromeAdapter.toResponse(text);
   }
 
@@ -120,8 +120,12 @@ export class ChromeAdapter {
       // TODO: normalize on-device params during construction.
       this.onDeviceParams || {}
     );
-    const messages = ChromeAdapter.toLanguageModelMessages(request.contents);
-    const stream = await session.promptStreaming(messages);
+    // TODO: support multiple content objects when Chrome supports
+    // sequence<LanguageModelMessage>
+    const contents = request.contents[0].parts.map(
+      ChromeAdapter.toLanguageModelMessageContent
+    );
+    const stream = await session.promptStreaming(contents);
     return ChromeAdapter.toStreamResponse(stream);
   }
 
@@ -131,8 +135,12 @@ export class ChromeAdapter {
       // TODO: normalize on-device params during construction.
       this.onDeviceParams || {}
     );
-    const messages = ChromeAdapter.toLanguageModelMessages(request.contents);
-    const tokenCount = await session.measureInputUsage(messages);
+    // TODO: support multiple content objects when Chrome supports
+    // sequence<LanguageModelMessage>
+    const contents = request.contents[0].parts.map(
+      ChromeAdapter.toLanguageModelMessageContent
+    );
+    const tokenCount = await session.measureInputUsage(contents);
     return {
       json: async () => ({
         totalTokens: tokenCount
@@ -188,25 +196,6 @@ export class ChromeAdapter {
       });
   }
 
-  /**
-   * Converts a Vertex role string to a Chrome role string.
-   */
-  private static toOnDeviceRole(role: Role): LanguageModelMessageRole {
-    return role === 'model' ? 'assistant' : 'user';
-  }
-
-  /**
-   * Converts a Vertex Content object to a Chrome LanguageModelMessage object.
-   */
-  private static toLanguageModelMessages(
-    contents: Content[]
-  ): LanguageModelMessage[] {
-    return contents.map(c => ({
-      role: ChromeAdapter.toOnDeviceRole(c.role),
-      content: c.parts.map(ChromeAdapter.toLanguageModelMessageContent)
-    }));
-  }
-
   /**
    * Converts a Vertex Part object to a Chrome LanguageModelMessageContent object.
    */

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

@@ -56,14 +56,12 @@ interface LanguageModelExpectedInput {
   type: LanguageModelMessageType;
   languages?: string[];
 }
-export type LanguageModelPrompt =
-  | LanguageModelMessage[]
-  | LanguageModelMessageShorthand[]
-  | string;
+// TODO: revert to type from Prompt API explainer once it's supported.
+export type LanguageModelPrompt = LanguageModelMessageContent[];
 type LanguageModelInitialPrompts =
   | LanguageModelMessage[]
   | LanguageModelMessageShorthand[];
-export interface LanguageModelMessage {
+interface LanguageModelMessage {
   role: LanguageModelMessageRole;
   content: LanguageModelMessageContent[];
 }
@@ -75,7 +73,7 @@ export interface LanguageModelMessageContent {
   type: LanguageModelMessageType;
   content: LanguageModelMessageContentValue;
 }
-export type LanguageModelMessageRole = 'system' | 'user' | 'assistant';
+type LanguageModelMessageRole = 'system' | 'user' | 'assistant';
 type LanguageModelMessageType = 'text' | 'image' | 'audio';
 type LanguageModelMessageContentValue =
   | ImageBitmapSource