diff --git a/packages/vertexai/src/methods/generate-content.ts b/packages/vertexai/src/methods/generate-content.ts index 2dee91f12e8..0944b38016a 100644 --- a/packages/vertexai/src/methods/generate-content.ts +++ b/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 }; diff --git a/packages/vertexai/src/requests/response-helpers.ts b/packages/vertexai/src/requests/response-helpers.ts index efff09591a6..3521dc1bc5f 100644 --- a/packages/vertexai/src/requests/response-helpers.ts +++ b/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). diff --git a/packages/vertexai/src/requests/stream-reader.ts b/packages/vertexai/src/requests/stream-reader.ts index c4163d26b60..da8eca71880 100644 --- a/packages/vertexai/src/requests/stream-reader.ts +++ b/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; } }