diffcontext 0.5.1__py3-none-any.whl

diffcontext/context/selector.py

@@ -0,0 +1,258 @@
+"""
+selector.py — Select which symbols to include in context, respecting token budget.
+
+Key fix over original:
+  - Changed symbols are the ONLY unconditional include.
+  - Score >= 80 threshold no longer bypasses the token budget.
+    That rule caused direct callees (score ~90 after structural bonus)
+    to consume the entire budget before any co-change sibling candidates
+    (score ~38-50) were even evaluated. Precision suffered badly.
+  - Instead: rank strictly by score, apply budget universally, with one
+    exception: changed symbols always fit (they're the reason we're here).
+  - Added a per-symbol token cap so one giant function can't crowd out
+    ten relevant small ones.
+
+Fix vs previous version:
+  The token cap logic was wrong: it counted a capped amount toward the
+  budget (250 tokens) but included the full symbol in the result, causing
+  silent overruns. The correct behavior is: if a symbol exceeds the cap,
+  SKIP IT ENTIRELY rather than including it at a lie. This means the
+  selector tries smaller candidates next instead of filling context with
+  one huge function and then claiming there's room for more.
+
+Fix vs previous version (token-accounting mismatch):
+  The selector used to budget on token_count(symbol.code) — the bare
+  function body — while the compiler renders each symbol with a FILE:/
+  FUNCTION: header and a CALLERS/CALLEES relationship block on top of the
+  code, and reports tokens over that full rendered block. The gap between
+  what was budgeted and what was emitted produced a systematic 25-41%
+  overshoot of --max-tokens (reproduced on psf/black at every budget from
+  500 to 8000). Now, when the caller passes the call graph, each candidate
+  is measured with compiler.render_symbol_block() — the exact rendering the
+  compiler will emit — using a pessimistic empty selected_set so every
+  relationship entry counts the longer " [NOT IN CONTEXT]" tag. Without a
+  graph the old code-only behavior is preserved so existing library callers
+  don't silently change.
+"""
+
+from typing import Callable, Dict, List, Optional, Set, Tuple
+
+from ..models import Symbol
+from .compiler import build_reverse_graph, relationship_cap, render_symbol_block
+
+
+# A single symbol can burn at most this fraction of the total budget.
+# Prevents one huge function from crowding out ten relevant small ones.
+MAX_SINGLE_SYMBOL_FRACTION = 0.25
+
+# Largest-gap cutoff ("gap50" in the benchmarks): the relative-drop search
+# window and the score floor below which a candidate is not retrievable at
+# all. Both values are the exact ones measured in blend_loro.py
+# eval_cutoff_policies (benchmarks/RIGOR_REPORT_2026-07.md §7).
+GAP_CUTOFF_WINDOW = 50
+GAP_SCORE_EPSILON = 1e-12
+
+# Minimum ratio for the gap to fire. The original gap_cut_count always cut at
+# the largest relative drop, even when that drop was 1.10x (noise, not a real
+# break). On ContextBench, 99% of 624 missed gold symbols were gap_cut — many
+# at rank 3-4 with scores 75-84, cut by a trivial 1.10x drop at rank 2.
+# min_ratio=1.5 means the gap only fires when the drop is >= 50% — a real break.
+GAP_MIN_RATIO = 1.5
+
+# Minimum number of candidates the gap must keep. The original kept ~1-2 on
+# ContextBench (out of 7000+ scored symbols), over-pruning to the point where
+# direct callees at score 90 were cut. min_keep=10 ensures the gap never
+# prunes below 10 candidates; the budget controls context size after that.
+GAP_MIN_KEEP = 10
+
+
+def gap_cut_count(ranked_scores: List[float],
+                  min_ratio: float = 1.0, min_keep: int = 0) -> int:
+    """
+    How many leading candidates the largest-gap cutoff keeps.
+
+    `ranked_scores` must be positive scores sorted descending. The cut lands
+    at the largest relative drop (score[i] / score[i+1]) between consecutive
+    candidates within the first GAP_CUTOFF_WINDOW — the policy measured
+    F1-optimal on all five benchmark repos: ~4x the precision of fixed
+    top-20 at 6-9 retrieved symbols, for ~30% relative recall cost
+    (benchmarks/RIGOR_REPORT_2026-07.md §7). With fewer than 3 candidates
+    there is no distribution to read, so everything is kept.
+
+    Args:
+        min_ratio: only cut if the largest relative drop >= this ratio.
+            Default 1.0 = always cut (original behavior). 1.5 = only cut
+            when the drop is >= 50% (filters noise drops like 1.10x).
+        min_keep: always keep at least this many candidates, even if the
+            gap would cut earlier. Default 0 = no minimum (original
+            behavior). 10 = never prune below 10 candidates; the budget
+            controls context size after that.
+    """
+    n = len(ranked_scores)
+    if n < 3:
+        return n
+    head = ranked_scores[:GAP_CUTOFF_WINDOW]
+    best_i, best_ratio = 1, 0.0
+    for i in range(len(head) - 1):
+        ratio = head[i] / max(head[i + 1], GAP_SCORE_EPSILON)
+        if ratio > best_ratio:
+            best_ratio = ratio
+            best_i = i + 1
+    # Only apply the gap if the drop is significant enough to be a real break.
+    if best_ratio < min_ratio:
+        return n  # no significant gap — keep all, let budget/top_k decide
+    # Never prune below min_keep candidates.
+    return max(best_i, min_keep)
+
+
+def select_context(
+    symbols: Dict[str, Symbol],
+    scores: Dict[str, float],
+    changed: List[str],
+    max_tokens: Optional[int] = None,
+    token_counter: Optional[Callable[[str], int]] = None,
+    top_k: Optional[int] = None,
+    graph: Optional[Dict[str, List[str]]] = None,
+    reverse: Optional[Dict[str, Set[str]]] = None,
+    rel_cap: Optional[int] = None,
+    cutoff: Optional[str] = None,
+    gap_min_ratio: float = 1.0,
+    gap_min_keep: int = 0,
+) -> Tuple[List[str], List[str]]:
+    """
+    Select symbols for context based on scores and token budget.
+
+    Priority:
+      1. Changed symbols always included (no budget bypass for anything else)
+      2. All remaining symbols ranked by score, included until budget exhausted
+
+    Args:
+        token_counter: Optional callable text -> token count. Pass your
+            model's real tokenizer (e.g. tiktoken, Anthropic counting) when
+            enforcing a hard context-window limit; defaults to the
+            ~4-chars-per-token heuristic, which is approximate.
+        top_k: Optional cap on the number of NON-changed symbols included,
+            applied on top of the token budget. The eval_v2 benchmark found
+            retrieval recall plateaus around 20 symbols per changed symbol
+            while precision keeps degrading, so a caller optimizing for
+            signal-to-noise should pass ~20 * len(changed).
+        graph: Optional call graph (id -> [dep ids]). When provided, each
+            candidate is budgeted at its FULL rendered size (headers +
+            relationship annotations + code) via render_symbol_block, which
+            is what the compiler actually emits. Omit for the legacy
+            code-only accounting (kept for backward compatibility, but it
+            undercounts and the compiled output will overshoot the budget).
+        reverse: Optional precomputed reverse graph (callee -> callers).
+            Derived from `graph` when absent.
+        rel_cap: Relationship-block entry cap used for size measurement;
+            defaults to compiler.relationship_cap(max_tokens) so selector
+            and compiler always measure the same rendering.
+        cutoff: Optional dynamic cutoff policy applied to the score ranking
+            BEFORE top_k and the token budget (the order the benchmark
+            measured). "gap" cuts at the largest relative score drop
+            (see gap_cut_count) and additionally drops zero-score
+            candidates, which the policy never retrieves. None/"topk"
+            keeps today's recall-first behavior. Note: the benchmark
+            measured the policy per single-changed-symbol query; with
+            multiple changed symbols it applies to the merged ranking.
+
+    Returns:
+        (selected_ids, dropped_ids)
+        dropped_ids: scored symbols that exist in `symbols` but were cut by
+        the token budget. The LLM is told about these explicitly.
+    """
+    count = token_counter or _estimate_tokens
+
+    if cutoff not in (None, "topk", "gap"):
+        raise ValueError(
+            f"unknown cutoff policy {cutoff!r} — expected 'gap', 'topk', or None"
+        )
+
+    if not scores:
+        return list(changed), []
+
+    if graph is not None and reverse is None:
+        reverse = build_reverse_graph(graph)
+    if rel_cap is None:
+        rel_cap = relationship_cap(max_tokens)
+
+    def rendered_size(sym_id: str) -> int:
+        """Tokens this symbol will actually cost in the compiled output."""
+        if graph is None:
+            return count(symbols[sym_id].code)
+        assert reverse is not None  # derived from graph above when absent
+        # Empty selected_set = every relationship entry gets the longer
+        # " [NOT IN CONTEXT]" tag = safe upper bound on the real rendering.
+        return count(render_symbol_block(
+            sym_id, symbols, scores.get(sym_id, 0), graph, reverse,
+            set(), rel_cap=rel_cap,
+        ))
+
+    per_sym_cap = int(max_tokens * MAX_SINGLE_SYMBOL_FRACTION) if max_tokens else None
+
+    changed_set = set(changed)
+    result: List[str] = []
+    dropped: List[str] = []
+    current_tokens = 0
+
+    # ── Pass 1: changed symbols always in, no budget check ───────────────
+    for sym_id in changed:
+        if sym_id in symbols:
+            result.append(sym_id)
+            current_tokens += rendered_size(sym_id)
+
+    # ── Pass 2: everything else ranked by score, budget-gated ────────────
+    scored = sorted(
+        ((sid, sc) for sid, sc in scores.items() if sid not in changed_set),
+        key=lambda x: x[1],
+        reverse=True,
+    )
+
+    gap_kept: Optional[Set[str]] = None
+    if cutoff == "gap":
+        candidates = [
+            (sid, sc) for sid, sc in scored
+            if sid in symbols and sc > GAP_SCORE_EPSILON
+        ]
+        keep_n = gap_cut_count([sc for _, sc in candidates],
+                               min_ratio=gap_min_ratio, min_keep=gap_min_keep)
+        gap_kept = {sid for sid, _ in candidates[:keep_n]}
+
+    included_non_changed = 0
+    for sym_id, score in scored:
+        if sym_id not in symbols:
+            continue
+
+        if gap_kept is not None and sym_id not in gap_kept:
+            dropped.append(sym_id)
+            continue
+
+        if top_k is not None and included_non_changed >= top_k:
+            dropped.append(sym_id)
+            continue
+
+        sym_tokens = rendered_size(sym_id)
+
+        # FIX: if the symbol exceeds the per-symbol cap, skip it entirely.
+        # Previous code counted the capped amount toward the budget but
+        # still included the full symbol, silently overrunning the budget
+        # and then continuing to include more symbols as if there were room.
+        # Skipping is correct: the budget should gate what's actually included.
+        if per_sym_cap is not None and sym_tokens > per_sym_cap:
+            dropped.append(sym_id)
+            continue
+
+        if max_tokens is not None and current_tokens + sym_tokens > max_tokens:
+            dropped.append(sym_id)
+            continue
+
+        result.append(sym_id)
+        included_non_changed += 1
+        current_tokens += sym_tokens
+
+    return result, dropped
+
+
+def _estimate_tokens(text: str) -> int:
+    """~4 chars per token (GPT approximation). Add 20% buffer for safety."""
+    return max(1, int(len(text) / 4 * 1.2))

