Coverage for src/lilbee/cli/tui/widgets/message.py: 100%
148 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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 TYPE_CHECKING, 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
21if TYPE_CHECKING:
22 from lilbee.retrieval.query.formatting import SourceRef
24# Minimum interval (seconds) between markdown widget updates during streaming
25_MD_UPDATE_INTERVAL = 0.1
27_SPEAKER_YOU = "[bold $primary]you[/]"
28_SPEAKER_LILBEE = "[bold $success]lilbee[/]"
31class _AnswerMarkdownIt(MarkdownIt):
32 """Markdown parser for answers, additionally allowing ``file:`` links.
34 markdown-it rejects ``file:`` destinations by default (a web-context XSS
35 guard), which left the Sources block's citation links rendering as raw
36 ``[label](file://...)`` text. Answers link to the reader's own documents on
37 disk, so re-admit the scheme; everything else keeps the default validation.
38 """
40 @override
41 def validateLink(self, url: str) -> bool:
42 return url.startswith("file://") or super().validateLink(url)
45def _answer_markdown_parser() -> MarkdownIt:
46 """Textual's default gfm-like parser with ``file:`` links admitted."""
47 return _AnswerMarkdownIt("gfm-like")
50def _shown_text(text: str, sources: Sequence[SourceRef], markdown: bool) -> str:
51 """*text* followed by the Sources list of *sources*, as markdown with each label and
52 locator shown as written when *markdown*, and made safe for a terminal."""
53 # Lazy: formatting transitively imports the store stack, which widget
54 # import must not pay at TUI startup.
55 from lilbee.retrieval.query.formatting import (
56 append_sources_block,
57 markdown_source_line,
58 source_ref_line,
59 terminal_sources_block,
60 )
62 if sources:
63 render = markdown_source_line if markdown else source_ref_line
64 text = append_sources_block(text, sources, render)
65 return terminal_sources_block(text)
68_REASONING_BLOCK_CLASS = "reasoning-block"
69_REASONING_STREAMING_CLASS = "-streaming"
71_CSS_FILE = Path(__file__).parent / "message.tcss"
72_MESSAGE_CSS = _CSS_FILE.read_text(encoding="utf-8")
75class UserMessage(Vertical):
76 """A user's question in the chat log."""
78 DEFAULT_CSS: ClassVar[str] = _MESSAGE_CSS
80 def __init__(self, text: str) -> None:
81 super().__init__(classes="user-message")
82 self._text = text
84 def compose(self) -> ComposeResult:
85 yield Static(_SPEAKER_YOU, classes="speaker-label")
86 # Content() renders the question literally: a user asking about e.g. arr[0]
87 # or "[/]" must not have it parsed as console markup (which would crash).
88 yield Static(Content(self._text), classes="message-content")
91class AssistantMessage(Vertical):
92 """An assistant's response with streaming markdown, reasoning, and citations."""
94 DEFAULT_CSS: ClassVar[str] = _MESSAGE_CSS
96 def __init__(self, content: str = "", sources: Sequence[SourceRef] = ()) -> None:
97 """A live answer bubble, or a finished one restored from a saved session.
99 A restored turn's *content* must be passed here, not appended after
100 mounting: mount() is async, so compose has not run and appends no-op
101 against a still-None ``_content_widget``, silently dropping the text.
103 Sources render ONE way: the clickable numbered ``Sources:`` list under the
104 answer text, built from *sources* (or from ``show_sources`` for a live
105 answer) whatever the text says, each link's target the file of its ref's
106 own source.
107 """
108 super().__init__(classes="assistant-message")
109 self._reasoning_parts: list[str] = []
110 self._content_parts: list[str] = [content] if content else []
111 self._sources: Sequence[SourceRef] = sources
112 # A restored turn is finished by definition: it must not raise a spinner.
113 self._finished = bool(content)
114 self._content_widget: Markdown | Static | None = None
115 self._reasoning_widget: Collapsible | None = None
116 self._reasoning_static: Static | None = None
117 self._thinking_header: ThinkingHeader | None = None
118 self._last_md_update: float = 0.0
119 self._last_reasoning_update: float = 0.0
120 self._use_markdown: bool = cfg.markdown_rendering
122 def compose(self) -> ComposeResult:
123 yield Static(_SPEAKER_LILBEE, classes="speaker-label")
124 # Built with the restored text (empty for a live turn, which streams in).
125 self._content_widget = self._build_content_widget()
126 yield self._content_widget
128 def on_mount(self) -> None:
129 """Raise the thinking header, unless this turn is already finished.
131 ``compose`` populates ``_content_widget`` before this hook runs. A
132 restored turn's answer is already on screen, so a spinner would claim
133 it is still being written.
134 """
135 if self._content_widget is None or self._finished:
136 return
137 header = ThinkingHeader()
138 self._thinking_header = header
139 self.mount(header, before=self._content_widget)
141 def _shown(self, markdown: bool) -> str:
142 """The answer text and its Sources list, as markdown when *markdown*."""
143 return _shown_text("".join(self._content_parts), self._sources, markdown)
145 def _build_content_widget(self) -> Markdown | Static:
146 """Create the content widget based on the current rendering mode.
148 The text is passed at construction rather than via a later
149 ``update()``: Textual's Markdown re-renders from its constructor
150 argument on mount, discarding any pre-mount update.
152 ``open_links=False``: clicks route to the chat screen's link handler,
153 which opens ``file:`` citations with the OS opener instead of the
154 browser the default handling would use.
155 """
156 if self._use_markdown:
157 return Markdown(
158 self._shown(markdown=True),
159 classes="response-md",
160 parser_factory=_answer_markdown_parser,
161 open_links=False,
162 )
163 return Static(Content(self._shown(markdown=False)), classes="response-md")
165 @property
166 def use_markdown(self) -> bool:
167 """Whether this message is using Markdown rendering."""
168 return self._use_markdown
170 async def rebuild_content_widget(self, use_markdown: bool) -> None:
171 """Replace the content widget with a different rendering mode."""
172 if self._content_widget is None:
173 return
174 self._use_markdown = use_markdown
175 old = self._content_widget
176 new_widget = self._build_content_widget()
177 await self.mount(new_widget, after=old)
178 self._content_widget = new_widget
179 await old.remove()
181 def _set_content(self, widget: Markdown | Static) -> None:
182 """Update a content widget with the raw model text. A Markdown widget consumes
183 the raw markdown string, but a Static parses console markup -- so wrap the
184 text as literal Content. Otherwise a ``[..]`` in the answer (quoted code, an
185 option like ``[/path]``) raises MarkupError and crashes the whole TUI.
186 """
187 if isinstance(widget, Markdown):
188 widget.update(self._shown(markdown=True))
189 else:
190 widget.update(Content(self._shown(markdown=False)))
192 def show_sources(self, sources: Sequence[SourceRef]) -> None:
193 """Show *sources* as the Sources list under the answer."""
194 self._sources = sources
195 if self._content_widget is not None:
196 self._set_content(self._content_widget)
198 def append_reasoning(self, text: str) -> None:
199 """Append a reasoning token; debounced at ``_MD_UPDATE_INTERVAL``."""
200 first_token = not self._reasoning_parts
201 self._reasoning_parts.append(text)
202 if first_token and self._reasoning_widget is None:
203 self._mount_reasoning_collapsible()
204 now = time.monotonic()
205 ready = now - self._last_reasoning_update >= _MD_UPDATE_INTERVAL
206 if self._reasoning_static is not None and ready:
207 self._last_reasoning_update = now
208 self._reasoning_static.update(Content("".join(self._reasoning_parts)))
210 def set_thinking_status(self, detail: str) -> None:
211 """Show *detail* beside the thinking animator (e.g. an engine-load phase)."""
212 if self._thinking_header is not None:
213 self._thinking_header.set_status(detail)
215 def append_content(self, text: str) -> None:
216 """Append response content token (debounced markdown updates)."""
217 first_token = not self._content_parts
218 self._content_parts.append(text)
219 if first_token and not self._reasoning_parts:
220 # No reasoning ever arrived; drop the standalone header.
221 self._dismiss_thinking_header()
222 now = time.monotonic()
223 if self._content_widget is not None and now - self._last_md_update >= _MD_UPDATE_INTERVAL:
224 self._last_md_update = now
225 self._set_content(self._content_widget)
226 self.refresh()
228 def finish(self) -> None:
229 """Mark response as complete."""
230 self._finished = True
231 # Always retire the standalone header on finish; the reasoning fold
232 # (if mounted) carries the post-stream title.
233 self._dismiss_thinking_header()
234 if self._content_widget is not None and self._content_parts:
235 self._set_content(self._content_widget)
236 self.refresh()
237 if self._reasoning_widget is not None and self._reasoning_parts:
238 if self._reasoning_static is not None:
239 self._reasoning_static.update(Content("".join(self._reasoning_parts)))
240 token_count = len("".join(self._reasoning_parts).split())
241 self._reasoning_widget.remove_class(_REASONING_STREAMING_CLASS)
242 self._reasoning_widget.title = msg.CHAT_REASONING_FINISHED.format(tokens=token_count)
243 self._reasoning_widget.collapsed = True
245 def _mount_reasoning_collapsible(self) -> None:
246 """Mount the reasoning Collapsible with the streaming-state class.
248 Called from ``append_reasoning`` on the first reasoning token, after
249 the message itself is mounted. The Collapsible slots in beneath the
250 ``ThinkingHeader`` so the animator continues to drive the visual
251 weight while the toggle row is hidden by the ``-streaming`` rule.
252 """
253 classes = f"{_REASONING_BLOCK_CLASS} {_REASONING_STREAMING_CLASS}"
254 self._reasoning_static = Static("", classes="reasoning-text")
255 collapsible = Collapsible(
256 self._reasoning_static,
257 title=msg.CHAT_REASONING_FINISHED.format(tokens=0),
258 collapsed=False,
259 classes=classes,
260 )
261 self._reasoning_widget = collapsible
262 header = self._thinking_header
263 if header is not None and header.is_mounted:
264 self.mount(collapsible, after=header)
265 return
266 content = self._content_widget
267 if content is not None:
268 self.mount(collapsible, before=content)
270 def _dismiss_thinking_header(self) -> None:
271 """Stop the animator and remove the standalone header from the DOM."""
272 header = self._thinking_header
273 if header is None:
274 return
275 header.stop()
276 if header.is_mounted:
277 header.remove()
278 self._thinking_header = None