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 

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 

29def install_lancedb_thread_error_suppressor() -> None: 

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

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

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

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

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

35 bootstrap. 

36 """ 

37 original = threading.excepthook 

38 

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

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

41 return 

42 original(args) 

43 

44 threading.excepthook = _hook 

45 

46 

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

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

49 result = db.list_tables() 

50 try: 

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

52 except AttributeError: 

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

54 

55 

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

57 if name in table_names(db): 

58 return db.open_table(name) 

59 try: 

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

61 except ValueError: 

62 return db.open_table(name) 

63 

64 

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

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

67 

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

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

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

71 """ 

72 try: 

73 table.delete(predicate) 

74 return True 

75 except Exception: 

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

77 return False 

78 

79 

80def safe_delete( 

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

82) -> bool: 

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

84 

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

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

87 """ 

88 with write_lock(lancedb_dir): 

89 return _safe_delete_unlocked(table, predicate) 

90 

91 

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

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

94 

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

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

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

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

99 """ 

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

101 

102 

103def local_owner_predicate() -> str: 

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

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

106 

107 

108def human_recall_predicate() -> str: 

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

110 

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

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

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

114 """ 

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

116 

117 

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

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

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

121 

122 

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

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

125 

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

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

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

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

130 """ 

131 escaped = escape_sql_string(chunk_type) 

132 if chunk_type == ChunkType.RAW: 

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

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

135 

136 

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

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

139 try: 

140 for idx in table.list_indices(): 

141 if idx.index_type == "FTS" and column in idx.columns: 

142 return True 

143 except Exception: 

144 return False 

145 return False 

146 

147 

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

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

150 

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

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

153 """ 

154 try: 

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

156 except Exception: 

157 return False 

158 

159 

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

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

162 

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

164 family match is case-insensitive. 

165 """ 

166 try: 

167 for idx in table.list_indices(): 

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

169 return True 

170 except Exception: 

171 return False 

172 return False 

173 

174 

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

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

177 

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

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

180 because it is the ESCAPE character the predicate declares. 

181 """ 

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

183 

184 

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

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

187 

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

189 (pre-title stores lack it). 

190 """ 

191 if not search: 

192 return None 

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

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

195 if include_title: 

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

197 return clause 

198 

199 

200def refs_compatible( 

201 persisted_ref: str, 

202 current_ref: str, 

203 persisted_dim: int, 

204 current_dim: int, 

205) -> bool: 

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

207 

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

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

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

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

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

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

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

215 full-ref strict identity is preserved. 

216 """ 

217 if persisted_dim != current_dim: 

218 return False 

219 if persisted_ref == current_ref: 

220 return True 

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

222 return False 

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

224 return False 

225 return hf_repo_from_ref(current_ref) == persisted_ref