Coverage for src/lilbee/retrieval/concepts/graph.py: 100%
384 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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 ChunkType, ConceptRecords, SourceType, 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["source_id"], 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 sources = {row["chunk_source"] for row in records.chunk_concepts}
189 ids = self._store.source_ids(sources, SourceType.DOCUMENT)
190 chunk_concepts = [
191 {
192 "source_id": ids[row["chunk_source"]],
193 "chunk_index": row["chunk_index"],
194 "concept": row["concept"],
195 }
196 for row in records.chunk_concepts
197 if row["chunk_source"] in ids
198 ]
199 db = self._store.get_db() # opens the library, and so migrates it, outside the lock
200 with lock.write_lock(self._config.lancedb_dir):
201 # Always create tables so get_graph() returns True even when
202 # concept extraction yields no results for the current corpus.
203 nodes_tbl = data_store.ensure_table(db, CONCEPT_NODES_TABLE, _concept_nodes_schema())
204 edges_tbl = data_store.ensure_table(db, CONCEPT_EDGES_TABLE, _concept_edges_schema())
205 cc_tbl = data_store.ensure_table(db, CHUNK_CONCEPTS_TABLE, _chunk_concepts_schema())
206 if records.nodes:
207 nodes_tbl.add(records.nodes)
208 if records.edges:
209 edges_tbl.add(records.edges)
210 if chunk_concepts:
211 cc_tbl.add(chunk_concepts)
213 def boost_results(self, results: list[Any], query_concepts: list[str]) -> list[Any]:
214 """Boost search results whose chunks overlap with query concepts.
216 One batched chunk_concepts query serves the whole result set, grouped
217 back per chunk in Python, so the boost costs one query rather than one
218 per result. Concept rows are a document's, so a wiki chunk gets no boost.
219 """
220 if not query_concepts or not results:
221 return results
222 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
223 if table is None:
224 return results
225 query_set = set(query_concepts)
226 documents = {r.source for r in results if r.chunk_type != ChunkType.WIKI}
227 ids = self._store.source_ids(documents, SourceType.DOCUMENT)
228 keys = [
229 (ids.get(r.source, "") if r.chunk_type != ChunkType.WIKI else "", r.chunk_index)
230 for r in results
231 ]
232 concepts_by_chunk = self._chunk_concepts_batch(table, {key for key in keys if key[0]})
233 boosted: list[Any] = []
234 for r, key in zip(results, keys, strict=True):
235 chunk_concepts = concepts_by_chunk.get(key, set())
236 overlap = len(query_set & chunk_concepts)
237 if overlap > 0:
238 boost = (overlap / len(query_set)) * self._config.concept_boost_weight
239 r = r.model_copy()
240 if r.score is not None:
241 # Canonical [0, 1] space: the boost weight is directly
242 # comparable to an arm's fusion weight. (Added to a raw
243 # RRF score, whose whole range is ~0.017, the same 0.3
244 # default swamped hybrid ranking outright.)
245 r.score = min(1.0, r.score + boost)
246 boosted.append(r)
247 return boosted
249 def get_chunk_concepts(self, source: str, chunk_index: int) -> list[str]:
250 """Get concepts associated with a specific chunk."""
251 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
252 source_id = self._store.source_ids([source], None).get(source)
253 if table is None or source_id is None:
254 return []
255 try:
256 rows = (
257 table.search()
258 .where(f"source_id = '{source_id}' AND chunk_index = {int(chunk_index)}")
259 .to_list()
260 )
261 except Exception:
262 log.debug("get_chunk_concepts query failed for %r", source, exc_info=True)
263 return []
264 return [r["concept"] for r in rows]
266 @staticmethod
267 def _chunk_concepts_batch(
268 table: Any, chunks: set[tuple[str, int]]
269 ) -> dict[tuple[str, int], set[str]]:
270 """Fetch many chunks' concepts in one query, keyed by (source, index).
272 The predicate is the cross product of the distinct sources and
273 indexes -- a cheap superset -- and rows are filtered back to the
274 exact requested pairs in Python.
275 """
276 sources = ", ".join(f"'{escape_sql_string(s)}'" for s in sorted({s for s, _ in chunks}))
277 indexes = ", ".join(str(int(i)) for i in sorted({i for _, i in chunks}))
278 try:
279 rows = (
280 table.search()
281 .where(f"source_id IN ({sources}) AND chunk_index IN ({indexes})")
282 .to_list()
283 )
284 except Exception:
285 log.debug("chunk concepts batch query failed", exc_info=True)
286 return {}
287 concepts_by_chunk: dict[tuple[str, int], set[str]] = {}
288 for row in rows:
289 key = (row["source_id"], row["chunk_index"])
290 if key in chunks:
291 concepts_by_chunk.setdefault(key, set()).add(row["concept"])
292 return concepts_by_chunk
294 def expand_query(self, query: str) -> list[str]:
295 """Expand a query with related concepts from the graph."""
296 concepts = self.extract_concepts(query)
297 if not concepts:
298 return []
299 related: list[str] = []
300 seen = set(concepts)
301 for concept in concepts:
302 for neighbor in self.get_related_concepts(concept):
303 if neighbor not in seen:
304 related.append(neighbor)
305 seen.add(neighbor)
306 return related
308 def get_related_concepts(self, concept: str, depth: int = 1) -> list[str]:
309 """Find concepts related to *concept* via graph edges, up to *depth* hops.
311 One batched query per depth level: O(depth) DB round-trips,
312 independent of frontier size.
313 """
314 table = self._store.open_table(CONCEPT_EDGES_TABLE)
315 if table is None:
316 return []
317 visited: set[str] = {concept}
318 frontier: list[str] = [concept]
319 for _ in range(depth):
320 if not frontier:
321 break
322 escaped_list = ", ".join(f"'{escape_sql_string(n)}'" for n in frontier)
323 try:
324 rows = (
325 table.search()
326 .where(f"source IN ({escaped_list}) OR target IN ({escaped_list})")
327 .to_list()
328 )
329 except Exception:
330 log.debug(
331 "concept expand batch failed at frontier size %d",
332 len(frontier),
333 exc_info=True,
334 )
335 break
336 next_frontier: list[str] = []
337 for row in rows:
338 for endpoint in (row["source"], row["target"]):
339 if endpoint not in visited:
340 visited.add(endpoint)
341 next_frontier.append(endpoint)
342 frontier = next_frontier
343 return [c for c in visited if c != concept]
345 def top_communities(self, k: int = 10) -> list[Community]:
346 """Return the *k* largest concept communities.
348 Uses ``pyarrow.compute.value_counts`` to pick the top-k
349 cluster_ids in columnar memory, then materializes only those
350 clusters' members. Peak Python memory scales with members of
351 the top *k* clusters, not the total node count.
352 """
353 table = self._store.open_table(CONCEPT_NODES_TABLE)
354 if table is None:
355 return []
356 arrow_tbl = table.to_arrow()
357 if arrow_tbl.num_rows == 0:
358 return []
359 counts = pc.value_counts(arrow_tbl["cluster_id"]).to_pylist()
360 top = sorted(counts, key=lambda entry: entry["counts"], reverse=True)[:k]
361 top_ids = [entry["values"] for entry in top if entry["values"] is not None]
362 if not top_ids:
363 return []
364 member_rows = arrow_tbl.filter(
365 pc.is_in(arrow_tbl["cluster_id"], value_set=pa.array(top_ids))
366 ).to_pylist()
367 by_cluster: dict[int, list[str]] = {}
368 for row in member_rows:
369 by_cluster.setdefault(row["cluster_id"], []).append(row["concept"])
370 return [
371 Community(
372 cluster_id=cid,
373 size=len(by_cluster.get(cid, [])),
374 concepts=by_cluster.get(cid, []),
375 )
376 for cid in top_ids
377 if by_cluster.get(cid)
378 ]
380 def _corpus_pmi_inputs(
381 self,
382 ) -> _PmiInputs:
383 """Co-occurrence counts, concept document-frequencies, and chunk count,
384 all derived from the chunk_concepts table.
386 chunk_concepts is the ground-truth concept<->chunk map: it is source-scoped
387 (re-ingesting a source replaces its rows) and its schema is stable, so PMI
388 computed from it stays correct across re-ingests and version upgrades. The
389 edge table accrues per-file appends between rebuilds and those weights are
390 per-file co-occurrence counts, not corpus PMI, so it is not a safe source
391 for these corpus counts.
393 Concepts are de-duplicated per chunk, so a concept (or pair) counts once per
394 distinct chunk it appears in -- the document frequency PMI is defined on.
395 """
396 cooccurrences: Counter[tuple[str, str]] = Counter()
397 concept_counts: Counter[str] = Counter()
398 table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
399 if table is None:
400 return _PmiInputs(cooccurrences, concept_counts, 0, map_exists=False)
401 per_chunk: dict[tuple[str, int], set[str]] = {}
402 for rows in _iter_row_batches(table):
403 for row in rows:
404 key = (row["source_id"], row["chunk_index"])
405 per_chunk.setdefault(key, set()).add(row["concept"])
406 for concepts in per_chunk.values():
407 ordered = sorted(concepts)
408 for c in ordered:
409 concept_counts[c] += 1
410 for i, a in enumerate(ordered):
411 for b in ordered[i + 1 :]:
412 cooccurrences[(a, b)] += 1
413 return _PmiInputs(cooccurrences, concept_counts, len(per_chunk), map_exists=True)
415 def rebuild_clusters(self, added: Collection[str] = (), updated: Collection[str] = ()) -> None:
416 """Re-run Leiden clustering; scoped to changed sources when any are named.
418 Empty *added*/*updated* re-partitions the whole corpus. Otherwise only
419 the communities touching the changed sources' concepts are re-run and
420 merged back; untouched clusters keep their rows. Removals pass nothing:
421 their rows are already gone, so only a full pass clears their nodes.
422 Untouched edges keep their pre-partial weights and drift as the corpus
423 grows; the scoped partition can differ from a global pass on boundary
424 communities.
425 """
426 changed = set(added) | set(updated)
427 if changed and self._recluster_changed(changed, sweep_orphans=bool(updated)):
428 return
429 self._rebuild_all()
431 def _recluster_changed(self, changed: set[str], *, sweep_orphans: bool) -> bool:
432 """Re-partition the communities touching *changed*. False needs a full pass."""
433 nodes_table = self._store.open_table(CONCEPT_NODES_TABLE)
434 if nodes_table is None:
435 return False
436 node_rows = nodes_table.to_arrow().to_pylist()
437 if not node_rows:
438 return False
439 cc_table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
440 if cc_table is None:
441 return True
442 touched = self._touched_concepts(
443 cc_table, set(self._store.source_ids(changed, SourceType.DOCUMENT).values())
444 )
445 if touched and not self._recluster_affected(cc_table, node_rows, touched):
446 return False
447 if sweep_orphans:
448 self._sweep_orphan_concepts(cc_table, node_rows)
449 return True
451 @staticmethod
452 def _touched_concepts(cc_table: Any, changed: set[str]) -> set[str]:
453 """Distinct concepts the chunks of the sources *changed* (ids) currently carry."""
454 if not changed:
455 return set()
456 rows = cc_table.search().where(f"source_id IN ({_quoted(changed)})").to_list()
457 return {row["concept"] for row in rows}
459 def _recluster_affected(
460 self, cc_table: Any, node_rows: list[dict[str, Any]], touched: set[str]
461 ) -> bool:
462 """Re-partition the touched communities and merge them back. False needs a full pass."""
463 concept_to_cluster = {row["concept"]: row["cluster_id"] for row in node_rows}
464 affected_ids = {concept_to_cluster[c] for c in touched if c in concept_to_cluster}
465 affected = set(touched)
466 affected.update(c for c, cid in concept_to_cluster.items() if cid in affected_ids)
467 inputs = self._affected_pmi_inputs(cc_table, affected)
468 if inputs is None:
469 return False
470 counts, cooccurrences = inputs
471 pmi_weights = _compute_pmi(cooccurrences, counts, _distinct_chunk_count(cc_table))
472 node_records: list[dict[str, Any]] = []
473 edge_rows = [{"source": a, "target": b, "weight": w} for (a, b), w in pmi_weights.items()]
474 if pmi_weights:
475 partition, degree_map = _leiden_partition(edge_rows)
476 next_id = max(row["cluster_id"] for row in node_rows) + 1
477 node_records = [
478 {"concept": c, "cluster_id": next_id + cid, "degree": degree_map.get(c, 0)}
479 for c, cid in partition.items()
480 if c in affected
481 ]
482 nodes_predicate, edges_predicate = _affected_predicates(touched, affected_ids, affected)
483 self._store.clear_and_add(
484 CONCEPT_NODES_TABLE, _concept_nodes_schema(), node_records, nodes_predicate
485 )
486 self._store.clear_and_add(
487 CONCEPT_EDGES_TABLE, _concept_edges_schema(), edge_rows, edges_predicate
488 )
489 return True
491 def _affected_pmi_inputs(
492 self, cc_table: Any, affected: set[str]
493 ) -> tuple[Counter[str], Counter[tuple[str, str]]] | None:
494 """Corpus counts and co-occurrences for pairs touching *affected*.
496 A pair's chunks all hold its affected endpoint, so the neighborhood
497 carries every affected pair's co-occurrence count exactly; only the
498 outside neighbors' document frequencies need a second count. None
499 means the batch lookup failed and the caller must run a full pass.
500 """
501 chunk_rows = cc_table.search().where(f"concept IN ({_quoted(affected)})").to_list()
502 chunk_keys = {(row["source_id"], row["chunk_index"]) for row in chunk_rows}
503 concepts_by_chunk = self._chunk_concepts_batch(cc_table, chunk_keys)
504 if chunk_keys and not concepts_by_chunk:
505 # No rows for existing keys means the lookup failed; an empty
506 # map here would wipe the region, so fall back instead.
507 return None
508 counts: Counter[str] = Counter()
509 cooccurrences: Counter[tuple[str, str]] = Counter()
510 for concepts in concepts_by_chunk.values():
511 ordered = sorted(concepts)
512 for concept in ordered:
513 if concept in affected:
514 counts[concept] += 1
515 for i, first in enumerate(ordered):
516 for second in ordered[i + 1 :]:
517 if first in affected or second in affected:
518 cooccurrences[(first, second)] += 1
519 neighbors = {c for concepts in concepts_by_chunk.values() for c in concepts} - affected
520 if neighbors:
521 counts.update(self._neighbor_counts(cc_table, neighbors))
522 return counts, cooccurrences
524 @staticmethod
525 def _neighbor_counts(cc_table: Any, neighbors: set[str]) -> dict[str, int]:
526 """Distinct-chunk counts for concepts outside the affected set."""
527 count_rows = cc_table.search().where(f"concept IN ({_quoted(neighbors)})").to_list()
528 seen: dict[str, set[tuple[str, int]]] = {}
529 for row in count_rows:
530 seen.setdefault(row["concept"], set()).add((row["source_id"], row["chunk_index"]))
531 return {concept: len(keys) for concept, keys in seen.items()}
533 def _sweep_orphan_concepts(self, cc_table: Any, node_rows: list[dict[str, Any]]) -> None:
534 """Delete nodes and edges for concepts no chunk carries anymore."""
535 live = {row["concept"] for rows in _iter_row_batches(cc_table) for row in rows}
536 orphans = {row["concept"] for row in node_rows} - live
537 if orphans:
538 quoted = _quoted(orphans)
539 self._store.clear_and_add(
540 CONCEPT_NODES_TABLE, _concept_nodes_schema(), [], f"concept IN ({quoted})"
541 )
542 self._store.clear_and_add(
543 CONCEPT_EDGES_TABLE,
544 _concept_edges_schema(),
545 [],
546 f"source IN ({quoted}) OR target IN ({quoted})",
547 )
549 def _rebuild_all(self) -> None:
550 """Recompute corpus PMI from the chunk_concepts map, re-run Leiden, compact.
552 PMI is a corpus-level statistic, so it is computed once over corpus-wide
553 co-occurrence and concept counts (see :meth:`_corpus_pmi_inputs`) rather
554 than per file; summing per-file PMI would inflate pairs that recur across
555 many small files.
557 Both the nodes and the edges tables are replaced with the freshly
558 computed corpus graph. Per-file writes only ever append edges, so
559 without this rewrite the edges table grows monotonically across syncs
560 and expand_query keeps serving edges for concepts that left the corpus.
561 """
562 cooccurrences, concept_counts, total_chunks, map_exists = self._corpus_pmi_inputs()
563 if total_chunks == 0:
564 # The corpus was emptied: leaving the last graph in place would keep
565 # expansion serving concepts no document carries any more.
566 if map_exists:
567 self._clear_graph()
568 return
569 if not cooccurrences:
570 # Chunks remain but no concept pair co-occurs: the previous graph
571 # is stale, not still valid.
572 self._clear_graph()
573 return
574 pmi_weights = _compute_pmi(cooccurrences, concept_counts, total_chunks)
575 if not pmi_weights:
576 # Every pair co-occurred at or below chance: no edge set to
577 # cluster, and the previous graph no longer describes the corpus.
578 self._clear_graph()
579 return
580 edge_rows = [{"source": a, "target": b, "weight": w} for (a, b), w in pmi_weights.items()]
582 partition, degree_map = _leiden_partition(edge_rows)
584 node_records = [
585 {
586 "concept": node,
587 "cluster_id": cluster_id,
588 "degree": degree_map.get(node, 0),
589 }
590 for node, cluster_id in partition.items()
591 ]
593 # Delete the old rows and add the new ones under one lock per table so
594 # a reader never sees a table emptied while get_graph() still reports
595 # it present (which would blank top_communities / cluster labels).
596 self._store.clear_and_add(
597 CONCEPT_NODES_TABLE, _concept_nodes_schema(), node_records, "concept IS NOT NULL"
598 )
599 self._store.clear_and_add(
600 CONCEPT_EDGES_TABLE, _concept_edges_schema(), edge_rows, "source IS NOT NULL"
601 )
602 self.compact_tables()
604 def _clear_graph(self) -> None:
605 """Drop every node and edge, keeping both tables present for readers."""
606 self._store.clear_and_add(
607 CONCEPT_NODES_TABLE, _concept_nodes_schema(), [], "concept IS NOT NULL"
608 )
609 self._store.clear_and_add(
610 CONCEPT_EDGES_TABLE, _concept_edges_schema(), [], "source IS NOT NULL"
611 )
613 def compact_tables(self) -> None:
614 """Compact the concept tables; per-file adds otherwise accrete tiny versions."""
615 self._store.get_db() # opens the library, and so migrates it, outside the lock
616 with lock.write_lock(self._config.lancedb_dir):
617 for name in _CONCEPT_TABLES:
618 table = self._store.open_table(name)
619 if table is None:
620 continue
621 try:
622 table.optimize()
623 except Exception:
624 log.debug("Concept table optimize failed on '%s'", name, exc_info=True)
626 def get_cluster_sources(self, min_sources: int = 3) -> dict[int, set[str]]:
627 """Return clusters that span at least *min_sources* distinct sources.
628 Joins concept_nodes (concept -> cluster_id) with chunk_concepts
629 (concept -> source_id) to find which document sources each
630 cluster touches.
631 """
632 nodes_table = self._store.open_table(CONCEPT_NODES_TABLE)
633 cc_table = self._store.open_table(CHUNK_CONCEPTS_TABLE)
634 if nodes_table is None or cc_table is None:
635 return {}
637 concept_to_cluster: dict[str, int] = {}
638 for node_rows in _iter_row_batches(nodes_table):
639 for row in node_rows:
640 concept_to_cluster[row["concept"]] = row["cluster_id"]
642 cluster_sources: dict[int, set[str]] = {}
643 for cc_rows in _iter_row_batches(cc_table):
644 for row in cc_rows:
645 cid = concept_to_cluster.get(row["concept"])
646 if cid is None:
647 continue
648 cluster_sources.setdefault(cid, set()).add(row["source_id"])
650 wide = {cid: ids for cid, ids in cluster_sources.items() if len(ids) >= min_sources}
651 paths = self._store.source_paths({i for ids in wide.values() for i in ids})
652 return {cid: {paths[i] for i in ids if i in paths} for cid, ids in wide.items()}
654 def get_cluster_label(self, cluster_id: int) -> str:
655 """Return a human-readable label for *cluster_id* (highest-degree concept)."""
656 table = self._store.open_table(CONCEPT_NODES_TABLE)
657 if table is None:
658 return f"cluster-{cluster_id}"
659 try:
660 rows = table.search().where(f"cluster_id = {int(cluster_id)}").to_list()
661 except Exception:
662 log.debug("get_cluster_label query failed", exc_info=True)
663 return f"cluster-{cluster_id}"
664 if not rows:
665 return f"cluster-{cluster_id}"
666 best = max(rows, key=lambda r: r["degree"])
667 return str(best["concept"])
669 def get_graph(self) -> bool:
670 """Check whether a concept graph exists in the store."""
671 if not self._config.concept_graph:
672 return False
673 return self._store.open_table(CONCEPT_NODES_TABLE) is not None
675 def reset_nlp_cache(self) -> None:
676 """Clear the spaCy model cache. For testing only."""
677 self._nlp = None
678 self._nlp_unavailable = False