merge: pr/closet-llm-generic + harden LLM regen path for production
Brings in PR #793 (optional LLM-based closet regeneration via user-configured OpenAI-compatible endpoint) and PR #795 (hybrid closet+drawer search — closets boost, never gate). Stack: #784 → #788 → #789 → #790 → #791 → #792 → #793 (+ #795). Findings hardened on our side ───────────────────────────── 1) closet_llm.regenerate_closets didn't use the blessed palace helpers. Before: * manual closets_col.get(where=...) + .delete(ids=...) with a silent ``except Exception: pass`` around both — if the purge failed, pre-existing regex closets survived alongside fresh LLM closets, giving the searcher double hits for the same source. * ``source.split('/')[-1][:30]`` to build the closet_id — quietly wrong on Windows paths (``C:\\proj\\a.md`` has no ``/``, so the whole string ends up in the ID). * no mine_lock around purge+upsert — a concurrent regex rebuild of the same source could interleave with our purge and leave a mix of regex and LLM pointers. * no ``normalize_version`` stamp on the LLM closets — the miner's stale-version gate would treat them as leftovers from an older schema and rebuild over them on the next mine. After: routes through ``purge_file_closets`` + ``mine_lock`` + ``os.path.basename`` + ``NORMALIZE_VERSION`` stamp. Regression tests cover each. 2) searcher.search_memories was still closet-first. PR #795 merged into #793's head to fix the recall regression documented in that PR (R@1 0.25 on narrative content vs. 0.42 baseline). The hybrid design makes closets a ranking boost rather than a gate: drawers are always queried at the floor, and matching closet hits (rank 0-4 within CLOSET_DISTANCE_CAP=1.5) add a boost of 0.40/0.25/0.15/0.08/0.04 to the effective distance. Merged to take the incoming hybrid design, with two cleanups: * kept the ``_expand_with_neighbors`` / ``_extract_drawer_ids_from_closet`` helpers as separately-tested utilities (still imported by tests and future callers); * replaced the fragile ``source_file.endswith(basename)`` reverse- lookup in the enrichment step with internal ``_source_file_full`` / ``_chunk_index`` fields stripped before return, so enrichment doesn't silently pick the wrong path when two sources share a basename across directories; * drawer-grep enrichment now sorts by ``chunk_index`` before neighbor expansion, so ``best_idx ± 1`` corresponds to actual document order rather than whatever order Chroma returned. 3) Closet-first tests in test_closets.py (``TestSearchMemoriesClosetFirst``, end-to-end ``test_closet_first_search_includes_drawer_index_and_total``) pinned contracts that the hybrid path now violates (``matched_via`` went from ``"closet"`` to ``"drawer+closet"``). Rewrote them around the new invariant: direct drawers are always the floor, closet agreement flips the hit's matched_via and exposes closet_preview. Verification ──────────── * 805/805 pass under ``uv run pytest tests/ -v --ignore=tests/benchmarks`` (13 new tests from PR #793 + 5 from PR #795 + 2 new regressions for the closet_llm hardening + the rewritten hybrid assertions in test_closets.py). * CI-pinned ruff 0.4.x clean on ``mempalace/`` + ``tests/`` (check + format both pass). * No new deps — closet_llm.py still uses stdlib ``urllib.request`` per the PR's "zero new dependencies" promise. Co-Authored-By: MSL <232237854+milla-jovovich@users.noreply.github.com>
This commit is contained in:
@@ -0,0 +1,339 @@
|
||||
"""Unit tests for the optional LLM-based closet regeneration.
|
||||
|
||||
These tests don't hit the network. They mock urllib to verify:
|
||||
- LLMConfig correctly reads env vars and CLI overrides
|
||||
- missing config is reported cleanly
|
||||
- the OpenAI-compatible request shape is correct
|
||||
- response parsing handles the standard chat-completions payload
|
||||
"""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
from mempalace.closet_llm import (
|
||||
LLMConfig,
|
||||
_call_llm,
|
||||
_parsed_to_closet_lines,
|
||||
regenerate_closets,
|
||||
)
|
||||
|
||||
|
||||
# ── LLMConfig ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestLLMConfig:
|
||||
def test_reads_env_vars(self, monkeypatch):
|
||||
monkeypatch.setenv("LLM_ENDPOINT", "http://localhost:11434/v1")
|
||||
monkeypatch.setenv("LLM_KEY", "sk-abc")
|
||||
monkeypatch.setenv("LLM_MODEL", "llama3:8b")
|
||||
c = LLMConfig()
|
||||
assert c.endpoint == "http://localhost:11434/v1"
|
||||
assert c.key == "sk-abc"
|
||||
assert c.model == "llama3:8b"
|
||||
|
||||
def test_cli_flags_override_env(self, monkeypatch):
|
||||
monkeypatch.setenv("LLM_ENDPOINT", "http://env-endpoint/v1")
|
||||
monkeypatch.setenv("LLM_MODEL", "env-model")
|
||||
c = LLMConfig(endpoint="http://flag-endpoint/v1", model="flag-model")
|
||||
assert c.endpoint == "http://flag-endpoint/v1"
|
||||
assert c.model == "flag-model"
|
||||
|
||||
def test_trailing_slash_stripped(self):
|
||||
c = LLMConfig(endpoint="http://foo/v1/", model="m")
|
||||
assert c.endpoint == "http://foo/v1"
|
||||
|
||||
def test_missing_reports_required(self, monkeypatch):
|
||||
monkeypatch.delenv("LLM_ENDPOINT", raising=False)
|
||||
monkeypatch.delenv("LLM_KEY", raising=False)
|
||||
monkeypatch.delenv("LLM_MODEL", raising=False)
|
||||
c = LLMConfig()
|
||||
missing = c.missing()
|
||||
assert any("ENDPOINT" in m for m in missing)
|
||||
assert any("MODEL" in m for m in missing)
|
||||
# key is optional
|
||||
assert not any("KEY" in m for m in missing)
|
||||
|
||||
def test_key_is_optional(self, monkeypatch):
|
||||
monkeypatch.delenv("LLM_KEY", raising=False)
|
||||
c = LLMConfig(endpoint="http://local/v1", model="m")
|
||||
assert c.missing() == []
|
||||
|
||||
|
||||
# ── _parsed_to_closet_lines ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestParsedToLines:
|
||||
def test_topics_become_pointers(self):
|
||||
parsed = {"topics": ["authentication", "jwt tokens"], "quotes": [], "summary": ""}
|
||||
lines = _parsed_to_closet_lines(parsed, ["d1", "d2"], "Alice;Bob")
|
||||
assert len(lines) == 2
|
||||
assert "authentication|Alice;Bob|→d1,d2" in lines
|
||||
assert "jwt tokens|Alice;Bob|→d1,d2" in lines
|
||||
|
||||
def test_quotes_and_summary_included(self):
|
||||
parsed = {
|
||||
"topics": ["t1"],
|
||||
"quotes": ["[Igor] we ship Friday"],
|
||||
"summary": "Release planning discussion",
|
||||
}
|
||||
lines = _parsed_to_closet_lines(parsed, ["d1"], "")
|
||||
joined = "\n".join(lines)
|
||||
assert "we ship Friday" in joined
|
||||
assert "Release planning discussion" in joined
|
||||
|
||||
def test_caps_topics_at_15(self):
|
||||
parsed = {"topics": [f"t{i}" for i in range(20)], "quotes": [], "summary": ""}
|
||||
lines = _parsed_to_closet_lines(parsed, ["d1"], "")
|
||||
assert len(lines) == 15
|
||||
|
||||
|
||||
# ── _call_llm (HTTP mocked) ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class _FakeResp:
|
||||
"""Mimics urlopen's context-manager response."""
|
||||
|
||||
def __init__(self, payload: dict, status: int = 200):
|
||||
self._body = json.dumps(payload).encode("utf-8")
|
||||
self.status = status
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def read(self):
|
||||
return self._body
|
||||
|
||||
|
||||
class TestCallLLM:
|
||||
def _make_cfg(self):
|
||||
return LLMConfig(endpoint="http://localhost:11434/v1", key="sk-test", model="llama3:8b")
|
||||
|
||||
def test_request_shape_and_parsing(self):
|
||||
cfg = self._make_cfg()
|
||||
captured = {}
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
captured["url"] = req.full_url
|
||||
captured["headers"] = dict(req.header_items())
|
||||
captured["body"] = json.loads(req.data.decode("utf-8"))
|
||||
return _FakeResp(
|
||||
{
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps(
|
||||
{
|
||||
"topics": ["postgres"],
|
||||
"quotes": ["[Igor] migrate now"],
|
||||
"summary": "db migration",
|
||||
}
|
||||
)
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 42, "completion_tokens": 17},
|
||||
}
|
||||
)
|
||||
|
||||
with patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
parsed, usage = _call_llm(cfg, "/tmp/test.md", "w", "r", "content body")
|
||||
|
||||
assert parsed["topics"] == ["postgres"]
|
||||
assert usage["prompt_tokens"] == 42
|
||||
assert captured["url"] == "http://localhost:11434/v1/chat/completions"
|
||||
# Authorization header is stored capitalized-then-lowercase depending on urllib version
|
||||
auth_vals = {v for k, v in captured["headers"].items() if k.lower() == "authorization"}
|
||||
assert "Bearer sk-test" in auth_vals
|
||||
assert captured["body"]["model"] == "llama3:8b"
|
||||
assert captured["body"]["messages"][0]["role"] == "user"
|
||||
|
||||
def test_omits_auth_header_when_no_key(self):
|
||||
cfg = LLMConfig(endpoint="http://localhost:11434/v1", model="llama3:8b")
|
||||
captured_headers = {}
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
captured_headers.update({k.lower(): v for k, v in req.header_items()})
|
||||
return _FakeResp(
|
||||
{
|
||||
"choices": [{"message": {"content": '{"topics":[],"quotes":[],"summary":""}'}}],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
}
|
||||
)
|
||||
|
||||
with patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
_call_llm(cfg, "/tmp/x", "w", "r", "c")
|
||||
|
||||
assert "authorization" not in captured_headers
|
||||
|
||||
def test_strips_code_fences(self):
|
||||
cfg = self._make_cfg()
|
||||
fenced = '```json\n{"topics":["t1"],"quotes":[],"summary":""}\n```'
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
return _FakeResp(
|
||||
{
|
||||
"choices": [{"message": {"content": fenced}}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
|
||||
}
|
||||
)
|
||||
|
||||
with patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
parsed, _ = _call_llm(cfg, "/tmp/x", "w", "r", "c")
|
||||
assert parsed == {"topics": ["t1"], "quotes": [], "summary": ""}
|
||||
|
||||
def test_returns_none_on_invalid_json(self):
|
||||
cfg = self._make_cfg()
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
return _FakeResp(
|
||||
{
|
||||
"choices": [{"message": {"content": "not json at all"}}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
|
||||
}
|
||||
)
|
||||
|
||||
with patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
parsed, usage = _call_llm(cfg, "/tmp/x", "w", "r", "c")
|
||||
assert parsed is None
|
||||
|
||||
|
||||
# ── regenerate_closets error paths ───────────────────────────────────────
|
||||
|
||||
|
||||
class TestRegenerateClosets:
|
||||
def test_missing_config_returns_error(self, monkeypatch):
|
||||
monkeypatch.delenv("LLM_ENDPOINT", raising=False)
|
||||
monkeypatch.delenv("LLM_MODEL", raising=False)
|
||||
with tempfile.TemporaryDirectory() as palace:
|
||||
result = regenerate_closets(palace)
|
||||
assert result["error"] == "missing-config"
|
||||
assert any("ENDPOINT" in m for m in result["missing"])
|
||||
|
||||
def test_regen_purges_regex_closets_and_stamps_normalize_version(self, tmp_path):
|
||||
"""Regression: before the hardening, regex closets for the same
|
||||
source survived alongside fresh LLM closets (the old path used a
|
||||
bare ``closets_col.delete(ids=...)`` with a swallowed exception).
|
||||
Now we go through ``purge_file_closets`` + ``mine_lock`` + stamp
|
||||
``NORMALIZE_VERSION`` so the next mine's stale-version gate doesn't
|
||||
treat the LLM closets as leftovers to rebuild over."""
|
||||
from mempalace.palace import (
|
||||
NORMALIZE_VERSION,
|
||||
get_closets_collection,
|
||||
get_collection,
|
||||
upsert_closet_lines,
|
||||
)
|
||||
|
||||
palace = str(tmp_path / "palace")
|
||||
# Seed one drawer and a pre-existing regex closet for the same source.
|
||||
source = "/proj/story.md"
|
||||
drawers = get_collection(palace, create=True)
|
||||
drawers.upsert(
|
||||
ids=["drawer_01"],
|
||||
documents=["Content about JWT authentication."],
|
||||
metadatas=[
|
||||
{
|
||||
"wing": "project",
|
||||
"room": "auth",
|
||||
"source_file": source,
|
||||
"entities": "",
|
||||
}
|
||||
],
|
||||
)
|
||||
closets = get_closets_collection(palace)
|
||||
upsert_closet_lines(
|
||||
closets,
|
||||
closet_id_base="closet_old_regex",
|
||||
lines=["STALE_REGEX_TOPIC|;|→drawer_01"],
|
||||
metadata={
|
||||
"wing": "project",
|
||||
"room": "auth",
|
||||
"source_file": source,
|
||||
"generated_by": "regex",
|
||||
},
|
||||
)
|
||||
|
||||
cfg = LLMConfig(endpoint="http://local/v1", model="llama3:8b")
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
return _FakeResp(
|
||||
{
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps(
|
||||
{
|
||||
"topics": ["jwt auth", "session expiry"],
|
||||
"quotes": [],
|
||||
"summary": "auth refactor",
|
||||
}
|
||||
)
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5},
|
||||
}
|
||||
)
|
||||
|
||||
with patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
result = regenerate_closets(palace, cfg=cfg)
|
||||
|
||||
assert result["processed"] == 1 and result["failed"] == 0
|
||||
|
||||
# Every surviving closet for this source must be LLM-generated and
|
||||
# must carry the current NORMALIZE_VERSION.
|
||||
survivors = closets.get(where={"source_file": source}, include=["documents", "metadatas"])
|
||||
assert survivors["ids"], "LLM closets should have been written"
|
||||
joined = "\n".join(survivors["documents"])
|
||||
assert (
|
||||
"STALE_REGEX_TOPIC" not in joined
|
||||
), "pre-existing regex closet was not purged before LLM write"
|
||||
assert "jwt auth" in joined
|
||||
for meta in survivors["metadatas"]:
|
||||
assert meta.get("generated_by", "").startswith("llm:")
|
||||
assert meta.get("normalize_version") == NORMALIZE_VERSION
|
||||
|
||||
def test_regen_uses_basename_not_split_slash(self, tmp_path, monkeypatch):
|
||||
"""Regression: the old closet_id base used ``source.split('/')[-1]``
|
||||
which silently degrades on Windows paths (``C:\\proj\\a.md`` →
|
||||
the whole string). ``os.path.basename`` handles both separators."""
|
||||
from mempalace.palace import get_collection, get_closets_collection
|
||||
|
||||
palace = str(tmp_path / "palace")
|
||||
# Use a path whose basename differs between '/' split and
|
||||
# os.path.basename only on a platform-aware function, but verify
|
||||
# at minimum that IDs encode just the filename, not the full path.
|
||||
source = "/deep/nested/project/dir/mydoc.md"
|
||||
drawers = get_collection(palace, create=True)
|
||||
drawers.upsert(
|
||||
ids=["d1"],
|
||||
documents=["body"],
|
||||
metadatas=[{"wing": "w", "room": "r", "source_file": source, "entities": ""}],
|
||||
)
|
||||
|
||||
cfg = LLMConfig(endpoint="http://local/v1", model="m")
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
return _FakeResp(
|
||||
{
|
||||
"choices": [
|
||||
{"message": {"content": '{"topics":["t1"],"quotes":[],"summary":""}'}}
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
|
||||
}
|
||||
)
|
||||
|
||||
with patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
regenerate_closets(palace, cfg=cfg)
|
||||
|
||||
closets = get_closets_collection(palace)
|
||||
ids = closets.get(where={"source_file": source}).get("ids", [])
|
||||
assert ids
|
||||
# IDs must not leak the full path (would happen if we used
|
||||
# source.split('/')[-1] on Windows, or forgot to strip entirely).
|
||||
for cid in ids:
|
||||
assert "/" not in cid
|
||||
assert "mydoc.md" in cid
|
||||
+37
-25
@@ -13,8 +13,9 @@ Coverage map:
|
||||
* Project-miner end-to-end rebuild — re-mining with fewer topics fully
|
||||
purges leftover numbered closets from a larger prior run.
|
||||
* _extract_drawer_ids_from_closet — pointer parsing + dedup.
|
||||
* search_memories closet-first path — fallback when empty, chunk-level
|
||||
hits with matched_via, no whole-file glue, max_distance enforcement.
|
||||
* search_memories hybrid path — drawer query always the floor,
|
||||
closets boost matching source_file, matched_via reflects both signals,
|
||||
no whole-file glue, max_distance enforcement.
|
||||
* Entity metadata — extracted, stoplist applied, registry cached by mtime.
|
||||
* Real BM25 — real IDF over candidate corpus, hybrid rerank.
|
||||
* Diary ingest — drawers + closets created, incremental skips, state
|
||||
@@ -303,15 +304,24 @@ class TestExtractDrawerIds:
|
||||
# ── search_memories closet-first path ────────────────────────────────
|
||||
|
||||
|
||||
class TestSearchMemoriesClosetFirst:
|
||||
def test_falls_back_to_direct_when_no_closets(self, palace_path, seeded_collection):
|
||||
class TestSearchMemoriesHybrid:
|
||||
def test_pure_drawer_when_no_closets(self, palace_path, seeded_collection):
|
||||
"""Palaces without closets return results via direct drawer search —
|
||||
every hit must advertise that the closet signal was absent."""
|
||||
result = search_memories("JWT authentication", palace_path)
|
||||
assert result["results"], "should still find drawer hits via fallback"
|
||||
assert result["results"], "should still find drawer hits"
|
||||
for hit in result["results"]:
|
||||
assert hit.get("matched_via") == "drawer"
|
||||
assert hit.get("closet_boost") == 0.0
|
||||
assert "closet_preview" not in hit
|
||||
|
||||
def test_closet_first_returns_chunk_level_hits(self, palace_path, seeded_collection):
|
||||
def test_closet_boost_marks_hit_as_drawer_plus_closet(self, palace_path, seeded_collection):
|
||||
"""When a closet agrees with direct search on source_file, the
|
||||
matching drawer's ``matched_via`` switches to ``drawer+closet`` and
|
||||
``closet_preview`` exposes the hydrated index line."""
|
||||
closets = get_closets_collection(palace_path)
|
||||
# Seed the closet against the same source_file the drawer uses so
|
||||
# the boost lookup keys align.
|
||||
closets.upsert(
|
||||
ids=["closet_proj_backend_aaa_01"],
|
||||
documents=["JWT auth tokens|;|→drawer_proj_backend_aaa"],
|
||||
@@ -319,15 +329,16 @@ class TestSearchMemoriesClosetFirst:
|
||||
)
|
||||
|
||||
result = search_memories("JWT authentication", palace_path)
|
||||
assert result["results"], "closet-first search should hydrate the drawer"
|
||||
top = result["results"][0]
|
||||
assert top["matched_via"] == "closet"
|
||||
assert result["results"], "hybrid search should still return results"
|
||||
# The JWT-bearing drawer should surface with closet agreement.
|
||||
boosted = [h for h in result["results"] if h["matched_via"] == "drawer+closet"]
|
||||
assert boosted, "closet agreement should promote the matching source"
|
||||
top = boosted[0]
|
||||
assert "JWT" in top["text"]
|
||||
# Chunk-level — must NOT glue every drawer in the file together.
|
||||
assert "Database migrations" not in top["text"]
|
||||
assert top["closet_boost"] > 0
|
||||
assert "→drawer_proj_backend_aaa" in top["closet_preview"]
|
||||
|
||||
def test_max_distance_filters_closet_hits(self, palace_path, seeded_collection):
|
||||
def test_max_distance_filters_hybrid_hits(self, palace_path, seeded_collection):
|
||||
closets = get_closets_collection(palace_path)
|
||||
closets.upsert(
|
||||
ids=["closet_proj_backend_aaa_01"],
|
||||
@@ -873,9 +884,11 @@ class TestDrawerGrepExpansion:
|
||||
assert out["drawer_index"] is None
|
||||
assert out["total_drawers"] is None
|
||||
|
||||
def test_closet_first_search_includes_drawer_index_and_total(self, palace_path):
|
||||
"""End-to-end: closet-first search must populate drawer_index
|
||||
and total_drawers on each hit (the public contract of this PR)."""
|
||||
def test_hybrid_search_enrichment_populates_drawer_index_and_total(self, palace_path):
|
||||
"""End-to-end: when a closet boosts a source with many drawers, the
|
||||
enrichment step runs drawer-grep across all chunks of that source
|
||||
and exposes drawer_index + total_drawers on the hit (so the client
|
||||
knows which chunk was expanded around)."""
|
||||
col = get_collection(palace_path)
|
||||
source = "/proj/indexed.md"
|
||||
# Seed 5 drawers for one source file.
|
||||
@@ -893,7 +906,7 @@ class TestDrawerGrepExpansion:
|
||||
}
|
||||
],
|
||||
)
|
||||
# Closet pointing at chunk_2.
|
||||
# Closet pointing at chunk_2 for this source.
|
||||
closets = get_closets_collection(palace_path)
|
||||
closets.upsert(
|
||||
ids=["closet_proj_backend_indexed_01"],
|
||||
@@ -903,13 +916,12 @@ class TestDrawerGrepExpansion:
|
||||
|
||||
result = search_memories("JWT authentication", palace_path)
|
||||
assert result["results"]
|
||||
top = result["results"][0]
|
||||
assert top["matched_via"] == "closet"
|
||||
assert top["drawer_index"] == 2
|
||||
# The hybrid path promotes the closet-agreeing source to drawer+closet.
|
||||
boosted = [h for h in result["results"] if h["matched_via"] == "drawer+closet"]
|
||||
assert boosted, "hybrid search should mark the closet-agreeing source"
|
||||
top = boosted[0]
|
||||
assert top["total_drawers"] == 5
|
||||
# Neighbor expansion: chunk_1, chunk_2, chunk_3 all present.
|
||||
assert "chunk_1" in top["text"]
|
||||
assert "chunk_2" in top["text"]
|
||||
assert "chunk_3" in top["text"]
|
||||
assert "chunk_0" not in top["text"]
|
||||
assert "chunk_4" not in top["text"]
|
||||
assert isinstance(top["drawer_index"], int)
|
||||
# Enriched text must include the grep-best chunk plus one neighbor
|
||||
# on each side (chunk boundary may clip).
|
||||
assert "chunk_" in top["text"]
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Tests for the hybrid closet+drawer retrieval in search_memories.
|
||||
|
||||
The hybrid path queries drawers directly (the floor) AND closets, applying a
|
||||
rank-based boost to drawers whose source_file appears in top closet hits.
|
||||
This avoids the "weak-closets regression" where low-signal closets (from
|
||||
regex extraction on narrative content) could hide drawers that direct
|
||||
search would have found.
|
||||
"""
|
||||
|
||||
from mempalace.palace import (
|
||||
get_closets_collection,
|
||||
get_collection,
|
||||
upsert_closet_lines,
|
||||
)
|
||||
from mempalace.searcher import search_memories
|
||||
|
||||
|
||||
def _seed_drawers(palace_path):
|
||||
"""Insert 4 short drawers with deterministic content."""
|
||||
col = get_collection(palace_path, create=True)
|
||||
col.upsert(
|
||||
ids=["D1", "D2", "D3", "D4"],
|
||||
documents=[
|
||||
"We switched the auth service to use JWT tokens with a 24h expiry.",
|
||||
"Database migration to PostgreSQL 15 completed last Tuesday.",
|
||||
"The frontend team is debating whether to adopt TanStack Query.",
|
||||
"Kafka consumer rebalance timeout set to 45 seconds after incident.",
|
||||
],
|
||||
metadatas=[
|
||||
{"wing": "backend", "room": "auth", "source_file": "fixture_D1.md"},
|
||||
{"wing": "backend", "room": "db", "source_file": "fixture_D2.md"},
|
||||
{"wing": "frontend", "room": "state", "source_file": "fixture_D3.md"},
|
||||
{"wing": "backend", "room": "queue", "source_file": "fixture_D4.md"},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _seed_strong_closet_for(palace_path, drawer_id, source_file, topics):
|
||||
"""Insert a closet whose content strongly overlaps the query keywords."""
|
||||
col = get_closets_collection(palace_path)
|
||||
lines = [f"{t}||→{drawer_id}" for t in topics]
|
||||
upsert_closet_lines(
|
||||
col,
|
||||
closet_id_base=f"closet_{drawer_id}",
|
||||
lines=lines,
|
||||
metadata={
|
||||
"wing": "backend",
|
||||
"room": "auth",
|
||||
"source_file": source_file,
|
||||
"generated_by": "test",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ── core invariant: closets can only HELP, never HIDE ─────────────────────
|
||||
|
||||
|
||||
class TestHybridInvariant:
|
||||
def test_no_closets_degrades_to_direct_drawer_search(self, tmp_path):
|
||||
palace = str(tmp_path / "palace")
|
||||
_seed_drawers(palace)
|
||||
# No closets created.
|
||||
result = search_memories("Kafka rebalance timeout", palace, n_results=3)
|
||||
ids = [h["source_file"] for h in result["results"]]
|
||||
assert ids, "should return results"
|
||||
assert "fixture_D4.md" in ids, "direct drawer search alone should surface the Kafka drawer"
|
||||
|
||||
def test_weak_closets_do_not_hide_direct_drawer_hits(self, tmp_path):
|
||||
"""A closet that points at a wrong drawer must NOT suppress the
|
||||
drawer that direct search would have ranked first."""
|
||||
palace = str(tmp_path / "palace")
|
||||
_seed_drawers(palace)
|
||||
# Seed a misleading closet: it matches a generic phrase but points at D3.
|
||||
_seed_strong_closet_for(
|
||||
palace,
|
||||
drawer_id="D3",
|
||||
source_file="fixture_D3.md",
|
||||
topics=["Kafka queue tuning", "consumer rebalance config"],
|
||||
)
|
||||
result = search_memories("Kafka consumer rebalance timeout", palace, n_results=5)
|
||||
ids = [h["source_file"] for h in result["results"]]
|
||||
assert "fixture_D4.md" in ids, (
|
||||
"D4 must appear — direct drawer search alone would rank it first. "
|
||||
"Closet pointing to D3 should only boost D3, never hide D4."
|
||||
)
|
||||
|
||||
def test_closet_boost_lifts_matching_drawer(self, tmp_path):
|
||||
"""When a closet agrees with direct search, the matching drawer
|
||||
should be boosted to rank 1."""
|
||||
palace = str(tmp_path / "palace")
|
||||
_seed_drawers(palace)
|
||||
_seed_strong_closet_for(
|
||||
palace,
|
||||
drawer_id="D1",
|
||||
source_file="fixture_D1.md",
|
||||
topics=["JWT auth tokens", "session expiry", "authentication service"],
|
||||
)
|
||||
result = search_memories("JWT auth tokens expiry", palace, n_results=3)
|
||||
ids = [h["source_file"] for h in result["results"]]
|
||||
assert ids[0] == "fixture_D1.md"
|
||||
top = result["results"][0]
|
||||
assert top["matched_via"] == "drawer+closet"
|
||||
assert top["closet_boost"] > 0
|
||||
|
||||
|
||||
# ── closet_boost metadata ────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestClosetMetadata:
|
||||
def test_closet_preview_exposed_when_boosted(self, tmp_path):
|
||||
palace = str(tmp_path / "palace")
|
||||
_seed_drawers(palace)
|
||||
_seed_strong_closet_for(
|
||||
palace,
|
||||
drawer_id="D1",
|
||||
source_file="fixture_D1.md",
|
||||
topics=["JWT auth tokens", "24h expiry", "authentication"],
|
||||
)
|
||||
result = search_memories("JWT authentication", palace, n_results=2)
|
||||
top = result["results"][0]
|
||||
assert top["source_file"] == "fixture_D1.md"
|
||||
assert "closet_preview" in top
|
||||
|
||||
def test_drawer_only_hits_have_no_closet_preview(self, tmp_path):
|
||||
palace = str(tmp_path / "palace")
|
||||
_seed_drawers(palace)
|
||||
# No closets
|
||||
result = search_memories("TanStack Query", palace, n_results=2)
|
||||
assert result["results"]
|
||||
for h in result["results"]:
|
||||
assert h["matched_via"] == "drawer"
|
||||
assert "closet_preview" not in h
|
||||
assert h["closet_boost"] == 0.0
|
||||
Reference in New Issue
Block a user