ReMe/tests/unit/test_source_format_merge.py
xyf2020 630f26b119
feat(search): scoped dedup, session-chunk merge, and unified recall formatting (#384)
* feat(search): add tool_context-scoped chunk dedup with TTL

Introduce _ToolContextDedupMixin shared by search/vector_search/bm25_search
to skip already-seen chunks within one agent tool_context. Per-context state
lives in app_context.metadata with configurable TTL (default 24h).

* feat(search): unify chunk answer rendering with merge and explicit empty messages

- Refactor SearchStep/VectorSearchStep/Bm25SearchStep to share format_chunks_answer for consistent source rendering and adjacent session-chunk merging.

- Distinguish empty results: ALL_RETURNED_MESSAGE when dedup removes everything vs NO_RESULTS_MESSAGE when nothing matched.

- Bump JsonlFileChunker default max_chars to 4000.

- Add unit tests for source-format merge and empty-result messages.

* refactor(config): reorganize file_chunker components and move jsonl max_chars into config

- Register explicit markdown/json/jsonl chunkers in beam.yaml and lme.yaml with markdown options (embed_toc, max_ast_sections, frontmatter handling) and jsonl max_chars=4000.

- Restrict default chunker to txt/log extensions.

- Revert JsonlFileChunker code default max_chars back to 2000; the 4000 value now lives in config.

* chore(benchmark): increase longmemeval num_items from 64 to 500

* refactor(search): split SearchStep into simplified and v2 variants, extract counter utility

- Extract global_counter_next from ApplicationContext into reme/utils/counter.py
  as a standalone function operating on metadata dict with lazy initialization.

- Split SearchStep into two variants:
  - SearchStep (simplified): inline chunk.id dedup, single-branch vector/keyword
    optimization based on vector_weight, inline answer formatting.
  - SearchV2Step (full): preserves _ToolContextDedupMixin with interval-subset-aware
    dedup and format_chunks_answer with session-aware chunk merging.

- Update beam.yaml and lme.yaml to use search_v2_step for benchmark jobs.

- Rename existing search tests to test_search_v2_step_* and add new
  test_search_step_* tests covering the simplified variant.

* fix: normalise missing trailing newline in _build_union_chunk to prevent line collision

* refactor: lazy-init counter tree in ApplicationContext metadata

- Remove hardcoded _counter_tree and _counter_tree_lock initialization
  from ApplicationContext.metadata; rely on lazy initialization in
  reme.utils.counter.global_counter_next on first call
- Set longmemeval num_items back to 500
- Remove obsolete trailing-newline collision tests

---------

Co-authored-by: sa-buc <jiangniurou.xyf@dail-algo011164204033.ET135>
2026-07-22 17:17:23 +08:00

141 lines
6.1 KiB
Python

"""Unit tests for session-chunk merging in ``format_chunks_answer``.
Session chunks (``*.jsonl`` under the dialog dir) whose line ranges overlap,
contain one another, or are adjacent are merged into their union before
rendering. Bodies here are plain (non-``Msg``) text, so ``render_chunk_body``
falls back to the raw chunk text — letting these tests assert the union
content and line order directly.
"""
from reme.schema import FileChunk
from reme.steps.index._source_format import format_chunks_answer
_DIALOG_DIR = "session"
def _chunk(start: int, end: int, text: str, score: float = 1.0, path: str = "session/s1.jsonl") -> FileChunk:
return FileChunk(path=path, start_line=start, end_line=end, text=text, scores={"score": score})
def test_overlapping_session_chunks_merge_into_union_without_duplicates():
"""Overlapping ranges collapse to one passage; the shared line is shown once, in order."""
a = _chunk(1, 3, "L1\nL2\nL3\n")
b = _chunk(3, 5, "L3\nL4\nL5\n")
answer = format_chunks_answer([a, b], _DIALOG_DIR, include_source=False)
assert answer == "L1\nL2\nL3\nL4\nL5\n"
def test_contained_session_chunk_is_absorbed_by_the_larger_range():
"""When one range fully contains another, only the union (the larger) is shown."""
big = _chunk(1, 5, "L1\nL2\nL3\nL4\nL5\n")
small = _chunk(2, 4, "L2\nL3\nL4\n")
answer = format_chunks_answer([big, small], _DIALOG_DIR, include_source=False)
assert answer == "L1\nL2\nL3\nL4\nL5\n"
def test_adjacent_session_chunks_merge_end_plus_one_equals_next_start():
"""Gap-free consecutive chunks (prev.end + 1 == next.start) merge into one union."""
a = _chunk(1, 3, "L1\nL2\nL3\n")
b = _chunk(4, 6, "L4\nL5\nL6\n")
answer = format_chunks_answer([a, b], _DIALOG_DIR, include_source=False)
assert answer == "L1\nL2\nL3\nL4\nL5\nL6\n"
def test_session_chunks_with_a_gap_are_not_merged():
"""A missing line between ranges (start > end + 1) keeps the chunks separate."""
a = _chunk(1, 3, "L1\nL2\nL3\n")
b = _chunk(5, 6, "L5\nL6\n") # line 4 missing -> not adjacent
answer = format_chunks_answer([a, b], _DIALOG_DIR, include_source=False)
assert answer == "L1\nL2\nL3\n\n\nL5\nL6\n"
def test_session_chunks_from_different_files_are_not_merged():
"""Overlapping ranges in different session files must stay separate."""
a = _chunk(1, 3, "A1\nA2\nA3\n", path="session/s1.jsonl")
b = _chunk(2, 4, "B2\nB3\nB4\n", path="session/s2.jsonl")
answer = format_chunks_answer([a, b], _DIALOG_DIR, include_source=False)
assert answer == "A1\nA2\nA3\n\n\nB2\nB3\nB4\n"
def test_non_session_chunks_are_never_merged():
"""Non-transcript chunks (not ``*.jsonl`` under the dialog dir) pass through untouched."""
a = _chunk(1, 3, "M1\nM2\nM3\n", path="daily/a.md")
b = _chunk(2, 4, "M2\nM3\nM4\n", path="daily/a.md")
answer = format_chunks_answer([a, b], _DIALOG_DIR, include_source=False)
assert answer == "M1\nM2\nM3\n\n\nM2\nM3\nM4\n"
def test_merge_preserves_line_order_regardless_of_input_rank_order():
"""A later, higher-ranked chunk does not reorder union content; lines stay chronological."""
later = _chunk(3, 5, "L3\nL4\nL5\n", score=9.0)
earlier = _chunk(1, 3, "L1\nL2\nL3\n", score=1.0)
# Higher-scored later-range chunk is listed first (as a ranker would).
answer = format_chunks_answer([later, earlier], _DIALOG_DIR, include_source=False)
assert answer == "L1\nL2\nL3\nL4\nL5\n"
def test_merged_header_spans_the_union_range_and_keeps_best_score():
"""With source headers, the merged entry reports the union range and the top score."""
a = _chunk(1, 3, "L1\nL2\nL3\n", score=2.0)
b = _chunk(3, 5, "L3\nL4\nL5\n", score=7.0)
answer = format_chunks_answer([a, b], _DIALOG_DIR, include_source=True)
assert answer.count("==========") == 2 # exactly one header (open + close markers)
assert "session/s1.jsonl:1-5" in answer
assert "score=7.0000" in answer
def test_separate_intervals_in_same_file_stay_separate():
"""Two disjoint interval clusters in one file yield two merged units, ordered by line."""
a = _chunk(1, 2, "L1\nL2\n")
b = _chunk(3, 4, "L3\nL4\n") # adjacent to a -> merges with a into 1-4
c = _chunk(10, 11, "L10\nL11\n") # far away -> separate
answer = format_chunks_answer([a, b, c], _DIALOG_DIR, include_source=False)
assert answer == "L1\nL2\nL3\nL4\n\n\nL10\nL11\n"
def test_same_file_units_stay_adjacent_and_sorted_even_when_interleaved_by_rank():
"""Two disjoint units of one session file are grouped together and ordered by
``start_line``, even when a different file is ranked between them and the
lower interval was ranked last."""
s1_high = _chunk(10, 12, "S1x\nS1y\nS1z\n", score=9.0, path="session/s1.jsonl")
s2_mid = _chunk(1, 3, "S2a\nS2b\nS2c\n", score=5.0, path="session/s2.jsonl")
s1_low = _chunk(1, 3, "S1a\nS1b\nS1c\n", score=1.0, path="session/s1.jsonl")
# Rank order (as a ranker would emit, by score desc): s1[10-12], s2[1-3], s1[1-3].
answer = format_chunks_answer([s1_high, s2_mid, s1_low], _DIALOG_DIR, include_source=False)
# s1's two units are adjacent and sorted by start_line (1-3 before 10-12),
# placed at s1's earliest rank (0), so the whole s1 block precedes s2.
assert answer == "S1a\nS1b\nS1c\n\n\nS1x\nS1y\nS1z\n\n\nS2a\nS2b\nS2c\n"
def test_same_file_units_adjacency_with_source_headers():
"""Header view: same-file units are contiguous and ascending; other files follow."""
s1_high = _chunk(10, 12, "S1x\nS1y\nS1z\n", score=9.0, path="session/s1.jsonl")
s2_mid = _chunk(1, 3, "S2a\nS2b\nS2c\n", score=5.0, path="session/s2.jsonl")
s1_low = _chunk(1, 3, "S1a\nS1b\nS1c\n", score=1.0, path="session/s1.jsonl")
answer = format_chunks_answer([s1_high, s2_mid, s1_low], _DIALOG_DIR, include_source=True)
headers = [line for line in answer.splitlines() if line.startswith("==========")]
assert headers[0].split(" [")[0] == "========== session/s1.jsonl:1-3"
assert headers[1].split(" [")[0] == "========== session/s1.jsonl:10-12"
assert headers[2].split(" [")[0] == "========== session/s2.jsonl:1-3"