markdown-memory 0.1.0__py3-none-any.whl

markdown_memory/model_cache.py

@@ -0,0 +1,272 @@
+"""The on-disk model cache: where weights live, and how we know they are ours.
+
+The embedder that uses this is in ``embedders.py``; the cache is its own module because
+the pin, the manifest and the verification stamp are consulted from outside it too -
+``scripts/eval_cache.py`` keys the retrieval gate on them.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import fcntl
+import hashlib
+import json
+import os
+import shutil
+import stat
+from collections.abc import Iterator, Mapping
+from pathlib import Path
+
+from markdown_memory.exceptions import ModelLoadError
+
+BGE_SMALL_MODEL_NAME = "BAAI/bge-small-en-v1.5"
+GEMMA_REPOSITORY = "onnx-community/embeddinggemma-300m-ONNX"
+
+
+# Pinned so that an upstream re-export can never silently change stored vectors.
+GEMMA_REVISION = "5090578d9565bb06545b4552f76e6bc2c93e4a66"
+
+
+# The 4-bit graph, which quantizes the vocabulary table with `GatherBlockQuantized` and the
+# projections with `MatMulNBits`. Earlier versions ran `onnx/model_quantized.onnx` (int8)
+# and rewrote it on each machine to gather the vocabulary before dequantizing it; this
+# graph does that natively, in half the download and half the CPU per query. The file name
+# is part of the embedder's `model_name`, so changing it discards every stored vector -
+# which is correct, because the two graphs' vectors are not comparable.
+GEMMA_MODEL_FILE = "onnx/model_q4.onnx"
+
+
+# Size and sha256 of every file at GEMMA_REVISION, from the Hub's paths-info API. This is
+# what stands between a damaged cache and onnxruntime: huggingface_hub checks only the
+# size of what it downloads, and hands back a file that is already on disk without reading
+# it at all.
+GEMMA_MANIFEST: Mapping[str, tuple[int, str]] = {
+    GEMMA_MODEL_FILE: (
+        519_322,
+        "ad1dfee81a70f7944b9b9d1cc6e48075b832881cf33fab2f2b248be78f3f0043",
+    ),
+    GEMMA_MODEL_FILE + "_data": (
+        196_725_760,
+        "599962c3143b040de2dd05e5975be3e9091dd067cacc6a8f7186e3203bab9e02",
+    ),
+    "tokenizer.json": (
+        20_323_312,
+        "4dda02faaf32bc91031dc8c88457ac272b00c1016cc679757d1c441b248b9c47",
+    ),
+}
+
+
+GEMMA_FILES: tuple[str, ...] = tuple(GEMMA_MANIFEST)
+
+
+# The directory carries the revision, so moving the pin fetches the new weights instead of
+# serving the old ones under a name that claims to be the new ones. The graph file is not
+# in the path: a stamp recording one graph's files cannot match another's manifest, so a
+# folder holding the wrong graph is rejected and refetched rather than trusted.
+_GEMMA_DIR_PREFIX = "embeddinggemma-300m-onnx"
+
+
+# One lock for every revision, so two versions starting at once still exclude each other.
+_GEMMA_LOCK_NAME = f"{_GEMMA_DIR_PREFIX}.lock"
+
+
+_VERIFIED_STAMP = ".verified"
+
+
+def model_cache_root(cache_dir: Path | None = None) -> Path:
+    """Where every model this package downloads is kept."""
+    return cache_dir or Path.home() / ".cache" / "markdown-memory" / "models"
+
+
+def gemma_model_dir(cache_dir: Path | None = None) -> Path:
+    """The folder holding the pinned EmbeddingGemma revision."""
+    return model_cache_root(cache_dir) / f"{_GEMMA_DIR_PREFIX}-{GEMMA_REVISION[:12]}"
+
+
+def fastembed_model_dir(cache_dir: Path | None) -> Path | None:
+    """The folder fastembed keeps bge-small in, or None when it cannot be derived.
+
+    Not guessable from the model name: fastembed downloads its own re-export of the
+    model (`qdrant/bge-small-en-v1.5-onnx-q`), so the repository is read out of its
+    registry rather than assumed.
+    """
+    if cache_dir is None:
+        return None
+    from fastembed import TextEmbedding
+
+    for entry in TextEmbedding.list_supported_models():
+        if not isinstance(entry, dict) or entry.get("model") != BGE_SMALL_MODEL_NAME:
+            continue
+        sources = entry.get("sources")
+        repository = sources.get("hf") if isinstance(sources, dict) else None
+        if isinstance(repository, str) and repository:
+            return cache_dir / ("models--" + repository.replace("/", "--"))
+    return None
+
+
+def _cache_path(model_dir: Path, name: str) -> Path | None:
+    """``model_dir/name``, or None when any directory on the way there is a symlink.
+
+    `_file_identity` only ever looked at the last component, so an `onnx` that pointed
+    somewhere else was trusted - and then repaired, which meant deleting and overwriting
+    files outside the cache entirely.
+    """
+    current = model_dir
+    for part in Path(name).parts:
+        if current.is_symlink():
+            return None
+        current = current / part
+    return current
+
+
+def _file_identity(path: Path) -> dict[str, int] | None:
+    """What a stamp remembers about one model file; None when it is not a plain file.
+
+    `ctime_ns` earns its place: `cp -p`, `tar x` and `rsync --inplace` all rewrite a
+    file's contents and then restore its old mtime, so size, mtime and inode can agree
+    across different bytes. Nothing in user space can set ctime back.
+    """
+    try:
+        info = path.lstat()
+    except OSError:
+        return None
+    if not stat.S_ISREG(info.st_mode):
+        return None  # a symlink into a blob store is not a file this cache vouches for
+    return {
+        "size": info.st_size,
+        "mtime_ns": info.st_mtime_ns,
+        "ctime_ns": info.st_ctime_ns,
+        "inode": info.st_ino,
+        "device": info.st_dev,
+    }
+
+
+def _identity_in(model_dir: Path, name: str) -> dict[str, int] | None:
+    """The identity of one cached file, refusing a path that leaves the cache."""
+    path = _cache_path(model_dir, name)
+    return None if path is None else _file_identity(path)
+
+
+def _stamp_is_current(model_dir: Path) -> bool:
+    """Whether every file still looks exactly as it did when it was last verified.
+
+    Hashing 218 MB costs most of a second of one core, which is too much for every
+    server start when an editor starts one per session. This is a handful of `stat`
+    calls; anything that disagrees sends the files back to be hashed.
+    """
+    try:
+        stamp = json.loads((model_dir / _VERIFIED_STAMP).read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return False
+    if not isinstance(stamp, dict) or stamp.get("revision") != GEMMA_REVISION:
+        return False
+    recorded = stamp.get("files")
+    if not isinstance(recorded, dict) or set(recorded) != set(GEMMA_FILES):
+        # Also how a folder left behind by a different graph is rejected: its stamp names
+        # files this manifest does not, so it is refetched rather than half-trusted.
+        return False
+    if model_dir.is_symlink():
+        return False
+    for name in recorded:
+        identity = _identity_in(model_dir, name)
+        # `None` means the path is not a regular file - a symlink, most likely. Nothing
+        # should be able to stamp one (`_unverified` calls it wrong), and this is the
+        # line that makes a stamp that somehow recorded `None` stop matching `None`.
+        if identity is None or recorded[name] != identity:
+            return False
+    return True
+
+
+def _write_stamp(model_dir: Path) -> None:
+    """Record what was just verified. Atomically: a half-written stamp is a false claim."""
+    stamp = {
+        "revision": GEMMA_REVISION,
+        "files": {name: _file_identity(model_dir / name) for name in GEMMA_FILES},
+    }
+    temporary = model_dir / f"{_VERIFIED_STAMP}.{os.getpid()}"
+    _remove(temporary)  # a stamp left behind by a crash under this same pid
+    _write_new_file(temporary, json.dumps(stamp).encode("utf-8"))
+    stamped = model_dir / _VERIFIED_STAMP
+    if stamped.is_dir() and not stamped.is_symlink():
+        # `os.replace` will not put a file where a directory is, and the stamp is read
+        # back as simply invalid, so the repair it ends would run again on every start
+        # and end the same way. Nothing of ours is ever a directory here.
+        _remove(stamped)
+    os.replace(temporary, stamped)
+
+
+def _write_new_file(path: Path, payload: bytes) -> None:
+    """Create ``path`` with its contents, refusing to follow a symlink or reuse a file.
+
+    The temporary names these writes use are predictable (the pid), and a plain write
+    follows a symlink planted at one of them: the caller would then overwrite whatever it
+    points at, anywhere the user can write. `O_EXCL | O_NOFOLLOW` makes both refusals the
+    kernel's.
+    """
+    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
+    try:
+        with os.fdopen(descriptor, "wb") as handle:
+            handle.write(payload)
+    except BaseException:
+        with contextlib.suppress(OSError):
+            path.unlink()
+        raise
+
+
+def _remove(path: Path) -> None:
+    """Delete whatever sits at ``path``, file or directory.
+
+    A directory where a model file belongs is not something `unlink` can clear, and a
+    cache that cannot be repaired is a server that never starts again.
+    """
+    if path.is_dir() and not path.is_symlink():
+        shutil.rmtree(path, ignore_errors=True)
+    else:
+        path.unlink(missing_ok=True)
+    if path.exists() or path.is_symlink():
+        # Say so here, where the path is known, rather than failing three lines later on
+        # a rename that cannot explain itself.
+        raise ModelLoadError(f"Cannot clear {path} to repair the model cache")
+
+
+def _hash_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _unverified(model_dir: Path) -> list[str]:
+    """The model files that are missing, the wrong size, or the wrong bytes."""
+    problems: list[str] = []
+    for name, (size, checksum) in GEMMA_MANIFEST.items():
+        path = _cache_path(model_dir, name)
+        before = None if path is None else _file_identity(path)
+        if path is None or before is None or before["size"] != size:
+            problems.append(name)
+        elif _hash_file(path) != checksum or _file_identity(path) != before:
+            # Second identity: a file rewritten while it was being read was never hashed
+            # as it now stands, so the answer that came back means nothing.
+            problems.append(name)
+    return problems
+
+
+@contextlib.contextmanager
+def _model_cache_lock(cache_dir: Path | None, *, exclusive: bool) -> Iterator[None]:
+    """Serialise verification, repair and session construction across processes.
+
+    Shared while a verified cache is being opened, exclusive while it is being changed,
+    so nothing can repair files another process has verified but not yet handed to
+    onnxruntime. The kernel drops a `flock` when the process dies, so a crash leaves
+    nothing held. A model cache on NFS or SMB shared between machines is out of scope:
+    `flock` can be local-only there - the boundary SQLite's WAL already has.
+    """
+    root = model_cache_root(cache_dir)
+    root.mkdir(parents=True, exist_ok=True)
+    with (root / _GEMMA_LOCK_NAME).open("a") as handle:
+        fcntl.flock(handle, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
+        try:
+            yield
+        finally:
+            fcntl.flock(handle, fcntl.LOCK_UN)

markdown_memory/models.py

@@ -0,0 +1,339 @@
+"""Immutable, fully typed domain models.
+
+Every entity is a frozen, slotted dataclass. ``to_dict`` methods produce the
+JSON-serialisable payloads returned by the MCP tools.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TypeAlias
+
+from pydantic import JsonValue
+
+JsonDict: TypeAlias = dict[str, JsonValue]
+
+PREAMBLE_TITLE = "[Overview / Preamble]"
+PATH_SEPARATOR = " > "
+CHARS_PER_TOKEN = 4
+
+
+def estimate_tokens(text: str) -> int:
+    """Cheap token estimate (~4 characters per token), never below 1 for non-empty text."""
+    if not text:
+        return 0
+    return max(1, -(-len(text) // CHARS_PER_TOKEN))
+
+
+def part_path(base_path: str, part_index: int) -> str:
+    """Breadcrumb for one part of an oversized section, e.g. ``A > B (Part 2)``."""
+    return f"{base_path} (Part {part_index})"
+
+
+@dataclass(slots=True, frozen=True)
+class SectionDraft:
+    """A section produced by the parser, not yet persisted.
+
+    ``part_index`` is ``0`` for a section stored whole and ``1..n`` for the
+    sequential parts of an oversized section. ``base_path`` is the breadcrumb
+    without the ``(Part n)`` suffix. Line numbers are 1-based and inclusive.
+
+    ``units`` are the section's passages as plain text - one per paragraph, list item,
+    table row or code block. Each is embedded separately, because one vector for a
+    whole section dilutes a single relevant table row beyond recognition. A section
+    without units is heading-only and gets no vectors at all.
+    """
+
+    heading_title: str
+    heading_level: int
+    heading_path: str
+    base_path: str
+    content: str
+    start_line: int
+    end_line: int
+    part_index: int = 0
+    units: tuple[str, ...] = ()
+
+    @property
+    def embedding_text(self) -> str:
+        """Text for the section-level vector: breadcrumb, then the body without markup."""
+        return f"{self.heading_path}\n\n{' '.join(self.units)}"
+
+    @property
+    def unit_texts(self) -> tuple[str, ...]:
+        """Texts for the passage-level vectors, each carrying the breadcrumb for context."""
+        return tuple(f"{self.heading_path}: {unit}" for unit in self.units)
+
+
+@dataclass(slots=True, frozen=True)
+class SectionVectors:
+    """Embeddings of one section: ``section`` is ``None`` for a heading-only section."""
+
+    section: Sequence[float] | None
+    units: tuple[Sequence[float], ...] = ()
+
+
+@dataclass(slots=True, frozen=True)
+class ParsedDocument:
+    """Result of parsing one Markdown source."""
+
+    title: str
+    sections: tuple[SectionDraft, ...]
+    line_count: int
+
+
+@dataclass(slots=True, frozen=True)
+class Document:
+    """A persisted Markdown file."""
+
+    id: int
+    file_path: str
+    title: str
+    content_hash: str
+    last_modified: int
+
+
+@dataclass(slots=True, frozen=True)
+class DocumentSummary:
+    """A document together with the number of sections indexed for it."""
+
+    file_path: str
+    title: str
+    section_count: int
+    last_modified: int
+
+    def to_dict(self) -> JsonDict:
+        return {
+            "file_path": self.file_path,
+            "title": self.title,
+            "section_count": self.section_count,
+            "last_modified": self.last_modified,
+        }
+
+
+@dataclass(slots=True, frozen=True)
+class Section:
+    """A persisted section (or one part of an oversized section)."""
+
+    id: int
+    doc_id: int
+    heading_title: str
+    heading_level: int
+    heading_path: str
+    content: str
+    start_line: int
+    end_line: int
+    part_index: int
+
+    @property
+    def base_path(self) -> str:
+        """Breadcrumb with any ``(Part n)`` suffix removed."""
+        if self.part_index <= 0:
+            return self.heading_path
+        suffix = part_path("", self.part_index)
+        return self.heading_path.removesuffix(suffix)
+
+
+@dataclass(slots=True, frozen=True)
+class OutlineNode:
+    """One heading in a document's hierarchical table of contents."""
+
+    heading_title: str
+    heading_level: int
+    heading_path: str
+    start_line: int
+    end_line: int
+    token_estimate: int
+    part_count: int
+    children: tuple[OutlineNode, ...] = ()
+
+    def to_dict(self) -> JsonDict:
+        payload: JsonDict = {
+            "title": self.heading_title,
+            "level": self.heading_level,
+            "heading_path": self.heading_path,
+            "lines": f"{self.start_line}-{self.end_line}",
+            "tokens": self.token_estimate,
+        }
+        if self.part_count > 1:
+            payload["parts"] = self.part_count
+        if self.children:
+            payload["children"] = [child.to_dict() for child in self.children]
+        return payload
+
+
+@dataclass(slots=True, frozen=True)
+class SearchResult:
+    """A section matched by hybrid search, with its fused and per-index ranks."""
+
+    section_id: int
+    file_path: str
+    document_title: str
+    heading_title: str
+    heading_path: str
+    content: str
+    start_line: int
+    end_line: int
+    score: float
+    fts_rank: int | None
+    vec_rank: int | None
+    matched_passage: str | None = None
+
+    def to_dict(self) -> JsonDict:
+        payload: JsonDict = {
+            "file_path": self.file_path,
+            "document_title": self.document_title,
+            "heading_path": self.heading_path,
+            "heading_title": self.heading_title,
+            "lines": f"{self.start_line}-{self.end_line}",
+            "score": round(self.score, 6),
+            "fts_rank": self.fts_rank,
+            "vec_rank": self.vec_rank,
+            "tokens": estimate_tokens(self.content),
+            "content": self.content,
+        }
+        if self.matched_passage is not None:
+            payload["matched_passage"] = self.matched_passage
+        return payload
+
+
+@dataclass(slots=True, frozen=True)
+class FileFailure:
+    """A per-file failure recorded while indexing (the run itself continues)."""
+
+    file_path: str
+    message: str
+
+    def to_dict(self) -> JsonDict:
+        return {"file_path": self.file_path, "message": self.message}
+
+
+#: Failures carried on a search answer. Enough to act on, few enough not to bury the
+#: answer itself - the count in `message` says how many were left out.
+MAX_REPORTED_FAILURES = 20
+
+
+@dataclass(slots=True, frozen=True)
+class IndexStatus:
+    """Whether answers drawn from one tree can be trusted to be drawn from all of it.
+
+    `verified` means a full walk of this scope finished and read every file it found. It
+    is deliberately not a claim that the filesystem has stopped changing: a file created
+    after its directory was walked is not in the index and not in `failures`, and the next
+    run picks it up. A killed run and a tree nobody ever indexed both read unverified,
+    which is the same answer because it is the same situation - nothing walked it whole.
+    """
+
+    verified: bool
+    failures: tuple[FileFailure, ...] = ()
+    #: Set when the weights behind the model name changed under an existing index: the
+    #: stored vectors and the vectors a query would produce now come from different
+    #: models. Nothing is discarded, and nothing new is written, until it is resolved.
+    weights_mismatch: str | None = None
+    #: Documents under this scope whose vectors were built by an older pooling scheme.
+    #: They still answer, less well, and only a run over the directory holding them
+    #: rebuilds - a parent run prunes `.venv`, `node_modules` and the like, so one indexed
+    #: deliberately inside such a directory is never reached again.
+    stale_vectors: int = 0
+    #: Indexed documents a cheap probe could not confirm are still what was indexed: their
+    #: bytes differ, or they are gone, unreadable, or no longer a regular file. It is
+    #: best-effort in both directions. Counted over the rows the index holds, so a file
+    #: nobody has indexed yet is not in it - finding those needs the directory walk, which
+    #: is the expensive half. And the bytes are only read where the modification time moved,
+    #: so an edit that restores a file's own timestamp is not seen. Zero means nothing was
+    #: detected, not that every indexed file was hashed.
+    changed_files: int = 0
+    #: This server's own background run is indexing the tree right now. A hint about one
+    #: process only: another process's run shows as coverage withdrawn, as it always did.
+    indexing: bool = False
+
+    def to_dict(self) -> JsonDict:
+        shown = self.failures[:MAX_REPORTED_FAILURES]
+        return {
+            "coverage": "verified" if self.verified else "unknown",
+            "failures": [failure.to_dict() for failure in shown],
+            "changed_files": self.changed_files,
+            "indexing": self.indexing,
+            "message": self.message(),
+        }
+
+    def message(self) -> str | None:
+        """One sentence, or nothing at all when there is nothing to act on."""
+        # First, because it is the only one that says the answers themselves may be
+        # wrong rather than incomplete. Hoisted above `verified` rather than left below
+        # it: the two cannot both hold today, and a reader should not have to know that.
+        if self.weights_mismatch:
+            return self.weights_mismatch
+        if self.indexing:
+            # Before everything that ends in "run index_directory": that run would be
+            # refused as busy, and it would be refused for doing what is already being done.
+            return (
+                "An automatic index run is in progress, so an answer may be missing a file "
+                "changed or added since the last one finished; there is no need to run "
+                "index_directory."
+            )
+        if self.verified:
+            # A walk that finished still describes the moment it finished. Files edited
+            # since are the one thing a verified tree has left to say.
+            if self.changed_files:
+                return (
+                    f"The last full index completed, but {self.changed_files} indexed "
+                    "document(s) can no longer be confirmed to be what was indexed - "
+                    "changed, unreadable or gone - so an answer may quote text that is no "
+                    "longer there; run index_directory to refresh. The check is cheap and "
+                    "best-effort: files created since that scan are not counted, and an "
+                    "edit that puts a file's modification time back is not seen."
+                )
+            return None
+        if not self.failures:
+            if self.stale_vectors:
+                return (
+                    f"{self.stale_vectors} document(s) here were indexed by an older "
+                    "vector format and rank less well until the directory holding them is "
+                    "indexed again."
+                )
+            return (
+                "This documentation root has not been indexed end to end since it last "
+                "changed, so an answer may be missing part of it. Run index_directory."
+            )
+        hidden = len(self.failures) - MAX_REPORTED_FAILURES
+        more = f" (showing the first {MAX_REPORTED_FAILURES})" if hidden > 0 else ""
+        return (
+            f"{len(self.failures)} path(s) could not be indexed{more}; answers here are "
+            "drawn from a tree that is missing them."
+        )
+
+
+@dataclass(slots=True, frozen=True)
+class IndexReport:
+    """Outcome of one ``index_directory`` run."""
+
+    directory: str
+    files_scanned: int
+    files_indexed: int
+    files_unchanged: int
+    files_purged: int
+    sections_indexed: int
+    elapsed_seconds: float
+    passages_indexed: int = 0
+    errors: tuple[FileFailure, ...] = ()
+    notes: tuple[str, ...] = ()
+
+    def summary(self) -> str:
+        lines = [
+            f"Indexed {self.directory} in {self.elapsed_seconds:.2f}s: "
+            f"{self.files_scanned} scanned, {self.files_indexed} (re)indexed, "
+            f"{self.files_unchanged} unchanged, {self.files_purged} purged, "
+            f"{self.sections_indexed} sections embedded ({self.passages_indexed} passages)."
+        ]
+        lines.extend(f"NOTE {note}" for note in self.notes)
+        lines.extend(f"ERROR {error.file_path}: {error.message}" for error in self.errors)
+        if self.errors:
+            # Without this the run reads as a success with some noise attached, and an
+            # index missing part of its tree answers questions as if it were whole.
+            lines.append(
+                f"INCOMPLETE: {len(self.errors)} file(s) could not be indexed; "
+                "this documentation root is only partly searchable."
+            )
+        return "\n".join(lines)