diff --git a/graphify/__main__.py b/graphify/__main__.py index e122a5b6e..09e0e9755 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -4258,6 +4258,7 @@ def _parse_float(name: str, raw: str) -> float: # AST extraction on code files. Empty code list (docs-only corpus) is # the issue #698 case — skip cleanly instead of crashing inside extract(). ast_result: dict = {"nodes": [], "edges": [], "input_tokens": 0, "output_tokens": 0} + _stitch_new_ids: set[str] = set() if code_files: from graphify.extract import extract as _ast_extract # Anchor the cache at the output root, not the scanned project: @@ -4272,6 +4273,9 @@ def _parse_float(name: str, raw: str) -> float: except Exception as exc: print(f"[graphify extract] AST extraction failed: {exc}", file=sys.stderr) ast_result = {"nodes": [], "edges": [], "input_tokens": 0, "output_tokens": 0} + for _n in ast_result.get("nodes", []): + if _n.get("id"): + _stitch_new_ids.add(str(_n["id"])) # Semantic extraction on docs/papers/images. Check cache first. from graphify.cache import ( @@ -4284,6 +4288,7 @@ def _parse_float(name: str, raw: str) -> float: } sem_cache_hits = 0 sem_cache_misses = 0 + fresh: dict = {"nodes": [], "edges": [], "hyperedges": []} if semantic_files: sem_paths_str = [str(p) for p in semantic_files] cached_nodes, cached_edges, cached_hyperedges, uncached_paths = ( @@ -4352,6 +4357,32 @@ def _progress(idx: int, total: int, _result: dict) -> None: file=sys.stderr, ) sys.exit(1) + # Incremental: never prune/replace changed files unless re-extraction + # actually produced nodes/edges. Invalid JSON and connection failures + # that yield empty chunks must abort before merge (#incremental-safe). + if incremental_mode and uncached_paths: + from graphify.build import paths_missing_from_extraction as _paths_missing + + _missing = _paths_missing(fresh, uncached_paths, target) + if fresh.get("failed_chunks", 0) > 0 or _missing: + if fresh.get("failed_chunks", 0) > 0: + print( + f"[graphify extract] error: {fresh['failed_chunks']} semantic " + f"chunk(s) failed during incremental re-extraction.", + file=sys.stderr, + ) + for _path in _missing: + print( + f"[graphify extract] error: re-extraction produced no " + f"nodes/edges for {_path}", + file=sys.stderr, + ) + print( + "[graphify extract] incremental update aborted — existing " + "graph.json and manifest unchanged.", + file=sys.stderr, + ) + sys.exit(1) try: _save_semantic_cache( fresh.get("nodes", []), @@ -4366,6 +4397,9 @@ def _progress(idx: int, total: int, _result: dict) -> None: sem_result["hyperedges"].extend(fresh.get("hyperedges", [])) sem_result["input_tokens"] += fresh.get("input_tokens", 0) sem_result["output_tokens"] += fresh.get("output_tokens", 0) + for _n in fresh.get("nodes", []): + if _n.get("id"): + _stitch_new_ids.add(str(_n["id"])) pg_result: dict = {"nodes": [], "edges": []} if cli_postgres_dsn is not None: @@ -4422,6 +4456,21 @@ def _progress(idx: int, total: int, _result: dict) -> None: for ftype, flist in files_by_type.items() } + _incremental_prune: list[str] | None = None + _changed_for_stitch: list[str] = [] + if incremental_mode: + from graphify.build import path_covered_by_extraction as _path_covered + + _changed_code = [str(p) for p in code_files] + _changed_sem = [ + str(p) for p in semantic_files + if _path_covered(str(p), sem_result, target) + ] + _changed_for_stitch = list(dict.fromkeys(_changed_code + [str(p) for p in semantic_files])) + _incremental_prune = ( + list(dict.fromkeys(deleted_files + _changed_code + _changed_sem)) or None + ) + if no_cluster: # --no-cluster: dump the raw merged extraction as graph.json. # No NetworkX, no community detection, no analysis sidecar. @@ -4431,18 +4480,56 @@ def _progress(idx: int, total: int, _result: dict) -> None: # anchors emitted per importing file, #1327). from graphify.build import dedupe_edges as _dedupe_edges, dedupe_nodes as _dedupe_nodes from graphify.export import backup_if_protected as _backup - merged["nodes"] = _dedupe_nodes(merged["nodes"]) - merged["edges"] = _dedupe_edges(merged["edges"]) + from graphify.build import build_merge as _build_merge _backup(graphify_out) - graph_json_path.write_text( - json.dumps(merged, indent=2), encoding="utf-8" - ) + if incremental_mode: + G = _build_merge( + [merged], + graph_path=existing_graph_path, + prune_sources=_incremental_prune, + dedup=True, + dedup_llm_backend=backend if dedup_llm else None, + root=target, + ) + if _changed_for_stitch: + from graphify.stitch import stitch_incremental_links as _stitch_links + _stitch_links( + G, + _changed_for_stitch, + root=target, + new_node_ids=_stitch_new_ids, + ) + out_graph = { + "nodes": [{"id": n, **d} for n, d in G.nodes(data=True)], + "edges": [ + { + **{k: val for k, val in d.items() if k not in ("_src", "_tgt", "source", "target")}, + "source": d.get("_src", u), + "target": d.get("_tgt", v), + } + for u, v, d in G.edges(data=True) + ], + "hyperedges": list(G.graph.get("hyperedges", [])), + "input_tokens": merged["input_tokens"], + "output_tokens": merged["output_tokens"], + } + graph_json_path.write_text( + json.dumps(out_graph, indent=2), encoding="utf-8" + ) + node_count = len(out_graph["nodes"]) + edge_count = len(out_graph["edges"]) + else: + graph_json_path.write_text( + json.dumps(merged, indent=2), encoding="utf-8" + ) + node_count = len(merged["nodes"]) + edge_count = len(merged["edges"]) cost = _estimate_cost( backend, merged["input_tokens"], merged["output_tokens"] ) print( f"[graphify extract] wrote {graph_json_path} — " - f"{len(merged['nodes'])} nodes, {len(merged['edges'])} edges " + f"{node_count} nodes, {edge_count} edges " f"(no clustering)" ) if merged["input_tokens"] or merged["output_tokens"]: @@ -4484,11 +4571,19 @@ def _progress(idx: int, total: int, _result: dict) -> None: G = _build_merge( [merged], graph_path=existing_graph_path, - prune_sources=deleted_files or None, + prune_sources=_incremental_prune, dedup=True, dedup_llm_backend=dedup_backend, root=target, ) + if _changed_for_stitch: + from graphify.stitch import stitch_incremental_links as _stitch_links + _stitch_links( + G, + _changed_for_stitch, + root=target, + new_node_ids=_stitch_new_ids, + ) else: G = _build([merged], dedup=True, dedup_llm_backend=dedup_backend, root=target) if G.number_of_nodes() == 0: diff --git a/graphify/build.py b/graphify/build.py index 05e005dac..b5db491f6 100644 --- a/graphify/build.py +++ b/graphify/build.py @@ -532,6 +532,55 @@ def build_merge( return G +def source_path_aliases(p: str, root: str | Path | None = None) -> set[str]: + """All normalized forms of a scan/manifest path for matching source_file.""" + aliases: set[str] = set() + if not p: + return aliases + root_resolved = Path(root).resolve() if root else None + root_str = str(root_resolved) if root_resolved else None + aliases.add(p.replace("\\", "/")) + norm = _norm_source_file(p, root_str) + if norm: + aliases.add(norm) + if root_resolved is not None: + path = Path(p) + try: + abs_p = path if path.is_absolute() else root_resolved / path + aliases.add(abs_p.resolve().relative_to(root_resolved).as_posix()) + except ValueError: + pass + return aliases + + +def source_files_in_extraction(extraction: dict, root: str | Path | None = None) -> set[str]: + """Collect all source_file aliases present in an extraction dict.""" + found: set[str] = set() + for key in ("nodes", "edges", "hyperedges"): + for item in extraction.get(key, []): + sf = item.get("source_file") + if sf: + found |= source_path_aliases(str(sf), root) + return found + + +def path_covered_by_extraction(p: str, extraction: dict, root: str | Path | None = None) -> bool: + """True when extraction contains at least one node/edge/hyperedge for path p.""" + return bool(source_path_aliases(p, root) & source_files_in_extraction(extraction, root)) + + +def paths_missing_from_extraction( + extraction: dict, + paths: list[str], + root: str | Path | None = None, +) -> list[str]: + """Return paths whose re-extraction produced no nodes/edges/hyperedges.""" + if not paths: + return [] + extracted = source_files_in_extraction(extraction, root) + return [p for p in paths if not (source_path_aliases(p, root) & extracted)] + + def prefix_graph_for_global(G: nx.Graph, repo_tag: str) -> nx.Graph: """Return a copy of G with all node IDs prefixed with repo_tag::. diff --git a/graphify/file_slice.py b/graphify/file_slice.py new file mode 100644 index 000000000..d19c547ad --- /dev/null +++ b/graphify/file_slice.py @@ -0,0 +1,253 @@ +"""Intra-file slicing for semantic LLM extraction. + +Large markdown/text documents can exceed a model context window or graphify's +per-chunk token budget. FileSlice represents a byte range within one source +file; several slices still cache and increment under the parent file path. +""" +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path +import sys + +# Document-like extensions eligible for intra-file splitting. +_SPLITTABLE_SUFFIXES = frozenset({".md", ".mdx", ".txt", ".rst"}) + +_CHARS_PER_TOKEN = 4 +_PER_FILE_OVERHEAD_CHARS = 160 + +_HEADING_SPLIT = re.compile(r"(?=^#{1,6}\s)", re.MULTILINE) + + +@dataclass(frozen=True) +class FileSlice: + """A character range within a single on-disk file.""" + + path: Path + start_char: int + end_char: int + slice_index: int = 0 + slice_count: int = 1 + + def source_location(self) -> str: + return f"chars:{self.start_char}-{self.end_char}" + + +SemanticUnit = Path | FileSlice + + +def is_file_slice(unit: SemanticUnit) -> bool: + return isinstance(unit, FileSlice) + + +def unit_path(unit: SemanticUnit) -> Path: + return unit.path if is_file_slice(unit) else unit + + +def is_splittable_text(path: Path) -> bool: + return path.suffix.lower() in _SPLITTABLE_SUFFIXES + + +def read_unit_text(unit: SemanticUnit) -> str: + if is_file_slice(unit): + full = unit.path.read_text(encoding="utf-8", errors="replace") + return full[unit.start_char : unit.end_char] + return unit.read_text(encoding="utf-8", errors="replace") + + +def _count_tokens(text: str, tokenizer: object | None) -> int: + if tokenizer is not None: + return len(tokenizer.encode(text)) # type: ignore[attr-defined] + return len(text) // _CHARS_PER_TOKEN + + +def estimate_unit_tokens( + unit: SemanticUnit, + *, + tokenizer: object | None, + char_cap: int | None = None, +) -> int: + if is_file_slice(unit): + text = read_unit_text(unit) + return _count_tokens(text, tokenizer) + (_PER_FILE_OVERHEAD_CHARS // _CHARS_PER_TOKEN) + + path = unit + try: + text = path.read_text(encoding="utf-8", errors="replace") + if char_cap is not None: + text = text[:char_cap] + except OSError: + return 0 + return _count_tokens(text, tokenizer) + (_PER_FILE_OVERHEAD_CHARS // _CHARS_PER_TOKEN) + + +def _pick_ranges( + text: str, + max_tokens: int, + tokenizer: object | None, +) -> list[tuple[int, int]]: + """Split *text* into (start, end) char ranges each <= max_tokens.""" + if not text: + return [(0, 0)] + + total = _count_tokens(text, tokenizer) + if total <= max_tokens: + return [(0, len(text))] + + # Prefer markdown heading boundaries, then paragraph breaks, then newlines. + for pattern in (_HEADING_SPLIT, re.compile(r"\n\n+"), re.compile(r"\n")): + boundaries = [0] + for match in pattern.finditer(text): + pos = match.start() + if pos > boundaries[-1]: + boundaries.append(pos) + boundaries.append(len(text)) + ranges = _pack_boundaries(text, boundaries, max_tokens, tokenizer) + if ranges: + return ranges + + # Hard split by character budget. + ranges: list[tuple[int, int]] = [] + start = 0 + approx_chars = max(256, max_tokens * _CHARS_PER_TOKEN) + while start < len(text): + end = min(len(text), start + approx_chars) + while end > start and _count_tokens(text[start:end], tokenizer) > max_tokens: + end -= max(1, (end - start) // 8) + if end == start: + end = min(len(text), start + 1) + ranges.append((start, end)) + start = end + return ranges + + +def _pack_boundaries( + text: str, + boundaries: list[int], + max_tokens: int, + tokenizer: object | None, +) -> list[tuple[int, int]]: + ranges: list[tuple[int, int]] = [] + chunk_start = 0 + i = 1 + while i < len(boundaries): + end = boundaries[i] + if _count_tokens(text[chunk_start:end], tokenizer) <= max_tokens: + i += 1 + continue + if boundaries[i - 1] > chunk_start: + ranges.append((chunk_start, boundaries[i - 1])) + chunk_start = boundaries[i - 1] + continue + # Single boundary span still too large — give up on this pattern. + return [] + if chunk_start < len(text): + ranges.append((chunk_start, len(text))) + return ranges + + +def split_file_into_slices( + path: Path, + max_tokens: int, + *, + tokenizer: object | None, +) -> list[FileSlice]: + if max_tokens <= 0: + raise ValueError(f"max_tokens must be positive, got {max_tokens}") + try: + text = path.read_text(encoding="utf-8", errors="replace") + except OSError: + return [FileSlice(path=path, start_char=0, end_char=0)] + + ranges = _pick_ranges(text, max_tokens, tokenizer) + count = len(ranges) + return [ + FileSlice(path=path, start_char=start, end_char=end, slice_index=idx, slice_count=count) + for idx, (start, end) in enumerate(ranges) + ] + + +def bisect_file_slice(slice_: FileSlice, *, tokenizer: object | None) -> tuple[FileSlice, FileSlice]: + text = read_unit_text(slice_) + if len(text) < 2: + return slice_, slice_ + mid = len(text) // 2 + nl = text.rfind("\n", 0, mid) + if nl > 0: + mid = nl + 1 + left_end = slice_.start_char + mid + return ( + FileSlice(path=slice_.path, start_char=slice_.start_char, end_char=left_end), + FileSlice(path=slice_.path, start_char=left_end, end_char=slice_.end_char), + ) + + +def expand_oversized_files( + files: list[Path], + token_budget: int, + *, + tokenizer: object | None, + char_cap: int | None = None, +) -> list[SemanticUnit]: + """Replace oversized splittable text files with FileSlice units.""" + units: list[SemanticUnit] = [] + for path in files: + if not is_splittable_text(path): + units.append(path) + continue + cost = estimate_unit_tokens(path, tokenizer=tokenizer, char_cap=char_cap) + if cost <= token_budget: + units.append(path) + continue + slices = split_file_into_slices(path, token_budget, tokenizer=tokenizer) + if len(slices) <= 1: + units.append(path) + continue + print( + f"[graphify] split {path.name} into {len(slices)} slices " + f"(~{cost} tokens > budget {token_budget})", + file=sys.stderr, + ) + units.extend(slices) + return units + + +def split_unit_for_retry( + unit: SemanticUnit, + token_budget: int, + *, + tokenizer: object | None, +) -> list[SemanticUnit]: + """Break one unit into smaller pieces for adaptive retry.""" + if is_file_slice(unit): + left, right = bisect_file_slice(unit, tokenizer=tokenizer) + return [left, right] + if is_splittable_text(unit): + slices = split_file_into_slices(unit, max(256, token_budget // 2), tokenizer=tokenizer) + if len(slices) > 1: + return list(slices) + return [unit] + + +def unit_label(unit: SemanticUnit) -> str: + if is_file_slice(unit): + return f"{unit.path} [{unit.slice_index + 1}/{unit.slice_count}]" + return str(unit) + + +def split_chunk_for_retry( + chunk: list[SemanticUnit], + token_budget: int | None, + *, + tokenizer: object | None, +) -> list[list[SemanticUnit]] | None: + """Split *chunk* into smaller retry units, or None if unrecoverable.""" + if len(chunk) > 1: + mid = len(chunk) // 2 + return [chunk[:mid], chunk[mid:]] + if len(chunk) == 1 and token_budget is not None: + parts = split_unit_for_retry(chunk[0], token_budget, tokenizer=tokenizer) + if len(parts) > 1: + return [[part] for part in parts] + return None diff --git a/graphify/llm.py b/graphify/llm.py index 335e1b3fc..be7fc4edc 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -16,6 +16,18 @@ from dataclasses import dataclass, replace from pathlib import Path +from graphify.file_slice import ( + FileSlice, + SemanticUnit, + expand_oversized_files, + estimate_unit_tokens, + is_file_slice, + read_unit_text, + split_chunk_for_retry, + unit_label, + unit_path, +) + # `_read_files` truncates each file at this many characters before joining into # the user message. Token estimates use the same cap so packing matches reality. _FILE_CHAR_CAP = 20_000 @@ -83,7 +95,7 @@ def _get_tokenizer(): "gemini": { "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/", "default_model": "gemini-3-flash-preview", - "env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY"], + "env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY", "GOOGLE_BYOK"], "model_env_key": "GRAPHIFY_GEMINI_MODEL", "pricing": {"input": 0.50, "output": 3.00}, # USD per 1M tokens "temperature": 0, @@ -453,7 +465,7 @@ def _wrap_untrusted(rel: str, content: str) -> str: ) -def _read_files(paths: list[Path], root: Path) -> str: +def _read_files(units: list[SemanticUnit], root: Path) -> str: """Return file contents formatted for the extraction prompt. Each file is wrapped in an delimiter block and known @@ -461,16 +473,26 @@ def _read_files(paths: list[Path], root: Path) -> str: be confused with the trusted system instructions (see issue #1210). """ parts: list[str] = [] - for p in paths: + for unit in units: + p = unit_path(unit) try: rel = str(p.relative_to(root)) except ValueError: rel = str(p) try: - content = _file_to_text(p) + if is_file_slice(unit): + content = read_unit_text(unit) + sl = unit.source_location() + content = ( + f"[graphify slice {unit.slice_index + 1}/{unit.slice_count}, {sl}]\n" + + content + ) + else: + content = _file_to_text(p) + content = content[:_FILE_CHAR_CAP] except OSError: continue - parts.append(_wrap_untrusted(rel, content[:_FILE_CHAR_CAP])) + parts.append(_wrap_untrusted(rel, content)) return "\n\n".join(parts) @@ -538,11 +560,19 @@ def _is_vision_image(path: Path) -> bool: return path.suffix.lower() in _VISION_IMAGE_EXTENSIONS -def _partition_semantic_files(files: list[Path]) -> tuple[list[Path], list[Path]]: - """Split a chunk into (text-like files, raster-image files).""" - text_files = [f for f in files if not _is_vision_image(f)] - image_files = [f for f in files if _is_vision_image(f)] - return text_files, image_files +def _partition_semantic_files( + units: list[SemanticUnit], +) -> tuple[list[SemanticUnit], list[Path]]: + """Split a chunk into (text-like units, raster-image files).""" + text_units: list[SemanticUnit] = [] + image_files: list[Path] = [] + for unit in units: + p = unit_path(unit) + if _is_vision_image(p): + image_files.append(p) + else: + text_units.append(unit) + return text_units, image_files def _build_image_refs(image_files: list[Path], root: Path, *, read_bytes: bool = True) -> list[_ImageRef]: @@ -835,6 +865,12 @@ def _backend_pkg_hint(pkg: str, extra: str) -> str: ) +def _supports_reasoning_effort(model: str) -> bool: + """True when the model accepts OpenAI-style reasoning_effort / thinking level.""" + m = (model or "").lower().rsplit("/", 1)[-1] + return not m.startswith("gemma") + + def _call_openai_compat( base_url: str, api_key: str, @@ -1257,7 +1293,7 @@ def _call_bedrock(model: str, user_message: str, max_tokens: int = 8192, *, deep def extract_files_direct( - files: list[Path], + files: list[SemanticUnit], backend: str | None = None, api_key: str | None = None, model: str | None = None, @@ -1346,7 +1382,7 @@ def extract_files_direct( mdl, user_msg, temperature=_resolve_temperature(cfg.get("temperature", 0), mdl), - reasoning_effort=cfg.get("reasoning_effort"), + reasoning_effort=cfg.get("reasoning_effort") if _supports_reasoning_effort(mdl) else None, max_completion_tokens=_resolve_max_tokens(cfg.get("max_completion_tokens", 8192)), backend=backend, deep_mode=deep_mode, @@ -1356,63 +1392,45 @@ def extract_files_direct( def _estimate_file_tokens(path: Path) -> int: - """Estimate the prompt-token cost of a single file under `_read_files` rules. + """Estimate the prompt-token cost of a single file under `_read_files` rules.""" + return estimate_unit_tokens(path, tokenizer=_TOKENIZER, char_cap=_FILE_CHAR_CAP) - Uses tiktoken (`cl100k_base`) when available for accurate counts. Falls back - to the chars/4 heuristic if tiktoken is not installed. Both paths cap at - `_FILE_CHAR_CAP` to match `_read_files`'s truncation, plus a constant for - the `=== rel ===` separator. Returns 0 for unreadable paths so they don't - blow up packing. - """ - # Raster images are not read as text; a vision model bills them at a roughly - # fixed token cost, so estimate by image count rather than (binary) byte size. - if _is_vision_image(path): - return _IMAGE_TOKEN_ESTIMATE - if _TOKENIZER is None: - try: - size = path.stat().st_size - except OSError: - return 0 - chars = min(size, _FILE_CHAR_CAP) + _PER_FILE_OVERHEAD_CHARS - return chars // _CHARS_PER_TOKEN - try: - content = path.read_text(encoding="utf-8", errors="replace")[:_FILE_CHAR_CAP] - except OSError: - return 0 - return len(_TOKENIZER.encode(content)) + (_PER_FILE_OVERHEAD_CHARS // _CHARS_PER_TOKEN) +def _estimate_unit_tokens(unit: SemanticUnit) -> int: + if _is_vision_image(unit_path(unit)): + return _IMAGE_TOKEN_ESTIMATE + return estimate_unit_tokens(unit, tokenizer=_TOKENIZER, char_cap=_FILE_CHAR_CAP) def _pack_chunks_by_tokens( - files: list[Path], + units: list[SemanticUnit], token_budget: int, -) -> list[list[Path]]: +) -> list[list[SemanticUnit]]: """Greedily pack files into chunks that fit a token budget. Files are first grouped by parent directory so related artifacts share a chunk (cross-file edges are more likely to be extracted within a chunk than across chunks). Within each directory, files are added one at a time; a chunk is closed when adding the next file would exceed the - budget. A single file larger than the budget gets its own chunk and the - caller is expected to handle the API error if it actually overflows the - model's context window — packing can't shrink one big file. + budget. Oversized splittable text files should already be expanded into + :class:`~graphify.file_slice.FileSlice` units before packing. """ if token_budget <= 0: raise ValueError(f"token_budget must be positive, got {token_budget}") - by_dir: dict[Path, list[Path]] = {} - for f in files: - by_dir.setdefault(f.parent, []).append(f) + by_dir: dict[Path, list[SemanticUnit]] = {} + for unit in units: + by_dir.setdefault(unit_path(unit).parent, []).append(unit) - chunks: list[list[Path]] = [] - current: list[Path] = [] + chunks: list[list[SemanticUnit]] = [] + current: list[SemanticUnit] = [] current_tokens = 0 current_images = 0 for directory in sorted(by_dir): - for path in by_dir[directory]: - cost = _estimate_file_tokens(path) - is_image = _is_vision_image(path) + for unit in by_dir[directory]: + cost = _estimate_unit_tokens(unit) + is_image = _is_vision_image(unit_path(unit)) over_budget = current_tokens + cost > token_budget over_images = is_image and current_images >= _MAX_IMAGES_PER_CHUNK if current and (over_budget or over_images): @@ -1420,7 +1438,7 @@ def _pack_chunks_by_tokens( current = [] current_tokens = 0 current_images = 0 - current.append(path) + current.append(unit) current_tokens += cost current_images += is_image @@ -1458,8 +1476,25 @@ def _looks_like_context_exceeded(exc: BaseException) -> bool: return any(marker in msg for marker in _CONTEXT_EXCEEDED_MARKERS) +def _merge_extraction_results( + left: dict, + right: dict, + *, + model: str | None = None, +) -> dict: + return { + "nodes": left.get("nodes", []) + right.get("nodes", []), + "edges": left.get("edges", []) + right.get("edges", []), + "hyperedges": left.get("hyperedges", []) + right.get("hyperedges", []), + "input_tokens": left.get("input_tokens", 0) + right.get("input_tokens", 0), + "output_tokens": left.get("output_tokens", 0) + right.get("output_tokens", 0), + "model": model or right.get("model") or left.get("model"), + "finish_reason": "stop", + } + + def _extract_with_adaptive_retry( - chunk: list[Path], + chunk: list[SemanticUnit], backend: str, api_key: str | None, model: str | None, @@ -1468,6 +1503,7 @@ def _extract_with_adaptive_retry( _depth: int = 0, *, deep_mode: bool = False, + token_budget: int | None = None, ) -> dict: """Extract a chunk; if the response is truncated (`finish_reason="length"`) or the API rejects the prompt as too large for the model's context window, @@ -1497,9 +1533,34 @@ def _extract_with_adaptive_retry( still failing at the cap, we surface the (likely empty) result with a warning rather than infinite-loop. - A single-file chunk that overflows is unrecoverable here — we can't make - one file smaller than itself, so we return what we got and warn. + A single-file chunk that overflows may still be recoverable when the unit + is splittable text (``.md``, ``.txt``, etc.) — we bisect the file or slice + and recurse. Non-splittable single files (e.g. one huge ``.py``) cannot be + shrunk further; we return what we got and warn. """ + def _retry_halves() -> dict: + sub_chunks = split_chunk_for_retry(chunk, token_budget, tokenizer=_TOKENIZER) + if sub_chunks is None: + return {} + results = [ + _extract_with_adaptive_retry( + sub, + backend, + api_key, + model, + root, + max_depth, + _depth + 1, + deep_mode=deep_mode, + token_budget=token_budget, + ) + for sub in sub_chunks + ] + merged = results[0] + for part in results[1:]: + merged = _merge_extraction_results(merged, part, model=model) + return merged + try: result = extract_files_direct( chunk, backend=backend, api_key=api_key, model=model, root=root, deep_mode=deep_mode @@ -1507,87 +1568,57 @@ def _extract_with_adaptive_retry( except Exception as exc: # noqa: BLE001 — re-raise unless it's a known context overflow if not _looks_like_context_exceeded(exc): raise - if len(chunk) <= 1: + if _depth >= max_depth: print( - f"[graphify] single-file chunk {chunk[0]} exceeds model context " - f"and cannot be split further: {exc}", + f"[graphify] chunk of {len(chunk)} still overflows context at " + f"recursion depth {_depth} (max {max_depth}) — dropping", file=sys.stderr, ) return {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0, "model": model, "finish_reason": "stop"} - if _depth >= max_depth: + sub = split_chunk_for_retry(chunk, token_budget, tokenizer=_TOKENIZER) + if sub is None: print( - f"[graphify] chunk of {len(chunk)} still overflows context at " - f"recursion depth {_depth} (max {max_depth}) — dropping", + f"[graphify] single-file chunk {unit_label(chunk[0])} exceeds model context " + f"and cannot be split further: {exc}", file=sys.stderr, ) return {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0, "model": model, "finish_reason": "stop"} print( f"[graphify] chunk of {len(chunk)} exceeded context at depth " - f"{_depth} ({type(exc).__name__}); splitting in half and retrying", + f"{_depth} ({type(exc).__name__}); splitting and retrying", file=sys.stderr, ) - mid = len(chunk) // 2 - left = _extract_with_adaptive_retry( - chunk[:mid], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode - ) - right = _extract_with_adaptive_retry( - chunk[mid:], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode - ) - return { - "nodes": left.get("nodes", []) + right.get("nodes", []), - "edges": left.get("edges", []) + right.get("edges", []), - "hyperedges": left.get("hyperedges", []) + right.get("hyperedges", []), - "input_tokens": left.get("input_tokens", 0) + right.get("input_tokens", 0), - "output_tokens": left.get("output_tokens", 0) + right.get("output_tokens", 0), - "model": model, - "finish_reason": "stop", - } + return _retry_halves() if result.get("finish_reason") != "length": return result - if len(chunk) <= 1: + if _depth >= max_depth: print( - f"[graphify] single-file chunk {chunk[0]} truncated at " - f"max_completion_tokens — partial result kept", + f"[graphify] chunk of {len(chunk)} still truncated at recursion " + f"depth {_depth} (max {max_depth}) — partial result kept", file=sys.stderr, ) return result - if _depth >= max_depth: + sub = split_chunk_for_retry(chunk, token_budget, tokenizer=_TOKENIZER) + if sub is None: print( - f"[graphify] chunk of {len(chunk)} still truncated at recursion " - f"depth {_depth} (max {max_depth}) — partial result kept", + f"[graphify] single-file chunk {unit_label(chunk[0])} truncated at " + f"max_completion_tokens — partial result kept", file=sys.stderr, ) return result print( f"[graphify] chunk of {len(chunk)} truncated at depth {_depth}, " - f"splitting into halves of {len(chunk) // 2} and " - f"{len(chunk) - len(chunk) // 2}", + f"splitting into {len(sub)} sub-chunk(s) and retrying", file=sys.stderr, ) - mid = len(chunk) // 2 - left = _extract_with_adaptive_retry( - chunk[:mid], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode - ) - right = _extract_with_adaptive_retry( - chunk[mid:], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode - ) - - return { - "nodes": left.get("nodes", []) + right.get("nodes", []), - "edges": left.get("edges", []) + right.get("edges", []), - "hyperedges": left.get("hyperedges", []) + right.get("hyperedges", []), - "input_tokens": left.get("input_tokens", 0) + right.get("input_tokens", 0), - "output_tokens": left.get("output_tokens", 0) + right.get("output_tokens", 0), - "model": result.get("model"), - # Both halves either succeeded or have already surfaced their own - # truncation warning; the merged result is no longer truncated as a - # logical unit. - "finish_reason": "stop", - } + retried = _retry_halves() + if retried: + return retried + return result def extract_corpus_parallel( @@ -1638,7 +1669,13 @@ def extract_corpus_parallel( chunk does not abort the run. """ if token_budget is not None: - chunks = _pack_chunks_by_tokens(files, token_budget=token_budget) + units = expand_oversized_files( + files, + token_budget, + tokenizer=_TOKENIZER, + char_cap=_FILE_CHAR_CAP, + ) + chunks = _pack_chunks_by_tokens(units, token_budget=token_budget) else: chunks = [files[i:i + chunk_size] for i in range(0, len(files), chunk_size)] @@ -1649,7 +1686,7 @@ def extract_corpus_parallel( } total = len(chunks) - def _run_one(idx: int, chunk: list[Path]) -> tuple[int, dict | None, Exception | None]: + def _run_one(idx: int, chunk: list[SemanticUnit]) -> tuple[int, dict | None, Exception | None]: t0 = time.time() try: result = _extract_with_adaptive_retry( @@ -1660,6 +1697,7 @@ def _run_one(idx: int, chunk: list[Path]) -> tuple[int, dict | None, Exception | root=root, max_depth=max_retry_depth, deep_mode=deep_mode, + token_budget=token_budget, ) result["elapsed_seconds"] = round(time.time() - t0, 2) return idx, result, None @@ -1855,7 +1893,7 @@ def _call_llm( temperature = _resolve_temperature(cfg.get("temperature", 0), mdl) if temperature is not None: kwargs["temperature"] = temperature - if cfg.get("reasoning_effort"): + if cfg.get("reasoning_effort") and _supports_reasoning_effort(mdl): kwargs["reasoning_effort"] = cfg["reasoning_effort"] # Custom providers can override via providers.json `extra_body`; falls back # to the moonshot default to preserve existing behavior. diff --git a/graphify/stitch.py b/graphify/stitch.py new file mode 100644 index 000000000..5f0b7a826 --- /dev/null +++ b/graphify/stitch.py @@ -0,0 +1,288 @@ +"""Post-merge incremental stitch — wire changed-file subgraphs to the existing graph. + +After incremental ``build_merge``, re-extracted files often contain mentions of +symbols and paths elsewhere in the corpus, but the LLM chunk lacked that context. +This module adds conservative ``references`` edges by scanning changed files on +disk and resolving mentions against the full merged graph. +""" +from __future__ import annotations + +import re +import sys +from pathlib import Path + +import networkx as nx + +from graphify.build import source_path_aliases +from graphify.symbol_resolution import existing_edge_pairs, normalise_callable_label + +_BACKTICK = re.compile(r"`([^`]+)`") +_MD_LINK = re.compile(r"\[[^\]]*\]\(([^)]+)\)") +_PATH_EXT = re.compile( + r"\.(?:md|py|ts|tsx|js|jsx|java|go|rs|txt|rst|pdf)(?:$|[#?])", + re.IGNORECASE, +) +_MIN_SYMBOL_LEN = 3 + + +def _stitch_label_index(G: nx.Graph) -> dict[str, list[str]]: + """Map normalised label -> node ids (code, document, and concept nodes).""" + index: dict[str, list[str]] = {} + for nid, data in G.nodes(data=True): + if data.get("file_type") == "rationale": + continue + label = str(data.get("label", "")).strip().strip("()") + if not label or len(label) < _MIN_SYMBOL_LEN: + continue + if label.endswith((".py", ".md", ".ts", ".tsx", ".js", ".jsx")): + continue + key = normalise_callable_label(label) + if key: + index.setdefault(key, []).append(nid) + return index + + +def _nodes_for_source_file(G: nx.Graph, path: str, root: Path) -> list[str]: + aliases = source_path_aliases(path, root) + if not aliases: + return [] + found: list[str] = [] + for nid, data in G.nodes(data=True): + sf = data.get("source_file") + if sf and source_path_aliases(str(sf), root) & aliases: + found.append(nid) + return found + + +def _source_file_matches_path(source_file: str | None, rel: Path, root: Path) -> bool: + if not source_file: + return False + changed = source_path_aliases(rel.as_posix(), root) + return bool(source_path_aliases(str(source_file), root) & changed) + + +def _is_foreign_attribution(source_file: str | None, rel: Path, root: Path) -> bool: + """True when source_file points at a different existing file than *rel*.""" + if not source_file: + return False + if _source_file_matches_path(source_file, rel, root): + return False + disk = Path(source_file) + if not disk.is_absolute(): + disk = root / disk + return disk.is_file() + + +def _fallback_local_nodes( + G: nx.Graph, + new_ids: set[str], + rel: Path, + root: Path, +) -> list[str]: + """Fresh extraction nodes for stitch anchors when source_file was wrong. + + Keeps nodes attributed to missing/hallucinated paths; drops nodes the LLM + placed under a different on-disk file (e.g. calendar.md while stitching array.md). + """ + kept: list[str] = [] + for nid in sorted(new_ids): + if nid not in G: + continue + sf = G.nodes[nid].get("source_file") + if _is_foreign_attribution(str(sf) if sf else None, rel, root): + continue + kept.append(nid) + return kept + + +def _pick_path_target_anchor(node_ids: list[str], G: nx.Graph, rel_path: Path) -> str | None: + """Pick a file-level target for a path mention; never an arbitrary code symbol.""" + if not node_ids: + return None + stem = rel_path.stem.lower().replace("_", " ") + doc_nodes = [n for n in node_ids if n.endswith("_document")] + if doc_nodes: + return sorted(doc_nodes)[0] + for nid in sorted(node_ids): + data = G.nodes[nid] + if data.get("file_type") not in ("document", "concept", "paper"): + continue + label = str(data.get("label", "")).lower().replace("_", " ") + if stem in label or label in stem: + return nid + return None + + +def _pick_file_anchor(node_ids: list[str], G: nx.Graph, rel_path: Path) -> str | None: + if not node_ids: + return None + stem = rel_path.stem.lower().replace("_", " ") + doc_nodes = [n for n in node_ids if n.endswith("_document")] + if doc_nodes: + return sorted(doc_nodes)[0] + for nid in node_ids: + data = G.nodes[nid] + ft = data.get("file_type") + label = str(data.get("label", "")).lower().replace("_", " ") + if ft in ("document", "concept", "paper") and stem in label: + return nid + return sorted(node_ids)[0] + + +def _looks_like_path(token: str) -> bool: + token = token.strip() + if "/" in token or "\\" in token: + return True + return bool(_PATH_EXT.search(token)) + + +def _resolve_path_target(token: str, root: Path, G: nx.Graph) -> list[str]: + token = token.strip().split("#")[0].split("?")[0] + if not token: + return [] + candidates = [token] + if not Path(token).is_absolute(): + candidates.append(str(root / token)) + for cp in candidates: + nids = _nodes_for_source_file(G, cp, root) + if not nids: + continue + try: + rel = ( + Path(cp).resolve().relative_to(root.resolve()) + if Path(cp).is_absolute() + else Path(token) + ) + except ValueError: + rel = Path(token) + anchor = _pick_path_target_anchor(nids, G, rel) + return [anchor] if anchor else [] + return [] + + +def _resolve_symbol_targets( + token: str, + label_index: dict[str, list[str]], + local_ids: set[str], + G: nx.Graph, +) -> list[str]: + key = normalise_callable_label(token) + if not key or len(key) < _MIN_SYMBOL_LEN: + return [] + cands = label_index.get(key, []) + if not cands: + return [] + external = [c for c in cands if c not in local_ids] + if len(external) == 1: + return external + if len(external) > 1: + by_source: dict[str, list[str]] = {} + for nid in external: + sf = str(G.nodes[nid].get("source_file", "")) + by_source.setdefault(sf, []).append(nid) + if len(by_source) == 1: + return [sorted(by_source[next(iter(by_source))])[0]] + return [] + return [] + + +def _edge_triples_from_graph(G: nx.Graph) -> set[tuple[str, str, str]]: + edges = [ + { + "source": d.get("_src", u), + "target": d.get("_tgt", v), + "relation": d.get("relation", ""), + } + for u, v, d in G.edges(data=True) + ] + return existing_edge_pairs(edges) + + +def stitch_incremental_links( + G: nx.Graph, + changed_paths: list[str], + *, + root: str | Path, + new_node_ids: set[str] | frozenset[str] | None = None, +) -> int: + """Add ``references`` edges from changed files to the rest of the graph. + + Scans each changed file for backtick identifiers and markdown path links, + resolves them against *G*, and attaches edges from the changed file's anchor + node. Returns the number of edges added. + + When the LLM attributes re-extracted nodes to the wrong ``source_file``, + pass *new_node_ids* (node ids from the fresh extraction) so anchors can + still be chosen for stitch edges. + """ + if not changed_paths: + return 0 + root_path = Path(root).resolve() + new_ids = set(new_node_ids or ()) + label_index = _stitch_label_index(G) + known = _edge_triples_from_graph(G) + added = 0 + + for path_str in changed_paths: + disk_path = Path(path_str) + if not disk_path.is_absolute(): + disk_path = root_path / disk_path + if not disk_path.is_file(): + continue + try: + rel = disk_path.resolve().relative_to(root_path) + except ValueError: + rel = Path(path_str) + try: + text = disk_path.read_text(encoding="utf-8") + except OSError: + continue + + local_nids = _nodes_for_source_file(G, str(rel), root_path) + if not local_nids and new_ids: + local_nids = _fallback_local_nodes(G, new_ids, rel, root_path) + anchor = _pick_file_anchor(local_nids, G, rel) + if not anchor: + continue + + local_ids = set(local_nids) + seen_tokens: set[str] = set() + + for token in _BACKTICK.findall(text) + _MD_LINK.findall(text): + token = token.strip() + if not token or token in seen_tokens: + continue + seen_tokens.add(token) + + if _looks_like_path(token): + targets = _resolve_path_target(token, root_path, G) + else: + targets = _resolve_symbol_targets(token, label_index, local_ids, G) + + for tgt in targets: + if tgt == anchor: + continue + triple = (anchor, tgt, "references") + if triple in known: + continue + known.add(triple) + G.add_edge( + anchor, + tgt, + relation="references", + confidence="EXTRACTED", + confidence_score=1.0, + source_file=rel.as_posix(), + weight=1.0, + _src=anchor, + _tgt=tgt, + ) + added += 1 + + if added: + print( + f"[graphify] Stitched {added} cross-file reference edge(s) " + f"for {len(changed_paths)} changed file(s).", + file=sys.stderr, + ) + return added diff --git a/tests/test_build.py b/tests/test_build.py index 0a3fffa49..e7f86152f 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -2,7 +2,15 @@ from pathlib import Path import networkx as nx from networkx.readwrite import json_graph -from graphify.build import build_from_json, build, build_merge, edge_data, edge_datas, dedupe_edges, dedupe_nodes +from graphify.build import ( + build_from_json, + build, + build_merge, + edge_data, + edge_datas, + path_covered_by_extraction, + paths_missing_from_extraction, +) FIXTURES = Path(__file__).parent / "fixtures" @@ -549,6 +557,29 @@ def test_build_merge_prune_windows_backslash_paths(tmp_path): assert "parse_date" not in node_labels, "node should be pruned even with backslash path" +def test_paths_missing_from_extraction_flags_empty_reextract(tmp_path): + """Incremental update must detect a file whose chunk returned no nodes.""" + root = tmp_path / "corpus" + root.mkdir() + doc_a = root / "docs" / "a.md" + doc_b = root / "docs" / "b.md" + doc_a.parent.mkdir(parents=True) + doc_a.write_text("# A") + doc_b.write_text("# B") + + empty = {"nodes": [], "edges": [], "hyperedges": []} + partial = { + "nodes": [{"id": "n1", "label": "A", "file_type": "document", "source_file": "docs/a.md"}], + "edges": [], + "hyperedges": [], + } + assert paths_missing_from_extraction(empty, [str(doc_a)], root) == [str(doc_a)] + assert paths_missing_from_extraction(partial, [str(doc_a)], root) == [] + assert paths_missing_from_extraction(partial, [str(doc_b)], root) == [str(doc_b)] + assert path_covered_by_extraction(str(doc_a), partial, root) + assert not path_covered_by_extraction(str(doc_b), partial, root) + + def test_build_merge_rejects_oversized_existing_graph(monkeypatch, tmp_path): """#F4: build_merge must refuse to read an existing graph.json that exceeds the size cap, rather than json.loads-ing it into memory.""" diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 087464ab8..87ee5bb60 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -394,7 +394,7 @@ def always_truncate(chunk, **kwargs): def test_adaptive_retry_single_file_truncation_does_not_recurse(tmp_path, capsys): - """A single file that truncates can't be split further — surface a + """A single .py file that truncates can't be split further — surface a warning and return what we got. No infinite loop.""" from graphify.llm import _extract_with_adaptive_retry @@ -408,7 +408,8 @@ def stub(chunk, **kwargs): with patch("graphify.llm.extract_files_direct", side_effect=stub): _extract_with_adaptive_retry( - [f], backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3 + [f], backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3, + token_budget=1_000, ) assert calls == [1], f"single-file chunk recursed; calls = {calls}" @@ -416,6 +417,38 @@ def stub(chunk, **kwargs): assert "single-file chunk" in err and "truncated" in err +def test_adaptive_retry_splits_oversized_markdown_on_truncation(tmp_path): + """A single markdown slice that truncates should bisect and merge on retry.""" + from graphify.file_slice import FileSlice + from graphify.llm import _extract_with_adaptive_retry + + doc = tmp_path / "big.md" + text = "line\n" * 200 + doc.write_text(text) + whole = FileSlice(path=doc, start_char=0, end_char=len(text)) + + calls = [] + + def stub(chunk, **kwargs): + calls.append(len(chunk)) + finish = "length" if len(chunk) == 1 else "stop" + return _stub_with_finish(len(chunk), finish_reason=finish) + + with patch("graphify.llm.extract_files_direct", side_effect=stub): + result = _extract_with_adaptive_retry( + [whole], + backend="kimi", + api_key=None, + model=None, + root=tmp_path, + max_depth=3, + token_budget=2_000, + ) + + assert len(calls) > 1 + assert result["finish_reason"] == "stop" + + def test_corpus_parallel_uses_adaptive_retry(tmp_path): """End-to-end: extract_corpus_parallel routes through adaptive retry, so a chunk that truncates gets split and merged transparently before diff --git a/tests/test_file_slice.py b/tests/test_file_slice.py new file mode 100644 index 000000000..bcf9b956f --- /dev/null +++ b/tests/test_file_slice.py @@ -0,0 +1,94 @@ +"""Tests for intra-file slicing used by semantic extraction.""" +from pathlib import Path +from unittest.mock import patch + +import pytest + +from graphify.file_slice import ( + FileSlice, + bisect_file_slice, + expand_oversized_files, + is_splittable_text, + read_unit_text, + split_chunk_for_retry, + split_file_into_slices, + split_unit_for_retry, + unit_label, +) + + +def test_is_splittable_text_recognises_markdown(): + assert is_splittable_text(Path("doc.md")) + assert is_splittable_text(Path("doc.MDX")) + assert not is_splittable_text(Path("code.py")) + + +def test_split_file_into_slices_respects_token_budget(tmp_path): + doc = tmp_path / "big.md" + doc.write_text("# One\n\n" + ("paragraph.\n\n" * 200)) + + slices = split_file_into_slices(doc, max_tokens=50, tokenizer=None) + assert len(slices) > 1 + assert all(isinstance(s, FileSlice) for s in slices) + assert slices[0].slice_index == 0 + assert slices[-1].slice_count == len(slices) + rejoined = "".join(read_unit_text(s) for s in slices) + assert rejoined == doc.read_text(encoding="utf-8") + + +def test_expand_oversized_files_splits_markdown_only(tmp_path): + md = tmp_path / "big.md" + md.write_text("x" * 20_000) + py = tmp_path / "big.py" + py.write_text("x" * 20_000) + + with patch("graphify.file_slice._count_tokens", side_effect=lambda text, _tok: len(text) // 4): + units = expand_oversized_files([md, py], token_budget=1_000, tokenizer=None) + + md_units = [u for u in units if getattr(u, "path", u) == md or u == md] + py_units = [u for u in units if getattr(u, "path", u) == py or u == py] + assert len(md_units) > 1 + assert py_units == [py] + + +def test_bisect_file_slice_splits_on_newline(tmp_path): + doc = tmp_path / "doc.md" + doc.write_text("line one\nline two\nline three\n") + whole = FileSlice(path=doc, start_char=0, end_char=len(doc.read_text(encoding="utf-8"))) + left, right = bisect_file_slice(whole, tokenizer=None) + assert read_unit_text(left) + read_unit_text(right) == doc.read_text(encoding="utf-8") + assert left.end_char == right.start_char + + +def test_split_unit_for_retry_bisects_existing_slice(tmp_path): + doc = tmp_path / "doc.md" + text = "a" * 100 + doc.write_text(text) + sl = FileSlice(path=doc, start_char=0, end_char=len(text)) + parts = split_unit_for_retry(sl, token_budget=2_000, tokenizer=None) + assert len(parts) == 2 + + +def test_split_unit_for_retry_splits_whole_markdown(tmp_path): + doc = tmp_path / "doc.md" + doc.write_text("# Title\n\n" + ("body line\n" * 800)) + + with patch("graphify.file_slice._count_tokens", side_effect=lambda text, _tok: len(text) // 4): + parts = split_unit_for_retry(doc, token_budget=100, tokenizer=None) + + assert len(parts) > 1 + + +def test_split_chunk_for_retry_splits_multi_file_chunk(tmp_path): + a, b = tmp_path / "a.md", tmp_path / "b.md" + a.write_text("# A\n") + b.write_text("# B\n") + sub = split_chunk_for_retry([a, b], token_budget=1000, tokenizer=None) + assert sub is not None and len(sub) == 2 + + +def test_unit_label_for_slice(tmp_path): + doc = tmp_path / "doc.md" + doc.write_text("# x\n") + sl = FileSlice(path=doc, start_char=0, end_char=3, slice_index=1, slice_count=4) + assert unit_label(sl) == f"{doc} [2/4]" diff --git a/tests/test_stitch.py b/tests/test_stitch.py new file mode 100644 index 000000000..b696d8bf2 --- /dev/null +++ b/tests/test_stitch.py @@ -0,0 +1,184 @@ +"""Tests for incremental post-merge stitch pass.""" +from __future__ import annotations + +import json +from pathlib import Path + +import networkx as nx + +from graphify.build import build_from_json, build_merge +from graphify.stitch import stitch_incremental_links + + +def _calendar_graph() -> dict: + cal_sf = "BoT-calendars-shared/src/CALENDAR_CALCULATIONS.md" + return { + "nodes": [ + { + "id": "calendar_calculations_calculateworkingtimefromteamscalendar", + "label": "calculateWorkingTimeFromTeamsCalendar", + "file_type": "code", + "source_file": cal_sf, + }, + { + "id": "calendar_calculations_createassignmentcalendarsnapshot", + "label": "createAssignmentCalendarSnapshot", + "file_type": "code", + "source_file": cal_sf, + }, + ], + "edges": [], + "input_tokens": 0, + "output_tokens": 0, + } + + +def _array_graph() -> dict: + array_sf = "ARRAY_FIELDS_TYPED_APPROACH.md" + return { + "nodes": [ + { + "id": "array_fields_typed_approach_sprint", + "label": "Sprint", + "file_type": "code", + "source_file": array_sf, + }, + { + "id": "array_fields_typed_approach_document", + "label": "Typed Array Fields Implementation", + "file_type": "document", + "source_file": array_sf, + }, + ], + "edges": [], + "input_tokens": 0, + "output_tokens": 0, + } + + +def test_stitch_links_symbol_mentions_from_changed_markdown(tmp_path: Path) -> None: + root = tmp_path / "corpus" + root.mkdir() + array_md = root / "ARRAY_FIELDS_TYPED_APPROACH.md" + cal_dir = root / "BoT-calendars-shared" / "src" + cal_dir.mkdir(parents=True) + cal_md = cal_dir / "CALENDAR_CALCULATIONS.md" + cal_md.write_text("# Calendar\n", encoding="utf-8") + array_md.write_text( + "# Typed Array Fields Implementation\n\n" + "See `BoT-calendars-shared/src/CALENDAR_CALCULATIONS.md` and " + "`calculateWorkingTimeFromTeamsCalendar`.\n", + encoding="utf-8", + ) + + graph_path = tmp_path / "graph.json" + G0 = build_from_json(_calendar_graph(), root=root) + graph_path.write_text( + json.dumps(nx.node_link_data(G0, edges="edges")), + encoding="utf-8", + ) + + G = build_merge( + [_array_graph()], + graph_path=graph_path, + prune_sources=None, + dedup=False, + root=root, + ) + + added = stitch_incremental_links( + G, + [str(array_md.relative_to(root))], + root=root, + ) + assert added >= 1 + + refs = [ + (d.get("_src", u), d.get("_tgt", v)) + for u, v, d in G.edges(data=True) + if d.get("relation") == "references" + ] + anchor = "array_fields_typed_approach_document" + cal_fn = "calendar_calculations_calculateworkingtimefromteamscalendar" + assert (anchor, cal_fn) in refs + + +def test_stitch_uses_new_node_ids_when_source_file_hallucinated(tmp_path: Path) -> None: + """LLM may attribute nodes to wrong paths; stitch still wires via new_node_ids.""" + root = tmp_path / "corpus" + root.mkdir() + array_md = root / "ARRAY_FIELDS_TYPED_APPROACH.md" + cal_dir = root / "BoT-calendars-shared" / "src" + cal_dir.mkdir(parents=True) + (cal_dir / "CALENDAR_CALCULATIONS.md").write_text("# Calendar\n", encoding="utf-8") + array_md.write_text( + "# Doc\n\nUse `calculateWorkingTimeFromTeamsCalendar`.\n", + encoding="utf-8", + ) + + graph_path = tmp_path / "graph.json" + G0 = build_from_json(_calendar_graph(), root=root) + graph_path.write_text( + json.dumps(nx.node_link_data(G0, edges="edges")), + encoding="utf-8", + ) + + hallucinated = { + "nodes": [ + { + "id": "array_field_renderer_SprintArrayRenderer", + "label": "SprintArrayRenderer", + "file_type": "code", + "source_file": "nextjsapp/app/utils/components/tasks/ArrayFieldRenderer.tsx", + }, + { + "id": "calendar_calculations_document", + "label": "Calendar Calculations", + "file_type": "document", + "source_file": "BoT-calendars-shared/src/CALENDAR_CALCULATIONS.md", + }, + ], + "edges": [], + "input_tokens": 0, + "output_tokens": 0, + } + G = build_merge([hallucinated], graph_path=graph_path, prune_sources=None, dedup=False, root=root) + added = stitch_incremental_links( + G, + [str(array_md.relative_to(root))], + root=root, + new_node_ids={ + "array_field_renderer_SprintArrayRenderer", + "calendar_calculations_document", + }, + ) + assert added == 1 + refs = [ + (d.get("_src", u), d.get("_tgt", v)) + for u, v, d in G.edges(data=True) + if d.get("relation") == "references" + ] + assert ("array_field_renderer_SprintArrayRenderer", "calendar_calculations_calculateworkingtimefromteamscalendar") in refs + + +def test_stitch_skips_ambiguous_symbol(tmp_path: Path) -> None: + root = tmp_path / "corpus" + root.mkdir() + doc_a = root / "a.md" + doc_b = root / "b.md" + doc_a.write_text("Uses `handle`.\n", encoding="utf-8") + doc_b.write_text("# B\n", encoding="utf-8") + + extraction = { + "nodes": [ + {"id": "svc_a_handle", "label": "handle", "file_type": "code", "source_file": "svc/a.py"}, + {"id": "svc_b_handle", "label": "handle", "file_type": "code", "source_file": "svc/b.py"}, + {"id": "b_document", "label": "B", "file_type": "document", "source_file": "b.md"}, + ], + "edges": [], + "input_tokens": 0, + "output_tokens": 0, + } + G = build_from_json(extraction, root=root) + added = stitch_incremental_links(G, ["b.md"], root=root) + assert added == 0