LLeMbas/src/lembas/db/models/chat.py

"""Folders, chats and messages."""

from __future__ import annotations

from typing import Any

from sqlalchemy import Boolean, ForeignKey, Integer, String, Text
from sqlalchemy.orm import Mapped, mapped_column, relationship

from lembas.db.base import Base, Timestamps, UUIDPrimaryKey
from lembas.db.types import JSONDict, JSONList

ROLE_SYSTEM = "system"
ROLE_USER = "user"
ROLE_ASSISTANT = "assistant"
ROLE_TOOL = "tool"


class Folder(UUIDPrimaryKey, Timestamps, Base):
    """A user-owned, arbitrarily nested container for chats."""

    __tablename__ = "folders"

    user_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
    )
    parent_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("folders.id", ondelete="CASCADE")
    )
    name: Mapped[str] = mapped_column(String(200), nullable=False)
    position: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    collapsed: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    children: Mapped[list[Folder]] = relationship(
        back_populates="parent",
        cascade="all, delete-orphan",
        order_by="Folder.position, Folder.name",
    )
    parent: Mapped[Folder | None] = relationship(back_populates="children", remote_side="Folder.id")
    chats: Mapped[list[Chat]] = relationship(back_populates="folder")

    def __repr__(self) -> str:
        return f"<Folder {self.name}>"


class Chat(UUIDPrimaryKey, Timestamps, Base):
    __tablename__ = "chats"

    user_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
    )
    # Deleting a folder keeps its chats; they fall back to the unfiled list.
    folder_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("folders.id", ondelete="SET NULL"), index=True
    )

    title: Mapped[str] = mapped_column(String(300), default="New chat")
    # Set once the model writes the first reply, so auto-titling only runs once.
    title_generated: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    # Denormalised rather than a foreign key: chat history must survive an admin
    # deleting a connection or a model disappearing upstream.
    model_id: Mapped[str] = mapped_column(String(300), default="")
    connection_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("connections.id", ondelete="SET NULL")
    )

    system_prompt: Mapped[str] = mapped_column(Text, default="")
    params_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)

    pinned: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
    archived: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    folder: Mapped[Folder | None] = relationship(back_populates="chats")
    messages: Mapped[list[Message]] = relationship(
        back_populates="chat",
        cascade="all, delete-orphan",
        order_by="Message.created_at",
    )

    def __repr__(self) -> str:
        return f"<Chat {self.title!r}>"


class Message(UUIDPrimaryKey, Timestamps, Base):
    __tablename__ = "messages"

    chat_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("chats.id", ondelete="CASCADE"), nullable=False, index=True
    )

    # Reserved for conversation branching (edit a message, regenerate a reply
    # and keep both). Nothing reads it yet; it exists now because retrofitting a
    # column onto a live SQLite database without migrations is painful.
    parent_id: Mapped[str | None] = mapped_column(String(32), ForeignKey("messages.id"))

    role: Mapped[str] = mapped_column(String(16), nullable=False)
    content: Mapped[str] = mapped_column(Text, default="")

    # Reserved for multimodal turns: [{"type": "image_url", ...}, ...].
    # Plain-text messages leave this empty and use `content`.
    content_parts_json: Mapped[list[Any]] = mapped_column(JSONList, default=list)

    # A reasoning model's visible thinking, kept separate from the answer so it
    # can be collapsed, and so it is never fed back as context on the next turn
    # -- providers expect the answer alone, and replaying the thinking both
    # wastes the window and degrades the reply.
    reasoning: Mapped[str] = mapped_column(Text, default="")
    # Milliseconds spent producing the reasoning, for the "Thought for Xs" label.
    reasoning_ms: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

    model_id: Mapped[str] = mapped_column(String(300), default="")
    tool_calls_json: Mapped[list[Any]] = mapped_column(JSONList, default=list)
    usage_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)

    # Non-empty when generation failed. Rendered as a styled error in the
    # thread so a failed turn is never an unexplained blank bubble.
    error: Mapped[str] = mapped_column(Text, default="")
    # False while a reply is still streaming; flipped when the stream ends.
    complete: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

    chat: Mapped[Chat] = relationship(back_populates="messages")
    attachments: Mapped[list[Attachment]] = relationship(  # noqa: F821
        back_populates="message",
        cascade="all, delete-orphan",
        order_by="Attachment.created_at",
    )

    @property
    def images(self) -> list:
        return [a for a in self.attachments if a.is_image]

    @property
    def documents(self) -> list:
        return [a for a in self.attachments if not a.is_image]

    def __repr__(self) -> str:
        return f"<Message {self.role} {self.content[:40]!r}>"