Coverage for src/lilbee/retrieval/concepts/graph.py: 100%

1"""ConceptGraph: extracts, stores, and queries concept relationships."""

3from __future__ import annotations

5import logging

6import threading

7from collections import Counter

8from collections.abc import Iterator

9from typing import TYPE_CHECKING, Any, NamedTuple

11import pyarrow as pa

12import pyarrow.compute as pc

14if TYPE_CHECKING:

15 import lancedb.table

17from lilbee.core.config import (

18 CHUNK_CONCEPTS_TABLE,

19 CONCEPT_EDGES_TABLE,

20 CONCEPT_NODES_TABLE,

21 Config,

22)

23from lilbee.data import store as data_store

24from lilbee.data.store import ConceptRecords, Store, escape_sql_string

25from lilbee.retrieval.concepts.community import Community, _compute_pmi, _leiden_partition

26from lilbee.retrieval.concepts.nlp import _ensure_spacy_model, _filter_noun_chunks

27from lilbee.retrieval.concepts.schema import (

28 _chunk_concepts_schema,

29 _concept_edges_schema,

30 _concept_nodes_schema,

31)

32from lilbee.runtime import lock

34log = logging.getLogger(__name__)

36# Rows per record batch when scanning a concept table; bounds the Python-dict

37# working set while the columnar Arrow data stays compact.

38_TABLE_SCAN_BATCH_ROWS = 50_000

40_CONCEPT_TABLES = (CONCEPT_NODES_TABLE, CONCEPT_EDGES_TABLE, CHUNK_CONCEPTS_TABLE)

43def _iter_row_batches(table: lancedb.table.Table) -> Iterator[list[dict[str, Any]]]:

44 """Yield a table's rows as bounded-size lists of dicts."""

45 for batch in table.to_arrow().to_batches(max_chunksize=_TABLE_SCAN_BATCH_ROWS):

46 yield batch.to_pylist()

49class _PmiInputs(NamedTuple):

50 """Corpus-wide counts PMI is computed from (see _corpus_pmi_inputs)."""

52 cooccurrences: Counter[tuple[str, str]]

53 concept_counts: Counter[str]

54 total_chunks: int

55 # A missing map means concepts were never built; an existing but empty one

56 # means the corpus was emptied. Both give total_chunks == 0, and only the

57 # second should clear the graph.

58 map_exists: bool

61class ConceptGraph:

62 """Concept graph -- extracts, stores, and queries concept relationships."""

64 def __init__(self, config: Config, store: Store) -> None:

65 self._config = config

66 self._store = store

67 self._nlp: Any = None

68 self._nlp_unavailable: bool = False

69 # A spaCy Language is not safe for concurrent processing (shared Vocab /

70 # StringStore). ConceptGraph is a Services singleton, so serialize every

71 # nlp() / nlp.pipe() call on the shared daemon behind this lock.

72 self._nlp_lock = threading.Lock()

73 # Single-entry memo: one search() extracts concepts for the same query

74 # twice (expansion + boost), so cache the last (text, max) -> result to

75 # spare the second spaCy pass without unbounded growth.

76 self._last_extract: tuple[tuple[str, int], list[str]] | None = None

78 def _ensure_nlp(self) -> Any | None:

79 """Lazy-load and cache the spaCy model. Returns None if unavailable."""

80 if self._nlp is None and not self._nlp_unavailable:

81 # Double-checked under _nlp_lock so two concurrent first-callers don't

82 # each load en_core_web_sm (the loser would just be discarded).

83 with self._nlp_lock:

84 if self._nlp is None and not self._nlp_unavailable:

85 try:

86 self._nlp = _ensure_spacy_model()

87 except ImportError:

88 log.warning("Concept graph disabled: spaCy model unavailable")

89 self._nlp_unavailable = True

90 return self._nlp

92 def extract_concepts(self, text: str, max_concepts: int | None = None) -> list[str]:

93 """Extract noun-phrase concepts from text via spaCy."""

94 if max_concepts is None:

95 max_concepts = self._config.concept_max_per_chunk

96 if not text.strip():

97 return []

98 nlp = self._ensure_nlp()

99 if nlp is None:

100 return []

101 cache_key = (text, max_concepts)

