Merge dcf35e0 into dc211c6

Author: dlarocque

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

@@ -23,7 +23,7 @@ import {
   RequestOptions
 } from '../types';
 import { Task, makeRequest } from '../requests/request';
-import { addHelpers } from '../requests/response-helpers';
+import { createEnhancedContentResponse } from '../requests/response-helpers';
 import { processStream } from '../requests/stream-reader';
 import { ApiSettings } from '../types/internal';
 
@@ -59,7 +59,7 @@ export async function generateContent(
     requestOptions
   );
   const responseJson: GenerateContentResponse = await response.json();
-  const enhancedResponse = addHelpers(responseJson);
+  const enhancedResponse = createEnhancedContentResponse(responseJson);
   return {
     response: enhancedResponse
   };

packages/vertexai/src/requests/response-helpers.ts

@@ -25,6 +25,26 @@ import {
 } from '../types';
 import { VertexAIError } from '../errors';
 
+/**
+ * Creates an EnhancedGenerateContentResponse object that has helper functions and
+ * other modifications that improve usability.
+ */
+export function createEnhancedContentResponse(
+  response: GenerateContentResponse
+): EnhancedGenerateContentResponse {
+  /**
+   * The Vertex AI backend omits default values.
+   * This causes the `index` property to be omitted from the first candidate in the
+   * response, since it has index 0, and 0 is a default value.
+   */
+  if (response.candidates && !response.candidates[0].hasOwnProperty('index')) {
+    response.candidates[0].index = 0;
+  }
+
+  const responseWithHelpers = addHelpers(response);
+  return responseWithHelpers;
+}
+
 /**
  * Adds convenience helper methods to a response object, including stream
  * chunks (as long as each chunk is a complete GenerateContentResponse JSON).

packages/vertexai/src/requests/stream-reader.ts

@@ -24,7 +24,7 @@ import {
   VertexAIErrorCode
 } from '../types';
 import { VertexAIError } from '../errors';
-import { addHelpers } from './response-helpers';
+import { createEnhancedContentResponse } from './response-helpers';
 
 const responseLineRE = /^data\: (.*)(?:\n\n|\r\r|\r\n\r\n)/;
 
@@ -57,7 +57,10 @@ async function getResponsePromise(
   while (true) {
     const { done, value } = await reader.read();
     if (done) {
-      return addHelpers(aggregateResponses(allResponses));
+      const enhancedResponse = createEnhancedContentResponse(
+        aggregateResponses(allResponses)
+      );
+      return enhancedResponse;
     }
     allResponses.push(value);
   }
@@ -72,7 +75,9 @@ async function* generateResponseSequence(
     if (done) {
       break;
     }
-    yield addHelpers(value);
+
+    const enhancedResponse = createEnhancedContentResponse(value);
+    yield enhancedResponse;
   }
 }