211 lines
9.9 KiB
Python
211 lines
9.9 KiB
Python
"""
|
||
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 # 25
|
||
|
||
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)
|
||
|