LLeMbas/src/lembas/services/steps.py

"""A reply as the sequence of steps it actually was.

The three stores a reply writes into -- text, reasoning and tool events -- are
each append-only and each correct. What none of them records is *interleaving*:
where round three's thinking sat relative to round three's command and to the
sentence that came after it. So a bubble was rendered as three zones, all the
thinking, then all the tools, then all the prose, which reads fine on a two-round
answer and is unusable on a forty-round one.

The fix is a table of contents rather than a fourth copy of anything. A **mark**
is written when a round's contribution ends, holding the cumulative length of
each store at that moment; the text between two marks is one step's prose, and so
on. Nothing is duplicated, so `build_messages`, compaction, titling and the copy
button all still see `message.content` as the single string it always was.

**No marks means the old layout.** Every reply written before this existed reads
back an empty list, and `_build` answers that with thinking, then tools, then
text -- exactly what those bubbles have always shown. There is no version flag
and no branch in the template.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from lembas.services.markdown import open_fence, render_markdown

KIND_THINKING = "thinking"
KIND_TEXT = "text"
KIND_TOOLS = "tools"


@dataclass(frozen=True)
class Step:
    """One thing that happened, in the order it happened.

    `index` is the position of the mark this came from, and the trailing step --
    the one still being written -- takes the index one past the last mark. It is
    what every DOM id in the transcript is derived from, which is what lets an
    open block survive both a stream frame and the `done` frame that replaces the
    whole bubble: the marks are append-only, so index N always means the same
    step, live and afterwards alike.
    """

    index: int
    kind: str
    open: bool = False
    text: str = ""
    html: str = ""
    events: tuple[dict, ...] = field(default_factory=tuple)


def for_message(message: Any) -> list[Step]:
    """Every step of a finished reply, read off the row.

    A Jinja global (see `web/templating.py`) for the reason `tool_label` is: the
    bubble is rendered from four different handlers, and a fifth thing each of
    them had to remember to pass is a fifth thing one of them would forget.
    """
    return _build(
        text=message.content or "",
        thinking=(message.reasoning or "") if not message.error else "",
        events=list(message.tool_calls_json or []),
        marks=list(getattr(message, "steps_json", None) or []),
        # A stored reply has a trailing step and it is not being written. Those
        # are two facts and they used to be one flag: `include_open` both
        # emitted the tail and marked it live, so every finished bubble ending
        # in prose carried `msg__body--live` and blinked a caret for ever.
        live=False,
    )


def closed_from(generation: Any, since: int) -> list[Step]:
    """The finished steps of a running reply, from mark `since` onwards.

    Only the new ones, because `_follow` keeps what it has already rendered. A
    closed step never changes again -- that is what closing one means -- so the
    whole prefix does not have to be re-rendered twelve times a second, which is
    what the tool block used to cost on a long reply.
    """
    return _build(
        text=generation.text,
        thinking=generation.thinking,
        events=list(generation.tool_events),
        marks=list(generation.steps),
        since=since,
        include_open=False,
    )


def tail(generation: Any) -> tuple[str, str]:
    """What is being written right now: `(thinking, text)` past the last mark.

    The text carries a fence reopener where one is needed, so a code block
    started before the last tool call goes on rendering as a code block instead
    of the prose underneath it briefly becoming one.
    """
    marks = list(generation.steps)
    last = marks[-1] if marks else {}
    thinking = generation.thinking[_at(last, "thinking_to") :]
    text = generation.text[_at(last, "text_to") :]
    carry = _carry_before(generation.text, marks)
    return thinking, (f"{carry}\n{text}" if carry and text else text)


def _at(mark: dict, key: str) -> int:
    value = mark.get(key, 0)
    return value if isinstance(value, int) and value > 0 else 0


def _build(
    *,
    text: str,
    thinking: str,
    events: list[dict],
    marks: list[dict],
    since: int = 0,
    include_open: bool = True,
    live: bool = True,
) -> list[Step]:
    """The shared walk.

    Two flags, because they are two questions and conflating them put a blinking
    caret on every finished reply:

    * `include_open` -- emit the trailing step at all. `closed_from` says no,
      because the step still being written is carried by its own frames.
    * `live` -- mark that trailing step as still being written. Only ever true
      of a running generation. A stored reply has a tail and it is finished.

    Every offset is clamped and nothing here raises. A `steps_json` that
    disagrees with the three stores -- a row half-written when the process died,
    a hand-edited one -- has to degrade to a slightly odd order, never to a
    transcript that will not render at all.
    """
    steps: list[Step] = []

    if not marks:
        # The compatibility layout, and the layout of any reply that called
        # nothing: for that one the two orders are the same list, because there
        # are no tool blocks to sit between the prose.
        if thinking:
            steps.append(Step(index=0, kind=KIND_THINKING, text=thinking))
        if events:
            steps.append(Step(index=0, kind=KIND_TOOLS, events=tuple(events)))
        if text:
            steps.append(
                Step(index=0, kind=KIND_TEXT, open=live, html=render_markdown(text))
            )
        return steps

    thought_from = 0
    text_from = 0
    tools_from = 0
    carry = ""

    for index, mark in enumerate(marks):
        thought_to = min(max(_at(mark, "thinking_to"), thought_from), len(thinking))
        text_to = min(max(_at(mark, "text_to"), text_from), len(text))
        tools_to = min(max(_at(mark, "tools_to"), tools_from), len(events))

        thought = thinking[thought_from:thought_to]
        said = text[text_from:text_to]
        ran = events[tools_from:tools_to]

        # Computed for every step even when this one is not being returned:
        # `closed_from` renders a suffix, and whether a fence is open depends on
        # everything before it.
        source = f"{carry}\n{said}" if carry and said else said
        marker, info = open_fence(source)
        if marker:
            source = f"{source}\n{marker}"

        if index >= since:
            # Thinking, then prose, then tools -- the order a model emits them.
            if thought:
                steps.append(Step(index=index, kind=KIND_THINKING, text=thought))
            if said:
                steps.append(Step(index=index, kind=KIND_TEXT, html=render_markdown(source)))
            if ran:
                steps.append(Step(index=index, kind=KIND_TOOLS, events=tuple(ran)))

        carry = f"{marker}{info}" if marker else ""
        thought_from, text_from, tools_from = thought_to, text_to, tools_to

    if not include_open:
        return steps

    # Everything past the last mark. Implicit rather than written, in both the
    # live path and the stored one -- one rule instead of two that could drift.
    index = len(marks)
    if trailing_thought := thinking[thought_from:]:
        steps.append(Step(index=index, kind=KIND_THINKING, text=trailing_thought))
    if trailing_text := text[text_from:]:
        source = f"{carry}\n{trailing_text}" if carry else trailing_text
        steps.append(
            Step(index=index, kind=KIND_TEXT, open=live, html=render_markdown(source))
        )
    if trailing_events := events[tools_from:]:
        steps.append(Step(index=index, kind=KIND_TOOLS, events=tuple(trailing_events)))

    return steps


def _carry_before(text: str, marks: list[dict]) -> str:
    """The fence still open when the last mark was written, if any."""
    if not marks:
        return ""
    carry = ""
    start = 0
    for mark in marks:
        end = min(max(_at(mark, "text_to"), start), len(text))
        source = f"{carry}\n{text[start:end]}" if carry else text[start:end]
        marker, info = open_fence(source)
        carry = f"{marker}{info}" if marker else ""
        start = end
    return carry


__all__ = ["KIND_TEXT", "KIND_THINKING", "KIND_TOOLS", "Step", "closed_from", "for_message", "tail"]