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.config.enums import KvCacheType 

13from lilbee.core.system import LOCAL_ROOT_DIRNAME, default_data_dir 

14from lilbee.data.ingest.skip_marker import describe_skips, load_skip_markers 

15from lilbee.data.types import SkippedSource 

16 

17LILBEE_LABEL_MAX_LEN = 40 

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

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

20 

21STATUS_SKIPPED_LIMIT = 50 

22"""Cap on held-out files in one status response; ``skipped_total`` carries the real count.""" 

23 

24_ELLIPSIS = "…" 

25 

26 

27def _project_root() -> Path: 

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

29 root = cfg.data_root 

30 if root.name == LOCAL_ROOT_DIRNAME: 

31 return root.parent 

32 return root 

33 

34 

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

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

37 if len(leaf) <= max_len: 

38 return leaf 

39 if max_len <= 1: 

40 return _ELLIPSIS 

41 keep = max_len - 1 

42 head = keep // 2 

43 tail = keep - head 

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

45 

46 

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

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

49 home = str(Path.home()) 

50 if full == home: 

51 return "~" 

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

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

54 

55 

56def lilbee_label() -> str: 

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

58 

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

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

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

62 compact / "global" form. 

63 """ 

64 if cfg.lilbee_name: 

65 return cfg.lilbee_name 

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

67 if cfg.show_lilbee_path: 

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

69 if is_global: 

70 return "global" 

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

72 compact = _compact_path(full) 

73 if len(compact) <= LILBEE_LABEL_MAX_LEN: 

74 return compact 

75 leaf = _project_root().name or compact 

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

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

78 

79 

80class StatusConfig(BaseModel): 

81 """Configuration section of a status response. 

82 

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

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

85 what's active per role. 

86 """ 

87 

88 documents_dir: str 

89 data_dir: str 

90 chat_model: str 

91 embedding_model: str 

92 vision_model: str = "" 

93 reranker_model: str = "" 

94 enable_ocr: bool | None = None 

95 num_ctx: int | None = None 

96 num_ctx_max: int | None = None 

97 chat_n_ctx_target: int | None = None 

98 flash_attention: bool | None = None 

99 kv_cache_type: KvCacheType | None = None 

100 n_gpu_layers: int | None = None 

101 cpu_moe: bool | None = None 

102 n_cpu_moe: int | None = None 

103 main_gpu: int | None = None 

104 gpu_devices: str | None = None 

105 

106 

107class IndexStatus(BaseModel): 

108 """The embedder that built the persisted index.""" 

109 

110 embedding_model: str 

111 embedding_dim: int 

112 

113 

114class SourceInfo(BaseModel): 

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

116 

117 filename: str 

118 file_hash: str 

119 chunk_count: int 

120 ingested_at: str 

121 

122 

123class EntityStatus(BaseModel): 

124 """Entity-extraction section of a status response (present when enabled).""" 

125 

126 types: list[str] 

127 rows: int 

128 

129 

130class StatusResult(BaseModel): 

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

132 

133 command: str = "status" 

134 config: StatusConfig 

135 sources: list[SourceInfo] 

136 document_count: int 

137 total_chunks: int 

138 index: IndexStatus | None = None 

139 """The embedder that built the index; None before the first sync.""" 

140 entities: EntityStatus | None = None 

141 skipped: list[SkippedSource] = [] 

142 """Files a skip marker holds out of the index, capped at ``STATUS_SKIPPED_LIMIT``.""" 

143 skipped_total: int = 0 

144 

145 

146def _index_status() -> IndexStatus | None: 

147 """The persisted index identity, or None when nothing has been indexed.""" 

148 meta = get_services().store.get_meta() 

149 if meta is None: 

150 return None 

151 return IndexStatus(embedding_model=meta["embedding_model"], embedding_dim=meta["embedding_dim"]) 

152 

153 

154def gather_status() -> StatusResult: 

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

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

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

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

159 skipped, skipped_total = held_out_sources() 

160 return StatusResult( 

161 index=_index_status(), 

162 config=StatusConfig( 

163 documents_dir=str(cfg.documents_dir), 

164 data_dir=str(cfg.data_dir), 

165 chat_model=cfg.chat_model, 

166 embedding_model=cfg.embedding_model, 

167 vision_model=cfg.vision_model, 

168 reranker_model=cfg.reranker_model, 

169 enable_ocr=cfg.enable_ocr, 

170 num_ctx=cfg.num_ctx, 

171 num_ctx_max=cfg.num_ctx_max, 

172 chat_n_ctx_target=cfg.chat_n_ctx_target, 

173 flash_attention=cfg.flash_attention, 

174 kv_cache_type=cfg.kv_cache_type, 

175 n_gpu_layers=cfg.n_gpu_layers, 

176 cpu_moe=cfg.cpu_moe, 

177 n_cpu_moe=cfg.n_cpu_moe, 

178 main_gpu=cfg.main_gpu, 

179 gpu_devices=cfg.gpu_devices, 

180 ), 

181 sources=[ 

182 SourceInfo( 

183 filename=s["filename"], 

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

185 chunk_count=s["chunk_count"], 

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

187 ) 

188 for s in sorted_sources 

189 ], 

190 document_count=len(sources), 

191 total_chunks=total_chunks, 

192 entities=entity_status(), 

193 skipped=skipped, 

194 skipped_total=skipped_total, 

195 ) 

196 

197 

198def held_out_sources() -> tuple[list[SkippedSource], int]: 

199 """Held-out files with reasons, capped at ``STATUS_SKIPPED_LIMIT``, and the real count.""" 

200 held_out = sorted(load_skip_markers(cfg.data_root)) 

201 return describe_skips(cfg.data_root, held_out[:STATUS_SKIPPED_LIMIT]), len(held_out) 

202 

203 

204def entity_status() -> EntityStatus | None: 

205 """Entity types + extracted row count; None while the feature is off.""" 

206 if not cfg.entity_extraction: 

207 return None 

208 from lilbee.core.config import ENTITIES_TABLE 

209 from lilbee.retrieval.entities import load_schema 

210 

211 store = get_services().store 

212 schema = load_schema(store) 

213 table = store.open_table(ENTITIES_TABLE) 

214 rows = table.count_rows() if table is not None else 0 

215 return EntityStatus(types=[t.name for t in schema.types] if schema else [], rows=rows)