Coverage for src/lilbee/app/session_sources.py: 100%

1"""The Sources list of a saved chat turn, as every session surface returns it.""" 

2 

3from __future__ import annotations 

4 

5import re 

6from collections.abc import Sequence 

7from pathlib import Path 

8from typing import Any, NotRequired 

9from urllib.parse import urlparse 

10from urllib.request import url2pathname 

11 

12from typing_extensions import TypedDict 

13 

14from lilbee.app.titles import stored_web_titles 

15from lilbee.core import settings 

16from lilbee.core.config import active_config 

17from lilbee.data.store import is_memory_source 

18from lilbee.retrieval.query.formatting import ( 

19 FILE_LINK_RE, 

20 SourceRef, 

21 is_locator, 

22 source_label, 

23 split_sources_list, 

24) 

25from lilbee.sessions import MessageRole, SessionMessage, StoredSourceRef, message_with_refs 

26 

27 

28class SourceRefInput(TypedDict): 

29 """One entry of a turn's Sources list as a client saves it.""" 

30 

31 source: str 

32 # The ``[n]`` the answer cites the source by; absent or null when unknown. 

33 citation_number: NotRequired[int | None] 

34 # The page or line span the answer read, such as "page 3"; absent when none. 

35 locator: NotRequired[str] 

36 

37 

38LOCATOR_ERROR = ( 

39 "locator {locator!r} is not a page or line span: " 

40 "use 'page N', 'pages A-B', 'line N' or 'lines A-B', or leave it out" 

41) 

42# One numbered line of a stored Sources list; a hard-broken line ends in a backslash. 

43_LIST_LINE_RE = re.compile(r"^(?P<number>\d+)\\?\. (?P<item>.*?)\\?$") 

44 

45 

46class LocatorError(ValueError): 

47 """A posted Sources ref names a locator that lilbee never writes.""" 

48 

49 

50def message_sources(message: SessionMessage) -> tuple[str, list[SourceRef]]: 

51 """*message*'s content without lilbee's Sources list, and that list as refs. 

52 

53 A turn saved with ``source_refs`` gives its content whole and them, each 

54 labelled from the current stored title. An older turn whose content ends in 

55 lilbee's Sources list gives the refs read from that list: its numbers and its 

56 labels as written. Any other older turn gives its ``sources`` with no number, 

57 as their positions are not the answer's citation numbers. 

58 """ 

59 if message.source_refs is not None: 

60 return message.content, labelled(message.source_refs) 

61 text, listed = split_sources_list(message.content) 

62 refs = _listed_refs(listed) 

63 if refs: 

64 return text, refs 

65 return message.content, labelled(tuple(StoredSourceRef(s) for s in message.sources)) 

66 

67 

68def message_payload(message: SessionMessage) -> dict[str, Any]: 

69 """*message* as the session surfaces send it, its Sources list in ``source_refs``.""" 

70 content, refs = message_sources(message) 

71 return { 

72 "role": message.role.value, 

73 "content": content, 

74 "sources": list(message.sources), 

75 "ts": message.ts, 

76 "source_refs": [ 

77 { 

78 "source": r.source, 

79 "label": r.label, 

80 "citation_number": r.citation_number, 

81 "locator": r.locator, 

82 } 

83 for r in refs 

84 ], 

85 } 

86 

87 

88def posted_message( 

89 role: MessageRole, 

90 content: str, 

91 sources: Sequence[str], 

92 refs: Sequence[StoredSourceRef] | None, 

93) -> SessionMessage: 

94 """The turn a client appends: with its Sources list *refs* when given, ``sources`` 

95 then derived from it; else *sources* alone. 

96 

97 Raises LocatorError when a ref's locator is not one lilbee writes. 

98 """ 

99 if refs is None: 

100 return SessionMessage(role=role, content=content, sources=tuple(sources)) 

101 for ref in refs: 

102 if not is_locator(ref.locator): 

103 raise LocatorError(LOCATOR_ERROR.format(locator=ref.locator)) 

104 return message_with_refs(role, content, tuple(refs)) 

105 

106 

107def labelled(refs: tuple[StoredSourceRef, ...]) -> list[SourceRef]: 

108 """*refs* with the label of each source from its current stored title.""" 

109 titles = stored_web_titles({r.source for r in refs}) 

110 return [ 

111 SourceRef( 

112 r.source, source_label(r.source, titles.get(r.source, "")), r.citation_number, r.locator 

113 ) 

114 for r in refs 

115 ] 

116 

117 

118def _listed_refs(listed: str) -> list[SourceRef]: 

119 """The refs of a stored Sources list, one per numbered line.""" 

120 refs = [] 

121 for line in listed.split("\n"): 

122 match = _LIST_LINE_RE.match(line) 

123 if match is None: 

124 continue 

125 item = match["item"] 

126 link = FILE_LINK_RE.search(item) 

127 locator = "" 

128 if link is not None: 

129 label, source = link["label"], _source_of(link["url"]) 

130 # A linked line's locator follows its link: "[label](url), page 3". 

131 locator = item[link.end() :].removeprefix(", ") 

132 else: 

133 label, source = item, item if is_memory_source(item) else "" 

134 refs.append(SourceRef(source, label, int(match["number"]), locator)) 

135 return refs 

136 

137 

138def _source_of(url: str) -> str: 

139 """The source key of the file *url* names, or '' when no known root holds it.""" 

140 path = Path(url2pathname(urlparse(url).path)) 

141 config = active_config() 

142 settings.refresh_registry(config) 

143 # The reverse of ``resolve_source_path``: registered roots own their label's keys. 

144 roots: list[tuple[str, str | Path]] = [*config.linked_roots.items(), ("", config.documents_dir)] 

145 for label, root in roots: 

146 try: 

147 relative = path.relative_to(Path(root).resolve()).as_posix() 

148 except (OSError, ValueError): 

149 continue 

150 return "/".join(part for part in (label, relative) if part and part != ".") 

151 return ""