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

1"""Search-result post-processing shared by CLI, HTTP, and MCP.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6from dataclasses import replace 

7from typing import TYPE_CHECKING 

8 

9from lilbee.core.config import cfg 

10 

11if TYPE_CHECKING: 

12 from lilbee.data.store import SearchChunk 

13 from lilbee.retrieval.query.formatting import SourceRef 

14 

15 

16def resolve_vault_path(source_filename: str) -> str | None: 

17 """Return *source_filename* as a vault-relative path, or None if unresolvable. 

18 

19 Resolves symlinks on both sides and rejects ``..`` escapes from 

20 ``documents_dir``. 

21 """ 

22 if cfg.vault_base is None: 

23 return None 

24 from lilbee.data.ingest.discovery import resolve_source_path 

25 

26 try: 

27 vault_base = cfg.vault_base.resolve() 

28 documents_dir = cfg.documents_dir.resolve() 

29 source_path = resolve_source_path(source_filename).resolve() 

30 source_path.relative_to(documents_dir) 

31 relative_docs_dir = documents_dir.relative_to(vault_base) 

32 except (OSError, ValueError): 

33 return None 

34 if not source_path.is_file(): 

35 return None 

36 return (relative_docs_dir / source_path.relative_to(documents_dir)).as_posix() 

37 

38 

39def clean_result(result: SearchChunk, ref: SourceRef | None = None) -> dict: 

40 """Return SearchChunk as a JSON dict, stamping vault_path when resolvable and the 

41 ``label`` and ``citation_number`` of *ref*, the number only when it has one.""" 

42 payload = result.model_dump(exclude={"vector"}, exclude_none=True) 

43 vault_path = resolve_vault_path(result.source) 

44 if vault_path is not None: 

45 payload["vault_path"] = vault_path 

46 if ref is not None: 

47 payload["label"] = ref.label 

48 if ref.citation_number is not None: 

49 payload["citation_number"] = ref.citation_number 

50 return payload 

51 

52 

53def clean_results( 

54 chunks: Sequence[SearchChunk], answered_from: Sequence[SearchChunk] = () 

55) -> list[dict]: 

56 """*chunks* as JSON dicts, each with the label of its source. 

57 

58 *answered_from* is the retrieved list an answer was written from. A chunk whose 

59 source is in it also carries the ``[n]`` the answer cites that source by, so a 

60 cited subset keeps its true numbers. Search results pass none and carry no number. 

61 """ 

62 refs = _source_refs(chunks, answered_from) 

63 return [clean_result(c, refs[c.source]) for c in chunks] 

64 

65 

66def source_labels(chunks: Sequence[SearchChunk]) -> dict[str, str]: 

67 """The label of each source among *chunks*, as ``clean_results`` stamps it.""" 

68 return {source: ref.label for source, ref in _source_refs(chunks, ()).items()} 

69 

70 

71def _source_refs( 

72 chunks: Sequence[SearchChunk], answered_from: Sequence[SearchChunk] 

73) -> dict[str, SourceRef]: 

74 """The ref of each source among *chunks*, numbered only by *answered_from*.""" 

75 from lilbee.app.titles import stored_web_titles 

76 from lilbee.retrieval.query.formatting import source_refs 

77 

78 titles = stored_web_titles({c.source for c in (*chunks, *answered_from)}) 

79 refs = {r.source: replace(r, citation_number=None) for r in source_refs([*chunks], titles)} 

80 refs.update((r.source, r) for r in source_refs([*answered_from], titles)) 

81 return refs