""" Subscription and Usage Tracking Models Comprehensive models for usage-based subscription system with API cost tracking. """ from sqlalchemy import Column, Integer, String, DateTime, Float, Boolean, JSON, Text, ForeignKey, Enum from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import relationship from datetime import datetime, timedelta import enum from typing import Dict, Any, Optional Base = declarative_base() class SubscriptionTier(enum.Enum): FREE = "free" BASIC = "basic" PRO = "pro" ENTERPRISE = "enterprise" class UsageStatus(enum.Enum): ACTIVE = "active" WARNING = "warning" # 80% usage LIMIT_REACHED = "limit_reached" # 100% usage SUSPENDED = "suspended" class APIProvider(enum.Enum): GEMINI = "gemini" OPENAI = "openai" ANTHROPIC = "anthropic" MISTRAL = "mistral" TAVILY = "tavily" SERPER = "serper" METAPHOR = "metaphor" FIRECRAWL = "firecrawl" STABILITY = "stability" EXA = "exa" VIDEO = "video" IMAGE_EDIT = "image_edit" AUDIO = "audio" class BillingCycle(enum.Enum): MONTHLY = "monthly" YEARLY = "yearly" class SubscriptionPlan(Base): """Defines subscription tiers and their limits.""" __tablename__ = "subscription_plans" id = Column(Integer, primary_key=True) name = Column(String(50), nullable=False, unique=True) tier = Column(Enum(SubscriptionTier), nullable=False) price_monthly = Column(Float, nullable=False, default=0.0) price_yearly = Column(Float, nullable=False, default=0.0) # Unified AI Text Generation Call Limit (applies to all LLM providers: gemini, openai, anthropic, mistral) # Note: This column may not exist in older databases - use getattr() when accessing ai_text_generation_calls_limit = Column(Integer, default=0, nullable=True) # 0 = unlimited, None if column doesn't exist # Legacy per-provider limits (kept for backwards compatibility and analytics) gemini_calls_limit = Column(Integer, default=0) # 0 = unlimited (deprecated, use ai_text_generation_calls_limit) openai_calls_limit = Column(Integer, default=0) # (deprecated, use ai_text_generation_calls_limit) anthropic_calls_limit = Column(Integer, default=0) # (deprecated, use ai_text_generation_calls_limit) mistral_calls_limit = Column(Integer, default=0) # (deprecated, use ai_text_generation_calls_limit) # Other API Call Limits (non-LLM) tavily_calls_limit = Column(Integer, default=0) serper_calls_limit = Column(Integer, default=0) metaphor_calls_limit = Column(Integer, default=0) firecrawl_calls_limit = Column(Integer, default=0) stability_calls_limit = Column(Integer, default=0) # Image generation exa_calls_limit = Column(Integer, default=0) # Exa neural search video_calls_limit = Column(Integer, default=0) # AI video generation image_edit_calls_limit = Column(Integer, default=0) # AI image editing audio_calls_limit = Column(Integer, default=0) # AI audio generation (text-to-speech) # Token Limits (for LLM providers) gemini_tokens_limit = Column(Integer, default=0) openai_tokens_limit = Column(Integer, default=0) anthropic_tokens_limit = Column(Integer, default=0) mistral_tokens_limit = Column(Integer, default=0) # Cost Limits (in USD) monthly_cost_limit = Column(Float, default=0.0) # 0 = unlimited # Features features = Column(JSON, nullable=True) # JSON list of enabled features # Metadata description = Column(Text, nullable=True) is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) class UserSubscription(Base): """User's current subscription and billing information.""" __tablename__ = "user_subscriptions" id = Column(Integer, primary_key=True) user_id = Column(String(100), nullable=False, unique=True) plan_id = Column(Integer, ForeignKey('subscription_plans.id'), nullable=False) # Billing billing_cycle = Column(Enum(BillingCycle), default=BillingCycle.MONTHLY) current_period_start = Column(DateTime, nullable=False) current_period_end = Column(DateTime, nullable=False) # Status status = Column(Enum(UsageStatus), default=UsageStatus.ACTIVE) is_active = Column(Boolean, default=True) auto_renew = Column(Boolean, default=True) # Payment stripe_customer_id = Column(String(100), nullable=True) stripe_subscription_id = Column(String(100), nullable=True) payment_method = Column(String(50), nullable=True) # Metadata created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) # Relationships plan = relationship("SubscriptionPlan") class APIUsageLog(Base): """Detailed log of every API call for billing and monitoring.""" __tablename__ = "api_usage_logs" id = Column(Integer, primary_key=True) user_id = Column(String(100), nullable=False) # API Details provider = Column(Enum(APIProvider), nullable=False) endpoint = Column(String(200), nullable=False) method = Column(String(10), nullable=False) model_used = Column(String(100), nullable=True) # e.g., "gemini-2.5-flash" # Usage Metrics tokens_input = Column(Integer, default=0) tokens_output = Column(Integer, default=0) tokens_total = Column(Integer, default=0) # Cost Calculation cost_input = Column(Float, default=0.0) # Cost for input tokens cost_output = Column(Float, default=0.0) # Cost for output tokens cost_total = Column(Float, default=0.0) # Total cost for this call # Performance response_time = Column(Float, nullable=False) # Response time in seconds status_code = Column(Integer, nullable=False) # Request Details request_size = Column(Integer, nullable=True) # Request size in bytes response_size = Column(Integer, nullable=True) # Response size in bytes user_agent = Column(String(500), nullable=True) ip_address = Column(String(45), nullable=True) # Error Handling error_message = Column(Text, nullable=True) retry_count = Column(Integer, default=0) # Metadata timestamp = Column(DateTime, default=datetime.utcnow, nullable=False) billing_period = Column(String(20), nullable=False) # e.g., "2025-01" # Indexes for performance __table_args__ = ( {'mysql_engine': 'InnoDB'}, ) class UsageSummary(Base): """Aggregated usage statistics per user per billing period.""" __tablename__ = "usage_summaries" id = Column(Integer, primary_key=True) user_id = Column(String(100), nullable=False) billing_period = Column(String(20), nullable=False) # e.g., "2025-01" # API Call Counts gemini_calls = Column(Integer, default=0) openai_calls = Column(Integer, default=0) anthropic_calls = Column(Integer, default=0) mistral_calls = Column(Integer, default=0) tavily_calls = Column(Integer, default=0) serper_calls = Column(Integer, default=0) metaphor_calls = Column(Integer, default=0) firecrawl_calls = Column(Integer, default=0) stability_calls = Column(Integer, default=0) exa_calls = Column(Integer, default=0) video_calls = Column(Integer, default=0) # AI video generation image_edit_calls = Column(Integer, default=0) # AI image editing audio_calls = Column(Integer, default=0) # AI audio generation (text-to-speech) # Token Usage gemini_tokens = Column(Integer, default=0) openai_tokens = Column(Integer, default=0) anthropic_tokens = Column(Integer, default=0) mistral_tokens = Column(Integer, default=0) # Cost Tracking gemini_cost = Column(Float, default=0.0) openai_cost = Column(Float, default=0.0) anthropic_cost = Column(Float, default=0.0) mistral_cost = Column(Float, default=0.0) tavily_cost = Column(Float, default=0.0) serper_cost = Column(Float, default=0.0) metaphor_cost = Column(Float, default=0.0) firecrawl_cost = Column(Float, default=0.0) stability_cost = Column(Float, default=0.0) exa_cost = Column(Float, default=0.0) video_cost = Column(Float, default=0.0) # AI video generation image_edit_cost = Column(Float, default=0.0) # AI image editing audio_cost = Column(Float, default=0.0) # AI audio generation (text-to-speech) # Totals total_calls = Column(Integer, default=0) total_tokens = Column(Integer, default=0) total_cost = Column(Float, default=0.0) # Performance Metrics avg_response_time = Column(Float, default=0.0) error_rate = Column(Float, default=0.0) # Percentage # Status usage_status = Column(Enum(UsageStatus), default=UsageStatus.ACTIVE) warnings_sent = Column(Integer, default=0) # Number of warning emails sent # Metadata created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) # Unique constraint on user_id and billing_period __table_args__ = ( {'mysql_engine': 'InnoDB'}, ) class APIProviderPricing(Base): """Pricing configuration for different API providers.""" __tablename__ = "api_provider_pricing" id = Column(Integer, primary_key=True) provider = Column(Enum(APIProvider), nullable=False) model_name = Column(String(100), nullable=False) # Pricing per token (in USD) cost_per_input_token = Column(Float, default=0.0) cost_per_output_token = Column(Float, default=0.0) cost_per_request = Column(Float, default=0.0) # Fixed cost per API call # Pricing per unit for non-LLM APIs cost_per_search = Column(Float, default=0.0) # For search APIs cost_per_image = Column(Float, default=0.0) # For image generation cost_per_page = Column(Float, default=0.0) # For web crawling # Token conversion (tokens per unit) tokens_per_word = Column(Float, default=1.3) # Approximate tokens per word tokens_per_character = Column(Float, default=0.25) # Approximate tokens per character # Metadata description = Column(Text, nullable=True) is_active = Column(Boolean, default=True) effective_date = Column(DateTime, default=datetime.utcnow) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) # Unique constraint on provider and model __table_args__ = ( {'mysql_engine': 'InnoDB'}, ) class UsageAlert(Base): """Usage alerts and notifications.""" __tablename__ = "usage_alerts" id = Column(Integer, primary_key=True) user_id = Column(String(100), nullable=False) # Alert Details alert_type = Column(String(50), nullable=False) # "usage_warning", "limit_reached", "cost_warning" threshold_percentage = Column(Integer, nullable=False) # 80, 90, 100 provider = Column(Enum(APIProvider), nullable=True) # Specific provider or None for overall # Alert Content title = Column(String(200), nullable=False) message = Column(Text, nullable=False) severity = Column(String(20), default="info") # "info", "warning", "error" # Status is_sent = Column(Boolean, default=False) sent_at = Column(DateTime, nullable=True) is_read = Column(Boolean, default=False) read_at = Column(DateTime, nullable=True) # Metadata billing_period = Column(String(20), nullable=False) created_at = Column(DateTime, default=datetime.utcnow) class BillingHistory(Base): """Historical billing records.""" __tablename__ = "billing_history" id = Column(Integer, primary_key=True) user_id = Column(String(100), nullable=False) # Billing Period billing_period = Column(String(20), nullable=False) # e.g., "2025-01" period_start = Column(DateTime, nullable=False) period_end = Column(DateTime, nullable=False) # Subscription plan_name = Column(String(50), nullable=False) base_cost = Column(Float, default=0.0) # Usage Costs usage_cost = Column(Float, default=0.0) overage_cost = Column(Float, default=0.0) total_cost = Column(Float, default=0.0) # Payment payment_status = Column(String(20), default="pending") # "pending", "paid", "failed" payment_date = Column(DateTime, nullable=True) stripe_invoice_id = Column(String(100), nullable=True) # Usage Summary (snapshot) total_api_calls = Column(Integer, default=0) total_tokens = Column(Integer, default=0) usage_details = Column(JSON, nullable=True) # Detailed breakdown # Metadata created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) class SubscriptionRenewalHistory(Base): """Historical record of subscription renewals and expiration events.""" __tablename__ = "subscription_renewal_history" id = Column(Integer, primary_key=True) user_id = Column(String(100), nullable=False) # Subscription Details plan_id = Column(Integer, ForeignKey('subscription_plans.id'), nullable=False) plan_name = Column(String(50), nullable=False) plan_tier = Column(String(20), nullable=False) # e.g., "free", "basic", "pro", "enterprise" # Period Information previous_period_start = Column(DateTime, nullable=True) # Start of the previous period (if renewal) previous_period_end = Column(DateTime, nullable=True) # End of the previous period (when it expired) new_period_start = Column(DateTime, nullable=False) # Start of the new period (when renewed) new_period_end = Column(DateTime, nullable=False) # End of the new period # Billing Cycle billing_cycle = Column(Enum(BillingCycle), nullable=False) # "monthly" or "yearly" # Renewal Information renewal_type = Column(String(20), nullable=False) # "new", "renewal", "upgrade", "downgrade" renewal_count = Column(Integer, default=0) # Sequential renewal number (1st renewal, 2nd renewal, etc.) # Previous Subscription Snapshot (before renewal) previous_plan_name = Column(String(50), nullable=True) previous_plan_tier = Column(String(20), nullable=True) # Usage Summary Before Renewal (snapshot) usage_before_renewal = Column(JSON, nullable=True) # Snapshot of usage before renewal # Payment Information payment_amount = Column(Float, default=0.0) payment_status = Column(String(20), default="pending") # "pending", "paid", "failed" payment_date = Column(DateTime, nullable=True) stripe_invoice_id = Column(String(100), nullable=True) # Metadata created_at = Column(DateTime, default=datetime.utcnow) # Relationships plan = relationship("SubscriptionPlan") # Indexes for performance __table_args__ = ( {'mysql_engine': 'InnoDB'}, )