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

1"""Status snapshot of the local knowledge base.""" 

2 

3from __future__ import annotations 

4 

5import os 

6from pathlib import Path 

7 

8from pydantic import BaseModel 

9 

10from lilbee.app.services import get_services 

11from lilbee.core.config import cfg 

12from lilbee.core.system import LOCAL_ROOT_DIRNAME, default_data_dir 

13 

14LILBEE_LABEL_MAX_LEN = 40 

15"""Hard cap on the compact-label width. The leaf gets its own internal 

16ellipsis when even the leaf alone would breach the cap.""" 

17 

18_ELLIPSIS = "…" 

19 

20 

21def _project_root() -> Path: 

22 """Walk past a trailing ``.lilbee`` marker to the project dir that owns it.""" 

23 root = cfg.data_root 

24 if root.name == LOCAL_ROOT_DIRNAME: 

25 return root.parent 

26 return root 

27 

28 

29def _truncate_leaf(leaf: str, max_len: int) -> str: 

30 """Shrink an over-long leaf to fit a budget, with an internal ellipsis.""" 

31 if len(leaf) <= max_len: 

32 return leaf 

33 if max_len <= 1: 

34 return _ELLIPSIS 

35 keep = max_len - 1 

36 head = keep // 2 

37 tail = keep - head 

38 return f"{leaf[:head]}{_ELLIPSIS}{leaf[-tail:] if tail else ''}" 

39 

40 

41def _compact_path(full: str) -> str: 

42 """Render *full* with ``~`` substituted for ``$HOME`` when it leads.""" 

43 home = str(Path.home()) 

44 if full == home: 

45 return "~" 

46 home_prefix = f"{home}{os.sep}" 

47 return f"~{os.sep}{full[len(home_prefix) :]}" if full.startswith(home_prefix) else full 

48 

49 

50def lilbee_label() -> str: 

51 """Status-bar pill text for the active lilbee. 

52 

53 Precedence: ``lilbee_name`` override > ``"global"`` (when data_root 

54 is the platform default) > project path. ``show_lilbee_path`` 

55 (toggled by F4) returns the full absolute path instead of the 

56 compact / "global" form. 

57 """ 

58 if cfg.lilbee_name: 

59 return cfg.lilbee_name 

60 is_global = cfg.data_root.expanduser().resolve() == default_data_dir().resolve() 

61 if cfg.show_lilbee_path: 

62 return str(default_data_dir() if is_global else _project_root().expanduser().resolve()) 

63 if is_global: 

64 return "global" 

65 full = str(_project_root().expanduser().resolve()) 

66 compact = _compact_path(full) 

67 if len(compact) <= LILBEE_LABEL_MAX_LEN: 

68 return compact 

69 leaf = _project_root().name or compact 

70 leaf_budget = LILBEE_LABEL_MAX_LEN - 1 - len(os.sep) 

71 return f"{_ELLIPSIS}{os.sep}{_truncate_leaf(leaf, leaf_budget)}" 

72 

73 

74class StatusConfig(BaseModel): 

75 """Configuration section of a status response. 

76 

77 Exposes all four role-bound model fields (chat, embedding, vision, 

78 reranker) so the TUI status screen and plugin callers can show 

79 what's active per role. 

80 """ 

81 

82 documents_dir: str 

83 data_dir: str 

84 chat_model: str 

85 embedding_model: str 

86 vision_model: str = "" 

87 reranker_model: str = "" 

88 enable_ocr: bool | None = None 

89 

90 

91class SourceInfo(BaseModel): 

92 """A single indexed source in a status response.""" 

93 

94 filename: str 

95 file_hash: str 

96 chunk_count: int 

97 ingested_at: str 

98 

99 

100class StatusResult(BaseModel): 

101 """Full status response for the knowledge base.""" 

102 

103 command: str = "status" 

104 config: StatusConfig 

105 sources: list[SourceInfo] 

106 total_chunks: int 

107 

108 

109def gather_status() -> StatusResult: 

110 """Collect status data as a typed model (shared by human + JSON output).""" 

111 sources = get_services().store.get_sources() 

112 sorted_sources = sorted(sources, key=lambda x: x["filename"]) 

113 total_chunks = sum(s["chunk_count"] for s in sources) 

114 return StatusResult( 

115 config=StatusConfig( 

116 documents_dir=str(cfg.documents_dir), 

117 data_dir=str(cfg.data_dir), 

118 chat_model=cfg.chat_model, 

119 embedding_model=cfg.embedding_model, 

120 vision_model=cfg.vision_model, 

121 reranker_model=cfg.reranker_model, 

122 enable_ocr=cfg.enable_ocr, 

123 ), 

124 sources=[ 

125 SourceInfo( 

126 filename=s["filename"], 

127 file_hash=s["file_hash"][:12], 

128 chunk_count=s["chunk_count"], 

129 ingested_at=s["ingested_at"][:19], 

130 ) 

131 for s in sorted_sources 

132 ], 

133 total_chunks=total_chunks, 

134 )