Infer expected inputs from prompt

Author: erikeldridge

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

@@ -54,7 +54,7 @@ async function toStringArray(
 
 describe('ChromeAdapter', () => {
   describe('constructor', () => {
-    it('sets image as expected input type by default', async () => {
+    it('determines expected inputs by request inspection', async () => {
       const languageModelProvider = {
         availability: () => Promise.resolve(Availability.available)
       } as LanguageModel;
@@ -70,7 +70,11 @@ describe('ChromeAdapter', () => {
         contents: [
           {
             role: 'user',
-            parts: [{ text: 'hi' }]
+            parts: [
+              { text: 'hi' },
+              // Triggers image as expected type.
+              { inlineData: { mimeType: 'image/jpeg', data: 'asd' } }
+            ]
           }
         ]
       });

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

@@ -30,9 +30,11 @@ import {
 import {
   Availability,
   LanguageModel,
+  LanguageModelExpected,
   LanguageModelMessage,
   LanguageModelMessageContent,
-  LanguageModelMessageRole
+  LanguageModelMessageRole,
+  LanguageModelMessageType
 } from '../types/language-model';
 
 /**
@@ -48,13 +50,10 @@ export class ChromeAdapter {
   constructor(
     private languageModelProvider?: LanguageModel,
     private mode?: InferenceMode,
-    private onDeviceParams: OnDeviceParams = {
-      createOptions: {
-        // Defaults to support image inputs for convenience.
-        expectedInputs: [{ type: 'image' }]
-      }
-    }
-  ) {}
+    private onDeviceParams: OnDeviceParams = {}
+  ) {
+    this.onDeviceParams.createOptions ??= {};
+  }
 
   /**
    * Checks if a given request can be made on-device.
@@ -85,8 +84,10 @@ export class ChromeAdapter {
       return false;
     }
 
+    const expectedInputs = ChromeAdapter.extractExpectedInputs(request);
+
     // Triggers out-of-band download so model will eventually become available.
-    const availability = await this.downloadIfAvailable();
+    const availability = await this.downloadIfAvailable(expectedInputs);
 
     if (this.mode === 'only_on_device') {
       return true;
@@ -158,6 +159,33 @@ export class ChromeAdapter {
     );
   }
 
+  /**
+   * Maps
+   * <a href="https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference#blob">
+   * Vertex's input mime types</a> to
+   * <a href="https://github.com/webmachinelearning/prompt-api?tab=readme-ov-file#full-api-surface-in-web-idl">
+   * Chrome's expected types</a>.
+   *
+   * <p>Chrome's API checks availability by type. It's tedious to specify the types in advance, so
+   * this method infers the types.</p>
+   */
+  private static extractExpectedInputs(
+    request: GenerateContentRequest
+  ): LanguageModelExpected[] {
+    const inputSet = new Set<LanguageModelExpected>();
+    for (const content of request.contents) {
+      for (const part of content.parts) {
+        if (part.inlineData) {
+          const type = part.inlineData.mimeType.split(
+            '/'
+          )[0] as LanguageModelMessageType;
+          inputSet.add({ type });
+        }
+      }
+    }
+    return Array.from(inputSet);
+  }
+
   /**
    * Asserts inference for the given request can be performed by an on-device model.
    */
@@ -196,12 +224,21 @@ export class ChromeAdapter {
   /**
    * Encapsulates logic to get availability and download a model if one is downloadable.
    */
-  private async downloadIfAvailable(): Promise<Availability | undefined> {
+  private async downloadIfAvailable(
+    expectedInputs: LanguageModelExpected[]
+  ): Promise<Availability | undefined> {
+    // Side-effect: updates construction-time params with request-time params.
+    // This is required because params are referenced through multiple flows.
+    // TODO: remove this side effect, since we need to also pass options when creating a session.
+    Object.assign(this.onDeviceParams.createOptions!, { expectedInputs });
+
     const availability = await this.languageModelProvider?.availability(
       this.onDeviceParams.createOptions
     );
 
     if (availability === Availability.downloadable) {
+      // Side-effect: triggers out-of-band model download.
+      // This is required because Chrome manages the model download.
       this.download();
     }