typed information needs, hub suppression, seed weights, file discovery

Author: nikolay-e

src/treemapper/diffctx/__init__.py

@@ -24,13 +24,13 @@
     show_file_at_revision,
     split_diff_range,
 )
-from .graph import build_graph
+from .graph import Graph, build_graph
 from .languages import FILENAME_TO_LANGUAGE
 from .ppr import personalized_pagerank
 from .render import build_partial_tree
-from .select import lazy_greedy_select
+from .select import SelectionResult, lazy_greedy_select
 from .types import DiffHunk, Fragment, FragmentId, extract_identifiers
-from .utility import concepts_from_diff, concepts_from_diff_text
+from .utility import concepts_from_diff_text, needs_from_diff
 
 __all__ = ["GitError", "build_diff_context"]
 
@@ -172,25 +172,26 @@ def _select_with_ppr(
     alpha: float,
     tau: float,
     repo_root: Path | None = None,
+    seed_weights: dict[FragmentId, float] | None = None,
 ) -> tuple[list[Fragment], Any]:
     graph = build_graph(all_fragments, repo_root=repo_root)
-    rel_scores = personalized_pagerank(graph, core_ids, alpha=alpha)
+    rel_scores = personalized_pagerank(graph, core_ids, alpha=alpha, seed_weights=seed_weights)
 
-    concepts = concepts_from_diff(all_fragments, core_ids, graph, diff_text)
-    if not concepts:
-        concepts = concepts_from_diff_text(diff_text)
+    needs = needs_from_diff(all_fragments, core_ids, graph, diff_text)
 
     effective_budget = budget_tokens if budget_tokens is not None else _UNLIMITED_BUDGET
 
     result = lazy_greedy_select(
         fragments=all_fragments,
         core_ids=core_ids,
         rel=rel_scores,
-        concepts=concepts,
+        needs=needs,
         budget_tokens=effective_budget,
         tau=tau,
     )