102 with self._nlp_lock:

103 if self._last_extract is not None and self._last_extract[0] == cache_key:

104 return self._last_extract[1]

105 doc = nlp(text)

106 result = _filter_noun_chunks(doc, max_concepts)

107 self._last_extract = (cache_key, result)

108 return result

109

110 def extract_concepts_batch(self, texts: list[str]) -> list[list[str]]:

111 """Batch-extract concepts from multiple texts."""

112 if not texts:

113 return []

114 nlp = self._ensure_nlp()

115 if nlp is None:

116 return [[] for _ in texts]

117 max_concepts = self._config.concept_max_per_chunk

118 # Hold the lock across the full pipe iteration: nlp.pipe is lazy, so the

119 # actual parsing happens as the comprehension consumes it.

120 with self._nlp_lock:

121 return [_filter_noun_chunks(doc, max_concepts) for doc in nlp.pipe(texts)]

122

123 def build_concept_records(

124 self, chunk_ids: list[tuple[str, int]], concept_lists: list[list[str]]

125 ) -> ConceptRecords:

126 """Build co-occurrence graph rows from chunk concepts; no store access.

127

128 Edge weights are raw co-occurrence counts (the edges carry graph

129 connectivity), not PMI. Corpus PMI for clustering is recomputed from the

130 chunk_concepts map in :meth:`rebuild_clusters`: computing PMI per file (with

131 one file's chunk count as the denominator) and summing the per-file weights

132 inflates pairs that recur across many small files, which is not corpus PMI.

133 """

134 cooccurrences: Counter[tuple[str, str]] = Counter()

135 concept_counts: Counter[str] = Counter()

136 chunk_concept_records: list[dict[str, Any]] = []

137

138 for (source, idx), concepts in zip(chunk_ids, concept_lists, strict=True):

139 for c in concepts:

140 concept_counts[c] += 1

141 chunk_concept_records.append(

142 {"chunk_source": source, "chunk_index": idx, "concept": c}

143 )

144 for i, a in enumerate(concepts):

145 for b in concepts[i + 1 :]:

146 pair = (min(a, b), max(a, b))

147 cooccurrences[pair] += 1

148

149 return ConceptRecords(

150 nodes=[

151 {"concept": c, "cluster_id": 0, "degree": count}

152 for c, count in concept_counts.items()

153 ],

154 edges=[

155 {"source": a, "target": b, "weight": float(count)}

156 for (a, b), count in cooccurrences.items()

157 ],

158 chunk_concepts=chunk_concept_records,

159 )

160

161 def write_concept_records(self, records: ConceptRecords) -> None:

162 """Write batched concept rows: one lock acquisition, at most one add per table."""

163 with lock.write_lock(self._config.lancedb_dir):

164 db = self._store.get_db()

165 # Always create tables so get_graph() returns True even when

166 # concept extraction yields no results for the current corpus.

167 nodes_tbl = data_store.ensure_table(db, CONCEPT_NODES_TABLE, _concept_nodes_schema())

168 edges_tbl = data_store.ensure_table(db, CONCEPT_EDGES_TABLE, _concept_edges_schema())

169 cc_tbl = data_store.ensure_table(db, CHUNK_CONCEPTS_TABLE, _chunk_concepts_schema())

170 if records.nodes:

171 nodes_tbl.add(records.nodes)

172 if records.edges:

173 edges_tbl.add(records.edges)

174 if records.chunk_concepts:

175 cc_tbl.add(records.chunk_concepts)

176

177 def boost_results(self, results: list[Any], query_concepts: list[str]) -> list[Any]:

178 """Boost search results whose chunks overlap with query concepts.

179

180 One batched chunk_concepts query serves the whole result set, grouped

181 back per chunk in Python, so the boost costs one query rather than one

182 per result.

183 """

184 if not query_concepts or not results:

185 return results

186 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)

187 if table is None:

188 return results

189 query_set = set(query_concepts)

190 concepts_by_chunk = self._chunk_concepts_batch(

191 table, {(r.source, r.chunk_index) for r in results}

192 )

193 boosted: list[Any] = []

194 for r in results:

195 chunk_concepts = concepts_by_chunk.get((r.source, r.chunk_index), set())

