scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-08-28 05:25:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 11

feat(search): scoped dedup, session-chunk merge, and unified recall formatting (#384)

Browse source 
* 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>
This commit is contained in:
xyf2020 2026-07-22 17:17:23 +08:00 • committed by GitHub
parent 7b1da5a9ee
commit 630f26b119
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 1407 additions and 125 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
benchmark/longmemeval/config.yaml
View file

@ -4,7 +4,7 @@
dataset:
path: "benchmark/datasets/longmemeval/longmemeval_s_reme_cleaned.json"
start_index: 0 # first item index
num_items: 500 # how many items to evaluate (starting from start_index)
num_items: 500 # how many items to evaluate (starting from start_index)
max_sessions: 0 # 0 = all sessions; >0 = limit sessions per item for testing
question_types: [] # filter by question_type; empty list = no filtering (all types)
workspace_root: "benchmark/memory_workspaces/longmemeval-s" # workspace root for item workspaces

17
reme/config/beam.yaml
View file

@ -215,7 +215,7 @@ jobs:
required:
- query
steps:
- backend: search_step
- backend: search_v2_step
vector_weight: 0.7
candidate_multiplier: 5.0
expand_links: true
@ -567,10 +567,21 @@ components:
file_chunker:
markdown:
backend: markdown
supported_extensions: ["md"]
supported_extensions: [ "md" ]
embed_toc: true
max_ast_sections: 100
include_frontmatter_in_metadata: false
include_frontmatter_keys_in_metadata: [] # empty = all non-empty frontmatter keys
json:
backend: json
supported_extensions: [ "json" ]
jsonl:
backend: jsonl
supported_extensions: [ "jsonl" ] # noqa: keep #314 chunker scope intact after #325
max_chars: 4000
default:
backend: default
supported_extensions: ["jsonl"]
supported_extensions: ["txt","log"]
keyword_index:
default:

17
reme/config/lme.yaml
View file

@ -213,7 +213,7 @@ jobs:
required:
- query
steps:
- backend: search_step
- backend: search_v2_step
vector_weight: 0.7
candidate_multiplier: 5.0
expand_links: true
@ -564,10 +564,21 @@ components:
file_chunker:
markdown:
backend: markdown
supported_extensions: ["md"]
supported_extensions: [ "md" ]
embed_toc: true
max_ast_sections: 100
include_frontmatter_in_metadata: false
include_frontmatter_keys_in_metadata: [] # empty = all non-empty frontmatter keys
json:
backend: json
supported_extensions: [ "json" ]
jsonl:
backend: jsonl
supported_extensions: [ "jsonl" ] # noqa: keep #314 chunker scope intact after #325
max_chars: 4000
default:
backend: default
supported_extensions: ["jsonl"]
supported_extensions: ["txt","log"]
keyword_index:
default:

33
reme/steps/benchmark/base/agentic_answer.py
View file

@ -1,26 +1,11 @@
"""Shared base class for benchmark agentic-answer steps."""
import os
import threading
from ...base_step import BaseStep
from ...index._dedup import _ToolContextDedupMixin
from ....enumeration import ChunkEnum
# ---------------------------------------------------------------------------
# Process-safe & thread-safe counter for unique tool_context_id.
# PID guarantees cross-process uniqueness (multiprocessing Pool);
# threading.Lock + monotonic counter guarantees thread safety within a process.
# ---------------------------------------------------------------------------
_TOOL_CTX_LOCK = threading.Lock()
_TOOL_CTX_SEQ = 0
def _next_tool_context_id(prefix: str) -> str:
global _TOOL_CTX_SEQ
with _TOOL_CTX_LOCK:
_TOOL_CTX_SEQ += 1
seq = _TOOL_CTX_SEQ
return f"{prefix}_{os.getpid()}_{seq}"
from ....utils.counter import global_counter_next
class BaseAgenticAnswerStep(BaseStep):
@ -41,7 +26,7 @@ class BaseAgenticAnswerStep(BaseStep):
"""
MAX_ITERATION = 10
TOOL_CONTEXT_PREFIX: str = "agentic_answer"
TOOL_CONTEXT_PREFIX: str = "content_agentic_answer"
async def execute(self):
assert self.context is not None
@ -58,11 +43,18 @@ class BaseAgenticAnswerStep(BaseStep):
if query_time:
sys_prompt += "\n" + self.prompt_format("temporal_hint", query_time=query_time)
if self.app_context is not None:
tool_context_id = (
f"{self.TOOL_CONTEXT_PREFIX}_{os.getpid()}_"
f"{global_counter_next(self.app_context.metadata, [self.TOOL_CONTEXT_PREFIX])}"
)
else:
tool_context_id = f"{self.TOOL_CONTEXT_PREFIX}_{os.getpid()}_local"
wrapper_kwargs = {
"system_prompt": sys_prompt,
"job_tools": ["search", "add_draft", "read_all_draft"],
"react_config": {"max_iters": self.MAX_ITERATION},
"tool_context_id": _next_tool_context_id(self.TOOL_CONTEXT_PREFIX),
"tool_context_id": tool_context_id,
}
if self.context.stream:
@ -83,6 +75,9 @@ class BaseAgenticAnswerStep(BaseStep):
"response": text,
},
)
if self.app_context is not None:
self.app_context.metadata.get(_ToolContextDedupMixin.TOOL_CONTEXTS_KEY, {}).pop(tool_context_id, None)
return self.context.response
async def _stream_reply(self, query: str, **wrapper_kwargs) -> str:

2
reme/steps/index/__init__.py
View file

@ -8,6 +8,7 @@ from .log_changes import LogChangesStep
from .node_search import NodeSearchStep
from .init_changes import InitChangesStep
from .search import SearchStep
from .search_v2 import SearchV2Step
from .traverse import TraverseStep
from .update_changes import ChangeApplyStep, UpdateCatalogStep, UpdateIndexStep
from .vector_search import VectorSearchStep
@ -33,6 +34,7 @@ __all__ = [
"NodeSearchStep",
"ReadAllDraftStep",
"SearchStep",
"SearchV2Step",
"TraverseStep",
"UpdateCatalogStep",
"UpdateIndexStep",

160
reme/steps/index/_dedup.py Normal file
View file

@ -0,0 +1,160 @@
"""Shared tool_context-scoped chunk dedup with TTL.
Used by ``search``/``vector_search``/``bm25_search`` to avoid returning the
same content twice within one agent tool_context. Per-context state lives at
``app_context.metadata["tool_contexts"][tool_context_id]["search_seen_chunk_ids"]``
as ``{path: [(start_line, end_line, timestamp), ...]}``; a chunk is skipped when
its ``[start_line, end_line]`` is fully covered by the union of seen entries
(merged overlapping/adjacent intervals) for the same ``path``. Entries older
than ``seen_ttl_hours`` are expired on each call.
When ``app_context`` is absent the same structure is mirrored under
``self.kwargs["tool_contexts"][tool_context_id]`` for unit tests.
"""
import datetime
from typing import TYPE_CHECKING, Any, Callable, Final
from ...schema import FileChunk
if TYPE_CHECKING:
from ...components import ApplicationContext
class _ToolContextDedupMixin:
"""Mixin providing tool_context-scoped chunk dedup with TTL.
Must be mixed into a ``BaseStep`` subclass (e.g. ``SearchStep``); it
cannot be instantiated or subclassed on its own. The mixin relies on
``app_context``/``kwargs`` from ``BaseStep`` and on ``seen_ttl_hours``
set by the host step's ``__init__`` (default 24h).
"""
TOOL_CONTEXTS_KEY: Final[str] = "tool_contexts"
SEARCH_SEEN_KEY: Final[str] = "search_seen_chunk_ids"
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
# Deferred import avoids a circular dependency at module load time.
from ..base_step import BaseStep
if not issubclass(cls, BaseStep):
raise TypeError(
f"{cls.__name__!r} mixes in _ToolContextDedupMixin but does not "
f"inherit from BaseStep. Mix it in alongside BaseStep, e.g. "
f"class {cls.__name__}(_ToolContextDedupMixin, BaseStep).",
)
def __new__(cls, *args, **kwargs):
if cls is _ToolContextDedupMixin:
raise TypeError(
"_ToolContextDedupMixin is a mixin and cannot be instantiated "
"directly. Mix it into a BaseStep subclass, e.g. "
"class SearchStep(_ToolContextDedupMixin, BaseStep).",
)
return super().__new__(cls, *args, **kwargs)
if TYPE_CHECKING:
# Declared by BaseStep; repeated here so static analysis resolves
# attribute access on the mixin without inheriting BaseStep.
app_context: "ApplicationContext | None"
kwargs: dict[str, Any]
seen_ttl_hours: float
def _tool_context_store(self, tool_context_id: str) -> dict:
"""Return the mutable state bucket for a tool context.
The bucket is created lazily on first access and lives at
``metadata["tool_contexts"][tool_context_id]`` (or the same path under
``kwargs`` when no ``app_context`` is available, e.g. in unit tests).
"""
if self.app_context is not None:
contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {})
else:
contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {})
return contexts.setdefault(tool_context_id, {})
@staticmethod
def _now_ts() -> float:
return datetime.datetime.now().timestamp()
def _dedupe_tool_context(
self,
chunks: list[FileChunk],
tool_context_id: str,
limit: int,
*,
clock: Callable[[], float] | None = None,
ttl_override: float | None = None,
) -> tuple[list[FileChunk], dict]:
"""Drop chunks whose line range is already covered by a previously
returned chunk for this tool_context within the TTL window.
Seen intervals per path are merged (overlapping or adjacent) into a
minimal set of disjoint ranges. A chunk is skipped when its
``[start_line, end_line]`` is a subset of any merged range for the
same ``path`` — multiple previously returned chunks can jointly cover
a new chunk even if no single entry does. Partial overlap (superset or
straddle) is NOT skipped — the chunk carries lines not yet returned,
so it is kept.
``clock`` (a zero-arg callable returning a float timestamp) and
``ttl_override`` (seconds) allow tests to inject deterministic time.
Returns ``(returned, stats)``; callers that don't need the stats may
discard the second element.
"""
now = (clock or self._now_ts)()
ttl = ttl_override if ttl_override is not None else float(self.seen_ttl_hours) * 60 * 60
store = self._tool_context_store(tool_context_id)
seen = store.get(self.SEARCH_SEEN_KEY, {})
# Normalize legacy in-memory formats to {path: [(s, e, t), ...]}.
# Older shapes ({chunk_id: timestamp} or a plain list of ids) cannot
# be migrated because chunk_id is an opaque hash; seen is a transient
# per-Application cache, so dropping it is safe.
if not isinstance(seen, dict) or (seen and all(not isinstance(v, list) for v in seen.values())):
seen = {}
before_expire = sum(len(v) for v in seen.values())
# Expire stale tuples across all paths.
seen = {path: [(s, e, t) for (s, e, t) in entries if now - t < ttl] for path, entries in seen.items()}
seen = {path: entries for path, entries in seen.items() if entries}
store[self.SEARCH_SEEN_KEY] = seen
seen_before = sum(len(v) for v in seen.values())
def _is_covered(chunk: FileChunk) -> bool:
entries = seen.get(chunk.path)
if not entries:
return False
# Merge overlapping/adjacent intervals so that multiple seen
# entries can jointly cover a new chunk (e.g. (1,10)+(11,20)
# merge into (1,20) and cover (5,15)).
intervals = sorted((s, e) for s, e, _ in entries)
merged: list[tuple[int, int]] = []
for s, e in intervals:
if merged and s <= merged[-1][1] + 1:
merged[-1] = (merged[-1][0], max(merged[-1][1], e))
else:
merged.append((s, e))
return any(s <= chunk.start_line and chunk.end_line <= e for s, e in merged)
unvisited = [chunk for chunk in chunks if not _is_covered(chunk)]
returned = unvisited[:limit]
for chunk in returned:
seen.setdefault(chunk.path, []).append((chunk.start_line, chunk.end_line, now))
# Reorder for readability: keep chunks of the same path adjacent and sorted by
# ascending start_line; order paths by where each first appears in the original
# sequence (the path owning the earliest-ranked chunk comes first).
path_order: dict[str, int] = {}
for idx, chunk in enumerate(returned):
path_order.setdefault(chunk.path, idx)
returned = sorted(returned, key=lambda c: (path_order[c.path], c.start_line))
return returned, {
"tool_context_id": tool_context_id,
"seen_before": seen_before,
"skipped_seen": len(chunks) - len(unvisited),
"seen_after": sum(len(v) for v in seen.values()),
"expired": before_expire - seen_before,
"ttl_seconds": ttl,
}

202
reme/steps/index/_source_format.py
View file

@ -1,29 +1,193 @@
"""Shared helper: render search results with each hit's originating session_id.
"""Shared helpers: render retrieved chunks and assemble search-step answers.
Plain ``vector_search``/``bm25_search`` return only ``chunk.text``. The
LongMemEval agentic-answer flow needs each hit's ``session_id`` so the agent can
pivot back to the raw session via ``extract_session_by_id``. This helper reads
that ``session_id`` from the note's frontmatter and prefixes it to the text.
Raw session transcripts (``*.jsonl`` under the dialog dir) store one serialized
``Msg`` per line. :func:`render_chunk_body` turns those back into a readable
dialog; all other chunks keep their raw ``text``. Used by
``search``/``vector_search``/``bm25_search`` so every step renders session hits
identically.
:func:`format_chunks_answer` assembles a full answer string from a list of
chunks, with optional per-chunk score formatting and link expansion. Raw
session chunks from the same file whose line ranges overlap, contain one
another, or are adjacent are merged into their union before rendering (see
:func:`_merge_session_chunk_intervals`) so a passage is never shown twice.
:data:`ALL_RETURNED_MESSAGE` is the English notice shown when tool_context dedup
removes every previously-returned result. :data:`NO_RESULTS_MESSAGE` is the
English notice shown when the search returned no results at all.
"""
from pathlib import Path
from typing import Callable, Final
import frontmatter
from agentscope.message import Msg
from ..evolve._evolve import format_history
from ...schema import FileChunk
from ...utils.link_expansion import render_expansion_lines
#: English message written to ``response.answer`` when tool_context dedup
#: removed every result that was already returned in previous responses.
ALL_RETURNED_MESSAGE: Final[str] = (
"All retrieved content has already been returned in previous responses; " "no new content was found."
)
#: English message written to ``response.answer`` when the search returned no
#: results at all (before dedup).
NO_RESULTS_MESSAGE: Final[str] = "No relevant information was found for the given query."
def render_with_source(chunks: list[FileChunk], workspace_path: Path) -> str:
"""Render each chunk as ``[session_id=<sid>]`` header + text."""
def is_session_chunk(chunk: FileChunk, dialog_dir: str) -> bool:
"""True if the chunk comes from a raw session transcript (a jsonl file under the dialog dir)."""
path = (chunk.path or "").strip().strip("/")
if not path.endswith(".jsonl"):
return False
dialog_dir = (dialog_dir or "").strip("/")
return path == dialog_dir or path.startswith(f"{dialog_dir}/")
def render_chunk_body(chunk: FileChunk, dialog_dir: str) -> str:
"""Render a chunk's body, compacting raw session transcripts into a readable form.
Session chunks are jsonl where each line is a serialized ``Msg``. Parse every line
and render via :func:`format_history`; on any parse error (or no usable messages),
fall back to the chunk's raw ``text``.
"""
if not is_session_chunk(chunk, dialog_dir):
return chunk.text
try:
messages: list[Msg] = []
for line in chunk.text.splitlines():
line = line.strip()
if not line:
continue
messages.append(Msg.model_validate_json(line))
if not messages:
return chunk.text
return format_history(messages)
except Exception:
return chunk.text
def _build_union_chunk(group: list[FileChunk]) -> FileChunk:
"""Fuse a set of same-file session chunks into one covering their union.
Each chunk's ``text`` is line-aligned: text line ``i`` maps to file line
``start_line + i`` (1-based). Lines are keyed by their absolute file line
number so overlapping regions collapse to a single copy, then emitted in
ascending line order — this preserves the original message chronology and
never reorders content within the merged passage. The highest-scoring chunk
is used as the template so retrieval scores are carried through the header.
Every entry in ``line_map`` is normalised to end with ``\n`` before joining
so that a chunk whose text lacks a trailing newline (e.g. the last line of
a file with no final newline) does not collide with the next line.
"""
rep = max(group, key=lambda c: c.score)
line_map: dict[int, str] = {}
for c in group:
for offset, line in enumerate(c.text.splitlines(keepends=True)):
line_map[c.start_line + offset] = line
parts = [line_map[k] for k in sorted(line_map)]
union_text = "".join(p if p.endswith("\n") else f"{p}\n" for p in parts)
return rep.model_copy(
update={
"start_line": min(c.start_line for c in group),
"end_line": max(c.end_line for c in group),
"text": union_text,
},
)
def _merge_session_chunk_intervals(chunks: list[FileChunk], dialog_dir: str) -> list[FileChunk]:
"""Merge raw session chunks from the same file into their line-range union.
Only chunks recognized as raw session transcripts (see
:func:`is_session_chunk`) are considered; every other chunk passes through
unchanged. Within one session file, chunks are grouped by ascending line
range and merged when the next chunk's ``start_line`` is ``<= end + 1`` of
the group so far — covering the three overlap relations:
* containment: one range fully inside another;
* intersection: ranges partially overlap;
* adjacency: ``prev.end_line + 1 == next.start_line`` (gap-free consecutive
chunks, per :class:`~reme.components.file_chunker.JsonlFileChunker`).
Each merged group renders once as its union. Ordering: all units belonging
to one session file are kept adjacent and sorted by ascending ``start_line``;
the file as a whole is placed at the rank of its earliest-appearing chunk,
and non-session chunks keep their original rank position.
"""
session_by_path: dict[str, list[tuple[int, FileChunk]]] = {}
# (order_key, start_line, chunk): order_key ties all of a session file's
# units to that file's earliest rank so they sort adjacently, while
# non-session chunks use their own rank and thus keep their position.
ordered: list[tuple[int, int, FileChunk]] = []
for idx, c in enumerate(chunks):
if is_session_chunk(c, dialog_dir):
session_by_path.setdefault(c.path, []).append((idx, c))
else:
ordered.append((idx, c.start_line, c))
for items in session_by_path.values():
path_rank = min(idx for idx, _ in items)
items.sort(key=lambda t: (t[1].start_line, t[1].end_line))
group: list[FileChunk] = []
group_end: int | None = None
for _, c in items:
if group and group_end is not None and c.start_line <= group_end + 1:
group.append(c)
group_end = max(group_end, c.end_line)
else:
if group:
ordered.append((path_rank, group[0].start_line, _finalize_group(group)))
group = [c]
group_end = c.end_line
if group:
ordered.append((path_rank, group[0].start_line, _finalize_group(group)))
ordered.sort(key=lambda t: (t[0], t[1]))
return [c for _, _, c in ordered]
def _finalize_group(group: list[FileChunk]) -> FileChunk:
"""Collapse a merge group into one chunk; a single member passes through unchanged."""
if len(group) == 1:
return group[0]
return _build_union_chunk(group)
def format_chunks_answer(
chunks: list[FileChunk],
dialog_dir: str,
*,
include_source: bool = True,
score_fn: Callable[[FileChunk], str] | None = None,
link_expansion: dict[str, dict] | None = None,
) -> str:
"""Render a list of chunks into a single answer string.
Raw session chunks from the same file that overlap, contain one another, or
are adjacent are merged into their union first (see
:func:`_merge_session_chunk_intervals`).
When *include_source* is ``True`` (default), each chunk is prefixed with a
source header showing path, line range, and score. When ``False``, only the
rendered bodies are included, separated by blank lines.
*score_fn* customizes the score string in the header (default
``"score={chunk.score:.4f}"``). *link_expansion* appends per-path expansion
lines after each chunk's header block (used by hybrid search).
"""
chunks = _merge_session_chunk_intervals(chunks, dialog_dir)
if not include_source:
return "\n\n".join(render_chunk_body(c, dialog_dir) for c in chunks)
fmt = score_fn or (lambda c: f"score={c.score:.4f}")
lines: list[str] = []
for c in chunks:
sid = ""
if c.path:
try:
post = frontmatter.loads((workspace_path / c.path).read_text(encoding="utf-8"))
sid = str((post.metadata or {}).get("session_id", "") or "").strip()
except Exception:
sid = ""
header = f"[session_id={sid}]" if sid else "[session_id: unknown]"
lines.append(f"{header}\n{c.text}")
return "\n\n".join(lines)
lines.append(
f"========== {c.path}:{c.start_line}-{c.end_line} "
f"[{fmt(c)}] ==========\n{render_chunk_body(c, dialog_dir)}",
)
if link_expansion:
lines.extend(render_expansion_lines(link_expansion.get(c.path, {})))
return "\n".join(lines)

48
reme/steps/index/bm25_search.py
View file

@ -1,51 +1,25 @@
"""``bm25_search_step`` — plain BM25 keyword search with tool_context dedup."""
import datetime
from typing import Final
from ._dedup import _ToolContextDedupMixin
from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, format_chunks_answer
from ..base_step import BaseStep
from ._source_format import render_with_source
from ...components import R
from ...schema import FileChunk
_MAX_CANDIDATES: Final = 200
_CANDIDATE_MULTIPLIER: Final = 10
@R.register("bm25_search_step")
class Bm25SearchStep(BaseStep):
class Bm25SearchStep(_ToolContextDedupMixin, BaseStep):
"""BM25-only search: retrieve, filter by min_score, dedup by tool_context, truncate."""
TOOL_CONTEXTS_KEY: Final[str] = "tool_contexts"
SEARCH_SEEN_KEY: Final[str] = "search_seen_chunk_ids"
def __init__(self, *args, seen_ttl_hours: float = 24, include_source: bool = True, **kwargs):
super().__init__(*args, **kwargs)
self.seen_ttl_hours = seen_ttl_hours
self.include_source = include_source
def _tool_context_store(self, tool_context_id: str) -> dict:
"""Return the mutable state bucket for a tool context."""
if self.app_context is not None:
contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {})
else:
contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {})
return contexts.setdefault(tool_context_id, {})
def _dedupe_tool_context(self, chunks: list[FileChunk], tool_context_id: str, limit: int) -> list[FileChunk]:
"""Drop chunks already returned for this tool_context within the TTL window."""
now = datetime.datetime.now().timestamp()
ttl = float(self.seen_ttl_hours) * 60 * 60
store = self._tool_context_store(tool_context_id)
seen: dict = store.get(self.SEARCH_SEEN_KEY, {})
seen = {cid: ts for cid, ts in seen.items() if now - float(ts) < ttl}
returned = [c for c in chunks if c.id not in seen][:limit]
for c in returned:
seen[c.id] = now
store[self.SEARCH_SEEN_KEY] = seen
return returned
async def execute(self):
assert self.context is not None
query: str = (self.context.get("query", "") or "").strip()
@ -66,15 +40,21 @@ class Bm25SearchStep(BaseStep):
if min_score > 0.0:
results = [chunk for chunk in results if chunk.score >= min_score]
pre_dedup_count = 0
if tool_context_id:
results = self._dedupe_tool_context(results, tool_context_id, limit)
pre_dedup_count = len(results)
results, _ = self._dedupe_tool_context(results, tool_context_id, limit)
else:
results = results[:limit]
if self.include_source:
self.context.response.answer = render_with_source(results, self.workspace_path)
else:
self.context.response.answer = "\n\n".join(c.text for c in results)
dialog_dir = self.config_value("dialog_dir")
self.context.response.answer = format_chunks_answer(
results,
dialog_dir,
include_source=self.include_source,
)
if not results:
self.context.response.answer = ALL_RETURNED_MESSAGE if pre_dedup_count > 0 else NO_RESULTS_MESSAGE
self.context.response.metadata["results"] = [
c.model_dump(exclude_none=True, exclude={"embedding"}) for c in results
]

