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
|
||||
Reference in New Issue
Block a user