196 overlap = len(query_set & chunk_concepts)

197 if overlap > 0:

198 boost = (overlap / len(query_set)) * self._config.concept_boost_weight

199 r = r.model_copy()

200 if r.score is not None:

201 # Canonical [0, 1] space: the boost weight is directly

202 # comparable to an arm's fusion weight. (Added to a raw

203 # RRF score, whose whole range is ~0.017, the same 0.3

204 # default swamped hybrid ranking outright.)

205 r.score = min(1.0, r.score + boost)

206 boosted.append(r)

207 return boosted

208

209 def get_chunk_concepts(self, source: str, chunk_index: int) -> list[str]:

210 """Get concepts associated with a specific chunk."""

211 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)

212 if table is None:

213 return []

214 escaped = escape_sql_string(source)

215 try:

216 rows = (

217 table.search()

218 .where(f"chunk_source = '{escaped}' AND chunk_index = {int(chunk_index)}")

219 .to_list()

220 )

221 except Exception:

222 log.debug("get_chunk_concepts query failed for %r", source, exc_info=True)

223 return []

224 return [r["concept"] for r in rows]

225

226 @staticmethod

227 def _chunk_concepts_batch(

228 table: Any, chunks: set[tuple[str, int]]

229 ) -> dict[tuple[str, int], set[str]]:

230 """Fetch many chunks' concepts in one query, keyed by (source, index).

231

232 The predicate is the cross product of the distinct sources and

233 indexes -- a cheap superset -- and rows are filtered back to the

234 exact requested pairs in Python.

235 """

236 sources = ", ".join(f"'{escape_sql_string(s)}'" for s in sorted({s for s, _ in chunks}))

237 indexes = ", ".join(str(int(i)) for i in sorted({i for _, i in chunks}))

238 try:

239 rows = (

240 table.search()

241 .where(f"chunk_source IN ({sources}) AND chunk_index IN ({indexes})")

242 .to_list()

243 )

244 except Exception:

245 log.debug("chunk concepts batch query failed", exc_info=True)

246 return {}

247 concepts_by_chunk: dict[tuple[str, int], set[str]] = {}

248 for row in rows:

249 key = (row["chunk_source"], row["chunk_index"])

250 if key in chunks:

251 concepts_by_chunk.setdefault(key, set()).add(row["concept"])

252 return concepts_by_chunk

253

254 def expand_query(self, query: str) -> list[str]:

255 """Expand a query with related concepts from the graph."""

256 concepts = self.extract_concepts(query)

257 if not concepts:

258 return []

259 related: list[str] = []

260 seen = set(concepts)

261 for concept in concepts:

262 for neighbor in self.get_related_concepts(concept):

263 if neighbor not in seen:

264 related.append(neighbor)

265 seen.add(neighbor)

266 return related

267

268 def get_related_concepts(self, concept: str, depth: int = 1) -> list[str]:

269 """Find concepts related to *concept* via graph edges, up to *depth* hops.

270

271 One batched query per depth level: O(depth) DB round-trips,

272 independent of frontier size.

273 """

274 table = self._store.open_table(CONCEPT_EDGES_TABLE)

275 if table is None:

276 return []

277 visited: set[str] = {concept}

278 frontier: list[str] = [concept]

279 for _ in range(depth):

280 if not frontier:

281 break

282 escaped_list = ", ".join(f"'{escape_sql_string(n)}'" for n in frontier)

283 try:

284 rows = (

285 table.search()

286 .where(f"source IN ({escaped_list}) OR target IN ({escaped_list})")

287 .to_list()

288 )

289 except Exception:

290 log.debug(

291 "concept expand batch failed at frontier size %d",

292 len(frontier),

293 exc_info=True,

294 )

295 break

296 next_frontier: list[str] = []

297 for row in rows:

298 for endpoint in (row["source"], row["target"]):

299 if endpoint not in visited:

300 visited.add(endpoint)

301 next_frontier.append(endpoint)

302 frontier = next_frontier

303 return [c for c in visited if c != concept]

304

305 def top_communities(self, k: int = 10) -> list[Community]:

306 """Return the *k* largest concept communities.

307

308 Uses ``pyarrow.compute.value_counts`` to pick the top-k

309 cluster_ids in columnar memory, then materializes only those

310 clusters' members. Peak Python memory scales with members of

311 the top *k* clusters, not the total node count.

312 """

