Assert options extracted in availability and generate content paths

Author: erikeldridge

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

@@ -54,34 +54,6 @@ async function toStringArray(
 
 describe('ChromeAdapter', () => {
   describe('constructor', () => {
-    it('determines expected inputs by request inspection', async () => {
-      const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.available)
-      } as LanguageModel;
-      const availabilityStub = stub(
-        languageModelProvider,
-        'availability'
-      ).resolves(Availability.available);
-      const adapter = new ChromeAdapter(
-        languageModelProvider,
-        'prefer_on_device'
-      );
-      await adapter.isAvailable({
-        contents: [
-          {
-            role: 'user',
-            parts: [
-              { text: 'hi' },
-              // Triggers image as expected type.
-              { inlineData: { mimeType: 'image/jpeg', data: 'asd' } }
-            ]
-          }
-        ]
-      });
-      expect(availabilityStub).to.have.been.calledWith({
-        expectedInputs: [{ type: 'image' }]
-      });
-    });
     it('honors explicitly set expected inputs', async () => {
       const languageModelProvider = {
         availability: () => Promise.resolve(Availability.available)
@@ -303,6 +275,34 @@ describe('ChromeAdapter', () => {
         })
       ).to.be.false;
     });
+    it('extracts expected inputs from the request', async () => {
+      const languageModelProvider = {
+        availability: () => Promise.resolve(Availability.available)
+      } as LanguageModel;
+      const availabilityStub = stub(
+        languageModelProvider,
+        'availability'
+      ).resolves(Availability.available);
+      const adapter = new ChromeAdapter(
+        languageModelProvider,
+        'prefer_on_device'
+      );
+      await adapter.isAvailable({
+        contents: [
+          {
+            role: 'user',
+            parts: [
+              { text: 'hi' },
+              // Triggers image as expected type.
+              { inlineData: { mimeType: 'image/jpeg', data: 'asd' } }
+            ]
+          }
+        ]
+      });
+      expect(availabilityStub).to.have.been.calledWith({
+        expectedInputs: [{ type: 'image' }]
+      });
+    });
   });
   describe('generateContent', () => {
     it('throws if Chrome API is undefined', async () => {
@@ -382,14 +382,9 @@ describe('ChromeAdapter', () => {
       );
       const promptOutput = 'hi';
       const promptStub = stub(languageModel, 'prompt').resolves(promptOutput);
-      const createOptions = {
-        systemPrompt: 'be yourself',
-        expectedInputs: [{ type: 'image' }]
-      } as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
-        'prefer_on_device',
-        { createOptions }
+        'prefer_on_device'
       );
       const request = {
         contents: [
@@ -409,7 +404,9 @@ describe('ChromeAdapter', () => {
       } as GenerateContentRequest;
       const response = await adapter.generateContent(request);
       // Asserts initialization params are proxied.
-      expect(createStub).to.have.been.calledOnceWith(createOptions);
+      expect(createStub).to.have.been.calledOnceWith({
+        expectedInputs: [{ type: 'image' }]
+      });
       // Asserts Vertex input type is mapped to Chrome type.
       expect(promptStub).to.have.been.calledOnceWith([
         {
@@ -610,13 +607,9 @@ describe('ChromeAdapter', () => {
           }
         })
       );
-      const createOptions = {
-        expectedInputs: [{ type: 'image' }]
-      } as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
-        'prefer_on_device',
-        { createOptions }
+        'prefer_on_device'
       );
       const request = {
         contents: [
@@ -635,7 +628,9 @@ describe('ChromeAdapter', () => {
         ]
       } as GenerateContentRequest;
       const response = await adapter.generateContentStream(request);
-      expect(createStub).to.have.been.calledOnceWith(createOptions);
+      expect(createStub).to.have.been.calledOnceWith({
+        expectedInputs: [{ type: 'image' }]
+      });
       expect(promptStub).to.have.been.calledOnceWith([
         {
           role: request.contents[0].role,

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

@@ -84,8 +84,8 @@ export class ChromeAdapter {
       return false;
     }
 
-    const requestOptions = this.inferCreateOptions(request);
-    const mergedOptions = this.mergeCreateOptions(requestOptions);
+    const extractedOptions = this.extractCreateOptions(request);
+    const mergedOptions = this.mergeCreateOptions(extractedOptions);
 
     // Triggers out-of-band download so model will eventually become available.
     const availability = await this.downloadIfAvailable(mergedOptions);
@@ -120,8 +120,8 @@ export class ChromeAdapter {
    * @returns {@link Response}, so we can reuse common response formatting.
    */
   async generateContent(request: GenerateContentRequest): Promise<Response> {
-    const requestOptions = this.inferCreateOptions(request);
-    const mergedOptions = this.mergeCreateOptions(requestOptions);
+    const extractedOptions = this.extractCreateOptions(request);
+    const mergedOptions = this.mergeCreateOptions(extractedOptions);
     const session = await this.createSession(mergedOptions);
     const contents = await Promise.all(
       request.contents.map(ChromeAdapter.toLanguageModelMessage)
@@ -144,8 +144,8 @@ export class ChromeAdapter {
   async generateContentStream(
     request: GenerateContentRequest
   ): Promise<Response> {
-    const inferredOptions = this.inferCreateOptions(request);
-    const mergedOptions = this.mergeCreateOptions(inferredOptions);
+    const extractedOptions = this.extractCreateOptions(request);
+    const mergedOptions = this.mergeCreateOptions(extractedOptions);
     const session = await this.createSession(mergedOptions);
     const contents = await Promise.all(
       request.contents.map(ChromeAdapter.toLanguageModelMessage)
@@ -165,16 +165,18 @@ export class ChromeAdapter {
   }
 
   /**
-   * Maps
+   * Extracts session creation options specified at request-time.
+   *
+   * <p>In particular, this method 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 input types</a>.
+   * Chrome's expected input types</a>.</p>
    *
    * <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 inferCreateOptions(
+  private extractCreateOptions(
     request: GenerateContentRequest
   ): LanguageModelCreateOptions {
     const inputSet = new Set<LanguageModelExpected>();