Refactor to emulate Vertex response

Author: erikeldridge

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

@@ -252,7 +252,15 @@ describe('ChromeAdapter', () => {
           { role: 'assistant', content: text[1] }
         ]
       });
-      expect(response.text()).to.equal(text[2]);
+      expect(await response.json()).to.deep.equal({
+        candidates: [
+          {
+            content: {
+              parts: [{ text: text[2] }]
+            }
+          }
+        ]
+      });
     });
     it('Extracts system prompt', async () => {
       const aiProvider = {
@@ -279,7 +287,15 @@ describe('ChromeAdapter', () => {
         initialPrompts: [],
         systemPrompt: onDeviceParams.systemPrompt
       });
-      expect(response.text()).to.equal(text);
+      expect(await response.json()).to.deep.equal({
+        candidates: [
+          {
+            content: {
+              parts: [{ text }]
+            }
+          }
+        ]
+      });
     });
   });
 });

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

@@ -18,7 +18,6 @@
 import { isChrome } from '@firebase/util';
 import {
   Content,
-  EnhancedGenerateContentResponse,
   GenerateContentRequest,
   InferenceMode,
   Role
@@ -75,19 +74,31 @@ export class ChromeAdapter {
   }
   async generateContentOnDevice(
     request: GenerateContentRequest
-  ): Promise<EnhancedGenerateContentResponse> {
+  ): Promise<Response> {
     const createOptions = this.onDeviceParams || {};
     createOptions.initialPrompts ??= [];
-    const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(request.contents);
+    const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(
+      request.contents
+    );
     // Assumes validation asserted there is at least one initial prompt.
     const prompt = extractedInitialPrompts.pop()!;
     createOptions.initialPrompts.push(...extractedInitialPrompts);
     const session = await this.session(createOptions);
     const result = await session.prompt(prompt.content);
+    return ChromeAdapter.toResponse(result);
+  }
+  private static toResponse(text: string): Response {
     return {
-      text: () => result,
-      functionCalls: () => undefined
-    };
+      json: async () => ({
+        candidates: [
+          {
+            content: {
+              parts: [{ text }]
+            }
+          }
+        ]
+      })
+    } as Response;
   }
   private static isOnDeviceRequest(request: GenerateContentRequest): boolean {
     // Returns false if the prompt is empty.

packages/vertexai/src/methods/generate-content.test.ts

@@ -291,24 +291,23 @@ describe('generateContent()', () => {
     expect(mockFetch).to.be.called;
   });
   it('on-device', async () => {
-    const expectedText = 'hi';
     const chromeAdapter = new ChromeAdapter();
     const mockIsAvailable = stub(chromeAdapter, 'isAvailable').resolves(true);
-    const mockGenerateContent = stub(
+    const mockResponse = getMockResponse(
+      'unary-success-basic-reply-short.json'
+    );
+    const makeRequestStub = stub(
       chromeAdapter,
       'generateContentOnDevice'
-    ).resolves({
-      text: () => expectedText,
-      functionCalls: () => undefined
-    });
+    ).resolves(mockResponse as Response);
     const result = await generateContent(
       fakeApiSettings,
       'model',
       fakeRequestParams,
       chromeAdapter
     );
-    expect(result.response.text()).to.equal(expectedText);
+    expect(result.response.text()).to.include('Mountain View, California');
     expect(mockIsAvailable).to.be.called;
-    expect(mockGenerateContent).to.be.calledWith(fakeRequestParams);
+    expect(makeRequestStub).to.be.calledWith(fakeRequestParams);
   });
 });

packages/vertexai/src/methods/generate-content.ts

@@ -16,7 +16,6 @@
  */
 
 import {
-  EnhancedGenerateContentResponse,
   GenerateContentRequest,
   GenerateContentResponse,
   GenerateContentResult,
@@ -51,18 +50,15 @@ async function generateContentOnCloud(
   model: string,
   params: GenerateContentRequest,
   requestOptions?: RequestOptions
-): Promise<EnhancedGenerateContentResponse> {
-  const response = await makeRequest(
+): Promise<Response> {
+  return makeRequest(
     model,
     Task.GENERATE_CONTENT,
     apiSettings,
     /* stream */ false,
     JSON.stringify(params),
     requestOptions
   );
-  const responseJson: GenerateContentResponse = await response.json();
-  const enhancedResponse = createEnhancedContentResponse(responseJson);
-  return enhancedResponse;
 }
 
 export async function generateContent(
@@ -72,17 +68,19 @@ export async function generateContent(
   chromeAdapter: ChromeAdapter,
   requestOptions?: RequestOptions
 ): Promise<GenerateContentResult> {
-  let enhancedResponse;
+  let response;
   if (await chromeAdapter.isAvailable(params)) {
-    enhancedResponse = await chromeAdapter.generateContentOnDevice(params);
+    response = await chromeAdapter.generateContentOnDevice(params);
   } else {
-    enhancedResponse = await generateContentOnCloud(
+    response = await generateContentOnCloud(
       apiSettings,
       model,
       params,
       requestOptions
     );
   }
+  const responseJson: GenerateContentResponse = await response.json();
+  const enhancedResponse = createEnhancedContentResponse(responseJson);
   return {
     response: enhancedResponse
   };