mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-10 22:41:06 +00:00
feat: add start_date/end_date time filter support for search job (#317)
* feat: add start_date/end_date time filter support for search job
- Add _extract_date_from_path to extract validated YYYY-MM-DD from chunk paths
- Add start_date/end_date filtering in _matches_search_filter
- Implement progressive recall in FaissLocalFileStore.vector_search
- Promote start_date/end_date from context to search_filter in SearchStep
- Add start_date/end_date parameters to search job in default.yaml
- Add unit tests for date filter functionality
* fix: validate/normalize date filters and harden _extract_date_from_path
Address three code-review comments on the time_filter search feature:
1. Validate/normalize start_date and end_date before string comparison.
_matches_search_filter does lexicographic comparison against path_date
(always canonical YYYY-MM-DD). Raw caller values like '2026-2-28' or
'abc' would produce silently wrong results. Now SearchStep normalizes
valid dates via extract_daily_date (with strptime fallback for
non-zero-padded input) and silently ignores invalid dates with a
logger.warning, removing them from the filter.
2. Clarify behavior for paths without embedded dates.
Added optional strict_date_filter parameter (default False). When True
and at least one date bound is active, chunks whose path yields no date
(e.g. digest/personal/topic.md) are excluded. When False (default),
the existing behavior is preserved — dateless paths pass through.
3. Harden _extract_date_from_path against non-standard suffixes.
Previously parts[1].split('.')[0] accepted '2026-05-18.anything' as a
valid date. Now only exact 'YYYY-MM-DD' (dir) and 'YYYY-MM-DD.md'
(day-index) forms are accepted.
This commit is contained in:
parent
f63165c66b
commit
43a407bc4f
7 changed files with 357 additions and 8 deletions
4
.gitignore
vendored
4
.gitignore
vendored
|
|
@ -3,6 +3,7 @@
|
|||
.idea/
|
||||
.vscode/
|
||||
*.code-workspace
|
||||
.qoder/
|
||||
|
||||
# Local environment
|
||||
.env
|
||||
|
|
@ -52,3 +53,6 @@ vault/
|
|||
# Documentation build outputs
|
||||
docs/_build/
|
||||
site/
|
||||
|
||||
evaluation/
|
||||
datasets/
|
||||
|
|
|
|||
|
|
@ -264,14 +264,24 @@ class FaissLocalFileStore(LocalFileStore):
|
|||
if query_embedding is None:
|
||||
return []
|
||||
|
||||
# Over-fetch by len(tombstones) so dropped rows can't starve the result set.
|
||||
q = self._prepare(query_embedding)
|
||||
if search_filter:
|
||||
k = self._faiss_index.ntotal
|
||||
else:
|
||||
k = min(self._faiss_index.ntotal, limit + len(self._tombstones))
|
||||
scores, rows = self._faiss_index.search(q, k)
|
||||
return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter)
|
||||
ntotal = self._faiss_index.ntotal
|
||||
|
||||
if not search_filter:
|
||||
# No filter: simple over-fetch to cover tombstones.
|
||||
k = min(ntotal, limit + len(self._tombstones))
|
||||
scores, rows = self._faiss_index.search(q, k)
|
||||
return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter)
|
||||
|
||||
# With filter: progressively increase k until we collect enough results
|
||||
# or exhaust the entire index.
|
||||
k = min(ntotal, 3 * limit)
|
||||
while True:
|
||||
scores, rows = self._faiss_index.search(q, k)
|
||||
results = self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter)
|
||||
if len(results) >= limit or k >= ntotal:
|
||||
return results
|
||||
k = min(ntotal, k * 2)
|
||||
|
||||
def _collect_hits(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""In-memory file store with compressed JSONL persistence on close."""
|
||||
|
||||
import datetime
|
||||
from contextlib import suppress
|
||||
|
||||
import numpy as np
|
||||
|
|
@ -336,6 +337,30 @@ class LocalFileStore(BaseFileStore):
|
|||
return actual in set(expected)
|
||||
return actual == expected
|
||||
|
||||
@staticmethod
|
||||
def _extract_date_from_path(path: str) -> str | None:
|
||||
"""Extract the date from a file path following the project path convention.
|
||||
|
||||
Paths with dates always place them at the 2nd segment:
|
||||
daily/2026-05-18/note.md
|
||||
resource/2026-06-06/report.pdf
|
||||
daily/2026-05-18.md (day index, date is stem of segment)
|
||||
|
||||
Returns a validated YYYY-MM-DD string, or None if no date is found.
|
||||
"""
|
||||
parts = path.split("/")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
# Accept only exact "YYYY-MM-DD" (dir) or "YYYY-MM-DD.md" (day index).
|
||||
segment = parts[1]
|
||||
candidate = segment if "." not in segment else segment.rsplit(".", 1)[0]
|
||||
if segment != candidate and not segment.endswith(".md"):
|
||||
return None
|
||||
try:
|
||||
return datetime.date.fromisoformat(candidate).isoformat()
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _matches_search_filter(cls, chunk: FileChunk, search_filter: dict | None) -> bool:
|
||||
"""Conservative post-filter shared by vector and keyword search."""
|
||||
|
|
@ -356,6 +381,20 @@ class LocalFileStore(BaseFileStore):
|
|||
if prefixes and not any(chunk.path.startswith(prefix) for prefix in prefixes):
|
||||
return False
|
||||
|
||||
# Date range filtering based on date embedded in chunk path.
|
||||
# strict_date_filter (default False): when True and at least one date
|
||||
# bound is set, chunks whose path yields no date are excluded.
|
||||
start_date = search_filter.get("start_date")
|
||||
end_date = search_filter.get("end_date")
|
||||
strict_date = bool(search_filter.get("strict_date_filter", False))
|
||||
if start_date or end_date:
|
||||
path_date = cls._extract_date_from_path(chunk.path)
|
||||
if not path_date:
|
||||
if strict_date:
|
||||
return False
|
||||
elif (start_date and path_date < start_date) or (end_date and path_date > end_date):
|
||||
return False
|
||||
|
||||
metadata_filter = dict(search_filter.get("metadata") or {})
|
||||
reserved = {
|
||||
"path",
|
||||
|
|
@ -365,6 +404,9 @@ class LocalFileStore(BaseFileStore):
|
|||
"prefix",
|
||||
"prefixes",
|
||||
"metadata",
|
||||
"start_date",
|
||||
"end_date",
|
||||
"strict_date_filter",
|
||||
}
|
||||
for key, value in search_filter.items():
|
||||
if key not in reserved:
|
||||
|
|
|
|||
|
|
@ -282,6 +282,12 @@ jobs:
|
|||
type: number
|
||||
description: "min fused score"
|
||||
default: 0.0
|
||||
start_date:
|
||||
type: string
|
||||
description: "optional inclusive start date filter (YYYY-MM-DD); results earlier than this date are excluded"
|
||||
end_date:
|
||||
type: string
|
||||
description: "optional inclusive end date filter (YYYY-MM-DD); results later than this date are excluded"
|
||||
required:
|
||||
- query
|
||||
steps:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
"""Hybrid search over file_store using RRF fusion of vector + keyword results."""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
|
||||
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, render_expansion_lines
|
||||
|
|
@ -68,6 +70,9 @@ class SearchStep(BaseStep):
|
|||
candidate_multiplier: float = float(self.kwargs.get("candidate_multiplier", 3.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))
|
||||
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
|
||||
|
|
@ -77,7 +82,44 @@ class SearchStep(BaseStep):
|
|||
assert limit > 0, f"limit must be positive, got {limit}"
|
||||
|
||||
candidates = min(_MAX_CANDIDATES, max(1, int(limit * candidate_multiplier)))
|
||||
search_filter: dict = self.context.get("search_filter", {}) or {}
|
||||
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),
|
||||
|
|
|
|||
|
|
@ -227,3 +227,111 @@ def test_faiss_rebuilds_stale_sidecar_and_updates_same_id_text():
|
|||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
# -- Date filter tests -------------------------------------------------------
|
||||
|
||||
|
||||
def test_date_filter_extract_and_match():
|
||||
"""_extract_date_from_path and _matches_search_filter date filtering."""
|
||||
# Extract from various path formats
|
||||
assert LocalFileStore._extract_date_from_path("daily/2026-05-18/note.md") == "2026-05-18"
|
||||
assert LocalFileStore._extract_date_from_path("daily/2026-05-18.md") == "2026-05-18"
|
||||
assert LocalFileStore._extract_date_from_path("resource/2026-06-06/report.pdf") == "2026-06-06"
|
||||
assert LocalFileStore._extract_date_from_path("digest/personal/topic.md") is None
|
||||
assert LocalFileStore._extract_date_from_path("daily/9999-99-99/note.md") is None
|
||||
assert LocalFileStore._extract_date_from_path("note.md") is None
|
||||
|
||||
# start_date / end_date boundary checks
|
||||
filt = {"start_date": "2026-02-01", "end_date": "2026-02-28"}
|
||||
assert LocalFileStore._matches_search_filter(chunk("a", "daily/2026-01-31/n.md", "t"), filt) is False
|
||||
assert LocalFileStore._matches_search_filter(chunk("b", "daily/2026-02-01/n.md", "t"), filt) is True
|
||||
assert LocalFileStore._matches_search_filter(chunk("c", "daily/2026-02-15/n.md", "t"), filt) is True
|
||||
assert LocalFileStore._matches_search_filter(chunk("d", "daily/2026-02-28/n.md", "t"), filt) is True
|
||||
assert LocalFileStore._matches_search_filter(chunk("e", "daily/2026-03-01/n.md", "t"), filt) is False
|
||||
|
||||
# No date in path → not excluded (non-strict, default)
|
||||
assert LocalFileStore._matches_search_filter(chunk("x", "digest/personal/topic.md", "t"), filt) is True
|
||||
|
||||
# strict_date_filter=True → no-date paths excluded when date filter is active
|
||||
strict_filt = {**filt, "strict_date_filter": True}
|
||||
assert LocalFileStore._matches_search_filter(chunk("x", "digest/personal/topic.md", "t"), strict_filt) is False
|
||||
assert LocalFileStore._matches_search_filter(chunk("b", "daily/2026-02-15/n.md", "t"), strict_filt) is True
|
||||
|
||||
# strict_date_filter=True but no date bounds → no-date paths still pass
|
||||
strict_no_bounds = {"strict_date_filter": True}
|
||||
assert LocalFileStore._matches_search_filter(chunk("x", "digest/personal/topic.md", "t"), strict_no_bounds) is True
|
||||
|
||||
# start_date/end_date stay in reserved, not leaked to metadata
|
||||
c = chunk("z", "daily/2026-05-18/note.md", "text")
|
||||
assert LocalFileStore._matches_search_filter(c, {"start_date": "2026-01-01", "end_date": "2026-12-31"}) is True
|
||||
|
||||
|
||||
def test_date_filter_with_vector_and_keyword_search():
|
||||
"""vector_search and keyword_search respect start_date/end_date filters."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
store = LocalFileStore(name="t_date_search", embedding_store="")
|
||||
await store.start()
|
||||
store.embedding_store = FakeEmbeddingStore()
|
||||
|
||||
await store.upsert(
|
||||
[
|
||||
(node("daily/2026-01-10/a.md"), [chunk("a", "daily/2026-01-10/a.md", "alpha topic")]),
|
||||
(node("daily/2026-02-15/b.md"), [chunk("b", "daily/2026-02-15/b.md", "alpha topic")]),
|
||||
(node("daily/2026-03-20/c.md"), [chunk("c", "daily/2026-03-20/c.md", "alpha topic")]),
|
||||
],
|
||||
)
|
||||
|
||||
filt = {"start_date": "2026-02-01", "end_date": "2026-02-28"}
|
||||
assert [c.id for c in await store.vector_search("alpha", 5, filt)] == ["b"]
|
||||
assert [c.id for c in await store.keyword_search("alpha", 5, filt)] == ["b"]
|
||||
|
||||
# start_date only
|
||||
assert sorted(c.id for c in await store.vector_search("alpha", 5, {"start_date": "2026-02-01"})) == [
|
||||
"b",
|
||||
"c",
|
||||
]
|
||||
# end_date only
|
||||
assert sorted(c.id for c in await store.keyword_search("alpha", 5, {"end_date": "2026-02-28"})) == [
|
||||
"a",
|
||||
"b",
|
||||
]
|
||||
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_faiss_date_filter_progressive_recall():
|
||||
"""FaissLocalFileStore progressive recall collects enough results with date filter."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
try:
|
||||
store = FaissLocalFileStore(name="t_faiss_date", embedding_store="")
|
||||
except ImportError:
|
||||
pytest.skip("faiss is not installed")
|
||||
await store.start()
|
||||
store.embedding_store = FakeEmbeddingStore()
|
||||
store._faiss_index = store._new_index()
|
||||
|
||||
files = []
|
||||
for i in range(10):
|
||||
day = f"2026-01-{i + 10:02d}"
|
||||
path = f"daily/{day}/note.md"
|
||||
files.append((node(path), [chunk(f"c{i}", path, "alpha topic")]))
|
||||
await store.upsert(files)
|
||||
|
||||
# Enough matches exist
|
||||
results = await store.vector_search("alpha", 3, {"start_date": "2026-01-15", "end_date": "2026-01-17"})
|
||||
assert len(results) == 3
|
||||
|
||||
# Fewer matches than limit → returns all matching
|
||||
results = await store.vector_search("alpha", 5, {"start_date": "2026-01-18", "end_date": "2026-01-19"})
|
||||
assert len(results) == 2
|
||||
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
|
|
|||
|
|
@ -143,3 +143,140 @@ def test_search_step_empty_query_fails_before_store_calls():
|
|||
assert not store.calls
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_search_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 = SearchStep(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_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 = SearchStep(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_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 = SearchStep(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_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 = SearchStep(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_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 = SearchStep(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_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 = SearchStep(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())
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue