from sqlalchemy import create_engine, Column, Integer, String, ForeignKey, DateTime, Text, Float, Boolean, JSON, UniqueConstraint from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker, relationship from datetime import datetime, UTC # Define the base class for declarative models Base = declarative_base() # Define the Users table class User(Base): __tablename__ = 'users' user_id = Column(Integer, primary_key=True, autoincrement=True) username = Column(String, nullable=False) # Removed unique=True email = Column(String, unique=True, nullable=False) created_at = Column(DateTime, default=lambda: datetime.now(UTC)) # Add relationship for cascade options chats = relationship("Chat", back_populates="user", cascade="all, delete-orphan") usage_records = relationship("ModelUsage", back_populates="user") deep_analysis_reports = relationship("DeepAnalysisReport", back_populates="user", cascade="all, delete-orphan") template_preferences = relationship("UserTemplatePreference", back_populates="user", cascade="all, delete-orphan") # Define the Chats table class Chat(Base): __tablename__ = 'chats' chat_id = Column(Integer, primary_key=True, autoincrement=True) user_id = Column(Integer, ForeignKey('users.user_id', ondelete="CASCADE"), nullable=True) title = Column(String, default='New Chat') created_at = Column(DateTime, default=lambda: datetime.now(UTC)) # Add relationships for cascade options user = relationship("User", back_populates="chats") messages = relationship("Message", back_populates="chat", cascade="all, delete-orphan") usage_records = relationship("ModelUsage", back_populates="chat") # Define the Messages table class Message(Base): __tablename__ = 'messages' message_id = Column(Integer, primary_key=True, autoincrement=True) chat_id = Column(Integer, ForeignKey('chats.chat_id', ondelete="CASCADE"), nullable=False) sender = Column(String, nullable=False) # 'user' or 'ai' content = Column(Text, nullable=False) timestamp = Column(DateTime, default=lambda: datetime.now(UTC)) # Add relationship for cascade options chat = relationship("Chat", back_populates="messages") feedback = relationship("MessageFeedback", back_populates="message", uselist=False, cascade="all, delete-orphan") # Define the Model Usage table class ModelUsage(Base): """Tracks AI model usage metrics for analytics and billing purposes.""" __tablename__ = 'model_usage' usage_id = Column(Integer, primary_key=True) user_id = Column(Integer, ForeignKey('users.user_id', ondelete="SET NULL"), nullable=True) chat_id = Column(Integer, ForeignKey('chats.chat_id', ondelete="SET NULL"), nullable=True) model_name = Column(String(100), nullable=False) provider = Column(String(50), nullable=False) prompt_tokens = Column(Integer, default=0) completion_tokens = Column(Integer, default=0) total_tokens = Column(Integer, default=0) query_size = Column(Integer, default=0) # Size in characters response_size = Column(Integer, default=0) # Size in characters cost = Column(Float, default=0.0) # Cost in USD timestamp = Column(DateTime, default=lambda: datetime.now(UTC)) is_streaming = Column(Boolean, default=False) request_time_ms = Column(Integer, default=0) # Request processing time in milliseconds # Add relationships user = relationship("User", back_populates="usage_records") chat = relationship("Chat", back_populates="usage_records") # Define the Code Execution table class CodeExecution(Base): """Tracks code execution attempts and results for analysis and debugging.""" __tablename__ = 'code_executions' execution_id = Column(Integer, primary_key=True, autoincrement=True) message_id = Column(Integer, ForeignKey('messages.message_id', ondelete="CASCADE"), nullable=True) chat_id = Column(Integer, ForeignKey('chats.chat_id', ondelete="CASCADE"), nullable=True) user_id = Column(Integer, ForeignKey('users.user_id', ondelete="SET NULL"), nullable=True) # Code tracking initial_code = Column(Text, nullable=True) # First version of code submitted latest_code = Column(Text, nullable=True) # Most recent version of code # Execution results is_successful = Column(Boolean, default=False) output = Column(Text, nullable=True) # Full output including errors # Model and agent information model_provider = Column(String(50), nullable=True) model_name = Column(String(100), nullable=True) model_temperature = Column(Float, nullable=True) model_max_tokens = Column(Integer, nullable=True) # Failure information failed_agents = Column(Text, nullable=True) # JSON list of agent names that failed error_messages = Column(Text, nullable=True) # JSON map of error messages by agent # Metadata created_at = Column(DateTime, default=lambda: datetime.now(UTC)) updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) class MessageFeedback(Base): """Tracks user feedback and model settings for each message.""" __tablename__ = 'message_feedback' feedback_id = Column(Integer, primary_key=True, autoincrement=True) message_id = Column(Integer, ForeignKey('messages.message_id', ondelete="CASCADE"), nullable=False) # User feedback rating = Column(Integer, nullable=True) # Star rating (1-5) # Model settings used for this message model_name = Column(String(100), nullable=True) model_provider = Column(String(50), nullable=True) temperature = Column(Float, nullable=True) max_tokens = Column(Integer, nullable=True) # Metadata created_at = Column(DateTime, default=lambda: datetime.now(UTC)) updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) # Relationship message = relationship("Message", back_populates="feedback") class DeepAnalysisReport(Base): """Stores deep analysis reports with comprehensive analysis data and metadata.""" __tablename__ = 'deep_analysis_reports' report_id = Column(Integer, primary_key=True, autoincrement=True) report_uuid = Column(String(100), unique=True, nullable=False) # Frontend generated ID user_id = Column(Integer, ForeignKey('users.user_id', ondelete="CASCADE"), nullable=True) # Analysis objective and status goal = Column(Text, nullable=False) # The analysis objective/question status = Column(String(20), nullable=False, default='pending') # 'pending', 'running', 'completed', 'failed' # Timing information start_time = Column(DateTime, default=lambda: datetime.now(UTC)) end_time = Column(DateTime, nullable=True) duration_seconds = Column(Integer, nullable=True) # Calculated duration # Analysis components (stored as text/JSON) deep_questions = Column(Text, nullable=True) # Generated analytical questions deep_plan = Column(Text, nullable=True) # Analysis plan summaries = Column(JSON, nullable=True) # Array of analysis summaries analysis_code = Column(Text, nullable=True) # Generated Python code plotly_figures = Column(JSON, nullable=True) # Array of Plotly figure data synthesis = Column(JSON, nullable=True) # Array of synthesis insights final_conclusion = Column(Text, nullable=True) # Final analysis conclusion # Report output html_report = Column(Text, nullable=True) # Complete HTML report report_summary = Column(Text, nullable=True) # Brief summary for listing # Execution tracking progress_percentage = Column(Integer, default=0) # Progress 0-100 steps_completed = Column(JSON, nullable=True) # Array of completed step names error_message = Column(Text, nullable=True) # Error details if failed # Model and cost tracking model_provider = Column(String(50), nullable=True) model_name = Column(String(100), nullable=True) total_tokens_used = Column(Integer, default=0) estimated_cost = Column(Float, default=0.0) # Cost in USD credits_consumed = Column(Integer, default=0) # Credits deducted for this analysis # Metadata created_at = Column(DateTime, default=lambda: datetime.now(UTC)) updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) # Relationships user = relationship("User", back_populates="deep_analysis_reports") class AgentTemplate(Base): """Stores predefined agent templates that users can enable/disable.""" __tablename__ = 'agent_templates' template_id = Column(Integer, primary_key=True, autoincrement=True) # Template definition template_name = Column(String(100), nullable=False, unique=True) # e.g., 'pytorch_specialist', 'data_cleaning_expert' display_name = Column(String(200), nullable=True) # User-friendly display name description = Column(Text, nullable=False) # Short description for template selection prompt_template = Column(Text, nullable=False) # Main prompt/instructions for agent behavior # Template appearance icon_url = Column(String(500), nullable=True) # URL to template icon (CDN, data URL, or relative path) # Template categorization category = Column(String(50), nullable=True) # 'Visualization', 'Modelling', 'Data Manipulation' is_premium_only = Column(Boolean, default=False) # True if template requires premium subscription # Agent variant support variant_type = Column(String(20), default='individual') # 'planner', 'individual', or 'both' base_agent = Column(String(100), nullable=True) # Base agent name for variants (e.g., 'preprocessing_agent') # Status and metadata is_active = Column(Boolean, default=True) # Timestamps created_at = Column(DateTime, default=lambda: datetime.now(UTC)) updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) # Relationships user_preferences = relationship("UserTemplatePreference", back_populates="template", cascade="all, delete-orphan") class UserTemplatePreference(Base): """Tracks user preferences and usage for agent templates.""" __tablename__ = 'user_template_preferences' preference_id = Column(Integer, primary_key=True, autoincrement=True) user_id = Column(Integer, ForeignKey('users.user_id', ondelete="CASCADE"), nullable=False) template_id = Column(Integer, ForeignKey('agent_templates.template_id', ondelete="CASCADE"), nullable=False) # User preferences is_enabled = Column(Boolean, default=True) # Whether user has this template enabled # Usage tracking usage_count = Column(Integer, default=0) # Track how many times user has used this template last_used_at = Column(DateTime, nullable=True) # Last time user used this template # Timestamps created_at = Column(DateTime, default=lambda: datetime.now(UTC)) updated_at = Column(DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) # Relationships user = relationship("User", back_populates="template_preferences") template = relationship("AgentTemplate", back_populates="user_preferences") # Constraints - user can only have one preference record per template __table_args__ = ( UniqueConstraint('user_id', 'template_id', name='unique_user_template_preference'), )