diff --git a/benchmark/longmemeval/config.yaml b/benchmark/longmemeval/config.yaml index bc62d7a7..ecc06f52 100644 --- a/benchmark/longmemeval/config.yaml +++ b/benchmark/longmemeval/config.yaml @@ -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 diff --git a/reme/config/beam.yaml b/reme/config/beam.yaml index 87cda210..805f8769 100644 --- a/reme/config/beam.yaml +++ b/reme/config/beam.yaml @@ -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: diff --git a/reme/config/lme.yaml b/reme/config/lme.yaml index e32d2037..a83d19e4 100644 --- a/reme/config/lme.yaml +++ b/reme/config/lme.yaml @@ -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: diff --git a/reme/steps/benchmark/base/agentic_answer.py b/reme/steps/benchmark/base/agentic_answer.py index 7eb45af1..1423b5e7 100644 --- a/reme/steps/benchmark/base/agentic_answer.py +++ b/reme/steps/benchmark/base/agentic_answer.py @@ -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: diff --git a/reme/steps/index/__init__.py b/reme/steps/index/__init__.py index 8edf250a..79881d08 100644 --- a/reme/steps/index/__init__.py +++ b/reme/steps/index/__init__.py @@ -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", diff --git a/reme/steps/index/_dedup.py b/reme/steps/index/_dedup.py new file mode 100644 index 00000000..e18d1c40 --- /dev/null +++ b/reme/steps/index/_dedup.py @@ -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, + } diff --git a/reme/steps/index/_source_format.py b/reme/steps/index/_source_format.py index 7e6e420a..0e42cb82 100644 --- a/reme/steps/index/_source_format.py +++ b/reme/steps/index/_source_format.py @@ -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=]`` 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) diff --git a/reme/steps/index/bm25_search.py b/reme/steps/index/bm25_search.py index c056a1b6..4fb84e55 100644 --- a/reme/steps/index/bm25_search.py +++ b/reme/steps/index/bm25_search.py @@ -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 ] diff --git a/reme/steps/index/search.py b/reme/steps/index/search.py index 64b3f485..9fee48a5 100644 --- a/reme/steps/index/search.py +++ b/reme/steps/index/search.py @@ -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} " diff --git a/reme/steps/index/search_v2.py b/reme/steps/index/search_v2.py new file mode 100644 index 00000000..9f2821b1 --- /dev/null +++ b/reme/steps/index/search_v2.py @@ -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 diff --git a/reme/steps/index/vector_search.py b/reme/steps/index/vector_search.py index f0d29216..e6d103c7 100644 --- a/reme/steps/index/vector_search.py +++ b/reme/steps/index/vector_search.py @@ -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 ] diff --git a/reme/utils/__init__.py b/reme/utils/__init__.py index b19d5984..69a16ed6 100644 --- a/reme/utils/__init__.py +++ b/reme/utils/__init__.py @@ -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", ] diff --git a/reme/utils/counter.py b/reme/utils/counter.py new file mode 100644 index 00000000..204d9cf6 --- /dev/null +++ b/reme/utils/counter.py @@ -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 diff --git a/tests/unit/test_search_step.py b/tests/unit/test_search_step.py index ac80ef7a..59bfa6c0 100644 --- a/tests/unit/test_search_step.py +++ b/tests/unit/test_search_step.py @@ -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.""" diff --git a/tests/unit/test_source_format_merge.py b/tests/unit/test_source_format_merge.py new file mode 100644 index 00000000..f259eaad --- /dev/null +++ b/tests/unit/test_source_format_merge.py @@ -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"