""" Pricing Service for API Usage Tracking Manages API pricing, cost calculation, and subscription limits. """ from typing import Dict, Any, Optional, List, Tuple, Union from decimal import Decimal, ROUND_HALF_UP from datetime import datetime, timedelta from sqlalchemy.orm import Session from sqlalchemy import text from loguru import logger from models.subscription_models import ( APIProviderPricing, SubscriptionPlan, UserSubscription, UsageSummary, APIUsageLog, APIProvider, SubscriptionTier ) class PricingService: """Service for managing API pricing and cost calculations.""" # Class-level cache shared across all instances (critical for cache invalidation on subscription renewal) # key: f"{user_id}:{provider}", value: { 'result': (bool, str, dict), 'expires_at': datetime } _limits_cache: Dict[str, Dict[str, Any]] = {} def __init__(self, db: Session): self.db = db self._pricing_cache = {} self._plans_cache = {} # Cache for schema feature detection (ai_text_generation_calls_limit column) self._ai_text_gen_col_checked: bool = False self._ai_text_gen_col_available: bool = False # ------------------- Billing period helpers ------------------- def _compute_next_period_end(self, start: datetime, cycle: str) -> datetime: """Compute the next period end given a start and billing cycle.""" try: cycle_value = cycle.value if hasattr(cycle, 'value') else str(cycle) except Exception: cycle_value = str(cycle) if cycle_value == 'yearly': return start + timedelta(days=365) return start + timedelta(days=30) def _ensure_subscription_current(self, subscription) -> bool: """Auto-advance subscription period if expired and auto_renew is enabled.""" if not subscription: return False now = datetime.utcnow() try: if subscription.current_period_end and subscription.current_period_end < now: if getattr(subscription, 'auto_renew', False): subscription.current_period_start = now subscription.current_period_end = self._compute_next_period_end(now, subscription.billing_cycle) # Keep status active if model enum else string try: subscription.status = subscription.status.ACTIVE # type: ignore[attr-defined] except Exception: setattr(subscription, 'status', 'active') self.db.commit() else: return False except Exception: self.db.rollback() return True def get_current_billing_period(self, user_id: str) -> Optional[str]: """Return current billing period key (YYYY-MM) after ensuring subscription is current.""" subscription = self.db.query(UserSubscription).filter( UserSubscription.user_id == user_id, UserSubscription.is_active == True ).first() # Ensure subscription is current (advance if auto_renew) self._ensure_subscription_current(subscription) # Continue to use YYYY-MM for summaries return datetime.now().strftime("%Y-%m") @classmethod def clear_user_cache(cls, user_id: str) -> int: """Clear all cached limit checks for a specific user. Returns number of entries cleared.""" keys_to_remove = [key for key in cls._limits_cache.keys() if key.startswith(f"{user_id}:")] for key in keys_to_remove: del cls._limits_cache[key] logger.info(f"Cleared {len(keys_to_remove)} cache entries for user {user_id}") return len(keys_to_remove) def initialize_default_pricing(self): """Initialize default pricing for all API providers.""" # Gemini API Pricing (Updated as of September 2025 - Official Google AI Pricing) # Source: https://ai.google.dev/gemini-api/docs/pricing gemini_pricing = [ # Gemini 2.5 Pro - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-2.5-pro", "cost_per_input_token": 0.00000125, # $1.25 per 1M input tokens (prompts <= 200k tokens) "cost_per_output_token": 0.00001, # $10.00 per 1M output tokens (prompts <= 200k tokens) "description": "Gemini 2.5 Pro - State-of-the-art multipurpose model for coding and complex reasoning" }, { "provider": APIProvider.GEMINI, "model_name": "gemini-2.5-pro-large", "cost_per_input_token": 0.0000025, # $2.50 per 1M input tokens (prompts > 200k tokens) "cost_per_output_token": 0.000015, # $15.00 per 1M output tokens (prompts > 200k tokens) "description": "Gemini 2.5 Pro - Large context model for prompts > 200k tokens" }, # Gemini 2.5 Flash - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-2.5-flash", "cost_per_input_token": 0.0000003, # $0.30 per 1M input tokens (text/image/video) "cost_per_output_token": 0.0000025, # $2.50 per 1M output tokens "description": "Gemini 2.5 Flash - Hybrid reasoning model with 1M token context window" }, { "provider": APIProvider.GEMINI, "model_name": "gemini-2.5-flash-audio", "cost_per_input_token": 0.000001, # $1.00 per 1M input tokens (audio) "cost_per_output_token": 0.0000025, # $2.50 per 1M output tokens "description": "Gemini 2.5 Flash - Audio input model" }, # Gemini 2.5 Flash-Lite - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-2.5-flash-lite", "cost_per_input_token": 0.0000001, # $0.10 per 1M input tokens (text/image/video) "cost_per_output_token": 0.0000004, # $0.40 per 1M output tokens "description": "Gemini 2.5 Flash-Lite - Smallest and most cost-effective model for at-scale usage" }, { "provider": APIProvider.GEMINI, "model_name": "gemini-2.5-flash-lite-audio", "cost_per_input_token": 0.0000003, # $0.30 per 1M input tokens (audio) "cost_per_output_token": 0.0000004, # $0.40 per 1M output tokens "description": "Gemini 2.5 Flash-Lite - Audio input model" }, # Gemini 1.5 Flash - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-1.5-flash", "cost_per_input_token": 0.000000075, # $0.075 per 1M input tokens (prompts <= 128k tokens) "cost_per_output_token": 0.0000003, # $0.30 per 1M output tokens (prompts <= 128k tokens) "description": "Gemini 1.5 Flash - Fast multimodal model with 1M token context window" }, { "provider": APIProvider.GEMINI, "model_name": "gemini-1.5-flash-large", "cost_per_input_token": 0.00000015, # $0.15 per 1M input tokens (prompts > 128k tokens) "cost_per_output_token": 0.0000006, # $0.60 per 1M output tokens (prompts > 128k tokens) "description": "Gemini 1.5 Flash - Large context model for prompts > 128k tokens" }, # Gemini 1.5 Flash-8B - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-1.5-flash-8b", "cost_per_input_token": 0.0000000375, # $0.0375 per 1M input tokens (prompts <= 128k tokens) "cost_per_output_token": 0.00000015, # $0.15 per 1M output tokens (prompts <= 128k tokens) "description": "Gemini 1.5 Flash-8B - Smallest model for lower intelligence use cases" }, { "provider": APIProvider.GEMINI, "model_name": "gemini-1.5-flash-8b-large", "cost_per_input_token": 0.000000075, # $0.075 per 1M input tokens (prompts > 128k tokens) "cost_per_output_token": 0.0000003, # $0.30 per 1M output tokens (prompts > 128k tokens) "description": "Gemini 1.5 Flash-8B - Large context model for prompts > 128k tokens" }, # Gemini 1.5 Pro - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-1.5-pro", "cost_per_input_token": 0.00000125, # $1.25 per 1M input tokens (prompts <= 128k tokens) "cost_per_output_token": 0.000005, # $5.00 per 1M output tokens (prompts <= 128k tokens) "description": "Gemini 1.5 Pro - Highest intelligence model with 2M token context window" }, { "provider": APIProvider.GEMINI, "model_name": "gemini-1.5-pro-large", "cost_per_input_token": 0.0000025, # $2.50 per 1M input tokens (prompts > 128k tokens) "cost_per_output_token": 0.00001, # $10.00 per 1M output tokens (prompts > 128k tokens) "description": "Gemini 1.5 Pro - Large context model for prompts > 128k tokens" }, # Gemini Embedding - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-embedding", "cost_per_input_token": 0.00000015, # $0.15 per 1M input tokens "cost_per_output_token": 0.0, # No output tokens for embeddings "description": "Gemini Embedding - Newest embeddings model with higher rate limits" }, # Grounding with Google Search - Standard Tier { "provider": APIProvider.GEMINI, "model_name": "gemini-grounding-search", "cost_per_request": 0.035, # $35 per 1,000 requests (after free tier) "cost_per_input_token": 0.0, # No additional token cost for grounding "cost_per_output_token": 0.0, # No additional token cost for grounding "description": "Grounding with Google Search - 1,500 RPD free, then $35/1K requests" } ] # OpenAI Pricing (estimated, will be updated) openai_pricing = [ { "provider": APIProvider.OPENAI, "model_name": "gpt-4o", "cost_per_input_token": 0.0000025, # $2.50 per 1M input tokens "cost_per_output_token": 0.00001, # $10.00 per 1M output tokens "description": "GPT-4o - Latest OpenAI model" }, { "provider": APIProvider.OPENAI, "model_name": "gpt-4o-mini", "cost_per_input_token": 0.00000015, # $0.15 per 1M input tokens "cost_per_output_token": 0.0000006, # $0.60 per 1M output tokens "description": "GPT-4o Mini - Cost-effective model" } ] # Anthropic Pricing (estimated, will be updated) anthropic_pricing = [ { "provider": APIProvider.ANTHROPIC, "model_name": "claude-3.5-sonnet", "cost_per_input_token": 0.000003, # $3.00 per 1M input tokens "cost_per_output_token": 0.000015, # $15.00 per 1M output tokens "description": "Claude 3.5 Sonnet - Anthropic's flagship model" } ] # Search API Pricing (estimated) search_pricing = [ { "provider": APIProvider.TAVILY, "model_name": "tavily-search", "cost_per_request": 0.001, # $0.001 per search "description": "Tavily AI Search API" }, { "provider": APIProvider.SERPER, "model_name": "serper-search", "cost_per_request": 0.001, # $0.001 per search "description": "Serper Google Search API" }, { "provider": APIProvider.METAPHOR, "model_name": "metaphor-search", "cost_per_request": 0.003, # $0.003 per search "description": "Metaphor/Exa AI Search API" }, { "provider": APIProvider.FIRECRAWL, "model_name": "firecrawl-extract", "cost_per_page": 0.002, # $0.002 per page crawled "description": "Firecrawl Web Extraction API" }, { "provider": APIProvider.STABILITY, "model_name": "stable-diffusion", "cost_per_image": 0.04, # $0.04 per image "description": "Stability AI Image Generation" }, { "provider": APIProvider.EXA, "model_name": "exa-search", "cost_per_request": 0.005, # $0.005 per search (1-25 results) "description": "Exa Neural Search API" } ] # Combine all pricing data all_pricing = gemini_pricing + openai_pricing + anthropic_pricing + search_pricing # Insert pricing data for pricing_data in all_pricing: existing = self.db.query(APIProviderPricing).filter( APIProviderPricing.provider == pricing_data["provider"], APIProviderPricing.model_name == pricing_data["model_name"] ).first() if not existing: pricing = APIProviderPricing(**pricing_data) self.db.add(pricing) self.db.commit() logger.debug("Default API pricing initialized") def initialize_default_plans(self): """Initialize default subscription plans.""" plans = [ { "name": "Free", "tier": SubscriptionTier.FREE, "price_monthly": 0.0, "price_yearly": 0.0, "gemini_calls_limit": 100, "openai_calls_limit": 0, "anthropic_calls_limit": 0, "mistral_calls_limit": 50, "tavily_calls_limit": 20, "serper_calls_limit": 20, "metaphor_calls_limit": 10, "firecrawl_calls_limit": 10, "stability_calls_limit": 5, "exa_calls_limit": 100, "gemini_tokens_limit": 100000, "monthly_cost_limit": 0.0, "features": ["basic_content_generation", "limited_research"], "description": "Perfect for trying out ALwrity" }, { "name": "Basic", "tier": SubscriptionTier.BASIC, "price_monthly": 29.0, "price_yearly": 290.0, "ai_text_generation_calls_limit": 10, # Unified limit for all LLM providers "gemini_calls_limit": 1000, # Legacy, kept for backwards compatibility (not used for enforcement) "openai_calls_limit": 500, "anthropic_calls_limit": 200, "mistral_calls_limit": 500, "tavily_calls_limit": 200, "serper_calls_limit": 200, "metaphor_calls_limit": 100, "firecrawl_calls_limit": 100, "stability_calls_limit": 5, "exa_calls_limit": 500, "gemini_tokens_limit": 20000, # Increased from 5000 for better stability "openai_tokens_limit": 20000, # Increased from 5000 for better stability "anthropic_tokens_limit": 20000, # Increased from 5000 for better stability "mistral_tokens_limit": 20000, # Increased from 5000 for better stability "monthly_cost_limit": 50.0, "features": ["full_content_generation", "advanced_research", "basic_analytics"], "description": "Great for individuals and small teams" }, { "name": "Pro", "tier": SubscriptionTier.PRO, "price_monthly": 79.0, "price_yearly": 790.0, "gemini_calls_limit": 5000, "openai_calls_limit": 2500, "anthropic_calls_limit": 1000, "mistral_calls_limit": 2500, "tavily_calls_limit": 1000, "serper_calls_limit": 1000, "metaphor_calls_limit": 500, "firecrawl_calls_limit": 500, "stability_calls_limit": 200, "exa_calls_limit": 2000, "gemini_tokens_limit": 5000000, "openai_tokens_limit": 2500000, "anthropic_tokens_limit": 1000000, "mistral_tokens_limit": 2500000, "monthly_cost_limit": 150.0, "features": ["unlimited_content_generation", "premium_research", "advanced_analytics", "priority_support"], "description": "Perfect for growing businesses" }, { "name": "Enterprise", "tier": SubscriptionTier.ENTERPRISE, "price_monthly": 199.0, "price_yearly": 1990.0, "gemini_calls_limit": 0, # Unlimited "openai_calls_limit": 0, "anthropic_calls_limit": 0, "mistral_calls_limit": 0, "tavily_calls_limit": 0, "serper_calls_limit": 0, "metaphor_calls_limit": 0, "firecrawl_calls_limit": 0, "stability_calls_limit": 0, "exa_calls_limit": 0, # Unlimited "gemini_tokens_limit": 0, "openai_tokens_limit": 0, "anthropic_tokens_limit": 0, "mistral_tokens_limit": 0, "monthly_cost_limit": 500.0, "features": ["unlimited_everything", "white_label", "dedicated_support", "custom_integrations"], "description": "For large organizations with high-volume needs" } ] for plan_data in plans: existing = self.db.query(SubscriptionPlan).filter( SubscriptionPlan.name == plan_data["name"] ).first() if not existing: plan = SubscriptionPlan(**plan_data) self.db.add(plan) self.db.commit() logger.debug("Default subscription plans initialized") def calculate_api_cost(self, provider: APIProvider, model_name: str, tokens_input: int = 0, tokens_output: int = 0, request_count: int = 1, **kwargs) -> Dict[str, float]: """Calculate cost for an API call.""" # Get pricing for the provider and model pricing = self.db.query(APIProviderPricing).filter( APIProviderPricing.provider == provider, APIProviderPricing.model_name == model_name, APIProviderPricing.is_active == True ).first() if not pricing: logger.warning(f"No pricing found for {provider.value}:{model_name}, using default estimates") # Use default estimates cost_input = tokens_input * 0.000001 # $1 per 1M tokens default cost_output = tokens_output * 0.000001 cost_total = (cost_input + cost_output) * request_count else: # Calculate based on actual pricing cost_input = tokens_input * pricing.cost_per_input_token cost_output = tokens_output * pricing.cost_per_output_token cost_request = request_count * pricing.cost_per_request # Handle special cases for non-LLM APIs cost_search = kwargs.get('search_count', 0) * pricing.cost_per_search cost_image = kwargs.get('image_count', 0) * pricing.cost_per_image cost_page = kwargs.get('page_count', 0) * pricing.cost_per_page cost_total = cost_input + cost_output + cost_request + cost_search + cost_image + cost_page # Round to 6 decimal places for precision return { 'cost_input': round(cost_input, 6), 'cost_output': round(cost_output, 6), 'cost_total': round(cost_total, 6) } def get_user_limits(self, user_id: str) -> Optional[Dict[str, Any]]: """Get usage limits for a user based on their subscription.""" # CRITICAL: Expire all objects first to ensure fresh data after renewal self.db.expire_all() subscription = self.db.query(UserSubscription).filter( UserSubscription.user_id == user_id, UserSubscription.is_active == True ).first() if not subscription: # Return free tier limits free_plan = self.db.query(SubscriptionPlan).filter( SubscriptionPlan.tier == SubscriptionTier.FREE ).first() if free_plan: return self._plan_to_limits_dict(free_plan) return None # Ensure current period before returning limits self._ensure_subscription_current(subscription) # CRITICAL: Refresh subscription to get latest plan_id, then refresh plan relationship self.db.refresh(subscription) # Re-query plan directly to ensure fresh data (bypass relationship cache) plan = self.db.query(SubscriptionPlan).filter( SubscriptionPlan.id == subscription.plan_id ).first() if not plan: logger.error(f"Plan not found for subscription plan_id={subscription.plan_id}") return None # Refresh plan to ensure fresh limits self.db.refresh(plan) return self._plan_to_limits_dict(plan) def _ensure_ai_text_gen_column_detection(self) -> None: """Detect at runtime whether ai_text_generation_calls_limit column exists and cache the result.""" if self._ai_text_gen_col_checked: return try: # Try to query the column - if it exists, this will work self.db.execute(text('SELECT ai_text_generation_calls_limit FROM subscription_plans LIMIT 0')) self._ai_text_gen_col_available = True except Exception: self._ai_text_gen_col_available = False finally: self._ai_text_gen_col_checked = True def _plan_to_limits_dict(self, plan: SubscriptionPlan) -> Dict[str, Any]: """Convert subscription plan to limits dictionary.""" # Detect if unified AI text generation limit column exists self._ensure_ai_text_gen_column_detection() # Use unified AI text generation limit if column exists and is set ai_text_gen_limit = None if self._ai_text_gen_col_available: try: ai_text_gen_limit = getattr(plan, 'ai_text_generation_calls_limit', None) # If 0, treat as not set (unlimited for Enterprise or use fallback) if ai_text_gen_limit == 0: ai_text_gen_limit = None except (AttributeError, Exception): # Column exists but access failed - use fallback ai_text_gen_limit = None return { 'plan_name': plan.name, 'tier': plan.tier.value, 'limits': { # Unified AI text generation limit (applies to all LLM providers) # If not set, fall back to first non-zero legacy limit for backwards compatibility 'ai_text_generation_calls': ai_text_gen_limit if ai_text_gen_limit is not None else ( plan.gemini_calls_limit if plan.gemini_calls_limit > 0 else plan.openai_calls_limit if plan.openai_calls_limit > 0 else plan.anthropic_calls_limit if plan.anthropic_calls_limit > 0 else plan.mistral_calls_limit if plan.mistral_calls_limit > 0 else 0 ), # Legacy per-provider limits (for backwards compatibility and analytics) 'gemini_calls': plan.gemini_calls_limit, 'openai_calls': plan.openai_calls_limit, 'anthropic_calls': plan.anthropic_calls_limit, 'mistral_calls': plan.mistral_calls_limit, # Other API limits 'tavily_calls': plan.tavily_calls_limit, 'serper_calls': plan.serper_calls_limit, 'metaphor_calls': plan.metaphor_calls_limit, 'firecrawl_calls': plan.firecrawl_calls_limit, 'stability_calls': plan.stability_calls_limit, # Token limits 'gemini_tokens': plan.gemini_tokens_limit, 'openai_tokens': plan.openai_tokens_limit, 'anthropic_tokens': plan.anthropic_tokens_limit, 'mistral_tokens': plan.mistral_tokens_limit, 'monthly_cost': plan.monthly_cost_limit }, 'features': plan.features or [] } def check_usage_limits(self, user_id: str, provider: APIProvider, tokens_requested: int = 0, actual_provider_name: Optional[str] = None) -> Tuple[bool, str, Dict[str, Any]]: """Check if user can make an API call within their limits. Delegates to LimitValidator for actual validation logic. Args: user_id: User ID provider: APIProvider enum (may be MISTRAL for HuggingFace) tokens_requested: Estimated tokens for the request actual_provider_name: Optional actual provider name (e.g., "huggingface" when provider is MISTRAL) Returns: (can_proceed, error_message, usage_info) """ from .limit_validation import LimitValidator validator = LimitValidator(self) return validator.check_usage_limits(user_id, provider, tokens_requested, actual_provider_name) def estimate_tokens(self, text: str, provider: APIProvider) -> int: """Estimate token count for text based on provider.""" # Get pricing info for token estimation pricing = self.db.query(APIProviderPricing).filter( APIProviderPricing.provider == provider, APIProviderPricing.is_active == True ).first() if pricing and pricing.tokens_per_word: # Use provider-specific conversion word_count = len(text.split()) return int(word_count * pricing.tokens_per_word) else: # Use default estimation (roughly 1.3 tokens per word for most models) word_count = len(text.split()) return int(word_count * 1.3) def get_pricing_info(self, provider: APIProvider, model_name: str = None) -> Optional[Dict[str, Any]]: """Get pricing information for a provider/model.""" query = self.db.query(APIProviderPricing).filter( APIProviderPricing.provider == provider, APIProviderPricing.is_active == True ) if model_name: query = query.filter(APIProviderPricing.model_name == model_name) pricing = query.first() if not pricing: return None # Return pricing info as dict return { 'provider': pricing.provider.value, 'model_name': pricing.model_name, 'cost_per_input_token': pricing.cost_per_input_token, 'cost_per_output_token': pricing.cost_per_output_token, 'cost_per_request': pricing.cost_per_request, 'description': pricing.description } def check_comprehensive_limits( self, user_id: str, operations: List[Dict[str, Any]] ) -> Tuple[bool, Optional[str], Optional[Dict[str, Any]]]: """ Comprehensive pre-flight validation that checks ALL limits before making ANY API calls. Delegates to LimitValidator for actual validation logic. This prevents wasteful API calls by validating that ALL subsequent operations will succeed before making the first external API call. Args: user_id: User ID operations: List of operations to validate, each with: - 'provider': APIProvider enum - 'tokens_requested': int (estimated tokens for LLM calls, 0 for non-LLM) - 'actual_provider_name': Optional[str] (e.g., "huggingface" when provider is MISTRAL) - 'operation_type': str (e.g., "google_grounding", "llm_call", "image_generation") Returns: (can_proceed, error_message, error_details) If can_proceed is False, error_message explains which limit would be exceeded """ from .limit_validation import LimitValidator validator = LimitValidator(self) return validator.check_comprehensive_limits(user_id, operations) def get_pricing_for_provider_model(self, provider: APIProvider, model_name: str) -> Optional[Dict[str, Any]]: """Get pricing configuration for a specific provider and model.""" pricing = self.db.query(APIProviderPricing).filter( APIProviderPricing.provider == provider, APIProviderPricing.model_name == model_name ).first() if not pricing: return None return { 'provider': pricing.provider.value, 'model_name': pricing.model_name, 'cost_per_input_token': pricing.cost_per_input_token, 'cost_per_output_token': pricing.cost_per_output_token, 'cost_per_request': pricing.cost_per_request, 'cost_per_search': pricing.cost_per_search, 'cost_per_image': pricing.cost_per_image, 'cost_per_page': pricing.cost_per_page, 'description': pricing.description }