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 typing import TYPE_CHECKING 

8 

9from lilbee.catalog.refs import hf_repo_from_ref 

10from lilbee.runtime.lock import write_lock 

11 

12from .types import LOCAL_OWNER, ChunkType 

13 

14if TYPE_CHECKING: 

15 from pathlib import Path 

16 

17 import lancedb 

18 import lancedb.table 

19 import pyarrow as pa 

20 from lancedb.index import IndexConfig 

21 

22log = logging.getLogger(__name__) 

23 

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

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

26# its sibling _TITLE_COLUMN. 

27_CHUNK_COLUMN = "chunk" 

28 

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

30_FTS_INDEX_TYPE = "FTS" 

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

32 

33 

34def install_lancedb_thread_error_suppressor() -> None: 

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

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

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

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

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

40 bootstrap. 

41 """ 

42 original = threading.excepthook 

43 

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

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

46 return 

47 original(args) 

48 

49 threading.excepthook = _hook 

50 

51 

52def table_names(db: lancedb.DBConnection) -> list[str]: 

53 """Get list of table names, handling the ListTablesResponse object.""" 

54 result = db.list_tables() 

55 try: 

56 return result.tables # type: ignore[no-any-return, union-attr] 

57 except AttributeError: 

58 return list(result) # type: ignore[arg-type] 

59 

60 

61def ensure_table(db: lancedb.DBConnection, name: str, schema: pa.Schema) -> lancedb.table.Table: 

62 if name in table_names(db): 

63 return db.open_table(name) 

64 try: 

65 return db.create_table(name, schema=schema) 

66 except ValueError: 

67 return db.open_table(name) 

68 

69 

70def _safe_delete_unlocked(table: lancedb.table.Table, predicate: str) -> bool: 

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

72 

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

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

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

76 """ 

77 try: 

78 table.delete(predicate) 

79 return True 

80 except Exception: 

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

82 return False 

83 

84 

85def safe_delete( 

86 table: lancedb.table.Table, predicate: str, lancedb_dir: Path | None = None 

87) -> bool: 

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

89 

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

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

92 """ 

93 with write_lock(lancedb_dir): 

94 return _safe_delete_unlocked(table, predicate) 

95 

96 

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

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

99 

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

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

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

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

104 """ 

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

106 

107 

108def local_owner_predicate() -> str: 

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

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

111 

112 

113def human_recall_predicate() -> str: 

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

115 

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

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

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

119 """ 

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

121 

122 

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

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

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

126 

127 

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

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

130 

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

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

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

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

135 """ 

136 escaped = escape_sql_string(chunk_type) 

137 if chunk_type == ChunkType.RAW: 

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

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

140 

141 

142def _has_fts_index(table: lancedb.table.Table, column: str = _CHUNK_COLUMN) -> bool: 

143 """Return True when an FTS index on *column* already exists.""" 

144 try: 

145 for idx in table.list_indices(): 

146 if idx.index_type == _FTS_INDEX_TYPE and column in idx.columns: 

147 return True 

148 except Exception: 

149 return False 

150 return False 

151 

152 

153def _dangling_indices(table: lancedb.table.Table, lancedb_dir: Path) -> list[IndexConfig]: 

154 """Registered indexes whose directory under ``_indices`` is gone. 

155 

156 LanceDB keeps the registration in the manifest after the index directory 

157 is removed, and every query through the index then fails on a missing 

158 file. ``optimize()`` and ``index_stats()`` do not notice; only the 

159 directory does. Empty when the manifest cannot be read. 

160 """ 

161 indices_dir = lancedb_dir / f"{table.name}.lance" / "_indices" 

162 try: 

163 registered = table.list_indices() 

164 except Exception: 

165 return [] 

166 return [idx for idx in registered if not (indices_dir / idx.index_uuid).is_dir()] 

167 

168 

169def _fts_index_dangling( 

170 table: lancedb.table.Table, lancedb_dir: Path, column: str = _CHUNK_COLUMN 

171) -> bool: 

172 """True when an FTS index on *column* is registered but its files are gone.""" 

173 return any( 

174 idx.index_type == _FTS_INDEX_TYPE and column in idx.columns 

175 for idx in _dangling_indices(table, lancedb_dir) 

176 ) 

177 

178 

179def _scalar_index_dangling(table: lancedb.table.Table, lancedb_dir: Path) -> list[str]: 

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

181 columns = [ 

182 column 

183 for idx in _dangling_indices(table, lancedb_dir) 

184 if idx.index_type.lower() in _SCALAR_INDEX_TYPES 

185 for column in idx.columns 

186 ] 

187 return list(dict.fromkeys(columns)) 

188 

189 

190def _has_scalar_index(table: lancedb.table.Table, column: str) -> bool: 

191 """Return True when a scalar index on *column* already exists. 

192 

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

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

195 """ 

196 try: 

197 return any(column in idx.columns for idx in table.list_indices()) 

198 except Exception: 

199 return False 

200 

201 

202def _has_vector_index(table: lancedb.table.Table) -> bool: 

203 """Return True when an ANN index on the vector column already exists. 

204 

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

206 family match is case-insensitive. 

207 """ 

208 try: 

209 for idx in table.list_indices(): 

210 if "IVF" in idx.index_type.upper() and "vector" in idx.columns: 

211 return True 

212 except Exception: 

213 return False 

214 return False 

215 

216 

217def _escape_like_wildcards(value: str) -> str: 

218 """Escape LIKE metacharacters so a search term matches literally. 

219 

220 ``%`` and ``_`` are wildcards inside a LIKE pattern; without escaping, a 

221 search for ``a_b`` would also match ``axb``. Backslash is escaped first 

222 because it is the ESCAPE character the predicate declares. 

223 """ 

224 return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") 

225 

226 

227def _sources_search_filter(search: str | None, *, include_title: bool = False) -> str | None: 

228 """Case-insensitive filename (and optionally title) WHERE clause, or ``None``. 

229 

230 *include_title* requires the caller to have checked the column exists 

231 (pre-title stores lack it). 

232 """ 

233 if not search: 

234 return None 

235 escaped = escape_sql_string(_escape_like_wildcards(search.lower())) 

236 clause = f"LOWER(filename) LIKE '%{escaped}%' ESCAPE '\\'" 

237 if include_title: 

238 clause = f"({clause} OR LOWER(title) LIKE '%{escaped}%' ESCAPE '\\')" 

239 return clause 

240 

241 

242def refs_compatible( 

243 persisted_ref: str, 

244 current_ref: str, 

245 persisted_dim: int, 

246 current_dim: int, 

247) -> bool: 

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

249 

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

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

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

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

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

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

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

257 full-ref strict identity is preserved. 

258 """ 

259 if persisted_dim != current_dim: 

260 return False 

261 if persisted_ref == current_ref: 

262 return True 

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

264 return False 

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

266 return False 

267 return hf_repo_from_ref(current_ref) == persisted_ref