reponeedle 1.1.0__py3-none-any.whl

reponeedle/__init__.py

@@ -0,0 +1,42 @@
+"""
+RepoNeedle: Token-Budget-Aware, Workflow-Adaptive Repository Context Retrieval for AI Coding Agents.
+"""
+
+from reponeedle.models import (
+    FileDoc,
+    Query,
+    RetrievalResult,
+    ScoredFile,
+    TaskType,
+    RankerType,
+    BudgetConfig,
+    ContextTier,
+)
+from reponeedle.engine import RepoNeedleEngine
+from reponeedle.core.change_impact import ChangeImpactRanker, ChangeImpactConfig
+from reponeedle.core.packer import (
+    CUDPackOptimizer,
+    BudgetOptimizer,
+    ExactTokenizer,
+    CodeSpanExtractor,
+)
+
+__version__ = "1.1.0"
+__all__ = [
+    "RepoNeedleEngine",
+    "FileDoc",
+    "Query",
+    "RetrievalResult",
+    "ScoredFile",
+    "TaskType",
+    "RankerType",
+    "BudgetConfig",
+    "ContextTier",
+    "ChangeImpactRanker",
+    "ChangeImpactConfig",
+    "CUDPackOptimizer",
+    "BudgetOptimizer",
+    "ExactTokenizer",
+    "CodeSpanExtractor",
+]
+

reponeedle/cli.py

@@ -0,0 +1,125 @@
+"""
+Command-line interface for RepoNeedle.
+"""
+
+import sys
+import argparse
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+
+from reponeedle.engine import RepoNeedleEngine
+from reponeedle.models import TaskType, BudgetConfig
+
+
+def main():
+    if hasattr(sys.stdout, "reconfigure"):
+        try:
+            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+        except Exception:
+            pass
+    parser = argparse.ArgumentParser(
+        prog="reponeedle",
+        description="Token-Budget-Aware, Workflow-Adaptive Repository Context Retrieval for AI Coding Agents.",
+    )
+    parser.add_argument("query", type=str, help="Coding task, issue description, or error traceback.")
+    parser.add_argument("--repo", "-r", type=str, default=".", help="Path to repository root (default: current directory).")
+    parser.add_argument("--budget", "-b", type=int, default=4000, help="Maximum context token budget (default: 4000).")
+    parser.add_argument(
+        "--task",
+        "-t",
+        type=str,
+        default="auto",
+        choices=["auto", "code2test", "comment2context", "trace2code", "edit2ripple", "selective"],
+        help="Explicit workflow task type (default: auto).",
+    )
+    parser.add_argument(
+        "--mode",
+        "-m",
+        type=str,
+        default="auto",
+        choices=["auto", "compact", "complete"],
+        help="CUD-Pack context delivery mode: compact (high-density spans) or complete (whole files). Default: auto.",
+    )
+    parser.add_argument(
+        "--serialize",
+        "-s",
+        action="store_true",
+        help="Print the full packaged markdown prompt to stdout.",
+    )
+    parser.add_argument(
+        "--output-file",
+        "-o",
+        type=str,
+        default=None,
+        help="Save the serialized markdown prompt context to a specified file.",
+    )
+    parser.add_argument("--no-skeletons", action="store_true", help="Disallow structural skeleton mode; require full files.")
+    parser.add_argument(
+        "--ranker",
+        "--ranking-engine",
+        dest="ranker",
+        type=str,
+        default="legacy",
+        choices=["legacy", "bm25", "arm101", "cascade", "change_impact"],
+        help="Retrieval ranker architecture (default: legacy).",
+    )
+
+    args = parser.parse_args()
+    console = Console()
+
+    console.print(Panel(f"[bold cyan]RepoNeedle Context Retrieval[/bold cyan]\nTarget Repository: [green]{args.repo}[/green]\nQuery: [yellow]\"{args.query}\"[/yellow]\nToken Budget: [magenta]{args.budget}[/magenta]\nRanker: [cyan]{args.ranker}[/cyan]", expand=False))
+
+    config = BudgetConfig(
+        max_tokens=args.budget,
+        allow_skeletons=not args.no_skeletons,
+        mode=args.mode,
+    )
+
+    with console.status("[bold green]Indexing repository..."):
+        engine = RepoNeedleEngine.from_directory(args.repo, budget_config=config)
+
+    task_type = TaskType(args.task) if args.task != "auto" else None
+    result = engine.retrieve(query=args.query, budget_tokens=args.budget, task_type=task_type, ranker=args.ranker)
+
+    if result.abstained:
+        console.print(Panel(f"[bold red]System Abstained[/bold red]\n{result.explanation}\nConfidence Score: {result.confidence_score:.3f}\nExecution Time: {result.execution_time_ms:.1f}ms", border_style="red"))
+        return
+
+    table = Table(title=f"Retrieved Context ({result.total_tokens}/{result.budget_tokens} tokens) - {result.execution_time_ms:.1f}ms")
+    table.add_column("Rank", justify="center", style="cyan")
+    table.add_column("File Path", style="green")
+    table.add_column("Tier", justify="center", style="magenta")
+    table.add_column("Tokens", justify="right", style="yellow")
+    table.add_column("Score", justify="right", style="blue")
+    table.add_column("Explanation", style="white")
+
+    for rank, sf in enumerate(result.selected_files, start=1):
+        table.add_row(
+            str(rank),
+            sf.path,
+            sf.tier.value.upper(),
+            str(sf.token_count),
+            f"{sf.score:.3f}",
+            sf.explanation,
+        )
+
+    console.print(table)
+
+    if args.output_file:
+        with open(args.output_file, "w", encoding="utf-8") as f:
+            f.write(result.rendered_context)
+        console.print(f"[bold green]Serialized prompt saved to:[/bold green] {args.output_file}")
+
+    if args.serialize:
+        console.print("\n[bold cyan]=== Serialized Context Prompt ===[/bold cyan]\n")
+        try:
+            print(result.rendered_context)
+        except UnicodeEncodeError:
+            enc = sys.stdout.encoding or "utf-8"
+            print(result.rendered_context.encode(enc, errors="replace").decode(enc))
+
+
+if __name__ == "__main__":
+    main()
+

