Coverage for src/lilbee/cli/tui/widgets/message.py: 100%
143 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +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 # Lazy: formatting transitively imports the store stack, which widget
93 # import must not pay at TUI startup.
94 from lilbee.retrieval.query.formatting import with_sources_block
96 content = with_sources_block(content, sources)
97 self._content_parts: list[str] = [content] if content else []
98 # A restored turn is finished by definition: it must not raise a spinner.
99 self._finished = bool(content)
100 self._content_widget: Markdown | Static | None = None
101 self._reasoning_widget: Collapsible | None = None
102 self._reasoning_static: Static | None = None
103 self._thinking_header: ThinkingHeader | None = None
104 self._last_md_update: float = 0.0
105 self._last_reasoning_update: float = 0.0
106 self._use_markdown: bool = cfg.markdown_rendering
108 def compose(self) -> ComposeResult:
109 yield Static(_SPEAKER_LILBEE, classes="speaker-label")
110 # Built with the restored text (empty for a live turn, which streams in).
111 self._content_widget = self._build_content_widget("".join(self._content_parts))
112 yield self._content_widget
114 def on_mount(self) -> None:
115 """Raise the thinking header, unless this turn is already finished.
117 ``compose`` populates ``_content_widget`` before this hook runs. A
118 restored turn's answer is already on screen, so a spinner would claim
119 it is still being written.
120 """
121 if self._content_widget is None or self._finished:
122 return
123 header = ThinkingHeader()
124 self._thinking_header = header
125 self.mount(header, before=self._content_widget)
127 def _build_content_widget(self, text: str = "") -> Markdown | Static:
128 """Create the content widget based on the current rendering mode.
130 *text* must be passed at construction rather than via a later
131 ``update()``: Textual's Markdown re-renders from its constructor
132 argument on mount, discarding any pre-mount update.
134 ``open_links=False``: clicks route to the chat screen's link handler,
135 which opens ``file:`` citations with the OS opener instead of the
136 browser the default handling would use.
137 """
138 if self._use_markdown:
139 return Markdown(
140 text,
141 classes="response-md",
142 parser_factory=_answer_markdown_parser,
143 open_links=False,
144 )
145 return Static(Content(text), classes="response-md")
147 @property
148 def use_markdown(self) -> bool:
149 """Whether this message is using Markdown rendering."""
150 return self._use_markdown
152 async def rebuild_content_widget(self, use_markdown: bool) -> None:
153 """Replace the content widget with a different rendering mode."""
154 if self._content_widget is None:
155 return
156 self._use_markdown = use_markdown
157 old = self._content_widget
158 new_widget = self._build_content_widget("".join(self._content_parts))
159 await self.mount(new_widget, after=old)
160 self._content_widget = new_widget
161 await old.remove()
163 @staticmethod
164 def _set_content(widget: Markdown | Static, text: str) -> None:
165 """Update a content widget with raw model text. A Markdown widget consumes
166 the raw markdown string, but a Static parses console markup -- so wrap the
167 text as literal Content. Otherwise a ``[..]`` in the answer (quoted code, an
168 option like ``[/path]``) raises MarkupError and crashes the whole TUI.
169 """
170 if isinstance(widget, Markdown):
171 widget.update(text)
172 else:
173 widget.update(Content(text))
175 def append_reasoning(self, text: str) -> None:
176 """Append a reasoning token; debounced at ``_MD_UPDATE_INTERVAL``."""
177 first_token = not self._reasoning_parts
178 self._reasoning_parts.append(text)
179 if first_token and self._reasoning_widget is None:
180 self._mount_reasoning_collapsible()
181 now = time.monotonic()
182 ready = now - self._last_reasoning_update >= _MD_UPDATE_INTERVAL
183 if self._reasoning_static is not None and ready:
184 self._last_reasoning_update = now
185 self._reasoning_static.update(Content("".join(self._reasoning_parts)))
187 def set_thinking_status(self, detail: str) -> None:
188 """Show *detail* beside the thinking animator (e.g. an engine-load phase)."""
189 if self._thinking_header is not None:
190 self._thinking_header.set_status(detail)
192 def append_content(self, text: str) -> None:
193 """Append response content token (debounced markdown updates)."""
194 first_token = not self._content_parts
195 self._content_parts.append(text)
196 if first_token and not self._reasoning_parts:
197 # No reasoning ever arrived; drop the standalone header.
198 self._dismiss_thinking_header()
199 now = time.monotonic()
200 if self._content_widget is not None and now - self._last_md_update >= _MD_UPDATE_INTERVAL:
201 self._last_md_update = now
202 self._set_content(self._content_widget, "".join(self._content_parts))
203 self.refresh()
205 def finish(self, sources: list[str] | None = None) -> None:
206 """Mark response as complete, folding any structured *sources* into the
207 answer's own ``Sources:`` list (live RAG answers already carry one)."""
208 self._finished = True
209 # Always retire the standalone header on finish; the reasoning fold
210 # (if mounted) carries the post-stream title.
211 self._dismiss_thinking_header()
212 if sources and self._content_parts:
213 # Lazy: formatting transitively imports the store stack, which widget
214 # import must not pay at TUI startup.
215 from lilbee.retrieval.query.formatting import with_sources_block
217 joined = with_sources_block("".join(self._content_parts), sources)
218 self._content_parts = [joined]
219 if self._content_widget is not None and self._content_parts:
220 self._set_content(self._content_widget, "".join(self._content_parts))
221 self.refresh()
222 if self._reasoning_widget is not None and self._reasoning_parts:
223 if self._reasoning_static is not None:
224 self._reasoning_static.update(Content("".join(self._reasoning_parts)))
225 token_count = len("".join(self._reasoning_parts).split())
226 self._reasoning_widget.remove_class(_REASONING_STREAMING_CLASS)
227 self._reasoning_widget.title = msg.CHAT_REASONING_FINISHED.format(tokens=token_count)
228 self._reasoning_widget.collapsed = True
230 def _mount_reasoning_collapsible(self) -> None:
231 """Mount the reasoning Collapsible with the streaming-state class.
233 Called from ``append_reasoning`` on the first reasoning token, after
234 the message itself is mounted. The Collapsible slots in beneath the
235 ``ThinkingHeader`` so the animator continues to drive the visual
236 weight while the toggle row is hidden by the ``-streaming`` rule.
237 """
238 classes = f"{_REASONING_BLOCK_CLASS} {_REASONING_STREAMING_CLASS}"
239 self._reasoning_static = Static("", classes="reasoning-text")
240 collapsible = Collapsible(
241 self._reasoning_static,
242 title=msg.CHAT_REASONING_FINISHED.format(tokens=0),
243 collapsed=False,
244 classes=classes,
245 )
246 self._reasoning_widget = collapsible
247 header = self._thinking_header
248 if header is not None and header.is_mounted:
249 self.mount(collapsible, after=header)
250 return
251 content = self._content_widget
252 if content is not None:
253 self.mount(collapsible, before=content)
255 def _dismiss_thinking_header(self) -> None:
256 """Stop the animator and remove the standalone header from the DOM."""
257 header = self._thinking_header
258 if header is None:
259 return
260 header.stop()
261 if header.is_mounted:
262 header.remove()
263 self._thinking_header = None