diff --git a/src/khoj/configure.py b/src/khoj/configure.py index 53e19f71..1ca7e041 100644 --- a/src/khoj/configure.py +++ b/src/khoj/configure.py @@ -42,6 +42,7 @@ from khoj.database.adapters import ( from khoj.database.models import ClientApplication, KhojUser, ProcessLock, Subscription from khoj.processor.embeddings import CrossEncoderModel, EmbeddingsModel from khoj.routers.api_content import configure_content, configure_search +from khoj.routers.helpers import update_telemetry_state from khoj.routers.twilio import is_twilio_enabled from khoj.utils import constants, state from khoj.utils.config import SearchType @@ -165,7 +166,15 @@ class UserAuthenticationBackend(AuthenticationBackend): create_if_not_exists = request.query_params.get("create_if_not_exists") if create_if_not_exists: - user = await aget_or_create_user_by_phone_number(phone_number) + user, is_new = await aget_or_create_user_by_phone_number(phone_number) + if user and is_new: + update_telemetry_state( + request=request, + telemetry_type="api", + api="create_user", + metadata={"user_id": str(user.uuid)}, + ) + logger.log(logging.INFO, f"🥳 New User Created: {user.uuid}") else: user = await aget_user_by_phone_number(phone_number) diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py index 9687ec01..76fd31ab 100644 --- a/src/khoj/database/adapters/__init__.py +++ b/src/khoj/database/adapters/__init__.py @@ -113,13 +113,15 @@ async def get_or_create_user(token: dict) -> KhojUser: return user -async def aget_or_create_user_by_phone_number(phone_number: str) -> KhojUser: +async def aget_or_create_user_by_phone_number(phone_number: str) -> tuple[KhojUser, bool]: + is_new = False if is_none_or_empty(phone_number): - return None + return None, is_new user = await aget_user_by_phone_number(phone_number) if not user: user = await acreate_user_by_phone_number(phone_number) - return user + is_new = True + return user, is_new async def aset_user_phone_number(user: KhojUser, phone_number: str) -> KhojUser: @@ -165,8 +167,10 @@ async def acreate_user_by_phone_number(phone_number: str) -> KhojUser: return user -async def aget_or_create_user_by_email(email: str) -> KhojUser: - user, _ = await KhojUser.objects.filter(email=email).aupdate_or_create(defaults={"username": email, "email": email}) +async def aget_or_create_user_by_email(email: str) -> tuple[KhojUser, bool]: + user, is_new = await KhojUser.objects.filter(email=email).aupdate_or_create( + defaults={"username": email, "email": email} + ) await user.asave() if user: @@ -177,7 +181,7 @@ async def aget_or_create_user_by_email(email: str) -> KhojUser: if not user_subscription: await Subscription.objects.acreate(user=user, type="trial") - return user + return user, is_new async def aget_user_validated_by_email_verification_code(code: str) -> KhojUser: @@ -248,9 +252,9 @@ def get_user_subscription(email: str) -> Optional[Subscription]: async def set_user_subscription( email: str, is_recurring=None, renewal_date=None, type="standard" -) -> Optional[Subscription]: +) -> tuple[Optional[Subscription], bool]: # Get or create the user object and their subscription - user = await aget_or_create_user_by_email(email) + user, is_new = await aget_or_create_user_by_email(email) user_subscription = await Subscription.objects.filter(user=user).afirst() # Update the user subscription state @@ -262,7 +266,7 @@ async def set_user_subscription( elif renewal_date is not None: user_subscription.renewal_date = renewal_date await user_subscription.asave() - return user_subscription + return user_subscription, is_new def subscription_to_state(subscription: Subscription) -> str: diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py index 4acefe30..b175ef8a 100644 --- a/src/khoj/routers/api_chat.py +++ b/src/khoj/routers/api_chat.py @@ -305,7 +305,7 @@ def get_shared_chat( update_telemetry_state( request=request, telemetry_type="api", - api="chat_history", + api="get_shared_chat_history", **common.__dict__, ) diff --git a/src/khoj/routers/auth.py b/src/khoj/routers/auth.py index 56116d25..98702041 100644 --- a/src/khoj/routers/auth.py +++ b/src/khoj/routers/auth.py @@ -80,11 +80,19 @@ async def login_magic_link(request: Request, form: MagicLinkForm): request.session.pop("user", None) email = form.email - user = await aget_or_create_user_by_email(email) + user, is_new = await aget_or_create_user_by_email(email) unique_id = user.email_verification_code if user: await send_magic_link_email(email, unique_id, request.base_url) + if is_new: + update_telemetry_state( + request=request, + telemetry_type="api", + api="create_user", + metadata={"user_id": str(user.uuid)}, + ) + logger.log(logging.INFO, f"🥳 New User Created: {user.uuid}") return Response(status_code=200) diff --git a/src/khoj/routers/subscription.py b/src/khoj/routers/subscription.py index 2730b775..abaac675 100644 --- a/src/khoj/routers/subscription.py +++ b/src/khoj/routers/subscription.py @@ -7,6 +7,7 @@ from fastapi import APIRouter, Request from starlette.authentication import requires from khoj.database import adapters +from khoj.routers.helpers import update_telemetry_state from khoj.utils import state # Stripe integration for Khoj Cloud Subscription @@ -48,6 +49,8 @@ async def subscribe(request: Request): customer_id = subscription["customer"] customer = stripe.Customer.retrieve(customer_id) customer_email = customer["email"] + user = None + is_new = False # Handle valid stripe webhook events success = True @@ -55,7 +58,9 @@ async def subscribe(request: Request): # Mark the user as subscribed and update the next renewal date on payment subscription = stripe.Subscription.list(customer=customer_id).data[0] renewal_date = datetime.fromtimestamp(subscription["current_period_end"], tz=timezone.utc) - user = await adapters.set_user_subscription(customer_email, is_recurring=True, renewal_date=renewal_date) + user, is_new = await adapters.set_user_subscription( + customer_email, is_recurring=True, renewal_date=renewal_date + ) success = user is not None elif event_type in {"customer.subscription.updated"}: user_subscription = await sync_to_async(adapters.get_user_subscription)(customer_email) @@ -63,15 +68,24 @@ async def subscribe(request: Request): if user_subscription and user_subscription.renewal_date: # Mark user as unsubscribed or resubscribed is_recurring = not subscription["cancel_at_period_end"] - updated_user = await adapters.set_user_subscription(customer_email, is_recurring=is_recurring) - success = updated_user is not None + user, is_new = await adapters.set_user_subscription(customer_email, is_recurring=is_recurring) + success = user is not None elif event_type in {"customer.subscription.deleted"}: # Reset the user to trial state - user = await adapters.set_user_subscription( + user, is_new = await adapters.set_user_subscription( customer_email, is_recurring=False, renewal_date=False, type="trial" ) success = user is not None + if user and is_new: + update_telemetry_state( + request=request, + telemetry_type="api", + api="create_user", + metadata={"user_id": str(user.user.uuid)}, + ) + logger.log(logging.INFO, f"🥳 New User Created: {user.user.uuid}") + logger.info(f'Stripe subscription {event["type"]} for {customer_email}') return {"success": success}