diff --git a/codebase_rag/constants/cli.py b/codebase_rag/constants/cli.py index 47dc874e48..33b10663bd 100644 --- a/codebase_rag/constants/cli.py +++ b/codebase_rag/constants/cli.py @@ -1101,6 +1101,14 @@ class DiffMarker: ) # context slice (issue #1536). CONTEXT_WHY_TARGET = "target" +# How many of the candidates left out of a context slice it names; the rest +# are only counted in `omitted_count` (issue #3243). +CONTEXT_OMITTED_SAMPLE = 10 +# How deep a context slice prints each piece: inside the slice and its +# `pieces` list, each level one MCP_JSON_INDENT more. +CONTEXT_PIECE_DEPTH = 2 +# Files whose decoded lines a context slice keeps between quotes of them. +CONTEXT_SOURCE_CACHE_FILES = 256 CONTEXT_WHY_CALLER = "direct caller: the call line" CONTEXT_WHY_CALLEE = "direct callee: its signature" CONTEXT_WHY_RETURNS = "type it returns" diff --git a/codebase_rag/context_slice.py b/codebase_rag/context_slice.py index e76796f5ba..29843944e5 100644 --- a/codebase_rag/context_slice.py +++ b/codebase_rag/context_slice.py @@ -16,8 +16,11 @@ from __future__ import annotations +import json import re +import textwrap from collections.abc import Callable, Iterable +from functools import lru_cache, partial from pathlib import Path from typing import NamedTuple, TypedDict @@ -27,7 +30,7 @@ from .graph_query import QueryFn, ReachIndex from .types_defs import ResultRow, SemanticSearchResult from .utils.path_utils import cached_resolve_posix -from .utils.source_extraction import extract_source_lines +from .utils.source_extraction import read_source_lines, slice_source_lines from .utils.token_utils import count_tokens Searcher = Callable[[str], list[SemanticSearchResult]] @@ -53,6 +56,8 @@ class Piece(TypedDict): span: list[int] why_included: str source: str + # What the piece costs against the budget as printed: its source and + # every field beside it (issue #3243). tokens: int @@ -62,7 +67,10 @@ class ContextSlice(TypedDict): budget_tokens: int used_tokens: int pieces: list[Piece] + # The best-ranked candidates that did not fit, at most + # CONTEXT_OMITTED_SAMPLE of them; `omitted_count` counts them all. omitted: list[str] + omitted_count: int truncated: bool @@ -75,6 +83,9 @@ class _Candidate(NamedTuple): span: tuple[int, int] why: str source: str + # Reads the source when the piece might still fit: a central symbol + # has thousands of callers and tests, nearly all of them left out. + loader: Callable[[], str] | None = None @property def order(self) -> tuple[int, int, float, str]: @@ -104,6 +115,15 @@ def _normalise(text: str) -> str: return text.replace("\r\n", "\n").replace("\r", "\n") +@lru_cache(maxsize=cs.CONTEXT_SOURCE_CACHE_FILES) +def _file_lines(path: Path, stamp: tuple[int, int, int, int]) -> tuple[str, ...]: + # Keyed by modification time and size as well as path, so an edited + # file is read again; a slice quotes one test file once per test in it. + # The inode and change time catch an edit or atomic replacement that + # keeps the old size and modification time. + return tuple(read_source_lines(path)) + + def _lines(repo_root: Path | None, path: str | None, start: int, end: int) -> str: if repo_root is None or not path or start < 1 or end < start: return "" @@ -113,7 +133,15 @@ def _lines(repo_root: Path | None, path: str | None, start: int, end: int) -> st candidate = (root / path).resolve() if not candidate.is_relative_to(root): return "" - return _normalise(extract_source_lines(candidate, start, end) or "") + try: + stat = candidate.stat() + lines = _file_lines( + candidate, + (stat.st_mtime_ns, stat.st_size, stat.st_ctime_ns, stat.st_ino), + ) + except (OSError, ValueError, LookupError): + return "" + return _normalise(slice_source_lines(lines, start, end, candidate) or "") def _line(repo_root: Path | None, path: str | None, line: int | None) -> str: @@ -210,7 +238,8 @@ def _caller_pieces( out: list[_Candidate] = [] for row in graph_query.callers(fetch_all, project, qn): line = row["line"] - source = _line(repo_root, row["path"], line) + source = "" + loader: Callable[[], str] | None = partial(_line, repo_root, row["path"], line) if line is None: definition = graph_query.definition( fetch_all, project, row["qualified_name"], repo_root @@ -219,6 +248,7 @@ def _caller_pieces( source = _header( repo_root, row["path"], definition["start_line"], definition["end_line"] ) + loader = None out.append( _Candidate( 1, @@ -229,6 +259,7 @@ def _caller_pieces( (line or 0, line or 0), cs.CONTEXT_WHY_CALLER, source, + loader, ) ) return out @@ -306,10 +337,9 @@ def _test_pieces( out: list[_Candidate] = [] reach = ReachIndex.build(fetch_all, project) for row in reach.tests_reaching(qn): - definition = graph_query.definition( - fetch_all, project, row["qualified_name"], repo_root - ) - start, end = definition["start_line"] or 0, definition["end_line"] or 0 + # The span is in the reach index already; the source is read only + # if the test can still fit, never one graph query per test. + start, end = reach.span(row["qualified_name"]) out.append( _Candidate( row["depth"] + 1, @@ -319,7 +349,8 @@ def _test_pieces( row["path"], (start, end), cs.CONTEXT_WHY_TEST.format(depth=row["depth"], through=row["through"]), - _normalise(definition["source"] or "").rstrip("\n"), + "", + partial(_lines, repo_root, row["path"], start, end), ) ) return out @@ -391,55 +422,155 @@ def _doc_section(repo_root: Path, doc_qn: str, row: ResultRow) -> _Candidate | N # --- the slice --------------------------------------------------------------------- -def _fit( - candidates: Iterable[_Candidate], budget: int -) -> tuple[list[Piece], list[str], int, bool]: +def _piece_tokens( + qualified_name: str, file: str | None, span: list[int], why: str, source: str +) -> int: + """What a piece costs as printed: its source and every field beside it. + + Counting the source alone let a slice of one-line call sites print ten + times its budget in names, paths and reasons (issue #3243). + """ + piece = Piece( + qualified_name=qualified_name, + file=file, + span=span, + why_included=why, + source=source, + tokens=0, + ) + # The piece prints with its own count where it was measured with 0, and + # a count of four digits takes more tokens than "0" does: measure again + # with the count in place until it stops growing. It is measured as + # the slice prints it, every line indented to its depth there: measured + # flush left, a 10,000 budget printed 10,478 tokens of pieces and fields. + margin = " " * (cs.CONTEXT_PIECE_DEPTH * cs.MCP_JSON_INDENT) + while ( + cost := count_tokens( + textwrap.indent(json.dumps(piece, indent=cs.MCP_JSON_INDENT), margin) + ) + ) > piece["tokens"]: + piece["tokens"] = cost + return piece["tokens"] + + +class _Fitted(NamedTuple): + pieces: list[Piece] + omitted: list[str] + omitted_count: int + used: int + truncated: bool + + +class _Omitted: + # The best-ranked candidates that did not fit: a sample to print, and + # the count of all of them. + __slots__ = ("count", "sample") + + def __init__(self) -> None: + self.sample: list[str] = [] + self.count = 0 + + def add(self, candidate: _Candidate) -> None: + self.count += 1 + if len(self.sample) < cs.CONTEXT_OMITTED_SAMPLE: + self.sample.append(f"{candidate.qualified_name} ({candidate.why})") + + +def _fit(candidates: Iterable[_Candidate], budget: int) -> _Fitted: pieces: list[Piece] = [] - omitted: list[str] = [] + omitted = _Omitted() used = 0 truncated = False seen: set[tuple[str, str, tuple[int, int]]] = set() for candidate in sorted(candidates, key=lambda c: c.order): key = (candidate.qualified_name, candidate.why, candidate.span) - if key in seen or not candidate.source: + if key in seen: continue seen.add(key) - source = candidate.source - tokens = count_tokens(source) - if used + tokens > budget: - if candidate.distance == 0 and not pieces: - # The target itself must fit: keep as many of its lines as - # the budget allows and say so. A budget too small for even - # its first line leaves it out rather than padding with - # nothing. - source, tokens = _trim(source, budget) - truncated = True - if not source: - omitted.append(f"{candidate.qualified_name} ({candidate.why})") - continue - else: - omitted.append(f"{candidate.qualified_name} ({candidate.why})") - continue + span = [candidate.span[0], candidate.span[1]] + is_target = candidate.distance == 0 and not pieces + if (sized := _sized(candidate, span, budget - used, is_target)) is None: + omitted.add(candidate) + continue + source, tokens, trimmed = sized + if trimmed: + truncated = True + if not source: + # No source to quote, or a target too big for even its first + # line: only the latter is a candidate the budget left out. + if trimmed: + omitted.add(candidate) + continue used += tokens pieces.append( Piece( qualified_name=candidate.qualified_name, file=candidate.file, - span=[candidate.span[0], candidate.span[1]], + span=span, why_included=candidate.why, source=source, tokens=tokens, ) ) - return pieces, omitted, used, truncated + return _Fitted(pieces, omitted.sample, omitted.count, used, truncated) + + +def _sized( + candidate: _Candidate, span: list[int], room: int, is_target: bool +) -> tuple[str, int, bool] | None: + # (source, tokens, trimmed) of the candidate within `room` tokens, or + # None when it does not fit. Its fields alone are a floor on what it + # costs: past it, the source is never read. + floor = _piece_tokens( + candidate.qualified_name, candidate.file, span, candidate.why, "" + ) + if floor > room and not is_target: + return None + source = candidate.source or (candidate.loader() if candidate.loader else "") + if not source: + return "", 0, False + tokens = _piece_tokens( + candidate.qualified_name, candidate.file, span, candidate.why, source + ) + if tokens <= room: + return source, tokens, False + if not is_target: + return None + # The target itself must fit: keep as many of its lines as the budget + # allows and say so. A budget too small for even its first line leaves + # it out rather than padding with nothing. + source, tokens = _trim(candidate, span, source, room) + return source, tokens, True + +def _trim( + candidate: _Candidate, span: list[int], source: str, room: int +) -> tuple[str, int]: + """The longest run of the source's leading lines whose piece fits `room`. -def _trim(source: str, budget: int) -> tuple[str, int]: + Bisected on the line count: a few dozen encodes where trimming a line + at a time re-encoded the whole remaining text once per line, minutes + for a class of thousands of lines (issue #3243). + """ lines = source.split("\n") - while lines and count_tokens("\n".join(lines)) > budget: - lines.pop() - text = "\n".join(lines) - return text, count_tokens(text) if text else 0 + + def cost(count: int) -> int: + return _piece_tokens( + candidate.qualified_name, + candidate.file, + span, + candidate.why, + "\n".join(lines[:count]), + ) + + low, high = 0, len(lines) + while low < high: + middle = (low + high + 1) // 2 + if cost(middle) <= room: + low = middle + else: + high = middle - 1 + return "\n".join(lines[:low]), cost(low) def context( @@ -461,6 +592,7 @@ def context( used_tokens=0, pieces=[], omitted=[], + omitted_count=0, truncated=False, ) definition = graph_query.definition(fetch_all, project, qn, repo_root) @@ -472,6 +604,7 @@ def context( used_tokens=0, pieces=[], omitted=[], + omitted_count=0, truncated=False, ) hotness = _hotness(fetch_all, project, qn) @@ -488,13 +621,14 @@ def context( candidates = [ c._replace(similarity=scores.get(c.qualified_name, 0.0)) for c in candidates ] - pieces, omitted, used, truncated = _fit(candidates, budget) + fitted = _fit(candidates, budget) return ContextSlice( target=target, resolved=qn, budget_tokens=budget, - used_tokens=used, - pieces=pieces, - omitted=omitted, - truncated=truncated, + used_tokens=fitted.used, + pieces=fitted.pieces, + omitted=fitted.omitted, + omitted_count=fitted.omitted_count, + truncated=fitted.truncated, ) diff --git a/codebase_rag/cypher_queries.py b/codebase_rag/cypher_queries.py index 4560b0d1bd..5646c0fc6a 100644 --- a/codebase_rag/cypher_queries.py +++ b/codebase_rag/cypher_queries.py @@ -902,6 +902,7 @@ def _any_label(var: str, labels: str) -> str: RETURN labels(n)[0] AS label, n.qualified_name AS qualified_name, n.name AS name, n.path AS path, n.namespace AS namespace, n.start_line AS start_line, n.end_line AS end_line, + n.decorated_start_line AS decorated_start_line, n.decorators AS decorators, n.is_exported AS is_exported, n.overrides_external AS overrides_external, n.is_program_main AS is_program_main, diff --git a/codebase_rag/dead_code.py b/codebase_rag/dead_code.py index d9ff02021b..a41c084b3a 100644 --- a/codebase_rag/dead_code.py +++ b/codebase_rag/dead_code.py @@ -1408,6 +1408,9 @@ def _node_props(row: ResultRow) -> PropertyDict: # non-int values coalesce to 0 so the containment checks skip them. cs.KEY_START_LINE: _as_line(row.get(cs.KEY_START_LINE)), cs.KEY_END_LINE: _as_line(row.get(cs.KEY_END_LINE)), + # The first decorator's line: a reaching test is quoted from it, as + # `definition()` quotes a decorated def (issue #2428). + cs.KEY_DECORATED_START_LINE: _as_line(row.get(cs.KEY_DECORATED_START_LINE)), cs.KEY_RUST_CFG_TEST_MODS: _as_str_list(row.get(cs.KEY_RUST_CFG_TEST_MODS)), cs.KEY_RUST_UNGATED_MODS: _as_str_list(row.get(cs.KEY_RUST_UNGATED_MODS)), # The signature the C# entry-point rule checks (issue #2471). diff --git a/codebase_rag/graph_query.py b/codebase_rag/graph_query.py index 946a1a5de6..e0cd3b4c50 100644 --- a/codebase_rag/graph_query.py +++ b/codebase_rag/graph_query.py @@ -562,7 +562,9 @@ def source_root_for( return local if Path(stored).resolve() == local else None -def _decorated_start(row: ResultRow, own_start: int | None) -> int | None: +def _decorated_start( + row: ResultRow | PropertyDict, own_start: int | None +) -> int | None: """The first decorator's line when the node records one above its own start, else that start (issue #2428). @@ -1174,6 +1176,20 @@ def _walk(self, qualified_name: str) -> tuple[dict[str, int], dict[str, str]]: frontier = next_frontier return depth_of, through_of + def span(self, qualified_name: str) -> tuple[int, int]: + """The recorded (start, end) lines of a node, 0 where unknown; the + start is the first decorator's line, as `definition()` answers.""" + entry = self._by_qn.get(qualified_name) + if entry is None: + return 0, 0 + props = entry[1] + own, end = props.get(cs.KEY_START_LINE), props.get(cs.KEY_END_LINE) + start = _decorated_start(props, own if isinstance(own, int) else None) + return ( + start if isinstance(start, int) else 0, + end if isinstance(end, int) else 0, + ) + def tests_reaching(self, qualified_name: str) -> list[TestReachRow]: depth_of, through_of = self._walk(qualified_name) out: list[TestReachRow] = [] diff --git a/codebase_rag/tests/test_context_budget_output.py b/codebase_rag/tests/test_context_budget_output.py new file mode 100644 index 0000000000..f835f8eac9 --- /dev/null +++ b/codebase_rag/tests/test_context_budget_output.py @@ -0,0 +1,361 @@ +"""The context slice's printed output stays within its token budget, fast. + +`context` counted only each piece's source against the budget, listed every +candidate that did not fit in `omitted`, fetched every reaching test's +definition before budgeting, and trimmed an oversized target one line and +one full re-encode at a time. A 2-line function reached by 2,000 tests +printed 163,019 tokens for a 4,000 budget after 21 s, and a 9,000-line class +took 87 s to trim (issue #3243). +""" + +from __future__ import annotations + +import json +import os +import re +import textwrap +import time +from pathlib import Path + +import pytest + +from codebase_rag import constants as cs +from codebase_rag import context_slice +from codebase_rag import cypher_queries as cq +from codebase_rag.context_slice import context +from codebase_rag.graph_query import DefinitionRow, QueryFn +from codebase_rag.graph_updater import GraphUpdater +from codebase_rag.parser_loader import load_parsers +from codebase_rag.types_defs import PropertyParams, ResultRow +from codebase_rag.utils.token_utils import count_tokens +from evals.cgr_graph import _StatefulIngestor + +PROJECT = "ctxb" +TEST_FILES = 6 +TESTS_PER_FILE = 50 +# The slice's own fields beside its pieces: target, resolved, the counts and +# a capped sample of what was left out. +ENVELOPE = 400 + + +def _repo(root: Path, files: dict[str, str]) -> _StatefulIngestor: + for rel, text in files.items(): + (root / rel).parent.mkdir(parents=True, exist_ok=True) + (root / rel).write_text(text, encoding="utf-8") + store = _StatefulIngestor() + parsers, queries = load_parsers() + GraphUpdater( + ingestor=store, + repo_path=root, + parsers=parsers, + queries=queries, + project_name=PROJECT, + ).run(force=True) + return store + + +def _central(root: Path) -> _StatefulIngestor: + files = {"app.py": "def core(x):\n return x + 1\n", "tests/__init__.py": ""} + for f in range(TEST_FILES): + body = "from app import core\n\n\n" + "".join( + f"def test_{f}_{i}():\n assert core({i}) == {i + 1}\n\n\n" + for i in range(TESTS_PER_FILE) + ) + files[f"tests/test_{f}.py"] = body + return _repo(root, files) + + +def _printed(slice_: context_slice.ContextSlice) -> int: + return count_tokens(json.dumps(slice_, indent=cs.MCP_JSON_INDENT)) + + +@pytest.mark.parametrize("budget", [20_000, 10_000, 4000, 1000, 300]) +def test_the_printed_slice_fits_the_budget(temp_repo: Path, budget: int) -> None: + store = _central(temp_repo) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", budget, temp_repo) + + total = TEST_FILES * TESTS_PER_FILE + assert slice_["used_tokens"] <= budget + assert slice_["used_tokens"] == sum(p["tokens"] for p in slice_["pieces"]) + assert _printed(slice_) <= budget + ENVELOPE, _printed(slice_) + # What was left out is counted, and named only in a short sample. + assert len(slice_["omitted"]) <= cs.CONTEXT_OMITTED_SAMPLE + assert slice_["omitted_count"] > 0 + assert slice_["omitted_count"] + len(slice_["pieces"]) >= total + + +def test_a_piece_is_charged_for_its_fields(temp_repo: Path) -> None: + store = _central(temp_repo) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + for piece in slice_["pieces"]: + assert piece["tokens"] > count_tokens(piece["source"]), piece + + +def test_each_quoted_file_is_read_once_and_no_test_is_fetched_by_name( + temp_repo: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + store = _central(temp_repo) + reads: list[Path] = [] + real = context_slice.read_source_lines + + def counting(path: Path) -> list[str]: + reads.append(path) + return real(path) + + definitions: list[str] = [] + real_definition = context_slice.graph_query.definition + + def counting_definition( + fetch_all: QueryFn, project: str, qualified_name: str, root: Path | None + ) -> DefinitionRow: + definitions.append(qualified_name) + return real_definition(fetch_all, project, qualified_name, root) + + monkeypatch.setattr(context_slice, "read_source_lines", counting) + monkeypatch.setattr(context_slice.graph_query, "definition", counting_definition) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + # 300 tests and 300 call sites in six files reach `core`: each file is + # decoded once, and only the target's definition comes from the graph. + assert len(reads) == len(set(reads)) <= TEST_FILES, reads + assert definitions == [f"{PROJECT}.app.core"] + assert slice_["pieces"] + + +def test_nothing_is_read_once_the_budget_is_spent( + temp_repo: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # A budget the target fills leaves every caller and test unread. + store = _central(temp_repo) + full = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + target_tokens = full["pieces"][0]["tokens"] + loads: list[tuple[str | None, int]] = [] + real = context_slice._lines + + def counting(root: Path | None, path: str | None, start: int, end: int) -> str: + loads.append((path, start)) + return real(root, path, start, end) + + monkeypatch.setattr(context_slice, "_lines", counting) + + slice_ = context( + store.fetch_all, PROJECT, f"{PROJECT}.app.core", target_tokens, temp_repo + ) + + assert [p["why_included"] for p in slice_["pieces"]] == [cs.CONTEXT_WHY_TARGET] + assert loads == [] + assert slice_["omitted_count"] == 2 * TEST_FILES * TESTS_PER_FILE + + +def test_a_large_target_trims_without_re_encoding_every_line( + temp_repo: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(1500) + ) + store = _repo(temp_repo, {"big.py": body}) + calls = 0 + real = context_slice.count_tokens + + def counting(text: str) -> int: + nonlocal calls + calls += 1 + return real(text) + + monkeypatch.setattr(context_slice, "count_tokens", counting) + started = time.perf_counter() + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 500, temp_repo) + + elapsed = time.perf_counter() - started + assert slice_["truncated"] + (target,) = slice_["pieces"] + assert target["source"].startswith("class Big:\n def m0(self, x):") + assert slice_["used_tokens"] <= 500 + # Bisecting the line count: a few dozen encodes, not one per line. + assert calls < 100, calls + assert elapsed < 5, elapsed + + +def test_the_trim_keeps_the_most_whole_lines_that_fit(temp_repo: Path) -> None: + # The bisected trim keeps exactly the lines a line-by-line trim kept. + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(60) + ) + store = _repo(temp_repo, {"big.py": body}) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 300, temp_repo) + + (target,) = slice_["pieces"] + lines = body.rstrip("\n").split("\n") + kept = target["source"].split("\n") + assert kept == lines[: len(kept)] + one_more = "\n".join(lines[: len(kept) + 1]) + assert ( + context_slice._piece_tokens( + f"{PROJECT}.big.Big", + "big.py", + [1, len(lines)], + cs.CONTEXT_WHY_TARGET, + one_more, + ) + > 300 + ) + + +def test_a_budget_that_exactly_fits_some_lines_keeps_them(temp_repo: Path) -> None: + # A prefix whose piece costs exactly the budget fits: the trim's bound + # is inclusive. + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(60) + ) + store = _repo(temp_repo, {"big.py": body}) + lines = body.rstrip("\n").split("\n") + full = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 10_000, temp_repo) + span = full["pieces"][0]["span"] + budget = context_slice._piece_tokens( + f"{PROJECT}.big.Big", + "big.py", + span, + cs.CONTEXT_WHY_TARGET, + "\n".join(lines[:7]), + ) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", budget, temp_repo) + + assert slice_["pieces"][0]["source"] == "\n".join(lines[:7]) + assert slice_["used_tokens"] == budget + + +def test_a_neighbourhood_that_fits_omits_nothing(temp_repo: Path) -> None: + # Negative: a small slice is whole, with nothing counted as left out. + store = _repo( + temp_repo, + { + "app.py": "def core(x):\n return x + 1\n\n\ndef use():\n return core(1)\n", + }, + ) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + assert slice_["omitted"] == [] + assert slice_["omitted_count"] == 0 + assert {p["why_included"] for p in slice_["pieces"]} == { + cs.CONTEXT_WHY_TARGET, + cs.CONTEXT_WHY_CALLER, + } + + +def test_a_piece_is_charged_for_the_count_it_prints(temp_repo: Path) -> None: + # A piece was measured with `tokens: 0` and printed with its own count, + # and a count of four digits prints in more tokens than "0" does. + body = "class Big:\n" + "".join( + f" def m{i}(self, x):\n return x * {i}\n\n" for i in range(150) + ) + store = _repo(temp_repo, {"big.py": body}) + + slice_ = context(store.fetch_all, PROJECT, f"{PROJECT}.big.Big", 10_000, temp_repo) + + (target,) = slice_["pieces"] + assert target["tokens"] >= 1000 + # As the slice prints it: two levels deep, inside the slice and its + # `pieces` list, every line indented by both levels. + printed = textwrap.indent( + json.dumps(target, indent=cs.MCP_JSON_INDENT), " " * 2 * cs.MCP_JSON_INDENT + ) + assert printed in json.dumps(slice_, indent=cs.MCP_JSON_INDENT) + assert count_tokens(printed) == target["tokens"] + + +def test_a_replaced_file_is_quoted_afresh(temp_repo: Path) -> None: + # An atomic replacement that keeps the old size and modification time + # is still a different file: its new lines are quoted, not the cached. + store = _repo( + temp_repo, + { + "app.py": "def core(x):\n return x + 1\n\n\ndef use():\n return core(1)\n", + }, + ) + app = temp_repo / "app.py" + before = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + old = app.stat() + replacement = temp_repo / "app.py.new" + replacement.write_text( + "def core(x):\n return x + 1\n\n\ndef use():\n return core(2)\n", + encoding="utf-8", + ) + os.utime(replacement, ns=(old.st_atime_ns, old.st_mtime_ns)) + os.replace(replacement, app) + assert (app.stat().st_mtime_ns, app.stat().st_size) == ( + old.st_mtime_ns, + old.st_size, + ) + + after = context(store.fetch_all, PROJECT, f"{PROJECT}.app.core", 4000, temp_repo) + + def caller_sources(slice_: context_slice.ContextSlice) -> list[str]: + return [ + p["source"] + for p in slice_["pieces"] + if p["why_included"] == cs.CONTEXT_WHY_CALLER + ] + + assert caller_sources(before) == ["return core(1)"] + assert caller_sources(after) == ["return core(2)"] + + +_COLUMN = re.compile(r"\bn\.(\w+) AS (\w+)") + + +def _as_memgraph_answers(store: _StatefulIngestor) -> QueryFn: + # The stand-in answers the dead-code node query with a fixed key set; + # Memgraph answers with every column its RETURN names. + def fetch(query: str, params: PropertyParams | None) -> list[ResultRow]: + rows = store.fetch_all(query, dict(params) if params is not None else None) + if query != cq.CYPHER_DEAD_CODE_NODES: + return rows + columns = _COLUMN.findall(query) + out: list[ResultRow] = [] + for row in rows: + node = (str(row[cs.KEY_LABEL]), str(row[cs.KEY_QUALIFIED_NAME])) + props = store.nodes[node] + out.append( + {cs.KEY_LABEL: row[cs.KEY_LABEL]} + | {alias: props.get(prop) for prop, alias in columns} + ) + return out + + return fetch + + +def test_a_decorated_test_is_quoted_with_its_decorators(temp_repo: Path) -> None: + # The reach index read the def line alone, so a decorated test lost its + # decorators; it is quoted from the first one, as `definition()` quotes + # a decorated def (issue #2428). + store = _repo( + temp_repo, + { + "app.py": "def core(x):\n return x + 1\n", + "tests/__init__.py": "", + "tests/test_app.py": ( + "import pytest\n\nfrom app import core\n\n\n" + "@pytest.mark.slow\n" + "def test_core():\n assert core(1) == 2\n" + ), + }, + ) + + slice_ = context( + _as_memgraph_answers(store), PROJECT, f"{PROJECT}.app.core", 4000, temp_repo + ) + + why = cs.CONTEXT_WHY_TEST.format(depth=1, through=f"{PROJECT}.app.core") + (test,) = [p for p in slice_["pieces"] if p["why_included"] == why] + assert test["span"] == [6, 8] + assert test["source"] == ( + "@pytest.mark.slow\ndef test_core():\n assert core(1) == 2" + ) diff --git a/codebase_rag/tests/test_context_slice.py b/codebase_rag/tests/test_context_slice.py index 7b7c71cd93..85bf54bf2f 100644 --- a/codebase_rag/tests/test_context_slice.py +++ b/codebase_rag/tests/test_context_slice.py @@ -10,7 +10,7 @@ import pytest from codebase_rag import constants as cs -from codebase_rag.context_slice import _doc_pieces, context +from codebase_rag.context_slice import _doc_pieces, _piece_tokens, context from codebase_rag.graph_updater import GraphUpdater from codebase_rag.parser_loader import load_parsers from codebase_rag.utils.token_utils import count_tokens @@ -161,12 +161,22 @@ def test_token_count_never_exceeds_the_budget( ) -> None: root, store, _updater = repo full = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), 4000, root) - for budget in (full["used_tokens"] - 1, 60, 25, 5): + # A piece's fields count against the budget too (issue #3243): the + # target with an empty source is the floor any slice of it starts at. + target = full["pieces"][0] + floor = _piece_tokens( + target["qualified_name"], + target["file"], + target["span"], + target["why_included"], + "", + ) + for budget in (full["used_tokens"] - 1, floor + 40, floor + 15, 5): slice_ = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), budget, root) assert slice_["used_tokens"] <= budget assert sum(count_tokens(p["source"]) for p in slice_["pieces"]) <= budget assert slice_["omitted"] or slice_["truncated"] - shrunk = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), 25, root) + shrunk = context(store.fetch_all, PROJECT, _qn("pkg.util.helper"), floor + 15, root) # The target keeps as many of its lines as fit and comes first. assert shrunk["truncated"] assert shrunk["pieces"][0]["why_included"] == cs.CONTEXT_WHY_TARGET diff --git a/codebase_rag/tests/test_context_slice_normalisation.py b/codebase_rag/tests/test_context_slice_normalisation.py index eb861a4f00..683dad5e05 100644 --- a/codebase_rag/tests/test_context_slice_normalisation.py +++ b/codebase_rag/tests/test_context_slice_normalisation.py @@ -112,7 +112,6 @@ def test_test_pieces_normalises_the_source_it_reads( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: from codebase_rag import context_slice as cs_mod - from codebase_rag.utils.source_extraction import extract_source_lines (tmp_path / "test_app.py").write_bytes( b"def test_run():\r\n assert run() == 6\r\n" @@ -134,26 +133,20 @@ def tests_reaching(_qn: str) -> list[dict[str, object]]: } ] - def _definition( - _fetch: object, _project: str, _qn: str, _root: Path | None - ) -> dict[str, object]: - return { - "qualified_name": "p.test_app.test_run", - "path": "test_app.py", - "start_line": 1, - "end_line": 2, - "source": extract_source_lines(tmp_path / "test_app.py", 1, 2), - } + @staticmethod + def span(_qn: str) -> tuple[int, int]: + # The reach index carries each test's span, so the source is + # read from disk without a definition query (issue #3243). + return 1, 2 monkeypatch.setattr(cs_mod, "ReachIndex", _Reach) - monkeypatch.setattr(cs_mod.graph_query, "definition", _definition) pieces = cs_mod._test_pieces(lambda *a, **k: [], "p", "p.app.run", tmp_path) assert len(pieces) == 1, repr(pieces) - assert pieces[0].source == "def test_run():\n assert run() == 6", repr( - pieces[0].source - ) + loader = pieces[0].loader + assert loader is not None + assert loader() == "def test_run():\n assert run() == 6", repr(loader()) # The indexer writes absolute_path through `cached_resolve_posix`, so the lookup diff --git a/codebase_rag/utils/source_extraction.py b/codebase_rag/utils/source_extraction.py index 21e746ec27..04f9c0ed5e 100644 --- a/codebase_rag/utils/source_extraction.py +++ b/codebase_rag/utils/source_extraction.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Sequence from pathlib import Path from loguru import logger @@ -11,6 +11,46 @@ from .source_encoding import decode_python_source, decode_wide_source +def read_source_lines(file_path: Path, encoding: str = ENCODING_UTF8) -> list[str]: + """The file's lines as the indexer read them, line endings kept.""" + raw_bytes = file_path.read_bytes() + # The indexer parsed a Python source in the encoding it declares, and + # the lines it recorded are lines of that text (issue #2445). + # Any other UTF-16/32 source with a BOM was parsed as its decoded text + # too (issue #3153). + declared = ( + decode_python_source(raw_bytes, file_path) + if file_path.suffix in PY_EXTENSIONS + else decode_wide_source(raw_bytes, file_path) + ) + text = raw_bytes.decode(encoding) if declared is None else declared + return text.splitlines(keepends=True) + + +def slice_source_lines( + lines: Sequence[str], start_line: int, end_line: int, file_path: Path +) -> str | None: + """Lines `start_line`..`end_line` (1-based) of an already-read file.""" + if not lines: + return None + + if start_line > len(lines) or end_line > len(lines): + logger.warning( + ls.SOURCE_RANGE_EXCEEDS.format( + start=start_line, + end=end_line, + length=len(lines), + path=file_path, + ) + ) + end_line = min(end_line, len(lines)) + if start_line > len(lines): + return None + + extracted_lines = lines[start_line - 1 : end_line] + return "".join(extracted_lines).strip() + + def extract_source_lines( file_path: Path, start_line: int, end_line: int, encoding: str = ENCODING_UTF8 ) -> str | None: @@ -26,38 +66,9 @@ def extract_source_lines( return _notebook_source_lines(file_path, start_line, end_line) try: - raw_bytes = file_path.read_bytes() - # The indexer parsed a Python source in the encoding it declares, and - # the lines it recorded are lines of that text (issue #2445). - # Any other UTF-16/32 source with a BOM was parsed as its decoded text - # too (issue #3153). - declared = ( - decode_python_source(raw_bytes, file_path) - if file_path.suffix in PY_EXTENSIONS - else decode_wide_source(raw_bytes, file_path) + return slice_source_lines( + read_source_lines(file_path, encoding), start_line, end_line, file_path ) - text = raw_bytes.decode(encoding) if declared is None else declared - lines = text.splitlines(keepends=True) - - if not lines: - return None - - if start_line > len(lines) or end_line > len(lines): - logger.warning( - ls.SOURCE_RANGE_EXCEEDS.format( - start=start_line, - end=end_line, - length=len(lines), - path=file_path, - ) - ) - end_line = min(end_line, len(lines)) - if start_line > len(lines): - return None - - extracted_lines = lines[start_line - 1 : end_line] - return "".join(extracted_lines).strip() - except Exception as e: logger.warning(ls.SOURCE_EXTRACT_FAILED.format(path=file_path, error=e)) return None diff --git a/docs/architecture/context-slice.md b/docs/architecture/context-slice.md index 081e693343..c115f3b051 100644 --- a/docs/architecture/context-slice.md +++ b/docs/architecture/context-slice.md @@ -46,11 +46,24 @@ Candidates are ordered by graph distance, then trace hotness (`dynamic_call_count` on the caller's edge, from `cgr trace ingest`), then embedding similarity for a free-text task, then name. Tokens are counted with the same encoder the query tools use, and pieces are taken in order -while they fit: the total never exceeds `budget_tokens`. The target is -special: when it alone exceeds the budget it is trimmed line by line to -fit and the slice is marked `truncated`; a budget too small for even its -first line names it in `omitted` instead. Every piece that did not fit is -listed in `omitted` with the reason it would have been included. +while they fit: the total never exceeds `budget_tokens`. A piece is +charged for what it prints, its source and the fields beside it (name, +file, span, reason) at the indentation the slice prints them with, so a +slice of one-line call sites cannot print many times its budget in +metadata (issue #3243). The target is special: when it +alone exceeds the budget it is trimmed to the most leading lines that fit +and the slice is marked `truncated`; a budget too small for even its first +line names it in `omitted` instead. `omitted_count` counts every +candidate that did not fit, and `omitted` names the best-ranked ten of +them with the reason each would have been included; beside the pieces, +the printed slice adds only those fields. + +A central symbol has thousands of callers and tests, nearly all of which a +budget leaves out. Their sources are read only for a candidate whose +fields still fit, each file the slice quotes is decoded once, and a test's +span comes from the reach index already built, never from a graph query +per test. That span starts at a decorated test's first decorator, as +`definition` quotes it. ## Output @@ -59,14 +72,15 @@ listed in `omitted` with the reason it would have been included. "target": "myproj.pkg.util.helper", "resolved": "myproj.pkg.util.helper", "budget_tokens": 4000, - "used_tokens": 312, + "used_tokens": 398, "pieces": [ {"qualified_name": "myproj.pkg.util.helper", "file": "pkg/util.py", "span": [8, 13], - "why_included": "target", "source": "def helper(items: list[int]) -> Total:\n ...", "tokens": 61}, + "why_included": "target", "source": "def helper(items: list[int]) -> Total:\n ...", "tokens": 104}, {"qualified_name": "myproj.pkg.app.run", "file": "pkg/app.py", "span": [5, 5], - "why_included": "direct caller: the call line", "source": "return helper([1, 2]).value", "tokens": 9} + "why_included": "direct caller: the call line", "source": "return helper([1, 2]).value", "tokens": 52} ], "omitted": [], + "omitted_count": 0, "truncated": false } ```