ai-lls-lib 1.0.0__py3-none-any.whl → 1.0.0rc1__py3-none-any.whl

ai_lls_lib/payment/credit_manager.py

@@ -0,0 +1,186 @@
+"""Credit balance management with DynamoDB."""
+
+import os
+from typing import Optional, Dict, Any
+from decimal import Decimal
+import logging
+from datetime import datetime
+
+try:
+    import boto3
+    from botocore.exceptions import ClientError
+except ImportError:
+    boto3 = None  # Handle gracefully for testing
+
+logger = logging.getLogger(__name__)
+
+
+class CreditManager:
+    """
+    Manages user credit balances in DynamoDB CreditsTable.
+    """
+
+    def __init__(self, table_name: Optional[str] = None):
+        """Initialize with DynamoDB table."""
+        if not boto3:
+            raise RuntimeError("boto3 is required for CreditManager")
+
+        self.dynamodb = boto3.resource("dynamodb")
+        self.table_name = table_name if table_name else os.environ['CREDITS_TABLE']
+
+        try:
+            self.table = self.dynamodb.Table(self.table_name)
+        except Exception as e:
+            logger.error(f"Failed to connect to DynamoDB table {self.table_name}: {e}")
+            self.table = None
+
+    def get_balance(self, user_id: str) -> int:
+        """Get current credit balance for a user."""
+        if not self.table:
+            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")
+
+        try:
+            response = self.table.get_item(Key={"user_id": user_id})
+            if "Item" in response:
+                return int(response["Item"].get("credits", 0))
+            return 0
+        except ClientError as e:
+            logger.error(f"Error getting balance for {user_id}: {e}")
+            return 0
+
+    def add_credits(self, user_id: str, amount: int) -> int:
+        """Add credits to user balance and return new balance."""
+        if not self.table:
+            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")
+
+        try:
+            response = self.table.update_item(
+                Key={"user_id": user_id},
+                UpdateExpression="ADD credits :amount SET updated_at = :now",
+                ExpressionAttributeValues={
+                    ":amount": Decimal(amount),
+                    ":now": datetime.utcnow().isoformat()
+                },
+                ReturnValues="ALL_NEW"
+            )
+            return int(response["Attributes"]["credits"])
+        except ClientError as e:
+            logger.error(f"Error adding credits for {user_id}: {e}")
+            raise
+
+    def deduct_credits(self, user_id: str, amount: int) -> bool:
+        """
+        Deduct credits from user balance.
+        Returns True if successful, False if insufficient balance.
+        """
+        if not self.table:
+            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")
+
+        try:
+            # Conditional update - only deduct if balance >= amount
+            self.table.update_item(
+                Key={"user_id": user_id},
+                UpdateExpression="ADD credits :negative_amount SET updated_at = :now",
+                ConditionExpression="credits >= :amount",
+                ExpressionAttributeValues={
+                    ":negative_amount": Decimal(-amount),
+                    ":amount": Decimal(amount),
+                    ":now": datetime.utcnow().isoformat()
+                }
+            )
+            return True
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
+                logger.info(f"Insufficient credits for {user_id}")
+                return False
+            logger.error(f"Error deducting credits for {user_id}: {e}")
+            raise
+
+    def set_subscription_state(
+        self,
+        user_id: str,
+        status: str,
+        stripe_customer_id: Optional[str] = None,
+        stripe_subscription_id: Optional[str] = None
+    ) -> None:
+        """Update subscription state in CreditsTable."""
+        if not self.table:
+            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")
+
+        try:
+            update_expr = "SET subscription_status = :status, updated_at = :now"
+            expr_values = {
+                ":status": status,
+                ":now": datetime.utcnow().isoformat()
+            }
+
+            if stripe_customer_id:
+                update_expr += ", stripe_customer_id = :customer_id"
+                expr_values[":customer_id"] = stripe_customer_id
+
+            if stripe_subscription_id:
+                update_expr += ", stripe_subscription_id = :subscription_id"
+                expr_values[":subscription_id"] = stripe_subscription_id
+
+            self.table.update_item(
+                Key={"user_id": user_id},
+                UpdateExpression=update_expr,
+                ExpressionAttributeValues=expr_values
+            )
+        except ClientError as e:
+            logger.error(f"Error updating subscription state for {user_id}: {e}")
+            raise
+
+    def get_user_payment_info(self, user_id: str) -> Dict[str, Any]:
+        """Get user's payment-related information."""
+        if not self.table:
+            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")
+
+        try:
+            response = self.table.get_item(Key={"user_id": user_id})
+            if "Item" in response:
+                item = response["Item"]
+                return {
+                    "credits": int(item.get("credits", 0)),
+                    "stripe_customer_id": item.get("stripe_customer_id"),
+                    "stripe_subscription_id": item.get("stripe_subscription_id"),
+                    "subscription_status": item.get("subscription_status")
+                }
+            return {
+                "credits": 0,
+                "stripe_customer_id": None,
+                "stripe_subscription_id": None,
+                "subscription_status": None
+            }
+        except ClientError as e:
+            logger.error(f"Error getting payment info for {user_id}: {e}")
+            return {
+                "credits": 0,
+                "stripe_customer_id": None,
+                "stripe_subscription_id": None,
+                "subscription_status": None
+            }
+
+    def has_unlimited_access(self, user_id: str) -> bool:
+        """Check if user has unlimited access via active subscription."""
+        info = self.get_user_payment_info(user_id)
+        return info.get("subscription_status") == "active"
+
+    def set_stripe_customer_id(self, user_id: str, stripe_customer_id: str) -> None:
+        """Store Stripe customer ID for a user."""
+        if not self.table:
+            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")
+
+        try:
+            self.table.update_item(
+                Key={"user_id": user_id},
+                UpdateExpression="SET stripe_customer_id = :customer_id, updated_at = :now",
+                ExpressionAttributeValues={
+                    ":customer_id": stripe_customer_id,
+                    ":now": datetime.utcnow().isoformat()
+                }
+            )
+            logger.info(f"Stored Stripe customer ID {stripe_customer_id} for user {user_id}")
+        except ClientError as e:
+            logger.error(f"Error storing Stripe customer ID for {user_id}: {e}")
+            raise

