response-helpers.ts

Author: russellwheatley

packages/ai/lib/requests/response-helpers.ts

@@ -21,9 +21,10 @@ import {
   FunctionCall,
   GenerateContentCandidate,
   GenerateContentResponse,
-  VertexAIErrorCode,
+  AIErrorCode,
+  InlineDataPart,
 } from '../types';
-import { VertexAIError } from '../errors';
+import { AIError } from '../errors';
 import { logger } from '../logger';
 
 /**
@@ -62,8 +63,8 @@ export function addHelpers(response: GenerateContentResponse): EnhancedGenerateC
         );
       }
       if (hadBadFinishReason(response.candidates[0]!)) {
-        throw new VertexAIError(
-          VertexAIErrorCode.RESPONSE_ERROR,
+        throw new AIError(
+          AIErrorCode.RESPONSE_ERROR,
           `Response error: ${formatBlockErrorMessage(
             response,
           )}. Response body stored in error.response`,
@@ -74,8 +75,8 @@ export function addHelpers(response: GenerateContentResponse): EnhancedGenerateC
       }
       return getText(response);
     } else if (response.promptFeedback) {
-      throw new VertexAIError(
-        VertexAIErrorCode.RESPONSE_ERROR,
+      throw new AIError(
+        AIErrorCode.RESPONSE_ERROR,
         `Text not available. ${formatBlockErrorMessage(response)}`,
         {
           response,
@@ -84,6 +85,40 @@ export function addHelpers(response: GenerateContentResponse): EnhancedGenerateC
     }
     return '';
   };
+  (response as EnhancedGenerateContentResponse).inlineDataParts = ():
+    | InlineDataPart[]
+    | undefined => {
+    if (response.candidates && response.candidates.length > 0) {
+      if (response.candidates.length > 1) {
+        logger.warn(
+          `This response had ${response.candidates.length} ` +
+            `candidates. Returning data from the first candidate only. ` +
+            `Access response.candidates directly to use the other candidates.`,
+        );
+      }
+      if (hadBadFinishReason(response.candidates[0]!)) {
+        throw new AIError(
+          AIErrorCode.RESPONSE_ERROR,
+          `Response error: ${formatBlockErrorMessage(
+            response,
+          )}. Response body stored in error.response`,
+          {
+            response,
+          },
+        );
+      }
+      return getInlineDataParts(response);
+    } else if (response.promptFeedback) {
+      throw new AIError(
+        AIErrorCode.RESPONSE_ERROR,
+        `Data not available. ${formatBlockErrorMessage(response)}`,
+        {
+          response,
+        },
+      );
+    }
+    return undefined;
+  };
   (response as EnhancedGenerateContentResponse).functionCalls = () => {
     if (response.candidates && response.candidates.length > 0) {
       if (response.candidates.length > 1) {
@@ -94,8 +129,8 @@ export function addHelpers(response: GenerateContentResponse): EnhancedGenerateC
         );
       }
       if (hadBadFinishReason(response.candidates[0]!)) {
-        throw new VertexAIError(
-          VertexAIErrorCode.RESPONSE_ERROR,
+        throw new AIError(
+          AIErrorCode.RESPONSE_ERROR,
           `Response error: ${formatBlockErrorMessage(
             response,
           )}. Response body stored in error.response`,
@@ -106,8 +141,8 @@ export function addHelpers(response: GenerateContentResponse): EnhancedGenerateC
       }
       return getFunctionCalls(response);
     } else if (response.promptFeedback) {
-      throw new VertexAIError(
-        VertexAIErrorCode.RESPONSE_ERROR,
+      throw new AIError(
+        AIErrorCode.RESPONSE_ERROR,
         `Function call not available. ${formatBlockErrorMessage(response)}`,
         {
           response,
@@ -125,7 +160,7 @@ export function addHelpers(response: GenerateContentResponse): EnhancedGenerateC
 export function getText(response: GenerateContentResponse): string {
   const textStrings = [];
   if (response.candidates?.[0]?.content?.parts) {
-    for (const part of response.candidates?.[0].content?.parts) {
+    for (const part of response.candidates?.[0]?.content?.parts) {
       if (part.text) {
         textStrings.push(part.text);
       }
@@ -139,7 +174,7 @@ export function getText(response: GenerateContentResponse): string {
 }
 
 /**
- * Returns <code>{@link FunctionCall}</code>s associated with first candidate.
+ * Returns {@link FunctionCall}s associated with first candidate.
  */
 export function getFunctionCalls(response: GenerateContentResponse): FunctionCall[] | undefined {
   const functionCalls: FunctionCall[] = [];
@@ -157,6 +192,31 @@ export function getFunctionCalls(response: GenerateContentResponse): FunctionCal
   }
 }
 
+/**
+ * Returns {@link InlineDataPart}s in the first candidate if present.
+ *
+ * @internal
+ */
+export function getInlineDataParts(
+  response: GenerateContentResponse,
+): InlineDataPart[] | undefined {
+  const data: InlineDataPart[] = [];
+
+  if (response.candidates?.[0]?.content?.parts) {
+    for (const part of response.candidates?.[0]?.content?.parts) {
+      if (part.inlineData) {
+        data.push(part);
+      }
+    }
+  }
+
+  if (data.length > 0) {
+    return data;
+  } else {
+    return undefined;
+  }
+}
+
 const badFinishReasons = [FinishReason.RECITATION, FinishReason.SAFETY];
 
 function hadBadFinishReason(candidate: GenerateContentCandidate): boolean {

packages/ai/lib/types/responses.ts

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-import { Content, FunctionCall } from './content';
+import { Content, FunctionCall, InlineDataPart } from './content';
 import { BlockReason, FinishReason, HarmCategory, HarmProbability, HarmSeverity } from './enums';
 
 /**
@@ -51,6 +51,15 @@ export interface EnhancedGenerateContentResponse extends GenerateContentResponse
    * Throws if the prompt or candidate was blocked.
    */
   text: () => string;
+  /**
+   * Aggregates and returns all {@link InlineDataPart}s from the {@link GenerateContentResponse}'s
+   * first candidate.
+   *
+   * @returns An array of {@link InlineDataPart}s containing data from the response, if available.
+   *
+   * @throws If the prompt or candidate was blocked.
+   */
+  inlineDataParts: () => InlineDataPart[] | undefined;
   functionCalls: () => FunctionCall[] | undefined;
 }