Coverage for src/lilbee/data/store/lance_helpers.py: 100%

1"""LanceDB plumbing helpers: table introspection, safe deletes, SQL escaping, error text.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6import threading 

7from pathlib import Path 

8from typing import TYPE_CHECKING 

9 

10from lilbee.catalog.refs import hf_repo_from_ref 

11from lilbee.runtime.lock import write_lock 

12 

13from .types import LOCAL_OWNER, ChunkType 

14 

15if TYPE_CHECKING: 

16 import pyarrow as pa 

17 from lancedb.db import LanceDBConnection 

18 from lancedb.index import IndexConfig 

19 from lancedb.table import LanceTable 

20 

21log = logging.getLogger(__name__) 

22 

23# The chunks table's body-text column: FTS-indexed and searched by the lexical 

24# arm, named here so the store's query and index code shares one spelling with 

25# its sibling _TITLE_COLUMN. 

26_CHUNK_COLUMN = "chunk" 

27 

28# Index types as LanceDB's IndexConfig reports them. 

29_FTS_INDEX_TYPE = "FTS" 

30_SCALAR_INDEX_TYPES = frozenset({"bitmap", "btree"}) 

31 

32 

33def install_lancedb_thread_error_suppressor() -> None: 

34 """Install a ``threading.excepthook`` that swallows lancedb shutdown noise. 

35 lancedb has no ``close()`` API and its internal event loop thread crashes 

36 during Python interpreter teardown. The exception is harmless (the process 

37 is exiting anyway) but pollutes CLI/TUI output. This is opt-in so importing 

38 ``lilbee.data.store`` has no hidden side effects; call it once from the CLI/TUI 

39 bootstrap. 

40 """ 

41 original = threading.excepthook 

42 

43 def _hook(args: threading.ExceptHookArgs) -> None: 

44 if args.thread and "LanceDB" in args.thread.name: 

45 return 

46 original(args) 

47 

48 threading.excepthook = _hook 

49 

50 

51def table_names(db: LanceDBConnection) -> list[str]: 

52 """Every table name in *db*, following ``list_tables`` pages to the end.""" 

53 names: list[str] = [] 

54 page_token: str | None = None 

55 while True: 

56 page = db.list_tables(page_token=page_token) 

57 names.extend(page.tables) 

58 page_token = page.page_token 

59 if not page_token: 

60 return names 

61 

62 

63def open_committed(db: LanceDBConnection, name: str) -> LanceTable | None: 

64 """Open table *name*, or ``None`` when it is absent or no process has committed it yet. 

65 

66 The listing names a table as soon as its folder exists, before its first 

67 version is written, so a table that another process is creating is listed 

68 but cannot be opened. It holds no rows yet. 

69 """ 

70 if name not in table_names(db): 

71 return None 

72 try: 

73 table: LanceTable = db.open_table(name) 

74 except ValueError: 

75 return None 

76 return table 

77 

78 

79def ensure_table(db: LanceDBConnection, name: str, schema: pa.Schema) -> LanceTable: 

80 table: LanceTable 

81 if name in table_names(db): 

82 table = db.open_table(name) 

83 return table 

84 try: 

85 table = db.create_table(name, schema=schema) 

86 except ValueError: 

87 table = db.open_table(name) 

88 return table 

89 

90 

91def _safe_delete_unlocked(table: LanceTable, predicate: str) -> bool: 

92 """Delete rows matching predicate. Caller must hold write lock. 

93 

94 Returns True when the delete succeeded, False when it raised (logged). The 

95 return lets delete-then-add callers avoid corrupting state on a swallowed 

96 failure (e.g. inserting a row whose stale predecessor was never removed). 

97 """ 

98 try: 

99 table.delete(predicate) 

100 return True 

101 except Exception: 

102 log.warning("Failed to delete rows matching: %s", predicate, exc_info=True) 

103 return False 

104 

105 

106def safe_delete(table: LanceTable, predicate: str, lancedb_dir: Path | None = None) -> bool: 

107 """Delete rows matching predicate, logging on failure. Returns success. 

108 

109 Pass the store's ``lancedb_dir`` so the write lock coordinates on that 

