diff --git a/README.md b/README.md index ea89d60..83e3e96 100644 --- a/README.md +++ b/README.md @@ -20,8 +20,10 @@ Built as a personal second-brain pipeline; published in case it's useful to anyo - **Capture anything** — text, URLs, photos, voice notes, PDFs, Word docs (`.docx`/`.doc`), plain-text/code files, forwarded posts, YouTube links. Media groups (multi-photo posts) are debounced and stitched into a single note. - **Read the contents, not just the message** — Claude vision OCR for photos, OpenAI Whisper for voice, PyPDF/python-docx for documents, YouTube Transcript API for videos, plain HTTP fetcher for web pages. - **AI routing** — Claude (Sonnet) picks a folder, writes a title, summarises the body, generates tags, and proposes up to 5 related notes from your existing vault. Falls back to an "Other" bucket and an inbox queue when confidence is low. -- **Smart dedupe** — incoming notes that match by URL, title, or semantic similarity (cosine on OpenAI embeddings) get *appended* to the existing note instead of creating a duplicate. +- **Merge-and-rewrite dedupe** — incoming notes that match by URL, title, or semantic similarity (cosine on OpenAI embeddings) are **merged into the canonical note**: Claude rewrites the page to integrate the new source, collapse redundancy, and reconcile contradictions inline (instead of stacking dated append blocks). The vault is git-snapshotted before every rewrite; if a rewrite would drop the frontmatter or an attachment it safely falls back to a plain append, so a capture is never lost. +- **Entity (wiki) pages** — every capture also grows typed backbone pages under `People/`, `Concepts/`, and `Projects/`. Each entity page accumulates one grounded observation per source note plus `[[backlinks]]`, turning the chronological capture stream into a navigable wiki. - **Vault-grounded Q&A** — `/ask` runs a hybrid keyword + embedding retrieval over your notes and answers with Claude, with multi-turn follow-ups via Telegram reply threads. +- **Retroactive rebuild** — `/rebuild` git-snapshots the vault, merges existing duplicate notes, and rebuilds all entity pages from scratch. One command turns an existing note pile into the wiki. - **Manual override** — every capture shows an inline-keyboard folder picker; misroutes are one tap away. `/redo`, `/edit`, `/undo`, and `/relink` cover the rest. - **Single-tenant by design** — an `ALLOWED_USER_IDS` allowlist gates every handler. Nobody else who finds your bot can use it. @@ -37,9 +39,14 @@ Telegram message ↓ Claude enrichment: title · summary · tags · folder · related notes · confidence ↓ -Dedupe check (URL → title → semantic) → append to existing OR create new +Dedupe check (URL → title → semantic) + ├─ match → git snapshot → Claude merge-rewrites the canonical note + └─ no match → create new note + ↓ +Entity pass: extract people/concepts/projects → grow typed wiki pages ↓ //.md with YAML frontmatter + [[backlinks]] + attachments/ +<vault>/{People,Concepts,Projects}/<Entity>.md accumulating observations + backlinks ``` ## Requirements @@ -109,11 +116,23 @@ All config is via environment variables (loaded from `.env` if present). See [.e | `/inbox` | List notes flagged for review (low-confidence routing) | | `/review` | Walk pending notes one at a time with move / mark-reviewed / delete buttons | | `/relink [folder]` | Refresh related-note backlinks. No arg = last capture; with arg = entire folder | +| `/rebuild` | Git-snapshot the vault, merge existing duplicate notes, and rebuild all `People`/`Concepts`/`Projects` entity pages from scratch. Destructive but reversible (see below) | | `/redo` | Reply with `/redo` to regenerate a capture using the higher-quality Opus model | | `/edit <text>` | Replace the source of the last capture and re-enrich | | `/undo` | Delete the last capture in this chat | | `/refresh` | Rescan the vault index (also runs automatically every 10 minutes) | +### Rolling back a merge or rebuild + +Every merge and every `/rebuild` commits the vault to a git repo (auto-initialised at `BASE_DIR` on first use) **before** writing. To undo the most recent rewrite: + +```bash +git -C "$BASE_DIR" log --oneline # find the engram: pre-* commit +git -C "$BASE_DIR" reset --hard HEAD~ # discard the last rewrite +``` + +> **iCloud note:** if your vault lives in an iCloud-synced folder, the `.git` directory is synced too. This works fine but can cause occasional sync churn; that is the documented trade-off for in-vault rollback. + The plain message path: send a message → tap a folder button → done. Send a photo without a caption and the bot OCRs it first so it can route by content. ## What a note looks like @@ -157,11 +176,12 @@ uv sync uv run pytest -v ``` -15 test modules cover the bot handlers, vault indexing, embeddings, dedupe, link enrichment, vision/whisper/youtube/pdf adapters, and the inbox/review flow. `pytest-asyncio` is in `auto` mode. CI runs on push and PR against Python 3.11 / 3.12 / 3.13. +The test suite covers the bot handlers, vault indexing, embeddings, dedupe, link enrichment, the merge-and-rewrite path, entity extraction/pages, the retroactive `/rebuild` flow, git snapshotting, vision/whisper/youtube/pdf adapters, and the inbox/review flow. `pytest-asyncio` is in `auto` mode. CI runs on push and PR against Python 3.11 / 3.12 / 3.13. ## Roadmap - **Local-model support** — swap Claude / OpenAI for Ollama or llama.cpp so the bot can run end-to-end without paid API keys. Embeddings first (cheapest win), then enrichment. Tracked in [#1](https://github.com/mishablank/Engram/issues/1) — help welcome. +- **Scheduled maintenance** — a nightly pass that re-synthesises entity leads, reconciles contradictions across notes, and DMs a digest of what changed in the vault overnight. ## Security model diff --git a/src/engram/bot.py b/src/engram/bot.py index 27c54ed..b22947b 100644 --- a/src/engram/bot.py +++ b/src/engram/bot.py @@ -24,17 +24,24 @@ filters, ) -from . import pending_store -from .config import DEFAULT_CATEGORY, Config, load_config +from . import gitsafe, pending_store +from .config import DEFAULT_CATEGORY, ENTITY_TYPE_FOLDERS, Config, load_config from .dedupe import find_semantic_duplicate from .embeddings import OpenAIEmbedder, SemanticIndex, hybrid_search +from .entities import extract_entities, rebuild_entity_pages, upsert_entity_page from .fetcher import fetch_urls from .inbox import clear_pending, find_pending, move_note, preview from .linker import answer_from_vault, enrich_note -from .note_writer import CapturedMessage, append_to_note, extract_urls, write_note +from .note_writer import ( + CapturedMessage, + append_to_note, + extract_urls, + merge_into_note, + write_note, +) from .pdf import extract_pdf_text from .relinker import relink_note -from .vault import VaultIndex, scan_vault, search_vault +from .vault import VaultIndex, load_note_body, scan_vault, search_vault from .vision import ocr_image from .whisper import transcribe @@ -441,8 +448,20 @@ async def _process_messages( original_message_id=fwd_msg_id, review_pending=review_pending, ) - path = append_to_note(duplicate, captured) - return path, f"{path.parent.name} (appended)" + # Karpathy-wiki merge: rewrite the canonical note to integrate the new + # source instead of appending a dated block. Snapshot first so a bad + # rewrite is one `git revert` away; merge_into_note falls back to a safe + # append if the LLM rewrite would drop frontmatter or an attachment. + gitsafe.snapshot( + state.config.base_dir, f"engram: pre-merge {duplicate.stem}" + ) + path, merged = await asyncio.to_thread( + merge_into_note, duplicate, captured, state.anthropic + ) + state.invalidate_vault_cache() + await asyncio.to_thread(_update_entities, state, path.stem, body_text) + verb = "merged" if merged else "appended" + return path, f"{path.parent.name} ({verb})" target_dir = state.config.base_dir / folder attachments_dir = target_dir / "attachments" @@ -467,9 +486,24 @@ async def _process_messages( ) path = write_note(captured, target_dir, related=related, tags=tags) state.invalidate_vault_cache() + await asyncio.to_thread(_update_entities, state, path.stem, body_text) return path, folder +def _update_entities(state: BotState, source_title: str, body: str) -> None: + """Extract entities from a captured note and grow their typed wiki pages. + + Best-effort: never raises into the capture path. + """ + if not body.strip(): + return + try: + for ent in extract_entities(state.anthropic, source_title, body): + upsert_entity_page(state.config.base_dir, ent, source_title) + except Exception: + log.exception("Entity update failed for %s", source_title) + + def _move_pending_files(srcs: list[Path], target_dir: Path) -> list[str]: if not srcs: return [] @@ -1228,14 +1262,67 @@ async def on_relink(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: f"Done. Updated {changed_count}/{len(notes)} note(s)." ) + async def on_rebuild(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + if not _is_authorized(state, update): + return + base_dir = state.config.base_dir + await update.effective_chat.send_message( + "Rebuilding vault: snapshotting to git, merging duplicates, " + "rebuilding entity pages… this may take a minute." + ) + committed = await asyncio.to_thread( + gitsafe.snapshot, base_dir, "engram: pre-rebuild snapshot" + ) + + merges = 0 + if state._semantic_index.enabled: + await asyncio.to_thread(state._semantic_index.refresh) + note_paths = _vault_note_paths(base_dir) + merges = await asyncio.to_thread( + merge_duplicate_notes, + base_dir, + state._semantic_index, + state.anthropic, + note_paths, + ) + + notes = [ + (p.stem, load_note_body(p)) for p in _vault_note_paths(base_dir) + ] + entity_pages = await asyncio.to_thread( + rebuild_entity_pages, base_dir, state.anthropic, notes + ) + state.invalidate_vault_cache() + await asyncio.to_thread( + gitsafe.snapshot, base_dir, "engram: post-rebuild" + ) + snap = "snapshot saved" if committed else "no snapshot (git unavailable or clean)" + await update.effective_chat.send_message( + f"Rebuild done. Merged {merges} duplicate note(s); " + f"wrote {entity_pages} entity page(s). Pre-rebuild {snap}; " + "revert with `git -C <vault> reset --hard HEAD~` if it looks wrong." + ) + return ( on_message, on_start, on_refresh, on_redo, on_folder_choice, on_search, on_ask, on_undo, on_edit, on_inbox, on_review, on_review_choice, - on_relink, + on_relink, on_rebuild, ) +def _vault_note_paths(base_dir: Path) -> list[Path]: + """All capture notes, excluding attachments and the entity (wiki) folders.""" + skip = {"attachments", *ENTITY_TYPE_FOLDERS.values()} + out: list[Path] = [] + for p in base_dir.rglob("*.md"): + if any(part in skip for part in p.relative_to(base_dir).parts): + continue + out.append(p) + # Deterministic order so rebuilds and their git diffs are reproducible. + return sorted(out) + + def main() -> None: config = load_config() log_path = Path(os.environ.get("LOG_FILE", DEFAULT_LOG_PATH)).expanduser() @@ -1251,7 +1338,7 @@ def main() -> None: on_message, on_start, on_refresh, on_redo, on_folder_choice, on_search, on_ask, on_undo, on_edit, on_inbox, on_review, on_review_choice, - on_relink, + on_relink, on_rebuild, ) = make_handlers(state) app.add_handler(CommandHandler("start", on_start)) @@ -1264,6 +1351,7 @@ def main() -> None: app.add_handler(CommandHandler("inbox", on_inbox)) app.add_handler(CommandHandler("review", on_review)) app.add_handler(CommandHandler("relink", on_relink)) + app.add_handler(CommandHandler("rebuild", on_rebuild)) app.add_handler(CallbackQueryHandler(on_folder_choice, pattern=r"^f\|")) app.add_handler(CallbackQueryHandler(on_review_choice, pattern=r"^r\|")) app.add_handler( diff --git a/src/engram/config.py b/src/engram/config.py index 5baa17b..16c8622 100644 --- a/src/engram/config.py +++ b/src/engram/config.py @@ -20,9 +20,18 @@ SUMMARY_MODEL = "claude-sonnet-4-6" REDO_MODEL = "claude-opus-4-7" VISION_MODEL = "claude-haiku-4-5" +MERGE_MODEL = "claude-sonnet-4-6" +ENTITY_MODEL = "claude-sonnet-4-6" EMBEDDING_MODEL = "text-embedding-3-small" EMBEDDING_DIM = 1536 +# Entity (wiki backbone) pages live in typed top-level folders. +ENTITY_TYPE_FOLDERS: dict[str, str] = { + "person": "People", + "concept": "Concepts", + "project": "Projects", +} + @dataclass class Config: diff --git a/src/engram/entities.py b/src/engram/entities.py new file mode 100644 index 0000000..7f75e58 --- /dev/null +++ b/src/engram/entities.py @@ -0,0 +1,262 @@ +from __future__ import annotations + +import json +import logging +import re +from collections import OrderedDict +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +from anthropic import Anthropic + +from .config import ENTITY_MODEL, ENTITY_TYPE_FOLDERS + +log = logging.getLogger(__name__) + +MAX_ENTITIES = 6 +_BAD_CHARS = re.compile(r'[\\/:*?"<>|#\[\]]') +_WS = re.compile(r"\s+") + +_EXTRACT_SYSTEM = ( + "You extract the salient named entities from a single note for a personal wiki. " + "Return ONLY a JSON array. Each element is an object with:\n" + '- "name": the canonical name (a person, a concept, or a project/company/product)\n' + '- "type": exactly one of "person", "concept", "project"\n' + '- "observation": ONE concrete sentence about this entity grounded in THIS note\n' + "Rules: include at most 6 entities. Only entities that are genuinely central to " + "the note. Skip generic words, the author's own filler, and anything you would not " + "want a dedicated wiki page for. Use the most common canonical name (e.g. 'Andrej " + 'Karpathy\', not \'Karpathy\' or \'@karpathy\'). Return [] if nothing qualifies.' +) + +_SYNTH_SYSTEM = ( + "You write the lead paragraph of a wiki page about one entity, synthesizing a list " + "of observations gathered from many notes. Write 1-3 sentences, concrete and " + "specific, present tense, no preamble. Do not use bullet points. Do not repeat the " + "entity name as a heading. Output only the paragraph." +) + + +@dataclass +class Entity: + name: str + type: str + observation: str + + +def _safe_name(name: str) -> str: + name = _BAD_CHARS.sub(" ", name) + name = _WS.sub(" ", name).strip() + return name or "entity" + + +def _normalize(name: str) -> str: + return _WS.sub(" ", _BAD_CHARS.sub(" ", name.lower())).strip() + + +def entity_path(base_dir: Path, entity: Entity) -> Path: + folder = ENTITY_TYPE_FOLDERS.get(entity.type, ENTITY_TYPE_FOLDERS["concept"]) + return base_dir / folder / f"{_safe_name(entity.name)}.md" + + +def _parse_entities(text: str) -> list[Entity]: + match = re.search(r"\[.*\]", text, re.DOTALL) + if not match: + return [] + try: + data = json.loads(match.group(0)) + except json.JSONDecodeError: + return [] + out: list[Entity] = [] + for item in data: + if not isinstance(item, dict): + continue + name = item.get("name") + etype = item.get("type") + obs = item.get("observation", "") + if not isinstance(name, str) or not name.strip(): + continue + if etype not in ENTITY_TYPE_FOLDERS: + continue + out.append( + Entity( + name=name.strip(), + type=etype, + observation=obs.strip() if isinstance(obs, str) else "", + ) + ) + if len(out) >= MAX_ENTITIES: + break + return out + + +def extract_entities(client: Anthropic, title: str, body: str) -> list[Entity]: + if not body.strip(): + return [] + user_msg = f"NOTE TITLE: {title}\n\nNOTE BODY:\n{body}\n\nReturn the JSON array." + try: + resp = client.messages.create( + model=ENTITY_MODEL, + max_tokens=1000, + system=[{"type": "text", "text": _EXTRACT_SYSTEM}], + messages=[{"role": "user", "content": user_msg}], + ) + text = "".join( + b.text for b in resp.content if getattr(b, "type", None) == "text" + ) + except Exception: + log.exception("Entity extraction failed for %s", title) + return [] + return _parse_entities(text) + + +_OBS_RE = re.compile(r"^- (?P<obs>.*?)(?: \(\[\[(?P<src>[^\]]+)\]\]\))?$") + + +def _build_page( + name: str, + etype: str, + observations: list[tuple[str, str]], + *, + lead: str = "", + created: str | None = None, +) -> str: + now = datetime.now().isoformat(timespec="seconds") + lines = ["---", f"entity-type: {etype}", f"created: {created or now}", f"updated: {now}", "---", ""] + lines.append(f"# {name}") + lines.append("") + if lead: + lines.append(lead.strip()) + lines.append("") + lines.append("## Observations") + seen_src: "OrderedDict[str, None]" = OrderedDict() + for obs, src in observations: + suffix = f" ([[{src}]])" if src else "" + lines.append(f"- {obs}{suffix}") + if src: + seen_src.setdefault(src, None) + lines.append("") + lines.append("## Mentioned in") + for src in seen_src: + lines.append(f"- [[{src}]]") + lines.append("") + return "\n".join(lines) + + +def _read_observations(text: str) -> list[tuple[str, str]]: + obs: list[tuple[str, str]] = [] + in_section = False + for line in text.splitlines(): + if line.strip() == "## Observations": + in_section = True + continue + if in_section and line.startswith("## "): + break + if in_section and line.startswith("- "): + m = _OBS_RE.match(line.rstrip()) + if m: + obs.append((m.group("obs").strip(), (m.group("src") or "").strip())) + return obs + + +def _read_field(text: str, key: str) -> str | None: + m = re.search(rf"^{re.escape(key)}:\s*(.+)$", text, re.MULTILINE) + return m.group(1).strip() if m else None + + +def upsert_entity_page(base_dir: Path, entity: Entity, source_title: str) -> Path: + """Create or grow the wiki page for `entity`, adding this source's observation.""" + path = entity_path(base_dir, entity) + path.parent.mkdir(parents=True, exist_ok=True) + new_obs = (entity.observation, source_title) + if path.exists(): + existing = path.read_text(encoding="utf-8") + obs = _read_observations(existing) + if new_obs not in obs and entity.observation: + obs.append(new_obs) + created = _read_field(existing, "created") + lead = _extract_lead(existing) + path.write_text( + _build_page(entity.name, entity.type, obs, lead=lead, created=created), + encoding="utf-8", + ) + else: + obs = [new_obs] if entity.observation else [] + path.write_text( + _build_page(entity.name, entity.type, obs), encoding="utf-8" + ) + return path + + +def _extract_lead(text: str) -> str: + """The prose between the `# Title` heading and the first `##` section.""" + lines = text.splitlines() + out: list[str] = [] + started = False + for line in lines: + if line.startswith("# ") and not started: + started = True + continue + if started: + if line.startswith("## "): + break + out.append(line) + return "\n".join(out).strip() + + +def _synthesize_lead(client: Anthropic, name: str, observations: list[str]) -> str: + if not observations: + return "" + joined = "\n".join(f"- {o}" for o in observations) + try: + resp = client.messages.create( + model=ENTITY_MODEL, + max_tokens=400, + system=[{"type": "text", "text": _SYNTH_SYSTEM}], + messages=[{"role": "user", "content": f"ENTITY: {name}\n\nOBSERVATIONS:\n{joined}"}], + ) + return "".join( + b.text for b in resp.content if getattr(b, "type", None) == "text" + ).strip() + except Exception: + log.exception("Lead synthesis failed for %s", name) + return "" + + +def rebuild_entity_pages( + base_dir: Path, + client: Anthropic, + notes: list[tuple[str, str]], +) -> int: + """Extract entities across all notes and rebuild typed pages from scratch. + + `notes` is a list of (title, body). Returns the number of entity pages written. + """ + # group_key -> (display_name, type, [(observation, source_title)]) + groups: "OrderedDict[tuple[str, str], tuple[str, str, list[tuple[str, str]]]]" = OrderedDict() + for title, body in notes: + for ent in extract_entities(client, title, body): + key = (ent.type, _normalize(ent.name)) + if key not in groups: + groups[key] = (ent.name, ent.type, []) + if ent.observation: + groups[key][2].append((ent.observation, title)) + + # Clear existing typed folders so the rebuild is authoritative. + for folder in set(ENTITY_TYPE_FOLDERS.values()): + d = base_dir / folder + if d.is_dir(): + for f in d.glob("*.md"): + f.unlink() + + count = 0 + for name, etype, obs in groups.values(): + lead = _synthesize_lead(client, name, [o for o, _ in obs]) + path = base_dir / ENTITY_TYPE_FOLDERS[etype] / f"{_safe_name(name)}.md" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + _build_page(name, etype, obs, lead=lead), encoding="utf-8" + ) + count += 1 + return count diff --git a/src/engram/gitsafe.py b/src/engram/gitsafe.py new file mode 100644 index 0000000..090b8ba --- /dev/null +++ b/src/engram/gitsafe.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import logging +import subprocess +from pathlib import Path + +log = logging.getLogger(__name__) + +_IDENTITY = ("-c", "user.name=Engram", "-c", "user.email=engram@local") + + +def _git(base_dir: Path, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["git", "-C", str(base_dir), *args], + capture_output=True, + text=True, + ) + + +def is_repo(base_dir: Path) -> bool: + try: + r = _git(base_dir, "rev-parse", "--is-inside-work-tree") + except FileNotFoundError: + return False + return r.returncode == 0 and r.stdout.strip() == "true" + + +def ensure_repo(base_dir: Path) -> bool: + """Make `base_dir` a git repo if it isn't already. Returns True if git is usable.""" + if not base_dir.is_dir(): + return False + try: + if is_repo(base_dir): + return True + r = _git(base_dir, "init") + except FileNotFoundError: + log.warning("git not available; vault snapshots disabled") + return False + if r.returncode != 0: + log.warning("git init failed in %s: %s", base_dir, r.stderr.strip()) + return False + return True + + +def snapshot(base_dir: Path, message: str) -> bool: + """Commit the current vault state. Returns True only when a commit was created. + + Safe no-op (returns False) when git is unavailable or there is nothing to commit. + Never raises — a failed snapshot must not block a capture. + """ + if not ensure_repo(base_dir): + return False + try: + add = _git(base_dir, "add", "-A") + if add.returncode != 0: + log.warning("git add failed in %s: %s", base_dir, add.stderr.strip()) + return False + # Nothing staged → nothing to commit. + if _git(base_dir, "diff", "--cached", "--quiet").returncode == 0: + return False + commit = _git(base_dir, *_IDENTITY, "commit", "-m", message) + except FileNotFoundError: + return False + if commit.returncode != 0: + log.warning("git commit failed in %s: %s", base_dir, commit.stderr.strip()) + return False + return True diff --git a/src/engram/merger.py b/src/engram/merger.py new file mode 100644 index 0000000..26ad97f --- /dev/null +++ b/src/engram/merger.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import logging + +from anthropic import Anthropic + +from .config import MERGE_MODEL + +log = logging.getLogger(__name__) + +_MERGE_SYSTEM = ( + "You maintain a personal Obsidian wiki. You are given an EXISTING note and a NEW " + "source that is about the same topic. Rewrite the note so it integrates the new " + "source into a single coherent page. Output the complete updated note in Markdown " + "and NOTHING else.\n\n" + "Rules:\n" + "- Keep the YAML frontmatter block (the part between the opening and closing --- " + "lines) at the very top. You may add new entries but never delete existing keys.\n" + "- Integrate the new information into the existing prose. Remove redundancy so the " + "note becomes sharper, not just longer. Do NOT keep dated '## Update' sections.\n" + "- When the new source contradicts the existing note, keep the newer claim and add " + "a short inline note of the change with its date, e.g. '(was X as of May 11; " + "updated June 3)'.\n" + "- Preserve every wikilink [[...]] and every embed ![[...]] exactly as written.\n" + "- Keep the trailing '*Related: ...*' line and the '#tags' line if present.\n" + "- Do not wrap the output in code fences." +) + + +def _strip_fences(text: str) -> str: + t = text.strip() + if t.startswith("```"): + lines = t.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].strip() == "```": + lines = lines[:-1] + t = "\n".join(lines) + return t.strip() + + +def merge_note( + client: Anthropic, + existing_text: str, + new_text: str, + *, + new_source_date: str | None = None, +) -> str: + """Return the full rewritten note that integrates `new_text` into `existing_text`. + + Returns "" on any failure so the caller can fall back to a safe append. + """ + if not existing_text.strip(): + return "" + date_clause = f" (captured {new_source_date})" if new_source_date else "" + user_msg = ( + f"EXISTING NOTE:\n{existing_text}\n\n" + f"NEW SOURCE{date_clause}:\n{new_text}\n\n" + "Return the complete merged note in Markdown." + ) + try: + resp = client.messages.create( + model=MERGE_MODEL, + max_tokens=4000, + system=[{"type": "text", "text": _MERGE_SYSTEM}], + messages=[{"role": "user", "content": user_msg}], + ) + text = "".join( + b.text for b in resp.content if getattr(b, "type", None) == "text" + ) + except Exception: + log.exception("Merge rewrite failed") + return "" + return _strip_fences(text) diff --git a/src/engram/note_writer.py b/src/engram/note_writer.py index 1d528f1..0c42abd 100644 --- a/src/engram/note_writer.py +++ b/src/engram/note_writer.py @@ -1,13 +1,23 @@ from __future__ import annotations +import logging import re from dataclasses import dataclass, field from datetime import datetime from pathlib import Path +from anthropic import Anthropic + +from .merger import merge_note +from .vault import FRONTMATTER_RE + +log = logging.getLogger(__name__) + URL_RE = re.compile(r"https?://\S+") SLUG_BAD_CHARS = re.compile(r'[\\/:*?"<>|#\[\]]') WHITESPACE = re.compile(r"\s+") +EMBED_RE = re.compile(r"!\[\[[^\]]+\]\]") +RELATED_BLOCK_RE = re.compile(r"\n---\n\*Related: .*?\*[ \t]*\n", re.DOTALL) @dataclass @@ -167,3 +177,65 @@ def append_to_note(path: Path, msg: CapturedMessage) -> Path: new_text = existing + "\n" + build_update_section(msg) path.write_text(new_text, encoding="utf-8") return path + + +def _embeds(text: str) -> set[str]: + return set(EMBED_RE.findall(text)) + + +def is_safe_merge(existing: str, merged: str) -> bool: + """The LLM rewrite must not drop frontmatter or any attachment embed.""" + if not merged.strip(): + return False + if FRONTMATTER_RE.match(existing) and not FRONTMATTER_RE.match(merged): + return False + if not _embeds(existing).issubset(_embeds(merged)): + return False + return True + + +def _insert_before_related(text: str, block: str) -> str: + """Insert `block` ahead of a trailing `*Related: ...*` section, else at the end.""" + m = RELATED_BLOCK_RE.search(text) + if m is not None: + return text[: m.start()] + "\n" + block.rstrip() + "\n" + text[m.start():] + return text.rstrip() + "\n" + block.rstrip() + "\n" + + +def _new_media_block(merged: str, msg: CapturedMessage) -> str: + lines: list[str] = [] + for url in msg.source_urls: + if url and url not in merged: + lines.append(f"> Source: {url}") + for embed in [*msg.images, *msg.attachments]: + token = f"![[{embed}]]" + if token not in merged: + lines.append(token) + return "\n".join(lines) + + +def merge_into_note( + path: Path, msg: CapturedMessage, client: Anthropic +) -> tuple[Path, bool]: + """Rewrite an existing note to integrate `msg`, the Karpathy-wiki way. + + Returns (path, merged). When the LLM rewrite is unsafe (dropped frontmatter or + an attachment), falls back to the mechanical append so a capture is never lost, + and returns merged=False. + """ + existing = path.read_text(encoding="utf-8") + merged = merge_note( + client, + existing, + msg.text or "", + new_source_date=msg.created.strftime("%Y-%m-%d"), + ) + if not is_safe_merge(existing, merged): + log.info("Merge unsafe for %s; falling back to append", path.name) + append_to_note(path, msg) + return path, False + block = _new_media_block(merged, msg) + if block: + merged = _insert_before_related(merged, block) + path.write_text(merged.rstrip() + "\n", encoding="utf-8") + return path, True diff --git a/src/engram/retro.py b/src/engram/retro.py new file mode 100644 index 0000000..7a946e3 --- /dev/null +++ b/src/engram/retro.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import logging +from pathlib import Path + +from anthropic import Anthropic + +from .dedupe import find_semantic_duplicate +from .embeddings import SemanticIndex +from .merger import merge_note +from .note_writer import is_safe_merge +from .vault import load_note_body + +log = logging.getLogger(__name__) + + +def merge_duplicate_notes( + base_dir: Path, + semantic_index: SemanticIndex | None, + client: Anthropic, + paths: list[Path], +) -> int: + """Collapse semantically-duplicate notes across the vault into canonical pages. + + For each note, find its best duplicate among the others (cosine + LLM judge, + same thresholds as live capture). When found, merge this note's body into the + canonical match and delete this note. Returns the number of merges performed. + + Destructive: callers must `gitsafe.snapshot` first. + """ + if semantic_index is None or not semantic_index.enabled: + return 0 + deleted: set[Path] = set() + merges = 0 + for note in paths: + if note in deleted or not note.exists(): + continue + body = load_note_body(note) + if not body.strip(): + continue + target = find_semantic_duplicate( + semantic_index, + client, + note.stem, + body, + skip_paths=deleted | {note}, + ) + if target is None or target in deleted or not target.exists(): + continue + # Keep the older note as canonical so its creation date and title + # survive; merge the newer one into it and delete the newer. + if note.stat().st_mtime <= target.stat().st_mtime: + canonical, victim = note, target + else: + canonical, victim = target, note + existing = canonical.read_text(encoding="utf-8") + merged = merge_note(client, existing, load_note_body(victim)) + if not is_safe_merge(existing, merged): + log.info( + "Retro merge unsafe: %s -> %s; skipping", victim.name, canonical.name + ) + continue + canonical.write_text(merged.rstrip() + "\n", encoding="utf-8") + victim.unlink() + deleted.add(victim) + merges += 1 + log.info("Retro merged %s into %s", victim.name, canonical.name) + return merges diff --git a/tests/test_bot_rebuild.py b/tests/test_bot_rebuild.py new file mode 100644 index 0000000..6bdeb14 --- /dev/null +++ b/tests/test_bot_rebuild.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from engram import bot as bot_module +from engram.vault import VaultIndex +from tests.test_bot import ( + _fake_callback_update, + _fake_message, + _fake_update, + _make_state, + _patch_enrich, +) + + +def _set_text(state, text: str) -> None: + state.anthropic.messages.create.return_value = SimpleNamespace( + content=[SimpleNamespace(type="text", text=text)] + ) + + +def _set_sequence(state, texts: list[str]) -> None: + state.anthropic.messages.create.side_effect = [ + SimpleNamespace(content=[SimpleNamespace(type="text", text=t)]) for t in texts + ] + + +async def test_duplicate_branch_rewrites_when_merge_safe(tmp_path, monkeypatch): + state = _make_state(tmp_path, monkeypatch) + _patch_enrich(monkeypatch, folder="AI", title="X", summary="body") + ai = tmp_path / "AI" + ai.mkdir() + existing = ai / "Canonical.md" + existing.write_text("---\nc: x\n---\n# Canonical\nold body\n", encoding="utf-8") + state._vault_index = VaultIndex() + monkeypatch.setattr(bot_module, "find_semantic_duplicate", lambda *a, **k: existing) + # The LLM returns a valid rewrite that preserves frontmatter → safe merge. + _set_text(state, "---\nc: x\n---\n# Canonical\nmerged new body\n") + + update = _fake_update(_fake_message(text="some related thing")) + handlers = bot_module.make_handlers(state) + on_message, on_folder_choice = handlers[0], handlers[4] + await on_message(update, SimpleNamespace(application=None)) + token = next(iter(state._pending)) + ai_idx = state.config.categories.index("AI") + await on_folder_choice( + _fake_callback_update(f"f|{token}|{ai_idx}"), SimpleNamespace() + ) + + text = existing.read_text(encoding="utf-8") + assert "merged new body" in text + assert "## Update" not in text # rewrite, not append + assert list(ai.glob("*.md")) == [existing] # no duplicate created + + +async def test_rebuild_builds_entity_pages(tmp_path, monkeypatch): + state = _make_state(tmp_path, monkeypatch) + state._semantic_index.embedder = None # disable → skip merge step + ai = tmp_path / "AI" + ai.mkdir() + (ai / "a.md").write_text("---\nc: x\n---\n# a\nabout karpathy\n", encoding="utf-8") + (ai / "b.md").write_text("---\nc: y\n---\n# b\nnothing notable\n", encoding="utf-8") + # extract(a) -> Karpathy; extract(b) -> []; then one synth lead. + _set_sequence( + state, + [ + '[{"name":"Andrej Karpathy","type":"person","observation":"works on AI"}]', + "[]", + "Andrej Karpathy is an AI researcher.", + ], + ) + + on_rebuild = bot_module.make_handlers(state)[13] + update = _fake_update(_fake_message(text="/rebuild")) + await on_rebuild(update, SimpleNamespace()) + + page = tmp_path / "People" / "Andrej Karpathy.md" + assert page.exists() + assert "works on AI ([[a]])" in page.read_text(encoding="utf-8") + last = update.effective_chat.send_message.call_args.args[0] + assert "Rebuild done" in last diff --git a/tests/test_entities.py b/tests/test_entities.py new file mode 100644 index 0000000..a80cf90 --- /dev/null +++ b/tests/test_entities.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +from engram import entities +from engram.entities import ( + Entity, + entity_path, + extract_entities, + rebuild_entity_pages, + upsert_entity_page, +) + + +def _client(text: str) -> MagicMock: + client = MagicMock() + client.messages.create.return_value = SimpleNamespace( + content=[SimpleNamespace(type="text", text=text)] + ) + return client + + +def _client_seq(texts: list[str]) -> MagicMock: + client = MagicMock() + client.messages.create.side_effect = [ + SimpleNamespace(content=[SimpleNamespace(type="text", text=t)]) for t in texts + ] + return client + + +def test_extract_parses_typed_entities() -> None: + payload = ( + '[{"name":"Andrej Karpathy","type":"person","observation":"argues for X"},' + '{"name":"Bitter Lesson","type":"concept","observation":"scaling wins"},' + '{"name":"junk","type":"animal","observation":"skip"}]' + ) + ents = extract_entities(_client(payload), "Title", "body") + assert [(e.name, e.type) for e in ents] == [ + ("Andrej Karpathy", "person"), + ("Bitter Lesson", "concept"), + ] + + +def test_extract_empty_body_skips_api() -> None: + client = MagicMock() + assert extract_entities(client, "t", " ") == [] + client.messages.create.assert_not_called() + + +def test_entity_path_uses_typed_folder(tmp_path: Path) -> None: + p = entity_path(tmp_path, Entity("Andrej Karpathy", "person", "x")) + assert p == tmp_path / "People" / "Andrej Karpathy.md" + c = entity_path(tmp_path, Entity("Bitter Lesson", "concept", "x")) + assert c == tmp_path / "Concepts" / "Bitter Lesson.md" + + +def test_upsert_creates_page(tmp_path: Path) -> None: + ent = Entity("Andrej Karpathy", "person", "coined software 2.0") + p = upsert_entity_page(tmp_path, ent, "Source Note") + text = p.read_text(encoding="utf-8") + assert "entity-type: person" in text + assert "# Andrej Karpathy" in text + assert "- coined software 2.0 ([[Source Note]])" in text + assert "- [[Source Note]]" in text + + +def test_upsert_accumulates_across_sources(tmp_path: Path) -> None: + upsert_entity_page(tmp_path, Entity("Karpathy", "person", "obs one"), "Note A") + p = upsert_entity_page( + tmp_path, Entity("Karpathy", "person", "obs two"), "Note B" + ) + text = p.read_text(encoding="utf-8") + assert "- obs one ([[Note A]])" in text + assert "- obs two ([[Note B]])" in text + assert text.count("## Observations") == 1 + assert "- [[Note A]]" in text and "- [[Note B]]" in text + + +def test_upsert_is_idempotent_for_same_source(tmp_path: Path) -> None: + upsert_entity_page(tmp_path, Entity("Karpathy", "person", "same obs"), "Note A") + p = upsert_entity_page( + tmp_path, Entity("Karpathy", "person", "same obs"), "Note A" + ) + text = p.read_text(encoding="utf-8") + assert text.count("- same obs ([[Note A]])") == 1 + + +def test_rebuild_groups_entities_and_writes_pages(tmp_path: Path) -> None: + # 2 extraction calls (one per note), then synthesis calls per unique entity. + client = _client_seq( + [ + '[{"name":"Karpathy","type":"person","observation":"obs from note 1"}]', + '[{"name":"karpathy","type":"person","observation":"obs from note 2"},' + '{"name":"RAG","type":"concept","observation":"retrieval"}]', + "Andrej Karpathy is an AI researcher.", # synth lead for person + "RAG combines retrieval and generation.", # synth lead for concept + ] + ) + notes = [("Note 1", "body one"), ("Note 2", "body two")] + + count = rebuild_entity_pages(tmp_path, client, notes) + + assert count == 2 # Karpathy (merged) + RAG + person = (tmp_path / "People" / "Karpathy.md").read_text(encoding="utf-8") + assert "obs from note 1 ([[Note 1]])" in person + assert "obs from note 2 ([[Note 2]])" in person # case-insensitive merge + assert "Andrej Karpathy is an AI researcher." in person + assert (tmp_path / "Concepts" / "RAG.md").exists() + + +def test_rebuild_clears_stale_pages(tmp_path: Path) -> None: + stale = tmp_path / "People" / "Old Person.md" + stale.parent.mkdir(parents=True) + stale.write_text("---\nentity-type: person\n---\n# Old Person\n", encoding="utf-8") + client = _client_seq( + ['[{"name":"New","type":"person","observation":"o"}]', "New is a person."] + ) + + rebuild_entity_pages(tmp_path, client, [("N", "b")]) + + assert not stale.exists() + assert (tmp_path / "People" / "New.md").exists() diff --git a/tests/test_gitsafe.py b/tests/test_gitsafe.py new file mode 100644 index 0000000..72da3fd --- /dev/null +++ b/tests/test_gitsafe.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import subprocess +from pathlib import Path + +from engram import gitsafe + + +def _log(base: Path) -> list[str]: + r = subprocess.run( + ["git", "-C", str(base), "log", "--pretty=%s"], + capture_output=True, text=True, + ) + return [line for line in r.stdout.splitlines() if line] + + +def test_snapshot_inits_repo_and_commits(tmp_path: Path) -> None: + (tmp_path / "note.md").write_text("hello\n", encoding="utf-8") + assert not gitsafe.is_repo(tmp_path) + + created = gitsafe.snapshot(tmp_path, "engram: first") + + assert created is True + assert gitsafe.is_repo(tmp_path) + assert _log(tmp_path) == ["engram: first"] + + +def test_snapshot_noop_when_nothing_changed(tmp_path: Path) -> None: + (tmp_path / "note.md").write_text("hello\n", encoding="utf-8") + assert gitsafe.snapshot(tmp_path, "engram: first") is True + + # No file changes since the last commit. + assert gitsafe.snapshot(tmp_path, "engram: second") is False + assert _log(tmp_path) == ["engram: first"] + + +def test_snapshot_commits_subsequent_change(tmp_path: Path) -> None: + note = tmp_path / "note.md" + note.write_text("hello\n", encoding="utf-8") + gitsafe.snapshot(tmp_path, "engram: first") + + note.write_text("hello world\n", encoding="utf-8") + created = gitsafe.snapshot(tmp_path, "engram: edit") + + assert created is True + assert _log(tmp_path) == ["engram: edit", "engram: first"] + + +def test_snapshot_returns_false_for_missing_dir(tmp_path: Path) -> None: + assert gitsafe.snapshot(tmp_path / "nope", "x") is False diff --git a/tests/test_merge_into_note.py b/tests/test_merge_into_note.py new file mode 100644 index 0000000..555168e --- /dev/null +++ b/tests/test_merge_into_note.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from datetime import datetime +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +from engram.note_writer import CapturedMessage, merge_into_note + + +def _client(text: str) -> MagicMock: + client = MagicMock() + client.messages.create.return_value = SimpleNamespace( + content=[SimpleNamespace(type="text", text=text)] + ) + return client + + +def _msg(**kw) -> CapturedMessage: + base = dict(text="new info", created=datetime(2026, 6, 3, 9, 0)) + base.update(kw) + return CapturedMessage(**base) + + +def test_merge_rewrites_note_in_place(tmp_path: Path) -> None: + note = tmp_path / "n.md" + note.write_text("---\ncreated: x\n---\n# T\nold body\n", encoding="utf-8") + merged_doc = "---\ncreated: x\n---\n# T\nintegrated body with new info\n" + client = _client(merged_doc) + + path, merged = merge_into_note(note, _msg(), client) + + assert merged is True + assert "integrated body with new info" in note.read_text(encoding="utf-8") + assert "## Update" not in note.read_text(encoding="utf-8") + + +def test_falls_back_to_append_when_frontmatter_dropped(tmp_path: Path) -> None: + note = tmp_path / "n.md" + note.write_text("---\ncreated: x\n---\n# T\nold body\n", encoding="utf-8") + # LLM returned a doc that lost the frontmatter → unsafe. + client = _client("# T\nintegrated but no frontmatter\n") + + path, merged = merge_into_note(note, _msg(), client) + + assert merged is False + body = note.read_text(encoding="utf-8") + assert body.startswith("---\ncreated: x\n---") # original preserved + assert "## Update" in body # append fallback fired + + +def test_falls_back_when_attachment_embed_dropped(tmp_path: Path) -> None: + note = tmp_path / "n.md" + note.write_text( + "---\nc: x\n---\n# T\nbody\n![[attachments/pic.jpg]]\n", encoding="utf-8" + ) + client = _client("---\nc: x\n---\n# T\nrewrote and dropped the image\n") + + _, merged = merge_into_note(note, _msg(), client) + + assert merged is False + assert "![[attachments/pic.jpg]]" in note.read_text(encoding="utf-8") + + +def test_new_media_and_url_appended_after_merge(tmp_path: Path) -> None: + note = tmp_path / "n.md" + note.write_text("---\nc: x\n---\n# T\nold\n", encoding="utf-8") + client = _client("---\nc: x\n---\n# T\nmerged prose\n") + + _, merged = merge_into_note( + note, + _msg(source_urls=["https://new.example/post"], images=["attachments/new.jpg"]), + client, + ) + + body = note.read_text(encoding="utf-8") + assert merged is True + assert "https://new.example/post" in body + assert "![[attachments/new.jpg]]" in body + + +def test_new_media_inserted_before_related_block(tmp_path: Path) -> None: + note = tmp_path / "n.md" + note.write_text("---\nc: x\n---\n# T\nold\n", encoding="utf-8") + merged_doc = ( + "---\nc: x\n---\n# T\nmerged prose\n\n---\n*Related: [[Other]]*\n\n#ai\n" + ) + client = _client(merged_doc) + + _, merged = merge_into_note( + note, _msg(images=["attachments/new.jpg"]), client + ) + + body = note.read_text(encoding="utf-8") + assert body.index("![[attachments/new.jpg]]") < body.index("*Related:") diff --git a/tests/test_merger.py b/tests/test_merger.py new file mode 100644 index 0000000..84ab557 --- /dev/null +++ b/tests/test_merger.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from engram.merger import merge_note + + +def _client(text: str) -> MagicMock: + client = MagicMock() + client.messages.create.return_value = SimpleNamespace( + content=[SimpleNamespace(type="text", text=text)] + ) + return client + + +def test_returns_merged_text() -> None: + merged = "---\ncreated: x\n---\n# T\nintegrated body\n" + client = _client(merged) + out = merge_note(client, "---\ncreated: x\n---\n# T\nold body\n", "new info") + assert out == merged.strip() + client.messages.create.assert_called_once() + + +def test_strips_code_fences() -> None: + client = _client("```markdown\n---\na: b\n---\n# T\nbody\n```") + out = merge_note(client, "---\na: b\n---\n# T\nold\n", "new") + assert out.startswith("---") + assert "```" not in out + + +def test_empty_existing_returns_empty_without_calling_api() -> None: + client = MagicMock() + assert merge_note(client, " ", "new") == "" + client.messages.create.assert_not_called() + + +def test_api_error_returns_empty() -> None: + client = MagicMock() + client.messages.create.side_effect = RuntimeError("boom") + assert merge_note(client, "existing body", "new") == "" diff --git a/tests/test_retro.py b/tests/test_retro.py new file mode 100644 index 0000000..5e6e1c2 --- /dev/null +++ b/tests/test_retro.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +from engram.embeddings import NearestHit +from engram.retro import merge_duplicate_notes + + +class FakeIndex: + """nearest() returns the configured hit for the given query note path.""" + + def __init__(self, mapping: dict[str, list[tuple[Path, float]]]): + # mapping: query note stem -> ranked (path, score) + self._mapping = mapping + self.enabled = True + + def nearest(self, query: str, k: int = 5): + first_line = query.splitlines()[0] if query else "" + for stem, ranked in self._mapping.items(): + if stem in query: + return [NearestHit(path=p, score=s) for p, s in ranked][:k] + return [] + + +def _client(merged_body: str) -> MagicMock: + client = MagicMock() + client.messages.create.return_value = SimpleNamespace( + content=[SimpleNamespace(type="text", text=merged_body)] + ) + return client + + +def test_disabled_index_is_noop(tmp_path: Path) -> None: + class Off: + enabled = False + + assert merge_duplicate_notes(tmp_path, Off(), MagicMock(), []) == 0 + + +def test_merges_duplicate_keeps_older_as_canonical(tmp_path: Path) -> None: + import os + + canonical = tmp_path / "AI" / "Bitter Lesson.md" + dup = tmp_path / "AI" / "scaling laws.md" + canonical.parent.mkdir(parents=True) + canonical.write_text("---\nc: x\n---\n# Bitter Lesson\ncanonical body\n", encoding="utf-8") + dup.write_text("---\nc: y\n---\n# scaling laws\nduplicate body\n", encoding="utf-8") + # Make canonical clearly older so it survives the merge. + os.utime(canonical, (1000, 1000)) + os.utime(dup, (2000, 2000)) + + idx = FakeIndex({"scaling laws": [(canonical, 0.95)], "Bitter Lesson": [(dup, 0.95)]}) + client = _client("---\nc: x\n---\n# Bitter Lesson\nmerged canonical + duplicate\n") + + merges = merge_duplicate_notes(tmp_path, idx, client, [canonical, dup]) + + assert merges == 1 + assert canonical.exists() and not dup.exists() + assert "merged canonical + duplicate" in canonical.read_text(encoding="utf-8") + + +def test_unsafe_merge_keeps_both(tmp_path: Path) -> None: + a = tmp_path / "AI" / "a.md" + b = tmp_path / "AI" / "b.md" + a.parent.mkdir(parents=True) + a.write_text("---\nc: x\n---\n# a\nbody a\n", encoding="utf-8") + b.write_text("---\nc: y\n---\n# b\nbody b\n", encoding="utf-8") + + idx = FakeIndex({"a": [(b, 0.95)], "b": [(a, 0.95)]}) + # Returned merge dropped frontmatter → unsafe → skip, nothing deleted. + client = _client("# a\nno frontmatter\n") + + merges = merge_duplicate_notes(tmp_path, idx, client, [a, b]) + + assert merges == 0 + assert a.exists() and b.exists() + + +def test_no_duplicate_is_noop(tmp_path: Path) -> None: + a = tmp_path / "AI" / "a.md" + a.parent.mkdir(parents=True) + a.write_text("---\nc: x\n---\n# a\nbody a\n", encoding="utf-8") + idx = FakeIndex({"a": [(a, 0.3)]}) # only itself, low score + + assert merge_duplicate_notes(tmp_path, idx, MagicMock(), [a]) == 0 + assert a.exists()