chat_agent/models/message.py

"""Message model for storing chat messages."""

from datetime import datetime
from sqlalchemy import Column, String, Text, DateTime, Integer, JSON
from sqlalchemy import ForeignKey
from sqlalchemy.orm import relationship
from .base import BaseModel, db, GUID


class Message(BaseModel):
    """Model for storing individual chat messages."""
    
    __tablename__ = 'messages'
    
    # Core message fields
    session_id = Column(GUID(), ForeignKey('chat_sessions.id'), nullable=False, index=True)
    role = Column(String(20), nullable=False)  # 'user' or 'assistant'
    content = Column(Text, nullable=False)
    language = Column(String(50), nullable=False, default='python')
    
    # Metadata and tracking
    timestamp = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)
    message_metadata = Column(JSON, default=dict)  # Additional context like token count, processing time, etc.
    
    # Relationships
    session = relationship("ChatSession", back_populates="messages")
    
    def __init__(self, session_id, role, content, language='python', message_metadata=None):
        """Initialize a new message."""
        super().__init__()
        self.session_id = session_id
        self.role = role
        self.content = content
        self.language = language
        self.timestamp = datetime.utcnow()
        self.message_metadata = message_metadata or {}
    
    def to_dict(self):
        """Convert message to dictionary with formatted timestamp."""
        data = super().to_dict()
        # Format timestamp as ISO string for JSON serialization
        if self.timestamp:
            data['timestamp'] = self.timestamp.isoformat()
        return data
    
    @classmethod
    def create_user_message(cls, session_id, content, language='python'):
        """Create a user message."""
        return cls(
            session_id=session_id,
            role='user',
            content=content,
            language=language
        )
    
    @classmethod
    def create_assistant_message(cls, session_id, content, language='python', message_metadata=None):
        """Create an assistant message."""
        return cls(
            session_id=session_id,
            role='assistant',
            content=content,
            language=language,
            message_metadata=message_metadata
        )
    
    def __repr__(self):
        """String representation of the message."""
        content_preview = self.content[:50] + "..." if len(self.content) > 50 else self.content
        return f"<Message(id={self.id}, role={self.role}, content='{content_preview}')>"