19
reme/steps/index/search.py
View file

@ -208,10 +208,21 @@ class SearchStep(BaseStep):
if strict_date_filter:
search_filter["strict_date_filter"] = True
vector_results, keyword_results = await asyncio.gather(
self.file_store.vector_search(query, candidates, search_filter),
self.file_store.keyword_search(query, candidates, search_filter),
)
text_weight = 1.0 - vector_weight
use_vector = vector_weight > 0.0
use_keyword = text_weight > 0.0
if use_vector and use_keyword:
vector_results, keyword_results = await asyncio.gather(
self.file_store.vector_search(query, candidates, search_filter),
self.file_store.keyword_search(query, candidates, search_filter),
)
elif use_vector:
vector_results = await self.file_store.vector_search(query, candidates, search_filter)
keyword_results = []
else:
vector_results = []
keyword_results = await self.file_store.keyword_search(query, candidates, search_filter)
self.logger.info(
f"[{self.name}] query={query!r} candidates={candidates} "

232
reme/steps/index/search_v2.py Normal file
View file

@ -0,0 +1,232 @@
"""Hybrid search (v2) over file_store using RRF fusion of vector + keyword results.
This is the local fork of the upstream search step. It uses
:class:`_ToolContextDedupMixin` for subset-aware interval-merging dedup and
:func:`format_chunks_answer` for session-aware chunk formatting with
:data:`ALL_RETURNED_MESSAGE` / :data:`NO_RESULTS_MESSAGE` notices.
"""
import asyncio
import datetime
import os
from typing import Final
from ._dedup import _ToolContextDedupMixin
from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, format_chunks_answer
from ..base_step import BaseStep
from ..file_io import extract_daily_date
from ...components import R
from ...schema import FileChunk
from ...utils import expand_links
_RRF_K: Final = 60
_MAX_CANDIDATES: Final = 200
_DEFAULT_LIMIT_ENV: Final = "REME_SEARCH_LIMIT"
_DEFAULT_LIMIT: Final = 5
def _default_limit() -> int:
value = os.getenv(_DEFAULT_LIMIT_ENV)
if value is None:
return _DEFAULT_LIMIT
try:
return int(value)
except ValueError:
return _DEFAULT_LIMIT
@R.register("search_v2_step")
class SearchV2Step(_ToolContextDedupMixin, BaseStep):
"""Hybrid search: run vector + keyword in parallel, fuse via RRF, filter, truncate."""
def __init__(
self,
*args,
seen_ttl_hours: float = 24,
**kwargs,
):
super().__init__(*args, **kwargs)
self.seen_ttl_hours = seen_ttl_hours
@staticmethod
def _rrf_merge(
vector: list[FileChunk],
keyword: list[FileChunk],
vector_weight: float,
) -> list[FileChunk]:
"""Fuse two ranked lists with Reciprocal Rank Fusion, keyed by chunk.id."""
text_weight = 1.0 - vector_weight
merged: dict[str, FileChunk] = {}
for rank, chunk in enumerate(vector, start=1):
contrib = vector_weight / (_RRF_K + rank)
c = chunk.model_copy(deep=False)
c.scores = {**chunk.scores, "vector": chunk.scores.get("vector", chunk.score), "score": contrib}
merged[c.id] = c
for rank, chunk in enumerate(keyword, start=1):
contrib = text_weight / (_RRF_K + rank)
existing = merged.get(chunk.id)
if existing is not None:
existing.scores = {
**existing.scores,
"keyword": chunk.scores.get("keyword", chunk.score),
"score": existing.scores["score"] + contrib,
}
else:
c = chunk.model_copy(deep=False)
c.scores = {**chunk.scores, "keyword": chunk.scores.get("keyword", chunk.score), "score": contrib}
merged[c.id] = c
results = list(merged.values())
results.sort(key=lambda r: r.score, reverse=True)
return results
@staticmethod
def _format_scores(scores: dict[str, float], hybrid: bool) -> str:
"""Format scores for the answer line: always show fused; show per-branch when hybrid."""
parts = [f"score={scores.get('score', 0.0):.4f}"]
if hybrid:
for k in ("vector", "keyword"):
v = scores.get(k)
parts.append(f"{k}={v:.4f}" if v is not None else f"{k}=-")
return " ".join(parts)
async def execute(self):
assert self.context is not None
query: str = (self.context.get("query", "") or "").strip()
limit: int = int(self.context.get("limit") or _default_limit())
min_score: float = float(self.context.get("min_score") or 0.0)
# vector_weight: prefer agent-supplied context value; fallback to YAML kwargs / default 0.7.
# Convertible numeric inputs are clipped to [0.0, 1.0]; non-numeric inputs are silently ignored.
raw_vw = self.context.get("vector_weight")
vector_weight: float | None = None
if raw_vw is not None:
try:
vector_weight = float(raw_vw)
except (TypeError, ValueError):
self.logger.warning(
f"[{self.name}] non-numeric vector_weight={raw_vw!r}; ignoring and using default 0.7",
)
vector_weight = None
if vector_weight is None:
vector_weight = float(self.kwargs.get("vector_weight", 0.7))
vector_weight = max(0.0, min(1.0, vector_weight))
candidate_multiplier: float = float(self.kwargs.get("candidate_multiplier", 5.0))
expand_links_enabled: bool = bool(self.kwargs.get("expand_links", True))
max_links_per_direction: int = int(self.kwargs.get("max_links_per_direction", 10))
tool_context_id: str = (self.context.get("tool_context_id", "") or "").strip()
strict_date_filter: bool = bool(
self.context.get("strict_date_filter") or self.kwargs.get("strict_date_filter", False),
)
if not query:
self.context.response.success = False
self.context.response.answer = "Error: query cannot be empty"
return self.context.response
assert limit > 0, f"limit must be positive, got {limit}"
candidates = min(_MAX_CANDIDATES, max(1, int(limit * candidate_multiplier)))
search_filter: dict = dict(self.context.get("search_filter", {}) or {})
# Promote top-level date parameters into search_filter for file_store.
for date_key in ("start_date", "end_date"):
value = self.context.get(date_key)
if value and date_key not in search_filter:
search_filter[date_key] = value
# Validate and normalize date filters before they reach file_store.
# _matches_search_filter does lexicographic string comparison against
# path_date (always a canonical YYYY-MM-DD), so raw caller values like
# "2026-2-28" or "abc" would produce silently wrong results.
for date_key in ("start_date", "end_date"):
raw = search_filter.get(date_key)
if raw is None:
continue
normalized = extract_daily_date(raw)
if normalized is None:
# Fallback: accept non-zero-padded dates like "2024-1-5".
try:
normalized = (
datetime.datetime.strptime(
str(raw).strip(),
"%Y-%m-%d",
)
.date()
.isoformat()
)
except ValueError:
self.logger.warning(
f"Ignoring invalid {date_key}={raw!r}; " f"expected a valid YYYY-MM-DD date.",
)
del search_filter[date_key]
continue
search_filter[date_key] = normalized
if strict_date_filter:
search_filter["strict_date_filter"] = True
vector_results, keyword_results = await asyncio.gather(
self.file_store.vector_search(query, candidates, search_filter),
self.file_store.keyword_search(query, candidates, search_filter),
)
self.logger.info(
f"[{self.name}] query={query!r} candidates={candidates} "
f"vector_hits={len(vector_results)} keyword_hits={len(keyword_results)}",
)
hybrid = bool(vector_results) and bool(keyword_results)
if not vector_results and not keyword_results:
fused: list[FileChunk] = []
elif not keyword_results:
fused = vector_results
elif not vector_results:
fused = keyword_results
else:
fused = self._rrf_merge(vector_results, keyword_results, vector_weight)
if min_score > 0.0:
fused = [c for c in fused if c.score >= min_score]
pre_dedup_count = 0
dedup: dict | None = None
if tool_context_id:
pre_dedup_count = len(fused)
fused, dedup = self._dedupe_tool_context(
fused,
tool_context_id,
limit,
clock=self.kwargs.get("clock"),
ttl_override=self.kwargs.get("tool_context_chunk_ttl_seconds"),
)
else:
fused = fused[:limit]
unique_paths = list(dict.fromkeys(c.path for c in fused))
link_expansion: dict[str, dict] = (
await expand_links(self.file_store, unique_paths, max_links_per_direction) if expand_links_enabled else {}
)
dialog_dir = self.config_value("dialog_dir")
self.context.response.answer = format_chunks_answer(
fused,
dialog_dir,
score_fn=lambda c: self._format_scores(c.scores, hybrid),
link_expansion=link_expansion,
)
if not fused:
self.context.response.answer = ALL_RETURNED_MESSAGE if pre_dedup_count > 0 else NO_RESULTS_MESSAGE
self.context.response.metadata["results"] = [
c.model_dump(exclude_none=True, exclude={"embedding"}) for c in fused
]
self.context.response.metadata["link_expansion"] = link_expansion
self.context.response.metadata["counts"] = {
"vector": len(vector_results),
"keyword": len(keyword_results),
"returned": len(fused),
"hybrid": hybrid,
}
if dedup is not None:
self.context.response.metadata["dedup"] = dedup
return self.context.response

48
reme/steps/index/vector_search.py
View file

@ -1,51 +1,25 @@
"""``vector_search_step`` — plain vector search with tool_context dedup."""
import datetime
from typing import Final
from ._dedup import _ToolContextDedupMixin
from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, format_chunks_answer
from ..base_step import BaseStep
from ._source_format import render_with_source
from ...components import R
from ...schema import FileChunk
_MAX_CANDIDATES: Final = 200
_CANDIDATE_MULTIPLIER: Final = 10
@R.register("vector_search_step")
class VectorSearchStep(BaseStep):
class VectorSearchStep(_ToolContextDedupMixin, BaseStep):
"""Vector-only search: retrieve, filter by min_score, dedup by tool_context, truncate."""
TOOL_CONTEXTS_KEY: Final[str] = "tool_contexts"
SEARCH_SEEN_KEY: Final[str] = "search_seen_chunk_ids"
def __init__(self, *args, seen_ttl_hours: float = 24, include_source: bool = True, **kwargs):
super().__init__(*args, **kwargs)
self.seen_ttl_hours = seen_ttl_hours
self.include_source = include_source
def _tool_context_store(self, tool_context_id: str) -> dict:
"""Return the mutable state bucket for a tool context."""
if self.app_context is not None:
contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {})
else:
contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {})
return contexts.setdefault(tool_context_id, {})
def _dedupe_tool_context(self, chunks: list[FileChunk], tool_context_id: str, limit: int) -> list[FileChunk]:
"""Drop chunks already returned for this tool_context within the TTL window."""
now = datetime.datetime.now().timestamp()
ttl = float(self.seen_ttl_hours) * 60 * 60
store = self._tool_context_store(tool_context_id)
seen: dict = store.get(self.SEARCH_SEEN_KEY, {})
seen = {cid: ts for cid, ts in seen.items() if now - float(ts) < ttl}
returned = [c for c in chunks if c.id not in seen][:limit]
for c in returned:
seen[c.id] = now
store[self.SEARCH_SEEN_KEY] = seen
return returned
async def execute(self):
assert self.context is not None
query: str = (self.context.get("query", "") or "").strip()
@ -66,15 +40,21 @@ class VectorSearchStep(BaseStep):
if min_score > 0.0:
results = [chunk for chunk in results if chunk.score >= min_score]
pre_dedup_count = 0
if tool_context_id:
results = self._dedupe_tool_context(results, tool_context_id, limit)
pre_dedup_count = len(results)
results, _ = self._dedupe_tool_context(results, tool_context_id, limit)
else:
results = results[:limit]
if self.include_source:
self.context.response.answer = render_with_source(results, self.workspace_path)
else:
self.context.response.answer = "\n\n".join(c.text for c in results)
dialog_dir = self.config_value("dialog_dir")
self.context.response.answer = format_chunks_answer(
results,
dialog_dir,
include_source=self.include_source,
)
if not results:
self.context.response.answer = ALL_RETURNED_MESSAGE if pre_dedup_count > 0 else NO_RESULTS_MESSAGE
self.context.response.metadata["results"] = [
c.model_dump(exclude_none=True, exclude={"embedding"}) for c in results
]

