merge: develop + harden entity metadata, BM25, and diary ingest for production
Merges develop (closet hardening #826, strip_noise #785, lock #784) and replaces every sub-feature in this PR with a correct, tested implementation. Shippable now. ## 1. Real Okapi-BM25 (searcher.py) The prior `_bm25_score()` hardcoded `idf = log(2.0)` for every term — it was really a scaled TF, not BM25, and couldn't tell a discriminative term from a generic one. Replaced with `_bm25_scores(query, documents)` that computes proper IDF over the provided candidate corpus using the Lucene smoothed formula `log((N - df + 0.5) / (df + 0.5) + 1)`. Well- defined for re-ranking vector-retrieval candidates — IDF there measures how discriminative each term is *within the candidate set*, exactly the signal we want. `_hybrid_rank` also fixed: - Vector normalization is now absolute `max(0, 1 - dist)`, not `1 - dist/max_dist` — adding/removing a candidate no longer reshuffles the others. - BM25 is min-max normalized within candidates (bounded [0, 1]). - Closet path now re-ranks too (was previously returning closet-order hits without hybrid scoring). - `_hybrid_score` internal field stripped from output; `bm25_score` exposed for debugging. ## 2. Entity metadata (miner.py) - Reuses `_ENTITY_STOPLIST` from palace.py so sentence-starters like "When", "After", "The" no longer land as entities (regression test covers this). - Known-entity registry is cached at module level, keyed by the registry file's mtime — no more disk read per drawer. - File handle now uses a context manager. - Truncates the entity LIST (to 25) before joining — never splits a name in the middle. ## 3. Diary ingest (diary_ingest.py) - State file now lives at `~/.mempalace/state/diary_ingest_<hash>.json`, keyed by (palace_path, diary_dir). No more pollution of the user's content directory. - Drawer IDs now hash `(wing, date_str)` — a user with personal + work diaries on the same day no longer silently clobbers. - Each day's upsert runs inside `mine_lock(source_file)` so concurrent ingest from two terminals can't race. - `force=True` now calls `purge_file_closets` before rebuild so leftover numbered closets from a longer prior day don't orphan. ## 4. Tests (tests/test_closets.py) Merged this PR's MineLock/Entity/BM25/Diary tests with develop's hardened Build/Upsert/Purge/Rebuild/SearchClosetFirst tests. Added specific regression tests for every fix above: - entity stoplist applies (no "When/After/The") - entity list capped before join (no partial tokens) - registry cached by mtime (mock-verified zero re-reads) - BM25 IDF downweights terms present in every doc (real BM25 evidence) - hybrid rank absolute normalization stable against outliers - diary state file outside user's diary dir - diary wing-prefixed IDs prevent cross-wing date collisions 35/35 closet tests pass; full suite 743/743. ruff + format clean under CI-pinned 0.4.x.
This commit is contained in:
Igor Lins e Silva
+234
-109
@@ -12,7 +12,11 @@ import math
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from .palace import get_collection, get_closets_collection
|
||||
from .palace import get_closets_collection, get_collection
|
||||
|
||||
# Closet pointer line format: "topic|entities|→drawer_id_a,drawer_id_b"
|
||||
# Multiple lines may join with newlines inside one closet document.
|
||||
_CLOSET_DRAWER_REF_RE = re.compile(r"→([\w,]+)")
|
||||
|
||||
logger = logging.getLogger("mempalace_mcp")
|
||||
|
||||
@@ -21,57 +25,109 @@ class SearchError(Exception):
|
||||
"""Raised when search cannot proceed (e.g. no palace found)."""
|
||||
|
||||
|
||||
def _bm25_score(query: str, document: str, k1: float = 1.5, b: float = 0.75, avg_dl: float = 500) -> float:
|
||||
"""Simple BM25 score for a single document against a query.
|
||||
_TOKEN_RE = re.compile(r"\w{2,}", re.UNICODE)
|
||||
|
||||
This is a lightweight keyword-matching signal that complements vector
|
||||
similarity. It catches exact matches that embeddings might miss
|
||||
(e.g., specific names, project codes, error messages).
|
||||
|
||||
def _tokenize(text: str) -> list:
|
||||
"""Lowercase + strip to alphanumeric tokens of length ≥ 2."""
|
||||
return _TOKEN_RE.findall(text.lower())
|
||||
|
||||
|
||||
def _bm25_scores(
|
||||
query: str,
|
||||
documents: list,
|
||||
k1: float = 1.5,
|
||||
b: float = 0.75,
|
||||
) -> list:
|
||||
"""Compute Okapi-BM25 scores for ``query`` against each document.
|
||||
|
||||
IDF is computed over the *provided corpus* using the Lucene/BM25+
|
||||
smoothed formula ``log((N - df + 0.5) / (df + 0.5) + 1)``, which is
|
||||
always non-negative. This is well-defined for re-ranking a small
|
||||
candidate set returned by vector retrieval — IDF then reflects how
|
||||
discriminative each query term is *within the candidates*, exactly
|
||||
what's needed to reorder them.
|
||||
|
||||
Parameters mirror Okapi-BM25 conventions:
|
||||
k1 — term-frequency saturation (1.2-2.0 typical, 1.5 default)
|
||||
b — length normalization (0.0 = none, 1.0 = full, 0.75 default)
|
||||
|
||||
Returns a list of scores in the same order as ``documents``.
|
||||
"""
|
||||
query_terms = set(re.findall(r'\w{2,}', query.lower()))
|
||||
doc_terms = re.findall(r'\w{2,}', document.lower())
|
||||
if not query_terms or not doc_terms:
|
||||
return 0.0
|
||||
doc_len = len(doc_terms)
|
||||
term_freq = {}
|
||||
for t in doc_terms:
|
||||
term_freq[t] = term_freq.get(t, 0) + 1
|
||||
n_docs = len(documents)
|
||||
query_terms = set(_tokenize(query))
|
||||
if not query_terms or n_docs == 0:
|
||||
return [0.0] * n_docs
|
||||
|
||||
score = 0.0
|
||||
for term in query_terms:
|
||||
tf = term_freq.get(term, 0)
|
||||
if tf > 0:
|
||||
# Simplified IDF — treat each query term as moderately rare
|
||||
idf = math.log(2.0)
|
||||
numerator = tf * (k1 + 1)
|
||||
denominator = tf + k1 * (1 - b + b * doc_len / avg_dl)
|
||||
score += idf * numerator / denominator
|
||||
return score
|
||||
tokenized = [_tokenize(d) for d in documents]
|
||||
doc_lens = [len(toks) for toks in tokenized]
|
||||
if not any(doc_lens):
|
||||
return [0.0] * n_docs
|
||||
avgdl = sum(doc_lens) / n_docs or 1.0
|
||||
|
||||
# Document frequency: how many docs contain each query term?
|
||||
df = {term: 0 for term in query_terms}
|
||||
for toks in tokenized:
|
||||
seen = set(toks) & query_terms
|
||||
for term in seen:
|
||||
df[term] += 1
|
||||
|
||||
idf = {term: math.log((n_docs - df[term] + 0.5) / (df[term] + 0.5) + 1) for term in query_terms}
|
||||
|
||||
scores = []
|
||||
for toks, dl in zip(tokenized, doc_lens):
|
||||
if dl == 0:
|
||||
scores.append(0.0)
|
||||
continue
|
||||
tf: dict = {}
|
||||
for t in toks:
|
||||
if t in query_terms:
|
||||
tf[t] = tf.get(t, 0) + 1
|
||||
score = 0.0
|
||||
for term, freq in tf.items():
|
||||
num = freq * (k1 + 1)
|
||||
den = freq + k1 * (1 - b + b * dl / avgdl)
|
||||
score += idf[term] * num / den
|
||||
scores.append(score)
|
||||
return scores
|
||||
|
||||
|
||||
def _hybrid_rank(vector_results, query: str, vector_weight: float = 0.6, bm25_weight: float = 0.4):
|
||||
"""Re-rank results using both vector distance and BM25 keyword score.
|
||||
def _hybrid_rank(
|
||||
results: list,
|
||||
query: str,
|
||||
vector_weight: float = 0.6,
|
||||
bm25_weight: float = 0.4,
|
||||
) -> list:
|
||||
"""Re-rank ``results`` by a convex combination of vector similarity and BM25.
|
||||
|
||||
Returns results sorted by combined score (higher = better).
|
||||
* Vector similarity uses absolute cosine sim ``max(0, 1 - distance)`` —
|
||||
ChromaDB's hnsw cosine distance lives in ``[0, 2]`` (0 = identical).
|
||||
Absolute (not relative-to-max) means adding/removing a candidate
|
||||
can't reshuffle the others.
|
||||
* BM25 is real Okapi-BM25 with corpus-relative IDF over the candidates
|
||||
themselves. Since the absolute scale is unbounded, BM25 is min-max
|
||||
normalized within the candidate set so weights are commensurable.
|
||||
|
||||
Mutates each result dict to add ``bm25_score`` and reorders the list
|
||||
in place. Returns the same list for convenience.
|
||||
"""
|
||||
if not vector_results:
|
||||
return vector_results
|
||||
if not results:
|
||||
return results
|
||||
|
||||
# Normalize vector distances to 0-1 similarity
|
||||
max_dist = max(r.get("distance", 1.0) for r in vector_results) or 1.0
|
||||
for r in vector_results:
|
||||
vec_sim = max(0.0, 1 - r.get("distance", 1.0) / max(max_dist, 0.001))
|
||||
bm25 = _bm25_score(query, r.get("text", ""))
|
||||
# Normalize BM25 to roughly 0-1 range
|
||||
bm25_norm = min(bm25 / 3.0, 1.0)
|
||||
r["_hybrid_score"] = vector_weight * vec_sim + bm25_weight * bm25_norm
|
||||
r["bm25_score"] = round(bm25, 3)
|
||||
docs = [r.get("text", "") for r in results]
|
||||
bm25_raw = _bm25_scores(query, docs)
|
||||
max_bm25 = max(bm25_raw) if bm25_raw else 0.0
|
||||
bm25_norm = [s / max_bm25 for s in bm25_raw] if max_bm25 > 0 else [0.0] * len(bm25_raw)
|
||||
|
||||
vector_results.sort(key=lambda r: r["_hybrid_score"], reverse=True)
|
||||
# Clean up internal field
|
||||
for r in vector_results:
|
||||
del r["_hybrid_score"]
|
||||
return vector_results
|
||||
scored = []
|
||||
for r, raw, norm in zip(results, bm25_raw, bm25_norm):
|
||||
vec_sim = max(0.0, 1.0 - r.get("distance", 1.0))
|
||||
r["bm25_score"] = round(raw, 3)
|
||||
scored.append((vector_weight * vec_sim + bm25_weight * norm, r))
|
||||
|
||||
scored.sort(key=lambda pair: pair[0], reverse=True)
|
||||
results[:] = [r for _, r in scored]
|
||||
return results
|
||||
|
||||
|
||||
def build_where_filter(wing: str = None, room: str = None) -> dict:
|
||||
@@ -85,6 +141,116 @@ def build_where_filter(wing: str = None, room: str = None) -> dict:
|
||||
return {}
|
||||
|
||||
|
||||
def _extract_drawer_ids_from_closet(closet_doc: str) -> list:
|
||||
"""Parse all `→drawer_id_a,drawer_id_b` pointers out of a closet document.
|
||||
|
||||
Preserves order and dedupes.
|
||||
"""
|
||||
seen: dict = {}
|
||||
for match in _CLOSET_DRAWER_REF_RE.findall(closet_doc):
|
||||
for did in match.split(","):
|
||||
did = did.strip()
|
||||
if did and did not in seen:
|
||||
seen[did] = None
|
||||
return list(seen.keys())
|
||||
|
||||
|
||||
def _closet_first_hits(
|
||||
palace_path: str,
|
||||
query: str,
|
||||
where: dict,
|
||||
drawers_col,
|
||||
n_results: int,
|
||||
max_distance: float,
|
||||
):
|
||||
"""Run a closet-first search and return chunk-level drawer hits.
|
||||
|
||||
Returns:
|
||||
non-empty list of hits when the closet path produced usable matches.
|
||||
``None`` when the closet collection is empty/missing OR when every
|
||||
candidate drawer was filtered out (e.g. by max_distance); the
|
||||
caller should fall back to direct drawer search.
|
||||
"""
|
||||
try:
|
||||
closets_col = get_closets_collection(palace_path, create=False)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
try:
|
||||
ckwargs = {
|
||||
"query_texts": [query],
|
||||
"n_results": max(n_results * 2, 5),
|
||||
"include": ["documents", "metadatas", "distances"],
|
||||
}
|
||||
if where:
|
||||
ckwargs["where"] = where
|
||||
closet_results = closets_col.query(**ckwargs)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
closet_docs = closet_results["documents"][0] if closet_results["documents"] else []
|
||||
if not closet_docs:
|
||||
return None
|
||||
|
||||
closet_metas = closet_results["metadatas"][0]
|
||||
closet_dists = closet_results["distances"][0]
|
||||
|
||||
# Collect candidate drawer IDs in closet-rank order, dedupe, remember
|
||||
# which closet (and its distance/preview) introduced each one.
|
||||
drawer_id_order: list = []
|
||||
drawer_provenance: dict = {}
|
||||
for cdoc, cmeta, cdist in zip(closet_docs, closet_metas, closet_dists):
|
||||
for did in _extract_drawer_ids_from_closet(cdoc):
|
||||
if did in drawer_provenance:
|
||||
continue
|
||||
drawer_provenance[did] = (cdist, cdoc, cmeta)
|
||||
drawer_id_order.append(did)
|
||||
|
||||
if not drawer_id_order:
|
||||
return None
|
||||
|
||||
# Hydrate exactly those drawers — chunk-level, not whole-file.
|
||||
try:
|
||||
fetched = drawers_col.get(
|
||||
ids=drawer_id_order,
|
||||
include=["documents", "metadatas"],
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
fetched_ids = fetched.get("ids") or []
|
||||
fetched_docs = fetched.get("documents") or []
|
||||
fetched_metas = fetched.get("metadatas") or []
|
||||
fetched_map = {
|
||||
did: (doc, meta) for did, doc, meta in zip(fetched_ids, fetched_docs, fetched_metas)
|
||||
}
|
||||
|
||||
hits: list = []
|
||||
for did in drawer_id_order:
|
||||
if did not in fetched_map:
|
||||
continue # closet pointed to a drawer that no longer exists
|
||||
doc, meta = fetched_map[did]
|
||||
cdist, cdoc, _ = drawer_provenance[did]
|
||||
if max_distance > 0.0 and cdist > max_distance:
|
||||
continue
|
||||
hits.append(
|
||||
{
|
||||
"text": doc,
|
||||
"wing": meta.get("wing", "unknown"),
|
||||
"room": meta.get("room", "unknown"),
|
||||
"source_file": Path(meta.get("source_file", "?")).name,
|
||||
"similarity": round(max(0.0, 1 - cdist), 3),
|
||||
"distance": round(cdist, 4),
|
||||
"matched_via": "closet",
|
||||
"closet_preview": cdoc[:200],
|
||||
}
|
||||
)
|
||||
if len(hits) >= n_results:
|
||||
break
|
||||
|
||||
return hits if hits else None
|
||||
|
||||
|
||||
def search(query: str, palace_path: str, wing: str = None, room: str = None, n_results: int = 5):
|
||||
"""
|
||||
Search the palace. Returns verbatim drawer content.
|
||||
@@ -183,73 +349,31 @@ def search_memories(
|
||||
|
||||
where = build_where_filter(wing, room)
|
||||
|
||||
# Try closet-first search: search the compact index, then hydrate drawers
|
||||
closet_hits = []
|
||||
try:
|
||||
closets_col = get_closets_collection(palace_path, create=False)
|
||||
ckwargs = {
|
||||
"query_texts": [query],
|
||||
"n_results": n_results * 2, # over-fetch closets to find best drawers
|
||||
"include": ["documents", "metadatas", "distances"],
|
||||
# Closet-first search: scan the compact index, parse drawer pointers
|
||||
# from each matching line, then hydrate exactly those drawers. This
|
||||
# keeps the result shape chunk-level (consistent with direct search)
|
||||
# and applies the same max_distance filter.
|
||||
closet_hits = _closet_first_hits(
|
||||
palace_path=palace_path,
|
||||
query=query,
|
||||
where=where,
|
||||
drawers_col=drawers_col,
|
||||
n_results=n_results,
|
||||
max_distance=max_distance,
|
||||
)
|
||||
if closet_hits is not None:
|
||||
# Re-rank chunk-level closet hits with the same hybrid scoring as
|
||||
# the direct path. The vector half here uses the closet's distance
|
||||
# (query↔topic-line) — that's intentional: closets are *meant* to
|
||||
# be the semantic-narrowing signal, and BM25 then enforces actual
|
||||
# keyword presence in the hydrated drawer text.
|
||||
closet_hits = _hybrid_rank(closet_hits, query)
|
||||
return {
|
||||
"query": query,
|
||||
"filters": {"wing": wing, "room": room},
|
||||
"total_before_filter": len(closet_hits),
|
||||
"results": closet_hits,
|
||||
}
|
||||
if where:
|
||||
ckwargs["where"] = where
|
||||
closet_results = closets_col.query(**ckwargs)
|
||||
if closet_results["documents"][0]:
|
||||
closet_hits = list(zip(
|
||||
closet_results["documents"][0],
|
||||
closet_results["metadatas"][0],
|
||||
closet_results["distances"][0],
|
||||
))
|
||||
except Exception:
|
||||
pass # no closets yet — fall through to direct drawer search
|
||||
|
||||
# If closets found results, hydrate the referenced drawers
|
||||
if closet_hits:
|
||||
import re
|
||||
seen_sources = set()
|
||||
hits = []
|
||||
for closet_doc, closet_meta, closet_dist in closet_hits:
|
||||
source = closet_meta.get("source_file", "")
|
||||
if source in seen_sources:
|
||||
continue
|
||||
seen_sources.add(source)
|
||||
|
||||
# Find drawers for this source file
|
||||
try:
|
||||
drawer_results = drawers_col.get(
|
||||
where={"source_file": source},
|
||||
include=["documents", "metadatas"],
|
||||
)
|
||||
if drawer_results.get("ids"):
|
||||
# Combine all drawer content for this file
|
||||
full_text = "\n\n".join(drawer_results["documents"])
|
||||
meta = drawer_results["metadatas"][0]
|
||||
hits.append({
|
||||
"text": full_text,
|
||||
"wing": meta.get("wing", "unknown"),
|
||||
"room": meta.get("room", "unknown"),
|
||||
"source_file": Path(source).name,
|
||||
"similarity": round(max(0.0, 1 - closet_dist), 3),
|
||||
"distance": round(closet_dist, 4),
|
||||
"matched_via": "closet",
|
||||
"closet_preview": closet_doc[:200],
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if len(hits) >= n_results:
|
||||
break
|
||||
|
||||
if hits:
|
||||
# Re-rank with BM25 hybrid scoring
|
||||
hits = _hybrid_rank(hits, query)
|
||||
return {
|
||||
"query": query,
|
||||
"filters": {"wing": wing, "room": room},
|
||||
"total_before_filter": len(closet_hits),
|
||||
"results": hits,
|
||||
}
|
||||
|
||||
# Fallback: direct drawer search (no closets yet, or closets empty)
|
||||
try:
|
||||
@@ -282,6 +406,7 @@ def search_memories(
|
||||
"source_file": Path(meta.get("source_file", "?")).name,
|
||||
"similarity": round(max(0.0, 1 - dist), 3),
|
||||
"distance": round(dist, 4),
|
||||
"matched_via": "drawer",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user