from __future__ import annotations from datetime import datetime, timezone from sqlalchemy.sql import func from extensions import db, INDUSTRIAL_COPILOT import enum from sqlalchemy import desc import secrets import string class SenderTypeEnum(enum.Enum): USER = "user" BOT = "bot" # Function to generate a random and unique session token (NOT a full path) def generate_session_url(): # Generate a random string of length 16 consisting of uppercase, lowercase letters and digits session_code = ''.join(secrets.choice(string.ascii_letters + string.digits) for _ in range(16)) # IMPORTANT: store ONLY the token in DB (no "/chat/") return session_code class ChatSession(db.Model): __tablename__ = "chat_sessions" id = db.Column(db.Integer, primary_key=True) session_url = db.Column(db.String(255), unique=True, nullable=False) # Save the session URL here analysis_id = db.Column( db.String(36), db.ForeignKey( "incident_analyses.id", ondelete="SET NULL" ), nullable=True, index=True ) user_id = db.Column(db.Integer, db.ForeignKey("users.id", ondelete="CASCADE"), nullable=False) title = db.Column(db.String(140), nullable=False) created_at = db.Column(db.DateTime, default=datetime.now(timezone.utc), nullable=False) updated_at = db.Column(db.DateTime, default=datetime.now(timezone.utc), onupdate=datetime.now(timezone.utc)) # Keeping the sharing functionality intact is_shared = db.Column(db.Boolean, nullable=False, default=False, index=True) share_token = db.Column(db.String(64), unique=True, index=True, nullable=True) shared_at = db.Column(db.DateTime, nullable=True) user = db.relationship("User", backref=db.backref("chat_sessions", lazy="dynamic")) analysis = db.relationship( "IncidentAnalysisModel", backref=db.backref( "chat_sessions", lazy="dynamic" ) ) def __init__(self, user_id, title): self.user_id = user_id self.title = title self.session_url = generate_session_url() # Generate the unique session URL here # Share functionality methods (keeping this intact) def ensure_share_token(self): if not self.share_token: # 32 bytes => token طولانی و امن self.share_token = secrets.token_urlsafe(32) self.is_shared = True self.shared_at = datetime.now(timezone.utc) def get_incident_id(self,): if not INDUSTRIAL_COPILOT: return None else: return self.analysis_id def __repr__(self): return f"" @staticmethod def list_user_sessions(user_id: int, limit: int = 200): """ Return sessions ordered by most recently updated. """ return ( ChatSession.query .filter_by(user_id=user_id) .order_by(desc(ChatSession.updated_at), desc(ChatSession.id)) .limit(limit) .all() ) # -------------------- # Session helpers # -------------------- @staticmethod def get_or_create_latest_session(user_id: int) -> ChatSession: """ Return latest session for the user. If none exists, create one. """ sess = ( ChatSession.query .filter_by(user_id=user_id) .order_by(desc(ChatSession.updated_at), desc(ChatSession.id)) .first() ) if sess: return sess sess = ChatSession(user_id=user_id, title="New chat") db.session.add(sess) db.session.commit() return sess @staticmethod def delete_session(user_id: int, session_id: int) -> bool: """ Delete a single session and everything inside. """ sess = ChatSession.query.filter_by(user_id=user_id, id=session_id).first() if not sess: return False db.session.delete(sess) db.session.commit() return True @staticmethod def delete_session_by_url(user_id: int, session_url: str) -> bool: """ Delete a session by its URL, along with all related messages and images. """ # Find the session by session_url and user_id session = ChatSession.query.filter_by(session_url=session_url, user_id=user_id).first() if not session: return False # Session not found # Delete all messages associated with this session messages = ChatMessage.query.filter_by(session_id=session.id).all() for msg in messages: # Delete images associated with each message for img in msg.images: db.session.delete(img) db.session.delete(msg) # Delete the session itself db.session.delete(session) db.session.commit() return True # Successfully deleted the session and its resources class MessageLanguageEnum(enum.Enum): EN = "en" FA = "fa" DE = "de" FR = "fr" AR = "ar" TR = "tr" UN = 'un' class ChatMessage(db.Model): __tablename__ = "chat_messages" id = db.Column(db.Integer, primary_key=True) user_id = db.Column( db.Integer, db.ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True ) # Each message belongs to a session session_id = db.Column( db.Integer, db.ForeignKey("chat_sessions.id", ondelete="CASCADE"), nullable=False, index=True ) sender_type = db.Column( db.Enum(SenderTypeEnum), nullable=False, default=SenderTypeEnum.USER ) message = db.Column(db.Text, nullable=False) language = db.Column(db.String(10), nullable=False, default="en") created_at = db.Column( db.DateTime, default=datetime.utcnow, nullable=False ) user = db.relationship( "User", backref=db.backref("chat_messages", lazy="dynamic") ) session = db.relationship( "ChatSession", backref=db.backref("messages", lazy="dynamic", cascade="all, delete-orphan") ) def __repr__(self): return ( f"" ) # -------------------- # Message helpers # -------------------- @staticmethod def get_user_chat_history(user_id: int, limit: int = 100, session_id: int | None = None): """ Backward-compatible method: - If session_id is None, returns history of the latest session (creates one if missing). - Otherwise returns history for the requested session. Ordered oldest -> newest. """ if session_id is None: session_id = ChatMessage.get_or_create_latest_session(user_id).id messages = ( ChatMessage.query .filter_by(user_id=user_id, session_id=session_id) .order_by(desc(ChatMessage.id)) .limit(limit) .all() ) messages.reverse() return messages @staticmethod def reset_history(user_id: int): """ Delete ALL sessions (and their messages/images) for the user. This preserves the existing "clear history" functionality. """ ChatSession.query.filter_by(user_id=user_id).delete(synchronize_session=False) db.session.commit() class ChatMessageImage(db.Model): __tablename__ = "chat_message_images" id = db.Column(db.Integer, primary_key=True) chat_message_id = db.Column( db.Integer, db.ForeignKey("chat_messages.id", ondelete="CASCADE"), nullable=False, index=True ) image_url = db.Column(db.String(255), nullable=False) created_at = db.Column( db.DateTime, nullable=False, server_default=func.now() ) chat_message = db.relationship( "ChatMessage", backref=db.backref("images", lazy="dynamic", cascade="all, delete-orphan") )