llm-ingestion-okf/src/llm_ingestion_okf/bundlemap.py

"""The map of a bundle: its documents and their titles, in its own words (v1.1 C5).

WHY IT EXISTS. The ranking matches words, and a question put in words the
collection does not use finds little however good the ranking is -- a question
asked in one language of a collection written in another most of all. The
reader closes that gap by rewriting the question into two to four
sub-questions in the collection's OWN words, and the one place those words
are listed is the collection itself. This module lists them, compactly enough
to be read before the first question: one line per source document, its name
and then the titles of its concepts in document order.

A SERIES IS ONE LINE. Documents whose names differ only in their numbers -- a
changelog per release, a note per week -- are one kind of document, and four
hundred lines saying so crowd out everything else a reader needs. They are
written as one line: the name with every number as `#`, how many documents,
the first and the last by natural order, and the titles across the series that
are words (a title that is only a version number names nothing).

DERIVED, NEVER STORED, like the card that carries it: the map is recomputed
from the bundle on every call, so it cannot disagree with the bytes beside it.
Deterministic: every order is by name, by position or by a count with the
name breaking ties.
"""

from __future__ import annotations

import re
from collections import Counter
from collections.abc import Sequence
from pathlib import Path

from .consume import (
    Concept,
    enumerate_concepts,
    inherit_table_titles,
    link_parents,
    read_concept,
    read_path_in_bundle,
    root_bundle_id_of,
)
from .profiles import BundleProfile

#: How many documents sharing one name template make a series. Below it the
#: documents are listed one by one: two or three dated notes are still worth
#: their own lines, and a template shared by chance should not hide them.
SERIES_MIN = 5

#: The most titles one line lists before it says how many it left out. A
#: document is a handful of sections as a rule; a few are hundreds, and one
#: of those must not cost the whole map its room.
TITLES_PER_LINE = 24

#: The most bytes the map's lines take, together. A client keeps a tool reply
#: of 25 000 tokens (Claude Code's MCP output limit); at a pessimistic two
#: bytes a token that is 50 000 bytes, and the rest of the card needs a few
#: thousand. The largest bundle this was measured on stays under it, so the
#: ceiling is a guard for a larger one. Lines past it are counted in
#: `lines_truncated`, never dropped silently.
MAP_MAX_BYTES = 48_000

_DIGITS = re.compile(r"\d+")
_SPLIT = re.compile(r"(\d+)")
_LETTER = re.compile(r"[^\W\d_]")

#: The locators a concept's place in its document is read off, one per
#: document and never mixed (`consume.inherit_table_titles` reads the same).
_POSITION_KEYS = ("source_offset", "source_lines")
_FIRST_NUMBER = re.compile(r"\s*\[\s*(\d+)")


def _stem(source_file: str) -> str:
    return source_file.removesuffix(".md")


def _natural(name: str) -> tuple[tuple[int, str], ...]:
    """Numbers compared as numbers: `v1-2` before `v1-13`."""
    return tuple(
        (int(part), "") if part.isdigit() else (-1, part) for part in _SPLIT.split(name) if part
    )


def _position(concept: Concept, key: str) -> int | None:
    match = _FIRST_NUMBER.match(concept.locators.get(key, ""))
    return int(match.group(1)) if match else None


def _in_document_order(concepts: Sequence[Concept]) -> list[Concept]:
    for key in _POSITION_KEYS:
        positions = [_position(concept, key) for concept in concepts]
        if all(position is not None for position in positions):
            return [
                concept
                for _, _, concept in sorted(
                    (position, index, concept)
                    for index, (position, concept) in enumerate(
                        zip(positions, concepts, strict=True)
                    )
                )
            ]
    return list(concepts)


def _titled(titles: Sequence[str]) -> str:
    kept = titles[:TITLES_PER_LINE]
    text = " · ".join(kept)
    if len(titles) > len(kept):
        text += f" · (+{len(titles) - len(kept)} more)"
    return text


def build_map(concepts: Sequence[Concept]) -> dict[str, object]:
    """The map of `concepts`: one line per document, one per series."""
    by_document: dict[str, list[Concept]] = {}
    for concept in concepts:
        by_document.setdefault(_stem(concept.source_file), []).append(concept)
    by_template: dict[str, list[str]] = {}
    for document in by_document:
        by_template.setdefault(_DIGITS.sub("#", document), []).append(document)

    entries: list[tuple[str, str]] = []
    for template, documents in by_template.items():
        if len(documents) >= SERIES_MIN:
            ordered = sorted(documents, key=_natural)
            counts: Counter[str] = Counter(
                title
                for document in documents
                for title in dict.fromkeys(concept.title for concept in by_document[document])
                if _LETTER.search(title)
            )
            titles = sorted(counts, key=lambda title: (-counts[title], title))
            line = f"{template} ({len(documents)} documents: {ordered[0]} … {ordered[-1]})"
            if titles:
                line += f": {_titled(titles)}"
            entries.append((template, line))
            continue
        for document in documents:
            titles = list(
                dict.fromkeys(
                    concept.title for concept in _in_document_order(by_document[document])
                )
            )
            name = document or "(no source file)"
            entries.append((document, f"{name}: {_titled(titles)}"))
    lines = [line for _, line in sorted(entries, key=lambda entry: (_natural(entry[0]), entry[0]))]
    kept: list[str] = []
    spent = 0
    for line in lines:
        size = len(line.encode("utf-8"))
        if spent + size > MAP_MAX_BYTES:
            break
        kept.append(line)
        spent += size
    return {
        "documents": len(by_document),
        "concepts": len(concepts),
        "lines_count": len(lines),
        "lines_truncated": len(lines) - len(kept),
        "lines": kept,
    }


def read_concepts(bundle_root: Path, *, profile: BundleProfile) -> list[Concept]:
    """Every concept of the bundle, as `okf consume` reads them -- parents
    linked and a table fragment named by the heading above it."""
    bundle_id = root_bundle_id_of(bundle_root, profile=profile)
    return inherit_table_titles(
        link_parents(
            [
                read_concept(
                    read_path_in_bundle(bundle_root, f"{concept_id}{profile.paths.concept_suffix}"),
                    bundle_root=bundle_root,
                    root_bundle_id=bundle_id,
                )
                for concept_id in enumerate_concepts(bundle_root, profile=profile)
            ]
        )
    )


def bundle_map(bundle_root: Path, *, profile: BundleProfile) -> dict[str, object]:
    return build_map(read_concepts(bundle_root, profile=profile))