Coverage for src/lilbee/cli/tui/widgets/fork_picker.py: 100%

1"""Modal that picks the answer a fork of the conversation ends on.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6from typing import ClassVar 

7 

8from textual.app import ComposeResult 

9from textual.binding import Binding, BindingType 

10from textual.containers import Vertical 

11from textual.content import Content 

12from textual.screen import ModalScreen 

13from textual.widgets import OptionList, Static 

14from textual.widgets.option_list import Option 

15 

16from lilbee.cli.tui import messages as msg 

17from lilbee.cli.tui.widgets.clamped_option_list import ClampedOptionList 

18from lilbee.sessions import MessageRole, SessionMessage, derive_title 

19 

20 

21@dataclass(frozen=True, slots=True) 

22class ForkPoint: 

23 """A place to fork: after one answer, with the question it answers if there is one.""" 

24 

25 message_count: int 

26 answer: str 

27 question: str | None 

28 

29 

30def fork_points(messages: tuple[SessionMessage, ...]) -> list[ForkPoint]: 

31 """One point after each answer in the saved log, newest first.""" 

32 points: list[ForkPoint] = [] 

33 question: str | None = None 

34 for index, message in enumerate(messages): 

35 if message.role == MessageRole.USER: 

36 question = message.content 

37 else: 

38 points.append(ForkPoint(index + 1, message.content, question)) 

39 return points[::-1] 

40 

41 

42def _label(point: ForkPoint) -> Content: 

43 """The answer's first line, and under it the question it answers.""" 

44 answer = Content(derive_title(point.answer)) 

45 if point.question is None: 

46 return answer 

47 question = msg.FORK_PICKER_ANSWER_TO.format(question=derive_title(point.question)) 

48 return Content.assemble(answer, "\n", Content.styled(question, "$text-muted")) 

49 

50 

51class ForkPicker(ModalScreen[int | None]): 

52 """Dismisses with the number of leading messages to copy, or ``None`` on escape.""" 

53 

54 CSS_PATH = "fork_picker.tcss" 

55 

56 BINDINGS: ClassVar[list[BindingType]] = [ 

57 Binding("escape", "cancel", "Close", show=True), 

58 ] 

59 

60 def __init__(self, points: list[ForkPoint]) -> None: 

61 super().__init__() 

62 self._points = points 

63 

64 def compose(self) -> ComposeResult: 

65 with Vertical(id="fork-root"): 

66 yield Static(msg.FORK_PICKER_TITLE, id="fork-title") 

67 yield ClampedOptionList( 

68 *(Option(_label(point)) for point in self._points), id="fork-list" 

69 ) 

70 yield Static(msg.FORK_PICKER_HINT, id="fork-hint") 

71 

72 def on_mount(self) -> None: 

73 self.query_one("#fork-list", OptionList).focus() 

74 

75 def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: 

76 event.stop() 

77 self.dismiss(self._points[event.option_index].message_count) 

78 

79 def action_cancel(self) -> None: 

80 self.dismiss(None)