src/models/message.py

from datetime import datetime
from typing import TYPE_CHECKING, Optional
from uuid import UUID, uuid4
from enum import Enum

from sqlmodel import Field, SQLModel, Relationship

if TYPE_CHECKING:
    from .conversation import Conversation


class MessageRole(str, Enum):
    """Enum for message sender roles."""
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"


class MessageBase(SQLModel):
    """Base model for Message with common fields."""
    role: str = Field(max_length=20, nullable=False)
    content: str = Field(nullable=False)


class Message(MessageBase, table=True):
    """Message model for conversation messages."""
    id: UUID = Field(default_factory=uuid4, primary_key=True)
    conversation_id: int = Field(foreign_key="conversation.id", index=True, nullable=False)
    user_id: UUID = Field(nullable=False)
    created_at: datetime = Field(default_factory=datetime.utcnow, nullable=False, index=True)

    # Relationship to conversation
    conversation: Optional["Conversation"] = Relationship(back_populates="messages")