Merge remote-tracking branch 'origin/pr/drawer-grep-search' into pr/fact-checker
This commit is contained in:
+23
-14
@@ -32,13 +32,11 @@ Topics are never split across closets. If adding a topic would exceed 1,500 char
|
||||
|
||||
### When do closets update?
|
||||
|
||||
When a file is re-mined (content changed), its drawers are replaced and new closets are built from the fresh content. The old closet content is replaced via upsert.
|
||||
When a file is re-mined (content changed, or `NORMALIZE_VERSION` was bumped), the miner first deletes every closet for that source file (`purge_file_closets`) and then writes a fresh set. Stale topics from the prior mine are gone — closets are always a snapshot of the current content, never an accumulation across runs.
|
||||
|
||||
### What about stale topics?
|
||||
|
||||
If a file's content changes and a topic no longer exists, the closet is rebuilt entirely from the new content — stale topics are gone. Closets are tied to source files, not to individual topics.
|
||||
|
||||
If you add content to an existing file (e.g., a daily diary growing throughout the day), new topics are appended to the existing closet until the 1,500-char limit, then a new closet is created.
|
||||
There are no stale topics: each re-mine is a clean rebuild for that source file. If a file gets larger and produces fewer or more closets than last time, the leftover numbered closets from the larger run are still purged because the delete is done by `source_file`, not by ID.
|
||||
|
||||
### Do closets survive palace rebuilds?
|
||||
|
||||
@@ -49,31 +47,42 @@ Closets are stored in the `mempalace_closets` ChromaDB collection alongside `mem
|
||||
```
|
||||
Query → search mempalace_closets (fast, small documents)
|
||||
↓
|
||||
top closet hits → extract drawer IDs from pointer lines
|
||||
top closet hits → parse `→drawer_id_a,drawer_id_b` pointers
|
||||
↓
|
||||
fetch drawers from mempalace_drawers (full verbatim content)
|
||||
fetch exactly those drawers from mempalace_drawers (verbatim content)
|
||||
↓
|
||||
BM25 hybrid re-rank (keyword match + vector similarity)
|
||||
apply max_distance filter
|
||||
↓
|
||||
return results to user
|
||||
return chunk-level results (same shape as direct search)
|
||||
```
|
||||
|
||||
If no closets exist (palace created before this feature), search falls back to direct drawer search. Closets are created on next mine.
|
||||
Hits carry `matched_via: "closet"` (or `"drawer"` for the fallback path) plus a `closet_preview` field showing the line that surfaced them.
|
||||
|
||||
If no closets exist (palace created before this feature) — or all closet hits get filtered out by `max_distance` — search falls back to direct drawer search. Closets are created on next mine.
|
||||
|
||||
> **BM25 hybrid re-rank** is on the roadmap (deferred to a follow-up PR alongside generic `LLM_*` env-var support); the current closet search ranks purely by ChromaDB cosine distance against the closet text.
|
||||
|
||||
## Limits
|
||||
|
||||
| Setting | Value | Reason |
|
||||
|---------|-------|--------|
|
||||
| Max closet size | 1,500 chars | Leaves buffer under ChromaDB's working limit |
|
||||
| Max closet size | 1,500 chars (`CLOSET_CHAR_LIMIT`) | Leaves buffer under ChromaDB's working limit |
|
||||
| Source content scanned | 5,000 chars (`CLOSET_EXTRACT_WINDOW`) | Caps regex extraction cost on long files; back-of-file content is currently invisible to closet extraction (tracked for follow-up) |
|
||||
| Max topics per file | 12 | Keeps closets focused |
|
||||
| Max quotes per file | 3 | Most relevant only |
|
||||
| Max entities per pointer | 5 | Top names by frequency |
|
||||
| Max response chars | 10,000 | Prevents hydration blowup on large files |
|
||||
| Max entities per pointer | 5 | Top names by frequency, after stoplist filtering |
|
||||
|
||||
## For developers
|
||||
|
||||
Closet functions live in `mempalace/palace.py`:
|
||||
- `get_closets_collection()` — get the closets ChromaDB collection
|
||||
- `build_closet_lines()` — extract topics/entities/quotes into pointer lines
|
||||
- `upsert_closet_lines()` — write lines to closets respecting the char limit
|
||||
- `CLOSET_CHAR_LIMIT` — the 1,500 char limit constant
|
||||
- `upsert_closet_lines()` — write lines to closets respecting the char limit (overwrites existing IDs; does not append — call `purge_file_closets` first when re-mining)
|
||||
- `purge_file_closets()` — delete every closet for a given source file before rebuild
|
||||
- `CLOSET_CHAR_LIMIT` / `CLOSET_EXTRACT_WINDOW` — size constants
|
||||
|
||||
The closet-first search path lives in `mempalace/searcher.py`:
|
||||
- `_extract_drawer_ids_from_closet()` — parse `→drawer_a,drawer_b` pointers out of a closet document
|
||||
- `_closet_first_hits()` — query closets, parse pointers, hydrate matching drawers, return chunk-level hits or `None` to fall back
|
||||
|
||||
Note: only the project miner (`miner.py::process_file`) builds closets today. Conversation-mined wings (Claude Code JSONL, ChatGPT export, etc.) will keep using direct drawer search via the searcher fallback until the convo-closet PR lands.
|
||||
|
||||
+5
-1
@@ -133,6 +133,10 @@ Example output:
|
||||
[14:40:01] Session abc123: 18 exchanges, 3 since last save
|
||||
```
|
||||
|
||||
## Known Limitations
|
||||
|
||||
**Hooks require session restart after install.** Claude Code loads hooks from `settings.json` at session start only. If you run `mempalace init` or manually edit hook config mid-session, the hooks won't fire until you restart Claude Code. This is a Claude Code limitation.
|
||||
|
||||
## Cost
|
||||
|
||||
**Zero extra tokens.** The hooks are bash scripts that run locally. They don't call any API. The only "cost" is the AI spending a few seconds organizing memories at each checkpoint — and it's doing that with context it already has loaded.
|
||||
**Zero extra tokens.** The hooks notify the AI that saves happened in the background — the AI doesn't need to write anything in the chat. All filing is handled automatically. Previous versions asked the AI to write diary entries and drawer content in the chat window, which cost ~$1/session in retransmitted tokens.
|
||||
|
||||
@@ -68,10 +68,10 @@ if [ -n "$MEMPAL_DIR" ] && [ -d "$MEMPAL_DIR" ]; then
|
||||
python3 -m mempalace mine "$MEMPAL_DIR" >> "$STATE_DIR/hook.log" 2>&1
|
||||
fi
|
||||
|
||||
# Always block — compaction = save everything
|
||||
# Notify — compaction is about to happen but filing is handled in background
|
||||
cat << 'HOOKJSON'
|
||||
{
|
||||
"decision": "block",
|
||||
"reason": "COMPACTION IMMINENT. Save ALL topics, decisions, quotes, code, and important context from this session to your memory system. Be thorough — after compaction, detailed context will be lost. Organize into appropriate categories. Use verbatim quotes where possible. Save everything, then allow compaction to proceed."
|
||||
"decision": "allow",
|
||||
"reason": "MemPalace pre-compaction save. Your full conversation has been saved verbatim in the background — no action needed. Compaction can proceed safely."
|
||||
}
|
||||
HOOKJSON
|
||||
|
||||
@@ -140,12 +140,15 @@ if [ "$SINCE_LAST" -ge "$SAVE_INTERVAL" ] && [ "$EXCHANGE_COUNT" -gt 0 ]; then
|
||||
python3 -m mempalace mine "$MEMPAL_DIR" >> "$STATE_DIR/hook.log" 2>&1 &
|
||||
fi
|
||||
|
||||
# Block the AI and tell it to save
|
||||
# The "reason" becomes a system message the AI sees and acts on
|
||||
# Notify the AI that a checkpoint happened — but do NOT ask it to write
|
||||
# anything in chat. All filing happens in the background via the pipeline.
|
||||
# The old version asked the agent to write diary entries, add drawers, and
|
||||
# add KG triples in the chat window — that cost ~$1/session in retransmitted
|
||||
# tokens and cluttered the conversation.
|
||||
cat << 'HOOKJSON'
|
||||
{
|
||||
"decision": "block",
|
||||
"reason": "AUTO-SAVE checkpoint. Save key topics, decisions, quotes, and code from this session to your memory system. Organize into appropriate categories. Use verbatim quotes where possible. Continue conversation after saving."
|
||||
"decision": "allow",
|
||||
"reason": "MemPalace auto-save checkpoint. Your conversation is being saved verbatim in the background — no action needed from you. Continue working."
|
||||
}
|
||||
HOOKJSON
|
||||
else
|
||||
|
||||
+74
-35
@@ -16,7 +16,13 @@ from datetime import datetime
|
||||
from collections import defaultdict
|
||||
|
||||
from .normalize import normalize
|
||||
from .palace import SKIP_DIRS, get_collection, file_already_mined, mine_lock
|
||||
from .palace import (
|
||||
NORMALIZE_VERSION,
|
||||
SKIP_DIRS,
|
||||
file_already_mined,
|
||||
get_collection,
|
||||
mine_lock,
|
||||
)
|
||||
|
||||
|
||||
# File types that might contain conversations
|
||||
@@ -51,6 +57,7 @@ def _register_file(collection, source_file: str, wing: str, agent: str):
|
||||
"added_by": agent,
|
||||
"filed_at": datetime.now().isoformat(),
|
||||
"ingest_mode": "registry",
|
||||
"normalize_version": NORMALIZE_VERSION,
|
||||
}
|
||||
],
|
||||
)
|
||||
@@ -272,6 +279,62 @@ def scan_convos(convo_dir: str) -> list:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _file_chunks_locked(collection, source_file, chunks, wing, room, agent, extract_mode):
|
||||
"""Lock the source file, purge stale drawers, and upsert fresh chunks.
|
||||
|
||||
Combines the per-file serialization that prevents concurrent agents from
|
||||
duplicating work (via mine_lock) with the normalize-version rebuild
|
||||
contract (purge-before-insert so pre-v2 drawers don't survive).
|
||||
|
||||
Returns (drawers_added, room_counts_delta, skipped).
|
||||
"""
|
||||
room_counts_delta: dict = defaultdict(int)
|
||||
drawers_added = 0
|
||||
with mine_lock(source_file):
|
||||
# Re-check after lock — another agent may have just finished this file
|
||||
# at the current schema. A stale-version hit here returns False, so we
|
||||
# still fall through to the purge+rebuild path below.
|
||||
if file_already_mined(collection, source_file):
|
||||
return 0, room_counts_delta, True
|
||||
|
||||
# Purge stale drawers first. When the normalize schema bumps,
|
||||
# file_already_mined() returned False for pre-v2 drawers — clean
|
||||
# them out so the source doesn't end up with mixed old/new drawers.
|
||||
try:
|
||||
collection.delete(where={"source_file": source_file})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for chunk in chunks:
|
||||
chunk_room = chunk.get("memory_type", room) if extract_mode == "general" else room
|
||||
if extract_mode == "general":
|
||||
room_counts_delta[chunk_room] += 1
|
||||
drawer_id = f"drawer_{wing}_{chunk_room}_{hashlib.sha256((source_file + str(chunk['chunk_index'])).encode()).hexdigest()[:24]}"
|
||||
try:
|
||||
collection.upsert(
|
||||
documents=[chunk["content"]],
|
||||
ids=[drawer_id],
|
||||
metadatas=[
|
||||
{
|
||||
"wing": wing,
|
||||
"room": chunk_room,
|
||||
"source_file": source_file,
|
||||
"chunk_index": chunk["chunk_index"],
|
||||
"added_by": agent,
|
||||
"filed_at": datetime.now().isoformat(),
|
||||
"ingest_mode": "convos",
|
||||
"extract_mode": extract_mode,
|
||||
"normalize_version": NORMALIZE_VERSION,
|
||||
}
|
||||
],
|
||||
)
|
||||
drawers_added += 1
|
||||
except Exception as e:
|
||||
if "already exists" not in str(e).lower():
|
||||
raise
|
||||
return drawers_added, room_counts_delta, False
|
||||
|
||||
|
||||
def mine_convos(
|
||||
convo_dir: str,
|
||||
palace_path: str,
|
||||
@@ -375,40 +438,16 @@ def mine_convos(
|
||||
if extract_mode != "general":
|
||||
room_counts[room] += 1
|
||||
|
||||
# File each chunk — lock to prevent concurrent agents duplicating
|
||||
drawers_added = 0
|
||||
with mine_lock(source_file):
|
||||
# Re-check after lock — another agent may have just finished this file
|
||||
if file_already_mined(collection, source_file):
|
||||
files_skipped += 1
|
||||
continue
|
||||
|
||||
for chunk in chunks:
|
||||
chunk_room = chunk.get("memory_type", room) if extract_mode == "general" else room
|
||||
if extract_mode == "general":
|
||||
room_counts[chunk_room] += 1
|
||||
drawer_id = f"drawer_{wing}_{chunk_room}_{hashlib.sha256((source_file + str(chunk['chunk_index'])).encode()).hexdigest()[:24]}"
|
||||
try:
|
||||
collection.upsert(
|
||||
documents=[chunk["content"]],
|
||||
ids=[drawer_id],
|
||||
metadatas=[
|
||||
{
|
||||
"wing": wing,
|
||||
"room": chunk_room,
|
||||
"source_file": source_file,
|
||||
"chunk_index": chunk["chunk_index"],
|
||||
"added_by": agent,
|
||||
"filed_at": datetime.now().isoformat(),
|
||||
"ingest_mode": "convos",
|
||||
"extract_mode": extract_mode,
|
||||
}
|
||||
],
|
||||
)
|
||||
drawers_added += 1
|
||||
except Exception as e:
|
||||
if "already exists" not in str(e).lower():
|
||||
raise
|
||||
# Lock + purge stale + file fresh chunks. Lock serializes concurrent
|
||||
# agents; purge removes pre-v2 drawers so the schema bump applies.
|
||||
drawers_added, room_delta, skipped = _file_chunks_locked(
|
||||
collection, source_file, chunks, wing, room, agent, extract_mode
|
||||
)
|
||||
if skipped:
|
||||
files_skipped += 1
|
||||
continue
|
||||
for r, n in room_delta.items():
|
||||
room_counts[r] += n
|
||||
|
||||
total_drawers += drawers_added
|
||||
print(f" ✓ [{i:4}/{len(files)}] {filepath.name[:50]:50} +{drawers_added}")
|
||||
|
||||
+110
-74
@@ -2,10 +2,14 @@
|
||||
diary_ingest.py — Ingest daily summary files into the palace.
|
||||
|
||||
Architecture:
|
||||
- ONE drawer per day — full verbatim content, upserted as the day grows
|
||||
- Closets pack topics up to 1500 chars, never split mid-topic
|
||||
- Only new entries are processed (tracks entry count in state file)
|
||||
- Entities extracted and stamped on metadata for filterable search
|
||||
- ONE drawer per (wing, day) — full verbatim content, upserted as the day grows.
|
||||
- Closets pack topics up to CLOSET_CHAR_LIMIT, never split mid-topic.
|
||||
- A re-ingest fully purges the prior day's closets before rebuilding so a
|
||||
shorter day never leaves orphans behind.
|
||||
- Only new entries are processed by default (tracks entry count in a state
|
||||
file under ``~/.mempalace/state/`` — never inside the user's diary dir).
|
||||
- Per-file ``mine_lock`` so concurrent ingest from two terminals can't race.
|
||||
- Entities extracted and stamped on metadata for filterable search.
|
||||
|
||||
Usage:
|
||||
python -m mempalace.diary_ingest --dir ~/daily_summaries --palace ~/.mempalace/palace
|
||||
@@ -19,19 +23,32 @@ import re
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
from .palace import (
|
||||
get_collection,
|
||||
get_closets_collection,
|
||||
build_closet_lines,
|
||||
upsert_closet_lines,
|
||||
CLOSET_CHAR_LIMIT,
|
||||
)
|
||||
from .miner import _extract_entities_for_metadata
|
||||
|
||||
from .palace import (
|
||||
build_closet_lines,
|
||||
get_closets_collection,
|
||||
get_collection,
|
||||
mine_lock,
|
||||
purge_file_closets,
|
||||
upsert_closet_lines,
|
||||
)
|
||||
|
||||
DIARY_ENTRY_RE = re.compile(r"^## .+", re.MULTILINE)
|
||||
|
||||
|
||||
def _state_file_for(palace_path: str, diary_dir: Path) -> Path:
|
||||
"""Return the per-(palace, diary-dir) state-file path under ~/.mempalace/state.
|
||||
|
||||
Keyed by sha256 of (palace_path, diary_dir) so multiple diary folders
|
||||
pointing at the same palace each get an independent state file. The
|
||||
state file is *never* written inside the user's diary directory.
|
||||
"""
|
||||
state_root = Path(os.path.expanduser("~")) / ".mempalace" / "state"
|
||||
state_root.mkdir(parents=True, exist_ok=True)
|
||||
key = hashlib.sha256(f"{palace_path}|{diary_dir}".encode()).hexdigest()[:24]
|
||||
return state_root / f"diary_ingest_{key}.json"
|
||||
|
||||
|
||||
def _split_entries(text):
|
||||
"""Split diary text into (header, body) pairs per ## entry."""
|
||||
parts = DIARY_ENTRY_RE.split(text)
|
||||
@@ -43,6 +60,18 @@ def _split_entries(text):
|
||||
return entries
|
||||
|
||||
|
||||
def _diary_drawer_id(wing: str, date_str: str) -> str:
|
||||
"""Stable, wing-scoped drawer ID. Two diaries (e.g. 'work' vs 'personal')
|
||||
sharing the same date never collide."""
|
||||
suffix = hashlib.sha256(f"{wing}|{date_str}".encode()).hexdigest()[:24]
|
||||
return f"drawer_diary_{suffix}"
|
||||
|
||||
|
||||
def _diary_closet_id_base(wing: str, date_str: str) -> str:
|
||||
suffix = hashlib.sha256(f"{wing}|{date_str}".encode()).hexdigest()[:24]
|
||||
return f"closet_diary_{suffix}"
|
||||
|
||||
|
||||
def ingest_diaries(
|
||||
diary_dir,
|
||||
palace_path,
|
||||
@@ -51,24 +80,29 @@ def ingest_diaries(
|
||||
):
|
||||
"""Ingest daily summary files into the palace.
|
||||
|
||||
Each date file gets ONE drawer (upserted as day grows) and
|
||||
closets that pack topics atomically up to 1500 chars.
|
||||
Each date file gets ONE drawer keyed by ``(wing, date)`` and closets that
|
||||
pack topics atomically up to ``CLOSET_CHAR_LIMIT``. ``force=True`` rebuilds
|
||||
every entry's closets from scratch (purging stale ones); the default
|
||||
incremental mode only processes entries appended since the last run.
|
||||
"""
|
||||
diary_dir = Path(diary_dir).expanduser().resolve()
|
||||
if not diary_dir.exists():
|
||||
print(f"Diary directory not found: {diary_dir}")
|
||||
return
|
||||
return {"days_updated": 0, "closets_created": 0}
|
||||
|
||||
diary_files = sorted(diary_dir.glob("*.md"))
|
||||
if not diary_files:
|
||||
print(f"No .md files in {diary_dir}")
|
||||
return
|
||||
return {"days_updated": 0, "closets_created": 0}
|
||||
|
||||
# State tracks which entries have been closeted per file
|
||||
state_file = diary_dir / ".diary_ingest_state.json"
|
||||
state = {} if force else (
|
||||
json.loads(state_file.read_text()) if state_file.exists() else {}
|
||||
)
|
||||
state_file = _state_file_for(str(palace_path), diary_dir)
|
||||
if force or not state_file.exists():
|
||||
state: dict = {}
|
||||
else:
|
||||
try:
|
||||
state = json.loads(state_file.read_text())
|
||||
except Exception:
|
||||
state = {}
|
||||
|
||||
drawers_col = get_collection(palace_path)
|
||||
closets_col = get_closets_collection(palace_path)
|
||||
@@ -87,70 +121,72 @@ def ingest_diaries(
|
||||
date_str = date_match.group(1)
|
||||
|
||||
# Skip if content hasn't changed
|
||||
prev_size = state.get(diary_path.name, {}).get("size", 0)
|
||||
state_key = f"{wing}|{diary_path.name}"
|
||||
prev_size = state.get(state_key, {}).get("size", 0)
|
||||
curr_size = len(text)
|
||||
if curr_size == prev_size and not force:
|
||||
continue
|
||||
|
||||
now_iso = datetime.now(timezone.utc).isoformat()
|
||||
drawer_id = f"drawer_diary_{date_str}"
|
||||
|
||||
# Extract entities from full day text
|
||||
drawer_id = _diary_drawer_id(wing, date_str)
|
||||
entities = _extract_entities_for_metadata(text)
|
||||
source_file = str(diary_path)
|
||||
|
||||
# UPSERT the day's drawer (full verbatim, replaces as day grows)
|
||||
drawer_meta = {
|
||||
"date": date_str,
|
||||
"wing": wing,
|
||||
"room": "daily",
|
||||
"source_file": str(diary_path),
|
||||
"source_session": "daily_diary",
|
||||
"filed_at": now_iso,
|
||||
}
|
||||
if entities:
|
||||
drawer_meta["entities"] = entities
|
||||
drawers_col.upsert(
|
||||
documents=[text],
|
||||
ids=[drawer_id],
|
||||
metadatas=[drawer_meta],
|
||||
)
|
||||
# Serialize per source — two terminals running ingest at once must
|
||||
# not interleave the upsert + closet-rebuild.
|
||||
with mine_lock(source_file):
|
||||
drawer_meta = {
|
||||
"date": date_str,
|
||||
"wing": wing,
|
||||
"room": "daily",
|
||||
"source_file": source_file,
|
||||
"source_session": "daily_diary",
|
||||
"filed_at": now_iso,
|
||||
}
|
||||
if entities:
|
||||
drawer_meta["entities"] = entities
|
||||
drawers_col.upsert(
|
||||
documents=[text],
|
||||
ids=[drawer_id],
|
||||
metadatas=[drawer_meta],
|
||||
)
|
||||
|
||||
# Split into entries and find new ones
|
||||
entries = _split_entries(text)
|
||||
prev_entry_count = state.get(diary_path.name, {}).get("entry_count", 0)
|
||||
new_entries = entries[prev_entry_count:] if not force else entries
|
||||
entries = _split_entries(text)
|
||||
prev_entry_count = state.get(state_key, {}).get("entry_count", 0)
|
||||
new_entries = entries if force else entries[prev_entry_count:]
|
||||
|
||||
if new_entries:
|
||||
# Build closet lines from new entries
|
||||
all_lines = []
|
||||
for header, body in new_entries:
|
||||
entry_text = f"{header}\n{body}"
|
||||
entry_lines = build_closet_lines(
|
||||
str(diary_path), [drawer_id], entry_text, wing, "daily"
|
||||
)
|
||||
all_lines.extend(entry_lines)
|
||||
if new_entries:
|
||||
all_lines = []
|
||||
for header, body in new_entries:
|
||||
entry_text = f"{header}\n{body}"
|
||||
entry_lines = build_closet_lines(
|
||||
source_file, [drawer_id], entry_text, wing, "daily"
|
||||
)
|
||||
all_lines.extend(entry_lines)
|
||||
|
||||
if all_lines:
|
||||
closet_id_base = f"closet_diary_{date_str}"
|
||||
closet_meta = {
|
||||
"date": date_str,
|
||||
"wing": wing,
|
||||
"room": "daily",
|
||||
"source_file": str(diary_path),
|
||||
"filed_at": now_iso,
|
||||
}
|
||||
if entities:
|
||||
closet_meta["entities"] = entities
|
||||
n = upsert_closet_lines(
|
||||
closets_col, closet_id_base, all_lines, closet_meta
|
||||
)
|
||||
closets_created += n
|
||||
if all_lines:
|
||||
closet_id_base = _diary_closet_id_base(wing, date_str)
|
||||
closet_meta = {
|
||||
"date": date_str,
|
||||
"wing": wing,
|
||||
"room": "daily",
|
||||
"source_file": source_file,
|
||||
"filed_at": now_iso,
|
||||
}
|
||||
if entities:
|
||||
closet_meta["entities"] = entities
|
||||
# On a force rebuild, wipe any leftover numbered closets
|
||||
# from a longer prior run before re-writing.
|
||||
if force:
|
||||
purge_file_closets(closets_col, source_file)
|
||||
n = upsert_closet_lines(closets_col, closet_id_base, all_lines, closet_meta)
|
||||
closets_created += n
|
||||
|
||||
state[diary_path.name] = {
|
||||
"size": curr_size,
|
||||
"entry_count": len(entries),
|
||||
"ingested_at": now_iso,
|
||||
}
|
||||
state[state_key] = {
|
||||
"size": curr_size,
|
||||
"entry_count": len(entries),
|
||||
"ingested_at": now_iso,
|
||||
}
|
||||
days_updated += 1
|
||||
|
||||
state_file.write_text(json.dumps(state, indent=2))
|
||||
|
||||
@@ -893,7 +893,10 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general"):
|
||||
return _no_palace()
|
||||
|
||||
now = datetime.now()
|
||||
entry_id = f"diary_{wing}_{now.strftime('%Y%m%d_%H%M%S')}_{hashlib.sha256(entry[:50].encode()).hexdigest()[:12]}"
|
||||
entry_id = (
|
||||
f"diary_{wing}_{now.strftime('%Y%m%d_%H%M%S%f')}_"
|
||||
f"{hashlib.sha256(entry.encode()).hexdigest()[:12]}"
|
||||
)
|
||||
|
||||
_wal_log(
|
||||
"diary_write",
|
||||
|
||||
+85
-27
@@ -16,8 +16,15 @@ from datetime import datetime
|
||||
from collections import defaultdict
|
||||
|
||||
from .palace import (
|
||||
SKIP_DIRS, get_collection, get_closets_collection,
|
||||
file_already_mined, mine_lock, build_closet_lines, upsert_closet_lines,
|
||||
NORMALIZE_VERSION,
|
||||
SKIP_DIRS,
|
||||
build_closet_lines,
|
||||
file_already_mined,
|
||||
get_closets_collection,
|
||||
get_collection,
|
||||
mine_lock,
|
||||
purge_file_closets,
|
||||
upsert_closet_lines,
|
||||
)
|
||||
|
||||
READABLE_EXTENSIONS = {
|
||||
@@ -371,41 +378,86 @@ def chunk_text(content: str, source_file: str) -> list:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
_ENTITY_REGISTRY_PATH = os.path.join(os.path.expanduser("~"), ".mempalace", "known_entities.json")
|
||||
_ENTITY_REGISTRY_CACHE: dict = {"mtime": None, "names": frozenset()}
|
||||
_ENTITY_EXTRACT_WINDOW = 5000 # chars of content scanned for capitalized words
|
||||
_ENTITY_METADATA_LIMIT = 25 # max entities packed into the metadata field
|
||||
|
||||
|
||||
def _load_known_entities() -> frozenset:
|
||||
"""Load (and cache) the user's known-entity registry by mtime.
|
||||
|
||||
Reads ``~/.mempalace/known_entities.json``. The registry is shaped as
|
||||
``{"category": ["Name1", "Name2", ...], ...}``. Cached across calls
|
||||
in the same process; invalidated when the file's mtime changes.
|
||||
"""
|
||||
try:
|
||||
mtime = os.path.getmtime(_ENTITY_REGISTRY_PATH)
|
||||
except OSError:
|
||||
if _ENTITY_REGISTRY_CACHE["mtime"] is not None:
|
||||
_ENTITY_REGISTRY_CACHE["mtime"] = None
|
||||
_ENTITY_REGISTRY_CACHE["names"] = frozenset()
|
||||
return _ENTITY_REGISTRY_CACHE["names"]
|
||||
|
||||
if _ENTITY_REGISTRY_CACHE["mtime"] == mtime:
|
||||
return _ENTITY_REGISTRY_CACHE["names"]
|
||||
|
||||
names: set = set()
|
||||
try:
|
||||
import json
|
||||
|
||||
with open(_ENTITY_REGISTRY_PATH, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
for cat in data.values():
|
||||
if isinstance(cat, list):
|
||||
names.update(str(n) for n in cat if n)
|
||||
except Exception:
|
||||
names = set()
|
||||
|
||||
_ENTITY_REGISTRY_CACHE["mtime"] = mtime
|
||||
_ENTITY_REGISTRY_CACHE["names"] = frozenset(names)
|
||||
return _ENTITY_REGISTRY_CACHE["names"]
|
||||
|
||||
|
||||
def _extract_entities_for_metadata(content: str) -> str:
|
||||
"""Extract entity names from content for metadata tagging.
|
||||
|
||||
Returns semicolon-separated string of entity names found in the text,
|
||||
suitable for ChromaDB metadata filtering.
|
||||
Combines the user's known-entity registry (cached across calls) with
|
||||
capitalized words appearing ≥2 times in the first ``_ENTITY_EXTRACT_WINDOW``
|
||||
chars. Filters out the closet stoplist (``When``, ``After``, ``The``, …)
|
||||
so sentence-starters don't masquerade as proper nouns.
|
||||
|
||||
Returns semicolon-separated string suitable for ChromaDB metadata
|
||||
filtering. The list is truncated to ``_ENTITY_METADATA_LIMIT`` entries
|
||||
*before* joining so a name is never cut in half.
|
||||
"""
|
||||
import re
|
||||
# Load known entities from registry if available
|
||||
known_names = set()
|
||||
registry_path = os.path.join(os.path.expanduser("~"), ".mempalace", "known_entities.json")
|
||||
if os.path.exists(registry_path):
|
||||
try:
|
||||
import json
|
||||
kd = json.loads(open(registry_path).read())
|
||||
for cat in kd.values():
|
||||
if isinstance(cat, list):
|
||||
known_names.update(cat)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
matched = set()
|
||||
# Match known entities
|
||||
for name in known_names:
|
||||
if re.search(r'(?<!\w)' + re.escape(name) + r'(?!\w)', content):
|
||||
from .palace import _ENTITY_STOPLIST
|
||||
|
||||
matched: set = set()
|
||||
|
||||
known = _load_known_entities()
|
||||
for name in known:
|
||||
if re.search(r"(?<!\w)" + re.escape(name) + r"(?!\w)", content):
|
||||
matched.add(name)
|
||||
# Also catch capitalized words appearing 2+ times (likely proper nouns)
|
||||
words = re.findall(r"\b[A-Z][a-z]{2,}\b", content[:5000])
|
||||
freq = {}
|
||||
|
||||
window = content[:_ENTITY_EXTRACT_WINDOW]
|
||||
words = re.findall(r"\b[A-Z][a-z]{2,}\b", window)
|
||||
freq: dict = {}
|
||||
for w in words:
|
||||
if w in _ENTITY_STOPLIST:
|
||||
continue
|
||||
freq[w] = freq.get(w, 0) + 1
|
||||
for w, c in freq.items():
|
||||
if c >= 2 and len(w) > 2:
|
||||
matched.add(w)
|
||||
|
||||
return ";".join(sorted(matched))[:500] if matched else ""
|
||||
if not matched:
|
||||
return ""
|
||||
# Truncate the *list*, not the joined string — never split a name.
|
||||
capped = sorted(matched)[:_ENTITY_METADATA_LIMIT]
|
||||
return ";".join(capped)
|
||||
|
||||
|
||||
def add_drawer(
|
||||
@@ -421,6 +473,7 @@ def add_drawer(
|
||||
"chunk_index": chunk_index,
|
||||
"added_by": agent,
|
||||
"filed_at": datetime.now().isoformat(),
|
||||
"normalize_version": NORMALIZE_VERSION,
|
||||
}
|
||||
# Store file mtime so we can detect modifications later.
|
||||
try:
|
||||
@@ -511,15 +564,18 @@ def process_file(
|
||||
if added:
|
||||
drawers_added += 1
|
||||
|
||||
# Build closet — the searchable index pointing to these drawers
|
||||
# Each topic line is atomic — never split across closets
|
||||
# Build closet — the searchable index pointing to these drawers.
|
||||
# Purge first: a re-mine (mtime change or normalize_version bump) must
|
||||
# fully replace the prior closets, not append to them.
|
||||
if closets_col and drawers_added > 0:
|
||||
drawer_ids = [
|
||||
f"drawer_{wing}_{room}_{hashlib.sha256((source_file + str(c['chunk_index'])).encode()).hexdigest()[:24]}"
|
||||
for c in chunks
|
||||
]
|
||||
closet_lines = build_closet_lines(source_file, drawer_ids, content, wing, room)
|
||||
closet_id_base = f"closet_{wing}_{room}_{hashlib.sha256(source_file.encode()).hexdigest()[:24]}"
|
||||
closet_id_base = (
|
||||
f"closet_{wing}_{room}_{hashlib.sha256(source_file.encode()).hexdigest()[:24]}"
|
||||
)
|
||||
entities = _extract_entities_for_metadata(content)
|
||||
closet_meta = {
|
||||
"wing": wing,
|
||||
@@ -527,9 +583,11 @@ def process_file(
|
||||
"source_file": source_file,
|
||||
"drawer_count": drawers_added,
|
||||
"filed_at": datetime.now().isoformat(),
|
||||
"normalize_version": NORMALIZE_VERSION,
|
||||
}
|
||||
if entities:
|
||||
closet_meta["entities"] = entities
|
||||
purge_file_closets(closets_col, source_file)
|
||||
upsert_closet_lines(closets_col, closet_id_base, closet_lines, closet_meta)
|
||||
|
||||
return drawers_added, room
|
||||
|
||||
+93
-2
@@ -16,10 +16,93 @@ No API key. No internet. Everything local.
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# ─── Noise stripping ─────────────────────────────────────────────────────
|
||||
# Claude Code and other tools inject system tags, hook output, and UI chrome
|
||||
# into transcripts. These waste drawer space and pollute search results.
|
||||
#
|
||||
# Verbatim is sacred — every pattern here is anchored to line boundaries and
|
||||
# refuses to cross blank lines, so a stray unclosed tag in one message can
|
||||
# never eat content from neighboring messages. When in doubt, leave text
|
||||
# alone.
|
||||
|
||||
_NOISE_TAGS = (
|
||||
"system-reminder",
|
||||
"command-message",
|
||||
"command-name",
|
||||
"task-notification",
|
||||
"user-prompt-submit-hook",
|
||||
"hook_output",
|
||||
)
|
||||
|
||||
|
||||
def _tag_pattern(name: str) -> "re.Pattern[str]":
|
||||
# Opening tag must begin a line (optionally after a `> ` blockquote marker,
|
||||
# since _messages_to_transcript prefixes lines with `> `). Body is lazy but
|
||||
# forbidden from crossing a blank line, so a dangling open tag can't span
|
||||
# multiple messages. Closing tag eats optional trailing whitespace + newline.
|
||||
return re.compile(
|
||||
rf"(?m)^(?:> )?<{name}(?:\s[^>]*)?>" rf"(?:(?!\n\s*\n)[\s\S])*?" rf"</{name}>[ \t]*\n?"
|
||||
)
|
||||
|
||||
|
||||
_NOISE_TAG_PATTERNS = [_tag_pattern(t) for t in _NOISE_TAGS]
|
||||
|
||||
# Strings that identify an entire noise line when found at its start.
|
||||
# Matched case-sensitively and anchored to line-start so user prose mentioning
|
||||
# e.g. "current time:" in a sentence is untouched.
|
||||
_NOISE_LINE_PREFIXES = (
|
||||
"CURRENT TIME:",
|
||||
"VERIFIED FACTS (do not contradict)",
|
||||
"AGENT SPECIALIZATION:",
|
||||
"Checking verified facts...",
|
||||
"Injecting timestamp...",
|
||||
"Starting background pipeline...",
|
||||
"Checking emotional weights...",
|
||||
"Auto-save reminder...",
|
||||
"Checking pipeline...",
|
||||
"MemPalace auto-save checkpoint.",
|
||||
)
|
||||
|
||||
_NOISE_LINE_PATTERNS = [
|
||||
re.compile(rf"(?m)^(?:> )?{re.escape(p)}.*\n?") for p in _NOISE_LINE_PREFIXES
|
||||
]
|
||||
|
||||
# Claude Code TUI hook-run chrome, e.g. "Ran 2 Stop hook", "Ran 1 PreCompact hook".
|
||||
# Line-anchored, case-sensitive, explicit hook names — prose like
|
||||
# "our CI has a stop hook" stays intact.
|
||||
_HOOK_LINE_RE = re.compile(
|
||||
r"(?m)^(?:> )?Ran \d+ (?:Stop|PreCompact|PreToolUse|PostToolUse|UserPromptSubmit|Notification|SessionStart|SessionEnd) hook[s]?.*\n?"
|
||||
)
|
||||
|
||||
# "… +N lines" collapsed-output marker, line-anchored.
|
||||
_COLLAPSED_LINES_RE = re.compile(r"(?m)^(?:> )?…\s*\+\d+ lines.*\n?")
|
||||
|
||||
|
||||
def strip_noise(text: str) -> str:
|
||||
"""Remove system tags, hook output, and Claude Code UI chrome from text.
|
||||
|
||||
All patterns are line-anchored. User prose that happens to mention these
|
||||
strings inline (e.g., documenting them) is preserved verbatim.
|
||||
"""
|
||||
for pat in _NOISE_TAG_PATTERNS:
|
||||
text = pat.sub("", text)
|
||||
for pat in _NOISE_LINE_PATTERNS:
|
||||
text = pat.sub("", text)
|
||||
text = _HOOK_LINE_RE.sub("", text)
|
||||
text = _COLLAPSED_LINES_RE.sub("", text)
|
||||
# Strip the Claude Code collapsed-output chrome "[N tokens] (ctrl+o to expand)".
|
||||
# Narrow shape — a bare "(ctrl+o to expand)" in user prose stays intact.
|
||||
text = re.sub(r"\s*\[\d+\s+tokens?\]\s*\(ctrl\+o to expand\)", "", text)
|
||||
# Collapse runs of blank lines created by the removals
|
||||
text = re.sub(r"\n{4,}", "\n\n\n", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def normalize(filepath: str) -> str:
|
||||
"""
|
||||
Load a file and normalize to transcript format if it's a chat export.
|
||||
@@ -40,12 +123,14 @@ def normalize(filepath: str) -> str:
|
||||
if not content.strip():
|
||||
return content
|
||||
|
||||
# Already has > markers — pass through
|
||||
# Already has > markers — pass through unchanged.
|
||||
lines = content.split("\n")
|
||||
if sum(1 for line in lines if line.strip().startswith(">")) >= 3:
|
||||
return content
|
||||
|
||||
# Try JSON normalization
|
||||
# Try JSON normalization. strip_noise is applied inside the Claude Code
|
||||
# JSONL parser (the only format that injects system tags/hook chrome);
|
||||
# other formats pass through verbatim.
|
||||
ext = Path(filepath).suffix.lower()
|
||||
if ext in (".json", ".jsonl") or content.strip()[:1] in ("{", "["):
|
||||
normalized = _try_normalize_json(content)
|
||||
@@ -112,6 +197,10 @@ def _try_claude_code_jsonl(content: str) -> Optional[str]:
|
||||
isinstance(b, dict) and b.get("type") == "tool_result" for b in msg_content
|
||||
)
|
||||
text = _extract_content(msg_content, tool_use_map=tool_use_map)
|
||||
# Strip Claude Code system-injected noise per message, never across
|
||||
# message boundaries — prevents span-eating.
|
||||
if text:
|
||||
text = strip_noise(text)
|
||||
if text:
|
||||
if is_tool_only and messages and messages[-1][0] == "assistant":
|
||||
# Append tool results to the previous assistant message
|
||||
@@ -121,6 +210,8 @@ def _try_claude_code_jsonl(content: str) -> Optional[str]:
|
||||
messages.append(("user", text))
|
||||
elif msg_type == "assistant":
|
||||
text = _extract_content(msg_content, tool_use_map=tool_use_map)
|
||||
if text:
|
||||
text = strip_noise(text)
|
||||
if text:
|
||||
# If previous message is also assistant (multi-turn tool loop),
|
||||
# merge into the same assistant turn
|
||||
|
||||
+117
-27
@@ -38,6 +38,16 @@ SKIP_DIRS = {
|
||||
|
||||
_DEFAULT_BACKEND = ChromaBackend()
|
||||
|
||||
# Schema version for drawer normalization. Bump when the normalization
|
||||
# pipeline changes in a way that existing drawers should be rebuilt to pick up
|
||||
# (e.g., new noise-stripping rules). `file_already_mined` treats drawers with
|
||||
# a missing or stale `normalize_version` as "not mined", so the next mine pass
|
||||
# silently rebuilds them — users don't need to manually erase + re-mine.
|
||||
#
|
||||
# v2 (2026-04): introduced strip_noise() for Claude Code JSONL; previous
|
||||
# drawers stored system tags / hook chrome verbatim.
|
||||
NORMALIZE_VERSION = 2
|
||||
|
||||
|
||||
def get_collection(
|
||||
palace_path: str,
|
||||
@@ -58,6 +68,66 @@ def get_closets_collection(palace_path: str, create: bool = True):
|
||||
|
||||
|
||||
CLOSET_CHAR_LIMIT = 1500 # fill closet until ~1500 chars, then start a new one
|
||||
CLOSET_EXTRACT_WINDOW = 5000 # how many chars of source content to scan for entities/topics
|
||||
|
||||
# Common capitalized words that look like proper nouns but are usually
|
||||
# sentence-starters or filler. Filtered out of entity extraction.
|
||||
_ENTITY_STOPLIST = frozenset(
|
||||
{
|
||||
"The",
|
||||
"This",
|
||||
"That",
|
||||
"These",
|
||||
"Those",
|
||||
"When",
|
||||
"Where",
|
||||
"What",
|
||||
"Why",
|
||||
"Who",
|
||||
"Which",
|
||||
"How",
|
||||
"After",
|
||||
"Before",
|
||||
"Then",
|
||||
"Now",
|
||||
"Here",
|
||||
"There",
|
||||
"And",
|
||||
"But",
|
||||
"Or",
|
||||
"Yet",
|
||||
"So",
|
||||
"If",
|
||||
"Else",
|
||||
"Yes",
|
||||
"No",
|
||||
"Maybe",
|
||||
"Okay",
|
||||
"User",
|
||||
"Assistant",
|
||||
"System",
|
||||
"Tool",
|
||||
"Monday",
|
||||
"Tuesday",
|
||||
"Wednesday",
|
||||
"Thursday",
|
||||
"Friday",
|
||||
"Saturday",
|
||||
"Sunday",
|
||||
"January",
|
||||
"February",
|
||||
"March",
|
||||
"April",
|
||||
"May",
|
||||
"June",
|
||||
"July",
|
||||
"August",
|
||||
"September",
|
||||
"October",
|
||||
"November",
|
||||
"December",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def build_closet_lines(source_file, drawer_ids, content, wing, room):
|
||||
@@ -72,11 +142,15 @@ def build_closet_lines(source_file, drawer_ids, content, wing, room):
|
||||
from pathlib import Path
|
||||
|
||||
drawer_ref = ",".join(drawer_ids[:3])
|
||||
window = content[:CLOSET_EXTRACT_WINDOW]
|
||||
|
||||
# Extract proper nouns (capitalized words, 2+ occurrences)
|
||||
words = re.findall(r"\b[A-Z][a-z]{2,}\b", content[:5000])
|
||||
# Extract proper nouns (capitalized words, 2+ occurrences). Filter out
|
||||
# common sentence-starters that aren't real entities.
|
||||
words = re.findall(r"\b[A-Z][a-z]{2,}\b", window)
|
||||
word_freq = {}
|
||||
for w in words:
|
||||
if w in _ENTITY_STOPLIST:
|
||||
continue
|
||||
word_freq[w] = word_freq.get(w, 0) + 1
|
||||
entities = sorted(
|
||||
[w for w, c in word_freq.items() if c >= 2],
|
||||
@@ -89,15 +163,15 @@ def build_closet_lines(source_file, drawer_ids, content, wing, room):
|
||||
for pattern in [
|
||||
r"(?:built|fixed|wrote|added|pushed|tested|created|decided|migrated|reviewed|deployed|configured|removed|updated)\s+[\w\s]{3,40}",
|
||||
]:
|
||||
topics.extend(re.findall(pattern, content[:5000], re.IGNORECASE))
|
||||
topics.extend(re.findall(pattern, window, re.IGNORECASE))
|
||||
# Also grab section headers if present
|
||||
for header in re.findall(r"^#{1,3}\s+(.{5,60})$", content[:5000], re.MULTILINE):
|
||||
for header in re.findall(r"^#{1,3}\s+(.{5,60})$", window, re.MULTILINE):
|
||||
topics.append(header.strip())
|
||||
# Dedupe preserving order
|
||||
topics = list(dict.fromkeys(t.strip().lower() for t in topics))[:12]
|
||||
|
||||
# Extract quotes
|
||||
quotes = re.findall(r'"([^"]{15,150})"', content[:5000])
|
||||
quotes = re.findall(r'"([^"]{15,150})"', window)
|
||||
|
||||
# Build pointer lines — each one is atomic, never split
|
||||
lines = []
|
||||
@@ -114,17 +188,31 @@ def build_closet_lines(source_file, drawer_ids, content, wing, room):
|
||||
return lines
|
||||
|
||||
|
||||
def upsert_closet_lines(closets_col, closet_id_base, lines, metadata):
|
||||
"""Add topic lines to closets. Never splits a topic mid-line.
|
||||
def purge_file_closets(closets_col, source_file: str) -> None:
|
||||
"""Delete every closet associated with ``source_file``.
|
||||
|
||||
If adding a line WHOLE would exceed CLOSET_CHAR_LIMIT, a new closet
|
||||
is created. Some closets may have less than 1500 chars — that's fine.
|
||||
Every topic is complete and readable.
|
||||
Call this before ``upsert_closet_lines`` on a re-mine so stale topics
|
||||
from a prior schema/version don't survive in the closet collection.
|
||||
Mirrors the drawer-purge step in process_file().
|
||||
"""
|
||||
try:
|
||||
closets_col.delete(where={"source_file": source_file})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def upsert_closet_lines(closets_col, closet_id_base, lines, metadata):
|
||||
"""Write topic lines to closets, packed greedily without splitting a line.
|
||||
|
||||
Closets are deterministically numbered (``..._01``, ``..._02``, …) and
|
||||
each ``upsert`` fully overwrites the prior content at that ID. Callers
|
||||
are expected to ``purge_file_closets`` first when re-mining a source
|
||||
file so stale-numbered closets from larger prior runs don't leak.
|
||||
|
||||
Returns the number of closets written.
|
||||
"""
|
||||
closet_num = 1
|
||||
current_lines = []
|
||||
current_lines: list = []
|
||||
current_chars = 0
|
||||
closets_written = 0
|
||||
|
||||
@@ -134,17 +222,6 @@ def upsert_closet_lines(closets_col, closet_id_base, lines, metadata):
|
||||
return
|
||||
closet_id = f"{closet_id_base}_{closet_num:02d}"
|
||||
text = "\n".join(current_lines)
|
||||
|
||||
# Check if closet already has content — append if room
|
||||
try:
|
||||
existing = closets_col.get(ids=[closet_id])
|
||||
if existing.get("ids") and existing["documents"][0]:
|
||||
old = existing["documents"][0]
|
||||
if len(old) + len(text) + 1 <= CLOSET_CHAR_LIMIT:
|
||||
text = old + "\n" + text
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
closets_col.upsert(documents=[text], ids=[closet_id], metadatas=[metadata])
|
||||
closets_written += 1
|
||||
|
||||
@@ -152,7 +229,6 @@ def upsert_closet_lines(closets_col, closet_id_base, lines, metadata):
|
||||
line_len = len(line)
|
||||
# Would this line fit whole in the current closet?
|
||||
if current_chars > 0 and current_chars + line_len + 1 > CLOSET_CHAR_LIMIT:
|
||||
# Doesn't fit — flush current closet, start new one
|
||||
_flush()
|
||||
closet_num += 1
|
||||
current_lines = []
|
||||
@@ -182,18 +258,22 @@ def mine_lock(source_file: str):
|
||||
try:
|
||||
if os.name == "nt":
|
||||
import msvcrt
|
||||
|
||||
msvcrt.locking(lf.fileno(), msvcrt.LK_LOCK, 1)
|
||||
else:
|
||||
import fcntl
|
||||
|
||||
fcntl.flock(lf, fcntl.LOCK_EX)
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
if os.name == "nt":
|
||||
import msvcrt
|
||||
|
||||
msvcrt.locking(lf.fileno(), msvcrt.LK_UNLCK, 1)
|
||||
else:
|
||||
import fcntl
|
||||
|
||||
fcntl.flock(lf, fcntl.LOCK_UN)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -203,16 +283,26 @@ def mine_lock(source_file: str):
|
||||
def file_already_mined(collection, source_file: str, check_mtime: bool = False) -> bool:
|
||||
"""Check if a file has already been filed in the palace.
|
||||
|
||||
When check_mtime=True (used by project miner), returns False if the file
|
||||
has been modified since it was last mined, so it gets re-mined.
|
||||
When check_mtime=False (used by convo miner), just checks existence.
|
||||
Returns False (so the file gets re-mined) when:
|
||||
- no drawers exist for this source_file
|
||||
- the stored `normalize_version` is missing or older than the current
|
||||
schema (triggers silent rebuild after a normalization upgrade)
|
||||
- `check_mtime=True` and the file's mtime differs from the stored one
|
||||
|
||||
When check_mtime=True (used by project miner), also re-mines on content
|
||||
change. When check_mtime=False (used by convo miner), transcripts are
|
||||
assumed immutable, so only the version gate triggers a rebuild.
|
||||
"""
|
||||
try:
|
||||
results = collection.get(where={"source_file": source_file}, limit=1)
|
||||
if not results.get("ids"):
|
||||
return False
|
||||
stored_meta = results.get("metadatas", [{}])[0] or {}
|
||||
# Pre-v2 drawers have no version field — treat them as stale.
|
||||
stored_version = stored_meta.get("normalize_version", 1)
|
||||
if stored_version < NORMALIZE_VERSION:
|
||||
return False
|
||||
if check_mtime:
|
||||
stored_meta = results.get("metadatas", [{}])[0]
|
||||
stored_mtime = stored_meta.get("source_mtime")
|
||||
if stored_mtime is None:
|
||||
return False
|
||||
|
||||
+129
-62
@@ -18,11 +18,12 @@ No external graph DB needed — built from ChromaDB metadata.
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from collections import defaultdict, Counter
|
||||
from datetime import datetime
|
||||
from collections import Counter, defaultdict
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .config import MempalaceConfig
|
||||
from .palace import get_collection as _get_palace_collection
|
||||
from .palace import mine_lock
|
||||
|
||||
|
||||
def _get_collection(config=None):
|
||||
@@ -249,20 +250,66 @@ _TUNNEL_FILE = os.path.join(os.path.expanduser("~"), ".mempalace", "tunnels.json
|
||||
|
||||
|
||||
def _load_tunnels():
|
||||
"""Load explicit tunnels from disk."""
|
||||
if os.path.exists(_TUNNEL_FILE):
|
||||
try:
|
||||
return json.loads(open(_TUNNEL_FILE).read())
|
||||
except Exception:
|
||||
pass
|
||||
return []
|
||||
"""Load explicit tunnels from disk.
|
||||
|
||||
Returns an empty list if the file is missing or corrupt (e.g. truncated
|
||||
by a crash mid-write on a system that lacks atomic-rename semantics).
|
||||
"""
|
||||
if not os.path.exists(_TUNNEL_FILE):
|
||||
return []
|
||||
try:
|
||||
with open(_TUNNEL_FILE, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except Exception:
|
||||
return []
|
||||
return data if isinstance(data, list) else []
|
||||
|
||||
|
||||
def _save_tunnels(tunnels):
|
||||
"""Save explicit tunnels to disk."""
|
||||
"""Persist explicit tunnels atomically.
|
||||
|
||||
Writes to ``tunnels.json.tmp`` then ``os.replace``s it into place, so
|
||||
a crash mid-write can never leave a partial/empty tunnels.json that
|
||||
silently wipes every tunnel on next read.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(_TUNNEL_FILE), exist_ok=True)
|
||||
with open(_TUNNEL_FILE, "w") as f:
|
||||
tmp_path = _TUNNEL_FILE + ".tmp"
|
||||
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||
json.dump(tunnels, f, indent=2)
|
||||
f.flush()
|
||||
try:
|
||||
os.fsync(f.fileno())
|
||||
except OSError:
|
||||
# Not all filesystems (or Windows file handles) support fsync — tolerate.
|
||||
pass
|
||||
os.replace(tmp_path, _TUNNEL_FILE)
|
||||
|
||||
|
||||
def _endpoint_key(wing: str, room: str) -> str:
|
||||
return f"{wing}/{room}"
|
||||
|
||||
|
||||
def _canonical_tunnel_id(
|
||||
source_wing: str, source_room: str, target_wing: str, target_room: str
|
||||
) -> str:
|
||||
"""Compute a symmetric tunnel ID.
|
||||
|
||||
Tunnels are conceptually undirected — "auth relates to users" is the
|
||||
same connection as "users relates to auth". Sort the two endpoints
|
||||
before hashing so ``create_tunnel(A, B)`` and ``create_tunnel(B, A)``
|
||||
resolve to the same ID and dedup into one record.
|
||||
"""
|
||||
src = _endpoint_key(source_wing, source_room)
|
||||
tgt = _endpoint_key(target_wing, target_room)
|
||||
a, b = sorted((src, tgt))
|
||||
return hashlib.sha256(f"{a}↔{b}".encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
def _require_name(value: str, field: str) -> str:
|
||||
"""Reject empty / non-string endpoint identifiers."""
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"{field} must be a non-empty string")
|
||||
return value.strip()
|
||||
|
||||
|
||||
def create_tunnel(
|
||||
@@ -274,72 +321,88 @@ def create_tunnel(
|
||||
source_drawer_id: str = None,
|
||||
target_drawer_id: str = None,
|
||||
):
|
||||
"""Create an explicit tunnel between two locations in the palace.
|
||||
"""Create an explicit (symmetric) tunnel between two locations in the palace.
|
||||
|
||||
Use when an agent notices a connection between two projects/wings
|
||||
that wouldn't be found by passive room-name matching.
|
||||
Tunnels are undirected: ``create_tunnel(A, B)`` and ``create_tunnel(B, A)``
|
||||
resolve to the same canonical ID. A second call with the same endpoints
|
||||
updates the stored label (and drawer IDs, if provided) rather than
|
||||
creating a duplicate.
|
||||
|
||||
The ``source`` / ``target`` fields on the returned dict preserve the
|
||||
argument order the caller used, so callers can display it directionally
|
||||
if they like. The ID and dedup are symmetric.
|
||||
|
||||
Args:
|
||||
source_wing: Wing of the source (e.g., "project_api")
|
||||
source_room: Room in the source wing
|
||||
target_wing: Wing of the target (e.g., "project_database")
|
||||
target_room: Room in the target wing
|
||||
label: Description of the connection
|
||||
source_drawer_id: Optional specific drawer ID
|
||||
target_drawer_id: Optional specific drawer ID
|
||||
source_wing: Wing of the source (e.g., "project_api").
|
||||
source_room: Room in the source wing.
|
||||
target_wing: Wing of the target (e.g., "project_database").
|
||||
target_room: Room in the target wing.
|
||||
label: Description of the connection.
|
||||
source_drawer_id: Optional specific drawer ID.
|
||||
target_drawer_id: Optional specific drawer ID.
|
||||
|
||||
Returns:
|
||||
The created tunnel dict.
|
||||
The stored tunnel dict.
|
||||
|
||||
Raises:
|
||||
ValueError: if any wing or room is empty or non-string.
|
||||
"""
|
||||
tunnel_id = hashlib.sha256(
|
||||
f"{source_wing}/{source_room}↔{target_wing}/{target_room}".encode()
|
||||
).hexdigest()[:16]
|
||||
source_wing = _require_name(source_wing, "source_wing")
|
||||
source_room = _require_name(source_room, "source_room")
|
||||
target_wing = _require_name(target_wing, "target_wing")
|
||||
target_room = _require_name(target_room, "target_room")
|
||||
|
||||
tunnel_id = _canonical_tunnel_id(source_wing, source_room, target_wing, target_room)
|
||||
|
||||
tunnel = {
|
||||
"id": tunnel_id,
|
||||
"source": {"wing": source_wing, "room": source_room},
|
||||
"target": {"wing": target_wing, "room": target_room},
|
||||
"label": label,
|
||||
"created_at": datetime.now().isoformat(),
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
if source_drawer_id:
|
||||
tunnel["source"]["drawer_id"] = source_drawer_id
|
||||
if target_drawer_id:
|
||||
tunnel["target"]["drawer_id"] = target_drawer_id
|
||||
|
||||
tunnels = _load_tunnels()
|
||||
|
||||
# Dedup — don't create if same endpoints already linked
|
||||
for existing in tunnels:
|
||||
if existing.get("id") == tunnel_id:
|
||||
existing.update(tunnel) # update label/drawers
|
||||
_save_tunnels(tunnels)
|
||||
return existing
|
||||
|
||||
tunnels.append(tunnel)
|
||||
_save_tunnels(tunnels)
|
||||
# Serialize the load → mutate → save cycle. Without this, two concurrent
|
||||
# create_tunnel calls can both read the same snapshot and the later
|
||||
# writer silently drops the earlier writer's tunnel.
|
||||
with mine_lock(_TUNNEL_FILE):
|
||||
tunnels = _load_tunnels()
|
||||
for existing in tunnels:
|
||||
if existing.get("id") == tunnel_id:
|
||||
# Preserve original creation timestamp on label updates.
|
||||
tunnel["created_at"] = existing.get("created_at", tunnel["created_at"])
|
||||
tunnel["updated_at"] = datetime.now(timezone.utc).isoformat()
|
||||
existing.clear()
|
||||
existing.update(tunnel)
|
||||
_save_tunnels(tunnels)
|
||||
return existing
|
||||
tunnels.append(tunnel)
|
||||
_save_tunnels(tunnels)
|
||||
return tunnel
|
||||
|
||||
|
||||
def list_tunnels(wing: str = None):
|
||||
"""List all explicit tunnels, optionally filtered by wing.
|
||||
|
||||
Returns tunnels where the wing appears as either source or target.
|
||||
Returns tunnels where ``wing`` appears as either source or target
|
||||
(tunnels are symmetric, so either endpoint is a valid filter match).
|
||||
"""
|
||||
tunnels = _load_tunnels()
|
||||
if wing:
|
||||
tunnels = [
|
||||
t for t in tunnels
|
||||
if t["source"]["wing"] == wing or t["target"]["wing"] == wing
|
||||
]
|
||||
tunnels = [t for t in tunnels if t["source"]["wing"] == wing or t["target"]["wing"] == wing]
|
||||
return tunnels
|
||||
|
||||
|
||||
def delete_tunnel(tunnel_id: str):
|
||||
"""Delete an explicit tunnel by ID."""
|
||||
tunnels = _load_tunnels()
|
||||
tunnels = [t for t in tunnels if t.get("id") != tunnel_id]
|
||||
_save_tunnels(tunnels)
|
||||
"""Delete an explicit tunnel by ID. Returns ``{"deleted": <id>}``."""
|
||||
with mine_lock(_TUNNEL_FILE):
|
||||
tunnels = _load_tunnels()
|
||||
tunnels = [t for t in tunnels if t.get("id") != tunnel_id]
|
||||
_save_tunnels(tunnels)
|
||||
return {"deleted": tunnel_id}
|
||||
|
||||
|
||||
@@ -357,23 +420,27 @@ def follow_tunnels(wing: str, room: str, col=None, config=None):
|
||||
tgt = t["target"]
|
||||
|
||||
if src["wing"] == wing and src["room"] == room:
|
||||
connections.append({
|
||||
"direction": "outgoing",
|
||||
"connected_wing": tgt["wing"],
|
||||
"connected_room": tgt["room"],
|
||||
"label": t.get("label", ""),
|
||||
"drawer_id": tgt.get("drawer_id"),
|
||||
"tunnel_id": t["id"],
|
||||
})
|
||||
connections.append(
|
||||
{
|
||||
"direction": "outgoing",
|
||||
"connected_wing": tgt["wing"],
|
||||
"connected_room": tgt["room"],
|
||||
"label": t.get("label", ""),
|
||||
"drawer_id": tgt.get("drawer_id"),
|
||||
"tunnel_id": t["id"],
|
||||
}
|
||||
)
|
||||
elif tgt["wing"] == wing and tgt["room"] == room:
|
||||
connections.append({
|
||||
"direction": "incoming",
|
||||
"connected_wing": src["wing"],
|
||||
"connected_room": src["room"],
|
||||
"label": t.get("label", ""),
|
||||
"drawer_id": src.get("drawer_id"),
|
||||
"tunnel_id": t["id"],
|
||||
})
|
||||
connections.append(
|
||||
{
|
||||
"direction": "incoming",
|
||||
"connected_wing": src["wing"],
|
||||
"connected_room": src["room"],
|
||||
"label": t.get("label", ""),
|
||||
"drawer_id": src.get("drawer_id"),
|
||||
"tunnel_id": t["id"],
|
||||
}
|
||||
)
|
||||
|
||||
# If we have a collection, fetch drawer content for connected items
|
||||
if col and connections:
|
||||
|
||||
+305
-134
@@ -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,187 @@ 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 _expand_with_neighbors(drawers_col, matched_doc: str, matched_meta: dict, radius: int = 1):
|
||||
"""Expand a matched drawer with its ±radius sibling chunks in the same source file.
|
||||
|
||||
Motivation — "drawer-grep context" feature: a closet hit returns one
|
||||
drawer, but the chunk boundary may clip mid-thought (e.g., the matched
|
||||
chunk says "here's a breakdown:" and the actual breakdown lives in the
|
||||
next chunk). Fetching the small neighborhood around the match gives
|
||||
callers enough context without forcing a follow-up ``get_drawer`` call.
|
||||
|
||||
Returns a dict with:
|
||||
``text`` combined chunks in chunk_index order
|
||||
``drawer_index`` the matched chunk's index in the source file
|
||||
``total_drawers`` total drawer count for the source file (or None)
|
||||
|
||||
On any ChromaDB failure or missing metadata, falls back to returning the
|
||||
matched drawer alone so search never breaks because neighbor expansion
|
||||
failed.
|
||||
"""
|
||||
src = matched_meta.get("source_file")
|
||||
chunk_idx = matched_meta.get("chunk_index")
|
||||
if not src or not isinstance(chunk_idx, int):
|
||||
return {"text": matched_doc, "drawer_index": chunk_idx, "total_drawers": None}
|
||||
|
||||
target_indexes = [chunk_idx + offset for offset in range(-radius, radius + 1)]
|
||||
try:
|
||||
neighbors = drawers_col.get(
|
||||
where={
|
||||
"$and": [
|
||||
{"source_file": src},
|
||||
{"chunk_index": {"$in": target_indexes}},
|
||||
]
|
||||
},
|
||||
include=["documents", "metadatas"],
|
||||
)
|
||||
except Exception:
|
||||
return {"text": matched_doc, "drawer_index": chunk_idx, "total_drawers": None}
|
||||
|
||||
indexed_docs = []
|
||||
for doc, meta in zip(neighbors.get("documents") or [], neighbors.get("metadatas") or []):
|
||||
ci = meta.get("chunk_index")
|
||||
if isinstance(ci, int):
|
||||
indexed_docs.append((ci, doc))
|
||||
indexed_docs.sort(key=lambda pair: pair[0])
|
||||
|
||||
if not indexed_docs:
|
||||
combined_text = matched_doc
|
||||
else:
|
||||
combined_text = "\n\n".join(doc for _, doc in indexed_docs)
|
||||
|
||||
# Cheap total_drawers lookup: metadata-only scan of the source file.
|
||||
total_drawers = None
|
||||
try:
|
||||
all_meta = drawers_col.get(where={"source_file": src}, include=["metadatas"])
|
||||
ids = all_meta.get("ids") or []
|
||||
total_drawers = len(ids) if ids else None
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {
|
||||
"text": combined_text,
|
||||
"drawer_index": chunk_idx,
|
||||
"total_drawers": total_drawers,
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
# Expand with ±1 neighbor chunks from the same source file so a
|
||||
# closet hit that lands mid-thought still returns enough context to
|
||||
# be useful without a follow-up get_drawer call.
|
||||
expansion = _expand_with_neighbors(drawers_col, doc, meta, radius=1)
|
||||
hits.append(
|
||||
{
|
||||
"text": expansion["text"],
|
||||
"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],
|
||||
"drawer_index": expansion["drawer_index"],
|
||||
"total_drawers": expansion["total_drawers"],
|
||||
}
|
||||
)
|
||||
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,98 +420,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
|
||||
MAX_HYDRATION_CHARS = 10000 # cap to prevent blowup on large source files
|
||||
|
||||
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, grep for most relevant chunk
|
||||
try:
|
||||
drawer_results = drawers_col.get(
|
||||
where={"source_file": source},
|
||||
include=["documents", "metadatas"],
|
||||
)
|
||||
if drawer_results.get("ids"):
|
||||
# Drawer-grep: score each chunk against the query,
|
||||
# return the best-matching chunk first + surrounding context
|
||||
query_terms = set(re.findall(r'\w{2,}', query.lower()))
|
||||
best_idx = 0
|
||||
best_score = -1
|
||||
for idx, doc in enumerate(drawer_results["documents"]):
|
||||
doc_lower = doc.lower()
|
||||
score = sum(1 for t in query_terms if t in doc_lower)
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_idx = idx
|
||||
|
||||
# Build result: best chunk first, then neighbors
|
||||
docs = drawer_results["documents"]
|
||||
n_docs = len(docs)
|
||||
# Include best chunk + 1 before + 1 after for context
|
||||
start = max(0, best_idx - 1)
|
||||
end = min(n_docs, best_idx + 2)
|
||||
relevant_text = "\n\n".join(docs[start:end])
|
||||
|
||||
if len(relevant_text) > MAX_HYDRATION_CHARS:
|
||||
relevant_text = relevant_text[:MAX_HYDRATION_CHARS] + f"\n\n[...truncated. {n_docs} total drawers. Use mempalace_get_drawer for full content.]"
|
||||
|
||||
meta = drawer_results["metadatas"][best_idx]
|
||||
hits.append({
|
||||
"text": relevant_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],
|
||||
"drawer_index": best_idx,
|
||||
"total_drawers": n_docs,
|
||||
})
|
||||
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:
|
||||
@@ -307,6 +477,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",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
"""Single source of truth for the MemPalace package version."""
|
||||
|
||||
__version__ = "3.1.0"
|
||||
__version__ = "3.2.0"
|
||||
|
||||
+770
-117
File diff suppressed because it is too large
Load Diff
@@ -75,3 +75,86 @@ def test_mine_convos_does_not_reprocess_empty_chunk_files(capsys):
|
||||
assert "Files skipped (already filed): 1" in out2
|
||||
finally:
|
||||
shutil.rmtree(tmpdir, ignore_errors=True)
|
||||
|
||||
|
||||
def test_mine_convos_rebuilds_stale_drawers_after_schema_bump(capsys):
|
||||
"""When stored drawers have an older normalize_version, the next mine
|
||||
silently purges them and refiles — no manual erase required.
|
||||
|
||||
This is what makes the strip_noise upgrade apply to existing corpora:
|
||||
users just run `mempalace mine` again and old noise-filled drawers get
|
||||
replaced with clean ones."""
|
||||
from mempalace.palace import NORMALIZE_VERSION
|
||||
|
||||
tmpdir = tempfile.mkdtemp()
|
||||
try:
|
||||
convo_path = Path(tmpdir) / "chat.txt"
|
||||
convo_path.write_text(
|
||||
"> What is memory?\nMemory is persistence.\n\n"
|
||||
"> Why does it matter?\nIt enables continuity.\n\n"
|
||||
"> How do we build it?\nWith structured storage.\n"
|
||||
)
|
||||
palace_path = os.path.join(tmpdir, "palace")
|
||||
|
||||
# First mine — stamps drawers with NORMALIZE_VERSION
|
||||
mine_convos(tmpdir, palace_path, wing="test")
|
||||
capsys.readouterr()
|
||||
|
||||
client = chromadb.PersistentClient(path=palace_path)
|
||||
col = client.get_collection("mempalace_drawers")
|
||||
resolved = str(Path(tmpdir).resolve() / "chat.txt")
|
||||
first_pass = col.get(where={"source_file": resolved})
|
||||
first_ids = set(first_pass["ids"])
|
||||
assert first_ids, "first mine should produce drawers"
|
||||
for meta in first_pass["metadatas"]:
|
||||
assert meta.get("normalize_version") == NORMALIZE_VERSION
|
||||
|
||||
# Simulate pre-v2 drawers: rewrite metadata to an older version,
|
||||
# and replace content with "noise" so we can see it get cleaned up.
|
||||
stale_metas = []
|
||||
for meta in first_pass["metadatas"]:
|
||||
stale = dict(meta)
|
||||
stale["normalize_version"] = 1
|
||||
stale_metas.append(stale)
|
||||
col.update(
|
||||
ids=list(first_pass["ids"]),
|
||||
documents=["STALE NOISE"] * len(first_pass["ids"]),
|
||||
metadatas=stale_metas,
|
||||
)
|
||||
# Add an extra orphan drawer that should also be purged.
|
||||
col.add(
|
||||
ids=["orphan_drawer"],
|
||||
documents=["OLD ORPHAN"],
|
||||
metadatas=[
|
||||
{
|
||||
"wing": "test",
|
||||
"room": "default",
|
||||
"source_file": resolved,
|
||||
"chunk_index": 999,
|
||||
"normalize_version": 1,
|
||||
}
|
||||
],
|
||||
)
|
||||
del col, client
|
||||
|
||||
# Second mine — version gate should trigger rebuild
|
||||
mine_convos(tmpdir, palace_path, wing="test")
|
||||
out = capsys.readouterr().out
|
||||
assert (
|
||||
"Files skipped (already filed): 0" in out
|
||||
), "stale drawers should force a rebuild, not a skip"
|
||||
|
||||
client = chromadb.PersistentClient(path=palace_path)
|
||||
col = client.get_collection("mempalace_drawers")
|
||||
rebuilt = col.get(where={"source_file": resolved})
|
||||
# Orphan is gone
|
||||
assert "orphan_drawer" not in rebuilt["ids"]
|
||||
# No stale content survived
|
||||
assert all("STALE NOISE" not in d for d in rebuilt["documents"])
|
||||
assert all("OLD ORPHAN" not in d for d in rebuilt["documents"])
|
||||
# All rebuilt drawers carry the current version
|
||||
for meta in rebuilt["metadatas"]:
|
||||
assert meta.get("normalize_version") == NORMALIZE_VERSION
|
||||
del col, client
|
||||
finally:
|
||||
shutil.rmtree(tmpdir, ignore_errors=True)
|
||||
|
||||
@@ -6,6 +6,7 @@ dispatch layer (integration-level). Uses isolated palace + KG fixtures
|
||||
via monkeypatch to avoid touching real data.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
import json
|
||||
import sys
|
||||
|
||||
@@ -643,6 +644,48 @@ class TestDiaryTools:
|
||||
r = tool_diary_read(agent_name="Nobody")
|
||||
assert r["entries"] == []
|
||||
|
||||
def test_diary_write_same_second_shared_prefix_no_collision(
|
||||
self, monkeypatch, config, palace_path, kg
|
||||
):
|
||||
_patch_mcp_server(monkeypatch, config, kg)
|
||||
_client, _col = _get_collection(palace_path, create=True)
|
||||
del _client
|
||||
|
||||
from mempalace import mcp_server
|
||||
|
||||
class FrozenDateTime:
|
||||
calls = [
|
||||
datetime(2026, 4, 13, 22, 15, 30, 123456),
|
||||
datetime(2026, 4, 13, 22, 15, 30, 123457),
|
||||
]
|
||||
fallback = datetime(2026, 4, 13, 22, 15, 30, 123457)
|
||||
|
||||
@classmethod
|
||||
def now(cls):
|
||||
if cls.calls:
|
||||
return cls.calls.pop(0)
|
||||
return cls.fallback
|
||||
|
||||
monkeypatch.setattr(mcp_server, "datetime", FrozenDateTime)
|
||||
|
||||
from mempalace.mcp_server import tool_diary_read, tool_diary_write
|
||||
|
||||
entry1 = "A" * 50 + " entry one"
|
||||
entry2 = "A" * 50 + " entry two"
|
||||
|
||||
result1 = tool_diary_write(agent_name="TestAgent", entry=entry1, topic="status")
|
||||
result2 = tool_diary_write(agent_name="TestAgent", entry=entry2, topic="status")
|
||||
|
||||
assert result1["success"] is True
|
||||
assert result2["success"] is True
|
||||
assert result1["entry_id"] != result2["entry_id"]
|
||||
|
||||
read_result = tool_diary_read(agent_name="TestAgent")
|
||||
contents = [entry["content"] for entry in read_result["entries"]]
|
||||
assert read_result["total"] == 2
|
||||
assert entry1 in contents
|
||||
assert entry2 in contents
|
||||
|
||||
|
||||
# ── Cache Invalidation (inode/mtime) ──────────────────────────────────
|
||||
|
||||
|
||||
+90
-4
@@ -7,7 +7,7 @@ import chromadb
|
||||
import yaml
|
||||
|
||||
from mempalace.miner import mine, scan_project, status
|
||||
from mempalace.palace import file_already_mined
|
||||
from mempalace.palace import NORMALIZE_VERSION, file_already_mined
|
||||
|
||||
|
||||
def write_file(path: Path, content: str):
|
||||
@@ -227,11 +227,17 @@ def test_file_already_mined_check_mtime():
|
||||
assert file_already_mined(col, test_file) is False
|
||||
assert file_already_mined(col, test_file, check_mtime=True) is False
|
||||
|
||||
# Add it with mtime
|
||||
# Add it with mtime + current normalize_version
|
||||
col.add(
|
||||
ids=["d1"],
|
||||
documents=["hello world"],
|
||||
metadatas=[{"source_file": test_file, "source_mtime": str(mtime)}],
|
||||
metadatas=[
|
||||
{
|
||||
"source_file": test_file,
|
||||
"source_mtime": str(mtime),
|
||||
"normalize_version": NORMALIZE_VERSION,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
# Already mined (no mtime check)
|
||||
@@ -253,7 +259,12 @@ def test_file_already_mined_check_mtime():
|
||||
col.add(
|
||||
ids=["d2"],
|
||||
documents=["other"],
|
||||
metadatas=[{"source_file": "/fake/no_mtime.txt"}],
|
||||
metadatas=[
|
||||
{
|
||||
"source_file": "/fake/no_mtime.txt",
|
||||
"normalize_version": NORMALIZE_VERSION,
|
||||
}
|
||||
],
|
||||
)
|
||||
assert file_already_mined(col, "/fake/no_mtime.txt", check_mtime=True) is False
|
||||
finally:
|
||||
@@ -296,3 +307,78 @@ def test_status_missing_palace_does_not_create_empty_collection(tmp_path, capsys
|
||||
out = capsys.readouterr().out
|
||||
assert "No palace found" in out
|
||||
assert not palace_path.exists()
|
||||
|
||||
|
||||
# ── normalize_version schema gate ───────────────────────────────────────
|
||||
#
|
||||
# When the normalization pipeline changes shape (e.g., strip_noise lands),
|
||||
# `NORMALIZE_VERSION` is bumped so pre-existing drawers can be silently
|
||||
# rebuilt on the next mine. These tests pin that contract.
|
||||
|
||||
|
||||
def test_file_already_mined_returns_false_for_stale_normalize_version():
|
||||
"""Pre-v2 drawers (no field, or older integer) must not short-circuit."""
|
||||
tmpdir = tempfile.mkdtemp()
|
||||
try:
|
||||
palace_path = os.path.join(tmpdir, "palace")
|
||||
os.makedirs(palace_path)
|
||||
client = chromadb.PersistentClient(path=palace_path)
|
||||
col = client.get_or_create_collection("mempalace_drawers")
|
||||
|
||||
# Pre-v2 drawer: no normalize_version field at all
|
||||
col.add(
|
||||
ids=["d_old"],
|
||||
documents=["old"],
|
||||
metadatas=[{"source_file": "/fake/old.jsonl"}],
|
||||
)
|
||||
assert file_already_mined(col, "/fake/old.jsonl") is False
|
||||
|
||||
# Explicitly older version
|
||||
col.add(
|
||||
ids=["d_v1"],
|
||||
documents=["v1"],
|
||||
metadatas=[{"source_file": "/fake/v1.jsonl", "normalize_version": 1}],
|
||||
)
|
||||
assert file_already_mined(col, "/fake/v1.jsonl") is False
|
||||
|
||||
# Current version — short-circuits
|
||||
col.add(
|
||||
ids=["d_current"],
|
||||
documents=["cur"],
|
||||
metadatas=[
|
||||
{
|
||||
"source_file": "/fake/current.jsonl",
|
||||
"normalize_version": NORMALIZE_VERSION,
|
||||
}
|
||||
],
|
||||
)
|
||||
assert file_already_mined(col, "/fake/current.jsonl") is True
|
||||
finally:
|
||||
del col, client
|
||||
shutil.rmtree(tmpdir, ignore_errors=True)
|
||||
|
||||
|
||||
def test_add_drawer_stamps_normalize_version(tmp_path):
|
||||
"""Fresh drawers carry the current schema version so future upgrades work."""
|
||||
from mempalace.miner import add_drawer
|
||||
|
||||
palace_path = tmp_path / "palace"
|
||||
palace_path.mkdir()
|
||||
client = chromadb.PersistentClient(path=str(palace_path))
|
||||
col = client.get_or_create_collection("mempalace_drawers")
|
||||
try:
|
||||
added = add_drawer(
|
||||
collection=col,
|
||||
wing="test",
|
||||
room="notes",
|
||||
content="hello",
|
||||
source_file=str(tmp_path / "src.md"),
|
||||
chunk_index=0,
|
||||
agent="unit",
|
||||
)
|
||||
assert added is True
|
||||
stored = col.get(limit=1)
|
||||
meta = stored["metadatas"][0]
|
||||
assert meta["normalize_version"] == NORMALIZE_VERSION
|
||||
finally:
|
||||
del col, client
|
||||
|
||||
@@ -13,6 +13,7 @@ from mempalace.normalize import (
|
||||
_try_normalize_json,
|
||||
_try_slack_json,
|
||||
normalize,
|
||||
strip_noise,
|
||||
)
|
||||
|
||||
|
||||
@@ -1048,3 +1049,148 @@ def test_normalize_rejects_large_file():
|
||||
assert False, "Should have raised IOError"
|
||||
except IOError as e:
|
||||
assert "too large" in str(e).lower()
|
||||
|
||||
|
||||
# ── strip_noise() — verbatim-safety boundary tests ─────────────────────
|
||||
#
|
||||
# The "Verbatim always" design principle requires that we never delete
|
||||
# user-authored text. These tests pin down the boundary between system
|
||||
# noise (which we strip) and user prose that happens to mention the same
|
||||
# strings (which must survive untouched).
|
||||
|
||||
|
||||
class TestStripNoisePreservesUserContent:
|
||||
"""User prose that mentions noise strings inline must be preserved."""
|
||||
|
||||
def test_user_discusses_stop_hook_in_prose(self):
|
||||
# Regression: original regex with IGNORECASE + `.*\n?` ate the second
|
||||
# sentence from real user commentary.
|
||||
text = (
|
||||
"> User:\n"
|
||||
"> Our CI has a stop hook that rejects merges after 5pm. "
|
||||
"Ran 2 stop hooks last week.\n"
|
||||
"> Assistant:\n"
|
||||
"> Got it."
|
||||
)
|
||||
assert strip_noise(text) == text.strip()
|
||||
|
||||
def test_user_mentions_system_reminder_inline(self):
|
||||
# Inline <system-reminder> tags inside user prose (e.g. documenting
|
||||
# Claude Code behavior) must not be stripped.
|
||||
text = (
|
||||
"> User:\n"
|
||||
"> Here is what Claude Code emits: "
|
||||
"<system-reminder>Auto-save reminder...</system-reminder>"
|
||||
" — I want to ignore it."
|
||||
)
|
||||
assert strip_noise(text) == text.strip()
|
||||
|
||||
def test_ctrl_o_hint_in_prose_preserved(self):
|
||||
# Regression: original `.*\(ctrl\+o to expand\).*\n?` nuked the whole
|
||||
# line whenever a user documented the TUI shortcut.
|
||||
text = (
|
||||
"> User:\n"
|
||||
"> In the TUI you hit (ctrl+o to expand) to see more. "
|
||||
"That is the shortcut I want to document."
|
||||
)
|
||||
assert strip_noise(text) == text.strip()
|
||||
|
||||
def test_current_time_inline_in_prose(self):
|
||||
text = "> User:\n> At CURRENT TIME: the meeting starts, not before."
|
||||
assert strip_noise(text) == text.strip()
|
||||
|
||||
def test_plus_n_lines_marker_inline(self):
|
||||
text = "> User:\n> The log showed … +50 lines of stack trace, useful."
|
||||
assert strip_noise(text) == text.strip()
|
||||
|
||||
def test_dangling_open_tag_does_not_span_messages(self):
|
||||
# THE span-eating bug: a stray unclosed <system-reminder> in one
|
||||
# message must NOT merge with a closing tag in another message and
|
||||
# silently delete everything in between.
|
||||
text = (
|
||||
"> User 1: normal content <system-reminder>A\n"
|
||||
"> Assistant: reply\n"
|
||||
"> User 2: more content</system-reminder> tail"
|
||||
)
|
||||
out = strip_noise(text)
|
||||
assert "Assistant: reply" in out
|
||||
assert "User 2: more content" in out
|
||||
assert "User 1: normal content" in out
|
||||
|
||||
|
||||
class TestStripNoiseRemovesSystemChrome:
|
||||
"""System-injected noise with standalone/line-anchored shape must be stripped."""
|
||||
|
||||
def test_strips_line_anchored_system_reminder_block(self):
|
||||
text = (
|
||||
"> User:\n"
|
||||
"<system-reminder>\n"
|
||||
"Auto-save reminder...\n"
|
||||
"</system-reminder>\n"
|
||||
"> Real message."
|
||||
)
|
||||
out = strip_noise(text)
|
||||
assert "system-reminder" not in out
|
||||
assert "Auto-save reminder" not in out
|
||||
assert "Real message." in out
|
||||
|
||||
def test_strips_system_reminder_with_blockquote_prefix(self):
|
||||
# _messages_to_transcript prefixes lines with "> ", so the line
|
||||
# anchor must also accept that shape.
|
||||
text = "> User:\n" "> <system-reminder>Injected noise</system-reminder>\n" "> Real message."
|
||||
out = strip_noise(text)
|
||||
assert "Injected noise" not in out
|
||||
assert "Real message." in out
|
||||
|
||||
def test_strips_standalone_ran_hook_line(self):
|
||||
text = "Ran 2 Stop hook\n> User: real content"
|
||||
out = strip_noise(text)
|
||||
assert "Ran 2 Stop hook" not in out
|
||||
assert "real content" in out
|
||||
|
||||
def test_strips_known_hook_names(self):
|
||||
for hook in ("Stop", "PreCompact", "PreToolUse", "PostToolUse", "UserPromptSubmit"):
|
||||
text = f"Ran 1 {hook} hook\n> User: content"
|
||||
assert hook not in strip_noise(text)
|
||||
|
||||
def test_strips_current_time_standalone(self):
|
||||
text = "CURRENT TIME: 2026-04-13 10:00 UTC\n> User: Hello"
|
||||
out = strip_noise(text)
|
||||
assert "CURRENT TIME" not in out
|
||||
assert "Hello" in out
|
||||
|
||||
def test_strips_collapsed_lines_marker(self):
|
||||
text = "… +42 lines\n> User: Hello"
|
||||
out = strip_noise(text)
|
||||
assert "+42 lines" not in out
|
||||
assert "Hello" in out
|
||||
|
||||
def test_strips_token_count_ctrl_o_chrome(self):
|
||||
# Claude Code's actual collapsed-output chrome: "[N tokens] (ctrl+o to expand)"
|
||||
text = "> Assistant: some output [5 tokens] (ctrl+o to expand)\n> User: ok"
|
||||
out = strip_noise(text)
|
||||
assert "(ctrl+o to expand)" not in out
|
||||
assert "[5 tokens]" not in out
|
||||
assert "some output" in out
|
||||
|
||||
def test_strips_each_known_noise_tag(self):
|
||||
for tag in (
|
||||
"system-reminder",
|
||||
"command-message",
|
||||
"command-name",
|
||||
"task-notification",
|
||||
"user-prompt-submit-hook",
|
||||
"hook_output",
|
||||
):
|
||||
text = f"> User:\n<{tag}>junk</{tag}>\n> Real."
|
||||
out = strip_noise(text)
|
||||
assert tag not in out, f"{tag} leaked into output"
|
||||
assert "Real." in out
|
||||
|
||||
def test_collapses_excessive_blank_lines(self):
|
||||
text = "line one\n\n\n\n\n\nline two"
|
||||
out = strip_noise(text)
|
||||
assert "line one" in out
|
||||
assert "line two" in out
|
||||
# Should collapse to no more than 3 newlines
|
||||
assert "\n\n\n\n" not in out
|
||||
|
||||
Reference in New Issue
Block a user