-    return result.selected, result
+
+    selected = _coherence_post_pass(result, all_fragments, graph, effective_budget)
+    return selected.selected, selected
 
 
 def build_diff_context(
@@ -253,10 +254,16 @@ def build_diff_context(
 
     core_ids = _identify_core_fragments(hunks, all_fragments)
 
+    signature_frags = _generate_signature_variants(all_fragments, core_ids)
+    for frag in signature_frags:
+        frag.token_count = count_tokens(frag.content).count + _OVERHEAD_PER_FRAGMENT
+    all_fragments.extend(signature_frags)
+
     if full:
         selected = _select_full_mode(all_fragments, changed_files)
         _log_full_mode(selected)
     else:
+        seed_weights = _compute_seed_weights(hunks, core_ids, all_fragments)
         selected, result = _select_with_ppr(
             all_fragments,
             core_ids,
@@ -265,6 +272,7 @@ def build_diff_context(
             alpha,
             tau,
             repo_root=root_dir,
+            seed_weights=seed_weights,
         )
         _log_ppr_mode(selected, core_ids, budget_tokens, result, alpha, tau)
 
@@ -288,7 +296,110 @@ def _validate_inputs(root_dir: Path, alpha: float, tau: float, budget_tokens: in
         raise ValueError(f"budget_tokens must be > 0, got {budget_tokens}")
 
 
+def _coherence_post_pass(
+    result: SelectionResult,
+    all_fragments: list[Fragment],
+    graph: Graph,
+    budget: int,
+) -> SelectionResult:
+    selected_ids = {f.id for f in result.selected}
+    remaining = budget - result.used_tokens
+
+    name_to_frags: dict[str, list[Fragment]] = {}
+    for f in all_fragments:
+        if f.symbol_name:
+            name_to_frags.setdefault(f.symbol_name.lower(), []).append(f)
+
+    frag_by_id: dict[FragmentId, Fragment] = {f.id: f for f in all_fragments}
+
+    dangling_names: set[str] = set()
+    for frag in result.selected:
+        for nbr_id in graph.neighbors(frag.id):
+            if nbr_id in selected_ids:
+                continue
+            cat = graph.edge_categories.get((frag.id, nbr_id), "")
+            if cat == "semantic":
+                nbr_frag = frag_by_id.get(nbr_id)
+                if nbr_frag and nbr_frag.symbol_name:
+                    dangling_names.add(nbr_frag.symbol_name.lower())
+
+    added: list[Fragment] = []
+    for name in dangling_names:
+        candidates = name_to_frags.get(name, [])
+        for c in candidates:
+            if c.id in selected_ids:
+                break
+        else:
+            sig_candidates = [f for f in candidates if "_signature" in f.kind]
+            full_candidates = [f for f in candidates if "_signature" not in f.kind]
+            pick = sig_candidates[0] if sig_candidates else (full_candidates[0] if full_candidates else None)
+            if pick and pick.token_count <= remaining and pick.id not in selected_ids:
+                added.append(pick)
+                selected_ids.add(pick.id)
+                remaining -= pick.token_count
+
+    if not added:
+        return result
+
+    return SelectionResult(
+        selected=result.selected + added,
+        reason=result.reason,
+        used_tokens=result.used_tokens + sum(f.token_count for f in added),
+        utility=result.utility,
+    )
+
+
+def _compute_seed_weights(
+    hunks: list[DiffHunk],
+    core_ids: set[FragmentId],
+    all_fragments: list[Fragment],
+) -> dict[FragmentId, float]:
+    frag_hunk_lines: dict[FragmentId, float] = {}
+    for h in hunks:
+        h_start, h_end = h.core_selection_range
+        hunk_size = max(1, h_end - h_start + 1)
+        for frag in all_fragments:
+            if frag.id not in core_ids or frag.path != h.path:
+                continue
+            if frag.start_line <= h_end and frag.end_line >= h_start:
+                frag_hunk_lines[frag.id] = frag_hunk_lines.get(frag.id, 0) + hunk_size
+    if not frag_hunk_lines:
+        return {}
+    return frag_hunk_lines
+
+
 _CONTAINER_FRAGMENT_KINDS = frozenset({"class", "interface", "struct"})
+_SIGNATURE_ELIGIBLE_KINDS = frozenset({"function", "class", "method", "struct", "interface", "enum"})
+_MIN_LINES_FOR_SIGNATURE = 5
+
+
+def _generate_signature_variants(fragments: list[Fragment], core_ids: set[FragmentId]) -> list[Fragment]:
+    signatures: list[Fragment] = []
+    seen: set[FragmentId] = set()
+    for frag in fragments:
+        if frag.id in core_ids:
+            continue
+        if frag.kind not in _SIGNATURE_ELIGIBLE_KINDS:
+            continue
+        if frag.line_count < _MIN_LINES_FOR_SIGNATURE:
+            continue
+        lines = frag.content.splitlines()
+        sig_end = min(2, len(lines))
+        sig_content = "\n".join(lines[:sig_end])
+        sig_id = FragmentId(frag.path, frag.start_line, frag.start_line + sig_end - 1)
+        if sig_id in seen:
+            continue
+        seen.add(sig_id)
+        signatures.append(
+            Fragment(
+                id=sig_id,
+                kind=f"{frag.kind}_signature",
+                content=sig_content,
+                identifiers=frag.identifiers,
+                symbol_name=frag.symbol_name,
+            )
+        )
+    return signatures
 
 
 def _identify_core_fragments(hunks: list[DiffHunk], all_fragments: list[Fragment]) -> set[FragmentId]:

src/treemapper/diffctx/edges/__init__.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+from ..types import FragmentId
 from .base import EdgeBuilder, EdgeDict
 from .config import get_config_builders
 from .document import get_document_builders
@@ -14,6 +16,17 @@
 if TYPE_CHECKING:
     from ..types import Fragment
 
+EdgeCategories = dict[tuple[FragmentId, FragmentId], str]
+
+_BUILDER_CATEGORIES: list[tuple[str, Callable[[], list[type[EdgeBuilder]]]]] = [
+    ("semantic", get_semantic_builders),
+    ("structural", get_structural_builders),
+    ("config", get_config_builders),
+    ("document", get_document_builders),
+    ("similarity", get_similarity_builders),
+    ("history", get_history_builders),
+]
+
 
 def get_all_builders() -> list[EdgeBuilder]:
     all_builder_classes = (
@@ -27,12 +40,17 @@ def get_all_builders() -> list[EdgeBuilder]:
     return [cls() for cls in all_builder_classes]
 
 
-def collect_all_edges(fragments: list[Fragment], repo_root: Path | None = None) -> EdgeDict:
+def collect_all_edges(fragments: list[Fragment], repo_root: Path | None = None) -> tuple[EdgeDict, EdgeCategories]:
     all_edges: EdgeDict = {}
-    for builder in get_all_builders():
-        for (src, dst), weight in builder.build(fragments, repo_root).items():
-            all_edges[(src, dst)] = max(all_edges.get((src, dst), 0.0), weight)
-    return all_edges
+    edge_categories: EdgeCategories = {}
+    for category, get_builders in _BUILDER_CATEGORIES:
+        for cls in get_builders():
+            builder = cls()
+            for (src, dst), weight in builder.build(fragments, repo_root).items():
+                if weight > all_edges.get((src, dst), 0.0):
+                    all_edges[(src, dst)] = weight
+                    edge_categories[(src, dst)] = category
+    return all_edges, edge_categories
 
 
 def discover_all_related_files(
@@ -49,6 +67,7 @@ def discover_all_related_files(
 
 __all__ = [
     "EdgeBuilder",
+    "EdgeCategories",
     "EdgeDict",
     "collect_all_edges",
     "discover_all_related_files",

src/treemapper/diffctx/edges/semantic/python.py

@@ -5,15 +5,17 @@
 from collections import defaultdict
 from pathlib import Path
 
+from ...config.weights import LANG_WEIGHTS
 from ...python_semantics import PyFragmentInfo, analyze_python_fragment
 from ...types import Fragment, FragmentId
 from ..base import EdgeBuilder, EdgeDict, add_semantic_edges, path_to_module
 
 _PYTHON_EXTS = {".py", ".pyi", ".pyw"}
 
-_CALL_WEIGHT = 0.85
-_SYMBOL_REF_WEIGHT = 0.95
-_TYPE_REF_WEIGHT = 0.60
+_PY_WEIGHTS = LANG_WEIGHTS["python"]
+_CALL_WEIGHT = _PY_WEIGHTS.call
+_SYMBOL_REF_WEIGHT = _PY_WEIGHTS.symbol_ref
+_TYPE_REF_WEIGHT = _PY_WEIGHTS.type_ref
 
 _PY_IMPORT_RE = re.compile(r"(?:from\s{1,20}(\.{0,3}[\w.]{0,200})\s{1,20}import|import\s{1,20}([\w.]{1,200}))")
 

src/treemapper/diffctx/edges/structural/test.py

@@ -93,6 +93,39 @@ class TestEdgeBuilder(EdgeBuilder):
     weight_naming = EDGE_WEIGHTS["test_naming"].forward
     reverse_weight_factor = EDGE_WEIGHTS["test_reverse"].forward / EDGE_WEIGHTS["test_direct"].forward
 
+    def discover_related_files(
+        self,
+        changed_files: list[Path],
+        all_candidate_files: list[Path],
+        repo_root: Path | None = None,
+    ) -> list[Path]:
+        changed_set = set(changed_files)
+        discovered: list[Path] = []
+
+        candidate_by_stem: dict[str, list[Path]] = defaultdict(list)
+        for c in all_candidate_files:
+            if c not in changed_set:
+                candidate_by_stem[c.stem.lower()].append(c)
+
+        for changed in changed_files:
+            stem = changed.stem.lower()
+            suffix = changed.suffix.lower()
+
+            if _is_test_file(changed):
+                target = _extract_target_name_from_test(changed.stem)
+                if target:
+                    for cand in candidate_by_stem.get(target, []):
+                        if cand.suffix.lower() == suffix:
+                            discovered.append(cand)
+            else:
+                test_stems = [f"test_{stem}", f"{stem}_test"]
+                for ts in test_stems:
+                    for cand in candidate_by_stem.get(ts, []):
+                        if cand.suffix.lower() == suffix and _is_test_file(cand):
+                            discovered.append(cand)
+
+        return discovered
+
     def build(self, fragments: list[Fragment], repo_root: Path | None = None) -> EdgeDict:
         edges: EdgeDict = {}
 

src/treemapper/diffctx/graph.py

@@ -15,6 +15,7 @@
 class Graph:
     adjacency: dict[FragmentId, dict[FragmentId, float]] = field(default_factory=dict)
     nodes: set[FragmentId] = field(default_factory=set)
+    edge_categories: dict[tuple[FragmentId, FragmentId], str] = field(default_factory=dict)
 
     def add_node(self, node: FragmentId) -> None:
         self.nodes.add(node)
@@ -41,25 +42,36 @@ def build_graph(fragments: list[Fragment], repo_root: Path | None = None) -> Gra
         graph.nodes.add(frag.id)
 
     all_edges: dict[tuple[FragmentId, FragmentId], float] = {}
+    edge_categories: dict[tuple[FragmentId, FragmentId], str] = {}
 
-    plugin_edges = collect_all_edges(fragments, repo_root)
+    plugin_edges, plugin_categories = collect_all_edges(fragments, repo_root)
     for (src, dst), weight in plugin_edges.items():
-        all_edges[(src, dst)] = max(all_edges.get((src, dst), 0.0), weight)
+        if weight > all_edges.get((src, dst), 0.0):
+            all_edges[(src, dst)] = weight
+            edge_categories[(src, dst)] = plugin_categories.get((src, dst), "generic")
 
     embedding_edges = _build_embedding_edges(fragments, clamp_lexical_weight)
     for (src, dst), weight in embedding_edges.items():
-        all_edges[(src, dst)] = max(all_edges.get((src, dst), 0.0), weight)
+        if weight > all_edges.get((src, dst), 0.0):
+            all_edges[(src, dst)] = weight
+            edge_categories[(src, dst)] = "similarity"
 
-    all_edges = _apply_hub_suppression(all_edges)
+    all_edges = _apply_hub_suppression(all_edges, edge_categories)
 
     for (src, dst), weight in all_edges.items():
         graph.add_edge(src, dst, weight)
 
+    graph.edge_categories = edge_categories
+
     return graph
 
 
+_SUPPRESSION_EXEMPT = frozenset({"semantic", "structural", "config", "document"})
+
+
 def _apply_hub_suppression(
     edges: dict[tuple[FragmentId, FragmentId], float],
+    edge_categories: dict[tuple[FragmentId, FragmentId], str],
 ) -> dict[tuple[FragmentId, FragmentId], float]:
     if not edges:
         return edges
@@ -71,11 +83,16 @@ def _apply_hub_suppression(
     if not in_degree:
         return edges
 
+    degrees = sorted(in_degree.values())
+    mid = len(degrees) // 2
+    d_median = (degrees[mid] + degrees[~mid]) / 2.0
+
     suppressed: dict[tuple[FragmentId, FragmentId], float] = {}
     for (src, dst), weight in edges.items():
         deg = in_degree.get(dst, 0)
-        if deg > 0:
-            weight = weight / math.log(1 + deg)
+        cat = edge_categories.get((src, dst), "generic")
+        if deg > d_median and cat not in _SUPPRESSION_EXEMPT:
+            weight = weight / max(1.0, math.log(1 + deg))
         suppressed[(src, dst)] = weight
 
     return suppressed