2
reme/utils/__init__.py
View file

@ -15,6 +15,7 @@ from .service_utils import find_reme, locate_reme, precheck_start, cli_find_reme
from .similarity_utils import cosine_similarity, batch_cosine_similarity
from .token_utils import estimate_token_count
from .agent_state_io import AsStateHandler
from .counter import global_counter_next
__all__ = [
"hash_text",
@ -37,4 +38,5 @@ __all__ = [
"batch_cosine_similarity",
"estimate_token_count",
"AsStateHandler",
"global_counter_next",
]

45
reme/utils/counter.py Normal file
View file

@ -0,0 +1,45 @@
"""Thread-safe monotonic counter tree utility for shared application state."""
import threading
from typing import Any
COUNTER_TREE_KEY = "_counter_tree"
COUNTER_LOCK_KEY = "_counter_tree_lock"
def global_counter_next(metadata: dict[str, Any], key: list[str]) -> int:
"""Return the next monotonic value for ``key``, starting at 1.
Walks the counter tree stored in ``metadata`` along ``key``, creating
missing nodes on the way, then increments and returns the target node's
counter. An empty ``key`` increments the root node, which serves as a
process-wide thread-safe global counter.
The counter tree (``{"value": 0, "children": {}}``) and its
:class:`threading.Lock` are expected to live in ``metadata`` under
:data:`COUNTER_TREE_KEY` and :data:`COUNTER_LOCK_KEY` respectively.
If they are missing they are created lazily so the function is safe to
call with a plain ``dict``.
"""
lock = metadata.get(COUNTER_LOCK_KEY)
if lock is None:
lock = threading.Lock()
metadata[COUNTER_LOCK_KEY] = lock
with lock:
tree = metadata.get(COUNTER_TREE_KEY)
if tree is None:
tree = {"value": 0, "children": {}}
metadata[COUNTER_TREE_KEY] = tree
node: dict[str, Any] = tree
for part in key:
assert isinstance(part, str)
tmp = node["children"].get(part, None)
if tmp is None:
tmp = {"value": 0, "children": {}}
node["children"][part] = tmp
node = tmp
res = node["value"] + 1
node["value"] = res
return res

564
tests/unit/test_search_step.py
View file

@ -1,4 +1,4 @@
"""Unit tests for SearchStep without embedding or LLM dependencies."""
"""Unit tests for SearchV2Step without embedding or LLM dependencies."""
import asyncio
@ -7,11 +7,19 @@ from reme.components import ApplicationContext
from reme.components.runtime_context import RuntimeContext
from reme.enumeration import LinkScopeEnum
from reme.schema import FileChunk, FileLink, FileNode
from reme.steps.index import AddDraftStep, Bm25SearchStep, ReadAllDraftStep, SearchStep, VectorSearchStep
from reme.steps.index import (
AddDraftStep,
Bm25SearchStep,
ReadAllDraftStep,
SearchStep,
SearchV2Step,
VectorSearchStep,
)
from reme.steps.index._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE
class FakeSearchStore(BaseFileStore):
"""Minimal file_store for SearchStep: static search results and empty graph links."""
"""Minimal file_store for SearchV2Step: static search results and empty graph links."""
def __init__(
self,
@ -76,7 +84,7 @@ def _chunk(
)
def test_search_step_rrf_merges_vector_and_keyword_by_chunk_id():
def test_search_v2_step_rrf_merges_vector_and_keyword_by_chunk_id():
"""Hybrid search fuses same-id hits once and keeps per-branch scores in metadata."""
async def run():
@ -88,7 +96,7 @@ def test_search_step_rrf_merges_vector_and_keyword_by_chunk_id():
vector_results=[shared_v, vector_only],
keyword_results=[keyword_only, shared_k],
)
step = SearchStep(file_store=store, vector_weight=0.5, candidate_multiplier=2, expand_links=False)
step = SearchV2Step(file_store=store, vector_weight=0.5, candidate_multiplier=2, expand_links=False)
ctx = RuntimeContext(query="alpha", limit=3, search_filter={"path_prefix": "daily/"})
resp = await step(ctx)
@ -150,14 +158,14 @@ def test_draft_steps_accumulate_by_tool_context_id():
asyncio.run(run())
def test_search_step_keyword_only_uses_keyword_scores_and_min_score():
"""When vector has no hits, SearchStep returns keyword results directly and applies min_score."""
def test_search_v2_step_keyword_only_uses_keyword_scores_and_min_score():
"""When vector has no hits, SearchV2Step returns keyword results directly and applies min_score."""
async def run():
high = _chunk("high", "daily/high.md", "strong keyword hit", "keyword", 4.0)
low = _chunk("low", "daily/low.md", "weak keyword hit", "keyword", 0.2)
store = FakeSearchStore(keyword_results=[high, low])
step = SearchStep(file_store=store, expand_links=False)
step = SearchV2Step(file_store=store, expand_links=False)
ctx = RuntimeContext(query="keyword", limit=5, min_score=1.0)
resp = await step(ctx)
@ -203,6 +211,546 @@ def test_plain_search_steps_apply_min_score_before_truncation():
asyncio.run(run())
def test_search_v2_step_tool_context_deduplicates_returned_chunks_only():
"""When tool_context_id is supplied, repeated searches skip previously returned chunks."""
async def run():
chunks = [
_chunk("a", "daily/a.md", "first", "keyword", 5.0),
_chunk("b", "daily/b.md", "second", "keyword", 4.0),
_chunk("c", "daily/c.md", "third", "keyword", 3.0),
]
store = FakeSearchStore(keyword_results=chunks)
step = SearchV2Step(file_store=store, expand_links=False)
first = await step(RuntimeContext(query="alpha", limit=2, tool_context_id="ctx-1"))
second = await step(RuntimeContext(query="alpha", limit=2, tool_context_id="ctx-1"))
third = await step(RuntimeContext(query="alpha", limit=2))
assert [r["id"] for r in first.metadata["results"]] == ["a", "b"]
assert first.metadata["dedup"] == {
"tool_context_id": "ctx-1",
"seen_before": 0,
"skipped_seen": 0,
"seen_after": 2,
"expired": 0,
"ttl_seconds": 86400.0,
}
assert [r["id"] for r in second.metadata["results"]] == ["c"]
assert second.metadata["dedup"]["seen_before"] == 2
assert second.metadata["dedup"]["skipped_seen"] == 2
assert second.metadata["dedup"]["seen_after"] == 3
assert [r["id"] for r in third.metadata["results"]] == ["a", "b"]
assert "dedup" not in third.metadata
asyncio.run(run())
def test_search_v2_step_passes_metadata_filter_to_store():
"""Search filters can target chunk metadata such as conversation_date."""
async def run():
hit = _chunk("hit", "daily/2023-01-19/event.md", "historical hit", "keyword", 3.0)
store = FakeSearchStore(keyword_results=[hit])
step = SearchV2Step(file_store=store, expand_links=False)
search_filter = {"metadata": {"conversation_date": "2023-01-19"}}
resp = await step(RuntimeContext(query="Jon job", limit=5, search_filter=search_filter))
assert resp.success is True
assert resp.metadata["results"][0]["id"] == "hit"
assert all(call[3] == search_filter for call in store.calls)
asyncio.run(run())
def test_search_v2_step_tool_context_seen_chunks_expire_after_ttl():
"""Seen chunk ids under a tool_context_id are reusable after the configured TTL."""
async def run():
now = 1000.0
chunks = [
_chunk("a", "daily/a.md", "first", "keyword", 5.0),
_chunk("b", "daily/b.md", "second", "keyword", 4.0),
]
store = FakeSearchStore(keyword_results=chunks)
step = SearchV2Step(
file_store=store,
expand_links=False,
seen_ttl_hours=1,
clock=lambda: now,
)
first = await step(RuntimeContext(query="alpha", limit=1, tool_context_id="ctx-1"))
now = 4601.0
second = await step(RuntimeContext(query="alpha", limit=1, tool_context_id="ctx-1"))
assert [r["id"] for r in first.metadata["results"]] == ["a"]
assert [r["id"] for r in second.metadata["results"]] == ["a"]
assert second.metadata["dedup"]["expired"] == 1
assert second.metadata["dedup"]["seen_before"] == 0
assert second.metadata["dedup"]["ttl_seconds"] == 3600.0
asyncio.run(run())
def test_search_v2_step_tool_context_dedup_uses_subset_matching():
"""A chunk is skipped only when its line range is a subset of a seen entry.
Partial overlap (straddle/superset) and different paths are NOT skipped:
the chunk carries lines not yet returned.
"""
async def run():
app_context = ApplicationContext()
wide = FileChunk(
id="wide",
path="daily/a.md",
text="wide",
start_line=1,
end_line=20,
scores={"keyword": 5.0, "score": 5.0},
)
store = FakeSearchStore(keyword_results=[wide])
step = SearchV2Step(
file_store=store,
app_context=app_context,
expand_links=False,
)
first = await step(RuntimeContext(query="alpha", limit=1, tool_context_id="ctx-1"))
assert [r["id"] for r in first.metadata["results"]] == ["wide"]
# subset (5-10) -> subset of (1,20) -> skipped
# straddle (15-30) -> not a subset (30>20) -> kept
# far (40-50) -> not a subset -> kept
# other (b.md 1-5) -> different path -> kept
store.keyword_results = [
FileChunk(
id="subset",
path="daily/a.md",
text="subset",
start_line=5,
end_line=10,
scores={"keyword": 4.0, "score": 4.0},
),
FileChunk(
id="straddle",
path="daily/a.md",
text="straddle",
start_line=15,
end_line=30,
scores={"keyword": 3.0, "score": 3.0},
),
FileChunk(
id="far",
path="daily/a.md",
text="far",
start_line=40,
end_line=50,
scores={"keyword": 2.0, "score": 2.0},
),
FileChunk(
id="other",
path="daily/b.md",
text="other",
start_line=1,
end_line=5,
scores={"keyword": 1.0, "score": 1.0},
),
]
second = await step(RuntimeContext(query="alpha", limit=10, tool_context_id="ctx-1"))
ids = [r["id"] for r in second.metadata["results"]]
assert "subset" not in ids
assert "straddle" in ids
assert "far" in ids
assert "other" in ids
assert second.metadata["dedup"]["skipped_seen"] == 1
asyncio.run(run())
def test_search_v2_step_tool_context_dedup_merges_adjacent_seen_intervals():
"""Multiple seen entries that jointly cover a new chunk cause it to be skipped.
Adjacent intervals (1,10) and (11,20) merge into (1,20); a new chunk
(5,15) — not covered by any single entry — is skipped because it is
covered by the merged range.
"""
async def run():
app_context = ApplicationContext()
store = FakeSearchStore(
keyword_results=[
FileChunk(
id="c1",
path="daily/a.md",
text="c1",
start_line=1,
end_line=10,
scores={"keyword": 5.0, "score": 5.0},
),
],
)
step = SearchV2Step(
file_store=store,
app_context=app_context,
expand_links=False,
)
# Return (1,10) then (11,20) — adjacent, merge into (1,20).
await step(RuntimeContext(query="alpha", limit=1, tool_context_id="ctx-1"))
store.keyword_results = [
FileChunk(
id="c2",
path="daily/a.md",
text="c2",
start_line=11,
end_line=20,
scores={"keyword": 4.0, "score": 4.0},
),
]
await step(RuntimeContext(query="alpha", limit=1, tool_context_id="ctx-1"))
# (5,15) -> covered by merged (1,20) -> skipped
# (5,25) -> NOT covered (25 > 20) -> kept
# (0,5) -> NOT covered (0 < 1) -> kept
store.keyword_results = [
FileChunk(
id="bridge",
path="daily/a.md",
text="bridge",
start_line=5,
end_line=15,
scores={"keyword": 3.0, "score": 3.0},
),
FileChunk(
id="overshoot",
path="daily/a.md",
text="overshoot",
start_line=5,
end_line=25,
scores={"keyword": 2.0, "score": 2.0},
),
FileChunk(
id="undershoot",
path="daily/a.md",
text="undershoot",
start_line=0,
end_line=5,
scores={"keyword": 1.0, "score": 1.0},
),
]
resp = await step(RuntimeContext(query="alpha", limit=10, tool_context_id="ctx-1"))
ids = [r["id"] for r in resp.metadata["results"]]
assert "bridge" not in ids
assert "overshoot" in ids
assert "undershoot" in ids
assert resp.metadata["dedup"]["skipped_seen"] == 1
asyncio.run(run())
def test_search_v2_step_empty_query_fails_before_store_calls():
"""Empty queries fail fast and do not call file_store search methods."""
async def run():
store = FakeSearchStore()
step = SearchV2Step(file_store=store)
resp = await step(RuntimeContext(query=" ", limit=5))
assert resp.success is False
assert resp.answer == "Error: query cannot be empty"
assert not store.calls
asyncio.run(run())
def test_search_v2_step_start_end_date_promoted_into_search_filter():
"""start_date and end_date from context are promoted into search_filter passed to store."""
async def run():
hit = _chunk("hit", "daily/a.md", "some text", "keyword", 5.0)
store = FakeSearchStore(keyword_results=[hit])
step = SearchV2Step(file_store=store, expand_links=False)
ctx = RuntimeContext(
query="hello",
limit=5,
start_date="2024-01-01",
end_date="2024-06-30",
)
await step(ctx)
assert len(store.calls) == 2
for _, _, _, sf in store.calls:
assert sf["start_date"] == "2024-01-01"
assert sf["end_date"] == "2024-06-30"
asyncio.run(run())
def test_search_v2_step_invalid_date_is_ignored():
"""Invalid date strings are silently ignored (removed from filter) and search proceeds."""
async def run():
hit = _chunk("hit", "daily/a.md", "some text", "keyword", 5.0)
store = FakeSearchStore(keyword_results=[hit])
step = SearchV2Step(file_store=store, expand_links=False)
ctx = RuntimeContext(
query="hello",
limit=5,
start_date="abc",
)
resp = await step(ctx)
assert resp.success is True
assert len(store.calls) == 2
for _, _, _, sf in store.calls:
assert "start_date" not in sf
asyncio.run(run())
def test_search_v2_step_non_normalized_date_is_canonicalized():
"""Valid but non-canonical dates like '2024-1-5' are normalized to '2024-01-05'."""
async def run():
hit = _chunk("hit", "daily/a.md", "some text", "keyword", 5.0)
store = FakeSearchStore(keyword_results=[hit])
step = SearchV2Step(file_store=store, expand_links=False)
ctx = RuntimeContext(
query="hello",
limit=5,
start_date="2024-1-5",
end_date="2024-6-1",
)
await step(ctx)
for _, _, _, sf in store.calls:
assert sf["start_date"] == "2024-01-05"
assert sf["end_date"] == "2024-06-01"
asyncio.run(run())
def test_search_v2_step_date_in_search_filter_not_overridden_by_context():
"""Explicit search_filter dates take precedence over top-level context dates."""
async def run():
hit = _chunk("hit", "daily/a.md", "text", "keyword", 3.0)
store = FakeSearchStore(keyword_results=[hit])
step = SearchV2Step(file_store=store, expand_links=False)
ctx = RuntimeContext(
query="hello",
limit=5,
start_date="2024-01-01",
end_date="2024-06-30",
search_filter={"start_date": "2023-07-01", "end_date": "2023-12-31"},
)
await step(ctx)
for _, _, _, sf in store.calls:
assert sf["start_date"] == "2023-07-01"
assert sf["end_date"] == "2023-12-31"
asyncio.run(run())
def test_search_v2_step_strict_date_filter_propagated_to_search_filter():
"""strict_date_filter=True is passed through to file_store via search_filter."""
async def run():
hit = _chunk("hit", "daily/2024-03-01/a.md", "text", "keyword", 3.0)
store = FakeSearchStore(keyword_results=[hit])
step = SearchV2Step(file_store=store, expand_links=False, strict_date_filter=True)
ctx = RuntimeContext(
query="hello",
limit=5,
start_date="2024-01-01",
end_date="2024-06-30",
)
await step(ctx)
for _, _, _, sf in store.calls:
assert sf["strict_date_filter"] is True
asyncio.run(run())
def test_search_v2_step_non_strict_date_filter_not_in_search_filter():
"""strict_date_filter defaults to False and is not added to search_filter."""
async def run():
hit = _chunk("hit", "daily/2024-03-01/a.md", "text", "keyword", 3.0)
store = FakeSearchStore(keyword_results=[hit])
step = SearchV2Step(file_store=store, expand_links=False)
ctx = RuntimeContext(
query="hello",
limit=5,
start_date="2024-01-01",
)
await step(ctx)
for _, _, _, sf in store.calls:
assert "strict_date_filter" not in sf
asyncio.run(run())
def test_search_v2_step_all_deduped_shows_all_returned_message():
"""When tool_context dedup removes every result, the answer explains that all content was previously returned."""
async def run():
chunks = [
_chunk("a", "daily/a.md", "first", "keyword", 5.0),
_chunk("b", "daily/b.md", "second", "keyword", 4.0),
]
store = FakeSearchStore(keyword_results=chunks)
step = SearchV2Step(file_store=store, expand_links=False)
first = await step(RuntimeContext(query="alpha", limit=5, tool_context_id="ctx-1"))
second = await step(RuntimeContext(query="alpha", limit=5, tool_context_id="ctx-1"))
assert [r["id"] for r in first.metadata["results"]] == ["a", "b"]
assert first.answer != ""
assert second.metadata["results"] == []
assert second.answer == ALL_RETURNED_MESSAGE
assert second.metadata["counts"]["returned"] == 0
asyncio.run(run())
def test_plain_search_steps_all_deduped_shows_all_returned_message():
"""VectorSearchStep and Bm25SearchStep show the all-returned message when dedup empties results."""
async def run():
vector_store = FakeSearchStore(
vector_results=[
_chunk("a", "daily/a.md", "first", "vector", 5.0),
_chunk("b", "daily/b.md", "second", "vector", 4.0),
],
)
keyword_store = FakeSearchStore(
keyword_results=[
_chunk("a", "daily/a.md", "first", "keyword", 5.0),
_chunk("b", "daily/b.md", "second", "keyword", 4.0),
],
)
vector = VectorSearchStep(file_store=vector_store, include_source=False)
bm25 = Bm25SearchStep(file_store=keyword_store, include_source=False)
v_first = await vector(RuntimeContext(query="alpha", limit=5, tool_context_id="ctx-v"))
v_second = await vector(RuntimeContext(query="alpha", limit=5, tool_context_id="ctx-v"))
b_first = await bm25(RuntimeContext(query="alpha", limit=5, tool_context_id="ctx-b"))
b_second = await bm25(RuntimeContext(query="alpha", limit=5, tool_context_id="ctx-b"))
assert v_first.answer != ""
assert v_second.metadata["results"] == []
assert v_second.answer == ALL_RETURNED_MESSAGE
assert b_first.answer != ""
assert b_second.metadata["results"] == []
assert b_second.answer == ALL_RETURNED_MESSAGE
asyncio.run(run())
def test_search_steps_no_results_shows_no_results_message():
"""When there are no results at all (before dedup), the answer explains that nothing was found."""
async def run():
empty_store = FakeSearchStore()
hybrid = SearchV2Step(file_store=empty_store, expand_links=False)
vector = VectorSearchStep(file_store=empty_store, include_source=False)
bm25 = Bm25SearchStep(file_store=empty_store, include_source=False)
# With tool_context_id set (dedup path, but nothing to dedup)
for step in (hybrid, vector, bm25):
resp = await step(RuntimeContext(query="alpha", limit=5, tool_context_id="ctx-1"))
assert resp.metadata["results"] == []
assert resp.answer == NO_RESULTS_MESSAGE
# Without tool_context_id (plain truncation path)
for step in (hybrid, vector, bm25):
resp = await step(RuntimeContext(query="alpha", limit=5))
assert resp.metadata["results"] == []
assert resp.answer == NO_RESULTS_MESSAGE
asyncio.run(run())
# ---------------------------------------------------------------------------
# SearchStep tests — exercise the SearchStep (simple chunk.id dedup,
# inline answer formatting).
# ---------------------------------------------------------------------------
def test_search_step_rrf_merges_vector_and_keyword_by_chunk_id():
"""Hybrid search fuses same-id hits once and keeps per-branch scores in metadata."""
async def run():
shared_v = _chunk("shared", "daily/a.md", "shared vector text", "vector", 0.92, line=3)
vector_only = _chunk("vector-only", "daily/b.md", "vector text", "vector", 0.71)
keyword_only = _chunk("keyword-only", "digest/c.md", "keyword text", "keyword", 8.0)
shared_k = _chunk("shared", "daily/a.md", "shared keyword text", "keyword", 7.0, line=3)
store = FakeSearchStore(
vector_results=[shared_v, vector_only],
keyword_results=[keyword_only, shared_k],
)
step = SearchStep(file_store=store, vector_weight=0.5, candidate_multiplier=2, expand_links=False)
ctx = RuntimeContext(query="alpha", limit=3, search_filter={"path_prefix": "daily/"})
resp = await step(ctx)
assert resp.success is True
assert resp.metadata["counts"] == {"vector": 2, "keyword": 2, "returned": 3, "hybrid": True}
assert [r["id"] for r in resp.metadata["results"]] == ["shared", "keyword-only", "vector-only"]
shared = resp.metadata["results"][0]
assert shared["scores"]["vector"] == 0.92
assert shared["scores"]["keyword"] == 7.0
assert shared["scores"]["score"] > resp.metadata["results"][1]["scores"]["score"]
assert "daily/a.md:3-3" in resp.answer
assert "vector=0.9200" in resp.answer
assert "keyword=7.0000" in resp.answer
assert {call[0] for call in store.calls} == {"vector", "keyword"}
assert all(call[2] == 6 for call in store.calls)
assert all(call[3] == {"path_prefix": "daily/"} for call in store.calls)
asyncio.run(run())
def test_search_step_keyword_only_uses_keyword_scores_and_min_score():
"""When vector has no hits, SearchStep returns keyword results directly and applies min_score."""
async def run():
high = _chunk("high", "daily/high.md", "strong keyword hit", "keyword", 4.0)
low = _chunk("low", "daily/low.md", "weak keyword hit", "keyword", 0.2)
store = FakeSearchStore(keyword_results=[high, low])
step = SearchStep(file_store=store, expand_links=False)
ctx = RuntimeContext(query="keyword", limit=5, min_score=1.0)
resp = await step(ctx)
assert resp.metadata["counts"] == {"vector": 0, "keyword": 2, "returned": 1, "hybrid": False}
assert [r["id"] for r in resp.metadata["results"]] == ["high"]
assert "keyword=4.0000" not in resp.answer
assert "score=4.0000" in resp.answer
assert "daily/low.md" not in resp.answer
asyncio.run(run())
def test_search_step_tool_context_deduplicates_returned_chunks_only():
"""When tool_context_id is supplied, repeated searches skip previously returned chunks."""

141
tests/unit/test_source_format_merge.py Normal file
View file

@ -0,0 +1,141 @@
"""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"