From 7ce745f5daf1d80a532db93ffa787544171c11c0 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Tue, 7 Jan 2025 21:45:36 -0800
Subject: [PATCH 01/21] first pass of imagen

---
 .../firebase_vertexai/example/.gitignore      |   2 +
 .../firebase_vertexai/example/lib/main.dart   |  30 +++-
 .../example/macos/Runner/AppDelegate.swift    |   6 +-
 .../lib/firebase_vertexai.dart                |  10 ++
 .../firebase_vertexai/lib/src/base_model.dart | 111 ++++++++++++
 .../lib/src/firebase_vertexai.dart            |  17 ++
 .../firebase_vertexai/lib/src/imagen_api.dart | 168 ++++++++++++++++++
 .../lib/src/imagen_content.dart               | 112 ++++++++++++
 .../lib/src/imagen_model.dart                 | 108 +++++++++++
 .../firebase_vertexai/lib/src/model.dart      | 117 +++---------
 10 files changed, 587 insertions(+), 94 deletions(-)
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/.gitignore b/packages/firebase_vertexai/firebase_vertexai/example/.gitignore
index db070ebdbf4f..0498b592dfa0 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/.gitignore
+++ b/packages/firebase_vertexai/firebase_vertexai/example/.gitignore
@@ -5,9 +5,11 @@
 *.swp
 .DS_Store
 .atom/
+.build/
 .buildlog/
 .history
 .svn/
+.swiftpm/
 migrate_working_dir/
 
 # IntelliJ related
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
index 669252a77d82..f0ee767f8c2f 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -266,7 +266,8 @@ class _ChatWidgetState extends State<ChatWidget> {
                   tooltip: 'image prompt',
                   onPressed: !_loading
                       ? () async {
-                          await _sendImagePrompt(_textController.text);
+                          //await _sendImagePrompt(_textController.text);
+                          await _testImagen();
                         }
                       : null,
                   icon: Icon(
@@ -602,6 +603,33 @@ class _ChatWidgetState extends State<ChatWidget> {
     });
   }
 
+  Future<void> _testImagen() async {
+    var model = FirebaseVertexAI.instance.imageModel(
+      modelName: 'imagen-3.0-generate-001',
+      safetySettings: ImagenSafetySettings(
+        ImagenSafetyFilterLevel.blockLowAndAbove,
+        ImagenPersonFilterLevel.allowAdult,
+      ),
+    );
+
+    var prompt = 'An astronaut riding a horse.';
+    var generationConfig =
+        ImagenGenerationConfig('frog', 1, ImagenAspectRatio.square1x1);
+
+    var response = await model.generateImages(
+      prompt,
+      generationConfig: generationConfig,
+    );
+
+    if (response.images.isNotEmpty) {
+      var image = response.images[0];
+      // Process the image
+    } else {
+      // Handle the case where no images were generated
+      print('Error: No images were generated.');
+    }
+  }
+
   void _showError(String message) {
     showDialog<void>(
       context: context,
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift b/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift
index d53ef6437726..b3c176141221 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift
+++ b/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift
@@ -1,9 +1,13 @@
 import Cocoa
 import FlutterMacOS
 
-@NSApplicationMain
+@main
 class AppDelegate: FlutterAppDelegate {
   override func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool {
     return true
   }
+
+  override func applicationSupportsSecureRestorableState(_ app: NSApplication) -> Bool {
+    return true
+  }
 }
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
index a170d8deb969..017237731e86 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
@@ -28,6 +28,7 @@ export 'src/api.dart'
         PromptFeedback,
         SafetyRating,
         SafetySetting,
+        // TODO(cynthiajiang) remove in next breaking change.
         TaskType,
         UsageMetadata;
 export 'src/chat.dart' show ChatSession, StartChatExtension;
@@ -55,5 +56,14 @@ export 'src/function_calling.dart'
         FunctionDeclaration,
         Tool,
         ToolConfig;
+export 'src/imagen_api.dart'
+    show
+        ImagenModelConfig,
+        ImagenSafetySettings,
+        ImagenFormat,
+        ImagenSafetyFilterLevel,
+        ImagenPersonFilterLevel,
+        ImagenGenerationConfig,
+        ImagenAspectRatio;
 export 'src/model.dart' show GenerativeModel;
 export 'src/schema.dart' show Schema, SchemaType;
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
new file mode 100644
index 000000000000..6c0a8bf51342
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
@@ -0,0 +1,111 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import 'dart:async';
+
+import 'package:firebase_app_check/firebase_app_check.dart';
+import 'package:firebase_auth/firebase_auth.dart';
+import 'package:firebase_core/firebase_core.dart';
+
+import 'client.dart';
+import 'vertex_version.dart';
+
+/// [Task] enum class for [GenerativeModel] to make request.
+enum Task {
+  /// Request type to generate content.
+  generateContent,
+
+  /// Request type to stream content.
+  streamGenerateContent,
+
+  /// Request type to count token.
+  countTokens,
+
+  /// Imagen 3 task
+  predict,
+}
+
+abstract class BaseModel {
+  BaseModel({
+    required String model,
+    required String location,
+    required FirebaseApp app,
+    required ApiClient client,
+  })  : _model = normalizeModelName(model),
+        _app = app,
+        _location = location,
+        _projectUri = _vertexUri(app, location),
+        _client = client;
+
+  static const _baseUrl = 'firebasevertexai.googleapis.com';
+  static const _apiVersion = 'v1beta';
+
+  final ({String prefix, String name}) _model;
+  final FirebaseApp _app;
+  final String _location;
+  final Uri _projectUri;
+  final ApiClient _client;
+
+  ({String prefix, String name}) get model => _model;
+  ApiClient get client => _client;
+
+  /// Returns the model code for a user friendly model name.
+  ///
+  /// If the model name is already a model code (contains a `/`), use the parts
+  /// directly. Otherwise, return a `models/` model code.
+  static ({String prefix, String name}) normalizeModelName(String modelName) {
+    if (!modelName.contains('/')) return (prefix: 'models', name: modelName);
+    final parts = modelName.split('/');
+    return (prefix: parts.first, name: parts.skip(1).join('/'));
+  }
+
+  static Uri _vertexUri(FirebaseApp app, String location) {
+    var projectId = app.options.projectId;
+    return Uri.https(
+      _baseUrl,
+      '/$_apiVersion/projects/$projectId/locations/$location/publishers/google',
+    );
+  }
+
+  static FutureOr<Map<String, String>> Function() firebaseTokens(
+      FirebaseAppCheck? appCheck, FirebaseAuth? auth) {
+    return () async {
+      Map<String, String> headers = {};
+      // Override the client name in Google AI SDK
+      headers['x-goog-api-client'] =
+          'gl-dart/$packageVersion fire/$packageVersion';
+      if (appCheck != null) {
+        final appCheckToken = await appCheck.getToken();
+        if (appCheckToken != null) {
+          headers['X-Firebase-AppCheck'] = appCheckToken;
+        }
+      }
+      if (auth != null) {
+        final idToken = await auth.currentUser?.getIdToken();
+        if (idToken != null) {
+          headers['Authorization'] = 'Firebase $idToken';
+        }
+      }
+      return headers;
+    };
+  }
+
+  Uri taskUri(Task task) => _projectUri.replace(
+      pathSegments: _projectUri.pathSegments
+          .followedBy([_model.prefix, '${_model.name}:${task.name}']));
+
+  /// Make a unary request for [task] with JSON encodable [params].
+  Future<T> makeRequest<T>(Task task, Map<String, Object?> params,
+          T Function(Map<String, Object?>) parse) =>
+      _client.makeRequest(taskUri(task), params).then(parse);
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
index df9c26f9eaaf..b0cd87e854f5 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
@@ -21,6 +21,9 @@ import 'package:firebase_core_platform_interface/firebase_core_platform_interfac
 import 'api.dart';
 import 'content.dart';
 import 'function_calling.dart';
+import 'imagen_model.dart';
+import 'imagen_api.dart';
+import 'imagen_content.dart';
 import 'model.dart';
 
 const _defaultLocation = 'us-central1';
@@ -110,4 +113,18 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
       systemInstruction: systemInstruction,
     );
   }
+
+  ImagenModel imageModel(
+      {required String modelName,
+      ImagenModelConfig? modelConfig,
+      ImagenSafetySettings? safetySettings}) {
+    return createImagenModel(
+        app: app,
+        location: location,
+        modelName: modelName,
+        modelConfig: modelConfig,
+        safetySettings: safetySettings,
+        appCheck: appCheck,
+        auth: auth);
+  }
 }
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
new file mode 100644
index 000000000000..d56101a1331f
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
@@ -0,0 +1,168 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+///
+enum ImagenSafetyFilterLevel {
+  ///
+  blockLowAndAbove('block_low_and_above'),
+
+  ///
+  blockMediumAndAbove('block_medium_and_above'),
+
+  ///
+  blockOnlyHigh('block_only_high'),
+
+  ///
+  blockNone('block_none');
+
+  const ImagenSafetyFilterLevel(this._jsonString);
+
+  final String _jsonString;
+
+  /// Convert to json format
+  String toJson() => _jsonString;
+
+  // ignore: unused_element
+  static ImagenSafetyFilterLevel _parseValue(Object jsonObject) {
+    return switch (jsonObject) {
+      'block_low_and_above' => ImagenSafetyFilterLevel.blockLowAndAbove,
+      'block_medium_and_above' => ImagenSafetyFilterLevel.blockMediumAndAbove,
+      'block_only_high' => ImagenSafetyFilterLevel.blockOnlyHigh,
+      'block_none' => ImagenSafetyFilterLevel.blockNone,
+      _ => throw FormatException(
+          'Unhandled ImagenSafetyFilterLevel format', jsonObject),
+    };
+  }
+
+  @override
+  String toString() => name;
+}
+
+///
+enum ImagenPersonFilterLevel {
+  ///
+  blockAll('dont_allow'),
+
+  ///
+  allowAdult('allow_adult'),
+
+  ///
+  allowAll('allow_all');
+
+  const ImagenPersonFilterLevel(this._jsonString);
+
+  final String _jsonString;
+
+  /// Convert to json format
+  String toJson() => _jsonString;
+
+  // ignore: unused_element
+  static ImagenPersonFilterLevel _parseValue(Object jsonObject) {
+    return switch (jsonObject) {
+      'dont_allow' => ImagenPersonFilterLevel.blockAll,
+      'allow_adult' => ImagenPersonFilterLevel.allowAdult,
+      'allow_all' => ImagenPersonFilterLevel.allowAll,
+      _ => throw FormatException(
+          'Unhandled ImagenPersonFilterLevel format', jsonObject),
+    };
+  }
+
+  @override
+  String toString() => name;
+}
+
+///
+final class ImagenSafetySettings {
+  /// Constructor
+  ImagenSafetySettings(this.safetyFilterLevel, this.personFilterLevel);
+
+  ///
+  final ImagenSafetyFilterLevel? safetyFilterLevel;
+
+  ///
+  final ImagenPersonFilterLevel? personFilterLevel;
+
+  /// Convert to json format.
+  Object toJson() => {
+        if (safetyFilterLevel != null)
+          'safetySetting': safetyFilterLevel!.toJson(),
+        if (personFilterLevel != null)
+          'personGeneration': personFilterLevel!.toJson(),
+      };
+}
+
+///
+enum ImagenAspectRatio {
+  ///
+  square1x1('1:1'),
+
+  ///
+  portrait9x16('9:16'),
+
+  ///
+  landscape16x9('16:9'),
+
+  ///
+  portrait3x4('3:4'),
+
+  ///
+  landscape4x3('4:3');
+
+  const ImagenAspectRatio(this._jsonString);
+
+  final String _jsonString;
+
+  /// Convert to json format
+  String toJson() => _jsonString;
+
+  // ignore: unused_element
+  static ImagenAspectRatio _parseValue(Object jsonObject) {
+    return switch (jsonObject) {
+      '1:1' => ImagenAspectRatio.square1x1,
+      '9:16' => ImagenAspectRatio.portrait9x16,
+      '16:9' => ImagenAspectRatio.landscape16x9,
+      '3:4' => ImagenAspectRatio.portrait3x4,
+      '4:3' => ImagenAspectRatio.landscape4x3,
+      _ =>
+        throw FormatException('Unhandled ImagenAspectRatio format', jsonObject),
+    };
+  }
+
+  @override
+  String toString() => name;
+}
+
+final class ImagenGenerationConfig {
+  ImagenGenerationConfig(
+      this.negativePrompt, this.numberOfImages, this.aspectRatio);
+  final String? negativePrompt;
+  final int? numberOfImages;
+  final ImagenAspectRatio? aspectRatio;
+}
+
+final class ImagenFormat {
+  ImagenFormat(this.mimeType, this.compressionQuality);
+
+  ImagenFormat.png() : this("image/png", null);
+  ImagenFormat.jpeg({int? compressionQuality})
+      : this("image/jpeg", compressionQuality);
+  final String mimeType;
+  final int? compressionQuality;
+}
+
+final class ImagenModelConfig {
+  ImagenModelConfig(this.imagenFormat, this.addWatermark);
+  final ImagenFormat imagenFormat;
+  final bool? addWatermark;
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
new file mode 100644
index 000000000000..4289e6f1dd77
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
@@ -0,0 +1,112 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import 'dart:convert';
+import 'dart:typed_data';
+import 'error.dart';
+
+///
+sealed class ImagenImage {
+  final String mimeType;
+
+  /// Convert the [ImagenImage] content to json format.
+  Object toJson();
+
+  ImagenImage({required this.mimeType});
+}
+
+final class ImagenInlineImage implements ImagenImage {
+  /// Data contents in bytes.
+  final Uint8List bytesBase64Encoded;
+
+  @override
+  final String mimeType;
+
+  ImagenInlineImage({
+    required this.bytesBase64Encoded,
+    required this.mimeType,
+  });
+
+  @override
+  Object toJson() => {
+        'mimeType': mimeType,
+        'bytesBase64Encoded': bytesBase64Encoded,
+      };
+}
+
+final class ImagenGCSImage implements ImagenImage {
+  @override
+  final String mimeType;
+
+  final String gcsUri;
+
+  ImagenGCSImage({
+    required this.gcsUri,
+    required this.mimeType,
+  });
+
+  @override
+  Object toJson() => {
+        'mimeType': mimeType,
+        'gcsUri': gcsUri,
+      };
+}
+
+final class ImagenGenerationResponse<T extends ImagenImage> {
+  ImagenGenerationResponse({
+    required this.images,
+    this.filteredReason,
+  });
+
+  final List<T> images;
+  final String? filteredReason;
+
+  factory ImagenGenerationResponse.fromJson(Map<String, dynamic> json) {
+    final filteredReason = json['filteredReason'] as String?;
+    final imagesJson = json['images'] as List<dynamic>;
+
+    if (T == ImagenInlineImage) {
+      final images = imagesJson.map((imageJson) {
+        final mimeType = imageJson['mimeType'] as String;
+        final bytes = imageJson['bytesBase64Encoded'] as String;
+        final decodedBytes = base64Decode(bytes);
+        return ImagenInlineImage(
+          mimeType: mimeType,
+          bytesBase64Encoded: Uint8List.fromList(decodedBytes),
+        ) as T;
+      }).toList();
+      return ImagenGenerationResponse<T>(
+          images: images, filteredReason: filteredReason);
+    } else if (T == ImagenGCSImage) {
+      final images = imagesJson.map((imageJson) {
+        final mimeType = imageJson['mimeType'] as String;
+        final gcsUri = imageJson['gcsUri'] as String;
+        return ImagenGCSImage(
+          mimeType: mimeType,
+          gcsUri: gcsUri,
+        ) as T;
+      }).toList();
+      return ImagenGenerationResponse<T>(
+          images: images, filteredReason: filteredReason);
+    } else {
+      throw ArgumentError('Unsupported ImagenImage type: $T');
+    }
+  }
+}
+
+/// Parse the json to [ImagenGenerationResponse]
+ImagenGenerationResponse parseImagenGenerationResponse(Object jsonObject) {
+  if (jsonObject case {'error': final Object error}) throw parseError(error);
+  Map<String, dynamic> json = jsonObject as Map<String, dynamic>;
+  return ImagenGenerationResponse.fromJson(json);
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
new file mode 100644
index 000000000000..c25201436691
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -0,0 +1,108 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import 'imagen_api.dart';
+import 'imagen_content.dart';
+import 'base_model.dart';
+import 'client.dart';
+
+import 'package:firebase_app_check/firebase_app_check.dart';
+import 'package:firebase_auth/firebase_auth.dart';
+import 'package:firebase_core/firebase_core.dart';
+
+///
+final class ImagenModel extends BaseModel {
+  ImagenModel._(
+      {required FirebaseApp app,
+      required String modelName,
+      required String location,
+      FirebaseAppCheck? appCheck,
+      FirebaseAuth? auth,
+      ImagenModelConfig? modelConfig,
+      ImagenSafetySettings? safetySettings})
+      : _modelConfig = modelConfig,
+        _safetySettings = safetySettings,
+        super(
+            model: modelName,
+            app: app,
+            location: location,
+            client: HttpApiClient(
+                apiKey: app.options.apiKey,
+                requestHeaders: BaseModel.firebaseTokens(appCheck, auth)));
+
+  final ImagenModelConfig? _modelConfig;
+  final ImagenSafetySettings? _safetySettings;
+
+  Map<String, Object?> _generateImagenRequest(
+    String prompt, {
+    ImagenGenerationConfig? generationConfig,
+    String? gcsUri,
+  }) {
+    return {
+      'instances': [
+        {'prompt': prompt}
+      ],
+      'parameters': {
+        if (gcsUri != null) 'storageUri': gcsUri,
+        if (generationConfig != null && generationConfig.numberOfImages != null)
+          'sampleCount': generationConfig.numberOfImages,
+        if (generationConfig != null && generationConfig.aspectRatio != null)
+          'aspectRatio': generationConfig.aspectRatio,
+        if (generationConfig != null && generationConfig.negativePrompt != null)
+          'negativePrompt': generationConfig.negativePrompt,
+        if (_safetySettings != null &&
+            _safetySettings.personFilterLevel != null)
+          'personGeneration': _safetySettings.personFilterLevel,
+        if (_safetySettings != null &&
+            _safetySettings.safetyFilterLevel != null)
+          'safetySetting': _safetySettings.safetyFilterLevel,
+        if (_modelConfig != null && _modelConfig.addWatermark != null)
+          'addWatermark': _modelConfig.addWatermark,
+      }
+    };
+  }
+
+  Future<ImagenGenerationResponse> generateImages(
+    String prompt, {
+    ImagenGenerationConfig? generationConfig,
+    String? gcsUri,
+  }) =>
+      makeRequest(
+          Task.predict,
+          _generateImagenRequest(
+            prompt,
+            generationConfig: generationConfig,
+            gcsUri: gcsUri,
+          ),
+          parseImagenGenerationResponse);
+}
+
+/// Returns a [ImagenModel] using it's private constructor.
+ImagenModel createImagenModel({
+  required FirebaseApp app,
+  required String location,
+  required String modelName,
+  FirebaseAppCheck? appCheck,
+  FirebaseAuth? auth,
+  ImagenModelConfig? modelConfig,
+  ImagenSafetySettings? safetySettings,
+}) =>
+    ImagenModel._(
+      modelName: modelName,
+      app: app,
+      appCheck: appCheck,
+      auth: auth,
+      location: location,
+      safetySettings: safetySettings,
+      modelConfig: modelConfig,
+    );
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart
index 605adf9b4fa7..91a050cc2922 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart
@@ -27,33 +27,13 @@ import 'client.dart';
 import 'content.dart';
 import 'function_calling.dart';
 import 'vertex_version.dart';
-
-const _baseUrl = 'firebasevertexai.googleapis.com';
-const _apiVersion = 'v1beta';
-
-/// [Task] enum class for [GenerativeModel] to make request.
-enum Task {
-  /// Request type to generate content.
-  generateContent,
-
-  /// Request type to stream content.
-  streamGenerateContent,
-
-  /// Request type to count token.
-  countTokens,
-
-  /// Request type to embed content.
-  embedContent,
-
-  /// Request type to batch embed content.
-  batchEmbedContents;
-}
+import 'base_model.dart';
 
 /// A multimodel generative model (like Gemini).
 ///
 /// Allows generating content, creating embeddings, and counting the number of
 /// tokens in a piece of content.
-final class GenerativeModel {
+final class GenerativeModel extends BaseModel {
   /// Create a [GenerativeModel] backed by the generative model named [model].
   ///
   /// The [model] argument can be a model name (such as `'gemini-pro'`) or a
@@ -78,17 +58,19 @@ final class GenerativeModel {
     ToolConfig? toolConfig,
     Content? systemInstruction,
     http.Client? httpClient,
-  })  : _model = _normalizeModelName(model),
-        _baseUri = _vertexUri(app, location),
-        _safetySettings = safetySettings ?? [],
+  })  : _safetySettings = safetySettings ?? [],
         _generationConfig = generationConfig,
         _tools = tools,
         _toolConfig = toolConfig,
         _systemInstruction = systemInstruction,
-        _client = HttpApiClient(
-            apiKey: app.options.apiKey,
-            httpClient: httpClient,
-            requestHeaders: _firebaseTokens(appCheck, auth));
+        super(
+            model: model,
+            app: app,
+            location: location,
+            client: HttpApiClient(
+                apiKey: app.options.apiKey,
+                httpClient: httpClient,
+                requestHeaders: BaseModel.firebaseTokens(appCheck, auth)));
 
   GenerativeModel._constructTestModel({
     required String model,
@@ -102,79 +84,30 @@ final class GenerativeModel {
     ToolConfig? toolConfig,
     Content? systemInstruction,
     ApiClient? apiClient,
-  })  : _model = _normalizeModelName(model),
-        _baseUri = _vertexUri(app, location),
-        _safetySettings = safetySettings ?? [],
+  })  : _safetySettings = safetySettings ?? [],
         _generationConfig = generationConfig,
         _tools = tools,
         _toolConfig = toolConfig,
         _systemInstruction = systemInstruction,
-        _client = apiClient ??
-            HttpApiClient(
-                apiKey: app.options.apiKey,
-                requestHeaders: _firebaseTokens(appCheck, auth));
+        super(
+            model: model,
+            app: app,
+            location: location,
+            client: apiClient ??
+                HttpApiClient(
+                    apiKey: app.options.apiKey,
+                    requestHeaders: BaseModel.firebaseTokens(appCheck, auth)));
 
-  final ({String prefix, String name}) _model;
   final List<SafetySetting> _safetySettings;
   final GenerationConfig? _generationConfig;
   final List<Tool>? _tools;
-  final ApiClient _client;
-  final Uri _baseUri;
+
+  //final Uri _baseUri;
   final ToolConfig? _toolConfig;
   final Content? _systemInstruction;
 
   //static const _modelsPrefix = 'models/';
 
