feat(vertex): Add auth support in the vertexai (#12797)

* Add auth support in the vertexai * Update the example model name * revert to the previous color theme
firebase · May 21, 2024 · 3241c0b · 3241c0b
1 parent 8f16fd0
commit 3241c0b
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 10 deletions.
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/main.dart
@@ -88,10 +88,10 @@ class _ChatWidgetState extends State<ChatWidget> {
 
     initFirebase().then((value) {
       _model = FirebaseVertexAI.instance.generativeModel(
-        model: 'gemini-1.5-pro-preview-0409',
+        model: 'gemini-1.5-flash-preview-0514',
       );
       _functionCallModel = FirebaseVertexAI.instance.generativeModel(
-        model: 'gemini-1.5-pro-preview-0409',
+        model: 'gemini-1.5-flash-preview-0514',
         tools: [
           Tool(functionDeclarations: [exchangeRateTool]),
         ],

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart
@@ -19,6 +19,7 @@ import 'dart:convert';
 import 'dart:typed_data';
 
 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:firebase_core_platform_interface/firebase_core_platform_interface.dart'
     show FirebasePluginPlatform;

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/firebase_vertexai.dart
@@ -25,7 +25,8 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
       {required this.app,
       required this.options,
       required this.location,
-      this.appCheck})
+      this.appCheck,
+      this.auth})
       : super(app.name, 'plugins.flutter.io/firebase_vertexai');
 
   /// The [FirebaseApp] for this current [FirebaseVertexAI] instance.
@@ -35,6 +36,9 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
   /// https://firebase.google.com/docs/app-check
   FirebaseAppCheck? appCheck;
 
+  /// The optional [FirebaseAuth] for this current [FirebaseVertexAI] instance.
+  FirebaseAuth? auth;
+
   /// Configuration parameters for sending requests to the backend.
   RequestOptions options;
 
@@ -57,6 +61,7 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
   static FirebaseVertexAI instanceFor({
     FirebaseApp? app,
     FirebaseAppCheck? appCheck,
+    FirebaseAuth? auth,
     RequestOptions? options,
     String? location,
   }) {
@@ -72,7 +77,11 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
     location ??= _defaultLocation;
 
     FirebaseVertexAI newInstance = FirebaseVertexAI._(
-        app: app, options: options, location: location, appCheck: appCheck);
+        app: app,
+        options: options,
+        location: location,
+        appCheck: appCheck,
+        auth: auth);
     _cachedInstances[app.name] = newInstance;
 
     return newInstance;
@@ -100,6 +109,7 @@ class FirebaseVertexAI extends FirebasePluginPlatform {
         model: model,
         app: app,
         appCheck: appCheck,
+        auth: auth,
         location: location,
         safetySettings: safetySettings,
         generationConfig: generationConfig,

diff --git a/packages/firebase_vertexai/firebase_vertexai/lib/src/vertex_model.dart b/packages/firebase_vertexai/firebase_vertexai/lib/src/vertex_model.dart
@@ -41,6 +41,7 @@ final class GenerativeModel {
     required String location,
     required FirebaseApp app,
     FirebaseAppCheck? appCheck,
+    FirebaseAuth? auth,
     List<SafetySetting>? safetySettings,
     GenerationConfig? generationConfig,
     List<Tool>? tools,
@@ -51,7 +52,7 @@ final class GenerativeModel {
           model: _normalizeModelName(model),
           apiKey: app.options.apiKey,
           baseUri: _vertexUri(app, location),
-          requestHeaders: _appCheckToken(appCheck),
+          requestHeaders: _firebaseTokens(appCheck, auth),
           safetySettings: safetySettings != null
               ? safetySettings
                   .map((setting) => setting._toGoogleAISafetySetting())
@@ -90,16 +91,22 @@ final class GenerativeModel {
     }
   }
 
-  static FutureOr<Map<String, String>> Function() _appCheckToken(
-      FirebaseAppCheck? appCheck) {
+  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/flutter fire/$packageVersion';
       if (appCheck != null) {
-        final token = await appCheck.getToken();
-        if (token != null) {
-          headers['X-Firebase-AppCheck'] = token;
+        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'] = idToken;
         }
       }
       return headers;

diff --git a/packages/firebase_vertexai/firebase_vertexai/pubspec.yaml b/packages/firebase_vertexai/firebase_vertexai/pubspec.yaml
@@ -9,6 +9,7 @@ environment:
 
 dependencies:
   firebase_app_check: ^0.2.2+6
+  firebase_auth: ^4.19.5
   firebase_core: ^2.31.1
   firebase_core_platform_interface: ^5.0.0
   flutter: