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:
Igor Lins e Silva
2026-04-13 18:40:36 -03:00
5 changed files with 1012 additions and 181 deletions
+339
View File
@@ -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
View File
@@ -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"]
+133
View File
@@ -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