ReMe/reme/steps/index/vector_search.py
sa-buc db4763aead feat: add CompressorStep and search_v2 dual-mode session compression
- Add CompressorStep (reme/steps/evolve/compressor.py) for direct LLM
  text compression with optional query-guided relevance filtering
- Extend search_v2_step to support query-aware and query-independent
  session transcript compression via _compress injected kwargs
- Refactor _source_format.py: split into render_chunk_entries +
  join_chunk_entries; session chunks now render line-aligned with
  L<n>: prefixes for verbatim/compressed parity
- Add JOB_TOOLS and INJECTED_JOB_KWARGS to BaseAgenticAnswerStep for
  per-subclass tool and parameter injection
- LmeAgenticAnswerStep injects _search._compress payload to enable
  query-aware compression during benchmark evaluation
- Record compression ablation results in result-longmemeval.md
- Add unit tests for CompressorStep and search compression paths
2026-08-04 15:30:38 +08:00

63 lines
2.6 KiB
Python

"""``vector_search_step`` — plain vector search with tool_context dedup."""
from typing import Final
from ._dedup import _ToolContextDedupMixin
from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, join_chunk_entries
from ._source_format import merge_session_chunk_intervals, render_chunk_entries
from ..base_step import BaseStep
from ...components import R
_MAX_CANDIDATES: Final = 200
_CANDIDATE_MULTIPLIER: Final = 10
@R.register("vector_search_step")
class VectorSearchStep(_ToolContextDedupMixin, BaseStep):
"""Vector-only search: retrieve, filter by min_score, dedup by tool_context, truncate."""
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
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 5)
min_score: float = float(self.context.get("min_score") or 0.0)
tool_context_id: str = (self.context.get("tool_context_id", "") or "").strip()
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, limit * _CANDIDATE_MULTIPLIER))
results = await self.file_store.vector_search(query, candidates, {})
self.logger.info(f"[{self.name}] query={query!r} candidates={candidates} hits={len(results)}")
if min_score > 0.0:
results = [chunk for chunk in results if chunk.score >= min_score]
pre_dedup_count = 0
if tool_context_id:
pre_dedup_count = len(results)
results, _ = self._dedupe_tool_context(results, tool_context_id, limit)
else:
results = results[:limit]
dialog_dir = self.config_value("dialog_dir")
entries = render_chunk_entries(
merge_session_chunk_intervals(results, dialog_dir),
dialog_dir,
include_source=self.include_source,
)
self.context.response.answer = join_chunk_entries(entries)
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
]
return self.context.response