diffcontext/diff/__init__.py

@@ -0,0 +1 @@
+"""diff subpackage — git diff + state-based change detection."""

diffcontext/diff/git_diff.py

@@ -0,0 +1,298 @@
+"""
+git_diff.py — Extract changed files/symbols from git diff output.
+"""
+
+import os
+import subprocess
+from typing import List, Optional, Set
+
+
+def get_changed_files(
+    repo_path: str,
+    ref: str = "HEAD~1",
+    against: "Optional[str]" = None,
+) -> List[str]:
+    """
+    Get list of Python files changed between `ref` and `against`.
+
+    against=None (default): compare against the WORKING TREE, i.e. include
+    uncommitted changes (staged or not). This is `git diff <ref>` with no
+    second ref -- the same thing `git status`-style tools show you before
+    you commit.
+    against="HEAD" (or any other ref): compare two committed snapshots only;
+    uncommitted edits are invisible to this mode.
+
+    Returns list of relative paths like ["./src/auth.py", "./api/login.py"]
+    """
+    repo_path = os.path.abspath(repo_path)
+    try:
+        cmd = ["git", "diff", "--name-only", "--diff-filter=ACMR", ref]
+        if against is not None:
+            cmd.append(against)
+
+        result = subprocess.run(
+            cmd,
+            cwd=repo_path,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode != 0:
+            return []
+
+        from ..languages import indexable_extensions
+        exts = indexable_extensions()
+
+        files = []
+        for line in result.stdout.strip().split("\n"):
+            line = line.strip()
+            if line.endswith(exts):
+                files.append("./" + line)
+        return files
+
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        return []
+
+
+def get_changed_lines(
+    repo_path: str,
+    filepath: str,
+    ref: str = "HEAD~1",
+    against: "Optional[str]" = None,
+) -> Set[int]:
+    """
+    Get set of changed line numbers for a specific file.
+
+    against=None (default): compare `ref` against the working tree,
+    including uncommitted edits. See get_changed_files for details.
+
+    Returns set of 1-indexed line numbers that were added or modified.
+    """
+    repo_path = os.path.abspath(repo_path)
+    try:
+        cmd = ["git", "diff", "-U0", ref]
+        if against is not None:
+            cmd.append(against)
+        cmd += ["--", filepath.lstrip("./")]
+
+        result = subprocess.run(
+            cmd,
+            cwd=repo_path,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode != 0:
+            return set()
+
+        changed_lines: Set[int] = set()
+        for line in result.stdout.split("\n"):
+            if line.startswith("@@"):
+                # Parse @@ -old,count +new,count @@
+                parts = line.split()
+                for part in parts:
+                    if part.startswith("+") and "," in part:
+                        start_str, count_str = part[1:].split(",", 1)
+                        start = int(start_str)
+                        count = int(count_str)
+                        for i in range(start, start + max(count, 1)):
+                            changed_lines.add(i)
+                    elif part.startswith("+") and part[1:].isdigit():
+                        changed_lines.add(int(part[1:]))
+
+        return changed_lines
+
+    except (subprocess.TimeoutExpired, FileNotFoundError, ValueError):
+        return set()
+
+
+def get_patch_text(
+    repo_path: str,
+    filepath: str,
+    ref: str = "HEAD~1",
+    against: "Optional[str]" = None,
+    context_lines: int = 3,
+) -> str:
+    """
+    Return the raw unified-diff patch text for a single file between
+    `ref` and `against` (working tree if against=None).
+
+    Useful for files that no longer parse: you can't get a line-level
+    symbol diff out of AST comparison, but the raw patch still shows
+    exactly what text changed (e.g. a commented-out `class` line).
+
+    Returns "" if there's no diff, the file/ref doesn't exist, or git fails.
+    """
+    repo_path = os.path.abspath(repo_path)
+    try:
+        cmd = ["git", "diff", f"-U{context_lines}", ref]
+        if against is not None:
+            cmd.append(against)
+        cmd += ["--", filepath.lstrip("./")]
+
+        result = subprocess.run(
+            cmd,
+            cwd=repo_path,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode != 0:
+            return ""
+        return result.stdout
+
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        return ""
+
+
+def find_changed_symbols(
+    repo_path: str,
+    symbols: dict,
+    ref: str = "HEAD~1",
+    against: "Optional[str]" = None,
+    broken_files: Optional[List[str]] = None,
+    broken_file_patches: Optional[dict] = None,
+    known_broken_files: Optional[List[str]] = None,
+) -> List[str]:
+    """
+    Find which symbol IDs are affected by a git diff.
+
+    against=None (default): compares `ref` against the WORKING TREE, so
+    uncommitted edits are picked up. Pass against="HEAD" explicitly if you
+    only want committed changes.
+
+    Cross-references changed lines with symbol line ranges to find
+    exactly which functions/methods were modified.
+
+    `symbols` is the CURRENT index (built from the working tree / `against`
+    ref). A changed file that contributes zero entries to `symbols` is
+    ambiguous on its own -- it might have failed to parse (SyntaxError), or
+    it might just be a file with no function/method definitions (setup.py,
+    a constants module, an __init__.py with only imports, etc). To tell
+    these apart correctly, pass `known_broken_files` -- the ground-truth
+    list from extract_all_symbols()/RepositoryIndex.broken_files, which
+    only contains files that actually raised SyntaxError. Only files in
+    that list get the broken-file fallback treatment below; any other
+    zero-symbol file is treated as a normal (legitimately function-less)
+    file and simply contributes no changed symbols.
+
+    For files confirmed broken via `known_broken_files`:
+      - the file's relative path is appended to `broken_files` (if provided)
+      - we fall back to the symbol IDs that existed in the file at `ref`
+        (the prior, presumably-working revision) via `git show`, so the
+        change is still reported instead of silently disappearing.
+      - if `broken_file_patches` (a dict) is provided, it's populated with
+        {relative_file: raw_patch_text} -- the actual unified diff, since a
+        broken file can't be symbol-diffed and the patch is the only real
+        signal of what changed.
+    """
+    known_broken = set(known_broken_files or ())
+    changed_files = get_changed_files(repo_path, ref, against)
+    if not changed_files:
+        return []
+
+    changed_symbols = []
+
+    for sym_id, sym in symbols.items():
+        sym_file = "./" + os.path.relpath(sym.file, os.path.abspath(repo_path))
+        if sym_file not in changed_files:
+            continue
+
+        changed_lines = get_changed_lines(repo_path, sym_file, ref, against)
+        if not changed_lines:
+            continue
+
+        code_lines = sym.code.count("\n") + 1
+        sym_lines = set(range(sym.lineno, sym.lineno + code_lines))
+
+        if sym_lines & changed_lines:
+            changed_symbols.append(sym_id)
+
+    # Check for deleted symbols or broken files
+    for changed_file in changed_files:
+        # Handle broken files (SyntaxError)
+        if changed_file in known_broken:
+            if broken_files is not None and changed_file not in broken_files:
+                broken_files.append(changed_file)
+
+            if broken_file_patches is not None:
+                broken_file_patches[changed_file] = get_patch_text(
+                    repo_path, changed_file, ref, against
+                )
+
+            prior_ids = _symbol_ids_at_ref(repo_path, changed_file, ref)
+            for p_id in prior_ids:
+                if p_id not in changed_symbols:
+                    changed_symbols.append(p_id)
+            continue
+
+        # Handle valid files: find symbols that existed before but are gone now
+        prior_ids = _symbol_ids_at_ref(repo_path, changed_file, ref)
+        current_ids = {
+            sym_id for sym_id, sym in symbols.items()
+            if "./" + os.path.relpath(sym.file, os.path.abspath(repo_path)) == changed_file
+        }
+        
+        deleted_ids = set(prior_ids) - current_ids
+        for d_id in deleted_ids:
+            if d_id not in changed_symbols:
+                changed_symbols.append(d_id)
+
+    return changed_symbols
+
+
+def _symbol_ids_at_ref(repo_path: str, filepath: str, ref: str) -> List[str]:
+    """
+    Best-effort: list symbol IDs ("./file.py:Name") that existed in
+    `filepath` at git ref `ref`, by parsing the file's contents at that
+    revision. Returns [] if the file didn't exist, wasn't valid Python
+    at that ref either, or git/parse fails for any reason.
+    """
+    repo_path = os.path.abspath(repo_path)
+    try:
+        result = subprocess.run(
+            ["git", "show", f"{ref}:{filepath.lstrip('./')}"],
+            cwd=repo_path,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode != 0:
+            return []
+
+        import ast as _ast
+        try:
+            tree = _ast.parse(result.stdout)
+        except SyntaxError:
+            return []
+
+        class _FuncVisitor(_ast.NodeVisitor):
+            def __init__(self):
+                self.class_stack = []
+                self.names = []
+
+            def visit_ClassDef(self, node):
+                self.class_stack.append(node.name)
+                self.generic_visit(node)
+                self.class_stack.pop()
+
+            def visit_FunctionDef(self, node):
+                self._add(node)
+                self.generic_visit(node)
+
+            def visit_AsyncFunctionDef(self, node):
+                self._add(node)
+                self.generic_visit(node)
+
+            def _add(self, node):
+                if self.class_stack:
+                    self.names.append(f"{self.class_stack[-1]}.{node.name}")
+                else:
+                    self.names.append(node.name)
+
+        visitor = _FuncVisitor()
+        visitor.visit(tree)
+        return [f"{filepath}:{name}" for name in visitor.names]
+
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        return []

diffcontext/diff/state_manager.py

@@ -0,0 +1,75 @@
+"""
+state_manager.py — Snapshot-based change detection (no git required).
+"""
+
+import json
+import os
+from typing import Dict, Iterable, Optional
+
+from ..models import DiffResult
+
+
+def _symbol_file(fn_id: str) -> str:
+    """Extract the file portion of a 'file.py:Class.method' symbol id."""
+    return fn_id.split(":", 1)[0] if ":" in fn_id else fn_id
+
+
+def compare_states(
+    previous: Dict[str, dict],
+    current: Dict[str, dict],
+    broken_files: Optional[Iterable[str]] = None,
+) -> DiffResult:
+    """
+    Compare two snapshots of repository functions.
+
+    Each snapshot is: {function_id: {"code": "..."}}
+
+    broken_files: relative paths (e.g. "./objects.py") of files that failed
+    to parse (SyntaxError) when building `current`. Symbols that previously
+    existed in one of these files are reported as `modified` rather than
+    `deleted` -- the function wasn't removed, the file just can't be parsed
+    right now. The file itself is also recorded in `broken_files` so callers
+    can flag it distinctly from a normal diff.
+    """
+    broken_files = set(broken_files or ())
+
+    modified = []
+    added = []
+    deleted = []
+
+    for fn_id in previous:
+        if fn_id not in current:
+            if _symbol_file(fn_id) in broken_files:
+                # File failed to parse this run -- treat every symbol that
+                # used to live there as modified (not deleted), since the
+                # change (the syntax break) is exactly what needs review.
+                modified.append(fn_id)
+            else:
+                deleted.append(fn_id)
+        elif previous[fn_id] != current[fn_id]:
+            modified.append(fn_id)
+
+    for fn_id in current:
+        if fn_id not in previous:
+            added.append(fn_id)
+
+    return DiffResult(
+        modified=modified,
+        added=added,
+        deleted=deleted,
+        broken_files=sorted(broken_files),
+    )
+
+
+def save_state(state: Dict, path: str = "diffcontext_state.json"):
+    """Save function snapshot to disk."""
+    with open(path, "w") as f:
+        json.dump(state, f, indent=2)
+
+
+def load_state(path: str = "diffcontext_state.json") -> Dict:
+    """Load previous function snapshot from disk."""
+    if not os.path.exists(path):
+        return {}
+    with open(path, "r") as f:
+        return json.load(f)