From 3b6c6d05dd8f8bcd83f776cdc1c8fc64d3d85d13 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 19 May 2025 15:38:38 -0700 Subject: [PATCH] [Fix] Invitation Email does not include the invitation link (#10958) * fix: email invites should link to the invitation * fix: email invites should link to the invitation * fix: email invites should link to the invitation --- .../send_emails/base_email.py | 63 ++++++++++++++- ...odel_prices_and_context_window_backup.json | 14 ++++ .../send_emails/test_base_email.py | 79 +++++++++++++++++++ 3 files changed, 154 insertions(+), 2 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index d9f3ce46ba..9ea22074c0 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -18,6 +18,7 @@ from litellm.integrations.email_templates.user_invitation_email import ( ) from litellm.proxy._types import WebhookEvent from litellm.types.enterprise.enterprise_callbacks.send_emails import ( + EmailEvent, EmailParams, SendKeyCreatedEmailEvent, ) @@ -33,7 +34,9 @@ class BaseEmailLogger(CustomLogger): Send email to user after inviting them to the team """ email_params = await self._get_email_params( - user_id=event.user_id, user_email=getattr(event, "user_email", None) + email_event=EmailEvent.new_user_invitation, + user_id=event.user_id, + user_email=getattr(event, "user_email", None), ) # Implement invitation email logic using email_params @@ -68,6 +71,7 @@ class BaseEmailLogger(CustomLogger): email_params = await self._get_email_params( user_id=send_key_created_email_event.user_id, user_email=send_key_created_email_event.user_email, + email_event=EmailEvent.virtual_key_created, ) verbose_proxy_logger.debug( @@ -93,7 +97,10 @@ class BaseEmailLogger(CustomLogger): pass async def _get_email_params( - self, user_id: Optional[str] = None, user_email: Optional[str] = None + self, + email_event: EmailEvent, + user_id: Optional[str] = None, + user_email: Optional[str] = None, ) -> EmailParams: """ Get common email parameters used across different email sending methods @@ -113,6 +120,12 @@ class BaseEmailLogger(CustomLogger): f"User email not found for user_id: {user_id}. User email is required to send email." ) + # if user invited event then send invitation link + if email_event == EmailEvent.new_user_invitation: + base_url = await self._get_invitation_link( + user_id=user_id, base_url=base_url + ) + return EmailParams( logo_url=logo_url, support_contact=support_contact, @@ -148,6 +161,52 @@ class BaseEmailLogger(CustomLogger): return user_row.user_email return None + async def _get_invitation_link(self, user_id: Optional[str], base_url: str) -> str: + """ + Get invitation link for the user + """ + import asyncio + + from litellm.proxy.proxy_server import prisma_client + + ################################################################################ + ########## Sleep for 10 seconds to wait for the invitation link to be created ### + ################################################################################ + # The UI, calls /invitation/new to generate the invitation link + # We wait 10 seconds to ensure the link is created + ################################################################################ + await asyncio.sleep(10) + + if prisma_client is None: + verbose_proxy_logger.debug( + f"Prisma client not found. Unable to lookup user email for user_id: {user_id}" + ) + return base_url + + if user_id is None: + return base_url + + # get the latest invitation link for the user + invitation_rows = await prisma_client.db.litellm_invitationlink.find_many( + where={"user_id": user_id}, + orderBy={"created_at": "desc"}, + ) + if len(invitation_rows) > 0: + invitation_row = invitation_rows[0] + return self._construct_invitation_link( + invitation_id=invitation_row.id, base_url=base_url + ) + + return base_url + + def _construct_invitation_link(self, invitation_id: str, base_url: str) -> str: + """ + Construct invitation link for the user + + # http://localhost:4000/ui?invitation_id=7a096b3a-37c6-440f-9dd1-ba22e8043f6b + """ + return f"{base_url}/ui?invitation_id={invitation_id}" + async def send_email( self, from_email: str, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 649eda96e9..8ed64faceb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12731,5 +12731,19 @@ "/v1/images/generations" ], "source": "https://docs.nscale.com/docs/inference/serverless-models/current#image-models" + }, + "featherless_ai/featherless-ai/Qwerky-72B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "litellm_provider": "featherless_ai", + "mode": "chat" + }, + "featherless_ai/featherless-ai/Qwerky-QwQ-32B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "litellm_provider": "featherless_ai", + "mode": "chat" } } diff --git a/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 22d553e87c..d591a38dee 100644 --- a/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -14,6 +14,7 @@ from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( from litellm.proxy._types import Litellm_EntityType, WebhookEvent from litellm.types.enterprise.enterprise_callbacks.send_emails import ( + EmailEvent, SendKeyCreatedEmailEvent, ) @@ -199,3 +200,81 @@ async def test_send_key_created_email_no_email( # Test that it raises ValueError with pytest.raises(ValueError, match="User email not found"): await base_email_logger.send_key_created_email(event) + + +@pytest.mark.asyncio +async def test_get_invitation_link(base_email_logger): + # Mock prisma client and its response + mock_invitation_row = mock.MagicMock() + mock_invitation_row.id = "test-invitation-id" + + mock_prisma = mock.MagicMock() + + # Create an async mock for find_many + async def mock_find_many(*args, **kwargs): + return [mock_invitation_row] + + mock_prisma.db.litellm_invitationlink.find_many = mock_find_many + + with mock.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + # Test with valid user_id + result = await base_email_logger._get_invitation_link( + user_id="test-user", base_url="http://test.com" + ) + assert result == "http://test.com/ui?invitation_id=test-invitation-id" + + # Test with None user_id + result = await base_email_logger._get_invitation_link( + user_id=None, base_url="http://test.com" + ) + assert result == "http://test.com" + + # Test with no invitation links + async def mock_find_many_empty(*args, **kwargs): + return [] + + mock_prisma.db.litellm_invitationlink.find_many = mock_find_many_empty + result = await base_email_logger._get_invitation_link( + user_id="test-user", base_url="http://test.com" + ) + assert result == "http://test.com" + + +def test_construct_invitation_link(base_email_logger): + # Test invitation link construction + result = base_email_logger._construct_invitation_link( + invitation_id="test-id-123", base_url="http://test.com" + ) + assert result == "http://test.com/ui?invitation_id=test-id-123" + + +@pytest.mark.asyncio +async def test_get_email_params_user_invitation( + base_email_logger, mock_lookup_user_email +): + # Mock environment variables + with mock.patch.dict( + os.environ, + { + "EMAIL_LOGO_URL": "https://test-logo.com", + "EMAIL_SUPPORT_CONTACT": "support@test.com", + "PROXY_BASE_URL": "http://test.com", + }, + ): + # Mock invitation link + with mock.patch.object( + base_email_logger, + "_get_invitation_link", + return_value="http://test.com/ui?invitation_id=test-id", + ): + # Test with user invitation event + result = await base_email_logger._get_email_params( + email_event=EmailEvent.new_user_invitation, + user_id="test-user", + user_email="test@example.com", + ) + + assert result.logo_url == "https://test-logo.com" + assert result.support_contact == "support@test.com" + assert result.base_url == "http://test.com/ui?invitation_id=test-id" + assert result.recipient_email == "test@example.com"