110 instance's data dir; ``None`` falls back to the global config dir. 

111 """ 

112 with write_lock(lancedb_dir): 

113 return _safe_delete_unlocked(table, predicate) 

114 

115 

116def escape_sql_string(value: str) -> str: 

117 """Escape a value for a single-quoted SQL string literal in a LanceDB predicate. 

118 

119 LanceDB's Datafusion engine follows standard SQL: the only escape inside a 

120 ``'...'`` literal is doubling the single quote. Backslash is an ordinary 

121 character, so escaping it (``\\`` -> ``\\\\``) corrupts the literal and makes a 

122 value containing a backslash (e.g. a Windows path) never match. 

123 """ 

124 return value.replace("'", "''") 

125 

126 

127def local_owner_predicate() -> str: 

128 """SQL predicate selecting the local human's own memories.""" 

129 return f"owner = '{LOCAL_OWNER}'" 

130 

131 

132def human_recall_predicate() -> str: 

133 """SQL predicate for the human: own memories plus any an agent has shared. 

134 

135 The mirror of :func:`agent_recall_predicate`: ``shared=True`` on an agent 

136 memory means "expose to the human's TUI/CLI", so the human's view must 

137 include those rather than only ``owner = 'local'``. 

138 """ 

139 return f"owner = '{LOCAL_OWNER}' OR (shared = true AND owner != '{LOCAL_OWNER}')" 

140 

141 

142def agent_recall_predicate(owner: str) -> str: 

143 """SQL predicate for an agent: its own memories plus the human's shared ones.""" 

144 return f"owner = '{escape_sql_string(owner)}' OR (shared = true AND owner = '{LOCAL_OWNER}')" 

145 

146 

147def _chunk_type_predicate(chunk_type: ChunkType | str) -> str: 

148 """SQL predicate that matches ``chunk_type`` while tolerating NULL rows. 

149 

150 ``'raw'`` means document content, so it matches extracted table rows as 

151 well as raw ones, and the NULL rows written before the column existed. 

152 Scoping a search to the user's documents must not silently drop their 

153 tables. A ``'wiki'`` filter matches only generated pages. 

154 """ 

155 escaped = escape_sql_string(chunk_type) 

156 if chunk_type == ChunkType.RAW: 

157 return f"(chunk_type IN ('{ChunkType.RAW}', '{ChunkType.TABLE}') OR chunk_type IS NULL)" 

158 return f"chunk_type = '{escaped}'" 

159 

160 

161def _listed_indices(table: LanceTable) -> list[IndexConfig] | None: 

162 """The indexes lancedb lists on *table*, or ``None`` when the listing raises.""" 

163 try: 

164 return list(table.list_indices()) 

165 except Exception: 

166 return None 

167 

168 

169def _index_registry(table: LanceTable, lancedb_dir: Path) -> list[IndexConfig] | None: 

170 """The indexes registered on *table*, or ``None`` when an FTS index is unusable. 

171 

172 ``None`` means the listing raised, or an FTS index is registered but its 

173 directory under ``_indices`` is gone. lancedb lists such an index, and 

174 every keyword query through it fails on a missing file. An empty list 

175 means the table has no index. :func:`_dangling_indices` checks scalar 

176 and vector indexes. 

177 """ 

178 indices = _listed_indices(table) 

179 if indices is None: 

180 return None 

181 indices_dir = _indices_dir(table, lancedb_dir) 

182 if any( 

183 idx.index_type == _FTS_INDEX_TYPE and _index_files_gone(indices_dir, idx) for idx in indices 

184 ): 

185 return None 

186 return indices 

187 

188 

189def _indices_dir(table: LanceTable, lancedb_dir: Path) -> Path: 

190 """The ``_indices`` directory of *table* under the store's *lancedb_dir*.""" 

191 return lancedb_dir / f"{table.name}.lance" / "_indices" 

192 

193 

194def _index_files_gone(indices_dir: Path, idx: IndexConfig) -> bool: 

195 """True when *idx* is registered but its directory under *indices_dir* is gone. 

196 

197 An index without a uuid has no directory to check. 

198 """ 

