From 56d1dad442f1201a9597fb7e0296aa62c47940d1 Mon Sep 17 00:00:00 2001 From: himanshupatro-334 Date: Thu, 17 Sep 2026 01:06:26 +0530 Subject: [PATCH] Make community labeling size-aware --- graphify/llm.py | 225 ++++++++++++++++++++++++++++----- tests/test_labeling.py | 278 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 469 insertions(+), 34 deletions(-) diff --git a/graphify/llm.py b/graphify/llm.py index 781d30dd43..0d7efce67a 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -3228,36 +3228,145 @@ def _claude_cli_available() -> bool: _LABEL_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*|\s*```\s*$", re.IGNORECASE) _LABEL_MAX_COMMUNITIES = 200 # legacy soft-cap; kept for callers that pin it. -_LABEL_TOP_K = 12 # node labels sampled per community for the prompt +_LABEL_TOP_K = 12 # legacy fixed sample cap; kept for callers that pin it +_LABEL_MAX_ADAPTIVE_TOP_K = 32 # upper bound for adaptive community sampling (#3586) _LABEL_MAXLEN = 60 # truncate individual labels to keep the prompt small -_LABEL_BATCH_SIZE = 100 # communities per LLM call; sized for ~16k context windows +_LABEL_BATCH_SIZE = 100 # max communities per LLM call; bounded by _LABEL_MAX_PROMPT_TOKENS +_LABEL_MAX_PROMPT_TOKENS = 8000 # conservative input ceiling leaving room for completion + margin +_LABEL_PROMPT_PREAMBLE = ( + "You are naming clusters in a knowledge graph. For each community below, " + "return a concise 2-5 word plain-language name describing what it is about " + "(e.g. \"Order Management\", \"Payment Flow\", \"Auth Middleware\"). " + "Each input line is ': '. " + "Respond ONLY with a JSON object mapping the community id (as a string) to " + "its name - no prose, no markdown fences.\n\n" +) def _placeholder_community_labels(communities) -> dict[int, str]: return {int(cid): f"Community {cid}" for cid in communities} -def _community_label_lines(G, communities, gods, max_communities, top_k): - """One prompt line per community (largest first), sampling up to ``top_k`` - representative node labels (god nodes first). Returns (lines, labeled_cids); - skips communities with no resolvable nodes.""" +def _adaptive_sample_size(n: int) -> int: + """Size-aware representative sample count for community labeling (#3586). + + A bounded heuristic scaling sample size with community membership count: + small communities are sampled completely (<=12), while large communities + receive up to 32 representative nodes. Kept as a conservative step function + rather than a continuous or mathematically optimal curve to ensure + deterministic and predictable prompt budgets. + """ + if n <= 12: + return n + if n <= 50: + return 15 + if n <= 200: + return 20 + if n <= 1000: + return 26 + return 32 + + +def _community_label_lines( + G, + communities, + gods=None, + max_communities: int | None = None, + top_k: int | None = None, +) -> tuple[list[str], list[int]]: + """One prompt line per community (largest first), sampling representative + node labels. Returns (lines, labeled_cids); skips communities with no + resolvable nodes. + + When ``top_k`` is None, sample size is chosen adaptively via + ``_adaptive_sample_size(len(members))`` (#3586). When ``top_k`` is an int, + it acts as an explicit fixed cap for backwards compatibility. + + Representatives are ranked deterministically: + 1. Global god nodes in the community + 2. Internal degree within the community (edges to other community members) + 3. Global degree in the graph + 4. Node ID string (tie-breaker) + + A per-source-file diversity cap prevents any single file from consuming + the entire representative budget, followed by a backfill pass to ensure + the target quota is met even for single-file communities. + """ # gods may be node-id strings or god_nodes() dicts ({"id": ..., "label": ...}). god_set = {g["id"] if isinstance(g, dict) else g for g in (gods or [])} ordered = sorted(communities.items(), key=lambda kv: -len(kv[1])) lines: list[str] = [] labeled_cids: list[int] = [] - for cid, members in ordered[:max_communities]: - ranked = [m for m in members if m in god_set] + [m for m in members if m not in god_set] + + for cid, members in (ordered if max_communities is None else ordered[:max_communities]): + if not members: + continue + + target_k = _adaptive_sample_size(len(members)) if top_k is None else min(len(members), max(0, top_k)) + if target_k <= 0: + continue + + mset = set(members) + internal_deg: dict = {} + for n in members: + if n in G: + if hasattr(G, "is_directed") and G.is_directed(): + internal_deg[n] = ( + sum(1 for nb in G.successors(n) if nb in mset) + + sum(1 for nb in G.predecessors(n) if nb in mset) + ) + else: + internal_deg[n] = sum(1 for nb in G[n] if nb in mset) + else: + internal_deg[n] = 0 + + ranked = sorted( + members, + key=lambda n: ( + 0 if n in god_set else 1, + -internal_deg[n], + -(G.degree(n) if n in G else 0), + str(n), + ), + ) + + file_cap = max(2, (target_k + 2) // 3) + file_counts: dict[str, int] = {} + selected_nodes: set = set() names: list[str] = [] seen: set[str] = set() + + # Pass 1: select candidates respecting per-source-file cap for nid in ranked: + if nid in selected_nodes: + continue + sf = (G.nodes[nid].get("source_file") or "") if nid in G.nodes else "" + if file_counts.get(sf, 0) >= file_cap: + continue label = str(G.nodes[nid].get("label", nid)) if nid in G.nodes else str(nid) label = label.strip().strip("()")[:_LABEL_MAXLEN] if label and label.lower() not in seen: seen.add(label.lower()) names.append(label) - if len(names) >= top_k: + selected_nodes.add(nid) + file_counts[sf] = file_counts.get(sf, 0) + 1 + if len(names) >= target_k: break + + # Pass 2: backfill if diversity cap left vacancies (e.g. single-file community) + if len(names) < target_k: + for nid in ranked: + if nid in selected_nodes: + continue + label = str(G.nodes[nid].get("label", nid)) if nid in G.nodes else str(nid) + label = label.strip().strip("()")[:_LABEL_MAXLEN] + if label and label.lower() not in seen: + seen.add(label.lower()) + names.append(label) + selected_nodes.add(nid) + if len(names) >= target_k: + break + if names: # Bare id key, NOT "Community {cid}: ..." — that string doubles as the # placeholder sentinel (_placeholder_community_labels), so a model that @@ -3330,14 +3439,7 @@ def _label_batch_with_retry( missing config, programming bug) propagates unchanged — those are never split-retried. """ - prompt = ( - "You are naming clusters in a knowledge graph. For each community below, " - "return a concise 2-5 word plain-language name describing what it is about " - "(e.g. \"Order Management\", \"Payment Flow\", \"Auth Middleware\"). " - "Each input line is ': '. " - "Respond ONLY with a JSON object mapping the community id (as a string) to " - "its name - no prose, no markdown fences.\n\n" + "\n".join(batch_lines) - ) + prompt = _LABEL_PROMPT_PREAMBLE + "\n".join(batch_lines) # Budget generously: a 2-5 word name is ~10 tokens, but models (notably # gemini) often prepend a short preamble or reasoning that eats the # completion and truncates the JSON mid-object, which used to fail the whole @@ -3380,6 +3482,64 @@ def _label_batch_with_retry( return left | right +def _estimate_text_tokens(text: str) -> int: + """Estimate token count for a text string using tiktoken if available, + falling back to standard 4 chars/token heuristic.""" + if not text: + return 0 + if _TOKENIZER is not None: + try: + return len(_TOKENIZER.encode(text, disallowed_special=())) + except Exception: + pass + return max(1, len(text) // _CHARS_PER_TOKEN) + + +def _pack_label_batches( + labeled_cids: list[int], + lines: list[str], + batch_size: int, + max_prompt_tokens: int | None = None, +) -> list[tuple[list[int], list[str]]]: + """Greedily pack communities into batches bounded by community count + (batch_size) and estimated prompt tokens (max_prompt_tokens) (#3586). + + Maintains community ordering (largest first). If a single community line + exceeds max_prompt_tokens while the current batch is empty, it is accepted + in its own single-item batch rather than dropped or causing an infinite loop. + """ + if not lines or not labeled_cids: + return [] + + if max_prompt_tokens is None: + max_prompt_tokens = _LABEL_MAX_PROMPT_TOKENS + + batch_size = max(1, batch_size) + preamble_tokens = _estimate_text_tokens(_LABEL_PROMPT_PREAMBLE) + batches: list[tuple[list[int], list[str]]] = [] + cur_cids: list[int] = [] + cur_lines: list[str] = [] + cur_tokens = preamble_tokens + + for cid, line in zip(labeled_cids, lines): + line_tokens = _estimate_text_tokens("\n" + line if cur_lines else line) + if cur_cids and (len(cur_cids) >= batch_size or cur_tokens + line_tokens > max_prompt_tokens): + batches.append((cur_cids, cur_lines)) + cur_cids = [] + cur_lines = [] + cur_tokens = preamble_tokens + line_tokens = _estimate_text_tokens(line) + + cur_cids.append(cid) + cur_lines.append(line) + cur_tokens += line_tokens + + if cur_cids: + batches.append((cur_cids, cur_lines)) + + return batches + + def label_communities( G, communities, @@ -3388,19 +3548,17 @@ def label_communities( model: str | None = None, gods=None, max_communities: int | None = None, - top_k: int = _LABEL_TOP_K, + top_k: int | None = None, batch_size: int = _LABEL_BATCH_SIZE, max_concurrency: int = 4, usage_out: dict | None = None, ) -> dict[int, str]: """Return a complete ``{cid: name}`` map using ``backend`` for naming. - Communities are labeled in batches of ``batch_size`` so the prompt fits in a - 16k-token context window (which is enough for one batch of ~100 communities - × ``top_k`` node labels). With the previous hard cap of 200 communities in a - single call, self-hosted 16k models (Qwen3, Llama 3.1 8B-Instruct, etc.) - routinely overflowed context and dropped the entire labeling pass to - placeholders. + Communities are labeled in batches constrained by ``batch_size`` and + ``_LABEL_MAX_PROMPT_TOKENS`` so the prompt fits comfortably in a 16k-token + context window. When ``top_k`` is None, sample size is chosen adaptively + by community size (#3586). ``max_communities=None`` (the default) labels every community. Pass an integer to cap the total (the legacy 200 default preserved this behavior; @@ -3413,12 +3571,12 @@ def label_communities( :func:`generate_community_labels`. """ labels = _placeholder_community_labels(communities) - cap = len(communities) if max_communities is None else max_communities - lines, labeled_cids = _community_label_lines(G, communities, gods, cap, top_k) + lines, labeled_cids = _community_label_lines(G, communities, gods, max_communities, top_k) if not lines: return labels - n_batches = (len(labeled_cids) + batch_size - 1) // batch_size + batches = _pack_label_batches(labeled_cids, lines, batch_size, max_prompt_tokens=_LABEL_MAX_PROMPT_TOKENS) + n_batches = len(batches) # Mirror extract_corpus_parallel's backend guards: Ollama serves one request at # a time per loaded model (parallel batches cause VRAM pressure and hollow @@ -3432,8 +3590,7 @@ def label_communities( workers = max(1, min(max_concurrency, n_batches)) def _run_batch(batch_idx: int): - start = batch_idx * batch_size - end = min(start + batch_size, len(labeled_cids)) + batch_cids, batch_lines = batches[batch_idx] # Accumulate token usage into a per-batch dict so concurrent workers # never race on the shared accumulator; it is merged on the main thread # in _merge (#1694). @@ -3441,7 +3598,7 @@ def _run_batch(batch_idx: int): batch_kwargs = {"usage_out": batch_usage} if usage_out is not None else {} try: parsed = _label_batch_with_retry( - labeled_cids[start:end], lines[start:end], backend=backend, model=model, + batch_cids, batch_lines, backend=backend, model=model, **batch_kwargs, ) return batch_idx, parsed, None, batch_usage @@ -3460,11 +3617,10 @@ def _merge(batch_idx: int, parsed, exc, batch_usage=None) -> None: usage_out["output"] = usage_out.get("output", 0) + batch_usage.get("output", 0) if exc is not None: errors[batch_idx] = exc - start = batch_idx * batch_size - end = min(start + batch_size, len(labeled_cids)) + batch_cids, _ = batches[batch_idx] print( f"[graphify label] batch {batch_idx + 1}/{n_batches} " - f"({end - start} communities) failed: {exc}", + f"({len(batch_cids)} communities) failed: {exc}", file=sys.stderr, ) return @@ -3499,6 +3655,7 @@ def generate_community_labels( quiet: bool = False, max_concurrency: int = 4, batch_size: int = _LABEL_BATCH_SIZE, + top_k: int | None = None, usage_out: dict | None = None, ) -> tuple[dict[int, str], str]: """CLI entry point: resolve a backend, name communities, and degrade to @@ -3531,7 +3688,7 @@ def generate_community_labels( labels = label_communities( G, communities, backend=backend, model=model, gods=gods, max_concurrency=max_concurrency, batch_size=batch_size, - usage_out=usage_out, + top_k=top_k, usage_out=usage_out, ) return labels, "llm" except Exception as exc: diff --git a/tests/test_labeling.py b/tests/test_labeling.py index d142a9fef1..6d65f099c6 100644 --- a/tests/test_labeling.py +++ b/tests/test_labeling.py @@ -878,3 +878,281 @@ def fake_generate(G, communities, *, backend=None, model=None, gods=None, assert labels["0"] == "Order Management" # real name survives assert labels["5"] == "PaymentService" # sentinel echo dropped -> hub label assert labels["7"] == "ShippingService" # bare-key echo dropped -> hub label + + +# ── #3586: Size-aware community labels and token-bounded batching ───────────── + + +def test_adaptive_sample_size_boundaries(): + """Verify exact boundary transitions for size-aware adaptive sampling.""" + from graphify.llm import _adaptive_sample_size + + boundaries = { + 5: 5, + 12: 12, + 13: 15, + 50: 15, + 51: 20, + 200: 20, + 201: 26, + 1000: 26, + 1001: 32, + 2000: 32, + } + for n, expected in boundaries.items(): + assert _adaptive_sample_size(n) == expected, ( + f"Expected _adaptive_sample_size({n}) == {expected}, got {_adaptive_sample_size(n)}" + ) + + # Monotonicity and never sampling more nodes than available for small communities + for n in range(1, 13): + assert _adaptive_sample_size(n) == n + # Upper bound cap + assert _adaptive_sample_size(10_000) == 32 + + +def test_explicit_top_k_override_and_positional(): + """Explicit top_k caps representative sample, and positional signature is preserved.""" + from graphify.llm import _community_label_lines + + G = nx.Graph() + # Large community with 100 nodes + members = [f"node_{i:03d}" for i in range(100)] + for m in members: + G.add_node(m, label=f"Label_{m}") + communities = {0: members} + + # Adaptive default (top_k=None) for N=100 produces 20 representatives + lines_adaptive, cids_adaptive = _community_label_lines(G, communities) + assert cids_adaptive == [0] + names_adaptive = lines_adaptive[0].split(": ", 1)[1].split(", ") + assert len(names_adaptive) == 20 + + # Explicit top_k=7 caps representatives at 7 + lines_explicit, _ = _community_label_lines(G, communities, top_k=7) + names_explicit = lines_explicit[0].split(": ", 1)[1].split(", ") + assert len(names_explicit) == 7 + + # Preserves positional call: _community_label_lines(G, communities, gods, max_communities, 12) + lines_pos, _ = _community_label_lines(G, communities, None, 10, 12) + names_pos = lines_pos[0].split(": ", 1)[1].split(", ") + assert len(names_pos) == 12 + + +def test_label_communities_and_generate_forward_explicit_top_k(monkeypatch): + """label_communities and generate_community_labels forward explicit top_k.""" + G = nx.Graph() + members = [f"node_{i:03d}" for i in range(100)] + for m in members: + G.add_node(m, label=f"Label_{m}") + communities = {0: members} + + captured_prompts = [] + + def fake_call(prompt, *, backend, max_tokens=200, **kwargs): + captured_prompts.append(prompt) + return '{"0": "Custom Subsystem"}' + + monkeypatch.setattr("graphify.llm._call_llm", fake_call) + + # Explicit top_k=5 via label_communities + labels = label_communities(G, communities, backend="gemini", top_k=5) + assert labels == {0: "Custom Subsystem"} + line = captured_prompts[-1].strip().splitlines()[-1] + prompt_labels = line.split(": ", 1)[1].split(", ") + assert len(prompt_labels) == 5 + + # Forwarding via generate_community_labels + gen_labels, src = generate_community_labels(G, communities, backend="gemini", top_k=8) + assert src == "llm" + assert gen_labels == {0: "Custom Subsystem"} + line_gen = captured_prompts[-1].strip().splitlines()[-1] + prompt_labels_gen = line_gen.split(": ", 1)[1].split(", ") + assert len(prompt_labels_gen) == 8 + + +def test_representative_diversity_per_file_cap(): + """Representatives are not all taken from alphabetically first source file, + per-file cap is enforced, and internal hub & god nodes are prioritized.""" + from graphify.llm import _community_label_lines + + G = nx.Graph() + # 30 nodes in a_file.py, 10 nodes in b_file.py, 10 nodes in c_file.py + # N = 50 -> adaptive sample = 15. + # Per-file cap = max(2, (15 + 2) // 3) = 5. + members = [] + for i in range(30): + nid = f"a_node_{i:02d}" + G.add_node(nid, label=f"ALabel_{i:02d}", source_file="a_file.py") + members.append(nid) + for i in range(10): + nid = f"b_node_{i:02d}" + G.add_node(nid, label=f"BLabel_{i:02d}", source_file="b_file.py") + members.append(nid) + for i in range(10): + nid = f"c_node_{i:02d}" + G.add_node(nid, label=f"CLabel_{i:02d}", source_file="c_file.py") + members.append(nid) + + # Add community-internal edges to make b_node_05 a strong internal hub + for i in range(10): + if i != 5: + G.add_edge("b_node_05", f"b_node_{i:02d}") + + # Declare c_node_09 as a global god node + gods = [{"id": "c_node_09", "label": "CLabel_09"}] + + lines, _ = _community_label_lines(G, {0: members}, gods=gods) + labels = lines[0].split(": ", 1)[1].split(", ") + assert len(labels) == 15 + + # 1. God node is prioritized first + assert labels[0] == "CLabel_09" + + # 2. Internal hub b_node_05 is prioritized high (second, before low-degree nodes) + assert labels[1] == "BLabel_05" + + # 3. Alphabetically first file a_file.py does NOT dominate the sample: + a_labels = [l for l in labels if l.startswith("ALabel_")] + b_labels = [l for l in labels if l.startswith("BLabel_")] + c_labels = [l for l in labels if l.startswith("CLabel_")] + + assert len(a_labels) == 5, f"a_file.py exceeded per-file cap of 5: {a_labels}" + assert len(b_labels) == 5, f"b_file.py did not receive 5 representatives: {b_labels}" + assert len(c_labels) == 5, f"c_file.py did not receive 5 representatives: {c_labels}" + + +def test_single_file_community_receives_full_sample(): + """Single-file community receives its full adaptive sample via backfill.""" + from graphify.llm import _community_label_lines + + G = nx.Graph() + # 50 nodes all in one source file + members = [f"single_{i:02d}" for i in range(50)] + for m in members: + G.add_node(m, label=f"SingleLabel_{m}", source_file="only_one.py") + + # N = 50 -> adaptive sample = 15. File cap = 5, but backfill fills all 15. + lines, _ = _community_label_lines(G, {0: members}) + labels = lines[0].split(": ", 1)[1].split(", ") + assert len(labels) == 15 + + +def test_missing_or_empty_source_file_does_not_crash(): + """Missing or empty source_file attributes are handled gracefully.""" + from graphify.llm import _community_label_lines + + G = nx.Graph() + members = [f"n_{i:02d}" for i in range(40)] + for i, m in enumerate(members): + if i % 3 == 0: + G.add_node(m, label=f"Lbl_{m}") # no source_file + elif i % 3 == 1: + G.add_node(m, label=f"Lbl_{m}", source_file=None) + else: + G.add_node(m, label=f"Lbl_{m}", source_file="") + + # N = 40 -> adaptive sample = 15 + lines, _ = _community_label_lines(G, {0: members}) + labels = lines[0].split(": ", 1)[1].split(", ") + assert len(labels) == 15 + + +def test_community_smaller_than_target_uses_all_members(): + """Communities with fewer members than target sample use all available nodes.""" + from graphify.llm import _community_label_lines + + G = nx.Graph() + members = ["node_a", "node_b", "node_c"] + for m in members: + G.add_node(m, label=f"Label_{m}") + + lines, _ = _community_label_lines(G, {0: members}) + labels = lines[0].split(": ", 1)[1].split(", ") + assert len(labels) == 3 + assert set(labels) == {"Label_node_a", "Label_node_b", "Label_node_c"} + + +def test_token_aware_batching_splits_oversized_batches(monkeypatch): + """Batches respect both batch_size and prompt token ceilings.""" + from graphify.llm import _pack_label_batches, _estimate_text_tokens, _LABEL_PROMPT_PREAMBLE + + # Create 30 community lines of ~400 characters (~100 tokens) each + labeled_cids = list(range(30)) + lines = [f"{cid}: " + ", ".join(f"ComponentServiceHandler_{cid}_{i}" for i in range(15)) for cid in labeled_cids] + + # If max_prompt_tokens is small, e.g. 500 tokens: + batches = _pack_label_batches(labeled_cids, lines, batch_size=100, max_prompt_tokens=500) + assert len(batches) > 1, "Should split across multiple batches when tokens exceed limit" + + # Verify each batch respects constraints and all cids are preserved + seen_cids = [] + for b_cids, b_lines in batches: + assert len(b_cids) <= 100 + assert len(b_cids) == len(b_lines) + prompt = _LABEL_PROMPT_PREAMBLE + "\n".join(b_lines) + tokens = _estimate_text_tokens(prompt) + # Each batch must be within token budget unless it's a single oversized community line + if len(b_cids) > 1: + assert tokens <= 500 + seen_cids.extend(b_cids) + + assert seen_cids == labeled_cids + + +def test_token_aware_batching_accepts_single_oversized_line(): + """A community line exceeding max_prompt_tokens alone is accepted in its own batch.""" + from graphify.llm import _pack_label_batches + + labeled_cids = [0, 1, 2] + # Very long line + huge_line = "0: " + ", ".join(["LongRepresentativeNameHere"] * 100) + normal_line_1 = "1: ShortNameA, ShortNameB" + normal_line_2 = "2: ShortNameC, ShortNameD" + + lines = [huge_line, normal_line_1, normal_line_2] + # Set max_prompt_tokens smaller than huge_line + preamble + batches = _pack_label_batches(labeled_cids, lines, batch_size=10, max_prompt_tokens=50) + + # huge_line should be in its own batch, and normal lines in subsequent batch(es) + assert len(batches) >= 2 + assert batches[0][0] == [0] + all_packed = [cid for b_cids, _ in batches for cid in b_cids] + assert all_packed == [0, 1, 2] + + +def test_label_communities_end_to_end_with_token_batching(monkeypatch): + """label_communities processes all communities across token-split batches.""" + G = nx.Graph() + # 20 communities with long member names + communities = {} + for cid in range(20): + c_members = [f"c{cid}_node_{i:02d}" for i in range(30)] + for m in c_members: + G.add_node(m, label=f"VeryLongDescriptiveComponentName_{cid}_{m}") + communities[cid] = c_members + + batches_seen = [] + + def fake_call(prompt, *, backend, max_tokens=200, **kwargs): + batches_seen.append(prompt) + # Parse cids from prompt lines + results = {} + for line in prompt.splitlines(): + m = re.match(r"^(\d+):", line) + if m: + results[m.group(1)] = f"Community Name {m.group(1)}" + return json.dumps(results) + + monkeypatch.setattr("graphify.llm._call_llm", fake_call) + + # Force a tight token limit by monkeypatching _LABEL_MAX_PROMPT_TOKENS + monkeypatch.setattr("graphify.llm._LABEL_MAX_PROMPT_TOKENS", 400) + + labels = label_communities(G, communities, backend="gemini", batch_size=100) + assert len(labels) == 20 + for cid in range(20): + assert labels[cid] == f"Community Name {cid}" + + # Verify that splitting actually occurred because of the token limit + assert len(batches_seen) > 1