curu-memory 0.2.0__py3-none-any.whl

curu_memory/indexing.py

@@ -0,0 +1,440 @@
+"""The index use case: read the concept files, report, and build the index.
+
+`build_index` is the only path from the canonical files to the index (constitution,
+article I). It writes nothing until each file is read and each vector is made, so a
+failure keeps the old index.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from curu_memory.errors import EmbedderFailed, IndexRefused
+from curu_memory.media import MediaFile, read_media
+from curu_memory.okf import SUMMARY, ConceptFault, ParsedConcept, parse_concept, split_units
+from curu_memory.ports import (
+    GLOBAL,
+    ConceptMedia,
+    Embedder,
+    IndexMeta,
+    IndexStore,
+    Link,
+    Prompts,
+    Unit,
+    parse_space,
+)
+from curu_memory.store_sqlite import SqliteEmbeddingCache
+
+BATCH_SIZE = 32
+#: Fewer media files than texts go in one request, because a file can be 10 MB.
+MEDIA_BATCH_SIZE = 4
+
+
+@dataclass(frozen=True)
+class _PendingMedia:
+    """One accepted image of a concept, before the embedder makes its vector."""
+
+    concept_id: str
+    concept_path: str
+    #: The path as the concept wrote it, for a warning.
+    written: str
+    #: The place of the image in the `media:` list, from 1.
+    number: int
+    slug: str
+    ordinal: int
+    media: MediaFile
+
+
+@dataclass(frozen=True)
+class IndexReport:
+    files_read: int
+    concepts_indexed: int
+    #: Each (path, reason) of a file that is not in the index.
+    skipped: tuple[tuple[str, str], ...]
+    #: Each (path, reason) of a problem that did not stop the file.
+    warnings: tuple[tuple[str, str], ...]
+    #: Each (source id, target id, kind) of a link to a missing concept.
+    broken_links: tuple[tuple[str, str, str], ...]
+    #: The model identifier of the embedder, or "none".
+    embedder: str
+    #: The media units in the new index: images and audio.
+    media_embedded: int = 0
+    #: The media references that are not in the index. Each one is also a warning.
+    media_skipped: int = 0
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "files_read": self.files_read,
+            "concepts_indexed": self.concepts_indexed,
+            "skipped": [{"path": p, "reason": r} for p, r in self.skipped],
+            "warnings": [{"path": p, "reason": r} for p, r in self.warnings],
+            "broken_links": [
+                {"source_id": s, "target_id": t, "kind": k} for s, t, k in self.broken_links
+            ],
+            "embedder": self.embedder,
+            "media_embedded": self.media_embedded,
+            "media_skipped": self.media_skipped,
+        }
+
+    def render(self) -> str:
+        lines = [
+            f"files read: {self.files_read}",
+            f"concepts indexed: {self.concepts_indexed}",
+            f"skipped: {len(self.skipped)}",
+            f"warnings: {len(self.warnings)}",
+            f"broken links: {len(self.broken_links)}",
+            f"embedder: {self.embedder}",
+            f"media embedded: {self.media_embedded}",
+            f"media skipped: {self.media_skipped}",
+        ]
+        if self.skipped:
+            lines += ["", "Skipped:"] + [f"  {p}: {r}" for p, r in self.skipped]
+        if self.warnings:
+            lines += ["", "Warnings:"] + [f"  {p}: {r}" for p, r in self.warnings]
+        if self.broken_links:
+            lines += ["", "Broken links:"] + [
+                f"  {s} -> {t} ({k})" for s, t, k in self.broken_links
+            ]
+        return "\n".join(lines)
+
+
+def concept_files(root: Path) -> list[tuple[str, Path]]:
+    """Each (relative path, path) of a concept file, sorted by the relative path.
+
+    A file in a hidden directory, for example `.git` or `.obsidian`, is not a concept.
+    A file with the name `README.md` is not a concept.
+    """
+    found = []
+    for path in root.rglob("*.md"):
+        relative = path.relative_to(root)
+        if any(part.startswith(".") for part in relative.parts):
+            continue
+        if path.name.lower() == "readme.md" or not path.is_file():
+            continue
+        found.append((relative.as_posix(), path))
+    return sorted(found)
+
+
+def build_index(
+    concepts_dir: Path,
+    store: IndexStore,
+    embedder: Embedder | None = None,
+    *,
+    cache: SqliteEmbeddingCache | None = None,
+    tier_space: str = GLOBAL,
+    prompts: Prompts | None = None,
+) -> IndexReport:
+    """Index one directory into one space, under the lock of the data directory (ADR-016).
+
+    Raise `ValueError` for a bad space name, `IndexLocked` when another run holds the
+    lock and `IndexRefused` when the embedder model is not the model of another space.
+    """
+    tier_space = parse_space(tier_space)
+    with store.index_lock():
+        others = [space for space in store.spaces() if space != tier_space]
+        _check_model(store, others, embedder)
+        taken = {cid: space for space in others for cid in store.concept_ids(space)}
+        return _build(concepts_dir, store, embedder, cache, tier_space, prompts, taken)
+
+
+def _check_model(store: IndexStore, others: Sequence[str], embedder: Embedder | None) -> None:
+    """All spaces use one embedder model, so their vectors are comparable (ADR-016)."""
+    if embedder is None:
+        return
+    for space in others:
+        meta = store.meta(space)
+        model = meta.embedder_model if meta is not None else None
+        if model is not None and model != embedder.model:
+            raise IndexRefused(
+                f"The space {space} was indexed with the embedder {model!r}, but the embedder "
+                f"now is {embedder.model!r}. All spaces must use one embedder model. "
+                "Index each space again with one model."
+            )
+
+
+def _build(
+    concepts_dir: Path,
+    store: IndexStore,
+    embedder: Embedder | None,
+    cache: SqliteEmbeddingCache | None,
+    tier_space: str,
+    prompts: Prompts | None,
+    taken: dict[str, str],
+) -> IndexReport:
+    files = concept_files(concepts_dir)
+    skipped: list[tuple[str, str]] = []
+    warnings: list[tuple[str, str]] = []
+    parsed: list[ParsedConcept] = []
+    for relative, path in files:
+        try:
+            item = parse_concept(relative, path.read_text(encoding="utf-8"), tier_space)
+        except ConceptFault as fault:
+            reason = f"line {fault.line}: {fault.reason}" if fault.line else fault.reason
+            skipped.append((relative, reason))
+            continue
+        except UnicodeDecodeError:
+            skipped.append((relative, "the file is not UTF-8 text"))
+            continue
+        parsed.append(item)
+        warnings.extend((relative, w) for w in item.warnings)
+
+    by_id: dict[str, list[ParsedConcept]] = defaultdict(list)
+    for item in parsed:
+        by_id[item.concept.id].append(item)
+    kept: list[ParsedConcept] = []
+    for identifier, items in sorted(by_id.items()):
+        if len(items) == 1:
+            kept.append(items[0])
+            continue
+        paths = [i.concept.path for i in items]
+        for item in items:
+            others = ", ".join(p for p in paths if p != item.concept.path)
+            skipped.append(
+                (
+                    item.concept.path,
+                    f"the identifier {identifier!r} is also used by {others}, "
+                    "so neither file is indexed. Rename one of them.",
+                )
+            )
+
+    # An identifier is unique in the instance, so each link has one target (ADR-016).
+    unique: list[ParsedConcept] = []
+    for item in kept:
+        space = taken.get(item.concept.id)
+        if space is None:
+            unique.append(item)
+            continue
+        skipped.append(
+            (
+                item.concept.path,
+                f"the identifier {item.concept.id!r} is in the space {space}. "
+                "An identifier must be unique in the instance, so rename one of them.",
+            )
+        )
+    kept = unique
+
+    # A link to a concept of another space is not broken. A read resolves it again over
+    # the readable spaces of the caller (ADR-016).
+    known = {item.concept.id for item in kept} | set(taken)
+    links = [
+        Link(item.concept.id, target, kind, target in known)
+        for item in kept
+        for target, kind in item.links
+    ]
+    units = [unit for item in kept for unit in _units(item)]
+    titles = {item.concept.id: item.concept.title for item in kept}
+    units, meta = _embed(units, titles, embedder, cache, tier_space, prompts or Prompts())
+
+    paths = dict(files)
+    pending: list[_PendingMedia] = []
+    media_warnings: list[tuple[str, str]] = []
+    last_ordinal: dict[str, int] = defaultdict(int)
+    for unit in units:
+        last_ordinal[unit.concept_id] = max(last_ordinal[unit.concept_id], unit.ordinal)
+    for item in kept:
+        concept = item.concept
+        found, problems = _media(item, paths[concept.path], concepts_dir, last_ordinal[concept.id])
+        pending.extend(found)
+        media_warnings.extend(problems)
+    media_units, failures = _embed_media(pending, embedder, cache, meta.embedder_dimension)
+    media_warnings.extend(failures)
+    warnings.extend(media_warnings)
+    units = units + media_units
+
+    # Each accepted image is recorded with its bytes, with or without a vector (ADR-009).
+    media = [
+        ConceptMedia(p.concept_id, p.number, p.media.path, p.media.sha256, p.media.mime)
+        for p in pending
+    ]
+    blobs = [p.media for p in pending]
+    store.build(meta, [item.concept for item in kept], units, links, media, blobs)
+    return IndexReport(
+        files_read=len(files),
+        concepts_indexed=len(kept),
+        skipped=tuple(sorted(skipped)),
+        warnings=tuple(sorted(warnings)),
+        broken_links=tuple(
+            sorted((lk.source_id, lk.target_id, lk.kind) for lk in links if not lk.resolved)
+        ),
+        embedder=meta.embedder_model or "none",
+        media_embedded=len(media_units),
+        media_skipped=len(media_warnings),
+    )
+
+
+def _units(item: ParsedConcept) -> list[Unit]:
+    concept = item.concept
+    units = [Unit(concept.id, SUMMARY, 0, f"{concept.title}\n\n{concept.description}")]
+    for ordinal, (slug, text) in enumerate(split_units(concept.body), start=1):
+        units.append(Unit(concept.id, slug, ordinal, text))
+    return units
+
+
+def _media(
+    item: ParsedConcept, path: Path, concepts_dir: Path, ordinal_base: int
+) -> tuple[list[_PendingMedia], list[tuple[str, str]]]:
+    """The accepted images of one concept, and a warning for each problem (ADR-007).
+
+    The slug is `image.<n>` or `audio.<n>`, from its place in the `media:` list, so it does
+    not change when an image before it is missing. A section slug has no `.`, so the
+    two never meet. Its ordinal comes after the text units.
+    """
+    concept = item.concept
+    found: list[_PendingMedia] = []
+    problems: list[tuple[str, str]] = []
+    for number, written in enumerate(item.media, start=1):
+        files, refused = read_media(path, concepts_dir, [written])
+        problems.extend((concept.path, f"media {p.path!r}: {p.reason}") for p in refused)
+        for media in files:
+            found.append(
+                _PendingMedia(
+                    concept.id,
+                    concept.path,
+                    written,
+                    number,
+                    f"{media.kind}.{number}",
+                    ordinal_base + number,
+                    media,
+                )
+            )
+    return found, problems
+
+
+def _embed_media(
+    pending: list[_PendingMedia],
+    embedder: Embedder | None,
+    cache: SqliteEmbeddingCache | None,
+    dimension: int | None,
+) -> tuple[list[Unit], list[tuple[str, str]]]:
+    """The image units, and a warning for each image that has no vector.
+
+    The embedder gets each distinct file one time, with no prompt. The cache keys a
+    file by the model and `<kind>:<sha256>` (ADR-011). A failure gives a warning for
+    each file that fails alone, and does not stop the run (FR-006).
+    """
+    if embedder is None:
+        reason = "no embedder is set, so a search by meaning cannot find the {kind}"
+        return [], [
+            (p.concept_path, f"media {p.written!r}: " + reason.format(kind=p.media.kind))
+            for p in pending
+        ]
+    model = embedder.model
+    vectors: dict[str, tuple[float, ...]] = {}
+    failed: dict[str, str] = {}
+    missing: list[MediaFile] = []
+    for media in {p.media.sha256: p.media for p in pending}.values():
+        cached = cache.get(model, f"{media.kind}:{media.sha256}") if cache is not None else None
+        if cached is not None:
+            vectors[media.sha256] = cached
+        else:
+            missing.append(media)
+    made: list[tuple[MediaFile, list[float]]] = []
+    for start in range(0, len(missing), MEDIA_BATCH_SIZE):
+        batch = missing[start : start + MEDIA_BATCH_SIZE]
+        try:
+            made.extend(zip(batch, _call_media(embedder, batch), strict=True))
+            continue
+        except Exception:  # noqa: S110 - each image of the request goes again, alone.
+            pass
+        # One bad image fails its whole request, so each image goes again alone (FR-006).
+        for media in batch:
+            try:
+                made.extend(zip([media], _call_media(embedder, [media]), strict=True))
+            except Exception as error:
+                detail = str(error) or type(error).__name__
+                failed[media.sha256] = f"the embedder did not embed the {media.kind} ({detail})"
+    for media, vector in made:
+        if dimension is not None and len(vector) != dimension:
+            failed[media.sha256] = (
+                f"the embedder gave an image vector of {len(vector)} values, "
+                f"not {dimension} as for text"
+            )
+            continue
+        vectors[media.sha256] = tuple(vector)
+        if cache is not None:
+            cache.put(model, f"{media.kind}:{media.sha256}", vector)
+    units: list[Unit] = []
+    warnings: list[tuple[str, str]] = []
+    for p in pending:
+        sha256 = p.media.sha256
+        if sha256 in failed:
+            warnings.append((p.concept_path, f"media {p.written!r}: {failed[sha256]}"))
+            continue
+        units.append(
+            Unit(
+                p.concept_id,
+                p.slug,
+                p.ordinal,
+                p.media.path,
+                vectors[sha256],
+                kind=p.media.kind,
+                media_path=p.media.path,
+                media_sha256=sha256,
+            )
+        )
+    return units, warnings
+
+
+def _embed(
+    units: list[Unit],
+    titles: dict[str, str],
+    embedder: Embedder | None,
+    cache: SqliteEmbeddingCache | None,
+    tier_space: str,
+    prompts: Prompts,
+) -> tuple[list[Unit], IndexMeta]:
+    count = len(titles)
+    if embedder is None:
+        return units, IndexMeta(tier_space, count, None, None)
+    model = embedder.model
+    # The embedder gets each unit in the document prompt of the model (ADR-003).
+    inputs = {
+        (u.concept_id, u.slug): prompts.for_document(titles[u.concept_id], u.text) for u in units
+    }
+    vectors: dict[str, tuple[float, ...]] = {}
+    missing: list[str] = []
+    for text in dict.fromkeys(inputs.values()):
+        cached = cache.get(model, text) if cache is not None else None
+        if cached is not None:
+            vectors[text] = cached
+        else:
+            missing.append(text)
+    for start in range(0, len(missing), BATCH_SIZE):
+        batch = missing[start : start + BATCH_SIZE]
+        made = _call(embedder, batch)
+        for text, vector in zip(batch, made, strict=True):
+            vectors[text] = tuple(vector)
+            if cache is not None:
+                cache.put(model, text, vector)
+    dimensions = {len(v) for v in vectors.values()}
+    if len(dimensions) > 1:
+        raise EmbedderFailed(f"it gave vectors of different sizes: {sorted(dimensions)}")
+    dimension = dimensions.pop() if dimensions else embedder.dimension
+    embedded = [
+        Unit(u.concept_id, u.slug, u.ordinal, u.text, vectors[inputs[(u.concept_id, u.slug)]])
+        for u in units
+    ]
+    return embedded, IndexMeta(tier_space, count, model, dimension, prompts.fingerprint)
+
+
+def _call_media(embedder: Embedder, files: Sequence[MediaFile]) -> list[list[float]]:
+    made = embedder.embed_media(files)
+    if len(made) != len(files):
+        raise RuntimeError(f"it gave {len(made)} vectors for {len(files)} media files")
+    return made
+
+
+def _call(embedder: Embedder, texts: Sequence[str]) -> list[list[float]]:
+    try:
+        made = embedder.embed(texts)
+    except EmbedderFailed:
+        raise
+    except Exception as error:
+        raise EmbedderFailed(str(error) or type(error).__name__) from error
+    if len(made) != len(texts):
+        raise EmbedderFailed(f"it gave {len(made)} vectors for {len(texts)} texts")
+    return made

curu_memory/media.py

@@ -0,0 +1,199 @@
+"""The rules of the media files of a concept (ADR-007).
+
+A concept references media with a front-matter list `media:` of paths relative to the
+concept file. Each path must resolve inside the concept directory. The format comes from
+the first bytes of the file, not from its name.
+"""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import hashlib
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from curu_memory.errors import InvalidArgument
+
+MAX_BYTES = 10 * 1024 * 1024
+UNKNOWN_FORMAT = "the file is not an image, audio or video file of a known format"
+_SIGNATURES = (
+    (b"\x89PNG\r\n\x1a\n", "image/png"),
+    (b"\xff\xd8\xff", "image/jpeg"),
+    (b"fLaC", "audio/flac"),
+    (b"ID3", "audio/mpeg"),
+)
+
+
+@dataclass(frozen=True)
+class MediaFile:
+    #: The path relative to the concept directory, with `/` as separator.
+    path: str
+    mime: str
+    sha256: str
+    data: bytes = field(repr=False)
+
+    @property
+    def kind(self) -> str:
+        """`image`, `audio` or `video`, from the MIME type."""
+        return self.mime.partition("/")[0]
+
+    def data_uri(self) -> str:
+        return f"data:{self.mime};base64," + base64.b64encode(self.data).decode()
+
+
+@dataclass(frozen=True)
+class MediaProblem:
+    #: The path as the concept wrote it.
+    path: str
+    reason: str
+
+
+def read_media(
+    concept_file: Path, concepts_dir: Path, paths: Sequence[str]
+) -> tuple[list[MediaFile], list[MediaProblem]]:
+    """Read each media path of one concept. Give the accepted files and each problem."""
+    root = concepts_dir.resolve()
+    files: list[MediaFile] = []
+    problems: list[MediaProblem] = []
+    for written in paths:
+        if Path(written).is_absolute():
+            problems.append(MediaProblem(written, "the path is not relative to the concept file"))
+            continue
+        resolved = (concept_file.parent / written).resolve()
+        if not resolved.is_relative_to(root):
+            problems.append(MediaProblem(written, "the path is outside the concept directory"))
+            continue
+        if not resolved.exists():
+            problems.append(MediaProblem(written, "the file does not exist"))
+            continue
+        if not resolved.is_file():
+            problems.append(MediaProblem(written, "the path is not a file"))
+            continue
+        if resolved.stat().st_size > MAX_BYTES:
+            problems.append(MediaProblem(written, "the file is larger than 10 MB"))
+            continue
+        data = resolved.read_bytes()
+        mime = sniff_mime(data[:SNIFF_BYTES])
+        if mime is None:
+            problems.append(MediaProblem(written, UNKNOWN_FORMAT))
+            continue
+        files.append(
+            MediaFile(
+                path=resolved.relative_to(root).as_posix(),
+                mime=mime,
+                sha256=hashlib.sha256(data).hexdigest(),
+                data=data,
+            )
+        )
+    return files, problems
+
+
+#: The codec of the first packet of an Ogg audio file, at byte 28: Vorbis, Opus or FLAC.
+#: A video-only Ogg file, for example Theora, is not audio.
+_OGG_AUDIO = (b"\x01vorbis", b"OpusHead", b"\x7fFLAC")
+#: The bytes that `sniff_mime` reads. The Ogg codec ends at byte 36.
+SNIFF_BYTES = 64
+
+
+def sniff_mime(head: bytes) -> str | None:
+    """The MIME type of a media file from its first bytes, or None."""
+    for signature, mime in _SIGNATURES:
+        if head.startswith(signature):
+            return mime
+    if head[:4] == b"RIFF" and head[8:12] == b"WEBP":
+        return "image/webp"
+    if head[:4] == b"RIFF" and head[8:12] == b"WAVE":
+        return "audio/wav"
+    if head[:4] == b"OggS" and head[28:36].startswith(_OGG_AUDIO):
+        return "audio/ogg"
+    if _is_mp3_frame(head):
+        return "audio/mpeg"
+    return _video_mime(head)
+
+
+#: The ISO media brands of a video. HEIF and AVIF images and M4A audio use `ftyp` too,
+#: so a brand that is not in this set is not a video.
+_MP4_BRANDS = frozenset(
+    {b"isom", b"iso2", b"iso3", b"iso4", b"iso5", b"iso6", b"mp41", b"mp42", b"avc1",
+     b"M4V ", b"dash", b"mmp4", b"3gp4", b"3gp5", b"3gp6", b"3g2a"}
+)  # fmt: skip
+#: The EBML header of a Matroska or WebM file.
+_EBML = b"\x1a\x45\xdf\xa3"
+
+
+def _video_mime(head: bytes) -> str | None:
+    """The MIME type of an MP4, MOV, WebM or MKV video from its first bytes, or None."""
+    if head[4:8] == b"ftyp":
+        brand = head[8:12]
+        if brand == b"qt  ":
+            return "video/quicktime"
+        return "video/mp4" if brand in _MP4_BRANDS else None
+    if head.startswith(_EBML):
+        # The DocType element (0x42 0x82) names the kind of the Matroska file.
+        if b"\x42\x82\x84webm" in head:
+            return "video/webm"
+        if b"\x42\x82\x88matroska" in head:
+            return "video/x-matroska"
+    return None
+
+
+def _is_mp3_frame(head: bytes) -> bool:
+    """True for an MPEG audio frame header of Layer II or III, with no ID3 tag.
+
+    The header has 11 sync bits, a valid version, Layer II or III, a bitrate index from
+    1 to 14 and a valid sampling rate. So text, AAC and other data do not pass.
+    """
+    if len(head) < 3 or head[0] != 0xFF or head[1] & 0xE0 != 0xE0:
+        return False
+    version = (head[1] >> 3) & 0b11
+    layer = (head[1] >> 1) & 0b11
+    bitrate = head[2] >> 4
+    rate = (head[2] >> 2) & 0b11
+    return version != 0b01 and layer in (0b01, 0b10) and bitrate not in (0, 15) and rate != 3
+
+
+#: The most base64 characters that can hold MAX_BYTES, with some room for padding.
+_MAX_BASE64 = (MAX_BYTES + 2) // 3 * 4 + 4
+_TOO_LARGE = "The file is larger than 10 MB."
+
+
+def decode_media_argument(text: str) -> bytes:
+    """The bytes of a media argument: base64 text, or a base64 `data:` URI (ADR-012).
+
+    The length is checked before the decode, so a large argument costs no memory.
+    """
+    value = text.strip()
+    if value.startswith("data:"):
+        header, comma, value = value.partition(",")
+        if not comma or not header.endswith(";base64"):
+            raise _bad_media("The media is a data: URI with no base64 bytes.")
+    # base64 tools break the text into lines of 76 characters, with up to 2 more each.
+    if len(value) > _MAX_BASE64 + _MAX_BASE64 // 38 + 2:
+        raise _bad_media(_TOO_LARGE)
+    value = "".join(value.split())
+    if len(value) > _MAX_BASE64:
+        raise _bad_media(_TOO_LARGE)
+    try:
+        data = base64.b64decode(value, validate=True)
+    except (binascii.Error, ValueError):
+        # A character that is not ASCII gives ValueError, not binascii.Error.
+        data = b""
+    if not data:
+        raise _bad_media("The media is not base64 text or a data: URI.")
+    return data
+
+
+def read_query_media(data: bytes) -> MediaFile:
+    """A media question from its bytes: an image, a clip or a video, as `read_media` takes."""
+    if len(data) > MAX_BYTES:
+        raise _bad_media(_TOO_LARGE)
+    mime = sniff_mime(data[:SNIFF_BYTES])
+    if mime is None:
+        raise _bad_media(f"{UNKNOWN_FORMAT.capitalize()}.")
+    return MediaFile("question", mime, hashlib.sha256(data).hexdigest(), data)
+
+
+def _bad_media(message: str) -> InvalidArgument:
+    return InvalidArgument("media", message)