Add test for generateContentStream

Author: erikeldridge

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

@@ -407,7 +407,7 @@ describe('ChromeAdapter', () => {
         ]
       });
     });
-    it('honors response constraint', async () => {
+    it('supports structured output by mapping responseSchema to responseConstraint', async () => {
       const languageModel = {
         // eslint-disable-next-line @typescript-eslint/no-unused-vars
         prompt: (p: LanguageModelMessageContent[]) => Promise.resolve('')
@@ -430,7 +430,7 @@ describe('ChromeAdapter', () => {
         },
         contents: [{ role: 'user', parts: [{ text: 'anything' }] }]
       } as GenerateContentRequest;
-      const response = await adapter.generateContent(request);
+      await adapter.generateContent(request);
       expect(promptStub).to.have.been.calledOnceWith(
         [
           {
@@ -442,15 +442,6 @@ describe('ChromeAdapter', () => {
           responseConstraint: responseSchema
         }
       );
-      expect(await response.json()).to.deep.equal({
-        candidates: [
-          {
-            content: {
-              parts: [{ text: promptOutput }]
-            }
-          }
-        ]
-      });
     });
   });
   describe('countTokens', () => {
@@ -592,6 +583,43 @@ describe('ChromeAdapter', () => {
         `data: {"candidates":[{"content":{"role":"model","parts":[{"text":["${part}"]}]}}]}\n\n`
       ]);
     });
+    it('supports structured output by mapping responseSchema to responseConstraint', async () => {
+      const languageModel = {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        promptStreaming: p => new ReadableStream()
+      } as LanguageModel;
+      const languageModelProvider = {
+        create: () => Promise.resolve(languageModel)
+      } as LanguageModel;
+      const promptStub = stub(languageModel, 'promptStreaming').returns(
+        new ReadableStream()
+      );
+      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;
+      await adapter.generateContentStream(request);
+      expect(promptStub).to.have.been.calledOnceWith(
+        [
+          {
+            type: 'text',
+            content: request.contents[0].parts[0].text
+          }
+        ],
+        {
+          responseConstraint: responseSchema
+        }
+      );
+    });
   });
 });
 

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

@@ -107,27 +107,16 @@ 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 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.
    *
@@ -145,7 +134,8 @@ export class ChromeAdapter {
     const contents = await Promise.all(
       request.contents[0].parts.map(ChromeAdapter.toLanguageModelMessageContent)
     );
-    const stream = await session.promptStreaming(contents);
+    const options = ChromeAdapter.extractLanguageModelPromptOptions(request);
+    const stream = await session.promptStreaming(contents, options);
     return ChromeAdapter.toStreamResponse(stream);
   }
 
@@ -286,6 +276,18 @@ export class ChromeAdapter {
     return newSession;
   }
 
+  /**
+   * Extracts fields in common between {@link GenerateContentRequest} and
+   * {@link LanguageModelPromptOptions}.
+   */
+  private static extractLanguageModelPromptOptions(
+    request: GenerateContentRequest
+  ): LanguageModelPromptOptions {
+    return {
+      responseConstraint: request.generationConfig?.responseSchema
+    };
+  }
+
   private addImageTypeAsExpectedInput(): void {
     // Defaults to support image inputs for convenience.
     this.onDeviceParams.expectedInputs ??= [{ type: 'image' }];