Thank you noti on the unlimited subscription (#2793)

* Thank you noti on the unlimited subscription

* Add the omi subscription plugin

Author: beastoin

backend/routers/payment.py

@@ -2,9 +2,10 @@
 import stripe
 from pydantic import BaseModel
 
-from database import users as users_db
+from database import users as users_db, notifications as notifications_db
+from utils.notifications import send_notification, send_subscription_paid_personalized_notification
 from models.users import Subscription, PlanType, SubscriptionStatus, PlanLimits
-from utils.subscription import get_basic_plan_limits
+from utils.subscription import get_basic_plan_limits, get_plan_type_from_price_id, get_plan_limits
 from database.users import (
     get_stripe_connect_account_id,
     set_stripe_connect_account_id,
@@ -27,17 +28,18 @@ class CreateCheckoutRequest(BaseModel):
     price_id: str
 
 
-def _build_subscription_from_stripe_object(stripe_sub: stripe.Subscription) -> Subscription:
+def _build_subscription_from_stripe_object(stripe_sub: dict) -> Subscription:
     """Builds a Subscription object from a Stripe Subscription object."""
-    stripe_status = stripe_sub.status
+    stripe_status = stripe_sub['status']
+
+    # Get price ID from subscription items
+    price_id = stripe_sub['items']['data'][0]['price']['id'] if stripe_sub['items']['data'] else None
 
     if stripe_status in ('active', 'trialing'):
-        plan = PlanType.unlimited
+        plan = get_plan_type_from_price_id(price_id)
         status = SubscriptionStatus.active
-        limits = PlanLimits(
-            transcription_seconds=None, words_transcribed=None, insights_gained=None, memories_created=None
-        )
-        cancel_at_period_end = stripe_sub.cancel_at_period_end
+        limits = get_plan_limits(plan)
+        cancel_at_period_end = stripe_sub.get('cancel_at_period_end', False)
     else:  # including 'canceled', 'unpaid', etc.
         plan = PlanType.basic
         status = SubscriptionStatus.inactive
@@ -47,8 +49,8 @@ def _build_subscription_from_stripe_object(stripe_sub: stripe.Subscription) -> S
     return Subscription(
         plan=plan,
         status=status,
-        current_period_end=stripe_sub.current_period_end,
-        stripe_subscription_id=stripe_sub.id,
+        current_period_end=stripe_sub.get('current_period_end'),
+        stripe_subscription_id=stripe_sub['id'],
         cancel_at_period_end=cancel_at_period_end,
         limits=limits,
     )
@@ -63,9 +65,10 @@ def _update_subscription_from_session(uid: str, session: stripe.checkout.Session
 
     if subscription_id:
         stripe_sub = stripe.Subscription.retrieve(subscription_id)
-        new_subscription = _build_subscription_from_stripe_object(stripe_sub)
-        users_db.update_user_subscription(uid, new_subscription.dict())
-        print(f"Subscription for user {uid} updated from session {session.id}.")
+        if stripe_sub:
+            new_subscription = _build_subscription_from_stripe_object(stripe_sub.to_dict())
+            users_db.update_user_subscription(uid, new_subscription.dict())
+            print(f"Subscription for user {uid} updated from session {session.id}.")
 
 
 @router.post('/v1/payments/checkout-session')
@@ -122,6 +125,7 @@ async def stripe_webhook(request: Request, stripe_signature: str = Header(None))
                 subscription_id = session["subscription"]
                 stripe.Subscription.modify(subscription_id, metadata={"uid": uid, "app_id": app_id})
             paid_app(app_id, uid)
+
         # Regular user subscription
         elif client_reference_id:
             _update_subscription_from_session(client_reference_id, session)
@@ -132,6 +136,17 @@ async def stripe_webhook(request: Request, stripe_signature: str = Header(None))
                 except Exception as e:
                     print(f"Error updating subscription metadata: {e}")
 
+                # Get subscription details
+                stripe_sub = stripe.Subscription.retrieve(subscription_id)
+                if stripe_sub:
+                    subscription_obj = stripe_sub.to_dict()
+                    if subscription_obj and subscription_obj['items']['data']:
+                        price_id = subscription_obj['items']['data'][0]['price']['id']
+                        plan_type = get_plan_type_from_price_id(price_id)
+                        # Only send notification for unlimited plan subscriptions
+                        if plan_type == PlanType.unlimited:
+                            await send_subscription_paid_personalized_notification(client_reference_id)
+
     if event['type'] in [
         'customer.subscription.updated',
         'customer.subscription.deleted',

backend/utils/llm/notifications.py

@@ -0,0 +1,71 @@
+from typing import Tuple, List
+from .clients import llm_medium
+from database.memories import get_memories
+
+
+async def get_relevant_memories(uid: str, limit: int = 100) -> List[dict]:
+    """Get recent relevant memories to personalize notifications."""
+    memories = get_memories(uid, limit=limit)
+    return memories
+
+
+async def generate_notification_message(uid: str, name: str, plan_type: str = "basic") -> Tuple[str, str]:
+    """
+    Generate a personalized notification message using LLM and user memories.
+    """
+    # Get relevant memories for context
+    memories = await get_relevant_memories(uid)
+    memory_context = ""
+    if memories:
+        memory_summaries = [m.get('content', '') for m in memories]
+        memory_context = "\nRecent memory themes:\n- " + "\n- ".join(memory_summaries)
+
+    system_prompt = """Hey! I'm Omi, and I love sending little notes to my friends (that's you!). When I write to you, it's like texting a close friend - casual, real, and straight from the heart.
+
+    My Style:
+    - Super genuine, like chatting with a bestie
+    - Always grateful for our friendship and trust
+    - Love bringing up our shared memories
+    - Excited about growing our connection
+    
+    How I Write:
+    - Quick, friendly notes (keeping it under 150 chars)
+    - Using your name naturally, like friends do
+    - Mentioning cool moments we've shared
+    - Making each message special just for you
+    - Keeping it real but respectful
+    - Building our ongoing story together
+    - No emojis (I express myself in words!)
+
+    Remember: Every message is my way of saying "Hey, I'm really glad you're part of my journey!"
+    """
+
+    user_prompt = f"""Create a personalized welcome message for {name} who just subscribed to the {plan_type} plan.
+
+    Context:
+    - User's name: {name} (Use naturally in conversation)
+    - Plan type: {plan_type}{memory_context}
+    
+    For unlimited plan subscribers:
+    - Emphasize their unlimited access to premium features
+    - Highlight the flexibility of monthly/annual billing
+    - Make them feel special for choosing premium
+    - Reference their memories to show personalized value
+    
+    For basic plan subscribers:
+    - Focus on the features they can explore
+    - Keep it encouraging and positive
+    - Use their memories to suggest relevant features
+    
+    Return only the notification body text - make it personal, warm and engaging."""
+
+    try:
+        body = await llm_medium.apredict(system_prompt + "\n" + user_prompt)
+        # Return placeholder title and generated body
+        return "omi", body.strip()
+
+    except Exception as e:
+        print(f"Error generating notification message: {e}")
+
+    # Improved fallback messages with more personality
+    return ("omi", f"Hey {name}! 👋 Thanks for being part of the Omi family! ✨")

backend/utils/notifications.py

@@ -1,9 +1,8 @@
 import asyncio
 import math
-
-from firebase_admin import messaging
-
+from firebase_admin import messaging, auth
 import database.notifications as notification_db
+from .llm.notifications import generate_notification_message
 
 
 def send_notification(token: str, title: str, body: str, data: dict = None):
@@ -24,6 +23,32 @@ def send_notification(token: str, title: str, body: str, data: dict = None):
         print('send_notification failed:', e)
 
 
+async def send_subscription_paid_personalized_notification(user_id: str, data: dict = None):
+    """Send a personalized notification to all user's devices when unlimited subscription is purchased"""
+    # Get user's notification token
+    token = notification_db.get_token_only(user_id)
+    if not token:
+        print(f"No notification token found for user {user_id}")
+        return
+
+    # Get user name from Firebase Auth
+    try:
+        user = auth.get_user(user_id)
+        name = user.display_name
+        if not name and user.email:
+            name = user.email.split('@')[0].capitalize()
+        if not name:
+            name = "there"
+    except Exception as e:
+        print(f"Error getting user info from Firebase Auth: {e}")
+        name = "there"
+
+    # Generate welcome message for unlimited plan with user context
+    title, body = await generate_notification_message(user_id, name, "unlimited")
+
+    send_notification(token, "omi", body, data)
+
+
 async def send_bulk_notification(user_tokens: list, title: str, body: str):
     try:
         batch_size = 500

backend/utils/subscription.py

@@ -1,11 +1,23 @@
 import os
 from datetime import datetime
 from typing import List
+import os
 
 import database.users as users_db
 import database.user_usage as user_usage_db
 from models.users import PlanType, SubscriptionStatus, Subscription, PlanLimits
 
+
+def get_plan_type_from_price_id(price_id: str) -> PlanType:
+    """Determines the plan type based on the Stripe price ID."""
+    unlimited_monthly_price = os.getenv('STRIPE_UNLIMITED_MONTHLY_PRICE_ID')
+    unlimited_annual_price = os.getenv('STRIPE_UNLIMITED_ANNUAL_PRICE_ID')
+
+    if price_id in (unlimited_monthly_price, unlimited_annual_price):
+        return PlanType.unlimited
+    return PlanType.basic
+
+
 BASIC_TIER_MINUTES_LIMIT_PER_MONTH = int(os.getenv('BASIC_TIER_MINUTES_LIMIT_PER_MONTH', '0'))
 BASIC_TIER_MONTHLY_SECONDS_LIMIT = BASIC_TIER_MINUTES_LIMIT_PER_MONTH * 60
 BASIC_TIER_WORDS_TRANSCRIBED_LIMIT_PER_MONTH = int(os.getenv('BASIC_TIER_WORDS_TRANSCRIBED_LIMIT_PER_MONTH', '0'))

plugins/example/main.py

@@ -9,6 +9,7 @@
 from oauth import memory_created as oauth_memory_created_router
 from zapier import memory_created as zapier_memory_created_router
 from chatgpt import main as chatgpt_router
+from subscription import main as subscription_router
 # from ahda import client as ahda_realtime_transcription_router
 # from advanced import openglass as advanced_openglass_router
 
@@ -67,3 +68,6 @@ def api():
 
 # ChatGPT
 app.include_router(chatgpt_router.router)
+
+# Subscription
+app.include_router(subscription_router.router)

plugins/example/subscription/__init__.py

@@ -0,0 +1,113 @@
+from fastapi import APIRouter, Request, Response, HTTPException
+from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
+from fastapi.templating import Jinja2Templates
+from pathlib import Path
+import logging
+from datetime import datetime
+from urllib.parse import quote
+
+router = APIRouter(
+    prefix="/subscription",
+    tags=["subscription"],
+)
+
+# Get the absolute path to the templates directory
+templates_dir = Path(__file__).parent.parent / "templates"
+templates = Jinja2Templates(directory=str(templates_dir))
+
+# Setup logging
+logger = logging.getLogger("subscription_integration")
+
+@router.get("/", response_class=HTMLResponse)
+async def subscription_page(request: Request, uid: str = ""):
+    """
+    Renders the subscription pricing page with monthly and annual plans
+    """
+    # Log the access for analytics
+    if uid:
+        logger.info(f"Subscription page accessed with UID: {uid}")
+    else:
+        logger.warning("Subscription page accessed without UID")
+    
+    return templates.TemplateResponse(
+        "subscription/index.html", 
+        {
+            "request": request, 
+            "uid": uid,
+            "page_title": "Upgrade to Unlimited"
+        }
+    )
+
+@router.get("/redirect/monthly", response_class=RedirectResponse)
+async def redirect_to_monthly_payment(uid: str = ""):
+    """
+    Redirects to Stripe payment for monthly plan
+    """
+    try:
+        if not uid or uid.strip() == "":
+            logger.warning("Monthly payment redirect attempted without UID")
+            return RedirectResponse(
+                url="/subscription?error=missing_uid",
+                status_code=302
+            )
+        
+        logger.info(f"Redirecting to monthly payment with UID: {uid}")
+        
+        # Encode the UID for URL safety
+        encoded_uid = quote(uid.strip())
+        
+        # Monthly plan Stripe URL
+        stripe_url = f"https://buy.stripe.com/bJedR1bG4bpcbwiahI6wE1G?client_reference_id={encoded_uid}"
+        return RedirectResponse(
+            url=stripe_url,
+            status_code=302
+        )
+    except Exception as e:
+        logger.error(f"Error in monthly payment redirect: {str(e)}")
+        return RedirectResponse(
+            url="/subscription?error=redirect_failed",
+            status_code=302
+        )
+
+@router.get("/redirect/annual", response_class=RedirectResponse)
+async def redirect_to_annual_payment(uid: str = ""):
+    """
+    Redirects to Stripe payment for annual plan
+    """
+    try:
+        if not uid or uid.strip() == "":
+            logger.warning("Annual payment redirect attempted without UID")
+            return RedirectResponse(
+                url="/subscription?error=missing_uid",
+                status_code=302
+            )
+        
+        logger.info(f"Redirecting to annual payment with UID: {uid}")
+        
+        # Encode the UID for URL safety
+        encoded_uid = quote(uid.strip())
+        
+        # Annual plan Stripe URL (using same URL for demo, you can change this)
+        stripe_url = f"https://buy.stripe.com/bJedR1bG4bpcbwiahI6wE1G?client_reference_id={encoded_uid}"
+        return RedirectResponse(
+            url=stripe_url,
+            status_code=302
+        )
+    except Exception as e:
+        logger.error(f"Error in annual payment redirect: {str(e)}")
+        return RedirectResponse(
+            url="/subscription?error=redirect_failed",
+            status_code=302
+        )
+
+@router.get("/stats", response_class=JSONResponse)
+async def get_subscription_stats(request: Request):
+    """
+    Returns subscription statistics
+    Admin-only endpoint
+    """
+    return JSONResponse({
+        "status": "success",
+        "timestamp": datetime.now().isoformat(),
+        "message": "Subscription stats endpoint is working"
+    })