066b81dd2a
- CreditService.check_credits() now bypasses for: 1. System accounts (aws-admin, default-account, default) 2. Developer/admin users (if user provided) 3. Accounts with developer users (fallback for Celery tasks) - Updated check_credits_legacy() with same bypass logic - AIEngine credit check now uses updated CreditService - Fixes 52 console errors caused by credit checks blocking developers - Developers can now use AI functions without credit restrictions
313 lines
12 KiB
Python
313 lines
12 KiB
Python
"""
|
|
Credit Service for managing credit transactions and deductions
|
|
"""
|
|
from django.db import transaction
|
|
from django.utils import timezone
|
|
from .models import CreditTransaction, CreditUsageLog
|
|
from .constants import CREDIT_COSTS
|
|
from .exceptions import InsufficientCreditsError, CreditCalculationError
|
|
from igny8_core.auth.models import Account
|
|
|
|
|
|
class CreditService:
|
|
"""Service for managing credits"""
|
|
|
|
@staticmethod
|
|
def get_credit_cost(operation_type, amount=None):
|
|
"""
|
|
Get credit cost for operation.
|
|
|
|
Args:
|
|
operation_type: Type of operation (from CREDIT_COSTS)
|
|
amount: Optional amount (word count, image count, etc.)
|
|
|
|
Returns:
|
|
int: Number of credits required
|
|
|
|
Raises:
|
|
CreditCalculationError: If operation type is unknown
|
|
"""
|
|
base_cost = CREDIT_COSTS.get(operation_type, 0)
|
|
if base_cost == 0:
|
|
raise CreditCalculationError(f"Unknown operation type: {operation_type}")
|
|
|
|
# Variable cost operations
|
|
if operation_type == 'content_generation' and amount:
|
|
# Per 100 words
|
|
return max(1, int(base_cost * (amount / 100)))
|
|
elif operation_type == 'optimization' and amount:
|
|
# Per 200 words
|
|
return max(1, int(base_cost * (amount / 200)))
|
|
elif operation_type == 'image_generation' and amount:
|
|
# Per image
|
|
return base_cost * amount
|
|
elif operation_type == 'idea_generation' and amount:
|
|
# Per idea
|
|
return base_cost * amount
|
|
|
|
# Fixed cost operations
|
|
return base_cost
|
|
|
|
@staticmethod
|
|
def check_credits(account, operation_type, amount=None, user=None):
|
|
"""
|
|
Check if account has sufficient credits for an operation.
|
|
|
|
Args:
|
|
account: Account instance
|
|
operation_type: Type of operation
|
|
amount: Optional amount (word count, image count, etc.)
|
|
user: Optional user instance (for developer/admin bypass)
|
|
|
|
Raises:
|
|
InsufficientCreditsError: If account doesn't have enough credits
|
|
"""
|
|
# Bypass credit check for:
|
|
# 1. System accounts (aws-admin, default-account, default)
|
|
# 2. Developer/admin users (if user provided)
|
|
if account and account.is_system_account():
|
|
return True
|
|
|
|
if user:
|
|
try:
|
|
if hasattr(user, 'is_admin_or_developer') and user.is_admin_or_developer():
|
|
return True
|
|
if hasattr(user, 'is_system_account_user') and user.is_system_account_user():
|
|
return True
|
|
except (AttributeError, Exception):
|
|
pass
|
|
|
|
# Check if account has any developer users (fallback for Celery tasks without user context)
|
|
if account:
|
|
try:
|
|
from igny8_core.auth.models import User
|
|
if User.objects.filter(account=account, role='developer').exists():
|
|
return True
|
|
except (AttributeError, Exception):
|
|
pass
|
|
|
|
required = CreditService.get_credit_cost(operation_type, amount)
|
|
if account.credits < required:
|
|
raise InsufficientCreditsError(
|
|
f"Insufficient credits. Required: {required}, Available: {account.credits}"
|
|
)
|
|
return True
|
|
|
|
@staticmethod
|
|
def check_credits_legacy(account, required_credits, user=None):
|
|
"""
|
|
Legacy method: Check if account has enough credits (for backward compatibility).
|
|
|
|
Args:
|
|
account: Account instance
|
|
required_credits: Number of credits required
|
|
user: Optional user instance (for developer/admin bypass)
|
|
|
|
Raises:
|
|
InsufficientCreditsError: If account doesn't have enough credits
|
|
"""
|
|
# Bypass credit check for system accounts and developers
|
|
if account and account.is_system_account():
|
|
return
|
|
|
|
if user:
|
|
try:
|
|
if hasattr(user, 'is_admin_or_developer') and user.is_admin_or_developer():
|
|
return
|
|
if hasattr(user, 'is_system_account_user') and user.is_system_account_user():
|
|
return
|
|
except (AttributeError, Exception):
|
|
pass
|
|
|
|
# Check if account has any developer users (fallback for Celery tasks)
|
|
if account:
|
|
try:
|
|
from igny8_core.auth.models import User
|
|
if User.objects.filter(account=account, role='developer').exists():
|
|
return
|
|
except (AttributeError, Exception):
|
|
pass
|
|
|
|
if account.credits < required_credits:
|
|
raise InsufficientCreditsError(
|
|
f"Insufficient credits. Required: {required_credits}, Available: {account.credits}"
|
|
)
|
|
|
|
@staticmethod
|
|
@transaction.atomic
|
|
def deduct_credits(account, amount, operation_type, description, metadata=None, cost_usd=None, model_used=None, tokens_input=None, tokens_output=None, related_object_type=None, related_object_id=None):
|
|
"""
|
|
Deduct credits and log transaction.
|
|
|
|
Args:
|
|
account: Account instance
|
|
amount: Number of credits to deduct
|
|
operation_type: Type of operation (from CreditUsageLog.OPERATION_TYPE_CHOICES)
|
|
description: Description of the transaction
|
|
metadata: Optional metadata dict
|
|
cost_usd: Optional cost in USD
|
|
model_used: Optional AI model used
|
|
tokens_input: Optional input tokens
|
|
tokens_output: Optional output tokens
|
|
related_object_type: Optional related object type
|
|
related_object_id: Optional related object ID
|
|
|
|
Returns:
|
|
int: New credit balance
|
|
"""
|
|
# Check sufficient credits (legacy: amount is already calculated)
|
|
CreditService.check_credits_legacy(account, amount)
|
|
|
|
# Deduct from account.credits
|
|
account.credits -= amount
|
|
account.save(update_fields=['credits'])
|
|
|
|
# Create CreditTransaction
|
|
CreditTransaction.objects.create(
|
|
account=account,
|
|
transaction_type='deduction',
|
|
amount=-amount, # Negative for deduction
|
|
balance_after=account.credits,
|
|
description=description,
|
|
metadata=metadata or {}
|
|
)
|
|
|
|
# Create CreditUsageLog
|
|
CreditUsageLog.objects.create(
|
|
account=account,
|
|
operation_type=operation_type,
|
|
credits_used=amount,
|
|
cost_usd=cost_usd,
|
|
model_used=model_used or '',
|
|
tokens_input=tokens_input,
|
|
tokens_output=tokens_output,
|
|
related_object_type=related_object_type or '',
|
|
related_object_id=related_object_id,
|
|
metadata=metadata or {}
|
|
)
|
|
|
|
return account.credits
|
|
|
|
@staticmethod
|
|
@transaction.atomic
|
|
def deduct_credits_for_operation(account, operation_type, amount=None, description=None, metadata=None, cost_usd=None, model_used=None, tokens_input=None, tokens_output=None, related_object_type=None, related_object_id=None):
|
|
"""
|
|
Deduct credits for an operation (convenience method that calculates cost automatically).
|
|
|
|
Args:
|
|
account: Account instance
|
|
operation_type: Type of operation
|
|
amount: Optional amount (word count, image count, etc.)
|
|
description: Optional description (auto-generated if not provided)
|
|
metadata: Optional metadata dict
|
|
cost_usd: Optional cost in USD
|
|
model_used: Optional AI model used
|
|
tokens_input: Optional input tokens
|
|
tokens_output: Optional output tokens
|
|
related_object_type: Optional related object type
|
|
related_object_id: Optional related object ID
|
|
|
|
Returns:
|
|
int: New credit balance
|
|
"""
|
|
# Calculate credit cost
|
|
credits_required = CreditService.get_credit_cost(operation_type, amount)
|
|
|
|
# Check sufficient credits
|
|
CreditService.check_credits(account, operation_type, amount)
|
|
|
|
# Auto-generate description if not provided
|
|
if not description:
|
|
if operation_type == 'clustering':
|
|
description = f"Clustering operation"
|
|
elif operation_type == 'idea_generation':
|
|
description = f"Generated {amount or 1} idea(s)"
|
|
elif operation_type == 'content_generation':
|
|
description = f"Generated content ({amount or 0} words)"
|
|
elif operation_type == 'image_generation':
|
|
description = f"Generated {amount or 1} image(s)"
|
|
else:
|
|
description = f"{operation_type} operation"
|
|
|
|
return CreditService.deduct_credits(
|
|
account=account,
|
|
amount=credits_required,
|
|
operation_type=operation_type,
|
|
description=description,
|
|
metadata=metadata,
|
|
cost_usd=cost_usd,
|
|
model_used=model_used,
|
|
tokens_input=tokens_input,
|
|
tokens_output=tokens_output,
|
|
related_object_type=related_object_type,
|
|
related_object_id=related_object_id
|
|
)
|
|
|
|
@staticmethod
|
|
@transaction.atomic
|
|
def add_credits(account, amount, transaction_type, description, metadata=None):
|
|
"""
|
|
Add credits (purchase, subscription, etc.).
|
|
|
|
Args:
|
|
account: Account instance
|
|
amount: Number of credits to add
|
|
transaction_type: Type of transaction (from CreditTransaction.TRANSACTION_TYPE_CHOICES)
|
|
description: Description of the transaction
|
|
metadata: Optional metadata dict
|
|
|
|
Returns:
|
|
int: New credit balance
|
|
"""
|
|
# Add to account.credits
|
|
account.credits += amount
|
|
account.save(update_fields=['credits'])
|
|
|
|
# Create CreditTransaction
|
|
CreditTransaction.objects.create(
|
|
account=account,
|
|
transaction_type=transaction_type,
|
|
amount=amount, # Positive for addition
|
|
balance_after=account.credits,
|
|
description=description,
|
|
metadata=metadata or {}
|
|
)
|
|
|
|
return account.credits
|
|
|
|
@staticmethod
|
|
def calculate_credits_for_operation(operation_type, **kwargs):
|
|
"""
|
|
Calculate credits needed for an operation.
|
|
Legacy method - use get_credit_cost() instead.
|
|
|
|
Args:
|
|
operation_type: Type of operation
|
|
**kwargs: Operation-specific parameters
|
|
|
|
Returns:
|
|
int: Number of credits required
|
|
|
|
Raises:
|
|
CreditCalculationError: If calculation fails
|
|
"""
|
|
# Map legacy operation types
|
|
if operation_type == 'ideas':
|
|
operation_type = 'idea_generation'
|
|
elif operation_type == 'content':
|
|
operation_type = 'content_generation'
|
|
elif operation_type == 'images':
|
|
operation_type = 'image_generation'
|
|
|
|
# Extract amount from kwargs
|
|
amount = None
|
|
if 'word_count' in kwargs:
|
|
amount = kwargs.get('word_count')
|
|
elif 'image_count' in kwargs:
|
|
amount = kwargs.get('image_count')
|
|
elif 'idea_count' in kwargs:
|
|
amount = kwargs.get('idea_count')
|
|
|
|
return CreditService.get_credit_cost(operation_type, amount)
|
|
|