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 391ed28c4..9dd6a0754 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 @@ -454,7 +466,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 @@ -462,16 +474,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) @@ -539,11 +561,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]: @@ -788,7 +818,9 @@ def _response_is_hollow(raw_content: str | None, parsed: dict) -> bool: nodes = parsed.get("nodes") edges = parsed.get("edges") hyperedges = parsed.get("hyperedges") - return not nodes and not edges and not hyperedges + finish_reason = parsed.get("finish_reason") + # finish_reason stop needs to be checked in case empty graph is returned but the call was okay + return not nodes and not edges and not hyperedges and finish_reason != "stop" def _backend_env_keys(backend: str) -> list[str]: @@ -1266,7 +1298,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, @@ -1371,63 +1403,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): @@ -1435,7 +1449,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 @@ -1473,8 +1487,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, @@ -1483,6 +1514,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, @@ -1512,9 +1544,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 @@ -1522,87 +1579,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( @@ -1653,7 +1680,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)] @@ -1664,7 +1697,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( @@ -1675,6 +1708,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 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]"