ai_lls_lib/payment/models.py

@@ -0,0 +1,102 @@
+"""Payment data models with legacy shape compatibility."""
+
+from dataclasses import dataclass
+from typing import Optional, Dict, Any
+from enum import Enum
+
+
+class PlanType(Enum):
+    """Plan types matching legacy frontend expectations."""
+    PREPAID = "prepaid"
+    POSTPAID = "postpaid"
+    INTRO = "intro"
+
+
+class SubscriptionStatus(Enum):
+    """Subscription statuses."""
+    ACTIVE = "active"
+    PAUSED = "paused"
+    CANCELLED = "cancelled"
+    PAST_DUE = "past_due"
+
+
+@dataclass
+class Plan:
+    """
+    Plan model matching legacy frontend data structure.
+    Maps from Stripe Price/Product to legacy fields.
+    """
+    plan_reference: str  # Stripe price ID or legacy reference
+    plan_type: str  # prepaid, postpaid, intro
+    plan_name: str  # STANDARD, POWER, ELITE, UNLIMITED
+    plan_subtitle: str
+    plan_amount: float  # Price in USD
+    plan_credits: Optional[int]  # Number of credits or None for unlimited
+    plan_credits_text: str  # Display text like "5,000 credits"
+    percent_off: str  # Discount percentage display text
+
+    # Additional fields for internal use
+    stripe_price_id: Optional[str] = None
+    stripe_product_id: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        result = {
+            "plan_reference": self.plan_reference,
+            "plan_type": self.plan_type,
+            "plan_name": self.plan_name,
+            "plan_subtitle": self.plan_subtitle,
+            "plan_amount": self.plan_amount,
+            "plan_credits": self.plan_credits,
+            "plan_credits_text": self.plan_credits_text,
+            "percent_off": self.percent_off
+        }
+
+        # Add variable_amount flag for VARIABLE product
+        if self.plan_name == "VARIABLE":
+            result["variable_amount"] = True
+
+        return result
+
+    @classmethod
+    def from_stripe_price(cls, price: Dict[str, Any], product: Dict[str, Any]) -> "Plan":
+        """
+        Create Plan from Stripe Price and Product objects.
+        Maps Stripe metadata to legacy fields.
+        """
+        metadata = price.get("metadata", {})
+
+        # Determine plan type
+        if price.get("recurring"):
+            plan_type = "postpaid"
+        else:
+            plan_type = metadata.get("plan_type", "prepaid")
+
+        # Extract credits
+        credits_str = metadata.get("credits", "")
+        if credits_str.lower() == "unlimited":
+            plan_credits = None
+            plan_credits_text = "Unlimited"
+        elif credits_str:
+            try:
+                plan_credits = int(credits_str)
+                plan_credits_text = f"{plan_credits:,} credits"
+            except ValueError:
+                plan_credits = None
+                plan_credits_text = credits_str
+        else:
+            plan_credits = None
+            plan_credits_text = ""
+
+        return cls(
+            plan_reference=metadata.get("plan_reference", price["id"]),
+            plan_type=plan_type,
+            plan_name=metadata.get("tier", product.get("name", "")).upper(),
+            plan_subtitle=metadata.get("plan_subtitle", product.get("description", "")),
+            plan_amount=price["unit_amount"] / 100.0,  # Convert cents to dollars
+            plan_credits=plan_credits,
+            plan_credits_text=metadata.get("plan_credits_text", plan_credits_text),
+            percent_off=metadata.get("percent_off", ""),
+            stripe_price_id=price["id"],
+            stripe_product_id=product["id"]
+        )