Add support for Gemini (#1794)

py/core/database/users.py

@@ -253,7 +253,7 @@ async def create_user(
         is_superuser: bool = False,
         name: Optional[str] = None,
         bio: Optional[str] = None,
-        profile_picture: Optional[str] = None
+        profile_picture: Optional[str] = None,
     ) -> User:
         """Create a new user."""
         # 1) Check if a user with this email already exists

py/core/main/api/v3/users_router.py

@@ -151,7 +151,7 @@ def validate_password(password: str) -> bool:
                 password=password,
                 name=name,
                 bio=bio,
-                profile_picture=profile_picture
+                profile_picture=profile_picture,
             )
 
             return registration_response
@@ -1835,7 +1835,6 @@ async def google_callback(
                 oauth_id=google_id,
                 email=email,
             )
-            print("token_response = ", token_response)
 
             # 4. Return tokens or redirect to your front-end
             #   Some people store tokens in a cookie or redirect to a front-end route passing them as a query param.

py/core/main/services/auth_service.py

@@ -48,7 +48,7 @@ async def register(
             password=password,
             name=name,
             bio=bio,
-            profile_picture=profile_picture
+            profile_picture=profile_picture,
         )
 
     @telemetry_event("SendVerificationEmail")

py/core/providers/auth/r2r_auth.py

@@ -205,7 +205,7 @@ async def register(
         google_id: Optional[str] = None,
         name: Optional[str] = None,
         bio: Optional[str] = None,
-        profile_picture: Optional[str] = None
+        profile_picture: Optional[str] = None,
     ) -> User:
         if account_type == "password":
             if not password:
@@ -233,7 +233,7 @@ async def register(
             google_id=google_id,
             name=name,
             bio=bio,
-            profile_picture=profile_picture
+            profile_picture=profile_picture,
         )
         default_collection: CollectionResponse = (
             await self.database_provider.collections_handler.create_collection(
@@ -296,9 +296,9 @@ async def send_verification_email(
         )
 
         await self.email_provider.send_verification_email(
-            to_email=user.email, 
+            to_email=user.email,
             verification_code=verification_code,
-            dynamic_template_data={"first_name": first_name}
+            dynamic_template_data={"first_name": first_name},
         )
 
         return verification_code, expiry
@@ -428,11 +428,15 @@ async def change_password(
         )
         try:
             await self.email_provider.send_password_changed_email(
-                to_email=user.email, 
-                dynamic_template_data={"first_name": user.name.split(" ")[0] or 'User'}
+                to_email=user.email,
+                dynamic_template_data={
+                    "first_name": user.name.split(" ")[0] or "User"
+                },
             )
         except Exception as e:
-            logger.error(f"Failed to send password change notification: {str(e)}")
+            logger.error(
+                f"Failed to send password change notification: {str(e)}"
+            )
 
         return {"message": "Password changed successfully"}
 
@@ -456,9 +460,9 @@ async def request_password_reset(self, email: str) -> dict[str, str]:
                 user.name.split(" ")[0] if user.name else email.split("@")[0]
             )
             await self.email_provider.send_password_reset_email(
-                to_email=email, 
-                reset_token=reset_token, 
-                dynamic_template_data={"first_name": first_name}
+                to_email=email,
+                reset_token=reset_token,
+                dynamic_template_data={"first_name": first_name},
             )
 
             return {
@@ -494,18 +498,22 @@ async def confirm_password_reset(
         await self.database_provider.users_handler.remove_reset_token(
             id=user_id
         )
-         # Get the user information
+        # Get the user information
         user = await self.database_provider.users_handler.get_user_by_id(
             id=user_id
         )
 
         try:
             await self.email_provider.send_password_changed_email(
-                to_email=user.email, 
-                dynamic_template_data={"first_name": user.name.split(" ")[0] or 'User'}
+                to_email=user.email,
+                dynamic_template_data={
+                    "first_name": user.name.split(" ")[0] or "User"
+                },
             )
         except Exception as e:
-            logger.error(f"Failed to send password change notification: {str(e)}")
+            logger.error(
+                f"Failed to send password change notification: {str(e)}"
+            )
 
         return {"message": "Password reset successfully"}
 

py/core/providers/auth/supabase.py

@@ -74,12 +74,12 @@ async def decode_token(self, token: str) -> TokenData:
 
     async def register(
         self,
-        email: str, 
+        email: str,
         password: str,
         name: Optional[str] = None,
         bio: Optional[str] = None,
-        profile_picture: Optional[str] = None
-        ) -> User:  # type: ignore
+        profile_picture: Optional[str] = None,
+    ) -> User:  # type: ignore
         # Use Supabase client to create a new user
 
         if user := self.supabase.auth.sign_up(email=email, password=password):

py/core/providers/email/console_mock.py

@@ -67,7 +67,7 @@ async def send_password_changed_email(
             Subject: Your Password Has Been Changed
             Body:
             Your password has been successfully changed.
-            
+
             For security reasons, you will need to log in again on all your devices.
             -----------------------------
             """

py/core/providers/email/sendgrid.py

@@ -218,10 +218,13 @@ async def send_password_changed_email(
         to_email: str,
         dynamic_template_data: Optional[dict] = None,
         *args,
-        **kwargs
+        **kwargs,
     ) -> None:
         try:
-            if hasattr(self, 'password_changed_template_id') and self.password_changed_template_id:
+            if (
+                hasattr(self, "password_changed_template_id")
+                and self.password_changed_template_id
+            ):
                 await self.send_email(
                     to_email=to_email,
                     template_id=self.password_changed_template_id,
@@ -233,7 +236,7 @@ async def send_password_changed_email(
                 Your password has been successfully changed.
 
                 If you did not make this change, please contact support immediately and secure your account.
-                
+
                 """
                 html_body = """
                 <div style="font-family: Arial, sans-serif; max-width: 600px; margin: 0 auto;">

py/core/providers/email/smtp.py

@@ -152,10 +152,7 @@ async def send_password_reset_email(
         )
 
     async def send_password_changed_email(
-        self, 
-        to_email: str, 
-        *args, 
-        **kwargs
+        self, to_email: str, *args, **kwargs
     ) -> None:
         body = """
         Your password has been successfully changed.

py/r2r/compose.full.yaml

@@ -328,6 +328,9 @@ services:
       - VERTEX_PROJECT=${VERTEX_PROJECT:-}
       - VERTEX_LOCATION=${VERTEX_LOCATION:-}
 
+      # Google Gemini
+      - GEMINI_API_KEY=${GEMINI_API_KEY:-}
+
       # AWS Bedrock
       - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-}
       - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-}

py/r2r/compose.full_with_replicas.yaml

@@ -326,6 +326,9 @@ services:
       - VERTEX_PROJECT=${VERTEX_PROJECT:-}
       - VERTEX_LOCATION=${VERTEX_LOCATION:-}
 
+      # Google Gemini
+      - GEMINI_API_KEY=${GEMINI_API_KEY:-}
+
       # AWS Bedrock
       - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-}
       - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-}

py/r2r/compose.yaml

@@ -84,6 +84,9 @@ services:
       - VERTEX_PROJECT=${VERTEX_PROJECT:-}
       - VERTEX_LOCATION=${VERTEX_LOCATION:-}
 
+      # Google Gemini
+      - GEMINI_API_KEY=${GEMINI_API_KEY:-}
+
       # AWS Bedrock
       - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-}
       - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-}