Coverage for src/lilbee/app/session_export.py: 100%
143 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Render a saved chat session as a markdown document and write it to disk."""
3from __future__ import annotations
5import os
6import re
7from collections.abc import Iterator
8from functools import cache
9from pathlib import Path
11import yaml
12from markdown_it import MarkdownIt
13from markdown_it.rules_inline.html_inline import html_inline as _stock_html_inline
14from markdown_it.rules_inline.state_inline import StateInline
15from markdown_it.token import Token
17from lilbee.app.session_sources import message_sources
18from lilbee.core.security import write_private_text
19from lilbee.core.text import collapse_whitespace, make_slug
20from lilbee.retrieval.query.formatting import (
21 FILE_LINK_RE,
22 SourceRef,
23 append_sources_block,
24 close_open_fence,
25 located,
26 markdown_text,
27 split_sources_list,
28)
29from lilbee.sessions import MessageRole, Session, SessionMessage, SessionMeta
31_EXPORT_SUFFIX = ".md"
32SLUG_MAX_LEN = 60
33_FALLBACK_STEM = "chat"
34_ID_PREFIX_LEN = 8
35_FRONT_MATTER_FENCE = "---"
36_ROLE_HEADINGS: dict[MessageRole, str] = {
37 MessageRole.USER: "User",
38 MessageRole.ASSISTANT: "Assistant",
39}
40_ATX_MARKER = "#"
41_TURN_HEADING_LEVEL = 2
42_MAX_HEADING_LEVEL = 6
43_ESCAPE = "\\"
44# The only characters a markdown blank line may hold.
45_MARKDOWN_BLANK = " \t"
46_LINE_BREAK_RE = re.compile(r"\r\n?")
47# An HTML heading's ``<h``/``</h`` plus level; a boundary after the digit (or the
48# text simply ending there) is what excludes ``<header>``, ``<hr>`` and ``<h2o>``.
49_HTML_HEADING_RE = re.compile(r"<(?P<slash>/?)(?P<tag>[Hh])(?P<level>[1-6])(?=[\s/>]|$)")
52def session_markdown(session: Session) -> str:
53 """The session as markdown: front matter, a title heading, one section per turn."""
54 parts = [_front_matter(session.meta), f"# {collapse_whitespace(session.meta.title)}"]
55 parts.extend(_message_section(message) for message in session.messages)
56 return "\n\n".join(parts) + "\n"
59def _front_matter(meta: SessionMeta) -> str:
60 fields = {
61 "title": meta.title,
62 "session": meta.id,
63 "model": meta.model_ref,
64 "created": meta.created_at,
65 "updated": meta.updated_at,
66 }
67 if meta.forked_from:
68 fields["forked_from"] = meta.forked_from
69 dumped = yaml.safe_dump(fields, sort_keys=False, allow_unicode=True)
70 return f"{_FRONT_MATTER_FENCE}\n{dumped}{_FRONT_MATTER_FENCE}"
73def _message_section(message: SessionMessage) -> str:
74 """One turn: the message with its headings nested, then its Sources list without links."""
75 content = _LINE_BREAK_RE.sub("\n", message.content).rstrip()
76 if message.source_refs is None:
77 body = _legacy_body(content, message)
78 else:
79 body = _listed(_contained(content), message_sources(message)[1])
80 return f"## {_ROLE_HEADINGS[message.role]}\n\n{body}"
83def _legacy_body(content: str, message: SessionMessage) -> str:
84 """The body of a turn saved before ``source_refs``: its text, then the Sources
85 list it ends in with its links made plain, or the list of its ``sources``."""
86 text, stored_sources = split_sources_list(content)
87 body = _contained(text.rstrip())
88 if stored_sources:
89 return body + _contained(FILE_LINK_RE.sub(_plain_link, stored_sources))
90 return _listed(body, message_sources(message)[1])
93def _listed(body: str, refs: list[SourceRef]) -> str:
94 """*body* ending in the Sources list of *refs*, each named by its label and locator."""
95 if not refs:
96 return body
97 return append_sources_block(
98 body, refs, render=lambda ref: markdown_text(located(ref.label, ref))
99 )
102def _contained(text: str) -> str:
103 """*text* with open fences closed and headings nested, so it stays inside its turn."""
104 return _nest_headings(close_open_fence(text))
107def _plain_link(link: re.Match[str]) -> str:
108 """A stored Sources link as its label alone, as markdown that shows the label."""
109 return markdown_text(link["label"])
112def _html_inline_with_span(state: StateInline, silent: bool) -> bool:
113 """The stock ``html_inline`` rule, plus the ``state.src`` span it matched on ``token.meta``."""
114 start = state.pos
115 matched = _stock_html_inline(state, silent)
116 if matched and not silent and state.tokens and state.tokens[-1].type == "html_inline":
117 state.tokens[-1].meta = {"start": start, "end": state.pos}
118 return matched
121@cache
122def _heading_parser() -> MarkdownIt:
123 """A CommonMark parser whose ``html_inline`` tokens carry the source span they matched."""
124 md = MarkdownIt("commonmark")
125 md.inline.ruler.at("html_inline", _html_inline_with_span)
126 return md
129def _nest_headings(text: str) -> str:
130 """*text* with no heading, markdown or HTML, at or above the turn level.
132 The parser finds the headings, so a ``#`` in code or a ``#tag`` stays as
133 written, and an HTML heading tag inside a fence or an inline code span
134 stays as written too. A ``#`` heading drops two levels, to at most six;
135 an underlined heading cannot go below level two, so its underline is
136 escaped to text, and a blank line ends that text where the heading
137 ended. An HTML heading tag, open or close, drops the same two levels,
138 independently of whether it is ever closed.
139 """
140 lines = text.split("\n")
141 paragraph_ends: list[tuple[int, str | None]] = []
142 for token in _heading_parser().parse(text):
143 if token.map:
144 start, end = token.map
145 paragraph_ends.append((end - 1, _nest_block(lines, token, start, end)))
146 for last_line, blank_line in paragraph_ends:
147 if blank_line is not None:
148 lines[last_line] += "\n" + blank_line
149 return "\n".join(lines)
152def _nest_block(lines: list[str], token: Token, start: int, end: int) -> str | None:
153 """Rewrite *token*'s lines *start* to *end*; the blank line to add after them, if any."""
154 if token.type == "heading_open":
155 return _nest_heading(lines, token.markup, start, end)
156 if token.type == "html_block":
157 _demote_html_block(lines, start, end)
158 elif token.type == "inline":
159 _demote_html_inline(lines, token.content, token.children, start, end)
160 return None
163def _nest_heading(lines: list[str], markup: str, start: int, end: int) -> str | None:
164 """Rewrite the heading on lines *start* to *end*; the blank line to add after it, if any."""
165 if markup.startswith(_ATX_MARKER):
166 level = min(len(markup) + _TURN_HEADING_LEVEL, _MAX_HEADING_LEVEL)
167 lines[start] = lines[start].replace(markup, _ATX_MARKER * level, 1)
168 return None
169 return _escape_underline(lines, markup, end)
172def _escape_underline(lines: list[str], markup: str, end: int) -> str | None:
173 """Escape the underline before line *end*; the blank line that keeps the next block, if any."""
174 underline = end - 1
175 container, _, rest = lines[underline].partition(markup)
176 lines[underline] = container + _ESCAPE + markup + rest
177 if end < len(lines) and lines[end].strip(_MARKDOWN_BLANK):
178 return container.rstrip()
179 return None
182def _demote_html_tag(match: re.Match[str]) -> str:
183 """*match*, an HTML heading tag's opening chars, with its level raised by the turn offset."""
184 level = min(int(match["level"]) + _TURN_HEADING_LEVEL, _MAX_HEADING_LEVEL)
185 return f"<{match['slash']}{match['tag']}{level}"
188def _demote_html_block(lines: list[str], start: int, end: int) -> None:
189 """Rewrite any HTML heading tag's level in the raw HTML block spanning *start* to *end*."""
190 joined = "\n".join(lines[start:end])
191 lines[start:end] = _HTML_HEADING_RE.sub(_demote_html_tag, joined).split("\n")
194def _demote_html_inline(
195 lines: list[str], content: str, children: list[Token] | None, start: int, end: int
196) -> None:
197 """Rewrite each real HTML heading tag's level within the inline span *start* to *end*."""
198 spans = [child.meta for child in children or () if child.type == "html_inline" and child.meta]
199 if not spans:
200 return
201 real = _inside_spans(_HTML_HEADING_RE.finditer(content), spans)
202 if not any(real):
203 return
204 raw = "\n".join(lines[start:end])
205 pieces: list[str] = []
206 cursor = 0
207 # The parser strips only markers and whitespace and turns NUL into U+FFFD: matches pair up.
208 for match, is_real in zip(_HTML_HEADING_RE.finditer(raw), real, strict=True):
209 if is_real:
210 pieces += [raw[cursor : match.start()], _demote_html_tag(match)]
211 cursor = match.end()
212 pieces.append(raw[cursor:])
213 lines[start:end] = "".join(pieces).split("\n")
216def _inside_spans(matches: Iterator[re.Match[str]], spans: list[dict[str, int]]) -> list[bool]:
217 """Whether each of *matches*, in order, starts inside one of the ordered *spans*."""
218 inside = []
219 spans_left = iter(spans)
220 span = next(spans_left, None)
221 for match in matches:
222 while span is not None and span["end"] <= match.start():
223 span = next(spans_left, None)
224 inside.append(span is not None and span["start"] <= match.start())
225 return inside
228def default_export_name(meta: SessionMeta) -> str:
229 """``<title-slug>-<id prefix>.md``, or ``chat-<id prefix>.md`` when the title has no slug."""
230 slug = make_slug(meta.title)[:SLUG_MAX_LEN].strip("-") or _FALLBACK_STEM
231 return f"{slug}-{meta.id[:_ID_PREFIX_LEN]}{_EXPORT_SUFFIX}"
234def write_session_markdown(session: Session, destination: str) -> Path:
235 """Write *session* owner-only to *destination*, or to its default name inside
236 it when it is a directory or ends in a separator. Returns the absolute path."""
237 target = Path(destination).expanduser()
238 if target.is_dir() or destination.endswith(("/", os.sep)):
239 target = target / default_export_name(session.meta)
240 target = target.resolve()
241 write_private_text(target, session_markdown(session))
242 return target