Subscription dashboard improvements, AI text generation limit, and other fixes.

backend/scripts/add_ai_text_generation_limit_column.py

@@ -0,0 +1,146 @@
+"""
+Migration Script: Add ai_text_generation_calls_limit column to subscription_plans table.
+
+This adds the unified AI text generation limit column that applies to all LLM providers
+(gemini, openai, anthropic, mistral) instead of per-provider limits.
+"""
+
+import sys
+import os
+from pathlib import Path
+from datetime import datetime, timezone
+
+# Add the backend directory to Python path
+backend_dir = Path(__file__).parent.parent
+sys.path.insert(0, str(backend_dir))
+
+from sqlalchemy import create_engine, text, inspect
+from sqlalchemy.orm import sessionmaker
+from loguru import logger
+
+from models.subscription_models import SubscriptionPlan, SubscriptionTier
+from services.database import DATABASE_URL
+
+def add_ai_text_generation_limit_column():
+    """Add ai_text_generation_calls_limit column to subscription_plans table."""
+    
+    try:
+        engine = create_engine(DATABASE_URL, echo=False)
+        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+        db = SessionLocal()
+        
+        try:
+            # Check if column already exists
+            inspector = inspect(engine)
+            columns = [col['name'] for col in inspector.get_columns('subscription_plans')]
+            
+            if 'ai_text_generation_calls_limit' in columns:
+                logger.info("✅ Column 'ai_text_generation_calls_limit' already exists. Skipping migration.")
+                return True
+            
+            logger.info("📋 Adding 'ai_text_generation_calls_limit' column to subscription_plans table...")
+            
+            # Add the column (SQLite compatible)
+            alter_query = text("""
+                ALTER TABLE subscription_plans 
+                ADD COLUMN ai_text_generation_calls_limit INTEGER DEFAULT 0
+            """)
+            
+            db.execute(alter_query)
+            db.commit()
+            
+            logger.info("✅ Column added successfully!")
+            
+            # Update existing plans with unified limits based on their current limits
+            logger.info("\n🔄 Updating existing subscription plans with unified limits...")
+            
+            plans = db.query(SubscriptionPlan).all()
+            updated_count = 0
+            
+            for plan in plans:
+                # Use the first non-zero LLM provider limit as the unified limit
+                # Or use gemini_calls_limit as default
+                unified_limit = (
+                    plan.ai_text_generation_calls_limit or
+                    plan.gemini_calls_limit or
+                    plan.openai_calls_limit or
+                    plan.anthropic_calls_limit or
+                    plan.mistral_calls_limit or
+                    0
+                )
+                
+                # For Basic plan, ensure it's set to 10 (from our recent update)
+                if plan.tier == SubscriptionTier.BASIC:
+                    unified_limit = 10
+                
+                if plan.ai_text_generation_calls_limit != unified_limit:
+                    plan.ai_text_generation_calls_limit = unified_limit
+                    plan.updated_at = datetime.now(timezone.utc)
+                    updated_count += 1
+                    
+                    logger.info(f"  ✅ Updated {plan.name} ({plan.tier.value}): ai_text_generation_calls_limit = {unified_limit}")
+                else:
+                    logger.info(f"  ℹ️  {plan.name} ({plan.tier.value}): already set to {unified_limit}")
+            
+            if updated_count > 0:
+                db.commit()
+                logger.info(f"\n✅ Updated {updated_count} subscription plan(s)")
+            else:
+                logger.info("\nℹ️  No plans needed updating")
+            
+            # Display summary
+            logger.info("\n" + "="*60)
+            logger.info("MIGRATION SUMMARY")
+            logger.info("="*60)
+            
+            all_plans = db.query(SubscriptionPlan).all()
+            for plan in all_plans:
+                logger.info(f"\n{plan.name} ({plan.tier.value}):")
+                logger.info(f"  Unified AI Text Gen Limit: {plan.ai_text_generation_calls_limit if plan.ai_text_generation_calls_limit else 'Not set'}")
+                logger.info(f"  Legacy Limits: gemini={plan.gemini_calls_limit}, mistral={plan.mistral_calls_limit}")
+            
+            logger.info("\n" + "="*60)
+            logger.info("✅ Migration completed successfully!")
+            
+            return True
+            
+        except Exception as e:
+            db.rollback()
+            logger.error(f"❌ Error during migration: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            raise
+            
+        finally:
+            db.close()
+            
+    except Exception as e:
+        logger.error(f"❌ Failed to connect to database: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
+        return False
+
+if __name__ == "__main__":
+    logger.info("🚀 Starting ai_text_generation_calls_limit column migration...")
+    logger.info("="*60)
+    logger.info("This will add the unified AI text generation limit column")
+    logger.info("and update existing plans with appropriate values.")
+    logger.info("="*60)
+    
+    try:
+        success = add_ai_text_generation_limit_column()
+        
+        if success:
+            logger.info("\n✅ Script completed successfully!")
+            sys.exit(0)
+        else:
+            logger.error("\n❌ Script failed!")
+            sys.exit(1)
+            
+    except KeyboardInterrupt:
+        logger.info("\n⚠️  Script cancelled by user")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"\n❌ Unexpected error: {e}")
+        sys.exit(1)
+