Coverage for src/lilbee/core/results.py: 100%

1from __future__ import annotations 

2 

3from pydantic import BaseModel 

4 

5from lilbee.data.store import SearchChunk 

6 

7 

8class Excerpt(BaseModel): 

9 content: str 

10 page_start: int | None 

11 page_end: int | None 

12 line_start: int | None 

13 line_end: int | None 

14 relevance: float # 0.0-1.0 (1 = best match) 

15 

16 

17class DocumentResult(BaseModel): 

18 source: str 

19 content_type: str 

20 excerpts: list[Excerpt] 

21 best_relevance: float 

22 # Vault-relative path for clients to deep-link into the native UI. 

23 # ``None`` when the server can't resolve the source under ``cfg.vault_base``. 

24 vault_path: str | None = None 

25 # The plain one-line name of the source a client shows (``host · title`` for a 

26 # crawled page). Never escaped: a client shows it as text. 

27 label: str | None = None 

28 

29 

30def _zero_to_none(val: int) -> int | None: 

31 return None if val == 0 else val 

32 

33 

34def _to_excerpt(chunk: SearchChunk) -> Excerpt: 

35 # The canonical [0, 1] score is what every retrieval path stamps; the 

36 # distance fallback (which read keyword-only rows as a perfect 1.0) 

37 # covers only hand-built chunks that never went through retrieval. 

38 fallback = 1.0 / (1.0 + (chunk.distance or 0)) 

39 relevance = chunk.score if chunk.score is not None else fallback 

40 return Excerpt( 

41 content=chunk.chunk, 

42 page_start=_zero_to_none(chunk.page_start), 

43 page_end=_zero_to_none(chunk.page_end), 

44 line_start=_zero_to_none(chunk.line_start), 

45 line_end=_zero_to_none(chunk.line_end), 

46 relevance=relevance, 

47 ) 

48 

49 

50def _best_content_type(source_chunks: list[SearchChunk]) -> str: 

51 """Content type of the highest-scoring chunk for a source. 

52 

53 score is optional on the model, so an unscored chunk sorts last rather 

54 than raising. 

55 """ 

56 return max(source_chunks, key=lambda c: c.score if c.score is not None else -1.0).content_type 

57 

58 

59def group(chunks: list[SearchChunk]) -> list[DocumentResult]: 

60 """Group raw LanceDB chunks into document-centric results.""" 

61 from lilbee.app.search import resolve_vault_path 

62 

63 by_source: dict[str, list[SearchChunk]] = {} 

64 for chunk in chunks: 

65 source = chunk.source 

66 by_source.setdefault(source, []).append(chunk) 

67 

68 results: list[DocumentResult] = [] 

69 for source, source_chunks in by_source.items(): 

70 excerpts = sorted( 

71 [_to_excerpt(c) for c in source_chunks], 

72 key=lambda e: e.relevance, 

73 reverse=True, 

74 ) 

75 results.append( 

76 DocumentResult( 

77 source=source, 

78 # From the best-scoring chunk, not whichever the store returned 

79 # first: excerpts are already sorted by relevance and a source 

80 # can carry chunks of more than one type. 

81 content_type=_best_content_type(source_chunks), 

82 excerpts=excerpts, 

83 best_relevance=excerpts[0].relevance, 

84 vault_path=resolve_vault_path(source), 

85 ) 

86 ) 

87 

88 results.sort(key=lambda r: r.best_relevance, reverse=True) 

89 return results 

90 

91 

92def to_dicts(results: list[DocumentResult]) -> list[dict[str, object]]: 

93 """Serialize DocumentResults to JSON-safe dicts.""" 

94 return [r.model_dump() for r in results]