diff --git a/graphify/llm.py b/graphify/llm.py index 36b2afc48..e64c0c6c7 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -11,7 +11,7 @@ import re import sys import time -from collections.abc import Callable +from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, replace from pathlib import Path @@ -462,7 +462,7 @@ def _wrap_untrusted(rel: str, content: str) -> str: ) -def _read_files(units: "list[Path | FileSlice]", root: Path) -> str: +def _read_files(units: "Sequence[Path | FileSlice]", root: Path) -> str: """Return file/slice contents formatted for the extraction prompt. Each unit is wrapped in an delimiter block and known @@ -558,7 +558,7 @@ def _is_vision_image(path: Path) -> bool: def _partition_semantic_files( - units: "list[Path | FileSlice]", + units: "Sequence[Path | FileSlice]", ) -> tuple["list[Path | FileSlice]", list[Path]]: """Split a chunk into (text-like units, raster-image files). @@ -1290,7 +1290,7 @@ def _call_bedrock(model: str, user_message: str, max_tokens: int = 8192, *, deep def extract_files_direct( - files: list[Path], + files: Sequence[Path | FileSlice], backend: str | None = None, api_key: str | None = None, model: str | None = None, @@ -1516,7 +1516,7 @@ def _looks_like_context_exceeded(exc: BaseException) -> bool: def _extract_with_adaptive_retry( - chunk: list[Path], + chunk: Sequence[Path | FileSlice], backend: str, api_key: str | None, model: str | None, @@ -1690,7 +1690,7 @@ def _split_lone_slice() -> "tuple[FileSlice, FileSlice] | None": def extract_corpus_parallel( - files: list[Path], + files: Sequence[str | Path], backend: str = "kimi", api_key: str | None = None, model: str | None = None, @@ -1743,11 +1743,11 @@ def extract_corpus_parallel( # Split oversized splittable documents into slices that cover the whole file # before packing, so content past _FILE_CHAR_CAP is extracted instead of # silently dropped (#1369). Files at/under the cap pass through unchanged. - files = expand_oversized_files(files, _FILE_CHAR_CAP) + units = expand_oversized_files(files, _FILE_CHAR_CAP) if token_budget is not None: - chunks = _pack_chunks_by_tokens(files, token_budget=token_budget) + 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)] + chunks = [units[i:i + chunk_size] for i in range(0, len(units), chunk_size)] merged: dict = { "nodes": [], "edges": [], "hyperedges": [], @@ -1756,7 +1756,9 @@ 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: Sequence[Path | FileSlice] + ) -> tuple[int, dict | None, Exception | None]: t0 = time.time() try: result = _extract_with_adaptive_retry( diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 087464ab8..2f9e3dfc8 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -277,6 +277,33 @@ def record(chunk, **kwargs): assert chunks_seen[0] == 50 +def test_corpus_parallel_accepts_string_paths(tmp_path): + """String file paths should be normalised before slicing and chunking.""" + from graphify.llm import extract_corpus_parallel + + files = [] + for i in range(2): + f = tmp_path / f"f{i}.py" + f.write_text("x = 1\n") + files.append(f) + + chunks_seen = [] + + def record(chunk, **kwargs): + chunks_seen.append(chunk) + return _stub_chunk_result(len(chunk), len(chunks_seen)) + + with patch("graphify.llm.extract_files_direct", side_effect=record): + result = extract_corpus_parallel( + [str(f) for f in files], backend="kimi", max_concurrency=1 + ) + + assert len(chunks_seen) == 1 + assert chunks_seen[0] == files + assert all(isinstance(p, Path) for p in chunks_seen[0]) + assert len(result["nodes"]) == 2 + + # ---- Adaptive retry on truncation ------------------------------------------- def _stub_with_finish(file_count: int, finish_reason: str = "stop") -> dict: