Subscription dashboard improvements, AI text generation limit, and other fixes.
This commit is contained in:
ajaysi
146
backend/scripts/add_ai_text_generation_limit_column.py
Normal file
146
backend/scripts/add_ai_text_generation_limit_column.py
Normal file
@@ -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)
|
||||
|
||||
210
backend/scripts/cap_basic_plan_usage.py
Normal file
210
backend/scripts/cap_basic_plan_usage.py
Normal file
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
Standalone script to cap usage counters at new Basic plan limits.
|
||||
|
||||
This preserves historical usage data but caps it at the new limits so users
|
||||
can continue making new calls within their 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
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from loguru import logger
|
||||
|
||||
from models.subscription_models import SubscriptionPlan, SubscriptionTier, UserSubscription, UsageSummary, UsageStatus
|
||||
from services.database import DATABASE_URL
|
||||
from services.subscription import PricingService
|
||||
|
||||
def cap_basic_plan_usage():
|
||||
"""Cap usage counters at new Basic plan limits."""
|
||||
|
||||
try:
|
||||
engine = create_engine(DATABASE_URL, echo=False)
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
db = SessionLocal()
|
||||
|
||||
try:
|
||||
# Find Basic plan
|
||||
basic_plan = db.query(SubscriptionPlan).filter(
|
||||
SubscriptionPlan.tier == SubscriptionTier.BASIC
|
||||
).first()
|
||||
|
||||
if not basic_plan:
|
||||
logger.error("❌ Basic plan not found in database!")
|
||||
return False
|
||||
|
||||
# New limits
|
||||
new_call_limit = basic_plan.gemini_calls_limit # Should be 10
|
||||
new_token_limit = basic_plan.gemini_tokens_limit # Should be 2000
|
||||
new_image_limit = basic_plan.stability_calls_limit # Should be 5
|
||||
|
||||
logger.info(f"📋 Basic Plan Limits:")
|
||||
logger.info(f" Calls: {new_call_limit}")
|
||||
logger.info(f" Tokens: {new_token_limit}")
|
||||
logger.info(f" Images: {new_image_limit}")
|
||||
|
||||
# Get all Basic plan users
|
||||
user_subscriptions = db.query(UserSubscription).filter(
|
||||
UserSubscription.plan_id == basic_plan.id,
|
||||
UserSubscription.is_active == True
|
||||
).all()
|
||||
|
||||
logger.info(f"\n👥 Found {len(user_subscriptions)} Basic plan user(s)")
|
||||
|
||||
pricing_service = PricingService(db)
|
||||
capped_count = 0
|
||||
|
||||
for sub in user_subscriptions:
|
||||
try:
|
||||
# Get current billing period for this user
|
||||
current_period = pricing_service.get_current_billing_period(sub.user_id) or datetime.now(timezone.utc).strftime("%Y-%m")
|
||||
|
||||
# Find usage summary for current period
|
||||
usage_summary = db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == sub.user_id,
|
||||
UsageSummary.billing_period == current_period
|
||||
).first()
|
||||
|
||||
if usage_summary:
|
||||
# Store old values for logging
|
||||
old_gemini = usage_summary.gemini_calls or 0
|
||||
old_mistral = usage_summary.mistral_calls or 0
|
||||
old_openai = usage_summary.openai_calls or 0
|
||||
old_anthropic = usage_summary.anthropic_calls or 0
|
||||
old_tokens = max(
|
||||
usage_summary.gemini_tokens or 0,
|
||||
usage_summary.openai_tokens or 0,
|
||||
usage_summary.anthropic_tokens or 0,
|
||||
usage_summary.mistral_tokens or 0
|
||||
)
|
||||
old_images = usage_summary.stability_calls or 0
|
||||
|
||||
# Check if capping is needed
|
||||
needs_cap = (
|
||||
old_gemini > new_call_limit or
|
||||
old_mistral > new_call_limit or
|
||||
old_openai > new_call_limit or
|
||||
old_anthropic > new_call_limit or
|
||||
old_images > new_image_limit or
|
||||
old_tokens > new_token_limit
|
||||
)
|
||||
|
||||
if needs_cap:
|
||||
# Cap LLM provider counters at new limits
|
||||
usage_summary.gemini_calls = min(old_gemini, new_call_limit)
|
||||
usage_summary.mistral_calls = min(old_mistral, new_call_limit)
|
||||
usage_summary.openai_calls = min(old_openai, new_call_limit)
|
||||
usage_summary.anthropic_calls = min(old_anthropic, new_call_limit)
|
||||
|
||||
# Cap token counters at new limits
|
||||
usage_summary.gemini_tokens = min(usage_summary.gemini_tokens or 0, new_token_limit)
|
||||
usage_summary.openai_tokens = min(usage_summary.openai_tokens or 0, new_token_limit)
|
||||
usage_summary.anthropic_tokens = min(usage_summary.anthropic_tokens or 0, new_token_limit)
|
||||
usage_summary.mistral_tokens = min(usage_summary.mistral_tokens or 0, new_token_limit)
|
||||
|
||||
# Cap image counter at new limit
|
||||
usage_summary.stability_calls = min(old_images, new_image_limit)
|
||||
|
||||
# Recalculate totals based on capped values
|
||||
total_capped_calls = (
|
||||
usage_summary.gemini_calls +
|
||||
usage_summary.mistral_calls +
|
||||
usage_summary.openai_calls +
|
||||
usage_summary.anthropic_calls +
|
||||
usage_summary.stability_calls
|
||||
)
|
||||
total_capped_tokens = (
|
||||
usage_summary.gemini_tokens +
|
||||
usage_summary.mistral_tokens +
|
||||
usage_summary.openai_tokens +
|
||||
usage_summary.anthropic_tokens
|
||||
)
|
||||
|
||||
usage_summary.total_calls = total_capped_calls
|
||||
usage_summary.total_tokens = total_capped_tokens
|
||||
|
||||
# Reset status to active to allow new calls
|
||||
usage_summary.usage_status = UsageStatus.ACTIVE
|
||||
usage_summary.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
db.commit()
|
||||
capped_count += 1
|
||||
|
||||
logger.info(f"\n✅ Capped usage for user {sub.user_id} (period {current_period}):")
|
||||
logger.info(f" Gemini Calls: {old_gemini} → {usage_summary.gemini_calls} (limit: {new_call_limit})")
|
||||
logger.info(f" Mistral Calls: {old_mistral} → {usage_summary.mistral_calls} (limit: {new_call_limit})")
|
||||
logger.info(f" OpenAI Calls: {old_openai} → {usage_summary.openai_calls} (limit: {new_call_limit})")
|
||||
logger.info(f" Anthropic Calls: {old_anthropic} → {usage_summary.anthropic_calls} (limit: {new_call_limit})")
|
||||
logger.info(f" Tokens: {old_tokens} → {max(usage_summary.gemini_tokens, usage_summary.mistral_tokens)} (limit: {new_token_limit})")
|
||||
logger.info(f" Images: {old_images} → {usage_summary.stability_calls} (limit: {new_image_limit})")
|
||||
else:
|
||||
logger.info(f" ℹ️ User {sub.user_id} usage is within limits - no capping needed")
|
||||
else:
|
||||
logger.info(f" ℹ️ No usage summary found for user {sub.user_id} (period {current_period})")
|
||||
|
||||
except Exception as cap_error:
|
||||
logger.error(f" ❌ Error capping usage for user {sub.user_id}: {cap_error}")
|
||||
import traceback
|
||||
logger.error(traceback.format_exc())
|
||||
db.rollback()
|
||||
|
||||
if capped_count > 0:
|
||||
logger.info(f"\n✅ Successfully capped usage for {capped_count} user(s)")
|
||||
logger.info(" Historical usage preserved, but capped at new limits")
|
||||
logger.info(" Users can now make new calls within their limits")
|
||||
else:
|
||||
logger.info("\nℹ️ No usage counters needed capping")
|
||||
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("CAPPING COMPLETE")
|
||||
logger.info("="*60)
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"❌ Error capping usage: {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 Basic plan usage capping...")
|
||||
logger.info("="*60)
|
||||
logger.info("This will cap usage counters at new Basic plan limits")
|
||||
logger.info("while preserving historical usage data.")
|
||||
logger.info("="*60)
|
||||
|
||||
try:
|
||||
success = cap_basic_plan_usage()
|
||||
|
||||
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)
|
||||
|
||||
168
backend/scripts/reset_basic_plan_usage.py
Normal file
168
backend/scripts/reset_basic_plan_usage.py
Normal file
@@ -0,0 +1,168 @@
|
||||
"""
|
||||
Quick script to reset usage counters for Basic plan users.
|
||||
|
||||
This fixes the issue where plan limits were updated but old usage data remained.
|
||||
Resets all usage counters (calls, tokens, images) to 0 for the current billing period.
|
||||
"""
|
||||
|
||||
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
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from loguru import logger
|
||||
|
||||
from models.subscription_models import SubscriptionPlan, SubscriptionTier, UserSubscription, UsageSummary, UsageStatus
|
||||
from services.database import DATABASE_URL
|
||||
from services.subscription import PricingService
|
||||
|
||||
def reset_basic_plan_usage():
|
||||
"""Reset usage counters for all Basic plan users."""
|
||||
|
||||
try:
|
||||
engine = create_engine(DATABASE_URL, echo=False)
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
db = SessionLocal()
|
||||
|
||||
try:
|
||||
# Find Basic plan
|
||||
basic_plan = db.query(SubscriptionPlan).filter(
|
||||
SubscriptionPlan.tier == SubscriptionTier.BASIC
|
||||
).first()
|
||||
|
||||
if not basic_plan:
|
||||
logger.error("❌ Basic plan not found in database!")
|
||||
return False
|
||||
|
||||
# Get all Basic plan users
|
||||
user_subscriptions = db.query(UserSubscription).filter(
|
||||
UserSubscription.plan_id == basic_plan.id,
|
||||
UserSubscription.is_active == True
|
||||
).all()
|
||||
|
||||
logger.info(f"Found {len(user_subscriptions)} Basic plan user(s)")
|
||||
|
||||
pricing_service = PricingService(db)
|
||||
reset_count = 0
|
||||
|
||||
for sub in user_subscriptions:
|
||||
try:
|
||||
# Get current billing period for this user
|
||||
current_period = pricing_service.get_current_billing_period(sub.user_id) or datetime.now(timezone.utc).strftime("%Y-%m")
|
||||
|
||||
# Find usage summary for current period
|
||||
usage_summary = db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == sub.user_id,
|
||||
UsageSummary.billing_period == current_period
|
||||
).first()
|
||||
|
||||
if usage_summary:
|
||||
# Store old values for logging
|
||||
old_gemini = usage_summary.gemini_calls or 0
|
||||
old_mistral = usage_summary.mistral_calls or 0
|
||||
old_tokens = (usage_summary.mistral_tokens or 0) + (usage_summary.gemini_tokens or 0)
|
||||
old_images = usage_summary.stability_calls or 0
|
||||
old_total_calls = usage_summary.total_calls or 0
|
||||
old_total_tokens = usage_summary.total_tokens or 0
|
||||
|
||||
# Reset all LLM provider counters
|
||||
usage_summary.gemini_calls = 0
|
||||
usage_summary.openai_calls = 0
|
||||
usage_summary.anthropic_calls = 0
|
||||
usage_summary.mistral_calls = 0
|
||||
|
||||
# Reset all token counters
|
||||
usage_summary.gemini_tokens = 0
|
||||
usage_summary.openai_tokens = 0
|
||||
usage_summary.anthropic_tokens = 0
|
||||
usage_summary.mistral_tokens = 0
|
||||
|
||||
# Reset image counter
|
||||
usage_summary.stability_calls = 0
|
||||
|
||||
# Reset totals
|
||||
usage_summary.total_calls = 0
|
||||
usage_summary.total_tokens = 0
|
||||
usage_summary.total_cost = 0.0
|
||||
|
||||
# Reset status to active
|
||||
usage_summary.usage_status = UsageStatus.ACTIVE
|
||||
usage_summary.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
db.commit()
|
||||
reset_count += 1
|
||||
|
||||
logger.info(f"\n✅ Reset usage for user {sub.user_id} (period {current_period}):")
|
||||
logger.info(f" Calls: {old_gemini + old_mistral} (gemini: {old_gemini}, mistral: {old_mistral}) → 0")
|
||||
logger.info(f" Tokens: {old_tokens} → 0")
|
||||
logger.info(f" Images: {old_images} → 0")
|
||||
logger.info(f" Total Calls: {old_total_calls} → 0")
|
||||
logger.info(f" Total Tokens: {old_total_tokens} → 0")
|
||||
else:
|
||||
logger.info(f" ℹ️ No usage summary found for user {sub.user_id} (period {current_period}) - nothing to reset")
|
||||
|
||||
except Exception as reset_error:
|
||||
logger.error(f" ❌ Error resetting usage for user {sub.user_id}: {reset_error}")
|
||||
import traceback
|
||||
logger.error(traceback.format_exc())
|
||||
db.rollback()
|
||||
|
||||
if reset_count > 0:
|
||||
logger.info(f"\n✅ Successfully reset usage counters for {reset_count} user(s)")
|
||||
else:
|
||||
logger.info("\nℹ️ No usage counters to reset")
|
||||
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("RESET COMPLETE")
|
||||
logger.info("="*60)
|
||||
logger.info("\n💡 Usage counters have been reset. Users can now use their new limits.")
|
||||
logger.info(" Next API call will start counting from 0.")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"❌ Error resetting usage: {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 Basic plan usage counter reset...")
|
||||
logger.info("="*60)
|
||||
logger.info("This will reset all usage counters (calls, tokens, images) to 0")
|
||||
logger.info("for all Basic plan users in their current billing period.")
|
||||
logger.info("="*60)
|
||||
|
||||
try:
|
||||
success = reset_basic_plan_usage()
|
||||
|
||||
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)
|
||||
|
||||
279
backend/scripts/update_basic_plan_limits.py
Normal file
279
backend/scripts/update_basic_plan_limits.py
Normal file
@@ -0,0 +1,279 @@
|
||||
"""
|
||||
Script to update Basic plan subscription limits for testing rate limits and renewal flows.
|
||||
|
||||
Updates:
|
||||
- LLM Calls (all providers): 10 calls (was 500-1000)
|
||||
- LLM Tokens (all providers): 2000 tokens (was 200k-1M)
|
||||
- Images: 5 images (was 50)
|
||||
|
||||
This script updates the SubscriptionPlan table, which automatically applies to all users
|
||||
who have a Basic plan subscription via the plan_id foreign key.
|
||||
"""
|
||||
|
||||
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
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from loguru import logger
|
||||
|
||||
from models.subscription_models import SubscriptionPlan, SubscriptionTier, UserSubscription, UsageStatus
|
||||
from services.database import DATABASE_URL
|
||||
|
||||
def update_basic_plan_limits():
|
||||
"""Update Basic plan limits for testing rate limits and renewal."""
|
||||
|
||||
try:
|
||||
engine = create_engine(DATABASE_URL, echo=False)
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
db = SessionLocal()
|
||||
|
||||
try:
|
||||
# Find Basic plan
|
||||
basic_plan = db.query(SubscriptionPlan).filter(
|
||||
SubscriptionPlan.tier == SubscriptionTier.BASIC
|
||||
).first()
|
||||
|
||||
if not basic_plan:
|
||||
logger.error("❌ Basic plan not found in database!")
|
||||
return False
|
||||
|
||||
# Store old values for logging
|
||||
old_limits = {
|
||||
'gemini_calls': basic_plan.gemini_calls_limit,
|
||||
'mistral_calls': basic_plan.mistral_calls_limit,
|
||||
'gemini_tokens': basic_plan.gemini_tokens_limit,
|
||||
'mistral_tokens': basic_plan.mistral_tokens_limit,
|
||||
'stability_calls': basic_plan.stability_calls_limit,
|
||||
}
|
||||
|
||||
logger.info(f"📋 Current Basic plan limits:")
|
||||
logger.info(f" Gemini Calls: {old_limits['gemini_calls']}")
|
||||
logger.info(f" Mistral Calls: {old_limits['mistral_calls']}")
|
||||
logger.info(f" Gemini Tokens: {old_limits['gemini_tokens']}")
|
||||
logger.info(f" Mistral Tokens: {old_limits['mistral_tokens']}")
|
||||
logger.info(f" Images (Stability): {old_limits['stability_calls']}")
|
||||
|
||||
# Update unified AI text generation limit to 10
|
||||
basic_plan.ai_text_generation_calls_limit = 10
|
||||
|
||||
# Legacy per-provider limits (kept for backwards compatibility, but not used for enforcement)
|
||||
basic_plan.gemini_calls_limit = 1000
|
||||
basic_plan.openai_calls_limit = 500
|
||||
basic_plan.anthropic_calls_limit = 200
|
||||
basic_plan.mistral_calls_limit = 500
|
||||
|
||||
# Update all LLM provider token limits to 2000
|
||||
basic_plan.gemini_tokens_limit = 2000
|
||||
basic_plan.openai_tokens_limit = 2000
|
||||
basic_plan.anthropic_tokens_limit = 2000
|
||||
basic_plan.mistral_tokens_limit = 2000
|
||||
|
||||
# Update image generation limit to 5
|
||||
basic_plan.stability_calls_limit = 5
|
||||
|
||||
# Update timestamp
|
||||
basic_plan.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
logger.info("\n📝 New Basic plan limits:")
|
||||
logger.info(f" LLM Calls (all providers): 10")
|
||||
logger.info(f" LLM Tokens (all providers): 2000")
|
||||
logger.info(f" Images: 5")
|
||||
|
||||
# Count and get affected users
|
||||
user_subscriptions = db.query(UserSubscription).filter(
|
||||
UserSubscription.plan_id == basic_plan.id,
|
||||
UserSubscription.is_active == True
|
||||
).all()
|
||||
|
||||
affected_users = len(user_subscriptions)
|
||||
|
||||
logger.info(f"\n👥 Users affected: {affected_users}")
|
||||
|
||||
if affected_users > 0:
|
||||
logger.info("\n📋 Affected user IDs:")
|
||||
for sub in user_subscriptions:
|
||||
logger.info(f" - {sub.user_id}")
|
||||
else:
|
||||
logger.info(" (No active Basic plan subscriptions found)")
|
||||
|
||||
# Commit plan limit changes first
|
||||
db.commit()
|
||||
logger.info("\n✅ Basic plan limits updated successfully!")
|
||||
|
||||
# Cap usage at new limits for all affected users (preserve historical data, but cap enforcement)
|
||||
logger.info("\n🔄 Capping usage counters at new limits for Basic plan users...")
|
||||
logger.info(" (Historical usage preserved, but capped to allow new calls within limits)")
|
||||
from models.subscription_models import UsageSummary
|
||||
from services.subscription import PricingService
|
||||
|
||||
pricing_service = PricingService(db)
|
||||
capped_count = 0
|
||||
|
||||
# New limits - use unified AI text generation limit if available
|
||||
new_call_limit = getattr(basic_plan, 'ai_text_generation_calls_limit', None) or basic_plan.gemini_calls_limit
|
||||
new_token_limit = basic_plan.gemini_tokens_limit # 2000
|
||||
new_image_limit = basic_plan.stability_calls_limit # 5
|
||||
|
||||
for sub in user_subscriptions:
|
||||
try:
|
||||
# Get current billing period for this user
|
||||
current_period = pricing_service.get_current_billing_period(sub.user_id) or datetime.now(timezone.utc).strftime("%Y-%m")
|
||||
|
||||
# Find usage summary for current period
|
||||
usage_summary = db.query(UsageSummary).filter(
|
||||
UsageSummary.user_id == sub.user_id,
|
||||
UsageSummary.billing_period == current_period
|
||||
).first()
|
||||
|
||||
if usage_summary:
|
||||
# Store old values for logging
|
||||
old_gemini = usage_summary.gemini_calls or 0
|
||||
old_mistral = usage_summary.mistral_calls or 0
|
||||
old_openai = usage_summary.openai_calls or 0
|
||||
old_anthropic = usage_summary.anthropic_calls or 0
|
||||
old_tokens = max(
|
||||
usage_summary.gemini_tokens or 0,
|
||||
usage_summary.openai_tokens or 0,
|
||||
usage_summary.anthropic_tokens or 0,
|
||||
usage_summary.mistral_tokens or 0
|
||||
)
|
||||
old_images = usage_summary.stability_calls or 0
|
||||
|
||||
# Cap LLM provider counters at new limits (don't reset, just cap)
|
||||
# This allows historical data to remain but prevents blocking from old usage
|
||||
usage_summary.gemini_calls = min(old_gemini, new_call_limit)
|
||||
usage_summary.mistral_calls = min(old_mistral, new_call_limit)
|
||||
usage_summary.openai_calls = min(old_openai, new_call_limit)
|
||||
usage_summary.anthropic_calls = min(old_anthropic, new_call_limit)
|
||||
|
||||
# Cap token counters at new limits
|
||||
usage_summary.gemini_tokens = min(usage_summary.gemini_tokens or 0, new_token_limit)
|
||||
usage_summary.openai_tokens = min(usage_summary.openai_tokens or 0, new_token_limit)
|
||||
usage_summary.anthropic_tokens = min(usage_summary.anthropic_tokens or 0, new_token_limit)
|
||||
usage_summary.mistral_tokens = min(usage_summary.mistral_tokens or 0, new_token_limit)
|
||||
|
||||
# Cap image counter at new limit
|
||||
usage_summary.stability_calls = min(old_images, new_image_limit)
|
||||
|
||||
# Update totals based on capped values (approximate)
|
||||
# Recalculate total_calls and total_tokens based on capped provider values
|
||||
total_capped_calls = (
|
||||
usage_summary.gemini_calls +
|
||||
usage_summary.mistral_calls +
|
||||
usage_summary.openai_calls +
|
||||
usage_summary.anthropic_calls +
|
||||
usage_summary.stability_calls
|
||||
)
|
||||
total_capped_tokens = (
|
||||
usage_summary.gemini_tokens +
|
||||
usage_summary.mistral_tokens +
|
||||
usage_summary.openai_tokens +
|
||||
usage_summary.anthropic_tokens
|
||||
)
|
||||
|
||||
usage_summary.total_calls = total_capped_calls
|
||||
usage_summary.total_tokens = total_capped_tokens
|
||||
|
||||
# Reset status to active to allow new calls
|
||||
usage_summary.usage_status = UsageStatus.ACTIVE
|
||||
usage_summary.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
db.commit()
|
||||
capped_count += 1
|
||||
|
||||
logger.info(f" ✅ Capped usage for user {sub.user_id}:")
|
||||
logger.info(f" Gemini Calls: {old_gemini} → {usage_summary.gemini_calls} (limit: {new_call_limit})")
|
||||
logger.info(f" Mistral Calls: {old_mistral} → {usage_summary.mistral_calls} (limit: {new_call_limit})")
|
||||
logger.info(f" Tokens: {old_tokens} → {max(usage_summary.gemini_tokens, usage_summary.mistral_tokens)} (limit: {new_token_limit})")
|
||||
logger.info(f" Images: {old_images} → {usage_summary.stability_calls} (limit: {new_image_limit})")
|
||||
else:
|
||||
logger.info(f" ℹ️ No usage summary found for user {sub.user_id} (period {current_period})")
|
||||
|
||||
except Exception as cap_error:
|
||||
logger.error(f" ❌ Error capping usage for user {sub.user_id}: {cap_error}")
|
||||
import traceback
|
||||
logger.error(traceback.format_exc())
|
||||
db.rollback()
|
||||
|
||||
if capped_count > 0:
|
||||
logger.info(f"\n✅ Capped usage counters for {capped_count} user(s)")
|
||||
logger.info(" Historical usage preserved, but capped at new limits to allow new calls")
|
||||
else:
|
||||
logger.info("\nℹ️ No usage counters to cap")
|
||||
|
||||
# Note about cache clearing
|
||||
logger.info("\n🔄 Cache Information:")
|
||||
logger.info(" The subscription limits cache is per-instance and will refresh on next request.")
|
||||
logger.info(" No manual cache clearing needed - limits will be read from database on next check.")
|
||||
|
||||
# Display final summary
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("BASIC PLAN UPDATE SUMMARY")
|
||||
logger.info("="*60)
|
||||
logger.info(f"\nPlan: {basic_plan.name} ({basic_plan.tier.value})")
|
||||
logger.info(f"Price: ${basic_plan.price_monthly}/mo, ${basic_plan.price_yearly}/yr")
|
||||
logger.info(f"\nUpdated Limits:")
|
||||
logger.info(f" LLM Calls (gemini/openai/anthropic/mistral): {basic_plan.gemini_calls_limit}")
|
||||
logger.info(f" LLM Tokens (gemini/openai/anthropic/mistral): {basic_plan.gemini_tokens_limit}")
|
||||
logger.info(f" Images (stability): {basic_plan.stability_calls_limit}")
|
||||
logger.info(f"\nUsers Affected: {affected_users}")
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("\n💡 Note: These limits apply immediately to all Basic plan users.")
|
||||
logger.info(" Historical usage has been preserved but capped at new limits.")
|
||||
logger.info(" Users can continue making new calls up to the new limits.")
|
||||
logger.info(" Users will hit rate limits faster for testing purposes.")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"❌ Error updating Basic plan: {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 Basic plan limits update...")
|
||||
logger.info("="*60)
|
||||
logger.info("This will update Basic plan limits for testing rate limits:")
|
||||
logger.info(" - LLM Calls: 10 (all providers)")
|
||||
logger.info(" - LLM Tokens: 2000 (all providers)")
|
||||
logger.info(" - Images: 5")
|
||||
logger.info("="*60)
|
||||
|
||||
# Ask for confirmation in non-interactive mode, proceed directly
|
||||
# In interactive mode, you can add: input("\nPress Enter to continue or Ctrl+C to cancel...")
|
||||
|
||||
try:
|
||||
success = update_basic_plan_limits()
|
||||
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user