fix(ai): Add 'includeSafetyAttributes' to Predict request payloads

.changeset/good-cooks-deliver.md

@@ -0,0 +1,5 @@
+---
+'@firebase/ai': patch
+---
+
+Add 'includeSafetyAttributes' field to Predict request payloads.

packages/ai/src/requests/request-helpers.test.ts

@@ -214,6 +214,7 @@ describe('request formatting methods', () => {
       expect(body.instances[0].prompt).to.equal(prompt);
       expect(body.parameters.sampleCount).to.equal(1);
       expect(body.parameters.includeRaiReason).to.be.true;
+      expect(body.parameters.includeSafetyAttributes).to.be.true;
 
       // Parameters without default values should be undefined
       expect(body.parameters.storageUri).to.be.undefined;
@@ -258,6 +259,7 @@ describe('request formatting methods', () => {
       personGeneration: safetySettings.personFilterLevel,
       aspectRatio,
       includeRaiReason: true,
+      includeSafetyAttributes: true,
       storageUri: undefined
     });
   });

packages/ai/src/requests/request-helpers.ts

@@ -156,7 +156,8 @@ export function createPredictRequestBody(
       addWatermark,
       safetyFilterLevel,
       personGeneration: personFilterLevel,
-      includeRaiReason: true
+      includeRaiReason: true,
+      includeSafetyAttributes: true
     }
   };
   return body;

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

@@ -296,6 +296,8 @@ export async function handlePredictResponse<
         mimeType: prediction.mimeType,
         gcsURI: prediction.gcsUri
       } as T);
+    } else if (prediction.safetyAttributes) {
+      // Ignore safetyAttributes "prediction" to avoid throwing an error below.
     } else {
       throw new AIError(
         AIErrorCode.RESPONSE_ERROR,

packages/ai/src/types/imagen/internal.ts

@@ -61,6 +61,14 @@ export interface ImagenResponseInternal {
      * The reason why the image was filtered.
      */
     raiFilteredReason?: string;
+    /**
+     * The safety attributes.
+     *
+     * This type is currently unused in the SDK. It is sent back because our requests set
+     * `includeSafetyAttributes`. This property is currently only used to avoid throwing an error
+     * when encountering this unsupported prediction type.
+     */
+    safetyAttributes?: unknown;
   }>;
 }
 
@@ -84,6 +92,7 @@ export interface ImagenResponseInternal {
  *     "personGeneration": "allow_all",
  *     "sampleCount": 2,
  *     "includeRaiReason": true,
+ *     "includeSafetyAttributes": true,
  *     "aspectRatio": "9:16"
  *   }
  * }
@@ -111,6 +120,7 @@ export interface PredictRequestBody {
     safetyFilterLevel?: string;
     personGeneration?: string; // Maps to personFilterLevel
     includeRaiReason: boolean;
+    includeSafetyAttributes: boolean;
   };
 }