Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 12 additions & 10 deletions graphify/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 <untrusted_source> delimiter block and known
Expand Down Expand Up @@ -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).

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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": [],
Expand All @@ -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(
Expand Down
27 changes: 27 additions & 0 deletions tests/test_chunking.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down