[VertexAI] Add Auth support. Fix fileData input (#1236)

* [VertexAI] Add Auth support. Fix fileData input

* Update FirebaseInterops.cs

Author: a-maurice

vertexai/src/Internal/FirebaseInterops.cs

@@ -36,8 +36,20 @@ internal static class FirebaseInterops {
   // The header used by the AppCheck token.
   private const string appCheckHeader = "X-Firebase-AppCheck";
 
+  // The various Auth types needed to retrieve the token, cached via reflection on startup.
+  private static Type _authType;
+  private static MethodInfo _authGetAuthMethod;
+  private static PropertyInfo _authCurrentUserProperty;
+  private static MethodInfo _userTokenAsyncMethod;
+  private static PropertyInfo _userTokenTaskResultProperty;
+  // Used to determine if the Auth reflection initialized successfully, and should work.
+  private static bool _authReflectionInitialized = false;
+  // The header used by the AppCheck token.
+  private const string authHeader = "Authorization";
+
   static FirebaseInterops() {
     InitializeAppCheckReflection();
+    InitializeAuthReflection();
   }
 
   private static void LogError(string message) {
@@ -150,24 +162,144 @@ internal static async Task<string> GetAppCheckTokenAsync(FirebaseApp firebaseApp
     return null;
   }
 
+  // Cache the various types and methods needed for Auth token retrieval.
+  private static void InitializeAuthReflection() {
+    const string firebaseAuthTypeName = "Firebase.Auth.FirebaseAuth, Firebase.Auth";
+    const string getTokenMethodName = "TokenAsync";
+
+    try {
+      // Set this to false, to allow easy failing out via return.
+      _authReflectionInitialized = false;
+
+      _authType = Type.GetType(firebaseAuthTypeName);
+      if (_authType == null) {
+        // Auth assembly likely not present, fine to skip
+        return;
+      }
+
+      // Get the static method GetAuth(FirebaseApp app):
+      _authGetAuthMethod = _authType.GetMethod(
+          "GetAuth", BindingFlags.Static | BindingFlags.Public, null,
+          new Type[] { typeof(FirebaseApp) }, null);
+      if (_authGetAuthMethod == null) {
+        LogError("Could not find FirebaseAuth.GetAuth method via reflection.");
+        return;
+      }
+
+      // Get the CurrentUser property from FirebaseAuth instance
+      _authCurrentUserProperty = _authType.GetProperty("CurrentUser", BindingFlags.Instance | BindingFlags.Public);
+      if (_authCurrentUserProperty == null) {
+        LogError("Could not find FirebaseAuth.CurrentUser property via reflection.");
+        return;
+      }
+
+      // This should be FirebaseUser type
+      Type userType = _authCurrentUserProperty.PropertyType;
+
+      // Get the TokenAsync(bool) method from FirebaseUser
+      _userTokenAsyncMethod = userType.GetMethod(
+          getTokenMethodName, BindingFlags.Instance | BindingFlags.Public, null,
+          new Type[] { typeof(bool) }, null);
+      if (_userTokenAsyncMethod == null) {
+        LogError($"Could not find FirebaseUser.{getTokenMethodName}(bool) method via reflection.");
+        return;
+      }
+
+      // The return type is Task<string>
+      Type tokenTaskType = _userTokenAsyncMethod.ReturnType;
+
+      // Get the Result property from Task<string>
+      _userTokenTaskResultProperty = tokenTaskType.GetProperty("Result");
+      if (_userTokenTaskResultProperty == null) {
+        LogError("Could not find Result property on Auth token Task.");
+        return;
+      }
+
+      // Check if Result property is actually a string
+      if (_userTokenTaskResultProperty.PropertyType != typeof(string)) {
+          LogError("Auth token Task's Result property is not a string, " +
+              $"but is {_userTokenTaskResultProperty.PropertyType}");
+          return;
+      }
+
+      _authReflectionInitialized = true;
+    } catch (Exception e) {
+      LogError($"Exception during static initialization of Auth reflection in FirebaseInterops: {e}");
+      _authReflectionInitialized = false;
+    }
+  }
+
+  // Gets the Auth Token, assuming there is one. Otherwise, returns null.
+  internal static async Task<string> GetAuthTokenAsync(FirebaseApp firebaseApp) {
+    // If Auth reflection failed for any reason, nothing to do.
+    if (!_authReflectionInitialized) {
+      return null;
+    }
+
+    try {
+      // Get the FirebaseAuth instance for the given FirebaseApp.
+      object authInstance = _authGetAuthMethod.Invoke(null, new object[] { firebaseApp });
+      if (authInstance == null) {
+        LogError("Failed to get FirebaseAuth instance via reflection.");
+        return null;
+      }
+
+      // Get the CurrentUser property
+      object currentUser = _authCurrentUserProperty.GetValue(authInstance);
+      if (currentUser == null) {
+        // No user logged in, so no token
+        return null;
+      }
+
+      // Invoke TokenAsync(false) - returns a Task<string>
+      object taskObject = _userTokenAsyncMethod.Invoke(currentUser, new object[] { false });
+      if (taskObject is not Task tokenTask) {
+        LogError("Invoking TokenAsync did not return a Task.");
+        return null;
+      }
+
+      // Await the task to get the token result
+      await tokenTask;
+
+      // Check for exceptions in the task
+      if (tokenTask.IsFaulted) {
+        LogError($"Error getting Auth token: {tokenTask.Exception}");
+        return null;
+      }
+
+      // Get the Result property (which is the string token)
+      return _userTokenTaskResultProperty.GetValue(tokenTask) as string;
+    } catch (Exception e) {
+      // Log any exceptions during the reflection/invocation process
+      LogError($"An error occurred while trying to fetch Auth token: {e}");
+    }
+    return null;
+  }
+
   // Adds the other Firebase tokens to the HttpRequest, as available.
   internal static async Task AddFirebaseTokensAsync(HttpRequestMessage request, FirebaseApp firebaseApp) {
-    string tokenString = await GetAppCheckTokenAsync(firebaseApp);
+    string appCheckToken = await GetAppCheckTokenAsync(firebaseApp);
+    if (!string.IsNullOrEmpty(appCheckToken)) {
+      request.Headers.Add(appCheckHeader, appCheckToken);
+    }
 
-    // Add the header if the token is valid
-    if (!string.IsNullOrEmpty(tokenString)) {
-      request.Headers.Add(appCheckHeader, tokenString);
+    string authToken = await GetAuthTokenAsync(firebaseApp);
+    if (!string.IsNullOrEmpty(authToken)) {
+      request.Headers.Add(authHeader, $"Firebase {authToken}");
     }
   }
 
   // Adds the other Firebase tokens to the WebSocket, as available.
   internal static async Task AddFirebaseTokensAsync(ClientWebSocket socket, FirebaseApp firebaseApp) {
-    string tokenString = await GetAppCheckTokenAsync(firebaseApp);
+    string appCheckToken = await GetAppCheckTokenAsync(firebaseApp);
+    if (!string.IsNullOrEmpty(appCheckToken)) {
+      socket.Options.SetRequestHeader(appCheckHeader, appCheckToken);
+    }
 
-    // Add the header if the token is valid
-    if (!string.IsNullOrEmpty(tokenString)) {
-      socket.Options.SetRequestHeader(appCheckHeader, tokenString);
-    }   
+    string authToken = await GetAuthTokenAsync(firebaseApp);
+    if (!string.IsNullOrEmpty(authToken)) {
+      socket.Options.SetRequestHeader(authHeader, $"Firebase {authToken}");
+    }
   }
 
 }

vertexai/src/ModelContent.cs

@@ -208,7 +208,7 @@ Dictionary<string, object> Part.ToJson() {
 
     Dictionary<string, object> Part.ToJson() {
       return new Dictionary<string, object>() {
-        { "inlineData", new Dictionary<string, object>() {
+        { "fileData", new Dictionary<string, object>() {
             { "mimeType", MimeType },
             { "fileUri", Uri.AbsoluteUri }
           }

vertexai/testapp/Assets/Firebase/Sample/VertexAI/UIHandlerAutomated.cs

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Uncomment to include logic to sign in to Auth as part of the tests
+//#define INCLUDE_FIREBASE_AUTH
+
 namespace Firebase.Sample.VertexAI {
   using Firebase;
   using Firebase.Extensions;
@@ -23,6 +26,9 @@ namespace Firebase.Sample.VertexAI {
   using System.Threading.Tasks;
   using Google.MiniJSON;
   using UnityEngine;
+#if INCLUDE_FIREBASE_AUTH
+  using Firebase.Auth;
+#endif
 
   // An automated version of the UIHandler that runs tests on Vertex AI in Firebase.
   public class UIHandlerAutomated : UIHandler {
@@ -58,9 +64,12 @@ protected override void Start() {
         TestChatFunctionCalling,
         TestChatBasicTextStream,
         TestCountTokens,
+        TestYoutubeLink,
       };
       // Set of tests that only run the single time.
       Func<Task>[] singleTests = {
+        TestReadFile,
+        TestReadSecureFile,
         // Internal tests for Json parsing, requires using a source library.
         InternalTestBasicReplyShort,
         InternalTestCitations,
@@ -71,7 +80,7 @@ protected override void Start() {
         InternalTestCountTokenResponse,
         InternalTestBasicResponseLongUsageMetadata,
         InternalTestGoogleAIBasicReplyShort,
-        InternalTestGoogleAICitations
+        InternalTestGoogleAICitations,
       };
 
       // Create the set of tests, combining the above lists.
@@ -608,6 +617,75 @@ async Task TestCountTokens(Backend backend) {
       Assert($"CountToken Detail TokenCount {details.TokenCount}", details.TokenCount > 0);
     }
 
+    // Test being able to provide a Youtube link to the model.
+    async Task TestYoutubeLink(Backend backend) {
+      var model = CreateGenerativeModel(backend);
+
+      GenerateContentResponse response = await model.GenerateContentAsync(
+        ModelContent.Text("I am testing Youtube input. Can you give a short description of the video that I've linked you to?"),
+        ModelContent.FileData("video/mp4", new Uri($"https://www.youtube.com/watch?v=cEr8XCnoSVY"))
+      );
+
+      Assert("Response missing candidates.", response.Candidates.Any());
+
+      Assert($"Response should have included Firebase: {response.Text}",
+          response.Text.Contains("Firebase", StringComparison.OrdinalIgnoreCase));
+    }
+
+    // Test providing a file from a GCS bucket (Firebase Storage) to the model.
+    async Task TestReadFile() {
+      // GCS is currently only supported with VertexAI.
+      var model = CreateGenerativeModel(Backend.VertexAI);
+
+      GenerateContentResponse response = await model.GenerateContentAsync(
+        ModelContent.Text("I am testing File input. Can you describe the content in the attached file?"),
+        ModelContent.FileData("text/plain", new Uri($"gs://{FirebaseApp.DefaultInstance.Options.StorageBucket}/HelloWorld.txt"))
+      );
+
+      Assert("Response missing candidates.", response.Candidates.Any());
+
+      Assert($"Response should have included 'Hello World': {response.Text}",
+          response.Text.Contains("Hello World", StringComparison.OrdinalIgnoreCase));
+    }
+
+    // Test providing a file requiring authentication from a GCS bucket (Firebase Storage) to the model.
+    // Should pass if Auth is included or not. To turn Auth on, define INCLUDE_FIREBASE_AUTH at the top of the file.
+    async Task TestReadSecureFile() {
+      // GCS is currently only supported with VertexAI.
+      var model = CreateGenerativeModel(Backend.VertexAI);
+
+#if INCLUDE_FIREBASE_AUTH
+      var authResult = await FirebaseAuth.DefaultInstance.SignInAnonymouslyAsync();
+#endif
+
+      try {
+        GenerateContentResponse response = await model.GenerateContentAsync(
+          ModelContent.Text("I am testing File input. Can you describe the image in the attached file?"),
+          ModelContent.FileData("image/png", new Uri($"gs://{FirebaseApp.DefaultInstance.Options.StorageBucket}/FCMImages/mushroom.png"))
+        );
+
+        // Without Auth, the previous call should throw an exception.
+        // With Auth, we should be able to describe the image in the file.
+        Assert("Response missing candidates.", response.Candidates.Any());
+
+        Assert($"Response should have included mushroom: {response.Text}",
+            response.Text.Contains("mushroom", StringComparison.OrdinalIgnoreCase));
+      }
+#if !INCLUDE_FIREBASE_AUTH
+      catch (VertexAIException ex) {
+        AssertEq("Expected InnerException", ex.InnerException.GetType(), typeof(HttpRequestException));
+        // TODO: Improve how Http errors are passed back to users.
+        Assert("Missing Http Status Code 403", (ex.InnerException as HttpRequestException).Message.Contains("403"));
+      }
+#endif
+      finally {
+#if INCLUDE_FIREBASE_AUTH
+        // Clean up the created user.
+        await authResult.User.DeleteAsync();
+#endif
+      }
+    }
+
     // The url prefix to use when fetching test data to use from the separate GitHub repo.
     readonly string testDataUrl =
         "https://raw.githubusercontent.com/FirebaseExtended/vertexai-sdk-test-data/3737ae1fe9c5ecbd55abdeabc273ef4f392cbf19/mock-responses/";