-  /// Returns the model code for a user friendly model name.
-  ///
-  /// If the model name is already a model code (contains a `/`), use the parts
-  /// directly. Otherwise, return a `models/` model code.
-  static ({String prefix, String name}) _normalizeModelName(String modelName) {
-    if (!modelName.contains('/')) return (prefix: 'models', name: modelName);
-    final parts = modelName.split('/');
-    return (prefix: parts.first, name: parts.skip(1).join('/'));
-  }
-
-  static Uri _vertexUri(FirebaseApp app, String location) {
-    var projectId = app.options.projectId;
-    return Uri.https(
-      _baseUrl,
-      '/$_apiVersion/projects/$projectId/locations/$location/publishers/google',
-    );
-  }
-
-  static FutureOr<Map<String, String>> Function() _firebaseTokens(
-      FirebaseAppCheck? appCheck, FirebaseAuth? auth) {
-    return () async {
-      Map<String, String> headers = {};
-      // Override the client name in Google AI SDK
-      headers['x-goog-api-client'] =
-          'gl-dart/$packageVersion fire/$packageVersion';
-      if (appCheck != null) {
-        final appCheckToken = await appCheck.getToken();
-        if (appCheckToken != null) {
-          headers['X-Firebase-AppCheck'] = appCheckToken;
-        }
-      }
-      if (auth != null) {
-        final idToken = await auth.currentUser?.getIdToken();
-        if (idToken != null) {
-          headers['Authorization'] = 'Firebase $idToken';
-        }
-      }
-      return headers;
-    };
-  }
-
-  Uri _taskUri(Task task) => _baseUri.replace(
-      pathSegments: _baseUri.pathSegments
-          .followedBy([_model.prefix, '${_model.name}:${task.name}']));
-
-  /// Make a unary request for [task] with JSON encodable [params].
-  Future<T> makeRequest<T>(Task task, Map<String, Object?> params,
-          T Function(Map<String, Object?>) parse) =>
-      _client.makeRequest(_taskUri(task), params).then(parse);
-
   Map<String, Object?> _generateContentRequest(
     Iterable<Content> contents, {
     List<SafetySetting>? safetySettings,
@@ -187,7 +120,7 @@ final class GenerativeModel {
     tools ??= _tools;
     toolConfig ??= _toolConfig;
     return {
-      'model': '${_model.prefix}/${_model.name}',
+      'model': '${model.prefix}/${model.name}',
       'contents': contents.map((c) => c.toJson()).toList(),
       if (safetySettings.isNotEmpty)
         'safetySettings': safetySettings.map((s) => s.toJson()).toList(),
@@ -244,8 +177,8 @@ final class GenerativeModel {
       GenerationConfig? generationConfig,
       List<Tool>? tools,
       ToolConfig? toolConfig}) {
-    final response = _client.streamRequest(
-        _taskUri(Task.streamGenerateContent),
+    final response = client.streamRequest(
+        taskUri(Task.streamGenerateContent),
         _generateContentRequest(
           prompt,
           safetySettings: safetySettings,

From 7088f422b4ab305b89ed03268ed4f8c9c6e28d82 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Wed, 8 Jan 2025 15:46:27 -0800
Subject: [PATCH 02/21] working demo

---
 .../example/ios/Runner/AppDelegate.swift      |  2 +-
 .../firebase_vertexai/example/lib/main.dart   | 20 ++++--
 .../lib/firebase_vertexai.dart                |  1 +
 .../lib/src/imagen_content.dart               |  7 +-
 .../lib/src/imagen_model.dart                 | 66 ++++++++++++-------
 5 files changed, 64 insertions(+), 32 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/ios/Runner/AppDelegate.swift b/packages/firebase_vertexai/firebase_vertexai/example/ios/Runner/AppDelegate.swift
index 70693e4a8c12..b6363034812b 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/ios/Runner/AppDelegate.swift
+++ b/packages/firebase_vertexai/firebase_vertexai/example/ios/Runner/AppDelegate.swift
@@ -1,7 +1,7 @@
 import UIKit
 import Flutter
 
-@UIApplicationMain
+@main
 @objc class AppDelegate: FlutterAppDelegate {
   override func application(
     _ application: UIApplication,
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
index f0ee767f8c2f..99a5a22306f7 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -267,7 +267,7 @@ class _ChatWidgetState extends State<ChatWidget> {
                   onPressed: !_loading
                       ? () async {
                           //await _sendImagePrompt(_textController.text);
-                          await _testImagen();
+                          await _testImagen(_textController.text);
                         }
                       : null,
                   icon: Icon(
@@ -603,7 +603,10 @@ class _ChatWidgetState extends State<ChatWidget> {
     });
   }
 
-  Future<void> _testImagen() async {
+  Future<void> _testImagen(String prompt) async {
+    setState(() {
+      _loading = true;
+    });
     var model = FirebaseVertexAI.instance.imageModel(
       modelName: 'imagen-3.0-generate-001',
       safetySettings: ImagenSafetySettings(
@@ -612,7 +615,6 @@ class _ChatWidgetState extends State<ChatWidget> {
       ),
     );
 
-    var prompt = 'An astronaut riding a horse.';
     var generationConfig =
         ImagenGenerationConfig('frog', 1, ImagenAspectRatio.square1x1);
 
@@ -622,12 +624,22 @@ class _ChatWidgetState extends State<ChatWidget> {
     );
 
     if (response.images.isNotEmpty) {
-      var image = response.images[0];
+      var imagenImage = response.images[0] as ImagenInlineImage;
       // Process the image
+      _generatedContent.add(
+        (
+          image: Image.memory(imagenImage.bytesBase64Encoded),
+          text: prompt,
+          fromUser: false
+        ),
+      );
     } else {
       // Handle the case where no images were generated
       print('Error: No images were generated.');
     }
+    setState(() {
+      _loading = false;
+    });
   }
 
   void _showError(String message) {
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
index 017237731e86..f3e6026f4db4 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
@@ -65,5 +65,6 @@ export 'src/imagen_api.dart'
         ImagenPersonFilterLevel,
         ImagenGenerationConfig,
         ImagenAspectRatio;
+export 'src/imagen_content.dart' show ImagenInlineImage, ImagenGCSImage;
 export 'src/model.dart' show GenerativeModel;
 export 'src/schema.dart' show Schema, SchemaType;
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
index 4289e6f1dd77..e6a95e793efa 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
@@ -73,7 +73,7 @@ final class ImagenGenerationResponse<T extends ImagenImage> {
 
   factory ImagenGenerationResponse.fromJson(Map<String, dynamic> json) {
     final filteredReason = json['filteredReason'] as String?;
-    final imagesJson = json['images'] as List<dynamic>;
+    final imagesJson = json['predictions'] as List<dynamic>;
 
     if (T == ImagenInlineImage) {
       final images = imagesJson.map((imageJson) {
@@ -105,8 +105,9 @@ final class ImagenGenerationResponse<T extends ImagenImage> {
 }
 
 /// Parse the json to [ImagenGenerationResponse]
-ImagenGenerationResponse parseImagenGenerationResponse(Object jsonObject) {
+ImagenGenerationResponse parseImagenGenerationResponse<T extends ImagenImage>(
+    Object jsonObject) {
   if (jsonObject case {'error': final Object error}) throw parseError(error);
   Map<String, dynamic> json = jsonObject as Map<String, dynamic>;
-  return ImagenGenerationResponse.fromJson(json);
+  return ImagenGenerationResponse<T>.fromJson(json);
 }
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index c25201436691..9226fc4fdacb 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -48,27 +48,42 @@ final class ImagenModel extends BaseModel {
     ImagenGenerationConfig? generationConfig,
     String? gcsUri,
   }) {
+    final parameters = <String, Object?>{};
+
+    if (gcsUri != null) parameters['storageUri'] = gcsUri;
+
+    if (generationConfig != null) {
+      if (generationConfig.numberOfImages != null) {
+        parameters['sampleCount'] = generationConfig.numberOfImages;
+      }
+      if (generationConfig.aspectRatio != null) {
+        parameters['aspectRatio'] = generationConfig.aspectRatio!.toJson();
+      }
+      if (generationConfig.negativePrompt != null) {
+        parameters['negativePrompt'] = generationConfig.negativePrompt;
+      }
+    }
+
+    if (_safetySettings != null) {
+      if (_safetySettings.personFilterLevel != null) {
+        parameters['personGeneration'] =
+            _safetySettings.personFilterLevel!.toJson();
+      }
+      if (_safetySettings.safetyFilterLevel != null) {
+        parameters['safetySetting'] =
+            _safetySettings.safetyFilterLevel!.toJson();
+      }
+    }
+
+    if (_modelConfig != null && _modelConfig.addWatermark != null) {
+      parameters['addWatermark'] = _modelConfig.addWatermark;
+    }
+
     return {
       'instances': [
         {'prompt': prompt}
       ],
-      'parameters': {
-        if (gcsUri != null) 'storageUri': gcsUri,
-        if (generationConfig != null && generationConfig.numberOfImages != null)
-          'sampleCount': generationConfig.numberOfImages,
-        if (generationConfig != null && generationConfig.aspectRatio != null)
-          'aspectRatio': generationConfig.aspectRatio,
-        if (generationConfig != null && generationConfig.negativePrompt != null)
-          'negativePrompt': generationConfig.negativePrompt,
-        if (_safetySettings != null &&
-            _safetySettings.personFilterLevel != null)
-          'personGeneration': _safetySettings.personFilterLevel,
-        if (_safetySettings != null &&
-            _safetySettings.safetyFilterLevel != null)
-          'safetySetting': _safetySettings.safetyFilterLevel,
-        if (_modelConfig != null && _modelConfig.addWatermark != null)
-          'addWatermark': _modelConfig.addWatermark,
-      }
+      'parameters': parameters,
     };
   }
 
@@ -78,13 +93,16 @@ final class ImagenModel extends BaseModel {
     String? gcsUri,
   }) =>
       makeRequest(
-          Task.predict,
-          _generateImagenRequest(
-            prompt,
-            generationConfig: generationConfig,
-            gcsUri: gcsUri,
-          ),
-          parseImagenGenerationResponse);
+        Task.predict,
+        _generateImagenRequest(
+          prompt,
+          generationConfig: generationConfig,
+          gcsUri: gcsUri,
+        ),
+        (jsonObject) => gcsUri != null
+            ? parseImagenGenerationResponse<ImagenGCSImage>(jsonObject)
+            : parseImagenGenerationResponse<ImagenInlineImage>(jsonObject),
+      );
 }
 
 /// Returns a [ImagenModel] using it's private constructor.

From e826339fd1d5efe6dc8cd12d63fbe8a46dd35b58 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Thu, 9 Jan 2025 14:06:11 -0800
Subject: [PATCH 03/21] Separate generateImages and generateImagesGCS

---
 .../firebase_vertexai/example/lib/main.dart   |  2 +-
 .../lib/src/imagen_content.dart               |  4 ++--
 .../lib/src/imagen_model.dart                 | 23 +++++++++++++++----
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
index 99a5a22306f7..4804f63cfa8f 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -624,7 +624,7 @@ class _ChatWidgetState extends State<ChatWidget> {
     );
 
     if (response.images.isNotEmpty) {
-      var imagenImage = response.images[0] as ImagenInlineImage;
+      var imagenImage = response.images[0];
       // Process the image
       _generatedContent.add(
         (
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
index e6a95e793efa..49bc019ffae7 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
@@ -105,8 +105,8 @@ final class ImagenGenerationResponse<T extends ImagenImage> {
 }
 
 /// Parse the json to [ImagenGenerationResponse]
-ImagenGenerationResponse parseImagenGenerationResponse<T extends ImagenImage>(
-    Object jsonObject) {
+ImagenGenerationResponse<T>
+    parseImagenGenerationResponse<T extends ImagenImage>(Object jsonObject) {
   if (jsonObject case {'error': final Object error}) throw parseError(error);
   Map<String, dynamic> json = jsonObject as Map<String, dynamic>;
   return ImagenGenerationResponse<T>.fromJson(json);
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index 9226fc4fdacb..b8dade38f8af 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -87,10 +87,24 @@ final class ImagenModel extends BaseModel {
     };
   }
 
-  Future<ImagenGenerationResponse> generateImages(
+  Future<ImagenGenerationResponse<ImagenInlineImage>> generateImages(
     String prompt, {
     ImagenGenerationConfig? generationConfig,
-    String? gcsUri,
+  }) =>
+      makeRequest(
+        Task.predict,
+        _generateImagenRequest(
+          prompt,
+          generationConfig: generationConfig,
+        ),
+        (jsonObject) =>
+            parseImagenGenerationResponse<ImagenInlineImage>(jsonObject),
+      );
+
+  Future<ImagenGenerationResponse<ImagenGCSImage>> generateImagesGCS(
+    String prompt,
+    String gcsUri, {
+    ImagenGenerationConfig? generationConfig,
   }) =>
       makeRequest(
         Task.predict,
@@ -99,9 +113,8 @@ final class ImagenModel extends BaseModel {
           generationConfig: generationConfig,
           gcsUri: gcsUri,
         ),
-        (jsonObject) => gcsUri != null
-            ? parseImagenGenerationResponse<ImagenGCSImage>(jsonObject)
-            : parseImagenGenerationResponse<ImagenInlineImage>(jsonObject),
+        (jsonObject) =>
+            parseImagenGenerationResponse<ImagenGCSImage>(jsonObject),
       );
 }
 

From adb49f840b2fbfa95f9fcdacfdd6a95cbe3724d7 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Sun, 12 Jan 2025 19:55:23 -0800
Subject: [PATCH 04/21] expose ImagenImage basetype for future

---
 .../firebase_vertexai/lib/firebase_vertexai.dart               | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
index f3e6026f4db4..19bf7ed37043 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
@@ -65,6 +65,7 @@ export 'src/imagen_api.dart'
         ImagenPersonFilterLevel,
         ImagenGenerationConfig,
         ImagenAspectRatio;
-export 'src/imagen_content.dart' show ImagenInlineImage, ImagenGCSImage;
+export 'src/imagen_content.dart'
+    show ImagenInlineImage, ImagenGCSImage, ImagenImage;
 export 'src/model.dart' show GenerativeModel;
 export 'src/schema.dart' show Schema, SchemaType;

From 2a5b26c679c658c79f692340015a71c28269318b Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Mon, 13 Jan 2025 14:14:15 -0800
Subject: [PATCH 05/21] updates after api review

---
 .../firebase_vertexai/example/lib/main.dart   | 10 +++--
 .../lib/src/firebase_vertexai.dart            |  4 +-
 .../firebase_vertexai/lib/src/imagen_api.dart | 27 +++++++++---
 .../lib/src/imagen_model.dart                 | 42 ++++++++++---------
 4 files changed, 53 insertions(+), 30 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
index 4804f63cfa8f..6696b6ca7abc 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -157,7 +157,7 @@ class _ChatWidgetState extends State<ChatWidget> {
   Future<void> initFirebase() async {
     // ignore: avoid_redundant_argument_values
     await Firebase.initializeApp(options: options);
-    await FirebaseAuth.instance.signInAnonymously();
+    // await FirebaseAuth.instance.signInAnonymously();
   }
 
   void _scrollDown() {
@@ -609,14 +609,18 @@ class _ChatWidgetState extends State<ChatWidget> {
     });
     var model = FirebaseVertexAI.instance.imageModel(
       modelName: 'imagen-3.0-generate-001',
+      generationConfig: ImagenGenerationConfig(
+          imageFormat: ImagenFormat.jpeg(compressionQuality: 75)),
       safetySettings: ImagenSafetySettings(
         ImagenSafetyFilterLevel.blockLowAndAbove,
         ImagenPersonFilterLevel.allowAdult,
       ),
     );
 
-    var generationConfig =
-        ImagenGenerationConfig('frog', 1, ImagenAspectRatio.square1x1);
+    var generationConfig = ImagenGenerationConfig(
+        negativePrompt: 'frog',
+        numberOfImages: 1,
+        aspectRatio: ImagenAspectRatio.square1x1);
 
     var response = await model.generateImages(
       prompt,
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
index b0cd87e854f5..100036be6f90 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
@@ -116,13 +116,13 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
 
   ImagenModel imageModel(
       {required String modelName,
-      ImagenModelConfig? modelConfig,
+      ImagenGenerationConfig? generationConfig,
       ImagenSafetySettings? safetySettings}) {
     return createImagenModel(
         app: app,
         location: location,
         modelName: modelName,
-        modelConfig: modelConfig,
+        generationConfig: generationConfig,
         safetySettings: safetySettings,
         appCheck: appCheck,
         auth: auth);
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
index d56101a1331f..846dc7cd07dc 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
@@ -145,10 +145,25 @@ enum ImagenAspectRatio {
 
 final class ImagenGenerationConfig {
   ImagenGenerationConfig(
-      this.negativePrompt, this.numberOfImages, this.aspectRatio);
+      {this.negativePrompt,
+      this.numberOfImages,
+      this.aspectRatio,
+      this.imageFormat,
+      this.addWatermark});
   final String? negativePrompt;
   final int? numberOfImages;
   final ImagenAspectRatio? aspectRatio;
+  final ImagenFormat? imageFormat;
+  final bool? addWatermark;
+
+  /// Convert to json format.
+  Map<String, dynamic> toJson() => {
+        if (negativePrompt != null) 'negativePrompt': negativePrompt,
+        if (numberOfImages != null) 'numberOfImages': numberOfImages,
+        if (aspectRatio != null) 'aspectRatio': aspectRatio!.toJson(),
+        if (addWatermark != null) 'addWatermark': addWatermark,
+        if (imageFormat != null) 'outputOption': imageFormat!.toJson(),
+      };
 }
 
 final class ImagenFormat {
@@ -159,10 +174,10 @@ final class ImagenFormat {
       : this("image/jpeg", compressionQuality);
   final String mimeType;
   final int? compressionQuality;
-}
 
-final class ImagenModelConfig {
-  ImagenModelConfig(this.imagenFormat, this.addWatermark);
-  final ImagenFormat imagenFormat;
-  final bool? addWatermark;
+  Map<String, dynamic> toJson() => {
+        'mimeType': mimeType,
+        if (compressionQuality != null)
+          'compressionQuality': compressionQuality,
+      };
 }
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index b8dade38f8af..2eca0c77c8e7 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -28,9 +28,9 @@ final class ImagenModel extends BaseModel {
       required String location,
       FirebaseAppCheck? appCheck,
       FirebaseAuth? auth,
-      ImagenModelConfig? modelConfig,
+      ImagenGenerationConfig? generationConfig,
       ImagenSafetySettings? safetySettings})
-      : _modelConfig = modelConfig,
+      : _generationConfig = generationConfig,
         _safetySettings = safetySettings,
         super(
             model: modelName,
@@ -40,7 +40,7 @@ final class ImagenModel extends BaseModel {
                 apiKey: app.options.apiKey,
                 requestHeaders: BaseModel.firebaseTokens(appCheck, auth)));
 
-  final ImagenModelConfig? _modelConfig;
+  final ImagenGenerationConfig? _generationConfig;
   final ImagenSafetySettings? _safetySettings;
 
   Map<String, Object?> _generateImagenRequest(
@@ -52,16 +52,24 @@ final class ImagenModel extends BaseModel {
 
     if (gcsUri != null) parameters['storageUri'] = gcsUri;
 
-    if (generationConfig != null) {
-      if (generationConfig.numberOfImages != null) {
-        parameters['sampleCount'] = generationConfig.numberOfImages;
-      }
-      if (generationConfig.aspectRatio != null) {
-        parameters['aspectRatio'] = generationConfig.aspectRatio!.toJson();
-      }
-      if (generationConfig.negativePrompt != null) {
-        parameters['negativePrompt'] = generationConfig.negativePrompt;
-      }
+    // Merge generation configurations
+    final mergedConfig = {
+      ...(_generationConfig?.toJson() ?? {}),
+      ...(generationConfig?.toJson() ?? {}),
+    };
+
+    parameters['sampleCount'] = mergedConfig['numberOfImages'] ?? 1;
+    if (mergedConfig['aspectRatio'] != null) {
+      parameters['aspectRatio'] = mergedConfig['aspectRatio'];
+    }
+    if (mergedConfig['negativePrompt'] != null) {
+      parameters['negativePrompt'] = mergedConfig['negativePrompt'];
+    }
+    if (mergedConfig['addWatermark'] != null) {
+      parameters['addWatermark'] = mergedConfig['addWatermark'];
+    }
+    if (mergedConfig['outputOption'] != null) {
+      parameters['outputOption'] = mergedConfig['outputOption'];
     }
 
     if (_safetySettings != null) {
@@ -75,10 +83,6 @@ final class ImagenModel extends BaseModel {
       }
     }
 
-    if (_modelConfig != null && _modelConfig.addWatermark != null) {
-      parameters['addWatermark'] = _modelConfig.addWatermark;
-    }
-
     return {
       'instances': [
         {'prompt': prompt}
@@ -125,7 +129,7 @@ ImagenModel createImagenModel({
   required String modelName,
   FirebaseAppCheck? appCheck,
   FirebaseAuth? auth,
-  ImagenModelConfig? modelConfig,
+  ImagenGenerationConfig? generationConfig,
   ImagenSafetySettings? safetySettings,
 }) =>
     ImagenModel._(
@@ -135,5 +139,5 @@ ImagenModel createImagenModel({
       auth: auth,
       location: location,
       safetySettings: safetySettings,
-      modelConfig: modelConfig,
+      generationConfig: generationConfig,
     );

From cd59d7446c21b90992c7b5b0be74d55314295b7a Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Wed, 22 Jan 2025 22:07:48 -0800
Subject: [PATCH 06/21] organize example page and functions

---
 .../firebase_vertexai/example/.gitignore      |   2 +
 .../firebase_vertexai/example/lib/main.dart   | 694 +++---------------
 .../example/lib/pages/chat_page.dart          | 176 +++++
 .../lib/pages/function_calling_page.dart      | 186 +++++
 .../example/lib/pages/image_prompt_page.dart  | 183 +++++
 .../example/lib/pages/schema_page.dart        | 182 +++++
 .../example/lib/pages/storage_uri_page.dart   | 171 +++++
 .../example/lib/pages/token_count_page.dart   | 106 +++
 .../example/lib/widgets/message_widget.dart   |  68 ++
 .../example/macos/Runner/AppDelegate.swift    |   6 +-
 10 files changed, 1169 insertions(+), 605 deletions(-)
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/pages/chat_page.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/pages/function_calling_page.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/pages/schema_page.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/pages/token_count_page.dart
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/widgets/message_widget.dart

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/.gitignore b/packages/firebase_vertexai/firebase_vertexai/example/.gitignore
index db070ebdbf4f..0498b592dfa0 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/.gitignore
+++ b/packages/firebase_vertexai/firebase_vertexai/example/.gitignore
@@ -5,9 +5,11 @@
 *.swp
 .DS_Store
 .atom/
+.build/
 .buildlog/
 .history
 .svn/
+.swiftpm/
 migrate_working_dir/
 
 # IntelliJ related
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
index 669252a77d82..1c579d30c4bb 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -16,18 +16,33 @@ import 'package:firebase_core/firebase_core.dart';
 import 'package:firebase_auth/firebase_auth.dart';
 import 'package:firebase_vertexai/firebase_vertexai.dart';
 import 'package:flutter/material.dart';
-import 'package:flutter/services.dart';
-import 'package:flutter_markdown/flutter_markdown.dart';
+
+import 'pages/chat_page.dart';
+import 'pages/function_calling_page.dart';
+import 'pages/image_prompt_page.dart';
+import 'pages/token_count_page.dart';
+import 'pages/schema_page.dart';
+import 'pages/storage_uri_page.dart';
 
 // REQUIRED if you want to run on Web
 const FirebaseOptions? options = null;
 
-void main() {
-  runApp(const GenerativeAISample());
+void main() async {
+  WidgetsFlutterBinding.ensureInitialized();
+  await Firebase.initializeApp(options: options);
+  await FirebaseAuth.instance.signInAnonymously();
+
+  var vertex_instance =
+      FirebaseVertexAI.instanceFor(auth: FirebaseAuth.instance);
+  final model = vertex_instance.generativeModel(model: 'gemini-1.5-flash');
+
+  runApp(GenerativeAISample(model: model));
 }
 
 class GenerativeAISample extends StatelessWidget {
-  const GenerativeAISample({super.key});
+  final GenerativeModel model;
+
+  const GenerativeAISample({super.key, required this.model});
 
   @override
   Widget build(BuildContext context) {
@@ -40,632 +55,103 @@ class GenerativeAISample extends StatelessWidget {
         ),
         useMaterial3: true,
       ),
-      home: const ChatScreen(title: 'Flutter + Vertex AI'),
-    );
-  }
-}
-
-class ChatScreen extends StatefulWidget {
-  const ChatScreen({super.key, required this.title});
-
-  final String title;
-
-  @override
-  State<ChatScreen> createState() => _ChatScreenState();
-}
-
-class _ChatScreenState extends State<ChatScreen> {
-  @override
-  Widget build(BuildContext context) {
-    return Scaffold(
-      appBar: AppBar(
-        title: Text(widget.title),
-      ),
-      body: const ChatWidget(),
+      home: HomeScreen(model: model),
     );
   }
 }
 
-class ChatWidget extends StatefulWidget {
-  const ChatWidget({
-    super.key,
-  });
+class HomeScreen extends StatefulWidget {
+  final GenerativeModel model;
+  const HomeScreen({super.key, required this.model});
 
   @override
-  State<ChatWidget> createState() => _ChatWidgetState();
-}
-
-final class Location {
-  final String city;
-  final String state;
-
-  Location(this.city, this.state);
+  State<HomeScreen> createState() => _HomeScreenState();
 }
 
-class _ChatWidgetState extends State<ChatWidget> {
-  late final GenerativeModel _model;
-  late final GenerativeModel _functionCallModel;
-  ChatSession? _chat;
-  final ScrollController _scrollController = ScrollController();
-  final TextEditingController _textController = TextEditingController();
-  final FocusNode _textFieldFocus = FocusNode();
-  final List<({Image? image, String? text, bool fromUser})> _generatedContent =
-      <({Image? image, String? text, bool fromUser})>[];
-  bool _loading = false;
-
-  @override
-  void initState() {
-    super.initState();
+class _HomeScreenState extends State<HomeScreen> {
+  int _selectedIndex = 0;
+
+  List<Widget> get _pages => <Widget>[
+        // Build _pages dynamically
+        ChatPage(title: 'Chat', model: widget.model),
+        TokenCountPage(title: 'Token Count', model: widget.model),
+        const FunctionCallingPage(
+          title: 'Function Calling',
+        ), // function calling will initial its own model
+        ImagePromptPage(title: 'Image Prompt', model: widget.model),
+        StorageUriPromptPage(title: 'Storage URI Prompt', model: widget.model),
+        SchemaPromptPage(title: 'Schema Prompt', model: widget.model),
+      ];
 
-    initFirebase().then((value) {
-      var vertex_instance =
-          FirebaseVertexAI.instanceFor(auth: FirebaseAuth.instance);
-      _model = vertex_instance.generativeModel(
-        model: 'gemini-1.5-flash',
-      );
-      _functionCallModel = vertex_instance.generativeModel(
-        model: 'gemini-1.5-flash',
-        tools: [
-          Tool.functionDeclarations([fetchWeatherTool]),
-        ],
-      );
-      _chat = _model.startChat();
+  void _onItemTapped(int index) {
+    setState(() {
+      _selectedIndex = index;
     });
   }
 
-  // This is a hypothetical API to return a fake weather data collection for
-  // certain location
-  Future<Map<String, Object?>> fetchWeather(
-    Location location,
-    String date,
-  ) async {
-    // TODO(developer): Call a real weather API.
-    // Mock response from the API. In developer live code this would call the
-    // external API and return what that API returns.
-    final apiResponse = {
-      'temperature': 38,
-      'chancePrecipitation': '56%',
-      'cloudConditions': 'partly-cloudy',
-    };
-    return apiResponse;
-  }
-
-  /// Actual function to demonstrate the function calling feature.
-  final fetchWeatherTool = FunctionDeclaration(
-    'fetchWeather',
-    'Get the weather conditions for a specific city on a specific date.',
-    parameters: {
-      'location': Schema.object(
-        description: 'The name of the city and its state for which to get '
-            'the weather. Only cities in the USA are supported.',
-        properties: {
-          'city': Schema.string(
-            description: 'The city of the location.',
-          ),
-          'state': Schema.string(
-            description: 'The state of the location.',
-          ),
-        },
-      ),
-      'date': Schema.string(
-        description: 'The date for which to get the weather. '
-            'Date must be in the format: YYYY-MM-DD.',
-      ),
-    },
-  );
-
-  Future<void> initFirebase() async {
-    // ignore: avoid_redundant_argument_values
-    await Firebase.initializeApp(options: options);
-    await FirebaseAuth.instance.signInAnonymously();
-  }
-
-  void _scrollDown() {
-    WidgetsBinding.instance.addPostFrameCallback(
-      (_) => _scrollController.animateTo(
-        _scrollController.position.maxScrollExtent,
-        duration: const Duration(
-          milliseconds: 750,
-        ),
-        curve: Curves.easeOutCirc,
-      ),
-    );
-  }
-
   @override
   Widget build(BuildContext context) {
-    final textFieldDecoration = InputDecoration(
-      contentPadding: const EdgeInsets.all(15),
-      hintText: 'Enter a prompt...',
-      border: OutlineInputBorder(
-        borderRadius: const BorderRadius.all(
-          Radius.circular(14),
-        ),
-        borderSide: BorderSide(
-          color: Theme.of(context).colorScheme.secondary,
-        ),
+    return Scaffold(
+      appBar: AppBar(
+        title: const Text('Flutter + Vertex AI'),
       ),
-      focusedBorder: OutlineInputBorder(
-        borderRadius: const BorderRadius.all(
-          Radius.circular(14),
-        ),
-        borderSide: BorderSide(
-          color: Theme.of(context).colorScheme.secondary,
-        ),
+      body: Center(
+        child: _pages.elementAt(_selectedIndex),
       ),
-    );
-
-    return Padding(
-      padding: const EdgeInsets.all(8),
-      child: Column(
-        mainAxisAlignment: MainAxisAlignment.center,
-        crossAxisAlignment: CrossAxisAlignment.start,
-        children: [
-          Expanded(
-            child: ListView.builder(
-              controller: _scrollController,
-              itemBuilder: (context, idx) {
-                var content = _generatedContent[idx];
-                return MessageWidget(
-                  text: content.text,
-                  image: content.image,
-                  isFromUser: content.fromUser,
-                );
-              },
-              itemCount: _generatedContent.length,
+      bottomNavigationBar: BottomNavigationBar(
+        items: <BottomNavigationBarItem>[
+          BottomNavigationBarItem(
+            icon: Icon(
+              Icons.chat,
+              color: Theme.of(context).colorScheme.primary,
             ),
+            label: 'Chat',
+            tooltip: 'Chat',
           ),
-          Padding(
-            padding: const EdgeInsets.symmetric(
-              vertical: 25,
-              horizontal: 15,
-            ),
-            child: Row(
-              children: [
-                Expanded(
-                  child: TextField(
-                    autofocus: true,
-                    focusNode: _textFieldFocus,
-                    decoration: textFieldDecoration,
-                    controller: _textController,
-                    onSubmitted: _sendChatMessage,
-                  ),
-                ),
-                const SizedBox.square(
-                  dimension: 15,
-                ),
-                IconButton(
-                  tooltip: 'tokenCount Test',
-                  onPressed: !_loading
-                      ? () async {
-                          await _testCountToken();
-                        }
-                      : null,
-                  icon: Icon(
-                    Icons.numbers,
-                    color: _loading
-                        ? Theme.of(context).colorScheme.secondary
-                        : Theme.of(context).colorScheme.primary,
-                  ),
-                ),
-                IconButton(
-                  tooltip: 'function calling Test',
-                  onPressed: !_loading
-                      ? () async {
-                          await _testFunctionCalling();
-                        }
-                      : null,
-                  icon: Icon(
-                    Icons.functions,
-                    color: _loading
-                        ? Theme.of(context).colorScheme.secondary
-                        : Theme.of(context).colorScheme.primary,
-                  ),
-                ),
-                IconButton(
-                  tooltip: 'image prompt',
-                  onPressed: !_loading
-                      ? () async {
-                          await _sendImagePrompt(_textController.text);
-                        }
-                      : null,
-                  icon: Icon(
-                    Icons.image,
-                    color: _loading
-                        ? Theme.of(context).colorScheme.secondary
-                        : Theme.of(context).colorScheme.primary,
-                  ),
-                ),
-                IconButton(
-                  tooltip: 'storage prompt',
-                  onPressed: !_loading
-                      ? () async {
-                          await _sendStorageUriPrompt(_textController.text);
-                        }
-                      : null,
-                  icon: Icon(
-                    Icons.folder,
-                    color: _loading
-                        ? Theme.of(context).colorScheme.secondary
-                        : Theme.of(context).colorScheme.primary,
-                  ),
-                ),
-                IconButton(
-                  tooltip: 'schema prompt',
-                  onPressed: !_loading
-                      ? () async {
-                          await _promptSchemaTest(_textController.text);
-                        }
-                      : null,
-                  icon: Icon(
-                    Icons.schema,
-                    color: _loading
-                        ? Theme.of(context).colorScheme.secondary
-                        : Theme.of(context).colorScheme.primary,
-                  ),
-                ),
-                if (!_loading)
-                  IconButton(
-                    onPressed: () async {
-                      await _sendChatMessage(_textController.text);
-                    },
-                    icon: Icon(
-                      Icons.send,
-                      color: Theme.of(context).colorScheme.primary,
-                    ),
-                  )
-                else
-                  const CircularProgressIndicator(),
-              ],
+          BottomNavigationBarItem(
+            icon: Icon(
+              Icons.numbers,
+              color: Theme.of(context).colorScheme.primary,
             ),
+            label: 'Token Count',
+            tooltip: 'Token Count',
           ),
-          Padding(
-            padding: const EdgeInsets.only(
-              left: 15,
-              right: 15,
-              bottom: 25,
-            ),
-            child: Text(
-              'Total message count: ${_chat?.history.length ?? 0}',
+          BottomNavigationBarItem(
+            icon: Icon(
+              Icons.functions,
+              color: Theme.of(context).colorScheme.primary,
             ),
+            label: 'Function Calling',
+            tooltip: 'Function Calling',
           ),
-        ],
-      ),
-    );
-  }
-
-  Future<void> _promptSchemaTest(String subject) async {
-    setState(() {
-      _loading = true;
-    });
-    try {
-      final content = [
-        Content.text(
-          "For use in a children's card game, generate 10 animal-based "
-          'characters.',
-        ),
-      ];
-
-      final jsonSchema = Schema.object(
-        properties: {
-          'characters': Schema.array(
-            items: Schema.object(
-              properties: {
-                'name': Schema.string(),
-                'age': Schema.integer(),
-                'species': Schema.string(),
-                'accessory':
-                    Schema.enumString(enumValues: ['hat', 'belt', 'shoes']),
-              },
+          BottomNavigationBarItem(
+            icon: Icon(
+              Icons.image,
+              color: Theme.of(context).colorScheme.primary,
             ),
+            label: 'Image Prompt',
+            tooltip: 'Image Prompt',
           ),
-        },
-        optionalProperties: ['accessory'],
-      );
-
-      final response = await _model.generateContent(
-        content,
-        generationConfig: GenerationConfig(
-          responseMimeType: 'application/json',
-          responseSchema: jsonSchema,
-        ),
-      );
-
-      var text = response.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _sendStorageUriPrompt(String message) async {
-    setState(() {
-      _loading = true;
-    });
-    try {
-      final content = [
-        Content.multi([
-          TextPart(message),
-          FileData(
-            'image/jpeg',
-            'gs://vertex-ai-example-ef5a2.appspot.com/foodpic.jpg',
-          ),
-        ]),
-      ];
-      _generatedContent.add((image: null, text: message, fromUser: true));
-
-      var response = await _model.generateContent(content);
-      var text = response.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _sendImagePrompt(String message) async {
-    setState(() {
-      _loading = true;
-    });
-    try {
-      ByteData catBytes = await rootBundle.load('assets/images/cat.jpg');
-      ByteData sconeBytes = await rootBundle.load('assets/images/scones.jpg');
-      final content = [
-        Content.multi([
-          TextPart(message),
-          // The only accepted mime types are image/*.
-          InlineDataPart('image/jpeg', catBytes.buffer.asUint8List()),
-          InlineDataPart('image/jpeg', sconeBytes.buffer.asUint8List()),
-        ]),
-      ];
-      _generatedContent.add(
-        (
-          image: Image.asset('assets/images/cat.jpg'),
-          text: message,
-          fromUser: true
-        ),
-      );
-      _generatedContent.add(
-        (
-          image: Image.asset('assets/images/scones.jpg'),
-          text: null,
-          fromUser: true
-        ),
-      );
-
-      var response = await _model.generateContent(content);
-      var text = response.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _sendChatMessage(String message) async {
-    setState(() {
-      _loading = true;
-    });
-
-    try {
-      _generatedContent.add((image: null, text: message, fromUser: true));
-      var response = await _chat?.sendMessage(
-        Content.text(message),
-      );
-      var text = response?.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _testFunctionCalling() async {
-    setState(() {
-      _loading = true;
-    });
-    final functionCallChat = _functionCallModel.startChat();
-    const prompt = 'What is the weather like in Boston on 10/02 this year?';
-
-    // Send the message to the generative model.
-    var response = await functionCallChat.sendMessage(
-      Content.text(prompt),
-    );
-
-    final functionCalls = response.functionCalls.toList();
-    // When the model response with a function call, invoke the function.
-    if (functionCalls.isNotEmpty) {
-      final functionCall = functionCalls.first;
-      if (functionCall.name == 'fetchWeather') {
-        Map<String, dynamic> location =
-            functionCall.args['location']! as Map<String, dynamic>;
-        var date = functionCall.args['date']! as String;
-        var city = location['city'] as String;
-        var state = location['state'] as String;
-        final functionResult = await fetchWeather(Location(city, state), date);
-        // Send the response to the model so that it can use the result to
-        // generate text for the user.
-        response = await functionCallChat.sendMessage(
-          Content.functionResponse(functionCall.name, functionResult),
-        );
-      } else {
-        throw UnimplementedError(
-          'Function not declared to the model: ${functionCall.name}',
-        );
-      }
-    }
-    // When the model responds with non-null text content, print it.
-    if (response.text case final text?) {
-      _generatedContent.add((image: null, text: text, fromUser: false));
-      setState(() {
-        _loading = false;
-      });
-    }
-  }
-
-  Future<void> _testCountToken() async {
-    setState(() {
-      _loading = true;
-    });
-
-    const prompt = 'tell a short story';
-    var content = Content.text(prompt);
-    var tokenResponse = await _model.countTokens([content]);
-    final tokenResult = 'Count token: ${tokenResponse.totalTokens}, billable '
-        'characters: ${tokenResponse.totalBillableCharacters}';
-    _generatedContent.add((image: null, text: tokenResult, fromUser: false));
-
-    var contentResponse = await _model.generateContent([content]);
-    final contentMetaData = 'result metadata, promptTokenCount:'
-        '${contentResponse.usageMetadata!.promptTokenCount}, '
-        'candidatesTokenCount:'
-        '${contentResponse.usageMetadata!.candidatesTokenCount}, '
-        'totalTokenCount:'
-        '${contentResponse.usageMetadata!.totalTokenCount}';
-    _generatedContent
-        .add((image: null, text: contentMetaData, fromUser: false));
-    setState(() {
-      _loading = false;
-    });
-  }
-
-  void _showError(String message) {
-    showDialog<void>(
-      context: context,
-      builder: (context) {
-        return AlertDialog(
-          title: const Text('Something went wrong'),
-          content: SingleChildScrollView(
-            child: SelectableText(message),
-          ),
-          actions: [
-            TextButton(
-              onPressed: () {
-                Navigator.of(context).pop();
-              },
-              child: const Text('OK'),
+          BottomNavigationBarItem(
+            icon: Icon(
+              Icons.folder,
+              color: Theme.of(context).colorScheme.primary,
             ),
-          ],
-        );
-      },
-    );
-  }
-}
-
-class MessageWidget extends StatelessWidget {
-  final Image? image;
-  final String? text;
-  final bool isFromUser;
-
-  const MessageWidget({
-    super.key,
-    this.image,
-    this.text,
-    required this.isFromUser,
-  });
-
-  @override
-  Widget build(BuildContext context) {
-    return Row(
-      mainAxisAlignment:
-          isFromUser ? MainAxisAlignment.end : MainAxisAlignment.start,
-      children: [
-        Flexible(
-          child: Container(
-            constraints: const BoxConstraints(maxWidth: 600),
-            decoration: BoxDecoration(
-              color: isFromUser
-                  ? Theme.of(context).colorScheme.primaryContainer
-                  : Theme.of(context).colorScheme.surfaceContainerHighest,
-              borderRadius: BorderRadius.circular(18),
-            ),
-            padding: const EdgeInsets.symmetric(
-              vertical: 15,
-              horizontal: 20,
-            ),
-            margin: const EdgeInsets.only(bottom: 8),
-            child: Column(
-              children: [
-                if (text case final text?) MarkdownBody(data: text),
-                if (image case final image?) image,
-              ],
+            label: 'Storage URI Prompt',
+            tooltip: 'Storage URI Prompt',
+          ),
+          BottomNavigationBarItem(
+            icon: Icon(
+              Icons.schema,
+              color: Theme.of(context).colorScheme.primary,
             ),
+            label: 'Schema Prompt',
+            tooltip: 'Schema Prompt',
           ),
-        ),
-      ],
+        ],
+        currentIndex: _selectedIndex,
+        onTap: _onItemTapped,
+      ),
     );
   }
 }
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/chat_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/chat_page.dart
new file mode 100644
index 000000000000..33b6e1141b72
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/chat_page.dart
@@ -0,0 +1,176 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'package:flutter/material.dart';
+import 'package:firebase_vertexai/firebase_vertexai.dart';
+import '../widgets/message_widget.dart';
+
+class ChatPage extends StatefulWidget {
+  const ChatPage({super.key, required this.title, required this.model});
+
+  final String title;
+  final GenerativeModel model;
+
+  @override
+  State<ChatPage> createState() => _ChatPageState();
+}
+
+class _ChatPageState extends State<ChatPage> {
+  ChatSession? _chat;
+  final ScrollController _scrollController = ScrollController();
+  final TextEditingController _textController = TextEditingController();
+  final FocusNode _textFieldFocus = FocusNode();
+  final List<MessageData> _messages = <MessageData>[];
+  bool _loading = false;
+
+  @override
+  void initState() {
+    super.initState();
+    _chat = widget.model.startChat();
+  }
+
+  void _scrollDown() {
+    WidgetsBinding.instance.addPostFrameCallback(
+      (_) => _scrollController.animateTo(
+        _scrollController.position.maxScrollExtent,
+        duration: const Duration(
+          milliseconds: 750,
+        ),
+        curve: Curves.easeOutCirc,
+      ),
+    );
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    return Scaffold(
+      appBar: AppBar(
+        title: Text(widget.title),
+      ),
+      body: Padding(
+        padding: const EdgeInsets.all(8),
+        child: Column(
+          mainAxisAlignment: MainAxisAlignment.center,
+          crossAxisAlignment: CrossAxisAlignment.start,
+          children: [
+            Expanded(
+              child: ListView.builder(
+                controller: _scrollController,
+                itemBuilder: (context, idx) {
+                  return MessageWidget(
+                    text: _messages[idx].text,
+                    image: _messages[idx].image,
+                    isFromUser: _messages[idx].fromUser ?? false,
+                  );
+                },
+                itemCount: _messages.length,
+              ),
+            ),
+            Padding(
+              padding: const EdgeInsets.symmetric(
+                vertical: 25,
+                horizontal: 15,
+              ),
+              child: Row(
+                children: [
+                  Expanded(
+                    child: TextField(
+                      autofocus: true,
+                      focusNode: _textFieldFocus,
+                      controller: _textController,
+                      onSubmitted: _sendChatMessage,
+                    ),
+                  ),
+                  const SizedBox.square(
+                    dimension: 15,
+                  ),
+                  if (!_loading)
+                    IconButton(
+                      onPressed: () async {
+                        await _sendChatMessage(_textController.text);
+                      },
+                      icon: Icon(
+                        Icons.send,
+                        color: Theme.of(context).colorScheme.primary,
+                      ),
+                    )
+                  else
+                    const CircularProgressIndicator(),
+                ],
+              ),
+            ),
+          ],
+        ),
+      ),
+    );
+  }
+
+  Future<void> _sendChatMessage(String message) async {
+    setState(() {
+      _loading = true;
+    });
+
+    try {
+      _messages.add(MessageData(text: message, fromUser: true));
+      var response = await _chat?.sendMessage(
+        Content.text(message),
+      );
+      var text = response?.text;
+      _messages.add(MessageData(text: text, fromUser: false));
+
+      if (text == null) {
+        _showError('No response from API.');
+        return;
+      } else {
+        setState(() {
+          _loading = false;
+          _scrollDown();
+        });
+      }
+    } catch (e) {
+      _showError(e.toString());
+      setState(() {
+        _loading = false;
+      });
+    } finally {
+      _textController.clear();
+      setState(() {
+        _loading = false;
+      });
+      _textFieldFocus.requestFocus();
+    }
+  }
+
+  void _showError(String message) {
+    showDialog<void>(
+      context: context,
+      builder: (context) {
+        return AlertDialog(
+          title: const Text('Something went wrong'),
+          content: SingleChildScrollView(
+            child: SelectableText(message),
+          ),
+          actions: [
+            TextButton(
+              onPressed: () {
+                Navigator.of(context).pop();
+              },
+              child: const Text('OK'),
+            ),
+          ],
+        );
+      },
+    );
+  }
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/function_calling_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/function_calling_page.dart
new file mode 100644
index 000000000000..130afff5ce92
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/function_calling_page.dart
@@ -0,0 +1,186 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'package:flutter/material.dart';
+import 'package:firebase_vertexai/firebase_vertexai.dart';
+import 'package:firebase_auth/firebase_auth.dart';
+import '../widgets/message_widget.dart';
+
+class FunctionCallingPage extends StatefulWidget {
+  const FunctionCallingPage({super.key, required this.title});
+
+  final String title;
+
+  @override
+  State<FunctionCallingPage> createState() => _FunctionCallingPageState();
+}
+
+class Location {
+  final String city;
+  final String state;
+
+  Location(this.city, this.state);
+}
+
+class _FunctionCallingPageState extends State<FunctionCallingPage> {
+  late final GenerativeModel _functionCallModel;
+  final List<MessageData> _messages = <MessageData>[];
+  bool _loading = false;
+
+  @override
+  void initState() {
+    super.initState();
+    var vertex_instance =
+        FirebaseVertexAI.instanceFor(auth: FirebaseAuth.instance);
+    _functionCallModel = vertex_instance.generativeModel(
+      model: 'gemini-1.5-flash',
+      tools: [
+        Tool.functionDeclarations([fetchWeatherTool]),
+      ],
+    );
+  }
+
+  // This is a hypothetical API to return a fake weather data collection for
+  // certain location
+  Future<Map<String, Object?>> fetchWeather(
+    Location location,
+    String date,
+  ) async {
+    // TODO(developer): Call a real weather API.
+    // Mock response from the API. In developer live code this would call the
+    // external API and return what that API returns.
+    final apiResponse = {
+      'temperature': 38,
+      'chancePrecipitation': '56%',
+      'cloudConditions': 'partly-cloudy',
+    };
+    return apiResponse;
+  }
+
+  /// Actual function to demonstrate the function calling feature.
+  final fetchWeatherTool = FunctionDeclaration(
+    'fetchWeather',
+    'Get the weather conditions for a specific city on a specific date.',
+    parameters: {
+      'location': Schema.object(
+        description: 'The name of the city and its state for which to get '
+            'the weather. Only cities in the USA are supported.',
+        properties: {
+          'city': Schema.string(
+            description: 'The city of the location.',
+          ),
+          'state': Schema.string(
+            description: 'The state of the location.',
+          ),
+        },
+      ),
+      'date': Schema.string(
+        description: 'The date for which to get the weather. '
+            'Date must be in the format: YYYY-MM-DD.',
+      ),
+    },
+  );
+
+  @override
+  Widget build(BuildContext context) {
+    return Scaffold(
+      appBar: AppBar(
+        title: Text(widget.title),
+      ),
+      body: Padding(
+        padding: const EdgeInsets.all(8),
+        child: Column(
+          mainAxisAlignment: MainAxisAlignment.center,
+          crossAxisAlignment: CrossAxisAlignment.start,
+          children: [
+            Expanded(
+              child: ListView.builder(
+                itemBuilder: (context, idx) {
+                  return MessageWidget(
+                    text: _messages[idx].text,
+                    isFromUser: _messages[idx].fromUser ?? false,
+                  );
+                },
+                itemCount: _messages.length,
+              ),
+            ),
+            Padding(
+              padding: const EdgeInsets.symmetric(
+                vertical: 25,
+                horizontal: 15,
+              ),
+              child: Row(
+                children: [
+                  Expanded(
+                    child: ElevatedButton(
+                      onPressed: !_loading
+                          ? () async {
+                              await _testFunctionCalling();
+                            }
+                          : null,
+                      child: const Text('Test Function Calling'),
+                    ),
+                  ),
+                ],
+              ),
+            ),
+          ],
+        ),
+      ),
+    );
+  }
+
+  Future<void> _testFunctionCalling() async {
+    setState(() {
+      _loading = true;
+    });
+    final functionCallChat = _functionCallModel.startChat();
+    const prompt = 'What is the weather like in Boston on 10/02 this year?';
+
+    // Send the message to the generative model.
+    var response = await functionCallChat.sendMessage(
+      Content.text(prompt),
+    );
+
+    final functionCalls = response.functionCalls.toList();
+    // When the model response with a function call, invoke the function.
+    if (functionCalls.isNotEmpty) {
+      final functionCall = functionCalls.first;
+      if (functionCall.name == 'fetchWeather') {
+        Map<String, dynamic> location =
+            functionCall.args['location']! as Map<String, dynamic>;
+        var date = functionCall.args['date']! as String;
+        var city = location['city'] as String;
+        var state = location['state'] as String;
+        final functionResult = await fetchWeather(Location(city, state), date);
+        // Send the response to the model so that it can use the result to
+        // generate text for the user.
+        response = await functionCallChat.sendMessage(
+          Content.functionResponse(functionCall.name, functionResult),
+        );
+      } else {
+        throw UnimplementedError(
+          'Function not declared to the model: ${functionCall.name}',
+        );
+      }
+    }
+    // When the model responds with non-null text content, print it.
+    if (response.text case final text?) {
+      _messages.add(MessageData(text: text));
+      setState(() {
+        _loading = false;
+      });
+    }
+  }
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
new file mode 100644
index 000000000000..c110c5205757
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
@@ -0,0 +1,183 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'package:flutter/material.dart';
+import 'package:firebase_vertexai/firebase_vertexai.dart';
+import 'package:flutter/services.dart';
+import '../widgets/message_widget.dart';
+
+class ImagePromptPage extends StatefulWidget {
+  const ImagePromptPage({super.key, required this.title, required this.model});
+
+  final String title;
+  final GenerativeModel model;
+
+  @override
+  State<ImagePromptPage> createState() => _ImagePromptPageState();
+}
+
+class _ImagePromptPageState extends State<ImagePromptPage> {
+  final ScrollController _scrollController = ScrollController();
+  final TextEditingController _textController = TextEditingController();
+  final FocusNode _textFieldFocus = FocusNode();
+  final List<MessageData> _generatedContent = <MessageData>[];
+  bool _loading = false;
+
+  void _scrollDown() {
+    WidgetsBinding.instance.addPostFrameCallback(
+      (_) => _scrollController.animateTo(
+        _scrollController.position.maxScrollExtent,
+        duration: const Duration(
+          milliseconds: 750,
+        ),
+        curve: Curves.easeOutCirc,
+      ),
+    );
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    return Scaffold(
+      appBar: AppBar(
+        title: Text(widget.title),
+      ),
+      body: Padding(
+        padding: const EdgeInsets.all(8),
+        child: Column(
+          mainAxisAlignment: MainAxisAlignment.center,
+          crossAxisAlignment: CrossAxisAlignment.start,
+          children: [
+            Expanded(
+              child: ListView.builder(
+                controller: _scrollController,
+                itemBuilder: (context, idx) {
+                  var content = _generatedContent[idx];
+                  return MessageWidget(
+                    text: content.text,
+                    image: content.image,
+                    isFromUser: content.fromUser ?? false,
+                  );
+                },
+                itemCount: _generatedContent.length,
+              ),
+            ),
+            Padding(
+              padding: const EdgeInsets.symmetric(
+                vertical: 25,
+                horizontal: 15,
+              ),
+              child: Row(
+                children: [
+                  Expanded(
+                    child: TextField(
+                      autofocus: true,
+                      focusNode: _textFieldFocus,
+                      controller: _textController,
+                    ),
+                  ),
+                  const SizedBox.square(
+                    dimension: 15,
+                  ),
+                  ElevatedButton(
+                    onPressed: !_loading
+                        ? () async {
+                            await _sendImagePrompt(_textController.text);
+                          }
+                        : null,
+                    child: const Text('Send Image Prompt'),
+                  ),
+                ],
+              ),
+            ),
+          ],
+        ),
+      ),
+    );
+  }
+
+  Future<void> _sendImagePrompt(String message) async {
+    setState(() {
+      _loading = true;
+    });
+    try {
+      ByteData catBytes = await rootBundle.load('assets/images/cat.jpg');
+      ByteData sconeBytes = await rootBundle.load('assets/images/scones.jpg');
+      final content = [
+        Content.multi([
+          TextPart(message),
+          // The only accepted mime types are image/*.
+          InlineDataPart('image/jpeg', catBytes.buffer.asUint8List()),
+          InlineDataPart('image/jpeg', sconeBytes.buffer.asUint8List()),
+        ]),
+      ];
+      _generatedContent.add(
+        MessageData(
+            image: Image.asset('assets/images/cat.jpg'),
+            text: message,
+            fromUser: true),
+      );
+      _generatedContent.add(
+        MessageData(
+            image: Image.asset('assets/images/scones.jpg'), fromUser: true),
+      );
+
+      var response = await widget.model.generateContent(content);
+      var text = response.text;
+      _generatedContent.add(MessageData(text: text, fromUser: false));
+
+      if (text == null) {
+        _showError('No response from API.');
+        return;
+      } else {
+        setState(() {
+          _loading = false;
+          _scrollDown();
+        });
+      }
+    } catch (e) {
+      _showError(e.toString());
+      setState(() {
+        _loading = false;
+      });
+    } finally {
+      _textController.clear();
+      setState(() {
+        _loading = false;
+      });
+      _textFieldFocus.requestFocus();
+    }
+  }
+
+  void _showError(String message) {
+    showDialog<void>(
+      context: context,
+      builder: (context) {
+        return AlertDialog(
+          title: const Text('Something went wrong'),
+          content: SingleChildScrollView(
+            child: SelectableText(message),
+          ),
+          actions: [
+            TextButton(
+              onPressed: () {
+                Navigator.of(context).pop();
+              },
+              child: const Text('OK'),
+            ),
+          ],
+        );
+      },
+    );
+  }
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/schema_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/schema_page.dart
new file mode 100644
index 000000000000..71fff4269b1e
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/schema_page.dart
@@ -0,0 +1,182 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'package:flutter/material.dart';
+import 'package:firebase_vertexai/firebase_vertexai.dart';
+import '../widgets/message_widget.dart';
+
+class SchemaPromptPage extends StatefulWidget {
+  const SchemaPromptPage({super.key, required this.title, required this.model});
+
+  final String title;
+  final GenerativeModel model;
+
+  @override
+  State<SchemaPromptPage> createState() => _SchemaPromptPageState();
+}
+
+class _SchemaPromptPageState extends State<SchemaPromptPage> {
+  final ScrollController _scrollController = ScrollController();
+  final TextEditingController _textController = TextEditingController();
+  final FocusNode _textFieldFocus = FocusNode();
+  final List<MessageData> _messages = <MessageData>[];
+  bool _loading = false;
+
+  void _scrollDown() {
+    WidgetsBinding.instance.addPostFrameCallback(
+      (_) => _scrollController.animateTo(
+        _scrollController.position.maxScrollExtent,
+        duration: const Duration(
+          milliseconds: 750,
+        ),
+        curve: Curves.easeOutCirc,
+      ),
+    );
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    return Scaffold(
+      appBar: AppBar(
+        title: Text(widget.title),
+      ),
+      body: Padding(
+        padding: const EdgeInsets.all(8),
+        child: Column(
+          mainAxisAlignment: MainAxisAlignment.center,
+          crossAxisAlignment: CrossAxisAlignment.start,
+          children: [
+            Expanded(
+              child: ListView.builder(
+                controller: _scrollController,
+                itemBuilder: (context, idx) {
+                  return MessageWidget(
+                    text: _messages[idx].text,
+                    isFromUser: _messages[idx].fromUser ?? false,
+                  );
+                },
+                itemCount: _messages.length,
+              ),
+            ),
+            Padding(
+              padding: const EdgeInsets.symmetric(
+                vertical: 25,
+                horizontal: 15,
+              ),
+              child: Row(
+                children: [
+                  Expanded(
+                    child: ElevatedButton(
+                      onPressed: !_loading
+                          ? () async {
+                              await _promptSchemaTest();
+                            }
+                          : null,
+                      child: const Text('Schema Prompt'),
+                    ),
+                  ),
+                ],
+              ),
+            ),
+          ],
+        ),
+      ),
+    );
+  }
+
+  Future<void> _promptSchemaTest() async {
+    setState(() {
+      _loading = true;
+    });
+    try {
+      final content = [
+        Content.text(
+          "For use in a children's card game, generate 10 animal-based "
+          'characters.',
+        ),
+      ];
+
+      final jsonSchema = Schema.object(
+        properties: {
+          'characters': Schema.array(
+            items: Schema.object(
+              properties: {
+                'name': Schema.string(),
+                'age': Schema.integer(),
+                'species': Schema.string(),
+                'accessory':
+                    Schema.enumString(enumValues: ['hat', 'belt', 'shoes']),
+              },
+            ),
+          ),
+        },
+        optionalProperties: ['accessory'],
+      );
+
+      final response = await widget.model.generateContent(
+        content,
+        generationConfig: GenerationConfig(
+          responseMimeType: 'application/json',
+          responseSchema: jsonSchema,
+        ),
+      );
+
+      var text = response.text;
+      _messages.add(MessageData(text: text, fromUser: false));
+
+      if (text == null) {
+        _showError('No response from API.');
+        return;
+      } else {
+        setState(() {
+          _loading = false;
+          _scrollDown();
+        });
+      }
+    } catch (e) {
+      _showError(e.toString());
+      setState(() {
+        _loading = false;
+      });
+    } finally {
+      _textController.clear();
+      setState(() {
+        _loading = false;
+      });
+      _textFieldFocus.requestFocus();
+    }
+  }
+
+  void _showError(String message) {
+    showDialog<void>(
+      context: context,
+      builder: (context) {
+        return AlertDialog(
+          title: const Text('Something went wrong'),
+          content: SingleChildScrollView(
+            child: SelectableText(message),
+          ),
+          actions: [
+            TextButton(
+              onPressed: () {
+                Navigator.of(context).pop();
+              },
+              child: const Text('OK'),
+            ),
+          ],
+        );
+      },
+    );
+  }
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
new file mode 100644
index 000000000000..eb23c1950d40
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
@@ -0,0 +1,171 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'package:flutter/material.dart';
+import 'package:firebase_vertexai/firebase_vertexai.dart';
+import '../widgets/message_widget.dart';
+
+class StorageUriPromptPage extends StatefulWidget {
+  const StorageUriPromptPage(
+      {super.key, required this.title, required this.model});
+
+  final String title;
+  final GenerativeModel model;
+
+  @override
+  State<StorageUriPromptPage> createState() => _StorageUriPromptPageState();
+}
+
+class _StorageUriPromptPageState extends State<StorageUriPromptPage> {
+  final ScrollController _scrollController = ScrollController();
+  final TextEditingController _textController = TextEditingController();
+  final FocusNode _textFieldFocus = FocusNode();
+  final List<MessageData> _messages = <MessageData>[];
+  bool _loading = false;
+
+  void _scrollDown() {
+    WidgetsBinding.instance.addPostFrameCallback(
+      (_) => _scrollController.animateTo(
+        _scrollController.position.maxScrollExtent,
+        duration: const Duration(
+          milliseconds: 750,
+        ),
+        curve: Curves.easeOutCirc,
+      ),
+    );
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    return Scaffold(
+      appBar: AppBar(
+        title: Text(widget.title),
+      ),
+      body: Padding(
+        padding: const EdgeInsets.all(8),
+        child: Column(
+          mainAxisAlignment: MainAxisAlignment.center,
+          crossAxisAlignment: CrossAxisAlignment.start,
+          children: [
+            Expanded(
+              child: ListView.builder(
+                controller: _scrollController,
+                itemBuilder: (context, idx) {
+                  return MessageWidget(
+                    text: _messages[idx].text,
+                    isFromUser: _messages[idx].fromUser ?? false,
+                  );
+                },
+                itemCount: _messages.length,
+              ),
+            ),
+            Padding(
+              padding: const EdgeInsets.symmetric(
+                vertical: 25,
+                horizontal: 15,
+              ),
+              child: Row(
+                children: [
+                  Expanded(
+                    child: TextField(
+                      autofocus: true,
+                      focusNode: _textFieldFocus,
+                      controller: _textController,
+                    ),
+                  ),
+                  const SizedBox.square(
+                    dimension: 15,
+                  ),
+                  ElevatedButton(
+                    onPressed: !_loading
+                        ? () async {
+                            await _sendStorageUriPrompt(_textController.text);
+                          }
+                        : null,
+                    child: const Text('Send Storage URI Prompt'),
+                  ),
+                ],
+              ),
+            ),
+          ],
+        ),
+      ),
+    );
+  }
+
+  Future<void> _sendStorageUriPrompt(String message) async {
+    setState(() {
+      _loading = true;
+    });
+    try {
+      final content = [
+        Content.multi([
+          TextPart(message),
+          FileData(
+            'image/jpeg',
+            'gs://vertex-ai-example-ef5a2.appspot.com/foodpic.jpg',
+          ),
+        ]),
+      ];
+      _messages.add(MessageData(text: message, fromUser: true));
+
+      var response = await widget.model.generateContent(content);
+      var text = response.text;
+      _messages.add(MessageData(text: text, fromUser: false));
+
+      if (text == null) {
+        _showError('No response from API.');
+        return;
+      } else {
+        setState(() {
+          _loading = false;
+          _scrollDown();
+        });
+      }
+    } catch (e) {
+      _showError(e.toString());
+      setState(() {
+        _loading = false;
+      });
+    } finally {
+      _textController.clear();
+      setState(() {
+        _loading = false;
+      });
+      _textFieldFocus.requestFocus();
+    }
+  }
+
+  void _showError(String message) {
+    showDialog<void>(
+      context: context,
+      builder: (context) {
+        return AlertDialog(
+          title: const Text('Something went wrong'),
+          content: SingleChildScrollView(
+            child: SelectableText(message),
+          ),
+          actions: [
+            TextButton(
+              onPressed: () {
+                Navigator.of(context).pop();
+              },
+              child: const Text('OK'),
+            ),
+          ],
+        );
+      },
+    );
+  }
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/token_count_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/token_count_page.dart
new file mode 100644
index 000000000000..aa35483d5705
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/token_count_page.dart
@@ -0,0 +1,106 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'package:flutter/material.dart';
+import 'package:firebase_vertexai/firebase_vertexai.dart';
+import '../widgets/message_widget.dart';
+
+class TokenCountPage extends StatefulWidget {
+  const TokenCountPage({super.key, required this.title, required this.model});
+
+  final String title;
+  final GenerativeModel model;
+
+  @override
+  State<TokenCountPage> createState() => _TokenCountPageState();
+}
+
+class _TokenCountPageState extends State<TokenCountPage> {
+  final List<MessageData> _messages = <MessageData>[];
+  bool _loading = false;
+
+  @override
+  Widget build(BuildContext context) {
+    return Scaffold(
+      appBar: AppBar(
+        title: Text(widget.title),
+      ),
+      body: Padding(
+        padding: const EdgeInsets.all(8),
+        child: Column(
+          mainAxisAlignment: MainAxisAlignment.center,
+          crossAxisAlignment: CrossAxisAlignment.start,
+          children: [
+            Expanded(
+              child: ListView.builder(
+                itemBuilder: (context, idx) {
+                  return MessageWidget(
+                    text: _messages[idx].text,
+                    isFromUser: _messages[idx].fromUser ?? false,
+                  );
+                },
+                itemCount: _messages.length,
+              ),
+            ),
+            Padding(
+              padding: const EdgeInsets.symmetric(
+                vertical: 25,
+                horizontal: 15,
+              ),
+              child: Row(
+                children: [
+                  Expanded(
+                    child: ElevatedButton(
+                      onPressed: !_loading
+                          ? () async {
+                              await _testCountToken();
+                            }
+                          : null,
+                      child: const Text('Count Tokens'),
+                    ),
+                  ),
+                ],
+              ),
+            ),
+          ],
+        ),
+      ),
+    );
+  }
+
+  Future<void> _testCountToken() async {
+    setState(() {
+      _loading = true;
+    });
+
+    const prompt = 'tell a short story';
+    var content = Content.text(prompt);
+    var tokenResponse = await widget.model.countTokens([content]);
+    final tokenResult = 'Count token: ${tokenResponse.totalTokens}, billable '
+        'characters: ${tokenResponse.totalBillableCharacters}';
+    _messages.add(MessageData(text: tokenResult, fromUser: false));
+
+    var contentResponse = await widget.model.generateContent([content]);
+    final contentMetaData = 'result metadata, promptTokenCount:'
+        '${contentResponse.usageMetadata!.promptTokenCount}, '
+        'candidatesTokenCount:'
+        '${contentResponse.usageMetadata!.candidatesTokenCount}, '
+        'totalTokenCount:'
+        '${contentResponse.usageMetadata!.totalTokenCount}';
+    _messages.add(MessageData(text: contentMetaData, fromUser: false));
+    setState(() {
+      _loading = false;
+    });
+  }
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/widgets/message_widget.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/widgets/message_widget.dart
new file mode 100644
index 000000000000..b8a0f23ce03b
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/widgets/message_widget.dart
@@ -0,0 +1,68 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'package:flutter/material.dart';
+import 'package:flutter_markdown/flutter_markdown.dart';
+
+class MessageData {
+  MessageData({this.image, this.text, this.fromUser});
+  final Image? image;
+  final String? text;
+  final bool? fromUser;
+}
+
+class MessageWidget extends StatelessWidget {
+  final Image? image;
+  final String? text;
+  final bool isFromUser;
+
+  const MessageWidget({
+    super.key,
+    this.image,
+    this.text,
+    required this.isFromUser,
+  });
+
+  @override
+  Widget build(BuildContext context) {
+    return Row(
+      mainAxisAlignment:
+          isFromUser ? MainAxisAlignment.end : MainAxisAlignment.start,
+      children: [
+        Flexible(
+          child: Container(
+            constraints: const BoxConstraints(maxWidth: 600),
+            decoration: BoxDecoration(
+              color: isFromUser
+                  ? Theme.of(context).colorScheme.primaryContainer
+                  : Theme.of(context).colorScheme.surfaceContainerHighest,
+              borderRadius: BorderRadius.circular(18),
+            ),
+            padding: const EdgeInsets.symmetric(
+              vertical: 15,
+              horizontal: 20,
+            ),
+            margin: const EdgeInsets.only(bottom: 8),
+            child: Column(
+              children: [
+                if (text case final text?) MarkdownBody(data: text),
+                if (image case final image?) image,
+              ],
+            ),
+          ),
+        ),
+      ],
+    );
+  }
+}
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift b/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift
index d53ef6437726..b3c176141221 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift
+++ b/packages/firebase_vertexai/firebase_vertexai/example/macos/Runner/AppDelegate.swift
@@ -1,9 +1,13 @@
 import Cocoa
 import FlutterMacOS
 
-@NSApplicationMain
+@main
 class AppDelegate: FlutterAppDelegate {
   override func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool {
     return true
   }
+
+  override func applicationSupportsSecureRestorableState(_ app: NSApplication) -> Bool {
+    return true
+  }
 }

From 11e8876b73a97efccd765dbbcb079dd8406a1807 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Wed, 22 Jan 2025 23:27:33 -0800
Subject: [PATCH 07/21] fix analyzer

---
 .../firebase_vertexai/example/lib/main.dart           |  2 +-
 .../example/lib/pages/image_prompt_page.dart          | 11 +++++++----
 .../example/lib/pages/storage_uri_page.dart           |  7 +++++--
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
index 1c579d30c4bb..aa70b575c1dc 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -29,7 +29,7 @@ const FirebaseOptions? options = null;
 
 void main() async {
   WidgetsFlutterBinding.ensureInitialized();
-  await Firebase.initializeApp(options: options);
+  await Firebase.initializeApp();
   await FirebaseAuth.instance.signInAnonymously();
 
   var vertex_instance =
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
index c110c5205757..f8b111296287 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
@@ -123,13 +123,16 @@ class _ImagePromptPageState extends State<ImagePromptPage> {
       ];
       _generatedContent.add(
         MessageData(
-            image: Image.asset('assets/images/cat.jpg'),
-            text: message,
-            fromUser: true),
+          image: Image.asset('assets/images/cat.jpg'),
+          text: message,
+          fromUser: true,
+        ),
       );
       _generatedContent.add(
         MessageData(
-            image: Image.asset('assets/images/scones.jpg'), fromUser: true),
+          image: Image.asset('assets/images/scones.jpg'),
+          fromUser: true,
+        ),
       );
 
       var response = await widget.model.generateContent(content);
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
index eb23c1950d40..b1624d52efc6 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
@@ -17,8 +17,11 @@ import 'package:firebase_vertexai/firebase_vertexai.dart';
 import '../widgets/message_widget.dart';
 
 class StorageUriPromptPage extends StatefulWidget {
-  const StorageUriPromptPage(
-      {super.key, required this.title, required this.model});
+  const StorageUriPromptPage({
+    super.key,
+    required this.title,
+    required this.model,
+  });
 
   final String title;
   final GenerativeModel model;

From 8eefb6efc733804ba5a27c92e8faa4c414c1fc95 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Thu, 23 Jan 2025 21:13:54 -0800
Subject: [PATCH 08/21] imagen model working with new example layout

---
 .../firebase_vertexai/example/lib/main.dart   | 390 +-----------------
 .../example/lib/pages/image_prompt_page.dart  |  67 ++-
 ...storage_uri_page.dart => imagen_page.dart} | 114 ++---
 .../lib/firebase_vertexai.dart                |   2 +-
 4 files changed, 130 insertions(+), 443 deletions(-)
 rename packages/firebase_vertexai/firebase_vertexai/example/lib/pages/{storage_uri_page.dart => imagen_page.dart} (59%)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
index 120c1b656425..0813c334e14e 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -22,7 +22,7 @@ import 'pages/function_calling_page.dart';
 import 'pages/image_prompt_page.dart';
 import 'pages/token_count_page.dart';
 import 'pages/schema_page.dart';
-import 'pages/storage_uri_page.dart';
+import 'pages/imagen_page.dart';
 
 // REQUIRED if you want to run on Web
 const FirebaseOptions? options = null;
@@ -30,7 +30,7 @@ const FirebaseOptions? options = null;
 void main() async {
   WidgetsFlutterBinding.ensureInitialized();
   await Firebase.initializeApp();
-  await FirebaseAuth.instance.signInAnonymously();
+  //await FirebaseAuth.instance.signInAnonymously();
 
   var vertex_instance =
       FirebaseVertexAI.instanceFor(auth: FirebaseAuth.instance);
@@ -79,7 +79,7 @@ class _HomeScreenState extends State<HomeScreen> {
           title: 'Function Calling',
         ), // function calling will initial its own model
         ImagePromptPage(title: 'Image Prompt', model: widget.model),
-        StorageUriPromptPage(title: 'Storage URI Prompt', model: widget.model),
+        ImagenPage(title: 'Imagen Model', model: widget.model),
         SchemaPromptPage(title: 'Schema Prompt', model: widget.model),
       ];
 
@@ -134,11 +134,11 @@ class _HomeScreenState extends State<HomeScreen> {
           ),
           BottomNavigationBarItem(
             icon: Icon(
-              Icons.folder,
+              Icons.image_search,
               color: Theme.of(context).colorScheme.primary,
             ),
-            label: 'Storage URI Prompt',
-            tooltip: 'Storage URI Prompt',
+            label: 'Imagen Model',
+            tooltip: 'Imagen Model',
           ),
           BottomNavigationBarItem(
             icon: Icon(
@@ -154,382 +154,4 @@ class _HomeScreenState extends State<HomeScreen> {
       ),
     );
   }
-
-  Future<void> _promptSchemaTest(String subject) async {
-    setState(() {
-      _loading = true;
-    });
-    try {
-      final content = [
-        Content.text(
-          "For use in a children's card game, generate 10 animal-based "
-          'characters.',
-        ),
-      ];
-
-      final jsonSchema = Schema.object(
-        properties: {
-          'characters': Schema.array(
-            items: Schema.object(
-              properties: {
-                'name': Schema.string(),
-                'age': Schema.integer(),
-                'species': Schema.string(),
-                'accessory':
-                    Schema.enumString(enumValues: ['hat', 'belt', 'shoes']),
-              },
-            ),
-          ),
-        },
-        optionalProperties: ['accessory'],
-      );
-
-      final response = await _model.generateContent(
-        content,
-        generationConfig: GenerationConfig(
-          responseMimeType: 'application/json',
-          responseSchema: jsonSchema,
-        ),
-      );
-
-      var text = response.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _sendStorageUriPrompt(String message) async {
-    setState(() {
-      _loading = true;
-    });
-    try {
-      final content = [
-        Content.multi([
-          TextPart(message),
-          FileData(
-            'image/jpeg',
-            'gs://vertex-ai-example-ef5a2.appspot.com/foodpic.jpg',
-          ),
-        ]),
-      ];
-      _generatedContent.add((image: null, text: message, fromUser: true));
-
-      var response = await _model.generateContent(content);
-      var text = response.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _sendImagePrompt(String message) async {
-    setState(() {
-      _loading = true;
-    });
-    try {
-      ByteData catBytes = await rootBundle.load('assets/images/cat.jpg');
-      ByteData sconeBytes = await rootBundle.load('assets/images/scones.jpg');
-      final content = [
-        Content.multi([
-          TextPart(message),
-          // The only accepted mime types are image/*.
-          InlineDataPart('image/jpeg', catBytes.buffer.asUint8List()),
-          InlineDataPart('image/jpeg', sconeBytes.buffer.asUint8List()),
-        ]),
-      ];
-      _generatedContent.add(
-        (
-          image: Image.asset('assets/images/cat.jpg'),
-          text: message,
-          fromUser: true
-        ),
-      );
-      _generatedContent.add(
-        (
-          image: Image.asset('assets/images/scones.jpg'),
-          text: null,
-          fromUser: true
-        ),
-      );
-
-      var response = await _model.generateContent(content);
-      var text = response.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _sendChatMessage(String message) async {
-    setState(() {
-      _loading = true;
-    });
-
-    try {
-      _generatedContent.add((image: null, text: message, fromUser: true));
-      var response = await _chat?.sendMessage(
-        Content.text(message),
-      );
-      var text = response?.text;
-      _generatedContent.add((image: null, text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  Future<void> _testFunctionCalling() async {
-    setState(() {
-      _loading = true;
-    });
-    final functionCallChat = _functionCallModel.startChat();
-    const prompt = 'What is the weather like in Boston on 10/02 this year?';
-
-    // Send the message to the generative model.
-    var response = await functionCallChat.sendMessage(
-      Content.text(prompt),
-    );
-
-    final functionCalls = response.functionCalls.toList();
-    // When the model response with a function call, invoke the function.
-    if (functionCalls.isNotEmpty) {
-      final functionCall = functionCalls.first;
-      if (functionCall.name == 'fetchWeather') {
-        Map<String, dynamic> location =
-            functionCall.args['location']! as Map<String, dynamic>;
-        var date = functionCall.args['date']! as String;
-        var city = location['city'] as String;
-        var state = location['state'] as String;
-        final functionResult = await fetchWeather(Location(city, state), date);
-        // Send the response to the model so that it can use the result to
-        // generate text for the user.
-        response = await functionCallChat.sendMessage(
-          Content.functionResponse(functionCall.name, functionResult),
-        );
-      } else {
-        throw UnimplementedError(
-          'Function not declared to the model: ${functionCall.name}',
-        );
-      }
-    }
-    // When the model responds with non-null text content, print it.
-    if (response.text case final text?) {
-      _generatedContent.add((image: null, text: text, fromUser: false));
-      setState(() {
-        _loading = false;
-      });
-    }
-  }
-
-  Future<void> _testCountToken() async {
-    setState(() {
-      _loading = true;
-    });
-
-    const prompt = 'tell a short story';
-    var content = Content.text(prompt);
-    var tokenResponse = await _model.countTokens([content]);
-    final tokenResult = 'Count token: ${tokenResponse.totalTokens}, billable '
-        'characters: ${tokenResponse.totalBillableCharacters}';
-    _generatedContent.add((image: null, text: tokenResult, fromUser: false));
-
-    var contentResponse = await _model.generateContent([content]);
-    final contentMetaData = 'result metadata, promptTokenCount:'
-        '${contentResponse.usageMetadata!.promptTokenCount}, '
-        'candidatesTokenCount:'
-        '${contentResponse.usageMetadata!.candidatesTokenCount}, '
-        'totalTokenCount:'
-        '${contentResponse.usageMetadata!.totalTokenCount}';
-    _generatedContent
-        .add((image: null, text: contentMetaData, fromUser: false));
-    setState(() {
-      _loading = false;
-    });
-  }
-
-  Future<void> _testImagen(String prompt) async {
-    setState(() {
-      _loading = true;
-    });
-    var model = FirebaseVertexAI.instance.imageModel(
-      modelName: 'imagen-3.0-generate-001',
-      generationConfig: ImagenGenerationConfig(
-          imageFormat: ImagenFormat.jpeg(compressionQuality: 75)),
-      safetySettings: ImagenSafetySettings(
-        ImagenSafetyFilterLevel.blockLowAndAbove,
-        ImagenPersonFilterLevel.allowAdult,
-      ),
-    );
-
-    var generationConfig = ImagenGenerationConfig(
-        negativePrompt: 'frog',
-        numberOfImages: 1,
-        aspectRatio: ImagenAspectRatio.square1x1);
-
-    var response = await model.generateImages(
-      prompt,
-      generationConfig: generationConfig,
-    );
-
-    if (response.images.isNotEmpty) {
-      var imagenImage = response.images[0];
-      // Process the image
-      _generatedContent.add(
-        (
-          image: Image.memory(imagenImage.bytesBase64Encoded),
-          text: prompt,
-          fromUser: false
-        ),
-      );
-    } else {
-      // Handle the case where no images were generated
-      print('Error: No images were generated.');
-    }
-    setState(() {
-      _loading = false;
-    });
-  }
-
-  void _showError(String message) {
-    showDialog<void>(
-      context: context,
-      builder: (context) {
-        return AlertDialog(
-          title: const Text('Something went wrong'),
-          content: SingleChildScrollView(
-            child: SelectableText(message),
-          ),
-          actions: [
-            TextButton(
-              onPressed: () {
-                Navigator.of(context).pop();
-              },
-              child: const Text('OK'),
-            ),
-          ],
-        );
-      },
-    );
-  }
-}
-
-class MessageWidget extends StatelessWidget {
-  final Image? image;
-  final String? text;
-  final bool isFromUser;
-
-  const MessageWidget({
-    super.key,
-    this.image,
-    this.text,
-    required this.isFromUser,
-  });
-
-  @override
-  Widget build(BuildContext context) {
-    return Row(
-      mainAxisAlignment:
-          isFromUser ? MainAxisAlignment.end : MainAxisAlignment.start,
-      children: [
-        Flexible(
-          child: Container(
-            constraints: const BoxConstraints(maxWidth: 600),
-            decoration: BoxDecoration(
-              color: isFromUser
-                  ? Theme.of(context).colorScheme.primaryContainer
-                  : Theme.of(context).colorScheme.surfaceContainerHighest,
-              borderRadius: BorderRadius.circular(18),
-            ),
-            padding: const EdgeInsets.symmetric(
-              vertical: 15,
-              horizontal: 20,
-            ),
-            margin: const EdgeInsets.only(bottom: 8),
-            child: Column(
-              children: [
-                if (text case final text?) MarkdownBody(data: text),
-                if (image case final image?) image,
-              ],
-            ),
-          ),
-        ),
-      ],
-    );
-  }
 }
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
index f8b111296287..1a32cd370d84 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
@@ -89,13 +89,23 @@ class _ImagePromptPageState extends State<ImagePromptPage> {
                   const SizedBox.square(
                     dimension: 15,
                   ),
-                  ElevatedButton(
-                    onPressed: !_loading
-                        ? () async {
-                            await _sendImagePrompt(_textController.text);
-                          }
-                        : null,
-                    child: const Text('Send Image Prompt'),
+                  IconButton(
+                    onPressed: () async {
+                      await _sendImagePrompt(_textController.text);
+                    },
+                    icon: Icon(
+                      Icons.image,
+                      color: Theme.of(context).colorScheme.primary,
+                    ),
+                  ),
+                  IconButton(
+                    onPressed: () async {
+                      await _sendStorageUriPrompt(_textController.text);
+                    },
+                    icon: Icon(
+                      Icons.storage,
+                      color: Theme.of(context).colorScheme.primary,
+                    ),
                   ),
                 ],
               ),
@@ -162,6 +172,49 @@ class _ImagePromptPageState extends State<ImagePromptPage> {
     }
   }
 
+  Future<void> _sendStorageUriPrompt(String message) async {
+    setState(() {
+      _loading = true;
+    });
+    try {
+      final content = [
+        Content.multi([
+          TextPart(message),
+          FileData(
+            'image/jpeg',
+            'gs://vertex-ai-example-ef5a2.appspot.com/foodpic.jpg',
+          ),
+        ]),
+      ];
+      _generatedContent.add(MessageData(text: message, fromUser: true));
+
+      var response = await widget.model.generateContent(content);
+      var text = response.text;
+      _generatedContent.add(MessageData(text: text, fromUser: false));
+
+      if (text == null) {
+        _showError('No response from API.');
+        return;
+      } else {
+        setState(() {
+          _loading = false;
+          _scrollDown();
+        });
+      }
+    } catch (e) {
+      _showError(e.toString());
+      setState(() {
+        _loading = false;
+      });
+    } finally {
+      _textController.clear();
+      setState(() {
+        _loading = false;
+      });
+      _textFieldFocus.requestFocus();
+    }
+  }
+
   void _showError(String message) {
     showDialog<void>(
       context: context,
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
similarity index 59%
rename from packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
rename to packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
index b1624d52efc6..6e2dbb497241 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
@@ -16,8 +16,8 @@ import 'package:flutter/material.dart';
 import 'package:firebase_vertexai/firebase_vertexai.dart';
 import '../widgets/message_widget.dart';
 
-class StorageUriPromptPage extends StatefulWidget {
-  const StorageUriPromptPage({
+class ImagenPage extends StatefulWidget {
+  const ImagenPage({
     super.key,
     required this.title,
     required this.model,
@@ -27,15 +27,31 @@ class StorageUriPromptPage extends StatefulWidget {
   final GenerativeModel model;
 
   @override
-  State<StorageUriPromptPage> createState() => _StorageUriPromptPageState();
+  State<ImagenPage> createState() => _ImagenPageState();
 }
 
-class _StorageUriPromptPageState extends State<StorageUriPromptPage> {
+class _ImagenPageState extends State<ImagenPage> {
   final ScrollController _scrollController = ScrollController();
   final TextEditingController _textController = TextEditingController();
   final FocusNode _textFieldFocus = FocusNode();
-  final List<MessageData> _messages = <MessageData>[];
+  final List<MessageData> _generatedContent = <MessageData>[];
   bool _loading = false;
+  late final ImagenModel _imagenModel;
+
+  @override
+  void initState() {
+    super.initState();
+    _imagenModel = FirebaseVertexAI.instance.imageModel(
+      modelName: 'imagen-3.0-generate-001',
+      generationConfig: ImagenGenerationConfig(
+        imageFormat: ImagenFormat.jpeg(compressionQuality: 75),
+      ),
+      safetySettings: ImagenSafetySettings(
+        ImagenSafetyFilterLevel.blockLowAndAbove,
+        ImagenPersonFilterLevel.allowAdult,
+      ),
+    );
+  }
 
   void _scrollDown() {
     WidgetsBinding.instance.addPostFrameCallback(
@@ -66,11 +82,12 @@ class _StorageUriPromptPageState extends State<StorageUriPromptPage> {
                 controller: _scrollController,
                 itemBuilder: (context, idx) {
                   return MessageWidget(
-                    text: _messages[idx].text,
-                    isFromUser: _messages[idx].fromUser ?? false,
+                    text: _generatedContent[idx].text,
+                    image: _generatedContent[idx].image,
+                    isFromUser: _generatedContent[idx].fromUser ?? false,
                   );
                 },
-                itemCount: _messages.length,
+                itemCount: _generatedContent.length,
               ),
             ),
             Padding(
@@ -90,14 +107,18 @@ class _StorageUriPromptPageState extends State<StorageUriPromptPage> {
                   const SizedBox.square(
                     dimension: 15,
                   ),
-                  ElevatedButton(
-                    onPressed: !_loading
-                        ? () async {
-                            await _sendStorageUriPrompt(_textController.text);
-                          }
-                        : null,
-                    child: const Text('Send Storage URI Prompt'),
-                  ),
+                  if (!_loading)
+                    IconButton(
+                      onPressed: () async {
+                        await _testImagen(_textController.text);
+                      },
+                      icon: Icon(
+                        Icons.image_search,
+                        color: Theme.of(context).colorScheme.primary,
+                      ),
+                    )
+                  else
+                    const CircularProgressIndicator(),
                 ],
               ),
             ),
@@ -107,47 +128,38 @@ class _StorageUriPromptPageState extends State<StorageUriPromptPage> {
     );
   }
 
-  Future<void> _sendStorageUriPrompt(String message) async {
+  Future<void> _testImagen(String prompt) async {
     setState(() {
       _loading = true;
     });
-    try {
-      final content = [
-        Content.multi([
-          TextPart(message),
-          FileData(
-            'image/jpeg',
-            'gs://vertex-ai-example-ef5a2.appspot.com/foodpic.jpg',
-          ),
-        ]),
-      ];
-      _messages.add(MessageData(text: message, fromUser: true));
 
-      var response = await widget.model.generateContent(content);
-      var text = response.text;
-      _messages.add(MessageData(text: text, fromUser: false));
+    var generationConfig = ImagenGenerationConfig(
+        negativePrompt: 'frog',
+        numberOfImages: 1,
+        aspectRatio: ImagenAspectRatio.square1x1);
+
+    var response = await _imagenModel.generateImages(
+      prompt,
+      generationConfig: generationConfig,
+    );
 
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
+    if (response.images.isNotEmpty) {
+      var imagenImage = response.images[0];
+      // Process the image
+      _generatedContent.add(
+        MessageData(
+          image: Image.memory(imagenImage.bytesBase64Encoded),
+          text: prompt,
+          fromUser: false,
+        ),
+      );
+    } else {
+      // Handle the case where no images were generated
+      _showError('Error: No images were generated.');
     }
+    setState(() {
+      _loading = false;
+    });
   }
 
   void _showError(String message) {
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
index 19bf7ed37043..af79eed8efd1 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
@@ -58,7 +58,6 @@ export 'src/function_calling.dart'
         ToolConfig;
 export 'src/imagen_api.dart'
     show
-        ImagenModelConfig,
         ImagenSafetySettings,
         ImagenFormat,
         ImagenSafetyFilterLevel,
@@ -67,5 +66,6 @@ export 'src/imagen_api.dart'
         ImagenAspectRatio;
 export 'src/imagen_content.dart'
     show ImagenInlineImage, ImagenGCSImage, ImagenImage;
+export 'src/imagen_model.dart' show ImagenModel;
 export 'src/model.dart' show GenerativeModel;
 export 'src/schema.dart' show Schema, SchemaType;

From 8922f9b5941df58cb5fbcd152ddd35bcebed8cc5 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Thu, 23 Jan 2025 22:41:01 -0800
Subject: [PATCH 09/21] Finish up all functionality

---
 .../example/lib/pages/imagen_page.dart        | 20 ++---
 .../firebase_vertexai/lib/src/base_model.dart |  5 +-
 .../lib/src/firebase_vertexai.dart            |  4 +
 .../firebase_vertexai/lib/src/imagen_api.dart | 75 +++++++++++++------
 .../lib/src/imagen_content.dart               | 47 +++++++-----
 .../lib/src/imagen_model.dart                 | 59 +++++++--------
 .../firebase_vertexai/lib/src/model.dart      |  4 +-
 7 files changed, 124 insertions(+), 90 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
index 6e2dbb497241..e321976dd288 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
@@ -41,11 +41,15 @@ class _ImagenPageState extends State<ImagenPage> {
   @override
   void initState() {
     super.initState();
+    var generationConfig = ImagenGenerationConfig(
+      negativePrompt: 'frog',
+      numberOfImages: 1,
+      aspectRatio: ImagenAspectRatio.square1x1,
+      imageFormat: ImagenFormat.jpeg(compressionQuality: 75),
+    );
     _imagenModel = FirebaseVertexAI.instance.imageModel(
       modelName: 'imagen-3.0-generate-001',
-      generationConfig: ImagenGenerationConfig(
-        imageFormat: ImagenFormat.jpeg(compressionQuality: 75),
-      ),
+      generationConfig: generationConfig,
       safetySettings: ImagenSafetySettings(
         ImagenSafetyFilterLevel.blockLowAndAbove,
         ImagenPersonFilterLevel.allowAdult,
@@ -133,15 +137,7 @@ class _ImagenPageState extends State<ImagenPage> {
       _loading = true;
     });
 
-    var generationConfig = ImagenGenerationConfig(
-        negativePrompt: 'frog',
-        numberOfImages: 1,
-        aspectRatio: ImagenAspectRatio.square1x1);
-
-    var response = await _imagenModel.generateImages(
-      prompt,
-      generationConfig: generationConfig,
-    );
+    var response = await _imagenModel.generateImages(prompt);
 
     if (response.images.isNotEmpty) {
       var imagenImage = response.images[0];
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
index 6c0a8bf51342..805356d1909b 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
@@ -42,8 +42,6 @@ abstract class BaseModel {
     required FirebaseApp app,
     required ApiClient client,
   })  : _model = normalizeModelName(model),
-        _app = app,
-        _location = location,
         _projectUri = _vertexUri(app, location),
         _client = client;
 
@@ -51,8 +49,7 @@ abstract class BaseModel {
   static const _apiVersion = 'v1beta';
 
   final ({String prefix, String name}) _model;
-  final FirebaseApp _app;
-  final String _location;
+
   final Uri _projectUri;
   final ApiClient _client;
 
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
index 100036be6f90..19d37a947432 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
@@ -114,6 +114,10 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
     );
   }
 
+  /// Create a [ImagenModel].
+  ///
+  /// The optional [safetySettings] can be used to control and guide the
+  /// generation. See [ImagenSafetySettings] for details.
   ImagenModel imageModel(
       {required String modelName,
       ImagenGenerationConfig? generationConfig,
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
index 846dc7cd07dc..735eab7217d9 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
@@ -12,18 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-///
+/// Specifies the level of safety filtering for image generation.
+/// If not specified, default will be "block_medium_and_above".
 enum ImagenSafetyFilterLevel {
-  ///
+  /// Strongest filtering level, most strict blocking.
   blockLowAndAbove('block_low_and_above'),
 
-  ///
+  /// Block some problematic prompts and responses.
   blockMediumAndAbove('block_medium_and_above'),
 
-  ///
+  /// Reduces the number of requests blocked due to safety filters.
+  /// May increase objectionable content generated by Imagen.
   blockOnlyHigh('block_only_high'),
 
-  ///
+  /// Block very few problematic prompts and responses.
+  /// Access to this feature is restricted.
   blockNone('block_none');
 
   const ImagenSafetyFilterLevel(this._jsonString);
@@ -49,15 +52,16 @@ enum ImagenSafetyFilterLevel {
   String toString() => name;
 }
 
-///
+/// Allow generation of people by the model.
+/// If not specified, the default value is "allow_adult".
 enum ImagenPersonFilterLevel {
-  ///
+  /// Disallow the inclusion of people or faces in images.
   blockAll('dont_allow'),
 
-  ///
+  /// Allow generation of adults only.
   allowAdult('allow_adult'),
 
-  ///
+  /// Allow generation of people of all ages.
   allowAll('allow_all');
 
   const ImagenPersonFilterLevel(this._jsonString);
@@ -82,15 +86,17 @@ enum ImagenPersonFilterLevel {
   String toString() => name;
 }
 
+/// A class representing safety settings for image generation.
 ///
+/// It includes a safety filter level and a person filter level.
 final class ImagenSafetySettings {
   /// Constructor
   ImagenSafetySettings(this.safetyFilterLevel, this.personFilterLevel);
 
-  ///
+  /// The safety filter level
   final ImagenSafetyFilterLevel? safetyFilterLevel;
 
-  ///
+  /// The person filter level
   final ImagenPersonFilterLevel? personFilterLevel;
 
   /// Convert to json format.
@@ -102,21 +108,21 @@ final class ImagenSafetySettings {
       };
 }
 
-///
+/// The aspect ratio for the image. The default value is "1:1".
 enum ImagenAspectRatio {
-  ///
+  /// Square (1:1).
   square1x1('1:1'),
 
-  ///
+  /// Portrait (9:16).
   portrait9x16('9:16'),
 
-  ///
+  /// Landscape (16:9).
   landscape16x9('16:9'),
 
-  ///
+  /// Portrait (3:4).
   portrait3x4('3:4'),
 
-  ///
+  /// Landscape (4:3).
   landscape4x3('4:3');
 
   const ImagenAspectRatio(this._jsonString);
@@ -143,17 +149,31 @@ enum ImagenAspectRatio {
   String toString() => name;
 }
 
+/// Configuration options for image generation.
 final class ImagenGenerationConfig {
+  /// Constructor
   ImagenGenerationConfig(
-      {this.negativePrompt,
-      this.numberOfImages,
+      {this.numberOfImages,
+      this.negativePrompt,
       this.aspectRatio,
       this.imageFormat,
       this.addWatermark});
-  final String? negativePrompt;
+
+  /// The number of images to generate. Default is 1.
   final int? numberOfImages;
+
+  /// A description of what to discourage in the generated images.
+  final String? negativePrompt;
+
+  /// The aspect ratio for the image. The default value is "1:1".
   final ImagenAspectRatio? aspectRatio;
+
+  /// The image format of the generated images.
   final ImagenFormat? imageFormat;
+
+  /// Add an invisible watermark to the generated images.
+  /// Default value for each imagen model can be found in
+  /// https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api#generate_images
   final bool? addWatermark;
 
   /// Convert to json format.
@@ -166,15 +186,26 @@ final class ImagenGenerationConfig {
       };
 }
 
+/// Represents the image format and compression quality.
 final class ImagenFormat {
+  /// Constructor
   ImagenFormat(this.mimeType, this.compressionQuality);
 
-  ImagenFormat.png() : this("image/png", null);
+  /// Constructor for png
+  ImagenFormat.png() : this('image/png', null);
+
+  /// Constructor for jpeg
   ImagenFormat.jpeg({int? compressionQuality})
-      : this("image/jpeg", compressionQuality);
+      : this('image/jpeg', compressionQuality);
+
+  /// The MIME type of the image format. The default value is "image/png".
   final String mimeType;
+
+  /// The level of compression if the output type is "image/jpeg".
+  /// Accepted values are 0 through 100. The default value is 75.
   final int? compressionQuality;
 
+  /// Convert to json format.
   Map<String, dynamic> toJson() => {
         'mimeType': mimeType,
         if (compressionQuality != null)
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
index 49bc019ffae7..f959d8863a20 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
@@ -15,28 +15,32 @@ import 'dart:convert';
 import 'dart:typed_data';
 import 'error.dart';
 
-///
+/// Base type of Imagen Image.
 sealed class ImagenImage {
+  /// Constructor
+  ImagenImage({required this.mimeType});
+
+  /// The MIME type of the image format.
   final String mimeType;
 
   /// Convert the [ImagenImage] content to json format.
   Object toJson();
-
-  ImagenImage({required this.mimeType});
 }
 
+/// Represents an image stored as a base64-encoded string.
 final class ImagenInlineImage implements ImagenImage {
-  /// Data contents in bytes.
-  final Uint8List bytesBase64Encoded;
-
-  @override
-  final String mimeType;
-
+  /// Constructor
   ImagenInlineImage({
     required this.bytesBase64Encoded,
     required this.mimeType,
   });
 
+  /// The data contents in bytes, encoded as base64.
+  final Uint8List bytesBase64Encoded;
+
+  @override
+  final String mimeType;
+
   @override
   Object toJson() => {
         'mimeType': mimeType,
@@ -44,17 +48,20 @@ final class ImagenInlineImage implements ImagenImage {
       };
 }
 
+/// Represents an image stored in Google Cloud Storage.
 final class ImagenGCSImage implements ImagenImage {
-  @override
-  final String mimeType;
-
-  final String gcsUri;
-
+  /// Constructor
   ImagenGCSImage({
     required this.gcsUri,
     required this.mimeType,
   });
 
+  /// The storage URI of the image.
+  final String gcsUri;
+
+  @override
+  final String mimeType;
+
   @override
   Object toJson() => {
         'mimeType': mimeType,
@@ -62,15 +69,15 @@ final class ImagenGCSImage implements ImagenImage {
       };
 }
 
+/// Represents the response from an image generation request.
 final class ImagenGenerationResponse<T extends ImagenImage> {
+  /// Constructor
   ImagenGenerationResponse({
     required this.images,
     this.filteredReason,
   });
 
-  final List<T> images;
-  final String? filteredReason;
-
+  /// Factory method to create an ImagenGenerationResponse from a JSON object.
   factory ImagenGenerationResponse.fromJson(Map<String, dynamic> json) {
     final filteredReason = json['filteredReason'] as String?;
     final imagesJson = json['predictions'] as List<dynamic>;
@@ -102,6 +109,12 @@ final class ImagenGenerationResponse<T extends ImagenImage> {
       throw ArgumentError('Unsupported ImagenImage type: $T');
     }
   }
+
+  /// A list of generated images. The type of the images depends on the T parameter.
+  final List<T> images;
+
+  /// If the generation was filtered due to safety reasons, a message explaining the reason.
+  final String? filteredReason;
 }
 
 /// Parse the json to [ImagenGenerationResponse]
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index 2eca0c77c8e7..9d65fff160da 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -11,15 +11,15 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-import 'imagen_api.dart';
-import 'imagen_content.dart';
-import 'base_model.dart';
-import 'client.dart';
-
 import 'package:firebase_app_check/firebase_app_check.dart';
 import 'package:firebase_auth/firebase_auth.dart';
 import 'package:firebase_core/firebase_core.dart';
 
+import 'base_model.dart';
+import 'client.dart';
+import 'imagen_api.dart';
+import 'imagen_content.dart';
+
 ///
 final class ImagenModel extends BaseModel {
   ImagenModel._(
@@ -45,31 +45,26 @@ final class ImagenModel extends BaseModel {
 
   Map<String, Object?> _generateImagenRequest(
     String prompt, {
-    ImagenGenerationConfig? generationConfig,
     String? gcsUri,
   }) {
     final parameters = <String, Object?>{};
 
     if (gcsUri != null) parameters['storageUri'] = gcsUri;
 
-    // Merge generation configurations
-    final mergedConfig = {
-      ...(_generationConfig?.toJson() ?? {}),
-      ...(generationConfig?.toJson() ?? {}),
-    };
-
-    parameters['sampleCount'] = mergedConfig['numberOfImages'] ?? 1;
-    if (mergedConfig['aspectRatio'] != null) {
-      parameters['aspectRatio'] = mergedConfig['aspectRatio'];
-    }
-    if (mergedConfig['negativePrompt'] != null) {
-      parameters['negativePrompt'] = mergedConfig['negativePrompt'];
-    }
-    if (mergedConfig['addWatermark'] != null) {
-      parameters['addWatermark'] = mergedConfig['addWatermark'];
-    }
-    if (mergedConfig['outputOption'] != null) {
-      parameters['outputOption'] = mergedConfig['outputOption'];
+    parameters['sampleCount'] = _generationConfig?.numberOfImages ?? 1;
+    if (_generationConfig != null) {
+      if (_generationConfig.aspectRatio != null) {
+        parameters['aspectRatio'] = _generationConfig.aspectRatio;
+      }
+      if (_generationConfig.negativePrompt != null) {
+        parameters['negativePrompt'] = _generationConfig.negativePrompt;
+      }
+      if (_generationConfig.addWatermark != null) {
+        parameters['addWatermark'] = _generationConfig.addWatermark;
+      }
+      if (_generationConfig.imageFormat != null) {
+        parameters['outputOption'] = _generationConfig.imageFormat!.toJson();
+      }
     }
 
     if (_safetySettings != null) {
@@ -91,30 +86,30 @@ final class ImagenModel extends BaseModel {
     };
   }
 
+  /// Generates images with format of [ImagenInlineImage] based on the given
+  /// prompt.
   Future<ImagenGenerationResponse<ImagenInlineImage>> generateImages(
-    String prompt, {
-    ImagenGenerationConfig? generationConfig,
-  }) =>
+    String prompt,
+  ) =>
       makeRequest(
         Task.predict,
         _generateImagenRequest(
           prompt,
-          generationConfig: generationConfig,
         ),
         (jsonObject) =>
             parseImagenGenerationResponse<ImagenInlineImage>(jsonObject),
       );
 
+  /// Generates images with format of [ImagenGCSImage] based on the given
+  /// prompt.
   Future<ImagenGenerationResponse<ImagenGCSImage>> generateImagesGCS(
     String prompt,
-    String gcsUri, {
-    ImagenGenerationConfig? generationConfig,
-  }) =>
+    String gcsUri,
+  ) =>
       makeRequest(
         Task.predict,
         _generateImagenRequest(
           prompt,
-          generationConfig: generationConfig,
           gcsUri: gcsUri,
         ),
         (jsonObject) =>
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart
index 91a050cc2922..ddf0d1e5a37e 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart
@@ -19,15 +19,13 @@ import 'dart:async';
 import 'package:firebase_app_check/firebase_app_check.dart';
 import 'package:firebase_auth/firebase_auth.dart';
 import 'package:firebase_core/firebase_core.dart';
-
 import 'package:http/http.dart' as http;
 
 import 'api.dart';
+import 'base_model.dart';
 import 'client.dart';
 import 'content.dart';
 import 'function_calling.dart';
-import 'vertex_version.dart';
-import 'base_model.dart';
 
 /// A multimodel generative model (like Gemini).
 ///

From 9b8fda9b17c7d4eee7cfbd42ba7823eda651fed8 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Fri, 24 Jan 2025 11:21:59 -0800
Subject: [PATCH 10/21] remove the storage uri prompt page after merge

---
 .../example/lib/pages/storage_uri_page.dart   | 174 ------------------
 1 file changed, 174 deletions(-)
 delete mode 100644 packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
deleted file mode 100644
index b1624d52efc6..000000000000
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/storage_uri_page.dart
+++ /dev/null
@@ -1,174 +0,0 @@
-// Copyright 2025 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-import 'package:flutter/material.dart';
-import 'package:firebase_vertexai/firebase_vertexai.dart';
-import '../widgets/message_widget.dart';
-
-class StorageUriPromptPage extends StatefulWidget {
-  const StorageUriPromptPage({
-    super.key,
-    required this.title,
-    required this.model,
-  });
-
-  final String title;
-  final GenerativeModel model;
-
-  @override
-  State<StorageUriPromptPage> createState() => _StorageUriPromptPageState();
-}
-
-class _StorageUriPromptPageState extends State<StorageUriPromptPage> {
-  final ScrollController _scrollController = ScrollController();
-  final TextEditingController _textController = TextEditingController();
-  final FocusNode _textFieldFocus = FocusNode();
-  final List<MessageData> _messages = <MessageData>[];
-  bool _loading = false;
-
-  void _scrollDown() {
-    WidgetsBinding.instance.addPostFrameCallback(
-      (_) => _scrollController.animateTo(
-        _scrollController.position.maxScrollExtent,
-        duration: const Duration(
-          milliseconds: 750,
-        ),
-        curve: Curves.easeOutCirc,
-      ),
-    );
-  }
-
-  @override
-  Widget build(BuildContext context) {
-    return Scaffold(
-      appBar: AppBar(
-        title: Text(widget.title),
-      ),
-      body: Padding(
-        padding: const EdgeInsets.all(8),
-        child: Column(
-          mainAxisAlignment: MainAxisAlignment.center,
-          crossAxisAlignment: CrossAxisAlignment.start,
-          children: [
-            Expanded(
-              child: ListView.builder(
-                controller: _scrollController,
-                itemBuilder: (context, idx) {
-                  return MessageWidget(
-                    text: _messages[idx].text,
-                    isFromUser: _messages[idx].fromUser ?? false,
-                  );
-                },
-                itemCount: _messages.length,
-              ),
-            ),
-            Padding(
-              padding: const EdgeInsets.symmetric(
-                vertical: 25,
-                horizontal: 15,
-              ),
-              child: Row(
-                children: [
-                  Expanded(
-                    child: TextField(
-                      autofocus: true,
-                      focusNode: _textFieldFocus,
-                      controller: _textController,
-                    ),
-                  ),
-                  const SizedBox.square(
-                    dimension: 15,
-                  ),
-                  ElevatedButton(
-                    onPressed: !_loading
-                        ? () async {
-                            await _sendStorageUriPrompt(_textController.text);
-                          }
-                        : null,
-                    child: const Text('Send Storage URI Prompt'),
-                  ),
-                ],
-              ),
-            ),
-          ],
-        ),
-      ),
-    );
-  }
-
-  Future<void> _sendStorageUriPrompt(String message) async {
-    setState(() {
-      _loading = true;
-    });
-    try {
-      final content = [
-        Content.multi([
-          TextPart(message),
-          FileData(
-            'image/jpeg',
-            'gs://vertex-ai-example-ef5a2.appspot.com/foodpic.jpg',
-          ),
-        ]),
-      ];
-      _messages.add(MessageData(text: message, fromUser: true));
-
-      var response = await widget.model.generateContent(content);
-      var text = response.text;
-      _messages.add(MessageData(text: text, fromUser: false));
-
-      if (text == null) {
-        _showError('No response from API.');
-        return;
-      } else {
-        setState(() {
-          _loading = false;
-          _scrollDown();
-        });
-      }
-    } catch (e) {
-      _showError(e.toString());
-      setState(() {
-        _loading = false;
-      });
-    } finally {
-      _textController.clear();
-      setState(() {
-        _loading = false;
-      });
-      _textFieldFocus.requestFocus();
-    }
-  }
-
-  void _showError(String message) {
-    showDialog<void>(
-      context: context,
-      builder: (context) {
-        return AlertDialog(
-          title: const Text('Something went wrong'),
-          content: SingleChildScrollView(
-            child: SelectableText(message),
-          ),
-          actions: [
-            TextButton(
-              onPressed: () {
-                Navigator.of(context).pop();
-              },
-              child: const Text('OK'),
-            ),
-          ],
-        );
-      },
-    );
-  }
-}

From d95abfeafb19273bc09613fd02ed3b5c77b2de79 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Mon, 3 Feb 2025 21:33:53 -0800
Subject: [PATCH 11/21] Add example for gcs generation

---
 .../example/lib/pages/imagen_page.dart        | 45 +++++++++++++++++++
 .../firebase_vertexai/example/pubspec.yaml    |  1 +
 2 files changed, 46 insertions(+)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
index e321976dd288..9fa405837648 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
@@ -14,6 +14,7 @@
 
 import 'package:flutter/material.dart';
 import 'package:firebase_vertexai/firebase_vertexai.dart';
+import 'package:firebase_storage/firebase_storage.dart';
 import '../widgets/message_widget.dart';
 
 class ImagenPage extends StatefulWidget {
@@ -120,6 +121,20 @@ class _ImagenPageState extends State<ImagenPage> {
                         Icons.image_search,
                         color: Theme.of(context).colorScheme.primary,
                       ),
+                      tooltip: 'Imagen raw data',
+                    )
+                  else
+                    const CircularProgressIndicator(),
+                  if (!_loading)
+                    IconButton(
+                      onPressed: () async {
+                        await _testImagenGCS(_textController.text);
+                      },
+                      icon: Icon(
+                        Icons.imagesearch_roller,
+                        color: Theme.of(context).colorScheme.primary,
+                      ),
+                      tooltip: 'Imagen GCS',
                     )
                   else
                     const CircularProgressIndicator(),
@@ -158,6 +173,36 @@ class _ImagenPageState extends State<ImagenPage> {
     });
   }
 
+  Future<void> _testImagenGCS(String prompt) async {
+    setState(() {
+      _loading = true;
+    });
+    var gcsUrl = 'gs://vertex-ai-example-ef5a2.appspot.com/imagen';
+
+    var response = await _imagenModel.generateImagesGCS(prompt, gcsUrl);
+
+    if (response.images.isNotEmpty) {
+      var imagenImage = response.images[0];
+      final returnImageUri = imagenImage.gcsUri;
+      final reference = FirebaseStorage.instance.refFromURL(returnImageUri);
+      final downloadUrl = await reference.getDownloadURL();
+      // Process the image
+      _generatedContent.add(
+        MessageData(
+          image: Image(image: NetworkImage(downloadUrl)),
+          text: prompt,
+          fromUser: false,
+        ),
+      );
+    } else {
+      // Handle the case where no images were generated
+      _showError('Error: No images were generated.');
+    }
+    setState(() {
+      _loading = false;
+    });
+  }
+
   void _showError(String message) {
     showDialog<void>(
       context: context,
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/pubspec.yaml b/packages/firebase_vertexai/firebase_vertexai/example/pubspec.yaml
index 46a85c5c0825..8b123f56d255 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/pubspec.yaml
+++ b/packages/firebase_vertexai/firebase_vertexai/example/pubspec.yaml
@@ -20,6 +20,7 @@ dependencies:
   # Use with the CupertinoIcons class for iOS style icons.
   cupertino_icons: ^1.0.6
   firebase_core: ^3.10.1
+  firebase_storage: ^12.4.1
   firebase_vertexai: ^1.1.1
   flutter:
     sdk: flutter

From 74909e5de2822fe59197e8d371635d732710d678 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Mon, 3 Feb 2025 22:15:21 -0800
Subject: [PATCH 12/21] Add ImagenImagesBlockedException

---
 .../firebase_vertexai/lib/src/error.dart      | 15 ++++
 .../lib/src/imagen_content.dart               | 73 +++++++++++++------
 2 files changed, 64 insertions(+), 24 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart
index 66ab38e325c0..95979bd930ef 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart
@@ -111,6 +111,21 @@ final class VertexAISdkException implements Exception {
       'https://github.com/firebase/flutterfire/issues.';
 }
 
+/// Exception indicating all images filtered out.
+///
+/// This exception indicates all images were filtered out because they violated
+/// Vertex AI's usage guidelines.
+final class ImagenImagesBlockedException implements Exception {
+  /// Constructor
+  ImagenImagesBlockedException(this.message);
+
+  /// Message of the exception
+  final String message;
+
+  @override
+  String toString() => message;
+}
+
 /// Parse the error json object.
 VertexAIException parseError(Object jsonObject) {
   return switch (jsonObject) {
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
index f959d8863a20..a9a02b962766 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
@@ -35,6 +35,17 @@ final class ImagenInlineImage implements ImagenImage {
     required this.mimeType,
   });
 
+  /// Factory method to create an [ImagenInlineImage] from a JSON object.
+  factory ImagenInlineImage.fromJson(Map<String, dynamic> json) {
+    final mimeType = json['mimeType'] as String;
+    final bytes = json['bytesBase64Encoded'] as String;
+    final decodedBytes = base64Decode(bytes);
+    return ImagenInlineImage(
+      mimeType: mimeType,
+      bytesBase64Encoded: Uint8List.fromList(decodedBytes),
+    );
+  }
+
   /// The data contents in bytes, encoded as base64.
   final Uint8List bytesBase64Encoded;
 
@@ -56,6 +67,17 @@ final class ImagenGCSImage implements ImagenImage {
     required this.mimeType,
   });
 
+  /// Factory method to create an [ImagenGCSImage] from a JSON object.
+  factory ImagenGCSImage.fromJson(Map<String, dynamic> json) {
+    final mimeType = json['mimeType'] as String;
+    final uri = json['gcsUri'] as String;
+
+    return ImagenGCSImage(
+      mimeType: mimeType,
+      gcsUri: uri,
+    );
+  }
+
   /// The storage URI of the image.
   final String gcsUri;
 
@@ -77,37 +99,40 @@ final class ImagenGenerationResponse<T extends ImagenImage> {
     this.filteredReason,
   });
 
-  /// Factory method to create an ImagenGenerationResponse from a JSON object.
+  /// Factory method to create an [ImagenGenerationResponse] from a JSON object.
   factory ImagenGenerationResponse.fromJson(Map<String, dynamic> json) {
-    final filteredReason = json['filteredReason'] as String?;
-    final imagesJson = json['predictions'] as List<dynamic>;
+    final predictions = json['predictions'] as List<Map<String, dynamic>>;
+    List<T> images = [];
+    String? filteredReason;
 
     if (T == ImagenInlineImage) {
-      final images = imagesJson.map((imageJson) {
-        final mimeType = imageJson['mimeType'] as String;
-        final bytes = imageJson['bytesBase64Encoded'] as String;
-        final decodedBytes = base64Decode(bytes);
-        return ImagenInlineImage(
-          mimeType: mimeType,
-          bytesBase64Encoded: Uint8List.fromList(decodedBytes),
-        ) as T;
-      }).toList();
-      return ImagenGenerationResponse<T>(
-          images: images, filteredReason: filteredReason);
+      for (final prediction in predictions) {
+        if (prediction.containsKey('bytesBase64Encoded')) {
+          final image = ImagenInlineImage.fromJson(prediction) as T;
+          images.add(image);
+        } else if (prediction.containsKey('raiFilteredReason')) {
+          filteredReason = prediction['raiFilteredReason'] as String;
+        }
+      }
     } else if (T == ImagenGCSImage) {
-      final images = imagesJson.map((imageJson) {
-        final mimeType = imageJson['mimeType'] as String;
-        final gcsUri = imageJson['gcsUri'] as String;
-        return ImagenGCSImage(
-          mimeType: mimeType,
-          gcsUri: gcsUri,
-        ) as T;
-      }).toList();
-      return ImagenGenerationResponse<T>(
-          images: images, filteredReason: filteredReason);
+      for (final prediction in predictions) {
+        if (prediction.containsKey('gcsUri')) {
+          final image = ImagenGCSImage.fromJson(prediction) as T;
+          images.add(image);
+        } else if (prediction.containsKey('raiFilteredReason')) {
+          filteredReason = prediction['raiFilteredReason'] as String;
+        }
+      }
     } else {
       throw ArgumentError('Unsupported ImagenImage type: $T');
     }
+
+    if (images.isEmpty && filteredReason != null) {
+      throw ImagenImagesBlockedException(filteredReason);
+    }
+
+    return ImagenGenerationResponse<T>(
+        images: images, filteredReason: filteredReason);
   }
 
   /// A list of generated images. The type of the images depends on the T parameter.

From 01396d3046ec8bfc4a9baed605075c0d36ee13b9 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Mon, 3 Feb 2025 22:52:13 -0800
Subject: [PATCH 13/21] Add unit tests

---
 .../lib/src/imagen_content.dart               |   8 +-
 .../firebase_vertexai/test/imagen_test.dart   | 241 ++++++++++++++++++
 2 files changed, 247 insertions(+), 2 deletions(-)
 create mode 100644 packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
index a9a02b962766..9c17a04f4a64 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
@@ -55,7 +55,7 @@ final class ImagenInlineImage implements ImagenImage {
   @override
   Object toJson() => {
         'mimeType': mimeType,
-        'bytesBase64Encoded': bytesBase64Encoded,
+        'bytesBase64Encoded': base64Encode(bytesBase64Encoded),
       };
 }
 
@@ -101,7 +101,11 @@ final class ImagenGenerationResponse<T extends ImagenImage> {
 
   /// Factory method to create an [ImagenGenerationResponse] from a JSON object.
   factory ImagenGenerationResponse.fromJson(Map<String, dynamic> json) {
-    final predictions = json['predictions'] as List<Map<String, dynamic>>;
+    final predictions = json['predictions'];
+    if (predictions.isEmpty) {
+      throw ServerException('Got empty prediction with no reason');
+    }
+
     List<T> images = [];
     String? filteredReason;
 
diff --git a/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart b/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
new file mode 100644
index 000000000000..803234ad8199
--- /dev/null
+++ b/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
@@ -0,0 +1,241 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import 'dart:convert';
+import 'dart:typed_data';
+
+import 'package:firebase_vertexai/src/error.dart';
+import 'package:firebase_vertexai/src/imagen_content.dart';
+import 'package:flutter_test/flutter_test.dart';
+
+void main() {
+  group('ImagenInlineImage', () {
+    test('fromJson with valid base64', () {
+      final json = {
+        'mimeType': 'image/png',
+        'bytesBase64Encoded':
+            'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII='
+      };
+      final image = ImagenInlineImage.fromJson(json);
+      expect(image.mimeType, 'image/png');
+      expect(image.bytesBase64Encoded, isA<Uint8List>());
+      expect(image.bytesBase64Encoded, isNotEmpty);
+    });
+
+    test('fromJson with invalid base64', () {
+      final json = {
+        'mimeType': 'image/png',
+        'bytesBase64Encoded': 'invalid_base64_string'
+      };
+      // Expect that the constructor throws an exception.
+      expect(() => ImagenInlineImage.fromJson(json), throwsFormatException);
+    });
+
+    test('toJson', () {
+      final image = ImagenInlineImage(
+        mimeType: 'image/png',
+        bytesBase64Encoded: Uint8List.fromList(utf8.encode('Hello, world!')),
+      );
+      final json = image.toJson();
+      expect(json, {
+        'mimeType': 'image/png',
+        'bytesBase64Encoded': 'SGVsbG8sIHdvcmxkIQ==',
+      });
+    });
+  });
+
+  group('ImagenGCSImage', () {
+    test('fromJson', () {
+      final json = {
+        'mimeType': 'image/jpeg',
+        'gcsUri':
+            'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_0.jpg'
+      };
+      final image = ImagenGCSImage.fromJson(json);
+      expect(image.mimeType, 'image/jpeg');
+      expect(image.gcsUri,
+          'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_0.jpg');
+    });
+
+    test('toJson', () {
+      final image = ImagenGCSImage(
+        mimeType: 'image/jpeg',
+        gcsUri:
+            'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_0.jpg',
+      );
+      final json = image.toJson();
+      expect(json, {
+        'mimeType': 'image/jpeg',
+        'gcsUri':
+            'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_0.jpg',
+      });
+    });
+  });
+
+  group('ImagenGenerationResponse', () {
+    test('fromJson with gcsUri', () {
+      final json = {
+        'predictions': [
+          {
+            'mimeType': 'image/jpeg',
+            'gcsUri':
+                'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_0.jpg'
+          },
+          {
+            'mimeType': 'image/jpeg',
+            'gcsUri':
+                'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_1.jpg'
+          },
+          {
+            'mimeType': 'image/jpeg',
+            'gcsUri':
+                'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_2.jpg'
+          },
+          {
+            'mimeType': 'image/jpeg',
+            'gcsUri':
+                'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_3.jpg'
+          }
+        ]
+      };
+      final response = ImagenGenerationResponse<ImagenGCSImage>.fromJson(json);
+      expect(response.images, isA<List<ImagenGCSImage>>());
+      expect(response.images.length, 4);
+      expect(response.filteredReason, isNull);
+    });
+
+    test('fromJson with bytesBase64Encoded', () {
+      final json = {
+        'predictions': [
+          {
+            'mimeType': 'image/jpeg',
+            'bytesBase64Encoded': 'SGVsbG8sIHdvcmxkIQ=='
+          },
+          {
+            'mimeType': 'image/jpeg',
+            'bytesBase64Encoded': 'SGVsbG8sIHdvcmxkIQ=='
+          },
+          {
+            'mimeType': 'image/jpeg',
+            'bytesBase64Encoded': 'SGVsbG8sIHdvcmxkIQ=='
+          },
+          {
+            'mimeType': 'image/jpeg',
+            'bytesBase64Encoded': 'SGVsbG8sIHdvcmxkIQ=='
+          }
+        ]
+      };
+      final response =
+          ImagenGenerationResponse<ImagenInlineImage>.fromJson(json);
+      expect(response.images, isA<List<ImagenInlineImage>>());
+      expect(response.images.length, 4);
+      expect(response.filteredReason, isNull);
+    });
+
+    test('fromJson with bytesBase64Encoded and raiFilteredReason', () {
+      final json = {
+        'predictions': [
+          {
+            'mimeType': 'image/png',
+            'bytesBase64Encoded':
+                'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII='
+          },
+          {
+            'mimeType': 'image/png',
+            'bytesBase64Encoded':
+                'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII='
+          },
+          {
+            'raiFilteredReason':
+                'Your current safety filter threshold filtered out 2 generated images. You will not be charged for blocked images. Try rephrasing the prompt. If you think this was an error, send feedback.'
+          }
+        ]
+      };
+      final response =
+          ImagenGenerationResponse<ImagenInlineImage>.fromJson(json);
+      expect(response.images, isA<List<ImagenInlineImage>>());
+      expect(response.images.length, 2);
+      expect(response.filteredReason,
+          'Your current safety filter threshold filtered out 2 generated images. You will not be charged for blocked images. Try rephrasing the prompt. If you think this was an error, send feedback.');
+    });
+
+    test('fromJson with only raiFilteredReason', () {
+      final json = {
+        'predictions': [
+          {
+            'raiFilteredReason':
+                'Unable to show generated images. All images were filtered out because they violated Vertex AI\'s usage guidelines. You will not be charged for blocked images. Try rephrasing the prompt. If you think this was an error, send feedback. Support codes: 39322892, 29310472'
+          }
+        ]
+      };
+      // Expect that the constructor throws an exception.
+      expect(() => ImagenGenerationResponse<ImagenInlineImage>.fromJson(json),
+          throwsA(isA<ImagenImagesBlockedException>()));
+    });
+
+    test('fromJson with empty predictions', () {
+      final json = {'predictions': {}};
+      // Expect that the constructor throws an exception.
+      expect(() => ImagenGenerationResponse<ImagenInlineImage>.fromJson(json),
+          throwsA(isA<ServerException>()));
+    });
+
+    test('fromJson with unsupported type', () {
+      final json = {
+        'predictions': [
+          {
+            'mimeType': 'image/jpeg',
+            'gcsUri':
+                'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_0.jpg'
+          },
+        ]
+      };
+      // Expect that the constructor throws an exception.
+      expect(() => ImagenGenerationResponse<ImagenImage>.fromJson(json),
+          throwsA(isA<ArgumentError>()));
+    });
+  });
+
+  group('parseImagenGenerationResponse', () {
+    test('with valid response', () {
+      final json = {
+        'predictions': [
+          {
+            'mimeType': 'image/jpeg',
+            'gcsUri':
+                'gs://test-project-id-1234.firebasestorage.app/images/1234567890123/sample_0.jpg'
+          },
+        ]
+      };
+      final response = parseImagenGenerationResponse<ImagenGCSImage>(json);
+      expect(response.images, isA<List<ImagenGCSImage>>());
+      expect(response.images.length, 1);
+      expect(response.filteredReason, isNull);
+    });
+
+    test('with error', () {
+      final json = {
+        'error': {
+          'code': 400,
+          'message':
+              'Image generation failed with the following error: The prompt could not be submitted. This prompt contains sensitive words that violate Google\'s Responsible AI practices. Try rephrasing the prompt. If you think this was an error, send feedback. Support codes: 42876398',
+          'status': 'INVALID_ARGUMENT'
+        }
+      };
+      // Expect that the function throws an exception.
+      expect(() => parseImagenGenerationResponse<ImagenGCSImage>(json),
+          throwsA(isA<ServerException>()));
+    });
+  });
+}

From 232ac9c579377a8a47aee18eb171520848c1cceb Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Mon, 3 Feb 2025 22:52:29 -0800
Subject: [PATCH 14/21] update the year

---
 .../firebase_vertexai/firebase_vertexai/test/imagen_test.dart   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart b/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
index 803234ad8199..05a8736ec95f 100644
--- a/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
@@ -1,4 +1,4 @@
-// Copyright 2024 Google LLC
+// Copyright 2025 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

From 94653358543e43c8fc0b53012972aa0bcb81a205 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Fri, 7 Feb 2025 08:45:23 -0800
Subject: [PATCH 15/21] finalize with the name

---
 .../firebase_vertexai/example/lib/pages/imagen_page.dart  | 4 ++--
 .../firebase_vertexai/lib/src/firebase_vertexai.dart      | 6 +++---
 .../firebase_vertexai/lib/src/imagen_model.dart           | 8 ++++----
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
index 9fa405837648..757869e9a23e 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
@@ -48,8 +48,8 @@ class _ImagenPageState extends State<ImagenPage> {
       aspectRatio: ImagenAspectRatio.square1x1,
       imageFormat: ImagenFormat.jpeg(compressionQuality: 75),
     );
-    _imagenModel = FirebaseVertexAI.instance.imageModel(
-      modelName: 'imagen-3.0-generate-001',
+    _imagenModel = FirebaseVertexAI.instance.imagenModel(
+      model: 'imagen-3.0-generate-001',
       generationConfig: generationConfig,
       safetySettings: ImagenSafetySettings(
         ImagenSafetyFilterLevel.blockLowAndAbove,
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
index 19d37a947432..02786e6000d2 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
@@ -118,14 +118,14 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
   ///
   /// The optional [safetySettings] can be used to control and guide the
   /// generation. See [ImagenSafetySettings] for details.
-  ImagenModel imageModel(
-      {required String modelName,
+  ImagenModel imagenModel(
+      {required String model,
       ImagenGenerationConfig? generationConfig,
       ImagenSafetySettings? safetySettings}) {
     return createImagenModel(
         app: app,
         location: location,
-        modelName: modelName,
+        model: model,
         generationConfig: generationConfig,
         safetySettings: safetySettings,
         appCheck: appCheck,
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index 9d65fff160da..fe724dac65a7 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -24,7 +24,7 @@ import 'imagen_content.dart';
 final class ImagenModel extends BaseModel {
   ImagenModel._(
       {required FirebaseApp app,
-      required String modelName,
+      required String model,
       required String location,
       FirebaseAppCheck? appCheck,
       FirebaseAuth? auth,
@@ -33,7 +33,7 @@ final class ImagenModel extends BaseModel {
       : _generationConfig = generationConfig,
         _safetySettings = safetySettings,
         super(
-            model: modelName,
+            model: model,
             app: app,
             location: location,
             client: HttpApiClient(
@@ -121,14 +121,14 @@ final class ImagenModel extends BaseModel {
 ImagenModel createImagenModel({
   required FirebaseApp app,
   required String location,
-  required String modelName,
+  required String model,
   FirebaseAppCheck? appCheck,
   FirebaseAuth? auth,
   ImagenGenerationConfig? generationConfig,
   ImagenSafetySettings? safetySettings,
 }) =>
     ImagenModel._(
-      modelName: modelName,
+      model: model,
       app: app,
       appCheck: appCheck,
       auth: auth,

From 7117d53506ba1bedcacb3d5426a78559d72d4a43 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Tue, 11 Feb 2025 10:05:02 -0800
Subject: [PATCH 16/21] Hide gcs api and some review feedback

---
 .../example/lib/pages/imagen_page.dart        | 83 ++++++++++---------
 .../lib/firebase_vertexai.dart                |  5 +-
 .../firebase_vertexai/lib/src/chat.dart       |  2 +-
 .../lib/src/firebase_vertexai.dart            |  5 +-
 .../src/{model.dart => generative_model.dart} |  0
 .../lib/src/imagen_model.dart                 |  3 +-
 .../firebase_vertexai/test/chat_test.dart     |  2 +-
 .../firebase_vertexai/test/model_test.dart    |  2 +-
 8 files changed, 51 insertions(+), 51 deletions(-)
 rename packages/firebase_vertexai/firebase_vertexai/lib/src/{model.dart => generative_model.dart} (100%)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
index 757869e9a23e..42aacf6b8e36 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
@@ -125,19 +125,20 @@ class _ImagenPageState extends State<ImagenPage> {
                     )
                   else
                     const CircularProgressIndicator(),
-                  if (!_loading)
-                    IconButton(
-                      onPressed: () async {
-                        await _testImagenGCS(_textController.text);
-                      },
-                      icon: Icon(
-                        Icons.imagesearch_roller,
-                        color: Theme.of(context).colorScheme.primary,
-                      ),
-                      tooltip: 'Imagen GCS',
-                    )
-                  else
-                    const CircularProgressIndicator(),
+                  // NOTE: Keep this API private until future release.
+                  // if (!_loading)
+                  //   IconButton(
+                  //     onPressed: () async {
+                  //       await _testImagenGCS(_textController.text);
+                  //     },
+                  //     icon: Icon(
+                  //       Icons.imagesearch_roller,
+                  //       color: Theme.of(context).colorScheme.primary,
+                  //     ),
+                  //     tooltip: 'Imagen GCS',
+                  //   )
+                  // else
+                  //   const CircularProgressIndicator(),
                 ],
               ),
             ),
@@ -172,36 +173,36 @@ class _ImagenPageState extends State<ImagenPage> {
       _loading = false;
     });
   }
+  // NOTE: Keep this API private until future release.
+  // Future<void> _testImagenGCS(String prompt) async {
+  //   setState(() {
+  //     _loading = true;
+  //   });
+  //   var gcsUrl = 'gs://vertex-ai-example-ef5a2.appspot.com/imagen';
 
-  Future<void> _testImagenGCS(String prompt) async {
-    setState(() {
-      _loading = true;
-    });
-    var gcsUrl = 'gs://vertex-ai-example-ef5a2.appspot.com/imagen';
-
-    var response = await _imagenModel.generateImagesGCS(prompt, gcsUrl);
+  //   var response = await _imagenModel.generateImagesGCS(prompt, gcsUrl);
 
-    if (response.images.isNotEmpty) {
-      var imagenImage = response.images[0];
-      final returnImageUri = imagenImage.gcsUri;
-      final reference = FirebaseStorage.instance.refFromURL(returnImageUri);
-      final downloadUrl = await reference.getDownloadURL();
-      // Process the image
-      _generatedContent.add(
-        MessageData(
-          image: Image(image: NetworkImage(downloadUrl)),
-          text: prompt,
-          fromUser: false,
-        ),
-      );
-    } else {
-      // Handle the case where no images were generated
-      _showError('Error: No images were generated.');
-    }
-    setState(() {
-      _loading = false;
-    });
-  }
+  //   if (response.images.isNotEmpty) {
+  //     var imagenImage = response.images[0];
+  //     final returnImageUri = imagenImage.gcsUri;
+  //     final reference = FirebaseStorage.instance.refFromURL(returnImageUri);
+  //     final downloadUrl = await reference.getDownloadURL();
+  //     // Process the image
+  //     _generatedContent.add(
+  //       MessageData(
+  //         image: Image(image: NetworkImage(downloadUrl)),
+  //         text: prompt,
+  //         fromUser: false,
+  //       ),
+  //     );
+  //   } else {
+  //     // Handle the case where no images were generated
+  //     _showError('Error: No images were generated.');
+  //   }
+  //   setState(() {
+  //     _loading = false;
+  //   });
+  // }
 
   void _showError(String message) {
     showDialog<void>(
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
index af79eed8efd1..17b05428ed3a 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
@@ -56,6 +56,7 @@ export 'src/function_calling.dart'
         FunctionDeclaration,
         Tool,
         ToolConfig;
+export 'src/generative_model.dart' show GenerativeModel;
 export 'src/imagen_api.dart'
     show
         ImagenSafetySettings,
@@ -64,8 +65,6 @@ export 'src/imagen_api.dart'
         ImagenPersonFilterLevel,
         ImagenGenerationConfig,
         ImagenAspectRatio;
-export 'src/imagen_content.dart'
-    show ImagenInlineImage, ImagenGCSImage, ImagenImage;
+export 'src/imagen_content.dart' show ImagenInlineImage;
 export 'src/imagen_model.dart' show ImagenModel;
-export 'src/model.dart' show GenerativeModel;
 export 'src/schema.dart' show Schema, SchemaType;
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/chat.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/chat.dart
index fa9229f46552..553ccb4e05bd 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/chat.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/chat.dart
@@ -16,7 +16,7 @@ import 'dart:async';
 
 import 'api.dart';
 import 'content.dart';
-import 'model.dart';
+import 'generative_model.dart';
 import 'utils/mutex.dart';
 
 /// A back-and-forth chat with a generative model.
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
index 02786e6000d2..364c8d8f687b 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
@@ -21,10 +21,9 @@ import 'package:firebase_core_platform_interface/firebase_core_platform_interfac
 import 'api.dart';
 import 'content.dart';
 import 'function_calling.dart';
-import 'imagen_model.dart';
+import 'generative_model.dart';
 import 'imagen_api.dart';
-import 'imagen_content.dart';
-import 'model.dart';
+import 'imagen_model.dart';
 
 const _defaultLocation = 'us-central1';
 
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/generative_model.dart
similarity index 100%
rename from packages/firebase_vertexai/firebase_vertexai/lib/src/model.dart
rename to packages/firebase_vertexai/firebase_vertexai/lib/src/generative_model.dart
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index fe724dac65a7..19ed7d55b01a 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -102,7 +102,8 @@ final class ImagenModel extends BaseModel {
 
   /// Generates images with format of [ImagenGCSImage] based on the given
   /// prompt.
-  Future<ImagenGenerationResponse<ImagenGCSImage>> generateImagesGCS(
+  /// Note: Keep this API private until future release.
+  Future<ImagenGenerationResponse<ImagenGCSImage>> _generateImagesGCS(
     String prompt,
     String gcsUri,
   ) =>
diff --git a/packages/firebase_vertexai/firebase_vertexai/test/chat_test.dart b/packages/firebase_vertexai/firebase_vertexai/test/chat_test.dart
index 511c539f1a7e..a0f404cdac20 100644
--- a/packages/firebase_vertexai/firebase_vertexai/test/chat_test.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/test/chat_test.dart
@@ -13,7 +13,7 @@
 // limitations under the License.
 import 'package:firebase_core/firebase_core.dart';
 import 'package:firebase_vertexai/firebase_vertexai.dart';
-import 'package:firebase_vertexai/src/model.dart';
+import 'package:firebase_vertexai/src/generative_model.dart';
 import 'package:flutter_test/flutter_test.dart';
 
 import 'mock.dart';
diff --git a/packages/firebase_vertexai/firebase_vertexai/test/model_test.dart b/packages/firebase_vertexai/firebase_vertexai/test/model_test.dart
index f0c05f71e8dc..9dba6c498bb7 100644
--- a/packages/firebase_vertexai/firebase_vertexai/test/model_test.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/test/model_test.dart
@@ -13,7 +13,7 @@
 // limitations under the License.
 import 'package:firebase_core/firebase_core.dart';
 import 'package:firebase_vertexai/firebase_vertexai.dart';
-import 'package:firebase_vertexai/src/model.dart';
+import 'package:firebase_vertexai/src/generative_model.dart';
 import 'package:flutter_test/flutter_test.dart';
 
 import 'mock.dart';

From 257e2ec64f3c9d9d32b43eacac458ae68df85c36 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@users.noreply.github.com>
Date: Tue, 11 Feb 2025 10:50:15 -0800
Subject: [PATCH 17/21] Update
 packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart

Co-authored-by: Nate Bosch <nbosch@google.com>
---
 .../firebase_vertexai/lib/src/imagen_api.dart                  | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
index 735eab7217d9..366d9600cd61 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
@@ -171,7 +171,8 @@ final class ImagenGenerationConfig {
   /// The image format of the generated images.
   final ImagenFormat? imageFormat;
 
-  /// Add an invisible watermark to the generated images.
+  /// Whether to add an invisible watermark to generated images.
+  ///
   /// Default value for each imagen model can be found in
   /// https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api#generate_images
   final bool? addWatermark;

From fa45ad28d61827711193182815f2059504957720 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@users.noreply.github.com>
Date: Tue, 11 Feb 2025 10:51:13 -0800
Subject: [PATCH 18/21] Update
 packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart

Co-authored-by: Nate Bosch <nbosch@google.com>
---
 .../lib/src/imagen_model.dart                 | 46 +++++++------------
 1 file changed, 16 insertions(+), 30 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index 19ed7d55b01a..aa21d9d91301 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -47,36 +47,22 @@ final class ImagenModel extends BaseModel {
     String prompt, {
     String? gcsUri,
   }) {
-    final parameters = <String, Object?>{};
-
-    if (gcsUri != null) parameters['storageUri'] = gcsUri;
-
-    parameters['sampleCount'] = _generationConfig?.numberOfImages ?? 1;
-    if (_generationConfig != null) {
-      if (_generationConfig.aspectRatio != null) {
-        parameters['aspectRatio'] = _generationConfig.aspectRatio;
-      }
-      if (_generationConfig.negativePrompt != null) {
-        parameters['negativePrompt'] = _generationConfig.negativePrompt;
-      }
-      if (_generationConfig.addWatermark != null) {
-        parameters['addWatermark'] = _generationConfig.addWatermark;
-      }
-      if (_generationConfig.imageFormat != null) {
-        parameters['outputOption'] = _generationConfig.imageFormat!.toJson();
-      }
-    }
-
-    if (_safetySettings != null) {
-      if (_safetySettings.personFilterLevel != null) {
-        parameters['personGeneration'] =
-            _safetySettings.personFilterLevel!.toJson();
-      }
-      if (_safetySettings.safetyFilterLevel != null) {
-        parameters['safetySetting'] =
-            _safetySettings.safetyFilterLevel!.toJson();
-      }
-    }
+    final parameters = <String, Object?>{
+      if (gcsUri != null) 'storageUri': gcsUri,
+      'sampleCount': _generationConfig?.numberOfImages ?? 1,
+      if (_generationConfig?.aspectRatio case final aspectRatio?)
+        'aspectRatio': aspectRatio,
+      if (_generationConfig?.negativePrompt case final negativePrompt?)
+        'negativePrompt': negativePrompt,
+      if (_generationConfig?.addWatermark case final addWatermark?)
+        'addWatermark': addWatermark,
+      if (_generationConfig.imageFormat case final imageFormat?)
+        'outputOption': imageFormat.toJson(),
+      if (_safetySettings?.personFilterLevel case final personFilterLevel?)
+        'personGeneration': personFilterLevel.toJson(),
+      if (_safetySettings?.safetyFilterLevel case final safetyFilterLevel?)
+        'safetySetting': safetyFilterLevel.toJson(),
+    };
 
     return {
       'instances': [

From 75b331e016951fe1b318ccb91e527d19136be055 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Tue, 11 Feb 2025 11:01:38 -0800
Subject: [PATCH 19/21] review comments

---
 .../firebase_vertexai/lib/src/api.dart        | 20 +++++++++----------
 .../firebase_vertexai/lib/src/base_model.dart |  2 +-
 .../firebase_vertexai/lib/src/content.dart    | 12 +++++------
 .../firebase_vertexai/lib/src/error.dart      | 10 +++++-----
 .../lib/src/function_calling.dart             |  8 ++++----
 .../firebase_vertexai/lib/src/imagen_api.dart | 20 ++++++++++---------
 .../lib/src/imagen_content.dart               |  8 ++++----
 .../firebase_vertexai/lib/src/schema.dart     |  2 +-
 8 files changed, 42 insertions(+), 40 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/api.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/api.dart
index 4101364057a9..eae11b00072a 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/api.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/api.dart
@@ -18,7 +18,7 @@ import 'schema.dart';
 
 /// Response for Count Tokens
 final class CountTokensResponse {
-  /// Constructor
+  // ignore: public_member_api_docs
   CountTokensResponse(this.totalTokens, {this.totalBillableCharacters});
 
   /// The number of tokens that the `model` tokenizes the `prompt` into.
@@ -34,7 +34,7 @@ final class CountTokensResponse {
 
 /// Response from the model; supports multiple candidates.
 final class GenerateContentResponse {
-  /// Constructor
+  // ignore: public_member_api_docs
   GenerateContentResponse(this.candidates, this.promptFeedback,
       {this.usageMetadata});
 
@@ -108,7 +108,7 @@ final class GenerateContentResponse {
 
 /// Feedback metadata of a prompt specified in a [GenerativeModel] request.
 final class PromptFeedback {
-  /// Constructor
+  // ignore: public_member_api_docs
   PromptFeedback(this.blockReason, this.blockReasonMessage, this.safetyRatings);
 
   /// If set, the prompt was blocked and no candidates are returned.
@@ -127,7 +127,7 @@ final class PromptFeedback {
 
 /// Metadata on the generation request's token usage.
 final class UsageMetadata {
-  /// Constructor
+  // ignore: public_member_api_docs
   UsageMetadata._({
     this.promptTokenCount,
     this.candidatesTokenCount,
@@ -147,7 +147,7 @@ final class UsageMetadata {
 /// Response candidate generated from a [GenerativeModel].
 final class Candidate {
   // TODO: token count?
-  /// Constructor
+  // ignore: public_member_api_docs
   Candidate(this.content, this.safetyRatings, this.citationMetadata,
       this.finishReason, this.finishMessage);
 
@@ -212,7 +212,7 @@ final class Candidate {
 /// safety across a number of harm categories and the probability of the harm
 /// classification is included here.
 final class SafetyRating {
-  /// Constructor
+  // ignore: public_member_api_docs
   SafetyRating(this.category, this.probability,
       {this.probabilityScore,
       this.isBlocked,
@@ -406,7 +406,7 @@ enum HarmSeverity {
 
 /// Source attributions for a piece of content.
 final class CitationMetadata {
-  /// Constructor
+  // ignore: public_member_api_docs
   CitationMetadata(this.citations);
 
   /// Citations to sources for a specific response.
@@ -415,7 +415,7 @@ final class CitationMetadata {
 
 /// Citation to a source for a portion of a specific response.
 final class Citation {
-  /// Constructor
+  // ignore: public_member_api_docs
   Citation(this.startIndex, this.endIndex, this.uri, this.license);
 
   /// Start of segment of the response that is attributed to this source.
@@ -486,7 +486,7 @@ enum FinishReason {
 /// Passing a safety setting for a category changes the allowed probability that
 /// content is blocked.
 final class SafetySetting {
-  /// Constructor
+  // ignore: public_member_api_docs
   SafetySetting(this.category, this.threshold);
 
   /// The category for this setting.
@@ -542,7 +542,7 @@ enum HarmBlockThreshold {
 
 /// Configuration options for model generation and outputs.
 final class GenerationConfig {
-  /// Constructor
+  // ignore: public_member_api_docs
   GenerationConfig(
       {this.candidateCount,
       this.stopSequences,
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
index 805356d1909b..8347c49fddce 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
@@ -31,7 +31,7 @@ enum Task {
   /// Request type to count token.
   countTokens,
 
-  /// Imagen 3 task
+  /// Request type to talk to Prediction Services like Imagen.
   predict,
 }
 
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/content.dart
index 1fa79d2f5bf0..74a435fd6ef6 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/content.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/content.dart
@@ -18,7 +18,7 @@ import 'error.dart';
 
 /// The base structured datatype containing multi-part content of a message.
 final class Content {
-  /// Constructor
+  // ignore: public_member_api_docs
   Content(this.role, this.parts);
 
   /// The producer of the content.
@@ -112,7 +112,7 @@ sealed class Part {
 
 /// A [Part] with the text content.
 final class TextPart implements Part {
-  /// Constructor
+  // ignore: public_member_api_docs
   TextPart(this.text);
 
   /// The text content of the [Part]
@@ -123,7 +123,7 @@ final class TextPart implements Part {
 
 /// A [Part] with the byte content of a file.
 final class InlineDataPart implements Part {
-  /// Constructor
+  // ignore: public_member_api_docs
   InlineDataPart(this.mimeType, this.bytes);
 
   /// File type of the [InlineDataPart].
@@ -142,7 +142,7 @@ final class InlineDataPart implements Part {
 /// a string representing the `FunctionDeclaration.name` with the
 /// arguments and their values.
 final class FunctionCall implements Part {
-  /// Constructor
+  // ignore: public_member_api_docs
   FunctionCall(this.name, this.args);
 
   /// The name of the function to call.
@@ -160,7 +160,7 @@ final class FunctionCall implements Part {
 
 /// The response class for [FunctionCall]
 final class FunctionResponse implements Part {
-  /// Constructor
+  // ignore: public_member_api_docs
   FunctionResponse(this.name, this.response);
 
   /// The name of the function that was called.
@@ -180,7 +180,7 @@ final class FunctionResponse implements Part {
 
 /// A [Part] with Firebase Storage uri as prompt content
 final class FileData implements Part {
-  /// Constructor
+  // ignore: public_member_api_docs
   FileData(this.mimeType, this.fileUri);
 
   /// File type of the [FileData].
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart
index 95979bd930ef..ad4a5e09d9e2 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/error.dart
@@ -16,7 +16,7 @@
 ///
 /// The [message] may explain the cause of the failure.
 final class VertexAIException implements Exception {
-  /// Constructor
+  // ignore: public_member_api_docs
   VertexAIException(this.message);
 
   /// Message of the exception
@@ -28,7 +28,7 @@ final class VertexAIException implements Exception {
 
 /// Exception thrown when the server rejects the API key.
 final class InvalidApiKey implements VertexAIException {
-  /// Constructor
+  // ignore: public_member_api_docs
   InvalidApiKey(this.message);
   @override
   final String message;
@@ -81,7 +81,7 @@ final class QuotaExceeded implements VertexAIException {
 
 /// Exception thrown when the server failed to generate content.
 final class ServerException implements VertexAIException {
-  /// Constructor
+  // ignore: public_member_api_docs
   ServerException(this.message);
   @override
   final String message;
@@ -96,7 +96,7 @@ final class ServerException implements VertexAIException {
 /// as an inability to parse a new response format. Resolution paths may include
 /// updating to a new version of the SDK, or filing an issue.
 final class VertexAISdkException implements Exception {
-  /// Constructor
+  // ignore: public_member_api_docs
   VertexAISdkException(this.message);
 
   /// Message of the exception
@@ -116,7 +116,7 @@ final class VertexAISdkException implements Exception {
 /// This exception indicates all images were filtered out because they violated
 /// Vertex AI's usage guidelines.
 final class ImagenImagesBlockedException implements Exception {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenImagesBlockedException(this.message);
 
   /// Message of the exception
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/function_calling.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/function_calling.dart
index 5d552e33024e..f70bff0b3ff7 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/function_calling.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/function_calling.dart
@@ -20,7 +20,7 @@ import 'schema.dart';
 /// external systems to perform an action, or set of actions, outside of
 /// knowledge and scope of the model.
 final class Tool {
-  /// Constructor
+  // ignore: public_member_api_docs
   Tool._(this._functionDeclarations);
 
   /// Returns a [Tool] instance with list of [FunctionDeclaration].
@@ -54,7 +54,7 @@ final class Tool {
 /// FunctionDeclaration is a representation of a block of code that can be used
 /// as a `Tool` by the model and executed by the client.
 final class FunctionDeclaration {
-  /// Constructor
+  // ignore: public_member_api_docs
   FunctionDeclaration(this.name, this.description,
       {required Map<String, Schema> parameters,
       List<String> optionalParameters = const []})
@@ -82,7 +82,7 @@ final class FunctionDeclaration {
 
 /// Config for tools to use with model.
 final class ToolConfig {
-  /// Constructor
+  // ignore: public_member_api_docs
   ToolConfig({this.functionCallingConfig});
 
   /// Config for function calling.
@@ -98,7 +98,7 @@ final class ToolConfig {
 /// Configuration specifying how the model should use the functions provided as
 /// tools.
 final class FunctionCallingConfig {
-  /// Constructor
+  // ignore: public_member_api_docs
   FunctionCallingConfig._({this.mode, this.allowedFunctionNames});
 
   /// The mode in which function calling should execute.
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
index 735eab7217d9..ff7911b50c9f 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
@@ -33,7 +33,7 @@ enum ImagenSafetyFilterLevel {
 
   final String _jsonString;
 
-  /// Convert to json format
+  // ignore: public_member_api_docs
   String toJson() => _jsonString;
 
   // ignore: unused_element
@@ -68,7 +68,7 @@ enum ImagenPersonFilterLevel {
 
   final String _jsonString;
 
-  /// Convert to json format
+  // ignore: public_member_api_docs
   String toJson() => _jsonString;
 
   // ignore: unused_element
@@ -90,7 +90,7 @@ enum ImagenPersonFilterLevel {
 ///
 /// It includes a safety filter level and a person filter level.
 final class ImagenSafetySettings {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenSafetySettings(this.safetyFilterLevel, this.personFilterLevel);
 
   /// The safety filter level
@@ -151,7 +151,7 @@ enum ImagenAspectRatio {
 
 /// Configuration options for image generation.
 final class ImagenGenerationConfig {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenGenerationConfig(
       {this.numberOfImages,
       this.negativePrompt,
@@ -159,7 +159,9 @@ final class ImagenGenerationConfig {
       this.imageFormat,
       this.addWatermark});
 
-  /// The number of images to generate. Default is 1.
+  /// The number of images to generate.
+  ///
+  /// Default value is 1.
   final int? numberOfImages;
 
   /// A description of what to discourage in the generated images.
@@ -176,7 +178,7 @@ final class ImagenGenerationConfig {
   /// https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api#generate_images
   final bool? addWatermark;
 
-  /// Convert to json format.
+  // ignore: public_member_api_docs
   Map<String, dynamic> toJson() => {
         if (negativePrompt != null) 'negativePrompt': negativePrompt,
         if (numberOfImages != null) 'numberOfImages': numberOfImages,
@@ -188,13 +190,13 @@ final class ImagenGenerationConfig {
 
 /// Represents the image format and compression quality.
 final class ImagenFormat {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenFormat(this.mimeType, this.compressionQuality);
 
-  /// Constructor for png
+  // ignore: public_member_api_docs
   ImagenFormat.png() : this('image/png', null);
 
-  /// Constructor for jpeg
+  // ignore: public_member_api_docs
   ImagenFormat.jpeg({int? compressionQuality})
       : this('image/jpeg', compressionQuality);
 
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
index 9c17a04f4a64..71d16f9da704 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_content.dart
@@ -17,7 +17,7 @@ import 'error.dart';
 
 /// Base type of Imagen Image.
 sealed class ImagenImage {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenImage({required this.mimeType});
 
   /// The MIME type of the image format.
@@ -29,7 +29,7 @@ sealed class ImagenImage {
 
 /// Represents an image stored as a base64-encoded string.
 final class ImagenInlineImage implements ImagenImage {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenInlineImage({
     required this.bytesBase64Encoded,
     required this.mimeType,
@@ -61,7 +61,7 @@ final class ImagenInlineImage implements ImagenImage {
 
 /// Represents an image stored in Google Cloud Storage.
 final class ImagenGCSImage implements ImagenImage {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenGCSImage({
     required this.gcsUri,
     required this.mimeType,
@@ -93,7 +93,7 @@ final class ImagenGCSImage implements ImagenImage {
 
 /// Represents the response from an image generation request.
 final class ImagenGenerationResponse<T extends ImagenImage> {
-  /// Constructor
+  // ignore: public_member_api_docs
   ImagenGenerationResponse({
     required this.images,
     this.filteredReason,
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/schema.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/schema.dart
index 0dda0f564d8a..e73f44b355f3 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/schema.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/schema.dart
@@ -18,7 +18,7 @@
 /// Represents a select subset of an
 /// [OpenAPI 3.0 schema object](https://spec.openapis.org/oas/v3.0.3#schema).
 final class Schema {
-  /// Constructor
+  // ignore: public_member_api_docs
   Schema(
     this.type, {
     this.format,

From a3cfe59a357442cba4102d78fdaf17fc7ab37135 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Tue, 11 Feb 2025 15:26:52 -0800
Subject: [PATCH 20/21] More review comment update

---
 .../firebase_vertexai/lib/src/imagen_api.dart        | 12 ++++++++----
 .../firebase_vertexai/lib/src/imagen_model.dart      | 11 ++++++++++-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
index 4cfcd7bcc425..86ef4baae0d3 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_api.dart
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 /// Specifies the level of safety filtering for image generation.
+///
 /// If not specified, default will be "block_medium_and_above".
 enum ImagenSafetyFilterLevel {
   /// Strongest filtering level, most strict blocking.
@@ -53,6 +54,7 @@ enum ImagenSafetyFilterLevel {
 }
 
 /// Allow generation of people by the model.
+///
 /// If not specified, the default value is "allow_adult".
 enum ImagenPersonFilterLevel {
   /// Disallow the inclusion of people or faces in images.
@@ -99,7 +101,7 @@ final class ImagenSafetySettings {
   /// The person filter level
   final ImagenPersonFilterLevel? personFilterLevel;
 
-  /// Convert to json format.
+  // ignore: public_member_api_docs
   Object toJson() => {
         if (safetyFilterLevel != null)
           'safetySetting': safetyFilterLevel!.toJson(),
@@ -108,7 +110,9 @@ final class ImagenSafetySettings {
       };
 }
 
-/// The aspect ratio for the image. The default value is "1:1".
+/// The aspect ratio for the image.
+///
+/// The default value is "1:1".
 enum ImagenAspectRatio {
   /// Square (1:1).
   square1x1('1:1'),
@@ -129,7 +133,7 @@ enum ImagenAspectRatio {
 
   final String _jsonString;
 
-  /// Convert to json format
+  // ignore: public_member_api_docs
   String toJson() => _jsonString;
 
   // ignore: unused_element
@@ -208,7 +212,7 @@ final class ImagenFormat {
   /// Accepted values are 0 through 100. The default value is 75.
   final int? compressionQuality;
 
-  /// Convert to json format.
+  // ignore: public_member_api_docs
   Map<String, dynamic> toJson() => {
         'mimeType': mimeType,
         if (compressionQuality != null)
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index aa21d9d91301..9948a5ee3975 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -20,7 +20,16 @@ import 'client.dart';
 import 'imagen_api.dart';
 import 'imagen_content.dart';
 
+/// Represents a remote Imagen model with the ability to generate images using
+/// text prompts.
 ///
+/// See the [Cloud
+/// documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/image/generate-images)
+/// for more details about the image generation capabilities offered by the Imagen model.
+///
+/// > Warning: For Vertex AI in Firebase, image generation using Imagen 3 models
+/// is in Public Preview, which means that the feature is not subject to any SLA
+/// or deprecation policy and could change in backwards-incompatible ways.
 final class ImagenModel extends BaseModel {
   ImagenModel._(
       {required FirebaseApp app,
@@ -56,7 +65,7 @@ final class ImagenModel extends BaseModel {
         'negativePrompt': negativePrompt,
       if (_generationConfig?.addWatermark case final addWatermark?)
         'addWatermark': addWatermark,
-      if (_generationConfig.imageFormat case final imageFormat?)
+      if (_generationConfig?.imageFormat case final imageFormat?)
         'outputOption': imageFormat.toJson(),
       if (_safetySettings?.personFilterLevel case final personFilterLevel?)
         'personGeneration': personFilterLevel.toJson(),

From 8f217510fccec4fc5f04a7b25f68bd2f474e0cd2 Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Tue, 18 Feb 2025 11:17:47 -0800
Subject: [PATCH 21/21] fix analyzer

---
 .../example/lib/pages/image_prompt_page.dart  | 38 ++++++++++---------
 .../example/lib/pages/imagen_page.dart        |  3 +-
 .../firebase_vertexai/lib/src/base_model.dart |  9 +++++
 .../lib/src/imagen_model.dart                 |  1 +
 .../firebase_vertexai/test/imagen_test.dart   |  4 +-
 5 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
index 1a32cd370d84..0d84c5941c03 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/image_prompt_page.dart
@@ -89,24 +89,28 @@ class _ImagePromptPageState extends State<ImagePromptPage> {
                   const SizedBox.square(
                     dimension: 15,
                   ),
-                  IconButton(
-                    onPressed: () async {
-                      await _sendImagePrompt(_textController.text);
-                    },
-                    icon: Icon(
-                      Icons.image,
-                      color: Theme.of(context).colorScheme.primary,
+                  if (!_loading)
+                    IconButton(
+                      onPressed: () async {
+                        await _sendImagePrompt(_textController.text);
+                      },
+                      icon: Icon(
+                        Icons.image,
+                        color: Theme.of(context).colorScheme.primary,
+                      ),
                     ),
-                  ),
-                  IconButton(
-                    onPressed: () async {
-                      await _sendStorageUriPrompt(_textController.text);
-                    },
-                    icon: Icon(
-                      Icons.storage,
-                      color: Theme.of(context).colorScheme.primary,
-                    ),
-                  ),
+                  if (!_loading)
+                    IconButton(
+                      onPressed: () async {
+                        await _sendStorageUriPrompt(_textController.text);
+                      },
+                      icon: Icon(
+                        Icons.storage,
+                        color: Theme.of(context).colorScheme.primary,
+                      ),
+                    )
+                  else
+                    const CircularProgressIndicator(),
                 ],
               ),
             ),
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
index 42aacf6b8e36..3fa103419c41 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
@@ -14,7 +14,7 @@
 
 import 'package:flutter/material.dart';
 import 'package:firebase_vertexai/firebase_vertexai.dart';
-import 'package:firebase_storage/firebase_storage.dart';
+//import 'package:firebase_storage/firebase_storage.dart';
 import '../widgets/message_widget.dart';
 
 class ImagenPage extends StatefulWidget {
@@ -171,6 +171,7 @@ class _ImagenPageState extends State<ImagenPage> {
     }
     setState(() {
       _loading = false;
+      _scrollDown();
     });
   }
   // NOTE: Keep this API private until future release.
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
index 8347c49fddce..65e64550ebf1 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart
@@ -35,7 +35,11 @@ enum Task {
   predict,
 }
 
+/// Base class for models.
+///
+/// Do not instantiate directly.
 abstract class BaseModel {
+  // ignore: public_member_api_docs
   BaseModel({
     required String model,
     required String location,
@@ -53,7 +57,10 @@ abstract class BaseModel {
   final Uri _projectUri;
   final ApiClient _client;
 
+  /// The normalized model name.
   ({String prefix, String name}) get model => _model;
+
+  /// The API client.
   ApiClient get client => _client;
 
   /// Returns the model code for a user friendly model name.
@@ -74,6 +81,7 @@ abstract class BaseModel {
     );
   }
 
+  /// Returns a function that generates Firebase auth tokens.
   static FutureOr<Map<String, String>> Function() firebaseTokens(
       FirebaseAppCheck? appCheck, FirebaseAuth? auth) {
     return () async {
@@ -97,6 +105,7 @@ abstract class BaseModel {
     };
   }
 
+  /// Returns a URI for the given [task].
   Uri taskUri(Task task) => _projectUri.replace(
       pathSegments: _projectUri.pathSegments
           .followedBy([_model.prefix, '${_model.name}:${task.name}']));
diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
index 9948a5ee3975..631d7fa94849 100644
--- a/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart
@@ -98,6 +98,7 @@ final class ImagenModel extends BaseModel {
   /// Generates images with format of [ImagenGCSImage] based on the given
   /// prompt.
   /// Note: Keep this API private until future release.
+  // ignore: unused_element
   Future<ImagenGenerationResponse<ImagenGCSImage>> _generateImagesGCS(
     String prompt,
     String gcsUri,
diff --git a/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart b/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
index 05a8736ec95f..d030e6f89495 100644
--- a/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/test/imagen_test.dart
@@ -175,7 +175,7 @@ void main() {
         'predictions': [
           {
             'raiFilteredReason':
-                'Unable to show generated images. All images were filtered out because they violated Vertex AI\'s usage guidelines. You will not be charged for blocked images. Try rephrasing the prompt. If you think this was an error, send feedback. Support codes: 39322892, 29310472'
+                "Unable to show generated images. All images were filtered out because they violated Vertex AI's usage guidelines. You will not be charged for blocked images. Try rephrasing the prompt. If you think this was an error, send feedback. Support codes: 39322892, 29310472"
           }
         ]
       };
@@ -229,7 +229,7 @@ void main() {
         'error': {
           'code': 400,
           'message':
-              'Image generation failed with the following error: The prompt could not be submitted. This prompt contains sensitive words that violate Google\'s Responsible AI practices. Try rephrasing the prompt. If you think this was an error, send feedback. Support codes: 42876398',
+              "Image generation failed with the following error: The prompt could not be submitted. This prompt contains sensitive words that violate Google's Responsible AI practices. Try rephrasing the prompt. If you think this was an error, send feedback. Support codes: 42876398",
           'status': 'INVALID_ARGUMENT'
         }
       };