diff --git a/graphify/cache.py b/graphify/cache.py index 7122b5c58..c5edc0839 100644 --- a/graphify/cache.py +++ b/graphify/cache.py @@ -465,7 +465,8 @@ def cache_dir(root: Path = Path("."), kind: str = "ast", def load_cached(path: Path, root: Path = Path("."), kind: str = "ast", cache_root: Path | None = None, prompt: "str | Path | None" = None, prompt_file: "str | Path | None" = None, - allow_legacy: bool = True) -> dict | None: + allow_legacy: bool = True, + allow_partial: bool = False) -> dict | None: """Return cached extraction for this file if hash matches, else None. Cache key: SHA256 of file contents. @@ -514,6 +515,17 @@ def load_cached(path: Path, root: Path = Path("."), kind: str = "ast", result = json.loads(entry.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): return None + # A ``partial`` entry was produced from a truncated LLM response and + # covers only part of the file's symbols. Serving it as authoritative + # would return the incomplete node set forever until the file is + # re-extracted. Treat it as a cache MISS (the normal read path) so the + # file is re-dispatched and retried. Self-heals: a later complete + # extraction overwrites the same content-hash key with a non-partial + # entry. ``allow_partial`` is the one exception — the merge_existing + # checkpoint peeks at a partial prev so it can accumulate a file's slices + # across chunks without losing the truncated one (it stays partial). + if not allow_partial and isinstance(result, dict) and result.get("partial"): + return None if legacy_hit: _legacy_semantic_hits += 1 # Re-anchor relative source_file fields so callers see the same @@ -748,6 +760,23 @@ def check_semantic_cache( return cached_nodes, cached_edges, cached_hyperedges, uncached +def _group_has_partial_marker(group: dict) -> bool: + """True if any node/edge/hyperedge in a per-file group carries the internal + ``_partial`` truncation marker set by the adaptive-retry give-up sites. + + The marker rides the item dicts up through every chunk merge, so it reaches + ``save_semantic_cache`` on BOTH the incremental checkpoint path (llm.py) and + the final authoritative save (cli.py) without either caller having to thread + an extra argument — the final save would otherwise overwrite a checkpoint's + ``partial`` flag with a clean-looking entry. + """ + for bucket in ("nodes", "edges", "hyperedges"): + for item in group.get(bucket, []): + if isinstance(item, dict) and item.get("_partial"): + return True + return False + + def save_semantic_cache( nodes: list[dict], edges: list[dict], @@ -758,6 +787,7 @@ def save_semantic_cache( mode: str | None = None, prompt: "str | Path | None" = None, prompt_file: "str | Path | None" = None, + partial_source_files: Iterable[str | Path] | None = None, ) -> int: """Save semantic extraction results to cache, keyed by source_file. @@ -782,6 +812,13 @@ def save_semantic_cache( file, but a model must not be able to replace that file's complete cache entry unless the file was part of the current extraction batch (#1757). + When ``partial_source_files`` is provided, entries for those files are + stamped ``partial: True`` — the extraction was truncated, so the entry is + incomplete and :func:`load_cached` must treat it as a miss. Partial-ness is + ALSO detected intrinsically from a ``_partial`` marker on any grouped item, + so the flag survives even when a caller (e.g. cli.py's final save) does not + pass ``partial_source_files``. + ``prompt`` is the extraction prompt that produced these results — text, or a Path to the prompt file. It stamps entries into the p{fingerprint}/ namespace so a later run under a different prompt re-extracts rather than @@ -824,6 +861,10 @@ def save_semantic_cache( if allowed_source_files is not None: allowed_paths = {resolved_source_path(path) for path in allowed_source_files} + partial_paths = None + if partial_source_files is not None: + partial_paths = {resolved_source_path(path) for path in partial_source_files} + def group_skipped(fpath: str) -> bool: """Mirror the write-loop skip condition for one source_file group.""" p = resolved_source_path(fpath) @@ -898,14 +939,32 @@ def save_semantic_cache( # then stamp the result as current-vintage — the exact mixing # #1939 is about, made unfixable because the entry now claims a # prompt that only produced half of it. + # allow_partial=True: a file split into slices across chunks + # accumulates here; if an earlier slice truncated, keep its nodes + # in the union AND let the entry stay partial (the _partial + # markers ride through, so is_partial below re-detects it) rather + # than a later clean slice silently replacing it and promoting the + # half-file to complete. prev = load_cached(p, root, kind=kind, prompt=prompt, - prompt_file=prompt_file, allow_legacy=False) + prompt_file=prompt_file, allow_legacy=False, + allow_partial=True) if prev: result = { "nodes": (prev.get("nodes", []) or []) + result["nodes"], "edges": (prev.get("edges", []) or []) + result["edges"], "hyperedges": (prev.get("hyperedges", []) or []) + result["hyperedges"], } + # A file is partial if the caller named it OR any of its grouped + # items carries the intrinsic ``_partial`` marker. Stamp the flag + # load_cached keys off; copy so the caller's dict is never mutated. + # A later complete extraction overwrites this same content-hash key + # with a non-partial entry that then serves normally. + is_partial = ( + (partial_paths is not None and p in partial_paths) + or _group_has_partial_marker(result) + ) + if is_partial: + result = {**result, "partial": True} save_cached(p, result, root, kind=kind, prompt=prompt, prompt_file=prompt_file) saved += 1 return saved diff --git a/graphify/cli.py b/graphify/cli.py index abf278378..2c18e2e2e 100644 --- a/graphify/cli.py +++ b/graphify/cli.py @@ -58,12 +58,19 @@ def _stamped_manifest_files( files_by_type: dict[str, list[str]], sem_result: dict, root: Path, + partial_source_files: "set[str] | None" = None, ) -> dict[str, list[str]]: """Manifest-safe files dict: only stamp semantic files that actually produced output (cache hit or fresh extraction). Files whose chunk failed have no source_file entry in sem_result — leaving their semantic_hash empty so detect_incremental re-queues them (#933). + A file in ``partial_source_files`` DID produce output this run, but only a + truncated fragment of it, so it is excluded from stamping too — otherwise + detect_incremental would see it "done" and never re-dispatch it, leaving the + incomplete node set live forever on the warm-incremental path. Same #933 + mechanism: leave it unstamped and it is re-queued next run. + Both sides of the membership test are resolved against the scan ``root`` before comparing (#1897): node/edge/hyperedge ``source_file`` values are root-relative on a fresh extraction while ``files_by_type`` entries are @@ -95,11 +102,13 @@ def _stamped_manifest_files( sf = item.get("source_file", "") if sf: sem_extracted.add(_resolve(sf)) + partial_resolved = {_resolve(p) for p in (partial_source_files or set())} sem_types = {"document", "paper", "image"} return { ftype: [ f for f in flist - if ftype not in sem_types or _resolve(f) in sem_extracted + if ftype not in sem_types + or (_resolve(f) in sem_extracted and _resolve(f) not in partial_resolved) ] for ftype, flist in files_by_type.items() } @@ -2568,6 +2577,11 @@ def dispatch_command(cmd: str) -> None: "nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0, } + # Semantic files whose extraction truncated this run. They are left + # unstamped in the manifest so detect_incremental re-queues them next run + # (mirrors the #933 failed-chunk handling); captured below before the + # _partial markers are stripped from the corpus. + _partial_semantic_files: set[str] = set() sem_cache_hits = 0 sem_cache_misses = 0 # Deep mode uses its own namespace (cache/semantic-deep/) so deep and @@ -2673,6 +2687,17 @@ def dispatch_command(cmd: str) -> None: ) except Exception as exc: print(f"[graphify extract] warning: could not write semantic cache: {exc}", file=sys.stderr) + # Record which files truncated (before the markers are stripped) + # so they are left unstamped in the manifest and re-queued on the + # next incremental run. The save above consumed the marker to + # stamp affected cache entries partial: True; strip it before the + # corpus feeds the graph so it never leaks into graph.json. + from graphify.llm import ( + _partial_source_files as _partial_sf, + _strip_partial_markers as _strip_partial, + ) + _partial_semantic_files = set(_partial_sf(fresh)) + _strip_partial(fresh) sem_result["nodes"].extend(fresh.get("nodes", [])) sem_result["edges"].extend(fresh.get("edges", [])) sem_result["hyperedges"].extend(fresh.get("hyperedges", [])) @@ -2752,7 +2777,8 @@ def dispatch_command(cmd: str) -> None: # Path normalization against the scan root happens inside the helper # (#1897) so fresh root-relative source_files match detect()'s # absolute file lists. - _manifest_files = _stamped_manifest_files(files_by_type, sem_result, target) + _manifest_files = _stamped_manifest_files(files_by_type, sem_result, target, + partial_source_files=_partial_semantic_files) # Full-scan manifest saves prune rows for in-root files that left the # scan corpus but still exist on disk (#1908). The corpus must be the diff --git a/graphify/llm.py b/graphify/llm.py index 3ce10f283..3cbaf0eaa 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -1636,6 +1636,50 @@ def _looks_like_context_exceeded(exc: BaseException) -> bool: return any(marker in msg for marker in _CONTEXT_EXCEEDED_MARKERS) +def _mark_partial(result: dict) -> None: + """Tag every node/edge/hyperedge in a truncated chunk result with an internal + ``_partial`` marker. + + A chunk whose LLM response was truncated (`finish_reason="length"`) and could + not be recovered by splitting yields a PARTIAL node set. Left unmarked, that + set is checkpointed and (via the final save) written to the content-hash + semantic cache as authoritative, so it is served forever until the file + content changes or ``--force``. The marker rides these item dicts up through + every chunk merge (which concatenate the same object references) so it reaches + ``save_semantic_cache`` on both the checkpoint and the final-save paths, which + stamp the entry ``partial: True``; ``load_cached`` then treats it as a miss. + """ + for bucket in ("nodes", "edges", "hyperedges"): + for item in result.get(bucket, []): + if isinstance(item, dict): + item["_partial"] = True + + +def _partial_source_files(result: dict) -> list[str]: + """Source files that carry at least one ``_partial`` item in ``result``.""" + seen: set[str] = set() + for bucket in ("nodes", "edges", "hyperedges"): + for item in result.get(bucket, []): + if isinstance(item, dict) and item.get("_partial"): + sf = item.get("source_file") + if sf: + seen.add(str(sf)) + return sorted(seen) + + +def _strip_partial_markers(result: dict) -> None: + """Remove the internal ``_partial`` marker from every item in ``result``. + + Call this only AFTER the semantic cache has been saved (the save consumes the + marker to stamp affected entries ``partial: True``). Stripping it keeps the + internal flag out of the graph.json nodes/edges the corpus result feeds into. + """ + for bucket in ("nodes", "edges", "hyperedges"): + for item in result.get(bucket, []): + if isinstance(item, dict): + item.pop("_partial", None) + + def _extract_with_adaptive_retry( chunk: list[Path], backend: str, @@ -1769,17 +1813,25 @@ def _extract_with_adaptive_retry( return _merge_two([halves[0]], [halves[1]]) print( f"[graphify] single-file chunk {unit_path(chunk[0])} truncated at " - f"max_completion_tokens — partial result kept", + f"max_completion_tokens — partial result kept (not cached as complete)", file=sys.stderr, ) + # The node set is incomplete; mark it so it is not promoted to the + # semantic cache as authoritative and is re-dispatched next run. + _mark_partial(result) return result if _depth >= max_depth: print( f"[graphify] chunk of {len(chunk)} still truncated at recursion " - f"depth {_depth} (max {max_depth}) — partial result kept", + f"depth {_depth} (max {max_depth}) — partial result kept (not cached as complete)", file=sys.stderr, ) + # Conservative: this marks every file in the merged chunk partial, even + # ones that finished cleanly during recursion. Over-marking only costs a + # re-extraction next run; under-marking would serve a truncated file as + # complete, so err toward re-extraction. + _mark_partial(result) return result print( @@ -1939,6 +1991,10 @@ def extract_corpus_parallel( # that changes _EXTRACTION_SYSTEM re-extracts instead of replaying # this vintage forever (#1939). prompt=_extraction_system(deep=deep_mode), + # A truncated/partial chunk must not be checkpointed as + # authoritative: pass the partial file set so its entry is + # stamped ``partial: True`` and re-dispatched next run. + partial_source_files=_partial_source_files(result) or None, ) except Exception as _exc: # noqa: BLE001 — checkpoint is best-effort print(f"[graphify] incremental cache checkpoint failed: {_exc}", file=sys.stderr) diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 19e65efbc..dea9f68f2 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -329,6 +329,33 @@ def test_checkpoint_scopes_cache_writes_to_chunk_files(tmp_path): assert a_cache and any(n["id"] == "a_ok" for n in a_cache["nodes"]) +def test_truncated_chunk_is_cached_partial_and_missed_on_reload(tmp_path): + """A single-file chunk that stays truncated is checkpointed as a PARTIAL + entry, so reloading it is a cache miss (the file re-dispatches next run) + instead of serving the incomplete node set forever.""" + from graphify.llm import extract_corpus_parallel, _extraction_system + from graphify.cache import load_cached + + doc = tmp_path / "doc.md"; doc.write_text("# Heading\nlots of prose\n") + + def truncated(chunk, **kwargs): + return { + "nodes": [{"id": "n1", "source_file": "doc.md", "file_type": "document"}], + "edges": [], "hyperedges": [], + "input_tokens": 1, "output_tokens": 1, + "finish_reason": "length", + } + + with patch("graphify.llm.extract_files_direct", side_effect=truncated): + extract_corpus_parallel( + [doc], backend="kimi", root=tmp_path, + token_budget=None, chunk_size=1, max_concurrency=1, + ) + + # The entry was written but stamped partial, so load_cached rejects it. + assert load_cached(doc, tmp_path, kind="semantic", prompt=_extraction_system()) is None + + def test_checkpoint_writes_deep_namespace_in_deep_mode(tmp_path): """#1894: the per-chunk checkpoint must follow the run's mode — a deep_mode=True run checkpoints into cache/semantic-deep/, leaving the @@ -712,6 +739,68 @@ def test_adaptive_retry_single_file_truncation_does_not_recurse(tmp_path, capsys assert "single-file chunk" in err and "truncated" in err +def test_adaptive_retry_marks_single_file_truncation_partial(tmp_path): + """A non-splittable single-file truncation keeps its partial result but + marks every item ``_partial`` so it is not cached as complete.""" + from graphify.llm import _extract_with_adaptive_retry + + f = tmp_path / "huge.py"; f.write_text("x") + + def stub(chunk, **kwargs): + return _stub_with_finish(len(chunk), finish_reason="length") + + with patch("graphify.llm.extract_files_direct", side_effect=stub): + result = _extract_with_adaptive_retry( + [f], backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + assert result["nodes"], "the partial result should still be returned" + assert all(n.get("_partial") for n in result["nodes"]) + + +def test_adaptive_retry_marks_max_depth_giveup_partial(tmp_path): + """When recursion caps at max_depth with everything still truncated, the + merged partial result is marked ``_partial`` on every item.""" + from graphify.llm import _extract_with_adaptive_retry + + files = [tmp_path / f"f{i}.py" for i in range(8)] + for f in files: + f.write_text("x") + + def stub(chunk, **kwargs): + return _stub_with_finish(len(chunk), finish_reason="length") + + with patch("graphify.llm.extract_files_direct", side_effect=stub): + result = _extract_with_adaptive_retry( + files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=2 + ) + + assert result["nodes"] + assert all(n.get("_partial") for n in result["nodes"]) + + +def test_adaptive_retry_successful_split_is_not_marked_partial(tmp_path): + """A truncation that IS recovered by splitting yields a complete result — + it must NOT carry the partial marker.""" + from graphify.llm import _extract_with_adaptive_retry + + files = [tmp_path / f"f{i}.py" for i in range(4)] + for f in files: + f.write_text("x") + + def stub(chunk, **kwargs): + finish = "length" if len(chunk) == 4 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( + files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + assert result["nodes"] + assert not any(n.get("_partial") for n in result["nodes"]) + + 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_partial_cache.py b/tests/test_partial_cache.py new file mode 100644 index 000000000..9dd5be91d --- /dev/null +++ b/tests/test_partial_cache.py @@ -0,0 +1,141 @@ +"""Tests for partial-extraction cache promotion. + +A truncated LLM chunk (`finish_reason="length"` that could not be recovered by +splitting, or a max-depth adaptive-retry give-up) yields an incomplete node set. +It is tagged with an internal ``_partial`` marker; ``save_semantic_cache`` stamps +that file's entry ``partial: True``, and ``load_cached`` then treats a partial +entry as a cache MISS so the file is re-dispatched instead of served forever. +""" + +from graphify import llm +from graphify.cache import ( + save_semantic_cache, + load_cached, + _group_has_partial_marker, +) + + +def _doc(tmp_path): + doc = tmp_path / "doc.md" + doc.write_text("# Heading\nsome prose\n", encoding="utf-8") + return doc + + +def test_intrinsic_partial_marker_makes_entry_a_cache_miss(tmp_path): + doc = _doc(tmp_path) + nodes = [{"id": "n1", "label": "Heading", "source_file": "doc.md", "_partial": True}] + saved = save_semantic_cache(nodes, [], root=tmp_path, prompt="P") + assert saved == 1 + # The stamped entry is present on disk, but load_cached rejects it. + assert load_cached(doc, root=tmp_path, kind="semantic", prompt="P") is None + + +def test_partial_source_files_arg_stamps_entry(tmp_path): + doc = _doc(tmp_path) + # No intrinsic marker; partial-ness comes only from the explicit arg. + nodes = [{"id": "n1", "label": "Heading", "source_file": "doc.md"}] + save_semantic_cache(nodes, [], root=tmp_path, prompt="P", partial_source_files=["doc.md"]) + assert load_cached(doc, root=tmp_path, kind="semantic", prompt="P") is None + + +def test_non_partial_entry_loads_normally(tmp_path): + doc = _doc(tmp_path) + nodes = [{"id": "n1", "label": "Heading", "source_file": "doc.md"}] + save_semantic_cache(nodes, [], root=tmp_path, prompt="P") + loaded = load_cached(doc, root=tmp_path, kind="semantic", prompt="P") + assert loaded is not None + assert len(loaded["nodes"]) == 1 + + +def test_partial_entry_self_heals_on_complete_reextraction(tmp_path): + doc = _doc(tmp_path) + partial = [{"id": "n1", "source_file": "doc.md", "_partial": True}] + save_semantic_cache(partial, [], root=tmp_path, prompt="P") + assert load_cached(doc, root=tmp_path, kind="semantic", prompt="P") is None + # A later complete extraction overwrites the same content-hash key with a + # non-partial entry, which then serves normally. + complete = [ + {"id": "n1", "source_file": "doc.md"}, + {"id": "n2", "source_file": "doc.md"}, + ] + save_semantic_cache(complete, [], root=tmp_path, prompt="P") + loaded = load_cached(doc, root=tmp_path, kind="semantic", prompt="P") + assert loaded is not None + assert len(loaded["nodes"]) == 2 + + +def test_merge_existing_accumulates_slices_and_stays_partial(tmp_path): + """A file sliced across chunks: an earlier truncated slice must not be + dropped (nor the entry promoted to complete) by a later clean slice's + merge_existing checkpoint. The union keeps both slices and the entry stays + partial until a fully-clean re-extraction overwrites it.""" + doc = _doc(tmp_path) + partial = [{"id": "n1", "source_file": "doc.md", "_partial": True}] + save_semantic_cache(partial, [], root=tmp_path, prompt="P") + fresh = [{"id": "n2", "source_file": "doc.md"}] + save_semantic_cache(fresh, [], root=tmp_path, prompt="P", merge_existing=True) + # Normal read is still a miss: the file had a truncated slice, so it must be + # re-dispatched rather than served. + assert load_cached(doc, root=tmp_path, kind="semantic", prompt="P") is None + # But nothing was lost — both slices are present in the accumulated entry. + peek = load_cached(doc, root=tmp_path, kind="semantic", prompt="P", allow_partial=True) + assert peek is not None + assert {n["id"] for n in peek["nodes"]} == {"n1", "n2"} + + +def test_stamped_manifest_excludes_partial_files(): + """A truncated file produced output this run but is left unstamped in the + manifest (like a failed chunk) so detect_incremental re-queues it.""" + from pathlib import Path + from graphify.cli import _stamped_manifest_files + + files_by_type = {"document": ["a.md", "b.md"], "code": ["x.py"]} + sem_result = { + "nodes": [ + {"id": "1", "source_file": "a.md"}, + {"id": "2", "source_file": "b.md"}, + ], + "edges": [], "hyperedges": [], + } + out = _stamped_manifest_files(files_by_type, sem_result, Path("."), + partial_source_files={"b.md"}) + # a.md extracted cleanly -> stamped; b.md truncated -> excluded; code kept. + assert out["document"] == ["a.md"] + assert out["code"] == ["x.py"] + + +def test_group_has_partial_marker(): + assert _group_has_partial_marker({"nodes": [{"_partial": True}]}) is True + assert _group_has_partial_marker({"edges": [{"_partial": True}]}) is True + assert _group_has_partial_marker({"nodes": [{"id": "a"}], "edges": [], "hyperedges": []}) is False + assert _group_has_partial_marker({}) is False + + +def test_mark_partial_and_partial_source_files(): + result = { + "nodes": [{"id": "a", "source_file": "x.md"}], + "edges": [{"source": "a", "target": "b", "source_file": "x.md"}], + "hyperedges": [{"id": "h", "source_file": "y.md"}], + } + llm._mark_partial(result) + assert result["nodes"][0]["_partial"] is True + assert result["edges"][0]["_partial"] is True + assert result["hyperedges"][0]["_partial"] is True + assert llm._partial_source_files(result) == ["x.md", "y.md"] + + +def test_partial_source_files_empty_when_unmarked(): + result = {"nodes": [{"id": "a", "source_file": "x.md"}], "edges": [], "hyperedges": []} + assert llm._partial_source_files(result) == [] + + +def test_strip_partial_markers_removes_internal_key(): + result = { + "nodes": [{"id": "a", "_partial": True}], + "edges": [{"source": "a", "target": "b", "_partial": True}], + "hyperedges": [{"id": "h", "_partial": True}], + } + llm._strip_partial_markers(result) + assert "_partial" not in result["nodes"][0] + assert "_partial" not in result["edges"][0] + assert "_partial" not in result["hyperedges"][0]