Support structured output

Author: erikeldridge

e2e/sample-apps/modular.js

@@ -58,7 +58,7 @@ import {
   onValue,
   off
 } from 'firebase/database';
-import { getGenerativeModel, getVertexAI } from 'firebase/vertexai';
+import { getGenerativeModel, getVertexAI, Schema } from 'firebase/vertexai';
 import { getDataConnect, DataConnect } from 'firebase/data-connect';
 
 /**
@@ -313,23 +313,48 @@ function callPerformance(app) {
 async function callVertexAI(app) {
   console.log('[VERTEXAI] start');
   const vertexAI = getVertexAI(app);
-  const model = getGenerativeModel(vertexAI, {
-    mode: 'prefer_on_device'
+
+  const responseSchema = Schema.object({
+    properties: {
+      characters: Schema.array({
+        items: Schema.object({
+          properties: {
+            name: Schema.string(),
+            accessory: Schema.string(),
+            age: Schema.number(),
+            species: Schema.string()
+          },
+          optionalProperties: ['accessory']
+        })
+      })
+    }
   });
-  const singleResult = await model.generateContent([
-    { text: 'describe this 20 x 20 px image in two words' },
-    {
-      inlineData: {
-        mimeType: 'image/heic',
-        data: 'AAAAGGZ0eXBoZWljAAAAAGhlaWNtaWYxAAAB7G1ldGEAAAAAAAAAIWhkbHIAAAAAAAAAAHBpY3QAAAAAAAAAAAAAAAAAAAAAJGRpbmYAAAAcZHJlZgAAAAAAAAABAAAADHVybCAAAAABAAAADnBpdG0AAAAAAAEAAAA4aWluZgAAAAAAAgAAABVpbmZlAgAAAAABAABodmMxAAAAABVpbmZlAgAAAQACAABFeGlmAAAAABppcmVmAAAAAAAAAA5jZHNjAAIAAQABAAABD2lwcnAAAADtaXBjbwAAABNjb2xybmNseAACAAIABoAAAAAMY2xsaQDLAEAAAAAUaXNwZQAAAAAAAAAUAAAADgAAAChjbGFwAAAAFAAAAAEAAAANAAAAAQAAAAAAAAAB/8AAAACAAAAAAAAJaXJvdAAAAAAQcGl4aQAAAAADCAgIAAAAcWh2Y0MBA3AAAACwAAAAAAAe8AD8/fj4AAALA6AAAQAXQAEMAf//A3AAAAMAsAAAAwAAAwAecCShAAEAI0IBAQNwAAADALAAAAMAAAMAHqAUIEHAjw1iHuRZVNwICBgCogABAAlEAcBhcshAUyQAAAAaaXBtYQAAAAAAAAABAAEHgQIDhIUGhwAAACxpbG9jAAAAAEQAAAIAAQAAAAEAAAJsAAABDAACAAAAAQAAAhQAAABYAAAAAW1kYXQAAAAAAAABdAAAAAZFeGlmAABNTQAqAAAACAAEARIAAwAAAAEAAQAAARoABQAAAAEAAAA+ARsABQAAAAEAAABGASgAAwAAAAEAAgAAAAAAAAAAAEgAAAABAAAASAAAAAEAAAEIKAGvoR8wDimTiRYUbALiHkU3ZdZ8DXAcSrRB9GARtVQHvnCE0LEyBGAyb5P4eYr6JAK5UxNX10WNlARq3ZpcGeVD+Xom6LodYasuZKKtDHCz/xnswOtC/ksZzVKhtWQqGvkXcsJnLYqWevNkacnccQ95jbHJBg9nXub69jAAN3xhNOXxjGSxaG9QvES5R7sYICEojRjLF5OB5K3v+okQAwfgWpz/u21ayideOgOZQLAyBkKOv7ymLNCagiPWTlHAuy/3qR1Q7m2ERFaxKIAbLSkIVO/P8m8+anKxhzhC//L8NMAUoF+Sf3aEH9O41fwLc+PlcbrDrjgY2EboD3cn9DyN32Rum2Ym'
+
+  const model = getGenerativeModel(vertexAI, {
+    mode: 'prefer_on_device',
+    inCloudParams: {
+      generationConfig: {
+        responseSchema
       }
     }
-  ]);
-  console.log(`Generated text: ${singleResult.response.text()}`);
-  const chat = model.startChat();
-  let chatResult = await chat.sendMessage('describe red in two words');
-  chatResult = await chat.sendMessage('describe blue');
-  console.log('Chat history:', await chat.getHistory());
+  });
+
+  const singleResult = await model.generateContent({
+    generationConfig: {
+      responseSchema
+    },
+    contents: [
+      {
+        role: 'user',
+        parts: [
+          {
+            text: "For use in a children's card game, generate 10 animal-based characters."
+          }
+        ]
+      }
+    ]
+  });
+  console.log(`Generated text:`, JSON.parse(singleResult.response.text()));
   console.log(`[VERTEXAI] end`);
 }
 

e2e/webpack.config.js

@@ -88,7 +88,7 @@ module.exports = [
     stats: {
       colors: true
     },
-    devtool: 'source-map',
+    devtool: 'eval-source-map',
     devServer: {
       static: './build'
     }

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

@@ -28,6 +28,7 @@ import {
 } from '../types/language-model';
 import { match, stub } from 'sinon';
 import { GenerateContentRequest, AIErrorCode } from '../types';
+import { Schema } from '../api';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -406,6 +407,51 @@ describe('ChromeAdapter', () => {
         ]
       });
     });
+    it('honors response constraint', async () => {
+      const languageModel = {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        prompt: (p: LanguageModelMessageContent[]) => Promise.resolve('')
+      } as LanguageModel;
+      const languageModelProvider = {
+        create: () => Promise.resolve(languageModel)
+      } as LanguageModel;
+      const promptOutput = '{}';
+      const promptStub = stub(languageModel, 'prompt').resolves(promptOutput);
+      const adapter = new ChromeAdapter(
+        languageModelProvider,
+        'prefer_on_device'
+      );
+      const responseSchema = Schema.object({
+        properties: {}
+      });
+      const request = {
+        generationConfig: {
+          responseSchema
+        },
+        contents: [{ role: 'user', parts: [{ text: 'anything' }] }]
+      } as GenerateContentRequest;
+      const response = await adapter.generateContent(request);
+      expect(promptStub).to.have.been.calledOnceWith(
+        [
+          {
+            type: 'text',
+            content: request.contents[0].parts[0].text
+          }
+        ],
+        {
+          responseConstraint: responseSchema
+        }
+      );
+      expect(await response.json()).to.deep.equal({
+        candidates: [
+          {
+            content: {
+              parts: [{ text: promptOutput }]
+            }
+          }
+        ]
+      });
+    });
   });
   describe('countTokens', () => {
     it('counts tokens is not yet available', async () => {

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

@@ -28,7 +28,8 @@ import {
   Availability,
   LanguageModel,
   LanguageModelCreateOptions,
-  LanguageModelMessageContent
+  LanguageModelMessageContent,
+  LanguageModelPromptOptions
 } from '../types/language-model';
 
 /**
@@ -106,15 +107,27 @@ export class ChromeAdapter {
    */
   async generateContent(request: GenerateContentRequest): Promise<Response> {
     const session = await this.createSession();
+
     // TODO: support multiple content objects when Chrome supports
     // sequence<LanguageModelMessage>
     const contents = await Promise.all(
       request.contents[0].parts.map(ChromeAdapter.toLanguageModelMessageContent)
     );
-    const text = await session.prompt(contents);
+
+    const options = ChromeAdapter.extractLanguageModelPromptOptions(request);
+    
+    const text = await session.prompt(contents, options);
     return ChromeAdapter.toResponse(text);
   }
 
+  private static extractLanguageModelPromptOptions(
+    request: GenerateContentRequest
+  ): LanguageModelPromptOptions {
+    return {
+      responseConstraint: request.generationConfig?.responseSchema
+    };
+  }
+
   /**
    * Generates content stream on device.
    *

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

@@ -49,8 +49,9 @@ export interface LanguageModelCreateOptions
   systemPrompt?: string;
   initialPrompts?: LanguageModelInitialPrompts;
 }
-interface LanguageModelPromptOptions {
-  signal?: AbortSignal;
+export interface LanguageModelPromptOptions {
+  responseConstraint?: object;
+  // TODO: Restore AbortSignal once the API is defined.
 }
 interface LanguageModelExpectedInput {
   type: LanguageModelMessageType;