AI AUtomtaion, Schudelign and publishign fromt and backe end refoactr

backend/igny8_core/business/billing/models.py

@@ -828,6 +828,13 @@ class AIModelConfig(models.Model):
         help_text="basic / quality / premium - for image models"
     )
     
+    # Testing vs Live model designation
+    is_testing = models.BooleanField(
+        default=False,
+        db_index=True,
+        help_text="Testing model (cheap, for testing only). Only one per model_type can be is_testing=True."
+    )
+    
     # Image Size Configuration (for image models)
     landscape_size = models.CharField(
         max_length=20,
@@ -892,12 +899,18 @@ class AIModelConfig(models.Model):
         return self.display_name
     
     def save(self, *args, **kwargs):
-        """Ensure only one is_default per model_type"""
+        """Ensure only one is_default and one is_testing per model_type"""
         if self.is_default:
             AIModelConfig.objects.filter(
                 model_type=self.model_type,
                 is_default=True
             ).exclude(pk=self.pk).update(is_default=False)
+        if self.is_testing:
+            AIModelConfig.objects.filter(
+                model_type=self.model_type,
+                is_testing=True,
+                is_active=True
+            ).exclude(pk=self.pk).update(is_testing=False)
         super().save(*args, **kwargs)
     
     @classmethod
@@ -910,6 +923,25 @@ class AIModelConfig(models.Model):
         """Get the default image generation model"""
         return cls.objects.filter(model_type='image', is_default=True, is_active=True).first()
     
+    @classmethod
+    def get_testing_model(cls, model_type: str):
+        """Get the testing model for text or image"""
+        return cls.objects.filter(
+            model_type=model_type,
+            is_testing=True,
+            is_active=True
+        ).first()
+    
+    @classmethod
+    def get_live_model(cls, model_type: str):
+        """Get the live (default production) model for text or image"""
+        return cls.objects.filter(
+            model_type=model_type,
+            is_testing=False,
+            is_default=True,
+            is_active=True
+        ).first()
+    
     @classmethod
     def get_image_models_by_tier(cls):
         """Get all active image models grouped by quality tier"""
@@ -1044,3 +1076,121 @@ class WebhookEvent(models.Model):
         self.error_message = error_message
         self.retry_count += 1
         self.save(update_fields=['error_message', 'retry_count'])
+
+
+class SiteAIBudgetAllocation(AccountBaseModel):
+    """
+    Site-level AI budget allocation by function.
+    
+    Allows configuring what percentage of the site's credit budget
+    can be used for each AI function. This provides fine-grained
+    control over credit consumption during automation runs.
+    
+    Example: 40% content, 30% images, 20% clustering, 10% ideas
+    
+    When max_credits_per_run is set in AutomationConfig:
+    - Each function can only use up to its allocated % of that budget
+    - Prevents any single function from consuming all credits
+    """
+    
+    AI_FUNCTION_CHOICES = [
+        ('clustering', 'Keyword Clustering (Stage 1)'),
+        ('idea_generation', 'Ideas Generation (Stage 2)'),
+        ('content_generation', 'Content Generation (Stage 4)'),
+        ('image_prompt', 'Image Prompt Extraction (Stage 5)'),
+        ('image_generation', 'Image Generation (Stage 6)'),
+    ]
+    
+    site = models.ForeignKey(
+        'igny8_core_auth.Site',
+        on_delete=models.CASCADE,
+        related_name='ai_budget_allocations',
+        help_text="Site this allocation belongs to"
+    )
+    
+    ai_function = models.CharField(
+        max_length=50,
+        choices=AI_FUNCTION_CHOICES,
+        help_text="AI function to allocate budget for"
+    )
+    
+    allocation_percentage = models.PositiveIntegerField(
+        default=20,
+        validators=[MinValueValidator(0)],
+        help_text="Percentage of credit budget allocated to this function (0-100)"
+    )
+    
+    is_enabled = models.BooleanField(
+        default=True,
+        help_text="Whether this function is enabled for automation"
+    )
+    
+    created_at = models.DateTimeField(auto_now_add=True)
+    updated_at = models.DateTimeField(auto_now=True)
+    
+    class Meta:
+        app_label = 'billing'
+        db_table = 'igny8_site_ai_budget_allocations'
+        verbose_name = 'Site AI Budget Allocation'
+        verbose_name_plural = 'Site AI Budget Allocations'
+        unique_together = [['site', 'ai_function']]
+        ordering = ['site', 'ai_function']
+        indexes = [
+            models.Index(fields=['site', 'is_enabled']),
+            models.Index(fields=['account', 'site']),
+        ]
+    
+    def __str__(self):
+        return f"{self.site.name} - {self.get_ai_function_display()}: {self.allocation_percentage}%"
+    
+    @classmethod
+    def get_or_create_defaults_for_site(cls, site, account):
+        """
+        Get or create default allocations for a site.
+        Default: Equal distribution across all functions (20% each = 100%)
+        """
+        defaults = [
+            ('clustering', 15),
+            ('idea_generation', 10),
+            ('content_generation', 40),
+            ('image_prompt', 5),
+            ('image_generation', 30),
+        ]
+        
+        allocations = []
+        for ai_function, percentage in defaults:
+            allocation, _ = cls.objects.get_or_create(
+                account=account,
+                site=site,
+                ai_function=ai_function,
+                defaults={
+                    'allocation_percentage': percentage,
+                    'is_enabled': True,
+                }
+            )
+            allocations.append(allocation)
+        
+        return allocations
+    
+    @classmethod
+    def get_allocation_for_function(cls, site, ai_function) -> int:
+        """
+        Get allocation percentage for a specific AI function.
+        Returns 0 if not found or disabled.
+        """
+        try:
+            allocation = cls.objects.get(site=site, ai_function=ai_function)
+            if allocation.is_enabled:
+                return allocation.allocation_percentage
+            return 0
+        except cls.DoesNotExist:
+            # Return default percentage if no allocation exists
+            default_map = {
+                'clustering': 15,
+                'idea_generation': 10,
+                'content_generation': 40,
+                'image_prompt': 5,
+                'image_generation': 30,
+            }
+            return default_map.get(ai_function, 20)
+