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:
commit
7392ab8a25
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue