iamge genration credits fixing - not fixed

backend/igny8_core/ai/tasks.py

@@ -158,7 +158,8 @@ def process_image_generation_queue(self, image_ids: list, account_id: int = None
     from igny8_core.ai.ai_core import AICore
     from igny8_core.ai.prompts import PromptRegistry
     from igny8_core.business.billing.services.credit_service import CreditService
-    
+    from igny8_core.business.billing.exceptions import InsufficientCreditsError, CreditCalculationError
+
     logger.info("=" * 80)
     logger.info(f"process_image_generation_queue STARTED")
     logger.info(f"  - Task ID: {self.request.id}")
@@ -166,7 +167,7 @@ def process_image_generation_queue(self, image_ids: list, account_id: int = None
     logger.info(f"  - Account ID: {account_id}")
     logger.info(f"  - Content ID: {content_id}")
     logger.info("=" * 80)
-    
+
     account = None
     if account_id:
         from igny8_core.auth.models import Account
@@ -285,7 +286,38 @@ def process_image_generation_queue(self, image_ids: list, account_id: int = None
     
     # Initialize AICore
     ai_core = AICore(account=account)
-    
+
+    # Credit check before processing images
+    if account:
+        logger.info(f"[process_image_generation_queue] Step 3: Checking credits for {total_images} images")
+        try:
+            required_credits = CreditService.check_credits_for_image(
+                account=account,
+                model_name=model,
+                num_images=total_images
+            )
+            logger.info(f"[process_image_generation_queue] Credit check passed: {required_credits} credits required, {account.credits} available")
+        except InsufficientCreditsError as e:
+            error_msg = str(e)
+            logger.error(f"[process_image_generation_queue] Insufficient credits: {error_msg}")
+            return {
+                'success': False,
+                'error': error_msg,
+                'error_type': 'InsufficientCreditsError',
+                'total_images': total_images,
+                'completed': 0,
+                'failed': total_images,
+                'results': []
+            }
+        except CreditCalculationError as e:
+            # Model not found or no credits_per_image configured - log warning but continue
+            # This allows backward compatibility if model not configured
+            logger.warning(f"[process_image_generation_queue] Credit calculation warning: {e}")
+            logger.warning(f"[process_image_generation_queue] Proceeding without credit check (model may not be configured)")
+        except Exception as e:
+            # Don't fail for unexpected credit check errors - log and continue
+            logger.warning(f"[process_image_generation_queue] Unexpected credit check error: {e}")
+
     # Process each image sequentially
     for index, image_id in enumerate(image_ids, 1):
         try:
@@ -712,7 +744,7 @@ def process_image_generation_queue(self, image_ids: list, account_id: int = None
                 else:
                     # Deduct credits for successful image generation
                     credits_deducted = 0
-                    cost_usd = result.get('cost_usd', 0)
+                    cost_usd = result.get('cost', 0) or result.get('cost_usd', 0)  # generate_image returns 'cost'
                     if account:
                         try:
                             credits_deducted = CreditService.deduct_credits_for_image(
@@ -802,7 +834,66 @@ def process_image_generation_queue(self, image_ids: list, account_id: int = None
     logger.info(f"  - Completed: {completed}")
     logger.info(f"  - Failed: {failed}")
     logger.info("=" * 80)
-    
+
+    # Log to AITaskLog for consistency with other AI functions
+    if account:
+        try:
+            from igny8_core.ai.models import AITaskLog
+            import time
+
+            # Calculate total cost from results
+            total_cost = sum(r.get('cost', 0) for r in results if r.get('status') == 'completed')
+
+            AITaskLog.objects.create(
+                task_id=self.request.id,
+                function_name='generate_images',
+                account=account,
+                phase='DONE' if completed > 0 else 'ERROR',
+                message=f'Generated {completed} images ({failed} failed)' if completed > 0 else f'All {total_images} images failed',
+                status='success' if completed > 0 else 'error',
+                duration=0,  # Could track actual duration if needed
+                cost=total_cost,
+                tokens=0,  # Image generation doesn't use tokens
+                request_steps=[{
+                    'phase': 'IMAGE_GENERATION',
+                    'status': 'success' if completed > 0 else 'error',
+                    'message': f'Processed {total_images} images: {completed} completed, {failed} failed'
+                }],
+                response_steps=[],
+                error=None if completed > 0 else f'All {total_images} images failed',
+                payload={'image_ids': image_ids, 'content_id': content_id},
+                result={
+                    'total_images': total_images,
+                    'completed': completed,
+                    'failed': failed,
+                    'model': model,
+                    'provider': provider
+                }
+            )
+            logger.info(f"[process_image_generation_queue] AITaskLog entry created")
+        except Exception as log_error:
+            logger.warning(f"[process_image_generation_queue] Failed to create AITaskLog: {log_error}")
+
+    # Create notification for image generation completion
+    if account:
+        try:
+            from igny8_core.business.notifications.services import NotificationService
+
+            if completed > 0:
+                NotificationService.notify_images_complete(
+                    account=account,
+                    image_count=completed
+                )
+            elif failed > 0:
+                NotificationService.notify_images_failed(
+                    account=account,
+                    error=f'{failed} images failed to generate',
+                    image_count=failed
+                )
+            logger.info(f"[process_image_generation_queue] Notification created")
+        except Exception as notif_error:
+            logger.warning(f"[process_image_generation_queue] Failed to create notification: {notif_error}")
+
     return {
         'success': True,
         'total_images': total_images,

backend/igny8_core/business/billing/services/credit_service.py

@@ -285,13 +285,13 @@ class CreditService:
     def check_credits_for_tokens(account, operation_type, estimated_tokens_input, estimated_tokens_output):
         """
         Check if account has sufficient credits based on estimated token usage.
-        
+
         Args:
             account: Account instance
             operation_type: Type of operation
             estimated_tokens_input: Estimated input tokens
             estimated_tokens_output: Estimated output tokens
-        
+
         Raises:
             InsufficientCreditsError: If account doesn't have enough credits
         """
@@ -303,6 +303,36 @@ class CreditService:
                 f"Insufficient credits. Required: {required}, Available: {account.credits}"
             )
         return True
+
+    @staticmethod
+    def check_credits_for_image(account, model_name: str, num_images: int = 1):
+        """
+        Check if account has sufficient credits for image generation.
+
+        Args:
+            account: Account instance
+            model_name: AI model name (e.g., 'dall-e-3', 'runware:97@1')
+            num_images: Number of images to generate
+
+        Raises:
+            InsufficientCreditsError: If account doesn't have enough credits
+            CreditCalculationError: If model not found or has no credits_per_image
+
+        Returns:
+            int: Required credits for the operation
+        """
+        required = CreditService.calculate_credits_for_image(model_name, num_images)
+
+        if account.credits < required:
+            raise InsufficientCreditsError(
+                f"Insufficient credits for image generation. Required: {required}, Available: {account.credits}"
+            )
+
+        logger.info(
+            f"Credit check passed for image generation: "
+            f"{num_images} images with {model_name} = {required} credits (available: {account.credits})"
+        )
+        return required
     
     @staticmethod
     @transaction.atomic