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

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

from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING, Any

from sqlalchemy import Boolean, DateTime, 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

if TYPE_CHECKING:
    # Annotation only; SQLAlchemy resolves the name through its own registry at
    # runtime, so there is no import cycle. A bare `Mapped[list]` would be read
    # as a scalar and hand back None instead of [].
    from lembas.db.models.library import KnowledgeBase

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

# What a conversation is allowed to be. A plain chat can never act; an agent
# chat is pointed at a machine before it starts and stays pointed there.
KIND_CHAT = "chat"
KIND_AGENT = "agent"
KINDS = (KIND_CHAT, KIND_AGENT)

# Duplicated from services/agent/policy.py rather than imported: a model module
# importing a service would invert the dependency, and this is only the column
# default. policy.MODES is the vocabulary; this is what a row starts as.
MODE_MANUAL = "manual"


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 visible_chats(self, kind: str = "") -> list[Chat]:
        """The chats in this folder that belong in the sidebar.

        The relationship itself stays unfiltered -- back-population needs every
        row -- so the listing rule lives here rather than in the template, where
        the loop and the "Empty" check would have to agree by hand and already
        did not: archived chats have been showing inside folders since folders
        existed. The unfiled list has always filtered them (api/pages.py); the
        folder branch went through the relationship and filtered nothing.

        `kind` narrows to one side of the sidebar's Chat/Agent switch. Empty
        means both, which is what every caller outside the sidebar wants.

        Ordered like the unfiled list: pinned first, then most recently touched.
        """
        kept = [
            chat
            for chat in self.chats
            if not chat.archived and not chat.temporary and (not kind or chat.kind == kind)
        ]
        kept.sort(key=lambda chat: chat.updated_at, reverse=True)
        kept.sort(key=lambda chat: not chat.pinned)
        return kept

    def visible_children(self, kind: str = "") -> list[Folder]:
        """Sub-folders the sidebar should show on this side of the switch.

        Here rather than in the template because Jinja's `selectattr` names a
        test, it does not call a method -- so the filter would have to be spelled
        out as a loop appending to a list, in a template that already includes
        itself recursively.
        """
        return [child for child in self.children if child.shown_in(kind)]

    def holds(self, kind: str = "") -> bool:
        """Whether anything of this kind is anywhere under this folder.

        Recursive, because a folder's only matching chat may be three levels
        down and judging on its own contents alone would bury it.
        """
        if self.visible_chats(kind):
            return True
        return any(child.holds(kind) for child in self.children)

    def shown_in(self, kind: str = "") -> bool:
        """Whether this folder belongs on one side of the sidebar's switch.

        Two different reasons a folder can have nothing in it, and only one of
        them is a reason to hide it. A folder full of ordinary chats is noise on
        the Agent side and is dropped. A folder that is empty of *everything* is
        a container somebody just made and has not filled yet -- hiding that one
        means it can never be found again, let alone filed into, so it shows on
        both sides and says "Empty" for itself.
        """
        return self.holds(kind) or not self.holds()

    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)

    # Never listed in the sidebar, and swept a day after the last thing said in
    # it. A real row rather than something held in the browser, so a reload or a
    # dropped connection does not lose the conversation -- and `Keep` clears the
    # flag, because a temporary chat that turns out to matter must have a way
    # out. See services/chat.py:sweep_temporary.
    temporary: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    # A reply landed while nobody was watching this chat. Cleared when the chat
    # is next opened. `unread_notified` stops the same arrival being announced
    # on every poll.
    unread: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
    unread_notified: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    # --- Agent chats ---------------------------------------------------------
    # Whether this conversation may act, and where. Chosen on the new-chat
    # screen and fixed once there is a message: the harness, the tools offered
    # and the approval loop all differ, so a chat that changed kind halfway
    # would have a transcript whose earlier turns were produced under other
    # rules. The connection is locked with it -- a shell history and a project
    # directory do not transplant to another machine.
    kind: Mapped[str] = mapped_column(String(16), default=KIND_CHAT, nullable=False)
    # A plain id rather than a ForeignKey, for the reason `compacted_through_id`
    # below gives: migrations.py compiles only the column type, so a REFERENCES
    # clause would exist on a fresh database and not on an upgraded one.
    # Validated on read instead.
    ssh_profile_id: Mapped[str | None] = mapped_column(String(32))
    # Where commands start on the far side, and what file paths resolve against.
    project_dir: Mapped[str] = mapped_column(String(500), default="")
    # Which of the four permission modes is in force. The one agent field that
    # IS switchable mid-chat: it decides what gets asked about, not what the
    # conversation is.
    agent_mode: Mapped[str] = mapped_column(String(16), default=MODE_MANUAL, nullable=False)
    # Set when a turn was edited or regenerated in an agent chat. The project
    # directory is deliberately NOT rewound with the transcript -- it is
    # somebody's real working tree and deleting their work would be far worse
    # than an inconsistency -- so the harness says so instead.
    rewound_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    # Which message carries the plan currently in force. A plain id and not a
    # ForeignKey, for the reason `compacted_through_id` below gives; validated
    # on read. It exists so the harness can put the plan in front of the model
    # with one `db.get` by primary key rather than a scan for "the newest
    # message with a plan" -- `context_variables` is synchronous and on the
    # request path. A plan a model cannot see is a plan it cannot keep current.
    plan_message_id: Mapped[str | None] = mapped_column(String(32))
    # What this chat has switched off, narrowing what it is already allowed.
    # {"families": {"web_search": false}, "skills": {"weekly-report": false}}.
    # **Absent means on**, for every key -- the same convention
    # `McpServer.tool_overrides_json` uses, and for the same reason: two
    # representations of "on" makes "why is this off?" unanswerable.
    scope_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)

    # --- Compaction ----------------------------------------------------------
    # A summary of the turns up to `compacted_through_id`, sent in their place.
    # The messages themselves are kept and still shown; they simply stop being
    # part of the request. See services/compaction.py.
    compact_summary: Mapped[str] = mapped_column(Text, default="")
    # A plain id, deliberately not a ForeignKey: db/migrations.py compiles only
    # the column type, so a REFERENCES clause would exist on a freshly created
    # database and not on an upgraded one, and a constraint half the fleet has
    # is worse than none. It is validated on every read instead -- the same
    # reasoning `model_id` above carries.
    compacted_through_id: Mapped[str | None] = mapped_column(String(32))
    compacted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))

    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",
    )
    # Which knowledge bases this chat draws on. None means "everything its owner
    # can see"; naming some scopes the knowledge tool to those.
    knowledge_bases: Mapped[list[KnowledgeBase]] = relationship(
        "KnowledgeBase", secondary="chat_knowledge_bases"
    )

    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="")

    # What the model did before answering: one entry per tool call, with its
    # arguments and results. Shown in the transcript so the sources behind an
    # answer stay visible, and deliberately NOT replayed as context on the next
    # turn -- see services/generation.py for why.
    tool_calls_json: Mapped[list[Any]] = mapped_column(JSONList, default=list)
    usage_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)

    # A plan produced in Plan mode, or the state of one being carried out. See
    # services/plans.py for the shape. Marked on the row rather than parsed back
    # out of the prose, so the Execute button sends exactly what was proposed
    # and not an approximation of it. Read through the `plan` property below,
    # never directly: rows written before version 2 hold `{title, steps}`.
    plan_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)
    # True when the reader pressed Stop. Distinct from `error`: the text that
    # did arrive is kept and is perfectly usable, it is just cut short.
    stopped: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
    # Typed while a reply was still being written, and not yet handed to a
    # model. A row rather than something held in the browser: it survives a
    # restart, it is in the transcript the moment it is typed, and it can be
    # withdrawn before it is ever sent. `build_messages` skips it; delivery --
    # `generation._drain` at the end of a reply, or `_inject` between two rounds
    # of tool calls -- is the only thing that clears it.
    queued: Mapped[bool] = mapped_column(Boolean, default=False, 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]

    @property
    def plan(self) -> dict:
        """The plan, always in the current shape.

        A property for the reason `images` and `documents` are: a message bubble
        is rendered from four different handlers, and every one of them would
        otherwise have to remember to normalise. Rows written before version 2
        hold `{title, steps}` and come back through here as one phase.
        """
        from lembas.services import plans

        return plans.normalise(self.plan_json)

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