313 table = self._store.open_table(CONCEPT_NODES_TABLE)

314 if table is None:

315 return []

316 arrow_tbl = table.to_arrow()

317 if arrow_tbl.num_rows == 0:

318 return []

319 counts = pc.value_counts(arrow_tbl["cluster_id"]).to_pylist()

320 top = sorted(counts, key=lambda entry: entry["counts"], reverse=True)[:k]

321 top_ids = [entry["values"] for entry in top if entry["values"] is not None]

322 if not top_ids:

323 return []

324 member_rows = arrow_tbl.filter(

325 pc.is_in(arrow_tbl["cluster_id"], value_set=pa.array(top_ids))

326 ).to_pylist()

327 by_cluster: dict[int, list[str]] = {}

328 for row in member_rows:

329 by_cluster.setdefault(row["cluster_id"], []).append(row["concept"])

330 return [

331 Community(

332 cluster_id=cid,

333 size=len(by_cluster.get(cid, [])),

334 concepts=by_cluster.get(cid, []),

335 )

336 for cid in top_ids

337 if by_cluster.get(cid)

338 ]

339

340 def _corpus_pmi_inputs(

341 self,

342 ) -> _PmiInputs:

343 """Co-occurrence counts, concept document-frequencies, and chunk count,

344 all derived from the chunk_concepts table.

345

346 chunk_concepts is the ground-truth concept<->chunk map: it is source-scoped

347 (re-ingesting a source replaces its rows) and its schema is stable, so PMI

348 computed from it stays correct across re-ingests and version upgrades. The

349 edge table accrues per-file appends between rebuilds and those weights are

350 per-file co-occurrence counts, not corpus PMI, so it is not a safe source

351 for these corpus counts.

352

353 Concepts are de-duplicated per chunk, so a concept (or pair) counts once per

354 distinct chunk it appears in -- the document frequency PMI is defined on.

355 """

356 cooccurrences: Counter[tuple[str, str]] = Counter()

357 concept_counts: Counter[str] = Counter()

358 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)

359 if table is None:

360 return _PmiInputs(cooccurrences, concept_counts, 0, map_exists=False)

361 per_chunk: dict[tuple[str, int], set[str]] = {}

362 for rows in _iter_row_batches(table):

363 for row in rows:

364 key = (row["chunk_source"], row["chunk_index"])

365 per_chunk.setdefault(key, set()).add(row["concept"])

366 for concepts in per_chunk.values():

367 ordered = sorted(concepts)

368 for c in ordered:

369 concept_counts[c] += 1

370 for i, a in enumerate(ordered):

371 for b in ordered[i + 1 :]:

372 cooccurrences[(a, b)] += 1

373 return _PmiInputs(cooccurrences, concept_counts, len(per_chunk), map_exists=True)

374

375 def rebuild_clusters(self) -> None:

376 """Recompute corpus PMI from the chunk_concepts map, re-run Leiden, compact.

377

378 PMI is a corpus-level statistic, so it is computed once over corpus-wide

379 co-occurrence and concept counts (see :meth:`_corpus_pmi_inputs`) rather

380 than per file; summing per-file PMI would inflate pairs that recur across

381 many small files.

382

383 Both the nodes and the edges tables are replaced with the freshly

384 computed corpus graph. Per-file writes only ever append edges, so

385 without this rewrite the edges table grows monotonically across syncs

386 and expand_query keeps serving edges for concepts that left the corpus.

387 """

388 cooccurrences, concept_counts, total_chunks, map_exists = self._corpus_pmi_inputs()

389 if total_chunks == 0:

390 # The corpus was emptied: leaving the last graph in place would keep

391 # expansion serving concepts no document carries any more.

392 if map_exists:

393 self._clear_graph()

394 return

395 if not cooccurrences:

396 # Chunks remain but no concept pair co-occurs: the previous graph

397 # is stale, not still valid.

398 self._clear_graph()

399 return

400 pmi_weights = _compute_pmi(cooccurrences, concept_counts, total_chunks)

401 if not pmi_weights:

402 # Every pair co-occurred at or below chance: no edge set to

