From c1fe474db9196c1ac488d8423ce28cc8d3296cdd Mon Sep 17 00:00:00 2001 From: Nolan Tremelling <34580718+NolanTrem@users.noreply.github.com> Date: Thu, 9 Jan 2025 16:46:43 -0800 Subject: [PATCH] Add support for Gemini (#1794) --- py/core/database/users.py | 2 +- py/core/main/api/v3/users_router.py | 3 +-- py/core/main/services/auth_service.py | 2 +- py/core/providers/auth/r2r_auth.py | 36 +++++++++++++++---------- py/core/providers/auth/supabase.py | 6 ++--- py/core/providers/email/console_mock.py | 2 +- py/core/providers/email/sendgrid.py | 9 ++++--- py/core/providers/email/smtp.py | 5 +--- py/r2r/compose.full.yaml | 3 +++ py/r2r/compose.full_with_replicas.yaml | 3 +++ py/r2r/compose.yaml | 3 +++ 11 files changed, 45 insertions(+), 29 deletions(-) diff --git a/py/core/database/users.py b/py/core/database/users.py index 0a977782f..79fea5494 100644 --- a/py/core/database/users.py +++ b/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 diff --git a/py/core/main/api/v3/users_router.py b/py/core/main/api/v3/users_router.py index a78858425..5a3e39949 100644 --- a/py/core/main/api/v3/users_router.py +++ b/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. diff --git a/py/core/main/services/auth_service.py b/py/core/main/services/auth_service.py index 213916f11..b4628566a 100644 --- a/py/core/main/services/auth_service.py +++ b/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") diff --git a/py/core/providers/auth/r2r_auth.py b/py/core/providers/auth/r2r_auth.py index 5be2a5eff..e7c6094a8 100644 --- a/py/core/providers/auth/r2r_auth.py +++ b/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"} diff --git a/py/core/providers/auth/supabase.py b/py/core/providers/auth/supabase.py index d5c1e1d17..ad6a5b700 100644 --- a/py/core/providers/auth/supabase.py +++ b/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): diff --git a/py/core/providers/email/console_mock.py b/py/core/providers/email/console_mock.py index 37a3f5880..404d3b643 100644 --- a/py/core/providers/email/console_mock.py +++ b/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. ----------------------------- """ diff --git a/py/core/providers/email/sendgrid.py b/py/core/providers/email/sendgrid.py index 646c2a767..ff370ad38 100644 --- a/py/core/providers/email/sendgrid.py +++ b/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 = """
diff --git a/py/core/providers/email/smtp.py b/py/core/providers/email/smtp.py index e05eb6c57..79ac66e5c 100644 --- a/py/core/providers/email/smtp.py +++ b/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. diff --git a/py/r2r/compose.full.yaml b/py/r2r/compose.full.yaml index 2711c62a0..368906b82 100644 --- a/py/r2r/compose.full.yaml +++ b/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:-} diff --git a/py/r2r/compose.full_with_replicas.yaml b/py/r2r/compose.full_with_replicas.yaml index 1df3361ef..6fc7b5231 100644 --- a/py/r2r/compose.full_with_replicas.yaml +++ b/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:-} diff --git a/py/r2r/compose.yaml b/py/r2r/compose.yaml index 746bb13c5..9199a7c8c 100644 --- a/py/r2r/compose.yaml +++ b/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:-}