Payments and JWT integration

Author: satwikkansal

.env.example

@@ -17,6 +17,9 @@ IS_PROD=false
 WHITELIST_ENABLED=false
 X_API_KEY=your_secure_api_key_here
 
+# For /suggestions endpoint via chrome extension
+JWT_SECRET=your_jwt_secret_key_here
+
 # =============================================================================
 # 🤖 AI MODELS (Required)
 # =============================================================================
@@ -96,6 +99,17 @@ JINA_API_KEY=
 # RapidAPI (for external API data and other services)
 RAPIDAPI_KEY=
 
+# =============================================================================
+# 💳 PAYMENT INTEGRATION (Optional)
+# =============================================================================
+# Dodo Payments integration for user plan management
+
+# Dodo Payments API key for customer and subscription lookups
+DODO_API_KEY=your_dodo_payments_api_key_here
+
+# Product ID for PRO plan identification in Dodo Payments
+PRO_PLAN_PRODUCT_ID=your_pro_plan_product_id_here
+
 # =============================================================================
 # 📊 MONITORING & OBSERVABILITY (Optional)
 # =============================================================================

mxtoai/api.py

@@ -50,8 +50,10 @@ class ToolName(str, Enum):
 
 
 # Enum for Rate Limit Plans
-class RateLimitPlan(Enum):
+class UserPlan(Enum):
+    FREE = "free"
     BETA = "beta"
+    PRO = "pro"
 
 
 class EmailAttachment(BaseModel):

mxtoai/schemas.py

@@ -0,0 +1,170 @@
+"""
+User plan management and Dodo Payments integration.
+
+This module provides functionality to determine user subscription plans
+by integrating with the Dodo Payments API.
+"""
+
+import os
+from datetime import datetime
+from typing import Any
+
+import httpx
+from dotenv import load_dotenv
+
+from mxtoai._logging import get_logger
+from mxtoai.schemas import UserPlan
+
+# Load environment variables
+load_dotenv()
+
+# Configure logging
+logger = get_logger(__name__)
+
+# Dodo Payments API configuration
+DODO_API_KEY = os.getenv("DODO_API_KEY")
+PRO_PLAN_PRODUCT_ID = os.getenv("PRO_PLAN_PRODUCT_ID")
+DODO_API_BASE_URL = "https://live.dodopayments.com"
+
+# HTTP client timeout configuration
+REQUEST_TIMEOUT = 30.0
+
+# HTTP status codes
+HTTP_OK = 200
+
+
+async def get_user_plan(email: str) -> UserPlan:
+    """
+    Determine user plan based on Dodo Payments subscription status.
+
+    Args:
+        email: User's email address
+
+    Returns:
+        UserPlan: The user's subscription plan (PRO or BETA)
+
+    Note:
+        Falls back to UserPlan.BETA on any errors or missing configuration.
+
+    """
+    # Check if Dodo API key is configured
+    if not DODO_API_KEY:
+        logger.warning("DODO_API_KEY not configured, falling back to BETA plan for all users")
+        return UserPlan.BETA
+
+    try:
+        # Step 1: Look up customer by email
+        customer_id = await _get_customer_id_by_email(email)
+        if not customer_id:
+            logger.info(f"No customer found for email {email}, returning BETA plan")
+            return UserPlan.BETA
+
+        # Step 2: Get active subscriptions for the customer
+        latest_subscription = await _get_latest_active_subscription(customer_id)
+        if not latest_subscription:
+            logger.info(f"No active subscriptions found for customer {customer_id}, returning BETA plan")
+            return UserPlan.BETA
+
+        # Step 3: Check if subscription matches PRO plan product ID
+        product_id = latest_subscription.get("product_id")
+        if PRO_PLAN_PRODUCT_ID and product_id == PRO_PLAN_PRODUCT_ID:
+            logger.info(f"User {email} has PRO plan subscription (product_id: {product_id})")
+            return UserPlan.PRO
+        logger.info(
+            f"User {email} subscription does not match PRO plan (product_id: {product_id}), returning BETA plan"
+        )
+
+    except Exception as e:
+        logger.error(f"Error determining user plan for {email}: {e}")
+        logger.warning(f"Falling back to BETA plan for user {email} due to error")
+
+    return UserPlan.BETA
+
+
+async def _get_customer_id_by_email(email: str) -> str | None:
+    """
+    Look up customer ID by email address using Dodo Payments API.
+
+    Args:
+        email: Customer's email address
+
+    Returns:
+        str | None: Customer ID if found, None otherwise
+
+    """
+    try:
+        async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
+            response = await client.get(
+                f"{DODO_API_BASE_URL}/customers",
+                headers={"Authorization": f"Bearer {DODO_API_KEY}", "Content-Type": "application/json"},
+                params={"email": email},
+            )
+
+            if response.status_code == HTTP_OK:
+                data = response.json()
+                customers = data.get("items", [])
+
+                if customers:
+                    customer = customers[0]  # Take the first matching customer
+                    customer_id = customer.get("customer_id")
+                    logger.debug(f"Found customer {customer_id} for email {email}")
+                    return customer_id
+                logger.debug(f"No customers found for email {email}")
+                return None
+            logger.error(f"Dodo Payments API error for customer lookup: {response.status_code} - {response.text}")
+            return None
+
+    except httpx.TimeoutException:
+        logger.error(f"Timeout while looking up customer for email {email}")
+        return None
+    except Exception as e:
+        logger.error(f"Error looking up customer for email {email}: {e}")
+        return None
+
+
+async def _get_latest_active_subscription(customer_id: str) -> dict[str, Any] | None:
+    """
+    Get the latest active subscription for a customer.
+
+    Args:
+        customer_id: Customer's ID from Dodo Payments
+
+    Returns:
+        dict | None: Latest active subscription data if found, None otherwise
+
+    """
+    try:
+        async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
+            response = await client.get(
+                f"{DODO_API_BASE_URL}/subscriptions",
+                headers={"Authorization": f"Bearer {DODO_API_KEY}", "Content-Type": "application/json"},
+                params={"customer_id": customer_id, "status": "active"},
+            )
+
+            if response.status_code == HTTP_OK:
+                data = response.json()
+                subscriptions = data.get("items", [])
+
+                if subscriptions:
+                    # Sort by created_at to get the latest subscription
+                    sorted_subscriptions = sorted(
+                        subscriptions,
+                        key=lambda x: datetime.fromisoformat(x.get("created_at", "1970-01-01T00:00:00Z")),
+                        reverse=True,
+                    )
+                    latest_subscription = sorted_subscriptions[0]
+                    logger.debug(
+                        f"Found {len(subscriptions)} active subscriptions for customer {customer_id}, using latest: {latest_subscription.get('subscription_id')}"
+                    )
+                    return latest_subscription
+                logger.debug(f"No active subscriptions found for customer {customer_id}")
+                return None
+            logger.error(f"Dodo Payments API error for subscription lookup: {response.status_code} - {response.text}")
+            return None
+
+    except httpx.TimeoutException:
+        logger.error(f"Timeout while looking up subscriptions for customer {customer_id}")
+        return None
+    except Exception as e:
+        logger.error(f"Error looking up subscriptions for customer {customer_id}: {e}")
+        return None