reponeedle/core/__init__.py

@@ -0,0 +1,23 @@
+"""
+Core algorithms and sub-modules for RepoNeedle.
+"""
+
+from reponeedle.core.tokenizer import CodeTokenizer, estimate_tokens
+from reponeedle.core.lexical import BM25Index
+from reponeedle.core.semantic import SubwordSemanticIndex
+from reponeedle.core.graph import RepoGraph
+from reponeedle.core.router import WorkflowRouter
+from reponeedle.core.abstention import SelectiveAbstainer
+from reponeedle.core.budget import BudgetOptimizer
+
+__all__ = [
+    "CodeTokenizer",
+    "estimate_tokens",
+    "BM25Index",
+    "SubwordSemanticIndex",
+    "RepoGraph",
+    "WorkflowRouter",
+    "SelectiveAbstainer",
+    "BudgetOptimizer",
+]
+

reponeedle/core/abstention.py

@@ -0,0 +1,132 @@
+"""
+Calibrated selective abstention engine for no-gold and out-of-repository queries.
+Addresses the selective retrieval calibration gap identified in Agent Retrieval Bench.
+"""
+
+from typing import List, Dict, Set, Tuple
+import numpy as np
+
+from reponeedle.models import FileDoc
+from reponeedle.core.tokenizer import CodeTokenizer
+
+
+class SelectiveAbstainer:
+    """
+    Evaluates whether a query genuinely corresponds to evidence in the repository,
+    or whether the retriever should abstain to avoid polluting the agent context window.
+    """
+
+    def __init__(self, confidence_threshold: float = 0.38):
+        self.confidence_threshold = confidence_threshold
+        self.tokenizer = CodeTokenizer()
+        self.repo_vocabulary: Set[str] = set()
+        self._doc_vocab_cache: Dict[str, Tuple[str, Set[str]]] = {}
+
+    @staticmethod
+    def _stem(token: str) -> str:
+        w = token.lower()
+        for suf in ("ing", "ed", "es", "s"):
+            if len(w) > len(suf) + 2 and w.endswith(suf):
+                return w[:-len(suf)]
+        return w
+
+    def fit(self, docs: List[FileDoc]) -> "SelectiveAbstainer":
+        """
+        Builds repository vocabulary from paths, symbols, and contents.
+        """
+        self.repo_vocabulary = set()
+        for doc in docs:
+            cache_key = doc.path
+            sym_str = ",".join(doc.symbols) if doc.symbols else ""
+            content_sig = f"{doc.path}::{len(doc.content)}::{hash(doc.content)}::{sym_str}"
+            if cache_key in self._doc_vocab_cache and self._doc_vocab_cache[cache_key][0] == content_sig:
+                doc_vocab = self._doc_vocab_cache[cache_key][1]
+            else:
+                doc_vocab = set()
+                # Add path tokens
+                for t in self.tokenizer.tokenize_path(doc.path):
+                    tl = t.lower()
+                    doc_vocab.add(tl)
+                    doc_vocab.add(self._stem(tl))
+                # Add symbol tokens
+                for s in doc.symbols:
+                    for t in self.tokenizer.tokenize_text(s):
+                        tl = t.lower()
+                        doc_vocab.add(tl)
+                        doc_vocab.add(self._stem(tl))
+                # Add content identifiers if doc.content is present
+                if doc.content:
+                    for t in self.tokenizer.tokenize_text(doc.content[:3000]):
+                        if len(t) > 2:
+                            tl = t.lower()
+                            doc_vocab.add(tl)
+                            doc_vocab.add(self._stem(tl))
+                self._doc_vocab_cache[cache_key] = (content_sig, doc_vocab)
+
+            self.repo_vocabulary.update(doc_vocab)
+
+        return self
+
+    def evaluate_abstention(
+        self,
+        query: str,
+        fused_scores: np.ndarray,
+    ) -> Tuple[bool, float, str]:
+        """
+        Evaluates abstention criteria.
+        Returns:
+            abstained (bool): True if system should abstain.
+            confidence (float): Calibrated confidence score in [0, 1].
+            reason (str): Diagnostic explanation.
+        """
+        if len(fused_scores) == 0:
+            return True, 0.0, "Empty repository candidate set."
+
+        q_tokens = self.tokenizer.tokenize_text(query)
+        if not q_tokens:
+            return True, 0.0, "Empty or stop-word only query."
+
+        # 1. Lexical Overlap Ratio against repository symbol vocabulary
+        meaningful_tokens = [t for t in q_tokens if len(t) > 2]
+        if not meaningful_tokens:
+            overlap_ratio = 0.5
+        else:
+            in_repo_count = sum(
+                1 for t in meaningful_tokens
+                if t.lower() in self.repo_vocabulary or self._stem(t.lower()) in self.repo_vocabulary
+            )
+            overlap_ratio = in_repo_count / len(meaningful_tokens)
+
+        # 2. Peak-to-Noise Ratio (Z-score of top-1 file)
+        mean_score = float(np.mean(fused_scores))
+        std_score = float(np.std(fused_scores))
+        max_score = float(np.max(fused_scores))
+
+        if std_score > 1e-6:
+            z_score = (max_score - mean_score) / std_score
+        else:
+            z_score = 0.0
+
+        # Normalized Z-score component: z_score of 3.0 or higher maps to ~1.0
+        z_component = min(1.0, max(0.0, z_score / 3.0))
+
+        # 3. Score margin between top-1 and 50th percentile
+        median_score = float(np.median(fused_scores))
+        margin = max(0.0, max_score - median_score)
+
+        # Calibrated composite confidence score
+        confidence = (
+            0.30 * z_component +
+            0.45 * overlap_ratio +
+            0.25 * margin
+        )
+        confidence = float(np.clip(confidence, 0.0, 1.0))
+
+        # Decision
+        if overlap_ratio <= 0.10:
+            return True, confidence, f"Negligible query symbol overlap with repository vocabulary (overlap: {overlap_ratio:.2f}, z: {z_score:.2f})."
+
+        if confidence < self.confidence_threshold:
+            return True, confidence, f"Confidence {confidence:.3f} below threshold {self.confidence_threshold:.3f} (z: {z_score:.2f}, overlap: {overlap_ratio:.2f})."
+
+        return False, confidence, f"Confident repository match (confidence: {confidence:.3f}, z: {z_score:.2f}, overlap: {overlap_ratio:.2f})."

