Coverage for src/lilbee/retrieval/concepts/graph.py: 100%
373 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +0000
1"""ConceptGraph: extracts, stores, and queries concept relationships."""
3from __future__ import annotations
5import logging
6import threading
7from collections import Counter
8from collections.abc import Collection, Iterator
9from typing import TYPE_CHECKING, Any, NamedTuple
11import pyarrow as pa
12import pyarrow.compute as pc
14if TYPE_CHECKING:
15 from lancedb.table import LanceTable
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: LanceTable) -> 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()
49def _quoted(values: Collection[str]) -> str:
50 """SQL string list for an IN predicate, in stable order."""
51 return ", ".join(f"'{escape_sql_string(value)}'" for value in sorted(values))
54def _distinct_chunk_count(cc_table: Any) -> int:
55 """Chunks carrying at least one concept."""
56 keys: set[tuple[str, int]] = set()
57 for rows in _iter_row_batches(cc_table):
58 for row in rows:
59 keys.add((row["chunk_source"], row["chunk_index"]))
60 return len(keys)
63def _affected_predicates(
64 touched: set[str], affected_ids: set[int], affected: set[str]
65) -> tuple[str, str]:
66 """Scoped delete predicates covering the touched communities."""
67 parts = [f"concept IN ({_quoted(touched)})"]
68 if affected_ids:
69 parts.append(f"cluster_id IN ({', '.join(str(cid) for cid in sorted(affected_ids))})")
70 edge_list = _quoted(affected)
71 return " OR ".join(parts), f"source IN ({edge_list}) OR target IN ({edge_list})"
74class _PmiInputs(NamedTuple):
75 """Corpus-wide counts PMI is computed from (see _corpus_pmi_inputs)."""
77 cooccurrences: Counter[tuple[str, str]]
78 concept_counts: Counter[str]
79 total_chunks: int
80 # A missing map means concepts were never built; an existing but empty one
81 # means the corpus was emptied. Both give total_chunks == 0, and only the
82 # second should clear the graph.
83 map_exists: bool
86class ConceptGraph:
87 """Concept graph -- extracts, stores, and queries concept relationships."""
89 def __init__(self, config: Config, store: Store) -> None:
90 self._config = config
91 self._store = store
92 self._nlp: Any = None
93 self._nlp_unavailable: bool = False
94 # A spaCy Language is not safe for concurrent processing (shared Vocab /
95 # StringStore). ConceptGraph is a Services singleton, so serialize every
96 # nlp() / nlp.pipe() call on the shared daemon behind this lock.
97 self._nlp_lock = threading.Lock()
98 # Single-entry memo: one search() extracts concepts for the same query
99 # twice (expansion + boost), so cache the last (text, max) -> result to
100 # spare the second spaCy pass without unbounded growth.
101 self._last_extract: tuple[tuple[str, int], list[str]] | None = None
103 def _ensure_nlp(self) -> Any | None:
104 """Lazy-load and cache the spaCy model. Returns None if unavailable."""
105 if self._nlp is None and not self._nlp_unavailable:
106 # Double-checked under _nlp_lock so two concurrent first-callers don't
107 # each load en_core_web_sm (the loser would just be discarded).
108 with self._nlp_lock:
109 if self._nlp is None and not self._nlp_unavailable:
110 try:
111 self._nlp = _ensure_spacy_model()
112 except ImportError:
113 log.warning("Concept graph disabled: spaCy model unavailable")
114 self._nlp_unavailable = True
115 return self._nlp
117 def extract_concepts(self, text: str, max_concepts: int | None = None) -> list[str]:
118 """Extract noun-phrase concepts from text via spaCy."""
119 if max_concepts is None:
120 max_concepts = self._config.concept_max_per_chunk
121 if not text.strip():
122 return []
123 nlp = self._ensure_nlp()
124 if nlp is None:
125 return []
126 cache_key = (text, max_concepts)
127 with self._nlp_lock:
128 if self._last_extract is not None and self._last_extract[0] == cache_key:
129 return self._last_extract[1]
130 doc = nlp(text)
131 result = _filter_noun_chunks(doc, max_concepts)
132 self._last_extract = (cache_key, result)
133 return result
135 def extract_concepts_batch(self, texts: list[str]) -> list[list[str]]:
136 """Batch-extract concepts from multiple texts."""
137 if not texts:
138 return []
139 nlp = self._ensure_nlp()
140 if nlp is None:
141 return [[] for _ in texts]
142 max_concepts = self._config.concept_max_per_chunk
143 # Hold the lock across the full pipe iteration: nlp.pipe is lazy, so the
144 # actual parsing happens as the comprehension consumes it.
145 with self._nlp_lock:
146 return [_filter_noun_chunks(doc, max_concepts) for doc in nlp.pipe(texts)]
148 def build_concept_records(
149 self, chunk_ids: list[tuple[str, int]], concept_lists: list[list[str]]
150 ) -> ConceptRecords:
151 """Build co-occurrence graph rows from chunk concepts; no store access.
153 Edge weights are raw co-occurrence counts (the edges carry graph
154 connectivity), not PMI. Corpus PMI for clustering is recomputed from the
155 chunk_concepts map in :meth:`rebuild_clusters`: computing PMI per file (with
156 one file's chunk count as the denominator) and summing the per-file weights
157 inflates pairs that recur across many small files, which is not corpus PMI.
158 """
159 cooccurrences: Counter[tuple[str, str]] = Counter()
160 concept_counts: Counter[str] = Counter()
161 chunk_concept_records: list[dict[str, Any]] = []
163 for (source, idx), concepts in zip(chunk_ids, concept_lists, strict=True):
164 for c in concepts:
165 concept_counts[c] += 1
166 chunk_concept_records.append(
167 {"chunk_source": source, "chunk_index": idx, "concept": c}
168 )
169 for i, a in enumerate(concepts):
170 for b in concepts[i + 1 :]:
171 pair = (min(a, b), max(a, b))
172 cooccurrences[pair] += 1
174 return ConceptRecords(
175 nodes=[
176 {"concept": c, "cluster_id": 0, "degree": count}
177 for c, count in concept_counts.items()
178 ],
179 edges=[
180 {"source": a, "target": b, "weight": float(count)}
181 for (a, b), count in cooccurrences.items()
182 ],
183 chunk_concepts=chunk_concept_records,
184 )
186 def write_concept_records(self, records: ConceptRecords) -> None:
187 """Write batched concept rows: one lock acquisition, at most one add per table."""
188 with lock.write_lock(self._config.lancedb_dir):
189 db = self._store.get_db()
190 # Always create tables so get_graph() returns True even when
191 # concept extraction yields no results for the current corpus.
192 nodes_tbl = data_store.ensure_table(db, CONCEPT_NODES_TABLE, _concept_nodes_schema())
193 edges_tbl = data_store.ensure_table(db, CONCEPT_EDGES_TABLE, _concept_edges_schema())
194 cc_tbl = data_store.ensure_table(db, CHUNK_CONCEPTS_TABLE, _chunk_concepts_schema())
195 if records.nodes:
196 nodes_tbl.add(records.nodes)
197 if records.edges:
198 edges_tbl.add(records.edges)
199 if records.chunk_concepts:
200 cc_tbl.add(records.chunk_concepts)
202 def boost_results(self, results: list[Any], query_concepts: list[str]) -> list[Any]:
203 """Boost search results whose chunks overlap with query concepts.
205 One batched chunk_concepts query serves the whole result set, grouped
206 back per chunk in Python, so the boost costs one query rather than one
207 per result.
208 """
209 if not query_concepts or not results:
210 return results
211 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
212 if table is None:
213 return results
214 query_set = set(query_concepts)
215 concepts_by_chunk = self._chunk_concepts_batch(
216 table, {(r.source, r.chunk_index) for r in results}
217 )
218 boosted: list[Any] = []
219 for r in results:
220 chunk_concepts = concepts_by_chunk.get((r.source, r.chunk_index), set())
221 overlap = len(query_set & chunk_concepts)
222 if overlap > 0:
223 boost = (overlap / len(query_set)) * self._config.concept_boost_weight
224 r = r.model_copy()
225 if r.score is not None:
226 # Canonical [0, 1] space: the boost weight is directly
227 # comparable to an arm's fusion weight. (Added to a raw
228 # RRF score, whose whole range is ~0.017, the same 0.3
229 # default swamped hybrid ranking outright.)
230 r.score = min(1.0, r.score + boost)
231 boosted.append(r)
232 return boosted
234 def get_chunk_concepts(self, source: str, chunk_index: int) -> list[str]:
235 """Get concepts associated with a specific chunk."""
236 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
237 if table is None:
238 return []
239 escaped = escape_sql_string(source)
240 try:
241 rows = (
242 table.search()
243 .where(f"chunk_source = '{escaped}' AND chunk_index = {int(chunk_index)}")
244 .to_list()
245 )
246 except Exception:
247 log.debug("get_chunk_concepts query failed for %r", source, exc_info=True)
248 return []
249 return [r["concept"] for r in rows]
251 @staticmethod
252 def _chunk_concepts_batch(
253 table: Any, chunks: set[tuple[str, int]]
254 ) -> dict[tuple[str, int], set[str]]:
255 """Fetch many chunks' concepts in one query, keyed by (source, index).
257 The predicate is the cross product of the distinct sources and
258 indexes -- a cheap superset -- and rows are filtered back to the
259 exact requested pairs in Python.
260 """
261 sources = ", ".join(f"'{escape_sql_string(s)}'" for s in sorted({s for s, _ in chunks}))
262 indexes = ", ".join(str(int(i)) for i in sorted({i for _, i in chunks}))
263 try:
264 rows = (
265 table.search()
266 .where(f"chunk_source IN ({sources}) AND chunk_index IN ({indexes})")
267 .to_list()
268 )
269 except Exception:
270 log.debug("chunk concepts batch query failed", exc_info=True)
271 return {}
272 concepts_by_chunk: dict[tuple[str, int], set[str]] = {}
273 for row in rows:
274 key = (row["chunk_source"], row["chunk_index"])
275 if key in chunks:
276 concepts_by_chunk.setdefault(key, set()).add(row["concept"])
277 return concepts_by_chunk
279 def expand_query(self, query: str) -> list[str]:
280 """Expand a query with related concepts from the graph."""
281 concepts = self.extract_concepts(query)
282 if not concepts:
283 return []
284 related: list[str] = []
285 seen = set(concepts)
286 for concept in concepts:
287 for neighbor in self.get_related_concepts(concept):
288 if neighbor not in seen:
289 related.append(neighbor)
290 seen.add(neighbor)
291 return related
293 def get_related_concepts(self, concept: str, depth: int = 1) -> list[str]:
294 """Find concepts related to *concept* via graph edges, up to *depth* hops.
296 One batched query per depth level: O(depth) DB round-trips,
297 independent of frontier size.
298 """
299 table = self._store.open_table(CONCEPT_EDGES_TABLE)
300 if table is None:
301 return []
302 visited: set[str] = {concept}
303 frontier: list[str] = [concept]
304 for _ in range(depth):
305 if not frontier:
306 break
307 escaped_list = ", ".join(f"'{escape_sql_string(n)}'" for n in frontier)
308 try:
309 rows = (
310 table.search()
311 .where(f"source IN ({escaped_list}) OR target IN ({escaped_list})")
312 .to_list()
313 )
314 except Exception:
315 log.debug(
316 "concept expand batch failed at frontier size %d",
317 len(frontier),
318 exc_info=True,
319 )
320 break
321 next_frontier: list[str] = []
322 for row in rows:
323 for endpoint in (row["source"], row["target"]):
324 if endpoint not in visited:
325 visited.add(endpoint)
326 next_frontier.append(endpoint)
327 frontier = next_frontier
328 return [c for c in visited if c != concept]
330 def top_communities(self, k: int = 10) -> list[Community]:
331 """Return the *k* largest concept communities.
333 Uses ``pyarrow.compute.value_counts`` to pick the top-k
334 cluster_ids in columnar memory, then materializes only those
335 clusters' members. Peak Python memory scales with members of
336 the top *k* clusters, not the total node count.
337 """
338 table = self._store.open_table(CONCEPT_NODES_TABLE)
339 if table is None:
340 return []
341 arrow_tbl = table.to_arrow()
342 if arrow_tbl.num_rows == 0:
343 return []
344 counts = pc.value_counts(arrow_tbl["cluster_id"]).to_pylist()
345 top = sorted(counts, key=lambda entry: entry["counts"], reverse=True)[:k]
346 top_ids = [entry["values"] for entry in top if entry["values"] is not None]
347 if not top_ids:
348 return []
349 member_rows = arrow_tbl.filter(
350 pc.is_in(arrow_tbl["cluster_id"], value_set=pa.array(top_ids))
351 ).to_pylist()
352 by_cluster: dict[int, list[str]] = {}
353 for row in member_rows:
354 by_cluster.setdefault(row["cluster_id"], []).append(row["concept"])
355 return [
356 Community(
357 cluster_id=cid,
358 size=len(by_cluster.get(cid, [])),
359 concepts=by_cluster.get(cid, []),
360 )
361 for cid in top_ids
362 if by_cluster.get(cid)
363 ]
365 def _corpus_pmi_inputs(
366 self,
367 ) -> _PmiInputs:
368 """Co-occurrence counts, concept document-frequencies, and chunk count,
369 all derived from the chunk_concepts table.
371 chunk_concepts is the ground-truth concept<->chunk map: it is source-scoped
372 (re-ingesting a source replaces its rows) and its schema is stable, so PMI
373 computed from it stays correct across re-ingests and version upgrades. The
374 edge table accrues per-file appends between rebuilds and those weights are
375 per-file co-occurrence counts, not corpus PMI, so it is not a safe source
376 for these corpus counts.
378 Concepts are de-duplicated per chunk, so a concept (or pair) counts once per
379 distinct chunk it appears in -- the document frequency PMI is defined on.
380 """
381 cooccurrences: Counter[tuple[str, str]] = Counter()
382 concept_counts: Counter[str] = Counter()
383 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
384 if table is None:
385 return _PmiInputs(cooccurrences, concept_counts, 0, map_exists=False)
386 per_chunk: dict[tuple[str, int], set[str]] = {}
387 for rows in _iter_row_batches(table):
388 for row in rows:
389 key = (row["chunk_source"], row["chunk_index"])
390 per_chunk.setdefault(key, set()).add(row["concept"])
391 for concepts in per_chunk.values():
392 ordered = sorted(concepts)
393 for c in ordered:
394 concept_counts[c] += 1
395 for i, a in enumerate(ordered):
396 for b in ordered[i + 1 :]:
397 cooccurrences[(a, b)] += 1
398 return _PmiInputs(cooccurrences, concept_counts, len(per_chunk), map_exists=True)
400 def rebuild_clusters(self, added: Collection[str] = (), updated: Collection[str] = ()) -> None:
401 """Re-run Leiden clustering; scoped to changed sources when any are named.
403 Empty *added*/*updated* re-partitions the whole corpus. Otherwise only
404 the communities touching the changed sources' concepts are re-run and
405 merged back; untouched clusters keep their rows. Removals pass nothing:
406 their rows are already gone, so only a full pass clears their nodes.
407 Untouched edges keep their pre-partial weights and drift as the corpus
408 grows; the scoped partition can differ from a global pass on boundary
409 communities.
410 """
411 changed = set(added) | set(updated)
412 if changed and self._recluster_changed(changed, sweep_orphans=bool(updated)):
413 return
414 self._rebuild_all()
416 def _recluster_changed(self, changed: set[str], *, sweep_orphans: bool) -> bool:
417 """Re-partition the communities touching *changed*. False needs a full pass."""
418 nodes_table = self._store.open_table(CONCEPT_NODES_TABLE)
419 if nodes_table is None:
420 return False
421 node_rows = nodes_table.to_arrow().to_pylist()
422 if not node_rows:
423 return False
424 cc_table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
425 if cc_table is None:
426 return True
427 touched = self._touched_concepts(cc_table, changed)
428 if touched and not self._recluster_affected(cc_table, node_rows, touched):
429 return False
430 if sweep_orphans:
431 self._sweep_orphan_concepts(cc_table, node_rows)
432 return True
434 @staticmethod
435 def _touched_concepts(cc_table: Any, changed: set[str]) -> set[str]:
436 """Distinct concepts the changed sources' chunks currently carry."""
437 rows = cc_table.search().where(f"chunk_source IN ({_quoted(changed)})").to_list()
438 return {row["concept"] for row in rows}
440 def _recluster_affected(
441 self, cc_table: Any, node_rows: list[dict[str, Any]], touched: set[str]
442 ) -> bool:
443 """Re-partition the touched communities and merge them back. False needs a full pass."""
444 concept_to_cluster = {row["concept"]: row["cluster_id"] for row in node_rows}
445 affected_ids = {concept_to_cluster[c] for c in touched if c in concept_to_cluster}
446 affected = set(touched)
447 affected.update(c for c, cid in concept_to_cluster.items() if cid in affected_ids)
448 inputs = self._affected_pmi_inputs(cc_table, affected)
449 if inputs is None:
450 return False
451 counts, cooccurrences = inputs
452 pmi_weights = _compute_pmi(cooccurrences, counts, _distinct_chunk_count(cc_table))
453 node_records: list[dict[str, Any]] = []
454 edge_rows = [{"source": a, "target": b, "weight": w} for (a, b), w in pmi_weights.items()]
455 if pmi_weights:
456 partition, degree_map = _leiden_partition(edge_rows)
457 next_id = max(row["cluster_id"] for row in node_rows) + 1
458 node_records = [
459 {"concept": c, "cluster_id": next_id + cid, "degree": degree_map.get(c, 0)}
460 for c, cid in partition.items()
461 if c in affected
462 ]
463 nodes_predicate, edges_predicate = _affected_predicates(touched, affected_ids, affected)
464 self._store.clear_and_add(
465 CONCEPT_NODES_TABLE, _concept_nodes_schema(), node_records, nodes_predicate
466 )
467 self._store.clear_and_add(
468 CONCEPT_EDGES_TABLE, _concept_edges_schema(), edge_rows, edges_predicate
469 )
470 return True
472 def _affected_pmi_inputs(
473 self, cc_table: Any, affected: set[str]
474 ) -> tuple[Counter[str], Counter[tuple[str, str]]] | None:
475 """Corpus counts and co-occurrences for pairs touching *affected*.
477 A pair's chunks all hold its affected endpoint, so the neighborhood
478 carries every affected pair's co-occurrence count exactly; only the
479 outside neighbors' document frequencies need a second count. None
480 means the batch lookup failed and the caller must run a full pass.
481 """
482 chunk_rows = cc_table.search().where(f"concept IN ({_quoted(affected)})").to_list()
483 chunk_keys = {(row["chunk_source"], row["chunk_index"]) for row in chunk_rows}
484 concepts_by_chunk = self._chunk_concepts_batch(cc_table, chunk_keys)
485 if chunk_keys and not concepts_by_chunk:
486 # No rows for existing keys means the lookup failed; an empty
487 # map here would wipe the region, so fall back instead.
488 return None
489 counts: Counter[str] = Counter()
490 cooccurrences: Counter[tuple[str, str]] = Counter()
491 for concepts in concepts_by_chunk.values():
492 ordered = sorted(concepts)
493 for concept in ordered:
494 if concept in affected:
495 counts[concept] += 1
496 for i, first in enumerate(ordered):
497 for second in ordered[i + 1 :]:
498 if first in affected or second in affected:
499 cooccurrences[(first, second)] += 1
500 neighbors = {c for concepts in concepts_by_chunk.values() for c in concepts} - affected
501 if neighbors:
502 counts.update(self._neighbor_counts(cc_table, neighbors))
503 return counts, cooccurrences
505 @staticmethod
506 def _neighbor_counts(cc_table: Any, neighbors: set[str]) -> dict[str, int]:
507 """Distinct-chunk counts for concepts outside the affected set."""
508 count_rows = cc_table.search().where(f"concept IN ({_quoted(neighbors)})").to_list()
509 seen: dict[str, set[tuple[str, int]]] = {}
510 for row in count_rows:
511 seen.setdefault(row["concept"], set()).add((row["chunk_source"], row["chunk_index"]))
512 return {concept: len(keys) for concept, keys in seen.items()}
514 def _sweep_orphan_concepts(self, cc_table: Any, node_rows: list[dict[str, Any]]) -> None:
515 """Delete nodes and edges for concepts no chunk carries anymore."""
516 live = {row["concept"] for rows in _iter_row_batches(cc_table) for row in rows}
517 orphans = {row["concept"] for row in node_rows} - live
518 if orphans:
519 quoted = _quoted(orphans)
520 self._store.clear_and_add(
521 CONCEPT_NODES_TABLE, _concept_nodes_schema(), [], f"concept IN ({quoted})"
522 )
523 self._store.clear_and_add(
524 CONCEPT_EDGES_TABLE,
525 _concept_edges_schema(),
526 [],
527 f"source IN ({quoted}) OR target IN ({quoted})",
528 )
530 def _rebuild_all(self) -> None:
531 """Recompute corpus PMI from the chunk_concepts map, re-run Leiden, compact.
533 PMI is a corpus-level statistic, so it is computed once over corpus-wide
534 co-occurrence and concept counts (see :meth:`_corpus_pmi_inputs`) rather
535 than per file; summing per-file PMI would inflate pairs that recur across
536 many small files.
538 Both the nodes and the edges tables are replaced with the freshly
539 computed corpus graph. Per-file writes only ever append edges, so
540 without this rewrite the edges table grows monotonically across syncs
541 and expand_query keeps serving edges for concepts that left the corpus.
542 """
543 cooccurrences, concept_counts, total_chunks, map_exists = self._corpus_pmi_inputs()
544 if total_chunks == 0:
545 # The corpus was emptied: leaving the last graph in place would keep
546 # expansion serving concepts no document carries any more.
547 if map_exists:
548 self._clear_graph()
549 return
550 if not cooccurrences:
551 # Chunks remain but no concept pair co-occurs: the previous graph
552 # is stale, not still valid.
553 self._clear_graph()
554 return
555 pmi_weights = _compute_pmi(cooccurrences, concept_counts, total_chunks)
556 if not pmi_weights:
557 # Every pair co-occurred at or below chance: no edge set to
558 # cluster, and the previous graph no longer describes the corpus.
559 self._clear_graph()
560 return
561 edge_rows = [{"source": a, "target": b, "weight": w} for (a, b), w in pmi_weights.items()]
563 partition, degree_map = _leiden_partition(edge_rows)
565 node_records = [
566 {
567 "concept": node,
568 "cluster_id": cluster_id,
569 "degree": degree_map.get(node, 0),
570 }
571 for node, cluster_id in partition.items()
572 ]
574 # Delete the old rows and add the new ones under one lock per table so
575 # a reader never sees a table emptied while get_graph() still reports
576 # it present (which would blank top_communities / cluster labels).
577 self._store.clear_and_add(
578 CONCEPT_NODES_TABLE, _concept_nodes_schema(), node_records, "concept IS NOT NULL"
579 )
580 self._store.clear_and_add(
581 CONCEPT_EDGES_TABLE, _concept_edges_schema(), edge_rows, "source IS NOT NULL"
582 )
583 self.compact_tables()
585 def _clear_graph(self) -> None:
586 """Drop every node and edge, keeping both tables present for readers."""
587 self._store.clear_and_add(
588 CONCEPT_NODES_TABLE, _concept_nodes_schema(), [], "concept IS NOT NULL"
589 )
590 self._store.clear_and_add(
591 CONCEPT_EDGES_TABLE, _concept_edges_schema(), [], "source IS NOT NULL"
592 )
594 def compact_tables(self) -> None:
595 """Compact the concept tables; per-file adds otherwise accrete tiny versions."""
596 with lock.write_lock(self._config.lancedb_dir):
597 for name in _CONCEPT_TABLES:
598 table = self._store.open_table(name)
599 if table is None:
600 continue
601 try:
602 table.optimize()
603 except Exception:
604 log.debug("Concept table optimize failed on '%s'", name, exc_info=True)
606 def get_cluster_sources(self, min_sources: int = 3) -> dict[int, set[str]]:
607 """Return clusters that span at least *min_sources* distinct sources.
608 Joins concept_nodes (concept -> cluster_id) with chunk_concepts
609 (concept -> chunk_source) to find which document sources each
610 cluster touches.
611 """
612 nodes_table = self._store.open_table(CONCEPT_NODES_TABLE)
613 cc_table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
614 if nodes_table is None or cc_table is None:
615 return {}
617 concept_to_cluster: dict[str, int] = {}
618 for node_rows in _iter_row_batches(nodes_table):
619 for row in node_rows:
620 concept_to_cluster[row["concept"]] = row["cluster_id"]
622 cluster_sources: dict[int, set[str]] = {}
623 for cc_rows in _iter_row_batches(cc_table):
624 for row in cc_rows:
625 cid = concept_to_cluster.get(row["concept"])
626 if cid is None:
627 continue
628 cluster_sources.setdefault(cid, set()).add(row["chunk_source"])
630 return {
631 cid: sources for cid, sources in cluster_sources.items() if len(sources) >= min_sources
632 }
634 def get_cluster_label(self, cluster_id: int) -> str:
635 """Return a human-readable label for *cluster_id* (highest-degree concept)."""
636 table = self._store.open_table(CONCEPT_NODES_TABLE)
637 if table is None:
638 return f"cluster-{cluster_id}"
639 try:
640 rows = table.search().where(f"cluster_id = {int(cluster_id)}").to_list()
641 except Exception:
642 log.debug("get_cluster_label query failed", exc_info=True)
643 return f"cluster-{cluster_id}"
644 if not rows:
645 return f"cluster-{cluster_id}"
646 best = max(rows, key=lambda r: r["degree"])
647 return str(best["concept"])
649 def get_graph(self) -> bool:
650 """Check whether a concept graph exists in the store."""
651 if not self._config.concept_graph:
652 return False
653 return self._store.open_table(CONCEPT_NODES_TABLE) is not None
655 def reset_nlp_cache(self) -> None:
656 """Clear the spaCy model cache. For testing only."""
657 self._nlp = None
658 self._nlp_unavailable = False