mxtoai/user.py

@@ -12,7 +12,7 @@
 from mxtoai.config import MAX_ATTACHMENT_SIZE_MB, MAX_ATTACHMENTS_COUNT, MAX_TOTAL_ATTACHMENTS_SIZE_MB
 from mxtoai.dependencies import processing_instructions_resolver
 from mxtoai.email_sender import generate_message_id, send_email_reply
-from mxtoai.schemas import RateLimitPlan
+from mxtoai.schemas import UserPlan
 from mxtoai.whitelist import get_whitelist_signup_url, is_email_whitelisted, trigger_automatic_verification
 
 logger = get_logger(__name__)
@@ -23,16 +23,26 @@
 
 # Rate limit settings
 RATE_LIMITS_BY_PLAN = {
-    RateLimitPlan.BETA: {
-        "hour": {"limit": 20, "period_seconds": 3600, "expiry_seconds": 3600 * 2},  # 2hr expiry for 1hr window
-        "day": {"limit": 50, "period_seconds": 86400, "expiry_seconds": 86400 + 3600},  # 25hr expiry for 24hr window
+    UserPlan.BETA: {
+        "hour": {"limit": 10, "period_seconds": 3600, "expiry_seconds": 3600 * 2},  # 2hr expiry for 1hr window
+        "day": {"limit": 30, "period_seconds": 86400, "expiry_seconds": 86400 + 3600},  # 25hr expiry for 24hr window
         "month": {
-            "limit": 300,
+            "limit": 200,
             "period_seconds": 30 * 86400,
             "expiry_seconds": (30 * 86400) + 86400,
         },  # 31day expiry for 30day window
-    }
+    },
+    UserPlan.PRO: {
+        "hour": {"limit": 50, "period_seconds": 3600, "expiry_seconds": 3600 * 2},  # 2hr expiry for 1hr window
+        "day": {"limit": 100, "period_seconds": 86400, "expiry_seconds": 86400 + 3600},  # 25hr expiry for 24hr window
+        "month": {
+            "limit": 1000,
+            "period_seconds": 30 * 86400,
+            "expiry_seconds": (30 * 86400) + 86400,
+        },
+    },
 }
+
 RATE_LIMIT_PER_DOMAIN_HOUR = {  # Consistent structure for domain limits
     "hour": {"limit": 50, "period_seconds": 3600, "expiry_seconds": 3600 * 2}
 }
@@ -182,7 +192,7 @@ async def send_rate_limit_rejection_email(
 
 
 async def validate_rate_limits(
-    from_email: str, to: str, subject: str | None, message_id: str | None, plan: RateLimitPlan
+    from_email: str, to: str, subject: str | None, message_id: str | None, plan: UserPlan
 ) -> Response | None:
     """
     Validate incoming email against defined rate limits based on the plan, using Redis.

mxtoai/validators.py

@@ -56,6 +56,7 @@ dependencies = [
     "httpx (>=0.27.0,<1.0.0)",
     "httpx-aiohttp (>=0.1.6,<0.2.0)",
     "transformers (>=4.53.1,<5.0.0)",
+    "pyjwt[crypto] (>=2.8.0,<3.0.0)",
 ]
 
 [tool.ruff]

pyproject.toml

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+"""
+Generate a test JWT token for testing the /suggestions endpoint.
+"""
+
+import os
+from datetime import datetime, timedelta, timezone
+
+import jwt
+
+# JWT configuration
+JWT_SECRET = os.getenv("JWT_SECRET", "test_secret_key_for_development_only")
+JWT_ALGORITHM = "HS256"
+
+
+def generate_test_jwt(email: str = "test@example.com", user_id: str = "test_user_123") -> str:
+    """Generate a test JWT token."""
+    # Token expires in 1 hour
+    exp = datetime.now(timezone.utc) + timedelta(hours=1)
+
+    payload = {
+        "sub": user_id,  # Subject (user ID)
+        "email": email,
+        "exp": int(exp.timestamp()),
+        "iat": int(datetime.now(timezone.utc).timestamp()),  # Issued at
+    }
+
+    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)
+
+
+if __name__ == "__main__":
+    token = generate_test_jwt()