mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-11-10 08:28:22 +00:00
ff2956a41f
* refactor: share backend search, file responses, and vector digests * refactor(studio): remove unused scaffolding and share workspace updates * test: simplify async tests and share deterministic store fixtures
410 lines
15 KiB
Python
410 lines
15 KiB
Python
"""Tests for JsonlFileChunker."""
|
|
|
|
# pylint: disable=protected-access
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import tempfile
|
|
|
|
from reme.components.file_chunker import JsonlFileChunker
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run(coro):
|
|
return asyncio.run(coro)
|
|
|
|
|
|
def _write_jsonl(lines: list[str]) -> str:
|
|
"""Write lines to a temp .jsonl file and return its path."""
|
|
fd, path = tempfile.mkstemp(suffix=".jsonl")
|
|
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
|
for line in lines:
|
|
f.write(line if line.endswith("\n") else line + "\n")
|
|
return path
|
|
|
|
|
|
def _make_records(n: int, width: int = 20) -> list[str]:
|
|
"""Generate *n* JSONL lines, each roughly *width* chars of JSON."""
|
|
return [json.dumps({"id": i, "text": "x" * max(0, width - 20)}) for i in range(n)]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Basic behaviour
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_empty_file():
|
|
"""Empty file → zero chunks."""
|
|
path = _write_jsonl([])
|
|
try:
|
|
chunker = JsonlFileChunker()
|
|
node, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) == 0
|
|
assert node.links == []
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_blank_lines_only():
|
|
"""File with only blank lines → zero chunks."""
|
|
fd, path = tempfile.mkstemp(suffix=".jsonl")
|
|
with os.fdopen(fd, "w") as f:
|
|
f.write("\n\n\n")
|
|
try:
|
|
chunker = JsonlFileChunker()
|
|
_, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) == 0
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_single_line():
|
|
"""One line → one chunk."""
|
|
lines = _make_records(1)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker()
|
|
_, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) == 1
|
|
assert chunks[0].start_line == 1
|
|
assert chunks[0].end_line == 1
|
|
assert lines[0] in chunks[0].text
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_unicode_line_separator_inside_json_string_is_not_a_record_boundary():
|
|
"""U+2028 is valid JSON string content, not a JSONL physical newline."""
|
|
record = json.dumps({"id": 1, "text": "before\u2028\u2028after"}, ensure_ascii=False)
|
|
path = _write_jsonl([record])
|
|
try:
|
|
chunker = JsonlFileChunker(max_lines_per_chunk=1)
|
|
node, chunks = _run(chunker.chunk(path))
|
|
|
|
assert len(chunks) == 1
|
|
assert chunks[0].text == record + "\n"
|
|
assert chunks[0].start_line == chunks[0].end_line == 1
|
|
assert node.chunk_ids == [chunks[0].id]
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_all_lines_fit_in_one_chunk():
|
|
"""Small file fits entirely in one chunk."""
|
|
lines = _make_records(5, width=30)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=5000)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) == 1
|
|
assert chunks[0].start_line == 1
|
|
assert chunks[0].end_line == 5
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_max_lines_per_chunk_one():
|
|
"""max_lines_per_chunk=1 emits exactly one complete line per chunk."""
|
|
lines = _make_records(5, width=10)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=5000, max_lines_per_chunk=1)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) == len(lines)
|
|
for line_number, (line, chunk) in enumerate(zip(lines, chunks), start=1):
|
|
assert chunk.start_line == line_number
|
|
assert chunk.end_line == line_number
|
|
assert chunk.text == line + "\n"
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Multi-chunk splitting (no overlap)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_multiple_chunks_no_overlap():
|
|
"""Large file is split into multiple chunks with no overlap."""
|
|
# Each line ≈ 32 chars (with \n), 10 lines ≈ 320 chars.
|
|
lines = _make_records(10, width=30)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=100, max_overlap_chars=0)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) > 1, f"Expected >1 chunks, got {len(chunks)}"
|
|
# No overlap: consecutive chunks don't share lines.
|
|
for i in range(len(chunks) - 1):
|
|
assert chunks[i].end_line < chunks[i + 1].start_line
|
|
# All lines are covered.
|
|
assert chunks[0].start_line == 1
|
|
assert chunks[-1].end_line == 10
|
|
print(f" Created {len(chunks)} chunks")
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_line_aligned_no_intra_line_split():
|
|
"""Every chunk boundary falls on a line boundary."""
|
|
lines = _make_records(20, width=40)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=150, max_overlap_chars=0)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) > 1
|
|
for c in chunks:
|
|
# Each chunk text should be a concatenation of whole lines.
|
|
text_lines = c.text.splitlines()
|
|
assert len(text_lines) == (c.end_line - c.start_line + 1)
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Overlap behaviour
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_overlap_lines_shared():
|
|
"""Overlap lines appear at the end of chunk N and the start of chunk N+1."""
|
|
# Each line is ~32 chars. max_chars=100 → ~3 lines per chunk.
|
|
# max_overlap_chars=40 → ~1 line of overlap.
|
|
lines = _make_records(12, width=30)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=100, max_overlap_chars=40)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
assert len(chunks) > 1
|
|
# Verify overlap: chunk[i].end_line >= chunk[i+1].start_line
|
|
for i in range(len(chunks) - 1):
|
|
assert chunks[i].end_line >= chunks[i + 1].start_line, (
|
|
f"Chunk {i} end_line={chunks[i].end_line} < " f"chunk {i+1} start_line={chunks[i+1].start_line}"
|
|
)
|
|
# Last chunk still reaches the end of the file.
|
|
assert chunks[-1].end_line == 12
|
|
print(f" Created {len(chunks)} chunks with overlap")
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_zero_overlap():
|
|
"""max_overlap_chars=0 produces strictly non-overlapping chunks."""
|
|
lines = _make_records(10, width=30)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=100, max_overlap_chars=0)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
for i in range(len(chunks) - 1):
|
|
assert chunks[i].end_line < chunks[i + 1].start_line
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_overlap_capped_by_max_overlap_chars():
|
|
"""Overlap never exceeds max_overlap_chars in size."""
|
|
lines = _make_records(10, width=50)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
max_overlap = 60 # Should allow at most 1 line of overlap (~51 chars).
|
|
chunker = JsonlFileChunker(max_chars=200, max_overlap_chars=max_overlap)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
if len(chunks) > 1:
|
|
for i in range(len(chunks) - 1):
|
|
overlap_start = chunks[i + 1].start_line
|
|
overlap_end = chunks[i].end_line
|
|
if overlap_start <= overlap_end:
|
|
overlap_lines = list(range(overlap_start, overlap_end + 1))
|
|
overlap_text = "".join(lines[ln - 1] for ln in overlap_lines)
|
|
assert len(overlap_text) <= max_overlap, f"Overlap size {len(overlap_text)} > {max_overlap}"
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Oversized single line
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_single_line_exceeds_max():
|
|
"""A line larger than max_chars is emitted alone (no intra-line split)."""
|
|
long_line = json.dumps({"data": "A" * 200})
|
|
lines = ["short", long_line, "short"]
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=64, max_overlap_chars=0)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
# The long line must appear in exactly one chunk by itself (or with
|
|
# no split).
|
|
found = False
|
|
for c in chunks:
|
|
if long_line in c.text:
|
|
found = True
|
|
break
|
|
assert found, "Long line was not found in any chunk"
|
|
assert len(chunks) >= 2 # At least some splitting happened.
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Byte mode
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_byte_mode():
|
|
"""mode='bytes' counts byte length instead of char length."""
|
|
# CJK characters are 3 bytes each in UTF-8 but 1 char each.
|
|
# Use ensure_ascii=False so raw CJK chars appear in the file.
|
|
cjk_record = json.dumps({"text": "你" * 10}, ensure_ascii=False)
|
|
# The JSON string itself has CJK chars; after writing to file each line
|
|
# has ~13 ASCII chars ({"text": "..."}) + 10 CJK chars.
|
|
# char length ≈ 23, byte length ≈ 13 + 10*3 = 43.
|
|
lines = [cjk_record] * 6
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
# In char mode: each line ≈ 24 chars → 2 lines per chunk at max=50.
|
|
chunker_chars = JsonlFileChunker(max_chars=50, max_overlap_chars=0, mode="chars")
|
|
_, chunks_chars = _run(chunker_chars.chunk(path))
|
|
|
|
# In byte mode: each line ≈ 44 bytes → 1 line per chunk at max=50.
|
|
chunker_bytes = JsonlFileChunker(max_chars=50, max_overlap_chars=0, mode="bytes")
|
|
_, chunks_bytes = _run(chunker_bytes.chunk(path))
|
|
|
|
# Byte mode should produce more chunks since each line is larger in bytes.
|
|
assert len(chunks_bytes) > len(chunks_chars), (
|
|
f"Byte mode ({len(chunks_bytes)}) should produce more chunks "
|
|
f"than char mode ({len(chunks_chars)}) for CJK text"
|
|
)
|
|
print(f" char mode: {len(chunks_chars)} chunks, byte mode: {len(chunks_bytes)} chunks")
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_invalid_mode_raises():
|
|
"""An invalid mode string raises ValueError."""
|
|
try:
|
|
JsonlFileChunker(mode="words")
|
|
assert False, "Should have raised ValueError"
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FileNode / FileChunk properties
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_file_chunk_properties():
|
|
"""Each FileChunk has correct path, line range, text, and hash id."""
|
|
lines = _make_records(5, width=30)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=5000)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
c = chunks[0]
|
|
assert c.start_line >= 1
|
|
assert c.end_line >= c.start_line
|
|
assert c.id and len(c.id) > 0
|
|
assert c.text.strip()
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_file_node_properties():
|
|
"""FileNode carries path, st_mtime, chunk_ids, and empty links."""
|
|
lines = _make_records(3)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker()
|
|
node, chunks = _run(chunker.chunk(path))
|
|
assert node.st_mtime > 0
|
|
assert node.links == []
|
|
assert node.chunk_ids == [c.id for c in chunks]
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
def test_hash_id_deterministic():
|
|
"""Same file produces the same chunk hash ids on repeated calls."""
|
|
lines = _make_records(5, width=30)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
c1 = JsonlFileChunker(max_chars=100)
|
|
_, chunks1 = _run(c1.chunk(path))
|
|
c2 = JsonlFileChunker(max_chars=100)
|
|
_, chunks2 = _run(c2.chunk(path))
|
|
assert [c.id for c in chunks1] == [c.id for c in chunks2]
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Edge cases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_max_chars_floor():
|
|
"""max_chars is clamped to at least 64."""
|
|
c = JsonlFileChunker(max_chars=1)
|
|
assert c.max_chars == 64
|
|
|
|
|
|
def test_max_overlap_floor():
|
|
"""Negative max_overlap_chars is clamped to 0."""
|
|
c = JsonlFileChunker(max_overlap_chars=-10)
|
|
assert c.max_overlap_chars == 0
|
|
|
|
|
|
def test_full_coverage():
|
|
"""Every line of the file appears in at least one chunk."""
|
|
lines = _make_records(25, width=40)
|
|
path = _write_jsonl(lines)
|
|
try:
|
|
chunker = JsonlFileChunker(max_chars=150, max_overlap_chars=50)
|
|
_, chunks = _run(chunker.chunk(path))
|
|
covered: set[int] = set()
|
|
for c in chunks:
|
|
covered.update(range(c.start_line, c.end_line + 1))
|
|
assert covered == set(range(1, 26)), f"Missing lines: {set(range(1, 26)) - covered}"
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Internal algorithm
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_chunk_lines_algorithm_simple():
|
|
"""_chunk_lines produces correct ranges for a simple case."""
|
|
# max_chars floor is 64, so lines must be long enough.
|
|
# Each line = 20 chars + \n = 21 chars. 64 / 21 = 3 lines per chunk.
|
|
chunker = JsonlFileChunker(max_chars=64, max_overlap_chars=0)
|
|
lines = ["a" * 20 + "\n", "b" * 20 + "\n", "c" * 20 + "\n", "d" * 20 + "\n"]
|
|
# 3 lines = 63 chars ≤ 64 → chunk 1: (0,3). 4th line alone → chunk 2: (3,4).
|
|
ranges = chunker._chunk_lines(lines)
|
|
assert ranges == [(0, 3), (3, 4)], f"Got {ranges}"
|
|
|
|
|
|
def test_chunk_lines_algorithm_with_overlap():
|
|
"""_chunk_lines with overlap produces overlapping ranges."""
|
|
# Each line = 20 chars + \n = 21 chars. max_chars=64 → 3 lines per chunk.
|
|
# max_overlap_chars=22 → 1 line of overlap (21 chars ≤ 22).
|
|
chunker = JsonlFileChunker(max_chars=64, max_overlap_chars=22)
|
|
lines = ["a" * 20 + "\n"] * 6
|
|
ranges = chunker._chunk_lines(lines)
|
|
assert len(ranges) >= 2, f"Expected ≥2 ranges, got {ranges}"
|
|
# Verify overlap: range[i].end >= range[i+1].start
|
|
for i in range(len(ranges) - 1):
|
|
assert ranges[i][1] >= ranges[i + 1][0], f"No overlap between range {i} {ranges[i]} and {i+1} {ranges[i+1]}"
|
|
# Last range reaches end.
|
|
assert ranges[-1][1] == 6
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runner
|
|
# ---------------------------------------------------------------------------
|