Merge pull request #1817 from mvalentsev/fix/1815-source-file-filter

feat(search): add an optional source_file filter to mempalace_search (#1815)
This commit is contained in:
Igor Lins e Silva 2026-06-22 12:36:58 -03:00 committed by GitHub
commit 7392ab8a25
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 359 additions and 21 deletions

View File

@ -9,7 +9,7 @@ Tools (read):
mempalace_list_wings all wings with drawer counts
mempalace_list_rooms rooms within a wing
mempalace_get_taxonomy full wing room count tree
mempalace_search semantic search, optional wing/room filter
mempalace_search semantic search, optional wing/room/source_file filter
mempalace_check_duplicate check if content already exists before filing
Tools (write):
@ -1152,6 +1152,41 @@ def _sanitize_optional_name(value: str = None, field_name: str = "name") -> str:
return sanitize_name(value, field_name)
# Bounds the whole stored source_file string (often an absolute path), so it is
# Linux PATH_MAX rather than the 128-char wing/room NAME limit.
_MAX_SOURCE_FILE_LENGTH = 4096
def _sanitize_optional_source_file(value: str = None) -> str:
"""Validate an optional source_file search filter (#1815).
Unlike wing/room, a source_file is a path: ``/``, ``\\`` and ``.`` are
legal, so it is NOT run through ``sanitize_name`` (which rejects path
characters as traversal attempts). The value is matched verbatim as a
ChromaDB metadata-equality / parameterized-SQL value never used as a
filesystem path so there is no traversal risk to guard against. A null
byte or a pathological length can still upset the backend (chromadb
add/upsert chokes on null bytes / lone surrogates, #1235), so guard those
for parity with ``sanitize_name``. Blank / whitespace-only is "no filter".
"""
if value is None:
return None
if not isinstance(value, str):
raise ValueError("source_file must be a string")
value = value.strip()
if not value:
return None
if "\x00" in value:
raise ValueError("source_file contains null bytes")
if value != strip_lone_surrogates(value):
raise ValueError("source_file contains invalid surrogate characters")
if len(value) > _MAX_SOURCE_FILE_LENGTH:
raise ValueError(
f"source_file exceeds maximum length of {_MAX_SOURCE_FILE_LENGTH} characters"
)
return value
# ==================== READ TOOLS ====================
@ -1590,6 +1625,7 @@ def tool_search(
limit: int = 5,
wing: str = None,
room: str = None,
source_file: str = None,
max_distance: float = 1.5,
min_similarity: float = None,
context: str = None,
@ -1598,6 +1634,7 @@ def tool_search(
try:
wing = _sanitize_optional_name(wing, "wing")
room = _sanitize_optional_name(room, "room")
source_file = _sanitize_optional_source_file(source_file)
except ValueError as e:
return {"error": str(e)}
# Backwards compat: accept old name
@ -1616,6 +1653,7 @@ def tool_search(
palace_path=_config.palace_path,
wing=wing,
room=room,
source_file=source_file,
n_results=limit,
max_distance=dist,
vector_disabled=_vector_disabled,
@ -1633,6 +1671,7 @@ def tool_search(
palace_path=_config.palace_path,
wing=wing,
room=room,
source_file=source_file,
n_results=limit,
max_distance=dist,
vector_disabled=_vector_disabled,
@ -3483,6 +3522,15 @@ TOOLS = {
},
"wing": {"type": "string", "description": "Filter by wing (optional)"},
"room": {"type": "string", "description": "Filter by room (optional)"},
"source_file": {
"type": "string",
"description": (
"Filter to one exact source_file (optional). Matches the full "
"stored path exactly (leading/trailing whitespace trimmed); no "
"glob or basename matching. Pass the value from a result's "
"'source_path' field; the displayed 'source_file' is only a basename."
),
},
"max_distance": {
"type": "number",
"description": "Max cosine distance threshold (0=identical, 2=opposite). Results further than this are dropped. Lower = stricter. Default 1.5. Set to 0 to disable.",

View File

@ -226,15 +226,24 @@ def _hybrid_rank(
return results
def build_where_filter(wing: str = None, room: str = None) -> dict:
"""Build ChromaDB where filter for wing/room filtering."""
if wing and room:
return {"$and": [{"wing": wing}, {"room": room}]}
elif wing:
return {"wing": wing}
elif room:
return {"room": room}
return {}
def build_where_filter(wing: str = None, room: str = None, source_file: str = None) -> dict:
"""Build a ChromaDB where filter from optional wing/room/source_file.
ChromaDB needs a ``$and`` only when 2 clauses are present; a single
clause is returned bare and zero clauses yield an empty filter (#1815).
"""
clauses = []
if wing:
clauses.append({"wing": wing})
if room:
clauses.append({"room": room})
if source_file:
clauses.append({"source_file": source_file})
if not clauses:
return {}
if len(clauses) == 1:
return clauses[0]
return {"$and": clauses}
def _extract_drawer_ids_from_closet(closet_doc: str) -> list:
@ -476,6 +485,7 @@ def _bm25_only_via_sqlite(
palace_path: str,
wing: str = None,
room: str = None,
source_file: str = None,
n_results: int = 5,
max_candidates: int = 500,
_include_internal: bool = False,
@ -511,7 +521,7 @@ def _bm25_only_via_sqlite(
def _metadata_filter_sql(row_id_expr: str) -> tuple[str, list[str]]:
clauses = []
params = []
for key, value in (("wing", wing), ("room", room)):
for key, value in (("wing", wing), ("room", room), ("source_file", source_file)):
if not value:
continue
clauses.append(
@ -624,7 +634,7 @@ def _bm25_only_via_sqlite(
if not candidate_ids:
return {
"query": query,
"filters": {"wing": wing, "room": room},
"filters": {"wing": wing, "room": room, "source_file": source_file},
"total_before_filter": 0,
"results": [],
"fallback": "bm25_only_via_sqlite",
@ -660,6 +670,8 @@ def _bm25_only_via_sqlite(
continue
if room and meta.get("room") != room:
continue
if source_file and meta.get("source_file") != source_file:
continue
full_source = meta.get("source_file", "") or ""
candidates.append(
{
@ -667,6 +679,7 @@ def _bm25_only_via_sqlite(
"wing": meta.get("wing", "unknown"),
"room": meta.get("room", "unknown"),
"source_file": Path(full_source).name if full_source else "?",
"source_path": full_source,
"created_at": meta.get("filed_at", "unknown"),
# No vector distance available in BM25-only mode.
"similarity": None,
@ -702,7 +715,7 @@ def _bm25_only_via_sqlite(
return {
"query": query,
"filters": {"wing": wing, "room": room},
"filters": {"wing": wing, "room": room, "source_file": source_file},
"total_before_filter": len(candidates),
"results": hits,
"fallback": "bm25_only_via_sqlite",
@ -718,6 +731,7 @@ def _merge_bm25_union_candidates(
room: str,
n_results: int,
max_distance: float = 0.0,
source_file: str = None,
) -> None:
"""Append top-K backend lexical candidates into ``hits`` in place.
@ -744,7 +758,7 @@ def _merge_bm25_union_candidates(
if max_distance > 0.0:
return
where = build_where_filter(wing, room)
where = build_where_filter(wing, room, source_file)
try:
lexical = drawers_col.lexical_search(
query=query,
@ -767,6 +781,7 @@ def _merge_bm25_union_candidates(
"wing": meta.get("wing", "unknown"),
"room": meta.get("room", "unknown"),
"source_file": Path(full_source).name if full_source else "?",
"source_path": full_source,
"created_at": meta.get("filed_at", "unknown"),
"similarity": None,
"distance": None,
@ -832,6 +847,7 @@ def _apply_candidate_strategy(
room: str,
n_results: int,
max_distance: float = 0.0,
source_file: str = None,
) -> None:
"""Dispatch to the registered merger for ``strategy``.
@ -840,7 +856,16 @@ def _apply_candidate_strategy(
"""
merger = _CANDIDATE_MERGERS[strategy]
if merger is not None:
merger(hits, drawers_col, query, wing, room, n_results, max_distance=max_distance)
merger(
hits,
drawers_col,
query,
wing,
room,
n_results,
max_distance=max_distance,
source_file=source_file,
)
def _finalize_candidate_hits(
@ -853,6 +878,7 @@ def _finalize_candidate_hits(
room: str,
n_results: int,
max_distance: float,
source_file: str = None,
) -> tuple:
try:
_apply_candidate_strategy(
@ -864,6 +890,7 @@ def _finalize_candidate_hits(
room,
n_results,
max_distance=max_distance,
source_file=source_file,
)
except UnsupportedCapabilityError:
return [], {
@ -905,6 +932,7 @@ def _vector_disabled_search(
room: str,
n_results: int,
collection_name: str,
source_file: str = None,
) -> dict:
try:
backend_name = resolve_backend_name(palace_path)
@ -924,6 +952,7 @@ def _vector_disabled_search(
palace_path,
wing=wing,
room=room,
source_file=source_file,
n_results=n_results,
collection_name=collection_name,
)
@ -957,7 +986,9 @@ def _open_search_collection(palace_path: str, collection_name: str):
}
def _query_drawers_with_filter_fallback(drawers_col, dkwargs, query, n_results, wing, room):
def _query_drawers_with_filter_fallback(
drawers_col, dkwargs, query, n_results, wing, room, source_file=None
):
"""Run the filtered drawer query, falling back to an unfiltered query plus a
Python-side post-filter when ChromaDB raises on the filtered query.
@ -965,7 +996,7 @@ def _query_drawers_with_filter_fallback(drawers_col, dkwargs, query, n_results,
"Error finding id" even when unfiltered search works fine it happens when
drawers are ingested via two different paths (e.g. bulk import vs MCP tool
calls), leaving the vector index inconsistent with the metadata store. We
retry unfiltered (over-fetching) and re-apply the wing/room filter in Python.
retry unfiltered (over-fetching) and re-apply the wing/room/source_file filter in Python.
See #1245 / #1035.
"""
where = dkwargs.get("where")
@ -994,6 +1025,8 @@ def _query_drawers_with_filter_fallback(drawers_col, dkwargs, query, n_results,
continue
if room and meta.get("room") != room:
continue
if source_file and meta.get("source_file") != source_file:
continue
fdocs.append(doc)
fmetas.append(meta)
fdists.append(dist)
@ -1005,6 +1038,7 @@ def search_memories(
palace_path: str,
wing: str = None,
room: str = None,
source_file: str = None,
n_results: int = 5,
max_distance: float = 0.0,
vector_disabled: bool = False,
@ -1020,6 +1054,8 @@ def search_memories(
palace_path: Path to the ChromaDB palace directory.
wing: Optional wing filter.
room: Optional room filter.
source_file: Optional exact source_file filter. Matches the full
stored source_file value verbatim (#1815).
n_results: Max results to return.
max_distance: Max cosine distance threshold. The palace collection uses
cosine distance (hnsw:space=cosine) 0 = identical, 2 = opposite.
@ -1059,6 +1095,7 @@ def search_memories(
room=room,
n_results=n_results,
collection_name=collection_name,
source_file=source_file,
)
drawers_col, open_error = _open_search_collection(palace_path, collection_name)
@ -1066,7 +1103,7 @@ def search_memories(
return open_error
metric = _metric_for_collection(drawers_col)
where = build_where_filter(wing, room)
where = build_where_filter(wing, room, source_file)
# Hybrid retrieval: always query drawers directly (the floor), then use
# closet hits to boost rankings. Closets are a ranking SIGNAL, never a
@ -1084,7 +1121,7 @@ def search_memories(
if where:
dkwargs["where"] = where
drawer_results = _query_drawers_with_filter_fallback(
drawers_col, dkwargs, query, n_results, wing, room
drawers_col, dkwargs, query, n_results, wing, room, source_file
)
except Exception as e:
return {"error": f"Search error: {e}"}
@ -1156,7 +1193,10 @@ def search_memories(
"text": doc,
"wing": meta.get("wing", "unknown"),
"room": meta.get("room", "unknown"),
# source_file is the basename (display); source_path is the full
# stored value, the round-trippable key for the source_file filter.
"source_file": Path(source).name if source else "?",
"source_path": source,
"created_at": meta.get("filed_at", "unknown"),
"similarity": round(_distance_to_similarity(effective_dist, metric), 3),
"distance": round(dist, 4),
@ -1254,13 +1294,14 @@ def search_memories(
room=room,
n_results=n_results,
max_distance=max_distance,
source_file=source_file,
)
if strategy_error:
return strategy_error
return {
"query": query,
"filters": {"wing": wing, "room": room},
"filters": {"wing": wing, "room": room, "source_file": source_file},
"total_before_filter": len(_first_or_empty(drawer_results, "documents")),
"results": hits,
}

View File

@ -213,6 +213,26 @@ class TestCandidateUnion:
f"(basename collision would drop one); got sources={sources}"
)
def test_union_respects_source_file_filter(self, tmp_path):
"""Union pulls BM25 candidates from sqlite FTS5 directly; the
source_file filter must constrain that pool too, not just the vector
path otherwise union silently re-injects other sources (#1815)."""
palace = str(tmp_path / "palace")
_seed_drawers(palace)
result = search_memories(
_NARRATIVE_QUERY,
palace,
n_results=5,
candidate_strategy="union",
source_file="ticket_D2.md",
)
sources = {h["source_file"] for h in result["results"]}
assert sources <= {"ticket_D2.md"}, (
f"union must honor source_file on the BM25 pool; got {sources}"
)
# The BM25-strong brand-voice doc must NOT leak past the filter.
assert "brand_voice_D4.md" not in sources
class TestHybridRankTolerantOfMissingDistance:
"""``_hybrid_rank`` accepts ``distance=None`` — required for BM25-only

View File

@ -133,3 +133,43 @@ class TestClosetMetadata:
assert h["matched_via"] == "drawer"
assert "closet_preview" not in h
assert h["closet_boost"] == 0.0
# ── source_file filter scopes both drawer and closet queries (#1815) ──────
class TestSourceFileFilter:
def test_source_file_filter_excludes_other_sources(self, tmp_path):
palace = str(tmp_path / "palace")
_seed_drawers(palace)
result = search_memories(
"Kafka consumer rebalance timeout",
palace,
n_results=5,
source_file="fixture_D4.md",
)
ids = [h["source_file"] for h in result["results"]]
assert ids, "the matching source_file drawer should be returned"
assert set(ids) == {"fixture_D4.md"}
def test_source_file_filter_overrides_closet_boost_for_other_source(self, tmp_path):
# A strong closet pointing at D1 must NOT leak D1 in when the search
# is scoped to a different source_file — the where clause is applied
# to the closet query too, not just the drawer query.
palace = str(tmp_path / "palace")
_seed_drawers(palace)
_seed_strong_closet_for(
palace,
drawer_id="D1",
source_file="fixture_D1.md",
topics=["Kafka queue tuning", "consumer rebalance config"],
)
result = search_memories(
"Kafka consumer rebalance",
palace,
n_results=5,
source_file="fixture_D4.md",
)
ids = [h["source_file"] for h in result["results"]]
assert "fixture_D1.md" not in ids
assert set(ids) <= {"fixture_D4.md"}

View File

@ -1105,6 +1105,91 @@ class TestSearchTool:
result = tool_search(query="database", room="backend")
assert all(r["room"] == "backend" for r in result["results"])
def test_search_with_source_file_filter(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import tool_search
result = tool_search(query="authentication module", source_file="auth.py")
assert result["results"]
assert all(r["source_file"] == "auth.py" for r in result["results"])
assert result["filters"]["source_file"] == "auth.py"
def test_search_source_file_allows_path_separators(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
# Unlike wing/room, a source_file is a path — '/' must NOT be rejected
# as a path-traversal attempt the way sanitize_name() would.
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import tool_search
result = tool_search(query="authentication", source_file="/abs/path/to/auth.py")
assert "error" not in result
def test_search_blank_source_file_ignored(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import tool_search
result = tool_search(query="JWT authentication", source_file=" ")
assert "results" in result
assert result["filters"]["source_file"] is None
def test_search_rejects_null_byte_source_file(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
# A null byte in a metadata where-value can crash chromadb add/upsert
# (#1235 lineage); reject it cleanly the way sanitize_name does.
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import tool_search
result = tool_search(query="JWT", source_file="bad\x00null")
assert "error" in result
def test_search_rejects_overlong_source_file(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import tool_search
result = tool_search(query="JWT", source_file="x" * 5000)
assert "error" in result
def test_search_rejects_non_string_source_file(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
# A non-string source_file (e.g. a JSON number, which the schema's
# string type does not coerce) must yield a clean validation error,
# not an unhandled AttributeError from .strip().
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import tool_search
result = tool_search(query="JWT", source_file=42)
assert "error" in result
def test_search_rejects_lone_surrogate_source_file(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
# A lone UTF-16 surrogate can crash chromadb (#1235); reject it for
# parity with sanitize_name rather than letting it reach the backend.
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import tool_search
result = tool_search(query="JWT", source_file="bad\udc80surrogate")
assert "error" in result
def test_search_accepts_source_file_at_length_boundary(
self, monkeypatch, config, palace_path, seeded_collection, kg
):
# Exactly _MAX_SOURCE_FILE_LENGTH is allowed (the cap is a strict '>').
_patch_mcp_server(monkeypatch, config, kg)
from mempalace.mcp_server import _MAX_SOURCE_FILE_LENGTH, tool_search
result = tool_search(query="JWT", source_file="x" * _MAX_SOURCE_FILE_LENGTH)
assert "error" not in result
def test_search_min_similarity_backwards_compat(
self, monkeypatch, config, palace_path, seeded_collection, kg
):

View File

@ -9,7 +9,49 @@ from unittest.mock import MagicMock, patch
import pytest
from mempalace.searcher import SearchError, search, search_memories
from mempalace.searcher import SearchError, build_where_filter, search, search_memories
# ── build_where_filter (unit) ──────────────────────────────────────────
class TestBuildWhereFilter:
"""build_where_filter composes a ChromaDB where clause from optional
wing / room / source_file constraints (#1815). ChromaDB needs a ``$and``
only when 2 clauses are present; a single clause is returned bare and
zero clauses yield an empty filter."""
def test_no_filters_returns_empty(self):
assert build_where_filter() == {}
def test_wing_only(self):
assert build_where_filter(wing="backend") == {"wing": "backend"}
def test_room_only(self):
assert build_where_filter(room="auth") == {"room": "auth"}
def test_wing_and_room(self):
assert build_where_filter(wing="backend", room="auth") == {
"$and": [{"wing": "backend"}, {"room": "auth"}]
}
def test_source_file_only(self):
assert build_where_filter(source_file="auth.py") == {"source_file": "auth.py"}
def test_wing_and_source_file(self):
assert build_where_filter(wing="backend", source_file="auth.py") == {
"$and": [{"wing": "backend"}, {"source_file": "auth.py"}]
}
def test_room_and_source_file(self):
assert build_where_filter(room="auth", source_file="auth.py") == {
"$and": [{"room": "auth"}, {"source_file": "auth.py"}]
}
def test_wing_room_and_source_file(self):
assert build_where_filter(wing="backend", room="auth", source_file="auth.py") == {
"$and": [{"wing": "backend"}, {"room": "auth"}, {"source_file": "auth.py"}]
}
# ── search_memories (API) ──────────────────────────────────────────────
@ -34,6 +76,68 @@ class TestSearchMemories:
result = search_memories("code", palace_path, wing="project", room="frontend")
assert all(r["wing"] == "project" and r["room"] == "frontend" for r in result["results"])
def test_source_file_filter(self, palace_path, seeded_collection):
result = search_memories("authentication module", palace_path, source_file="auth.py")
assert result["results"], "exact source_file match should return its drawer"
assert all(r["source_file"] == "auth.py" for r in result["results"])
def test_source_file_with_wing_filter(self, palace_path, seeded_collection):
result = search_memories("database", palace_path, wing="project", source_file="db.py")
assert result["results"]
assert all(
r["source_file"] == "db.py" and r["wing"] == "project" for r in result["results"]
)
def test_nonmatching_source_file_returns_empty_not_error(self, palace_path, seeded_collection):
result = search_memories("authentication", palace_path, source_file="nope.md")
assert "error" not in result
assert result["results"] == []
def test_filters_envelope_includes_source_file(self, palace_path, seeded_collection):
result = search_memories("authentication", palace_path, source_file="auth.py")
assert result["filters"]["source_file"] == "auth.py"
def test_result_exposes_full_source_path(self, palace_path, seeded_collection):
# The displayed source_file is a basename; source_path carries the full
# stored value so a caller can round-trip it back into a source_file filter.
result = search_memories("authentication module", palace_path)
hit = result["results"][0]
assert hit["source_file"] == "auth.py"
assert hit["source_path"] == "auth.py"
def test_source_file_filter_matches_full_path_not_basename(self, palace_path):
from mempalace.palace import get_collection
col = get_collection(palace_path, create=True)
col.upsert(
ids=["fp1"],
documents=["The deploy script restarts the gunicorn workers nightly."],
metadatas=[{"wing": "ops", "room": "deploy", "source_file": "/srv/app/deploy.sh"}],
)
# The full stored path matches and round-trips via source_path.
hit = search_memories(
"deploy gunicorn workers", palace_path, source_file="/srv/app/deploy.sh"
)
assert [h["source_path"] for h in hit["results"]] == ["/srv/app/deploy.sh"]
assert [h["source_file"] for h in hit["results"]] == ["deploy.sh"]
# The basename does NOT match — exact full-path semantics only (issue v1).
miss = search_memories("deploy gunicorn workers", palace_path, source_file="deploy.sh")
assert miss["results"] == []
def test_source_file_filter_honored_in_bm25_fallback(self, palace_path, seeded_collection):
# vector_disabled routes through _bm25_only_via_sqlite (#1222); the
# source_file filter must hold there too, not silently no-op.
result = search_memories(
"authentication module",
palace_path,
source_file="auth.py",
vector_disabled=True,
collection_name="mempalace_drawers",
)
assert "error" not in result
assert result["results"], "BM25 fallback should still find the auth drawer"
assert all(r["source_file"] == "auth.py" for r in result["results"])
def test_n_results_limit(self, palace_path, seeded_collection):
result = search_memories("code", palace_path, n_results=2)
assert len(result["results"]) <= 2