diff --git a/.changeset/good-cooks-deliver.md b/.changeset/good-cooks-deliver.md new file mode 100644 index 00000000000..739a42f8574 --- /dev/null +++ b/.changeset/good-cooks-deliver.md @@ -0,0 +1,5 @@ +--- +'@firebase/ai': patch +--- + +Add 'includeSafetyAttributes' field to Predict request payloads. diff --git a/packages/ai/src/requests/request-helpers.test.ts b/packages/ai/src/requests/request-helpers.test.ts index d8337850925..993b0f1d3ae 100644 --- a/packages/ai/src/requests/request-helpers.test.ts +++ b/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 }); }); diff --git a/packages/ai/src/requests/request-helpers.ts b/packages/ai/src/requests/request-helpers.ts index c4cc1a20acc..ee80142481b 100644 --- a/packages/ai/src/requests/request-helpers.ts +++ b/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; diff --git a/packages/ai/src/requests/response-helpers.ts b/packages/ai/src/requests/response-helpers.ts index 2505b5c9276..32c0cdde6ce 100644 --- a/packages/ai/src/requests/response-helpers.ts +++ b/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, diff --git a/packages/ai/src/types/imagen/internal.ts b/packages/ai/src/types/imagen/internal.ts index 02a8a55e01c..1a34eb18f56 100644 --- a/packages/ai/src/types/imagen/internal.ts +++ b/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; }; }