reponeedle/core/arm101.py

@@ -0,0 +1,301 @@
+"""
+ARM-101 Ranker: Soft Lexical Damping with Verified Function-Call Graph Expansion.
+"""
+
+from __future__ import annotations
+
+import os
+import math
+from typing import List, Dict, Optional, Any
+import numpy as np
+
+from reponeedle.models import FileDoc
+from reponeedle.core.attenuation import compute_hub_attenuation
+from reponeedle.core.function_graph import FunctionGraphEngine
+
+
+class ARM101Ranker:
+    """
+    Controlled production ranker combining decoupled Okapi BM25, subword TF-IDF,
+    verified function-call graph diffusion, and Soft Lexical Damping (tau = 0.25).
+    """
+
+    def __init__(
+        self,
+        policy: str = "unattenuated",
+        soft_damping_tau: float = 0.25,
+        forward_weight: float = 0.25,
+        reverse_weight: float = 0.20,
+        graph_weight: float = 0.25,
+        class_to_test_weight: float = 0.35,
+        doc_filter_penalty: float = 0.50,
+    ):
+        self.policy = policy
+        self.soft_damping_tau = soft_damping_tau
+        self.forward_weight = forward_weight
+        self.reverse_weight = reverse_weight
+        self.graph_weight = graph_weight
+        self.class_to_test_weight = class_to_test_weight
+        self.doc_filter_penalty = doc_filter_penalty
+
+    def _is_pkg_compat(self, dir_a: str, dir_b: str) -> bool:
+        return dir_a == dir_b or dir_a.startswith(dir_b) or dir_b.startswith(dir_a)
+
+    def rank_with_scores(
+        self,
+        docs: List[FileDoc],
+        bm25_scores: np.ndarray,
+        semantic_scores: Optional[np.ndarray] = None,
+        fg: Optional[FunctionGraphEngine] = None,
+        repo_graph: Optional[Any] = None,
+    ) -> Tuple[List[int], Dict[int, float]]:
+        n_docs = len(docs)
+        if n_docs <= 1:
+            fused = {i: float(bm25_scores[i]) if len(bm25_scores) > i else 0.0 for i in range(n_docs)}
+            return list(range(n_docs)), fused
+
+        # 1. Deterministic BM25 sort order with path tie-breaking
+        bm25_order = sorted(range(n_docs), key=lambda i: (-float(bm25_scores[i]), docs[i].path))
+        seeds = bm25_order[:3]
+        graph_boost = np.zeros(n_docs, dtype=np.float32)
+
+        # 2. Base RepoGraph propagation
+        if repo_graph is not None and getattr(repo_graph, "graph", None) is not None and len(repo_graph.graph) > 0:
+            self._propagate_repograph(seeds, docs, bm25_scores, repo_graph, graph_boost)
+            self._propagate_class_to_test(seeds, docs, bm25_scores, repo_graph, graph_boost)
+
+        # 3. Function Graph propagation
+        if fg is not None and getattr(fg, "graph", None) is not None and len(fg.graph) > 0:
+            self._propagate_function_graph(seeds, docs, bm25_scores, fg, graph_boost)
+
+        # 4. Fused score combination
+        sem = semantic_scores if semantic_scores is not None else np.zeros(n_docs, dtype=np.float32)
+        fused = {}
+        for idx in range(n_docs):
+            score = 0.60 * float(bm25_scores[idx]) + 0.15 * float(sem[idx]) + self.graph_weight * float(graph_boost[idx])
+            p_lower = docs[idx].path.lower().replace("\\", "/")
+            if any(seg in p_lower for seg in ["/doc/", "/docs/", "doc/", "docs/", "documentation/", "guide/", "tutorial/", ".github/"]) or p_lower.endswith((".md", ".rst", ".txt", ".adoc", ".html")):
+                score *= self.doc_filter_penalty
+            fused[idx] = float(score)
+
+        sorted_by_fused = sorted(range(n_docs), key=lambda idx: (-fused[idx], docs[idx].path))
+
+        # 5. Soft Lexical Damping (Rank 1 Guard)
+        s1 = bm25_order[0]
+        s2 = bm25_order[1] if n_docs > 1 else s1
+        bm25_margin = max(0.0, float(bm25_scores[s1]) - float(bm25_scores[s2]))
+        threshold = self.soft_damping_tau * bm25_margin
+
+        best = sorted_by_fused[0]
+        if best == s1 or (fused[best] - fused[s1] >= threshold):
+            return sorted_by_fused, fused
+        else:
+            fused[s1] = fused[best] + 1e-4
+            return [s1] + [idx for idx in sorted_by_fused if idx != s1], fused
+
+    def rank(
+        self,
+        docs: List[FileDoc],
+        bm25_scores: np.ndarray,
+        semantic_scores: Optional[np.ndarray] = None,
+        fg: Optional[FunctionGraphEngine] = None,
+        repo_graph: Optional[Any] = None,
+    ) -> List[int]:
+        ranked_indices, _ = self.rank_with_scores(docs, bm25_scores, semantic_scores, fg, repo_graph)
+        return ranked_indices
+
+    def _propagate_function_graph(
+        self,
+        seeds: List[int],
+        docs: List[FileDoc],
+        bm25_scores: np.ndarray,
+        fg: FunctionGraphEngine,
+        graph_boost: np.ndarray,
+    ):
+        path_to_idx = fg.path_to_idx
+        allowed = {"verified_call", "unique_global_call", "intra_file_call"}
+
+        for s_idx in seeds:
+            seed_p = docs[s_idx].path
+            seed_score = float(bm25_scores[s_idx])
+            if seed_score <= 0.0 or seed_p not in path_to_idx:
+                continue
+
+            seed_dir = os.path.dirname(seed_p.replace("\\", "/"))
+
+            # Forward propagation (callees)
+            callees = fg.callees_map.get(seed_p, {})
+            out_deg = max(1, len(callees))
+            for callee_p, conf in callees.items():
+                if callee_p in path_to_idx:
+                    c_idx = path_to_idx[callee_p]
+                    edges = fg.file_to_file_edges.get((seed_p, callee_p), [])
+                    matching = [e for e in edges if e.edge_type in allowed]
+                    if not matching:
+                        continue
+                    edge_conf = max(e.confidence for e in matching)
+
+                    callee_dir = os.path.dirname(callee_p.replace("\\", "/"))
+                    is_test = docs[s_idx].is_test or docs[c_idx].is_test
+                    if not is_test and not self._is_pkg_compat(seed_dir, callee_dir):
+                        continue
+
+                    in_deg = max(1, len(fg.callers_map.get(callee_p, {})))
+                    callee_out_deg = max(1, len(fg.callees_map.get(callee_p, {})))
+                    norm = 1.0 / math.sqrt(out_deg * in_deg)
+                    boost = seed_score * self.forward_weight * edge_conf * norm
+
+                    w_hub = compute_hub_attenuation(callee_p, in_deg, callee_out_deg, self.policy)
+                    boost *= w_hub
+                    graph_boost[c_idx] = max(graph_boost[c_idx], float(boost))
+
+            # Reverse propagation (callers)
+            callers = fg.callers_map.get(seed_p, {})
+            in_deg_seed = max(1, len(callers))
+            for caller_p, conf in callers.items():
+                if caller_p in path_to_idx:
+                    c_idx = path_to_idx[caller_p]
+                    edges = fg.file_to_file_edges.get((caller_p, seed_p), [])
+                    matching = [e for e in edges if e.edge_type in allowed]
+                    if not matching:
+                        continue
+                    edge_conf = max(e.confidence for e in matching)
+
+                    caller_dir = os.path.dirname(caller_p.replace("\\", "/"))
+                    is_test = docs[s_idx].is_test or docs[c_idx].is_test
+                    if not is_test and not self._is_pkg_compat(seed_dir, caller_dir):
+                        continue
+
+                    caller_in_deg = max(1, len(fg.callers_map.get(caller_p, {})))
+                    caller_out_deg = max(1, len(fg.callees_map.get(caller_p, {})))
+                    norm = 1.0 / math.sqrt(in_deg_seed * caller_out_deg)
+                    boost = seed_score * self.reverse_weight * edge_conf * norm
+
+                    w_hub = compute_hub_attenuation(caller_p, caller_in_deg, caller_out_deg, self.policy)
+                    boost *= w_hub
+                    graph_boost[c_idx] = max(graph_boost[c_idx], float(boost))
+
+    def _propagate_repograph(
+        self,
+        seeds: List[int],
+        docs: List[FileDoc],
+        bm25_scores: np.ndarray,
+        repo_graph: Any,
+        graph_boost: np.ndarray,
+    ):
+        g = repo_graph.graph
+        path_to_idx = repo_graph.path_to_idx
+
+        for s_idx in seeds:
+            seed_p = docs[s_idx].path
+            seed_score = float(bm25_scores[s_idx])
+            if seed_score <= 0.0 or seed_p not in g:
+                continue
+
+            seed_dir = os.path.dirname(seed_p.replace("\\", "/"))
+            out_deg = max(1, g.out_degree(seed_p))
+            in_deg_seed = max(1, g.in_degree(seed_p))
+
+            for succ in g.successors(seed_p):
+                if succ in path_to_idx:
+                    n_idx = path_to_idx[succ]
+                    succ_dir = os.path.dirname(succ.replace("\\", "/"))
+                    if not self._is_pkg_compat(seed_dir, succ_dir):
+                        continue
+                    in_deg_succ = max(1, g.in_degree(succ))
+                    norm = 1.0 / math.sqrt(out_deg * in_deg_succ)
+                    boost = seed_score * 0.25 * norm
+                    if self.policy == "uniform_log2":
+                        boost *= (1.0 / math.log2(2.0 + in_deg_succ))
+                    graph_boost[n_idx] = max(graph_boost[n_idx], float(boost))
+
+            for pred in g.predecessors(seed_p):
+                if pred in path_to_idx:
+                    n_idx = path_to_idx[pred]
+                    pred_dir = os.path.dirname(pred.replace("\\", "/"))
+                    if not self._is_pkg_compat(seed_dir, pred_dir):
+                        continue
+                    out_deg_pred = max(1, g.out_degree(pred))
+                    in_deg_pred = max(1, g.in_degree(pred))
+                    norm = 1.0 / math.sqrt(in_deg_seed * out_deg_pred)
+                    boost = seed_score * 0.10 * norm
+                    if self.policy == "uniform_log2":
+                        boost *= (1.0 / math.log2(2.0 + in_deg_pred))
+                    graph_boost[n_idx] = max(graph_boost[n_idx], float(boost))
+
+            self._apply_test_discovery(seed_p, seed_score, path_to_idx, graph_boost)
+
+    def _apply_test_discovery(
+        self,
+        seed_path: str,
+        seed_score: float,
+        path_to_idx: Dict[str, int],
+        graph_boost: np.ndarray,
+    ):
+        norm_p = seed_path.replace("\\", "/")
+        stem = os.path.splitext(os.path.basename(norm_p))[0]
+        ext = os.path.splitext(norm_p)[1].lower()
+        cur_dir = os.path.dirname(norm_p)
+        tc_list = []
+
+        if ext == ".go":
+            tc_list.append(f"{cur_dir}/{stem}_test.go" if cur_dir else f"{stem}_test.go")
+        elif ext in (".py", ".pyi"):
+            tc_list.append(f"{cur_dir}/test_{stem}.py" if cur_dir else f"test_{stem}.py")
+            tc_list.append(f"tests/test_{stem}.py")
+            tc_list.append(f"tests/{stem}_test.py")
+        elif ext == ".rs":
+            tc_list.append(f"tests/{stem}.rs")
+            tc_list.append(f"tests/test_{stem}.rs")
+        elif ext in (".ts", ".tsx", ".js", ".jsx"):
+            tc_list.append(f"{cur_dir}/{stem}.test{ext}" if cur_dir else f"{stem}.test{ext}")
+            tc_list.append(f"{cur_dir}/__tests__/{stem}.test{ext}" if cur_dir else f"__tests__/{stem}.test{ext}")
+
+        for tc in tc_list:
+            if tc in path_to_idx:
+                idx = path_to_idx[tc]
+                graph_boost[idx] = max(graph_boost[idx], float(seed_score * 0.35))
+
+    def _propagate_class_to_test(
+        self,
+        seeds: List[int],
+        docs: List[FileDoc],
+        bm25_scores: np.ndarray,
+        repo_graph: Any,
+        graph_boost: np.ndarray,
+    ):
+        prod_to_tests = getattr(repo_graph, "prod_to_test_invocations", {})
+        test_to_prods = getattr(repo_graph, "test_to_prod_invocations", {})
+        path_to_idx = repo_graph.path_to_idx
+
+        for s_idx in seeds:
+            seed_p = docs[s_idx].path
+            seed_score = float(bm25_scores[s_idx])
+            if seed_score <= 0.0:
+                continue
+
+            if seed_p in prod_to_tests:
+                target_tests = prod_to_tests[seed_p]
+                out_deg = max(1, len(target_tests))
+                for test_p in target_tests:
+                    if test_p in path_to_idx:
+                        t_idx = path_to_idx[test_p]
+                        in_deg = max(1, len(test_to_prods.get(test_p, {seed_p})))
+                        norm = 1.0 / math.sqrt(out_deg * in_deg)
+                        boost = seed_score * self.class_to_test_weight * norm
+                        if self.policy == "uniform_log2":
+                            boost *= (1.0 / math.log2(2.0 + in_deg))
+                        graph_boost[t_idx] = max(graph_boost[t_idx], float(boost))
+
+            if seed_p in test_to_prods:
+                target_prods = test_to_prods[seed_p]
+                out_deg = max(1, len(target_prods))
+                for prod_p in target_prods:
+                    if prod_p in path_to_idx:
+                        p_idx = path_to_idx[prod_p]
+                        in_deg = max(1, len(prod_to_tests.get(prod_p, {seed_p})))
+                        norm = 1.0 / math.sqrt(out_deg * in_deg)
+                        boost = seed_score * self.class_to_test_weight * norm
+                        if self.policy == "uniform_log2":
+                            boost *= (1.0 / math.log2(2.0 + in_deg))
+                        graph_boost[p_idx] = max(graph_boost[p_idx], float(boost))