403 # cluster, and the previous graph no longer describes the corpus.

404 self._clear_graph()

405 return

406 edge_rows = [{"source": a, "target": b, "weight": w} for (a, b), w in pmi_weights.items()]

407

408 partition, degree_map = _leiden_partition(edge_rows)

409

410 node_records = [

411 {

412 "concept": node,

413 "cluster_id": cluster_id,

414 "degree": degree_map.get(node, 0),

415 }

416 for node, cluster_id in partition.items()

417 ]

418

419 # Delete the old rows and add the new ones under one lock per table so

420 # a reader never sees a table emptied while get_graph() still reports

421 # it present (which would blank top_communities / cluster labels).

422 self._store.clear_and_add(

423 CONCEPT_NODES_TABLE, _concept_nodes_schema(), node_records, "concept IS NOT NULL"

424 )

425 self._store.clear_and_add(

426 CONCEPT_EDGES_TABLE, _concept_edges_schema(), edge_rows, "source IS NOT NULL"

427 )

428 self.compact_tables()

429

430 def _clear_graph(self) -> None:

431 """Drop every node and edge, keeping both tables present for readers."""

432 self._store.clear_and_add(

433 CONCEPT_NODES_TABLE, _concept_nodes_schema(), [], "concept IS NOT NULL"

434 )

435 self._store.clear_and_add(

436 CONCEPT_EDGES_TABLE, _concept_edges_schema(), [], "source IS NOT NULL"

437 )

438

439 def compact_tables(self) -> None:

440 """Compact the concept tables; per-file adds otherwise accrete tiny versions."""

441 with lock.write_lock(self._config.lancedb_dir):

442 for name in _CONCEPT_TABLES:

443 table = self._store.open_table(name)

444 if table is None:

445 continue

446 try:

447 table.optimize()

448 except Exception:

449 log.debug("Concept table optimize failed on '%s'", name, exc_info=True)

450

451 def get_cluster_sources(self, min_sources: int = 3) -> dict[int, set[str]]:

452 """Return clusters that span at least *min_sources* distinct sources.

453 Joins concept_nodes (concept -> cluster_id) with chunk_concepts

454 (concept -> chunk_source) to find which document sources each

455 cluster touches.

456 """

457 nodes_table = self._store.open_table(CONCEPT_NODES_TABLE)

458 cc_table = self._store.open_table(CHUNK_CONCEPTS_TABLE)

459 if nodes_table is None or cc_table is None:

460 return {}

461

462 concept_to_cluster: dict[str, int] = {}

463 for node_rows in _iter_row_batches(nodes_table):

464 for row in node_rows:

465 concept_to_cluster[row["concept"]] = row["cluster_id"]

466

467 cluster_sources: dict[int, set[str]] = {}

468 for cc_rows in _iter_row_batches(cc_table):

469 for row in cc_rows:

470 cid = concept_to_cluster.get(row["concept"])

471 if cid is None:

472 continue

473 cluster_sources.setdefault(cid, set()).add(row["chunk_source"])

474

475 return {

476 cid: sources for cid, sources in cluster_sources.items() if len(sources) >= min_sources

477 }

478

479 def get_cluster_label(self, cluster_id: int) -> str:

480 """Return a human-readable label for *cluster_id* (highest-degree concept)."""

481 table = self._store.open_table(CONCEPT_NODES_TABLE)

482 if table is None:

483 return f"cluster-{cluster_id}"

484 try:

485 rows = table.search().where(f"cluster_id = {int(cluster_id)}").to_list()

486 except Exception:

487 log.debug("get_cluster_label query failed", exc_info=True)

488 return f"cluster-{cluster_id}"

489 if not rows:

490 return f"cluster-{cluster_id}"

491 best = max(rows, key=lambda r: r["degree"])

492 return str(best["concept"])

493

494 def get_graph(self) -> bool:

495 """Check whether a concept graph exists in the store."""

496 if not self._config.concept_graph:

497 return False

498 return self._store.open_table(CONCEPT_NODES_TABLE) is not None

499

500 def reset_nlp_cache(self) -> None:

501 """Clear the spaCy model cache. For testing only."""

502 self._nlp = None

503 self._nlp_unavailable = False