199 return idx.index_uuid is not None and not (indices_dir / idx.index_uuid).is_dir() 

200 

201 

202def _has_fts_index(indices: list[IndexConfig], column: str = _CHUNK_COLUMN) -> bool: 

203 """True when *indices* hold an FTS index on *column*.""" 

204 return any(idx.index_type == _FTS_INDEX_TYPE and column in idx.columns for idx in indices) 

205 

206 

207def _has_scalar_index(indices: list[IndexConfig], column: str) -> bool: 

208 """True when *indices* hold a scalar index on *column*. 

209 

210 lilbee builds only scalar indexes on the columns it prefilters by, never an 

211 FTS or vector index, so any index touching *column* is the scalar one. 

212 """ 

213 return any(column in idx.columns for idx in indices) 

214 

215 

216def _is_vector_index(idx: IndexConfig) -> bool: 

217 """True when *idx* is an ANN index on the vector column. 

218 

219 LanceDB reports IVF index types as ``IvfPq`` / ``IvfFlat`` etc., so the 

220 family match is case-insensitive. 

221 """ 

222 return "IVF" in idx.index_type.upper() and "vector" in idx.columns 

223 

224 

225def _has_vector_index(indices: list[IndexConfig]) -> bool: 

226 """True when *indices* hold an ANN index on the vector column.""" 

227 return any(_is_vector_index(idx) for idx in indices) 

228 

229 

230def _dangling_indices( 

231 table: LanceTable, indices: list[IndexConfig], lancedb_dir: Path 

232) -> list[IndexConfig]: 

233 """The indexes in *indices* whose directory under ``_indices`` is gone. 

234 

235 LanceDB keeps the registration after the index directory is removed, and 

236 every query through the index then fails on a missing file. The listing, 

237 ``optimize()`` and ``index_stats()`` do not notice for scalar and vector 

238 indexes; only the directory does. 

239 """ 

240 indices_dir = _indices_dir(table, lancedb_dir) 

241 return [idx for idx in indices if _index_files_gone(indices_dir, idx)] 

242 

243 

244def _scalar_index_dangling( 

245 table: LanceTable, indices: list[IndexConfig], lancedb_dir: Path 

246) -> list[str]: 

247 """Column names whose scalar (BITMAP/BTree) index is registered but its files are gone.""" 

248 columns = [ 

249 column 

250 for idx in _dangling_indices(table, indices, lancedb_dir) 

251 if idx.index_type.lower() in _SCALAR_INDEX_TYPES 

252 for column in idx.columns 

253 ] 

254 return list(dict.fromkeys(columns)) 

255 

256 

257def _vector_index_dangling( 

258 table: LanceTable, indices: list[IndexConfig], lancedb_dir: Path 

259) -> bool: 

260 """True when the ANN index on the vector column is registered but its files are gone.""" 

261 return any(_is_vector_index(idx) for idx in _dangling_indices(table, indices, lancedb_dir)) 

262 

263 

264def refs_compatible( 

265 persisted_ref: str, 

266 current_ref: str, 

267 persisted_dim: int, 

268 current_dim: int, 

269) -> bool: 

270 """Return True when *persisted_ref* and *current_ref* describe the same embedder. 

271 

272 Compatible iff dims match and either the raw refs are equal or the persisted 

273 ref is the legacy bare-repo form (``<org>/<repo>`` without a ``.gguf`` 

274 filename) whose repo matches the current canonical full ref. The legacy 

275 asymmetry exists because pre-canonical lilbee versions persisted only the 

276 repo; the current code persists the full ``<org>/<repo>/<filename>.gguf``. 

277 Two different ``.gguf`` files in the same repo are not lumped together 

278 (different quantizations can produce subtly different vectors), so both- 

279 full-ref strict identity is preserved. 

280 """ 

281 if persisted_dim != current_dim: 

282 return False 

283 if persisted_ref == current_ref: 

284 return True 

285 if persisted_ref.endswith(".gguf"): 

286 return False 

287 if not current_ref.endswith(".gguf"): 

288 return False 

289 return hf_repo_from_ref(current_ref) == persisted_ref