Subscription dashboard improvements, AI text generation limit, and other fixes.

backend/services/subscription/preflight_validator.py

@@ -0,0 +1,189 @@
+"""
+Pre-flight Validation Utility for Multi-Operation Workflows
+
+Provides transparent validation for operations that involve multiple API calls.
+Services can use this to validate entire workflows before making any external API calls.
+"""
+
+from typing import Dict, Any, List, Optional, Tuple
+from fastapi import HTTPException
+from loguru import logger
+
+from services.subscription.pricing_service import PricingService
+from models.subscription_models import APIProvider
+
+
+def validate_research_operations(
+    pricing_service: PricingService,
+    user_id: str,
+    gpt_provider: str = "google"
+) -> None:
+    """
+    Validate all operations for a research workflow before making ANY API calls.
+    
+    This prevents wasteful external API calls (like Google Grounding) if subsequent
+    LLM calls would fail due to token or call limits.
+    
+    Args:
+        pricing_service: PricingService instance
+        user_id: User ID for subscription checking
+        gpt_provider: GPT provider from env var (defaults to "google")
+        
+    Returns:
+        (can_proceed, error_message, error_details)
+        If can_proceed is False, raises HTTPException with 429 status
+    """
+    try:
+        # Determine actual provider for LLM calls based on GPT_PROVIDER env var
+        gpt_provider_lower = gpt_provider.lower()
+        if gpt_provider_lower == "huggingface":
+            llm_provider_enum = APIProvider.MISTRAL  # Maps to HuggingFace
+            llm_provider_name = "huggingface"
+        else:
+            llm_provider_enum = APIProvider.GEMINI
+            llm_provider_name = "gemini"
+        
+        # Estimate tokens for each operation in research workflow
+        # Google Grounding call: ~2000 tokens (input + output)
+        # Keyword analyzer: ~1000 tokens (input: 3000 chars research, output: structured JSON)
+        # Competitor analyzer: ~1000 tokens (input: 3000 chars research, output: structured JSON)
+        # Content angle generator: ~1000 tokens (input: 3000 chars research, output: list of angles)
+        
+        operations_to_validate = [
+            {
+                'provider': APIProvider.GEMINI,  # Google Grounding uses Gemini
+                'tokens_requested': 2000,
+                'actual_provider_name': 'gemini',
+                'operation_type': 'google_grounding'
+            },
+            {
+                'provider': llm_provider_enum,
+                'tokens_requested': 1000,
+                'actual_provider_name': llm_provider_name,
+                'operation_type': 'keyword_analysis'
+            },
+            {
+                'provider': llm_provider_enum,
+                'tokens_requested': 1000,
+                'actual_provider_name': llm_provider_name,
+                'operation_type': 'competitor_analysis'
+            },
+            {
+                'provider': llm_provider_enum,
+                'tokens_requested': 1000,
+                'actual_provider_name': llm_provider_name,
+                'operation_type': 'content_angle_generation'
+            }
+        ]
+        
+        logger.info(f"[Pre-flight Validator] 🚀 Starting Research Workflow Validation")
+        logger.info(f"   ├─ User: {user_id}")
+        logger.info(f"   ├─ LLM Provider: {llm_provider_name} (GPT_PROVIDER={gpt_provider})")
+        logger.info(f"   └─ Operations to validate: {len(operations_to_validate)}")
+        
+        can_proceed, message, error_details = pricing_service.check_comprehensive_limits(
+            user_id=user_id,
+            operations=operations_to_validate
+        )
+        
+        if not can_proceed:
+            usage_info = error_details.get('usage_info', {}) if error_details else {}
+            provider = usage_info.get('provider', llm_provider_name) if usage_info else llm_provider_name
+            operation_type = usage_info.get('operation_type', 'unknown')
+            
+            logger.error(f"[Pre-flight Validator] ❌ RESEARCH WORKFLOW BLOCKED")
+            logger.error(f"   ├─ User: {user_id}")
+            logger.error(f"   ├─ Blocked at: {operation_type}")
+            logger.error(f"   ├─ Provider: {provider}")
+            logger.error(f"   └─ Reason: {message}")
+            
+            # Raise HTTPException immediately - frontend gets immediate response, no API calls made
+            raise HTTPException(
+                status_code=429,
+                detail={
+                    'error': message,
+                    'message': message,
+                    'provider': provider,
+                    'usage_info': usage_info if usage_info else error_details
+                }
+            )
+        
+        logger.info(f"[Pre-flight Validator] ✅ RESEARCH WORKFLOW APPROVED")
+        logger.info(f"   ├─ User: {user_id}")
+        logger.info(f"   └─ All {len(operations_to_validate)} operations validated - proceeding with API calls")
+        # Validation passed - no return needed (function raises HTTPException if validation fails)
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"[Pre-flight Validator] Error validating research operations: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail={
+                'error': f"Failed to validate operations: {str(e)}",
+                'message': f"Failed to validate operations: {str(e)}"
+            }
+        )
+
+
+def validate_image_generation_operations(
+    pricing_service: PricingService,
+    user_id: str
+) -> None:
+    """
+    Validate image generation operation before making API calls.
+    
+    Args:
+        pricing_service: PricingService instance
+        user_id: User ID for subscription checking
+        
+    Returns:
+        (can_proceed, error_message, error_details)
+        If can_proceed is False, raises HTTPException with 429 status
+    """
+    try:
+        operations_to_validate = [
+            {
+                'provider': APIProvider.STABILITY,
+                'tokens_requested': 0,
+                'actual_provider_name': 'stability',
+                'operation_type': 'image_generation'
+            }
+        ]
+        
+        can_proceed, message, error_details = pricing_service.check_comprehensive_limits(
+            user_id=user_id,
+            operations=operations_to_validate
+        )
+        
+        if not can_proceed:
+            logger.error(f"[Pre-flight Validator] Image generation blocked for user {user_id}: {message}")
+            
+            usage_info = error_details.get('usage_info', {}) if error_details else {}
+            provider = usage_info.get('provider', 'stability') if usage_info else 'stability'
+            
+            raise HTTPException(
+                status_code=429,
+                detail={
+                    'error': message,
+                    'message': message,
+                    'provider': provider,
+                    'usage_info': usage_info if usage_info else error_details
+                }
+            )
+        
+        logger.info(f"[Pre-flight Validator] ✅ Image generation validated for user {user_id}")
+        # Validation passed - no return needed (function raises HTTPException if validation fails)
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"[Pre-flight Validator] Error validating image generation: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail={
+                'error': f"Failed to validate image generation: {str(e)}",
+                'message': f"Failed to validate image generation: {str(e)}"
+            }
+        )
+