Coverage for src/lilbee/cli/tui/widgets/message.py: 100%
147 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Chat message widgets: user and assistant bubbles."""
3from __future__ import annotations
5import time
6from collections.abc import Sequence
7from pathlib import Path
8from typing import ClassVar
10from markdown_it import MarkdownIt
11from textual.app import ComposeResult
12from textual.containers import Vertical
13from textual.content import Content
14from textual.widgets import Collapsible, Markdown, Static
15from typing_extensions import override
17from lilbee.cli.tui import messages as msg
18from lilbee.cli.tui.widgets.thinking_header import ThinkingHeader
19from lilbee.core.config import cfg
21# Minimum interval (seconds) between markdown widget updates during streaming
22_MD_UPDATE_INTERVAL = 0.1
24_SPEAKER_YOU = "[bold $primary]you[/]"
25_SPEAKER_LILBEE = "[bold $success]lilbee[/]"
28class _AnswerMarkdownIt(MarkdownIt):
29 """Markdown parser for answers, additionally allowing ``file:`` links.
31 markdown-it rejects ``file:`` destinations by default (a web-context XSS
32 guard), which left the Sources block's citation links rendering as raw
33 ``[label](file://...)`` text. Answers link to the reader's own documents on
34 disk, so re-admit the scheme; everything else keeps the default validation.
35 """
37 @override
38 def validateLink(self, url: str) -> bool:
39 return url.startswith("file://") or super().validateLink(url)
42def _answer_markdown_parser() -> MarkdownIt:
43 """Textual's default gfm-like parser with ``file:`` links admitted."""
44 return _AnswerMarkdownIt("gfm-like")
47_REASONING_BLOCK_CLASS = "reasoning-block"
48_REASONING_STREAMING_CLASS = "-streaming"
50_CSS_FILE = Path(__file__).parent / "message.tcss"
51_MESSAGE_CSS = _CSS_FILE.read_text(encoding="utf-8")
54class UserMessage(Vertical):
55 """A user's question in the chat log."""
57 DEFAULT_CSS: ClassVar[str] = _MESSAGE_CSS
59 def __init__(self, text: str) -> None:
60 super().__init__(classes="user-message")
61 self._text = text
63 def compose(self) -> ComposeResult:
64 yield Static(_SPEAKER_YOU, classes="speaker-label")
65 # Content() renders the question literally: a user asking about e.g. arr[0]
66 # or "[/]" must not have it parsed as console markup (which would crash).
67 yield Static(Content(self._text), classes="message-content")
70class AssistantMessage(Vertical):
71 """An assistant's response with streaming markdown, reasoning, and citations."""
73 DEFAULT_CSS: ClassVar[str] = _MESSAGE_CSS
75 def __init__(self, content: str = "", sources: Sequence[str] = ()) -> None:
76 """A live answer bubble, or a finished one restored from a saved session.
78 A restored turn's *content* must be passed here, not appended after
79 mounting: mount() is async, so compose has not run and appends no-op
80 against a still-None ``_content_widget``, silently dropping the text.
82 Sources render ONE way: the clickable numbered ``Sources:`` list a live
83 answer carries in its text. A turn arriving with structured *sources*
84 but no in-text list (seeded or written over HTTP/MCP, where the answer
85 text and the sources array are stored side by side) gets the same list
86 synthesized into its content, so a mixed transcript reads uniformly
87 instead of alternating between two citation styles.
88 """
89 super().__init__(classes="assistant-message")
90 self._reasoning_parts: list[str] = []
91 if content and sources:
92 content = _ensure_sources_block(content, sources)
93 self._content_parts: list[str] = [content] if content else []
94 # A restored turn is finished by definition: it must not raise a spinner.
95 self._finished = bool(content)
96 self._content_widget: Markdown | Static | None = None
97 self._reasoning_widget: Collapsible | None = None
98 self._reasoning_static: Static | None = None
99 self._thinking_header: ThinkingHeader | None = None
100 self._last_md_update: float = 0.0
101 self._last_reasoning_update: float = 0.0
102 self._use_markdown: bool = cfg.markdown_rendering
104 def compose(self) -> ComposeResult:
105 yield Static(_SPEAKER_LILBEE, classes="speaker-label")
106 # Built with the restored text (empty for a live turn, which streams in).
107 self._content_widget = self._build_content_widget("".join(self._content_parts))
108 yield self._content_widget
110 def on_mount(self) -> None:
111 """Raise the thinking header, unless this turn is already finished.
113 ``compose`` populates ``_content_widget`` before this hook runs. A
114 restored turn's answer is already on screen, so a spinner would claim
115 it is still being written.
116 """
117 if self._content_widget is None or self._finished:
118 return
119 header = ThinkingHeader()
120 self._thinking_header = header
121 self.mount(header, before=self._content_widget)
123 def _build_content_widget(self, text: str = "") -> Markdown | Static:
124 """Create the content widget based on the current rendering mode.
126 *text* must be passed at construction rather than via a later
127 ``update()``: Textual's Markdown re-renders from its constructor
128 argument on mount, discarding any pre-mount update.
130 ``open_links=False``: clicks route to the chat screen's link handler,
131 which opens ``file:`` citations with the OS opener instead of the
132 browser the default handling would use.
133 """
134 if self._use_markdown:
135 return Markdown(
136 text,
137 classes="response-md",
138 parser_factory=_answer_markdown_parser,
139 open_links=False,
140 )
141 return Static(Content(text), classes="response-md")
143 @property
144 def use_markdown(self) -> bool:
145 """Whether this message is using Markdown rendering."""
146 return self._use_markdown
148 async def rebuild_content_widget(self, use_markdown: bool) -> None:
149 """Replace the content widget with a different rendering mode."""
150 if self._content_widget is None:
151 return
152 self._use_markdown = use_markdown
153 old = self._content_widget
154 new_widget = self._build_content_widget("".join(self._content_parts))
155 await self.mount(new_widget, after=old)
156 self._content_widget = new_widget
157 await old.remove()
159 @staticmethod
160 def _set_content(widget: Markdown | Static, text: str) -> None:
161 """Update a content widget with raw model text. A Markdown widget consumes
162 the raw markdown string, but a Static parses console markup -- so wrap the
163 text as literal Content. Otherwise a ``[..]`` in the answer (quoted code, an
164 option like ``[/path]``) raises MarkupError and crashes the whole TUI.
165 """
166 if isinstance(widget, Markdown):
167 widget.update(text)
168 else:
169 widget.update(Content(text))
171 def append_reasoning(self, text: str) -> None:
172 """Append a reasoning token; debounced at ``_MD_UPDATE_INTERVAL``."""
173 first_token = not self._reasoning_parts
174 self._reasoning_parts.append(text)
175 if first_token and self._reasoning_widget is None:
176 self._mount_reasoning_collapsible()
177 now = time.monotonic()
178 ready = now - self._last_reasoning_update >= _MD_UPDATE_INTERVAL
179 if self._reasoning_static is not None and ready:
180 self._last_reasoning_update = now
181 self._reasoning_static.update(Content("".join(self._reasoning_parts)))
183 def set_thinking_status(self, detail: str) -> None:
184 """Show *detail* beside the thinking animator (e.g. an engine-load phase)."""
185 if self._thinking_header is not None:
186 self._thinking_header.set_status(detail)
188 def append_content(self, text: str) -> None:
189 """Append response content token (debounced markdown updates)."""
190 first_token = not self._content_parts
191 self._content_parts.append(text)
192 if first_token and not self._reasoning_parts:
193 # No reasoning ever arrived; drop the standalone header.
194 self._dismiss_thinking_header()
195 now = time.monotonic()
196 if self._content_widget is not None and now - self._last_md_update >= _MD_UPDATE_INTERVAL:
197 self._last_md_update = now
198 self._set_content(self._content_widget, "".join(self._content_parts))
199 self.refresh()
201 def finish(self, sources: list[str] | None = None) -> None:
202 """Mark response as complete, folding any structured *sources* into the
203 answer's own ``Sources:`` list (live RAG answers already carry one)."""
204 self._finished = True
205 # Always retire the standalone header on finish; the reasoning fold
206 # (if mounted) carries the post-stream title.
207 self._dismiss_thinking_header()
208 if sources and self._content_parts:
209 joined = _ensure_sources_block("".join(self._content_parts), sources)
210 self._content_parts = [joined]
211 if self._content_widget is not None and self._content_parts:
212 self._set_content(self._content_widget, "".join(self._content_parts))
213 self.refresh()
214 if self._reasoning_widget is not None and self._reasoning_parts:
215 if self._reasoning_static is not None:
216 self._reasoning_static.update(Content("".join(self._reasoning_parts)))
217 token_count = len("".join(self._reasoning_parts).split())
218 self._reasoning_widget.remove_class(_REASONING_STREAMING_CLASS)
219 self._reasoning_widget.title = msg.CHAT_REASONING_FINISHED.format(tokens=token_count)
220 self._reasoning_widget.collapsed = True
222 def _mount_reasoning_collapsible(self) -> None:
223 """Mount the reasoning Collapsible with the streaming-state class.
225 Called from ``append_reasoning`` on the first reasoning token, after
226 the message itself is mounted. The Collapsible slots in beneath the
227 ``ThinkingHeader`` so the animator continues to drive the visual
228 weight while the toggle row is hidden by the ``-streaming`` rule.
229 """
230 classes = f"{_REASONING_BLOCK_CLASS} {_REASONING_STREAMING_CLASS}"
231 self._reasoning_static = Static("", classes="reasoning-text")
232 collapsible = Collapsible(
233 self._reasoning_static,
234 title=msg.CHAT_REASONING_FINISHED.format(tokens=0),
235 collapsed=False,
236 classes=classes,
237 )
238 self._reasoning_widget = collapsible
239 header = self._thinking_header
240 if header is not None and header.is_mounted:
241 self.mount(collapsible, after=header)
242 return
243 content = self._content_widget
244 if content is not None:
245 self.mount(collapsible, before=content)
247 def _dismiss_thinking_header(self) -> None:
248 """Stop the animator and remove the standalone header from the DOM."""
249 header = self._thinking_header
250 if header is None:
251 return
252 header.stop()
253 if header.is_mounted:
254 header.remove()
255 self._thinking_header = None
258def _ensure_sources_block(content: str, sources: Sequence[str]) -> str:
259 """Append the same numbered, clickable ``Sources:`` list a live answer
260 carries, unless *content* already ends in one. One citation rendering
261 everywhere: a transcript mixing TUI-saved and API-saved turns must not
262 alternate styles."""
263 # Lazy: formatting transitively imports the store stack, which widget
264 # import must not pay at TUI startup.
265 from lilbee.retrieval.query.formatting import SOURCES_BLOCK_MARKER, source_markdown_link
267 if SOURCES_BLOCK_MARKER in content:
268 return content
269 lines = [f"{i}. {source_markdown_link(s)}" for i, s in enumerate(sources, 1)]
270 return content.rstrip() + SOURCES_BLOCK_MARKER + "\n" + "\n".join(lines)