diff --git a/mempalace/convo_miner.py b/mempalace/convo_miner.py index 24898d248c..02a5aca9ba 100644 --- a/mempalace/convo_miner.py +++ b/mempalace/convo_miner.py @@ -20,6 +20,7 @@ from .palace import ( NORMALIZE_VERSION, SKIP_DIRS, + _metadata_matches_extract_mode, file_already_mined, get_collection, mine_lock, @@ -70,14 +71,15 @@ def _detect_hall_cached(content: str) -> str: # use also scales with source size. -def _register_file(collection, source_file: str, wing: str, agent: str): +def _register_file(collection, source_file: str, wing: str, agent: str, extract_mode: str): """Write a sentinel so file_already_mined() returns True for 0-chunk files. Without this, files that normalize to nothing or produce zero chunks are re-read and re-processed on every mine run because nothing was written to ChromaDB on the first pass. """ - sentinel_id = f"_reg_{hashlib.sha256(source_file.encode()).hexdigest()[:24]}" + sentinel_key = f"{source_file}:{extract_mode}" + sentinel_id = f"_reg_{hashlib.sha256(sentinel_key.encode()).hexdigest()[:24]}" collection.upsert( documents=[f"[registry] {source_file}"], ids=[sentinel_id], @@ -89,12 +91,40 @@ def _register_file(collection, source_file: str, wing: str, agent: str): "added_by": agent, "filed_at": datetime.now().isoformat(), "ingest_mode": "registry", + "extract_mode": extract_mode, "normalize_version": NORMALIZE_VERSION, } ], ) +def _source_file_delete_ids(collection, source_file: str, extract_mode: str) -> list[str]: + """Collect drawer IDs for one source file and extraction mode. + + Legacy conversation drawers did not carry extract_mode; treat those as + exchange-mode rows so schema rebuilds can still clean them up without + deleting newer general-mode drawers for the same transcript. + """ + ids: list[str] = [] + offset = 0 + while True: + batch = collection.get( + where={"source_file": source_file}, + limit=1000, + offset=offset, + include=["metadatas"], + ) + batch_ids = batch.get("ids") or [] + metadatas = batch.get("metadatas") or [] + for drawer_id, meta in zip(batch_ids, metadatas): + if _metadata_matches_extract_mode(meta or {}, extract_mode): + ids.append(drawer_id) + if not batch_ids: + break + offset += len(batch_ids) + return ids + + # ============================================================================= # CHUNKING — exchange pairs for conversations # ============================================================================= @@ -358,14 +388,16 @@ def _file_chunks_locked(collection, source_file, chunks, wing, room, agent, extr # 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): + if file_already_mined(collection, source_file, extract_mode=extract_mode): 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}) + delete_ids = _source_file_delete_ids(collection, source_file, extract_mode) + if delete_ids: + collection.delete(ids=delete_ids) except Exception: logger.debug("Stale-drawer purge failed for %s", source_file, exc_info=True) @@ -382,7 +414,11 @@ def _file_chunks_locked(collection, source_file, chunks, wing, room, agent, extr 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]}" + drawer_key = f"{source_file}:{extract_mode}:{chunk['chunk_index']}" + drawer_id = ( + f"drawer_{wing}_{chunk_room}_" + f"{hashlib.sha256(drawer_key.encode()).hexdigest()[:24]}" + ) batch_docs.append(chunk["content"]) batch_ids.append(drawer_id) batch_metas.append( @@ -478,7 +514,9 @@ def mine_convos( # palace each per-file query costs ~2s, so a 2000-file sweep used to # spend >1h just deciding to skip. prefetch_mined_set() does the same # decisions in a single scan; loop body becomes an O(1) set check. - mined_set: set[str] = prefetch_mined_set(collection) if not dry_run else set() + mined_set: set[str] = ( + prefetch_mined_set(collection, extract_mode=extract_mode) if not dry_run else set() + ) total_drawers = 0 files_skipped = 0 @@ -497,12 +535,12 @@ def mine_convos( content = normalize(str(filepath)) except (OSError, ValueError): if not dry_run: - _register_file(collection, source_file, wing, agent) + _register_file(collection, source_file, wing, agent, extract_mode) continue if not content or len(content.strip()) < cfg_min_chunk_size: if not dry_run: - _register_file(collection, source_file, wing, agent) + _register_file(collection, source_file, wing, agent, extract_mode) continue # Chunk — either exchange pairs or general extraction @@ -520,7 +558,7 @@ def mine_convos( if not chunks: if not dry_run: - _register_file(collection, source_file, wing, agent) + _register_file(collection, source_file, wing, agent, extract_mode) continue # Detect room from content (general mode uses memory_type instead) diff --git a/mempalace/palace.py b/mempalace/palace.py index 38c3979763..3e3cfa7f26 100644 --- a/mempalace/palace.py +++ b/mempalace/palace.py @@ -531,7 +531,19 @@ def mine_palace_lock(palace_path: str): mine_global_lock = mine_palace_lock -def file_already_mined(collection, source_file: str, check_mtime: bool = False) -> bool: +def _metadata_matches_extract_mode(meta: dict, extract_mode: Optional[str]) -> bool: + if extract_mode is None: + return True + stored_mode = meta.get("extract_mode") + return stored_mode == extract_mode or (extract_mode == "exchange" and stored_mode is None) + + +def file_already_mined( + collection, + source_file: str, + check_mtime: bool = False, + extract_mode: Optional[str] = None, +) -> bool: """Check if a file has already been filed in the palace. Returns False (so the file gets re-mined) when: @@ -543,12 +555,43 @@ def file_already_mined(collection, source_file: str, check_mtime: bool = False) 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. + + When extract_mode is set (used by convo miner), idempotency is scoped to + that extraction mode so exchange-mode and general-mode drawers can coexist + for the same source transcript. Legacy drawers without extract_mode are + treated as exchange-mode drawers. """ try: - results = collection.get(where={"source_file": source_file}, limit=1) - if not results.get("ids"): + stored_meta = None + if extract_mode is None: + results = collection.get(where={"source_file": source_file}, limit=1) + if not results.get("ids"): + return False + stored_meta = results.get("metadatas", [{}])[0] or {} + else: + offset = 0 + while True: + results = collection.get( + where={"source_file": source_file}, + limit=1000, + offset=offset, + include=["metadatas"], + ) + ids = results.get("ids") or [] + metadatas = results.get("metadatas") or [] + stored_meta = next( + ( + meta or {} + for meta in metadatas + if _metadata_matches_extract_mode(meta or {}, extract_mode) + ), + None, + ) + if stored_meta is not None or not ids: + break + offset += len(ids) + if stored_meta is None: 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: @@ -593,7 +636,7 @@ def bulk_check_mined(collection) -> dict[str, float]: return mined -def prefetch_mined_set(collection) -> set[str]: +def prefetch_mined_set(collection, extract_mode: Optional[str] = None) -> set[str]: """Pre-fetch the set of source_files already mined at the current NORMALIZE_VERSION. Mirrors file_already_mined()'s version-gate semantics (check_mtime=False @@ -601,6 +644,9 @@ def prefetch_mined_set(collection) -> set[str]: Returns a set of source_file paths whose stored drawers are at or above NORMALIZE_VERSION; callers do `if path in result_set: skip`. + When extract_mode is set, mirrors file_already_mined(..., extract_mode=...) + so conversation mines skip per extraction mode rather than per source file. + The convo miner walks thousands of transcript files; per-file `collection.get(where={"source_file": X})` costs ~2s on a 150k-drawer palace, making a 2000-file sweep take >1h of pure skip-checking. This @@ -613,9 +659,12 @@ def prefetch_mined_set(collection) -> set[str]: while offset < total: batch = collection.get(limit=1000, offset=offset, include=["metadatas"]) for meta in batch["metadatas"]: + meta = meta or {} src = meta.get("source_file") if not src: continue + if not _metadata_matches_extract_mode(meta, extract_mode): + continue # Same default as file_already_mined: missing version == 1 version = meta.get("normalize_version", 1) if version >= NORMALIZE_VERSION: diff --git a/tests/test_convo_miner.py b/tests/test_convo_miner.py index 166644b002..015a9d3d40 100644 --- a/tests/test_convo_miner.py +++ b/tests/test_convo_miner.py @@ -77,6 +77,38 @@ def test_mine_convos_does_not_reprocess_empty_chunk_files(capsys): shutil.rmtree(tmpdir, ignore_errors=True) +def test_mine_convos_allows_general_after_exchange(capsys): + """A transcript mined as exchange can later be mined as general memories.""" + tmpdir = tempfile.mkdtemp() + try: + convo_path = Path(tmpdir) / "chat.txt" + convo_path.write_text( + "> What did we decide?\n" + "We decided to use SQLite because it keeps the local setup simple.\n\n" + "> What broke?\n" + "The search failed because the old index was stale, and the fix was rebuild.\n" + ) + palace_path = os.path.join(tmpdir, "palace") + + mine_convos(tmpdir, palace_path, wing="test", extract_mode="exchange") + capsys.readouterr() + mine_convos(tmpdir, palace_path, wing="test", extract_mode="general") + out = capsys.readouterr().out + + assert "Files skipped (already filed): 0" in out + + client = chromadb.PersistentClient(path=palace_path) + col = client.get_collection("mempalace_drawers") + resolved = str(Path(tmpdir).resolve() / "chat.txt") + rows = col.get(where={"source_file": resolved}, include=["metadatas"]) + modes = {meta.get("extract_mode") for meta in rows["metadatas"]} + assert {"exchange", "general"} <= modes + assert any(drawer_id.startswith("drawer_test_decision_") for drawer_id in rows["ids"]) + del col, client + 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. diff --git a/tests/test_convo_miner_unit.py b/tests/test_convo_miner_unit.py index 90ad573625..7ff7a91bc3 100644 --- a/tests/test_convo_miner_unit.py +++ b/tests/test_convo_miner_unit.py @@ -239,7 +239,7 @@ def upsert(self, documents, ids, metadatas): col = FakeCol() monkeypatch.setattr(convo_miner, "DRAWER_UPSERT_BATCH_SIZE", 2) monkeypatch.setattr( - convo_miner, "file_already_mined", lambda collection, source_file: False + convo_miner, "file_already_mined", lambda collection, source_file, **kwargs: False ) monkeypatch.setattr(convo_miner, "mine_lock", lambda source_file: contextlib.nullcontext()) monkeypatch.setattr(convo_miner, "_detect_hall_cached", lambda content: "conversations") diff --git a/tests/test_miner.py b/tests/test_miner.py index d58ef9201e..50535273bb 100644 --- a/tests/test_miner.py +++ b/tests/test_miner.py @@ -10,7 +10,7 @@ import yaml from mempalace.miner import detect_room, load_config, mine, scan_project, status -from mempalace.palace import NORMALIZE_VERSION, file_already_mined +from mempalace.palace import NORMALIZE_VERSION, file_already_mined, prefetch_mined_set def write_file(path: Path, content: str): @@ -406,6 +406,82 @@ def test_file_already_mined_check_mtime(): shutil.rmtree(tmpdir, ignore_errors=True) +def test_file_already_mined_scopes_convo_extract_mode(): + 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", metadata={"hnsw:space": "cosine"} + ) + + source_file = os.path.join(tmpdir, "chat.jsonl") + col.add( + ids=["exchange"], + documents=["exchange drawer"], + metadatas=[ + { + "source_file": source_file, + "extract_mode": "exchange", + "normalize_version": NORMALIZE_VERSION, + } + ], + ) + + assert file_already_mined(col, source_file, extract_mode="exchange") is True + assert file_already_mined(col, source_file, extract_mode="general") is False + assert source_file in prefetch_mined_set(col, extract_mode="exchange") + assert source_file not in prefetch_mined_set(col, extract_mode="general") + + col.add( + ids=["general"], + documents=["general drawer"], + metadatas=[ + { + "source_file": source_file, + "extract_mode": "general", + "normalize_version": NORMALIZE_VERSION, + } + ], + ) + + assert file_already_mined(col, source_file, extract_mode="general") is True + assert source_file in prefetch_mined_set(col, extract_mode="general") + finally: + del col, client + shutil.rmtree(tmpdir, ignore_errors=True) + + +def test_file_already_mined_extract_mode_paginates_large_sources(): + source_file = "/tmp/long-chat.jsonl" + metadatas = [ + { + "source_file": source_file, + "extract_mode": "exchange", + "normalize_version": NORMALIZE_VERSION, + } + for _ in range(1000) + ] + metadatas.append( + { + "source_file": source_file, + "extract_mode": "general", + "normalize_version": NORMALIZE_VERSION, + } + ) + + class FakeCollection: + def get(self, where=None, limit=1, offset=0, include=None): + batch = metadatas[offset : offset + limit] + return { + "ids": [f"id-{i}" for i in range(offset, offset + len(batch))], + "metadatas": batch, + } + + assert file_already_mined(FakeCollection(), source_file, extract_mode="general") is True + + def test_mine_dry_run_with_tiny_file_no_crash(): """Dry-run must not crash when process_file returns 0 drawers (room was None).""" tmpdir = tempfile.mkdtemp()