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 

26 

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

28 return None if val == 0 else val 

29 

30 

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

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

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

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

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

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

37 return Excerpt( 

38 content=chunk.chunk, 

39 page_start=_zero_to_none(chunk.page_start), 

40 page_end=_zero_to_none(chunk.page_end), 

41 line_start=_zero_to_none(chunk.line_start), 

42 line_end=_zero_to_none(chunk.line_end), 

43 relevance=relevance, 

44 ) 

45 

46 

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

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

49 

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

51 than raising. 

52 """ 

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

54 

55 

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

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

58 from lilbee.app.search import resolve_vault_path 

59 

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

61 for chunk in chunks: 

62 source = chunk.source 

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

64 

65 results: list[DocumentResult] = [] 

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

67 excerpts = sorted( 

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

69 key=lambda e: e.relevance, 

70 reverse=True, 

71 ) 

72 results.append( 

73 DocumentResult( 

74 source=source, 

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

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

77 # can carry chunks of more than one type. 

78 content_type=_best_content_type(source_chunks), 

79 excerpts=excerpts, 

80 best_relevance=excerpts[0].relevance, 

81 vault_path=resolve_vault_path(source), 

82 ) 

83 ) 

84 

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

86 return results 

87 

88 

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

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

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