docsgraph 0.1.0a2__py3-none-any.whl

cairn/entity/heuristic.py

@@ -0,0 +1,148 @@
+"""HeuristicExtractor — regex-based, no model needed.
+
+Covers two of the four ``EntityKind`` values:
+
+- **code**: identifiers inside fenced code blocks and inline ``` `code` ```
+  spans. Filters out common language keywords (Python, JS) and identifiers
+  shorter than three characters.
+- **defined**: text inside ``**bold**`` markdown markers that looks like a
+  term (no sentence-level punctuation, ≤ 80 chars).
+
+Span coordinates are offsets within the *section's* ``raw_text``. The
+:class:`cairn.index.entities.EntityBuilder` preserves this convention end
+to end.
+
+LLM-based extraction for ``term`` and ``proper`` kinds is the v0.2.1
+follow-up; this v0.2.0 extractor is fully offline.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable, Iterator
+from typing import Final
+
+from cairn.core.types import Document, Span
+from cairn.entity.base import ExtractionHit
+
+_FENCED_CODE = re.compile(r"```[^\n]*\n(.*?)\n```", re.DOTALL)
+_INLINE_CODE = re.compile(r"`([^`\n]+)`")
+_BOLD = re.compile(r"\*\*([^*\n]+)\*\*")
+_IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]{2,}")
+
+# Common stopwords that show up in code blocks but carry no entity meaning.
+# Conservative: only drops obvious language keywords + built-in names.
+_CODE_STOPWORDS: Final[frozenset[str]] = frozenset(
+    {
+        # Python keywords / builtins
+        "and", "as", "assert", "async", "await", "break", "class", "continue",
+        "def", "del", "elif", "else", "except", "finally", "for", "from",
+        "global", "if", "import", "in", "is", "lambda", "nonlocal", "not",
+        "or", "pass", "raise", "return", "try", "while", "with", "yield",
+        "True", "False", "None", "self", "cls",
+        "int", "str", "float", "bool", "list", "dict", "tuple", "set",
+        "type", "any", "all", "len", "range", "print",
+        # JS/TS keywords
+        "const", "let", "var", "function", "this", "new", "throw", "void",
+        "typeof", "instanceof", "switch", "case", "default", "extends",
+        # Common short words that masquerade as identifiers
+        "the", "that", "into",
+    }
+)
+
+# Maximum length of a "defined" entity in characters.
+_DEFINED_MAX_LEN: Final = 80
+
+# Tokens that disqualify a bold span from being a defined entity.
+_SENTENCE_MARKERS: Final = frozenset({".", ",", ";", ":", "!", "?"})
+
+
+class HeuristicExtractor:
+    """Regex-based entity extractor."""
+
+    name = "heuristic:regex-v1"
+
+    async def extract(self, document: Document) -> Iterable[ExtractionHit]:
+        hits: list[ExtractionHit] = []
+        for section in document.sections:
+            for hit in _scan_section(section.id, section.raw_text):
+                hits.append(hit)
+        return hits
+
+
+def _scan_section(section_id: str, text: str) -> Iterator[ExtractionHit]:
+    yield from _scan_code(section_id, text)
+    yield from _scan_defined(section_id, text)
+
+
+def _scan_code(section_id: str, text: str) -> Iterator[ExtractionHit]:
+    # Fenced code blocks.
+    for fence in _FENCED_CODE.finditer(text):
+        body = fence.group(1)
+        body_offset = fence.start(1)
+        for ident in _IDENTIFIER.finditer(body):
+            name = ident.group()
+            if name in _CODE_STOPWORDS:
+                continue
+            yield ExtractionHit(
+                section_id=section_id,
+                canonical=name,
+                surface_form=name,
+                kind="code",
+                span=Span(
+                    start=body_offset + ident.start(),
+                    end=body_offset + ident.end(),
+                ),
+            )
+
+    # Inline code spans. Skip those that fall inside a fenced block by
+    # checking offsets against fence ranges.
+    fence_ranges = [
+        (m.start(), m.end()) for m in _FENCED_CODE.finditer(text)
+    ]
+    for inline in _INLINE_CODE.finditer(text):
+        if _inside_any(inline.start(), fence_ranges):
+            continue
+        inner = inline.group(1)
+        inner_offset = inline.start(1)
+        # Only emit if the entire inline body looks like an identifier list.
+        # Skip prose-y `things like this`.
+        if " " in inner.strip() and len(inner.split()) > 1:
+            # Multi-word inline code → typically not a single identifier.
+            # Still scan for identifiers but treat them as separate hits.
+            pass
+        for ident in _IDENTIFIER.finditer(inner):
+            name = ident.group()
+            if name in _CODE_STOPWORDS:
+                continue
+            yield ExtractionHit(
+                section_id=section_id,
+                canonical=name,
+                surface_form=name,
+                kind="code",
+                span=Span(
+                    start=inner_offset + ident.start(),
+                    end=inner_offset + ident.end(),
+                ),
+            )
+
+
+def _scan_defined(section_id: str, text: str) -> Iterator[ExtractionHit]:
+    for bold in _BOLD.finditer(text):
+        inner_raw = bold.group(1)
+        inner = inner_raw.strip()
+        if not inner or len(inner) > _DEFINED_MAX_LEN:
+            continue
+        if any(marker in inner for marker in _SENTENCE_MARKERS):
+            continue
+        yield ExtractionHit(
+            section_id=section_id,
+            canonical=inner,
+            surface_form=inner,
+            kind="defined",
+            span=Span(start=bold.start(1), end=bold.end(1)),
+        )
+
+
+def _inside_any(pos: int, ranges: list[tuple[int, int]]) -> bool:
+    return any(start <= pos < end for start, end in ranges)

cairn/index/__init__.py

@@ -0,0 +1,39 @@
+"""Index layer — the five sub-indexes (Tree, Summaries, Entities, XRefs, Vectors)."""
+
+from cairn.index.entities import ENTITIES_FILENAME, Entities, EntityBuilder
+from cairn.index.summaries import (
+    SUMMARIES_FILENAME,
+    Summaries,
+    SummaryBuilder,
+    section_hash,
+)
+from cairn.index.tree import TREE_FILENAME, Tree, TreeBuilder
+from cairn.index.vectors import (
+    VECTORS_DB_DIRNAME,
+    VECTORS_MANIFEST_FILENAME,
+    VectorBuilder,
+    VectorHit,
+    Vectors,
+)
+from cairn.index.xrefs import XREFS_FILENAME, XRefBuilder, XRefs
+
+__all__ = [
+    "ENTITIES_FILENAME",
+    "SUMMARIES_FILENAME",
+    "TREE_FILENAME",
+    "VECTORS_DB_DIRNAME",
+    "VECTORS_MANIFEST_FILENAME",
+    "XREFS_FILENAME",
+    "Entities",
+    "EntityBuilder",
+    "Summaries",
+    "SummaryBuilder",
+    "Tree",
+    "TreeBuilder",
+    "VectorBuilder",
+    "VectorHit",
+    "Vectors",
+    "XRefBuilder",
+    "XRefs",
+    "section_hash",
+]

cairn/index/entities.py

@@ -0,0 +1,244 @@
+"""Entities sub-index — extractor outputs, deduplicated and persisted.
+
+The :class:`EntityBuilder` runs an :class:`cairn.entity.base.EntityExtractor`
+over a Document, aggregates ``ExtractionHit`` records by ``(canonical, kind)``
+into :class:`cairn.core.types.Entity`, and writes ``entities.json``.
+
+The :class:`Entities` reader exposes lookup by canonical / surface form /
+section, with optional ``kind`` filtering — exactly what the
+``find_mentions`` retrieval tool needs.
+"""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from collections.abc import Iterable, Iterator
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, Final
+
+from cairn.core.errors import IndexBuildError, IndexNotFoundError
+from cairn.core.types import Document, Entity, EntityKind, Mention, Span
+from cairn.entity.base import EntityExtractor, ExtractionHit
+
+ENTITIES_FILENAME: Final = "entities.json"
+ENTITIES_FORMAT_VERSION: Final = 1
+
+
+class EntityBuilder:
+    """Run an extractor, aggregate hits, persist ``entities.json``."""
+
+    def __init__(self, extractor: EntityExtractor) -> None:
+        self.extractor = extractor
+
+    async def build(self, document: Document, *, out_dir: Path) -> Path:
+        out_dir.mkdir(parents=True, exist_ok=True)
+        path = out_dir / ENTITIES_FILENAME
+
+        hits = await self.extractor.extract(document)
+        entities = _aggregate(hits)
+        now = datetime.now(UTC)
+
+        payload: dict[str, Any] = {
+            "format_version": ENTITIES_FORMAT_VERSION,
+            "doc_id": document.id,
+            "extractor": self.extractor.name,
+            "generated_at": now.isoformat(),
+            "entities": [_entity_to_dict(e) for e in entities],
+        }
+
+        with path.open("w", encoding="utf-8") as fh:
+            json.dump(payload, fh, ensure_ascii=False, indent=2)
+            fh.write("\n")
+        return path
+
+
+class Entities:
+    """Loaded entities sub-index. Read-only queries."""
+
+    def __init__(
+        self,
+        entities: tuple[Entity, ...],
+        *,
+        doc_id: str,
+        extractor: str,
+    ) -> None:
+        self._all = entities
+        self.doc_id = doc_id
+        self.extractor = extractor
+
+        # Indexes
+        self._by_canonical: dict[tuple[str, str], Entity] = {
+            (e.canonical, e.kind): e for e in entities
+        }
+        self._by_surface: dict[str, list[Entity]] = defaultdict(list)
+        self._by_section: dict[str, list[Entity]] = defaultdict(list)
+        for ent in entities:
+            for sf in ent.surface_forms:
+                self._by_surface[sf].append(ent)
+            for mention in ent.mentions:
+                self._by_section[mention.section_id].append(ent)
+
+    @classmethod
+    def load(cls, doc_dir: Path) -> Entities:
+        path = doc_dir / ENTITIES_FILENAME
+        if not path.exists():
+            msg = f"entities.json not found in {doc_dir}"
+            raise IndexNotFoundError(msg, details={"path": str(path)})
+
+        with path.open("r", encoding="utf-8") as fh:
+            payload = json.load(fh)
+
+        version = payload.get("format_version")
+        if version != ENTITIES_FORMAT_VERSION:
+            msg = (
+                f"unsupported entities format version: {version!r} "
+                f"(expected {ENTITIES_FORMAT_VERSION})"
+            )
+            raise IndexNotFoundError(msg, details={"path": str(path)})
+
+        entities = tuple(_entity_from_dict(d) for d in payload["entities"])
+        return cls(
+            entities,
+            doc_id=payload["doc_id"],
+            extractor=payload["extractor"],
+        )
+
+    # -- queries -----------------------------------------------------------
+
+    def __len__(self) -> int:
+        return len(self._all)
+
+    def __iter__(self) -> Iterator[Entity]:
+        return iter(self._all)
+
+    def lookup(
+        self,
+        name: str,
+        *,
+        kinds: tuple[EntityKind, ...] | None = None,
+    ) -> Entity | None:
+        """Return the first matching entity by canonical or any surface form.
+
+        Precedence: canonical match before surface-form match. When ``kinds``
+        is supplied, only entities of those kinds are considered.
+        """
+        for ent in self._candidates(name):
+            if kinds is None or ent.kind in kinds:
+                return ent
+        return None
+
+    def lookup_all(
+        self,
+        name: str,
+        *,
+        kinds: tuple[EntityKind, ...] | None = None,
+    ) -> list[Entity]:
+        """Return every matching entity (across kinds) for ``name``."""
+        seen: set[tuple[str, str]] = set()
+        out: list[Entity] = []
+        for ent in self._candidates(name):
+            key = (ent.canonical, ent.kind)
+            if key in seen:
+                continue
+            if kinds is None or ent.kind in kinds:
+                out.append(ent)
+                seen.add(key)
+        return out
+
+    def by_section(self, section_id: str) -> list[Entity]:
+        return list(self._by_section.get(section_id, ()))
+
+    def by_kind(self, kind: EntityKind) -> list[Entity]:
+        return [e for e in self._all if e.kind == kind]
+
+    def _candidates(self, name: str) -> Iterator[Entity]:
+        # Canonical hits first, across all kinds, in extractor order.
+        for ent in self._all:
+            if ent.canonical == name:
+                yield ent
+        # Then surface-form matches.
+        for ent in self._by_surface.get(name, ()):
+            if ent.canonical != name:
+                yield ent
+
+
+# ---------------------------------------------------------------------------
+# Aggregation
+# ---------------------------------------------------------------------------
+
+
+def _aggregate(hits: Iterable[ExtractionHit]) -> list[Entity]:
+    """Fold ``ExtractionHit`` stream into a flat list of :class:`Entity`."""
+    by_key: dict[tuple[str, str], _Acc] = {}
+    insertion_order: list[tuple[str, str]] = []
+
+    for hit in hits:
+        if not hit.canonical:
+            msg = "extractor emitted an empty canonical"
+            raise IndexBuildError(msg)
+        key = (hit.canonical, hit.kind)
+        acc = by_key.get(key)
+        if acc is None:
+            acc = _Acc(canonical=hit.canonical, kind=hit.kind)
+            by_key[key] = acc
+            insertion_order.append(key)
+        acc.surface_forms[hit.surface_form] = None
+        acc.mentions.append(Mention(section_id=hit.section_id, span=hit.span))
+
+    return [by_key[key].freeze() for key in insertion_order]
+
+
+class _Acc:
+    __slots__ = ("canonical", "kind", "mentions", "surface_forms")
+
+    def __init__(self, *, canonical: str, kind: EntityKind) -> None:
+        self.canonical = canonical
+        self.kind = kind
+        # dict[str, None] is the cheapest insertion-ordered set in Python.
+        self.surface_forms: dict[str, None] = {}
+        self.mentions: list[Mention] = []
+
+    def freeze(self) -> Entity:
+        return Entity(
+            canonical=self.canonical,
+            surface_forms=tuple(self.surface_forms.keys()),
+            kind=self.kind,
+            mentions=tuple(self.mentions),
+        )
+
+
+# ---------------------------------------------------------------------------
+# Serialization
+# ---------------------------------------------------------------------------
+
+
+def _entity_to_dict(e: Entity) -> dict[str, Any]:
+    return {
+        "canonical": e.canonical,
+        "surface_forms": list(e.surface_forms),
+        "kind": e.kind,
+        "mentions": [
+            {
+                "section_id": m.section_id,
+                "span": {"start": m.span.start, "end": m.span.end},
+            }
+            for m in e.mentions
+        ],
+    }
+
+
+def _entity_from_dict(d: dict[str, Any]) -> Entity:
+    return Entity(
+        canonical=d["canonical"],
+        surface_forms=tuple(d["surface_forms"]),
+        kind=d["kind"],
+        mentions=tuple(
+            Mention(
+                section_id=m["section_id"],
+                span=Span(start=m["span"]["start"], end=m["span"]["end"]),
+            )
+            for m in d["mentions"]
+        ),
+    )

cairn/index/summaries.py

@@ -0,0 +1,269 @@
+"""Summaries sub-index — per-section multi-granularity summaries.
+
+``SummaryBuilder`` runs at indexing time and writes ``summaries.json``;
+``Summaries`` loads that file and serves queries. Both honor the
+:class:`cairn.core.types.SummarySet` contract.
+
+The cache layer (``cairn.summarize.cache.SummaryCache``) is keyed by
+``(summarizer.name, level, section_hash)``. Switching to a different
+summarizer transparently invalidates prior cache entries; editing a section
+invalidates only that section's entries.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+from collections.abc import Awaitable, Callable, Iterator, Sequence
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, Final
+
+from cairn.core.errors import IndexBuildError, IndexNotFoundError
+from cairn.core.types import Document, SectionNode, SummarySet
+from cairn.summarize.base import Summarizer, SummaryLevel
+from cairn.summarize.cache import SummaryCache
+
+SUMMARIES_FILENAME: Final = "summaries.json"
+SUMMARIES_FORMAT_VERSION: Final = 1
+
+
+def section_hash(node: SectionNode) -> str:
+    """Cache-invalidating fingerprint of a section's content.
+
+    Uses ``sha256(title || NUL || raw_text)``. Changes in either invalidate
+    cached summaries for that section.
+    """
+    h = hashlib.sha256()
+    h.update(node.title.encode("utf-8"))
+    h.update(b"\x00")
+    h.update(node.raw_text.encode("utf-8"))
+    return h.hexdigest()
+
+
+class SummaryBuilder:
+    """Asynchronously generate and persist summaries for a Document."""
+
+    def __init__(
+        self,
+        summarizer: Summarizer,
+        *,
+        cache: SummaryCache | None = None,
+        concurrency: int = 4,
+        progress: Callable[[int, int], None] | None = None,
+    ) -> None:
+        if concurrency < 1:
+            msg = f"concurrency must be ≥1; got {concurrency}"
+            raise IndexBuildError(msg)
+        self.summarizer = summarizer
+        self.cache = cache
+        self.concurrency = concurrency
+        self.progress = progress
+
+    async def build(
+        self,
+        document: Document,
+        *,
+        out_dir: Path,
+        levels: Sequence[SummaryLevel] = (
+            SummaryLevel.GIST,
+            SummaryLevel.SYNOPSIS,
+            SummaryLevel.DIGEST,
+        ),
+    ) -> Path:
+        """Summarize every section in ``document`` and write ``summaries.json``.
+
+        Args:
+            document: The parsed document.
+            out_dir: Directory to write into. Created if absent.
+            levels: Which granularity levels to generate. Order is preserved
+                and duplicates are dropped. Defaults to all three levels
+                (gist, synopsis, digest) since v0.2.4.
+
+        Returns:
+            Path to the written ``summaries.json``.
+        """
+        ordered_levels = _dedupe_preserve_order(levels)
+        if not ordered_levels:
+            msg = "at least one SummaryLevel is required"
+            raise IndexBuildError(msg)
+
+        out_dir.mkdir(parents=True, exist_ok=True)
+        path = out_dir / SUMMARIES_FILENAME
+        now = datetime.now(UTC)
+        semaphore = asyncio.Semaphore(self.concurrency)
+        progress_lock = asyncio.Lock()
+        completed = 0
+        total = len(document.sections) * len(ordered_levels)
+
+        async def mark_progress() -> None:
+            nonlocal completed
+            async with progress_lock:
+                completed += 1
+                if self.progress is None:
+                    return
+                step = max(5, total // 20)
+                if completed != 1 and completed != total and completed % step != 0:
+                    return
+                self.progress(completed, total)
+
+        async def for_section(node: SectionNode) -> dict[str, Any]:
+            return await self._summarize_section(
+                node,
+                ordered_levels,
+                semaphore,
+                now,
+                mark_progress=mark_progress,
+            )
+
+        records = await asyncio.gather(
+            *(for_section(s) for s in document.sections)
+        )
+
+        payload: dict[str, Any] = {
+            "format_version": SUMMARIES_FORMAT_VERSION,
+            "doc_id": document.id,
+            "model": self.summarizer.name,
+            "levels": [lvl.value for lvl in ordered_levels],
+            "generated_at": now.isoformat(),
+            "summaries": list(records),
+        }
+
+        with path.open("w", encoding="utf-8") as fh:
+            json.dump(payload, fh, ensure_ascii=False, indent=2)
+            fh.write("\n")
+        return path
+
+    async def _summarize_section(
+        self,
+        node: SectionNode,
+        levels: Sequence[SummaryLevel],
+        semaphore: asyncio.Semaphore,
+        now: datetime,
+        mark_progress: Callable[[], Awaitable[None]],
+    ) -> dict[str, Any]:
+        sh = section_hash(node)
+        results: dict[str, str] = {}
+
+        for level in levels:
+            cache_key: str | None = None
+            cached: str | None = None
+            if self.cache is not None:
+                cache_key = SummaryCache.key(
+                    model=self.summarizer.name,
+                    level=level.value,
+                    section_hash=sh,
+                )
+                cached = self.cache.get(cache_key)
+
+            if cached is not None:
+                results[level.value] = cached
+                await mark_progress()
+                continue
+
+            async with semaphore:
+                text = await self.summarizer.summarize(
+                    title=node.title,
+                    body=node.raw_text,
+                    level=level,
+                )
+            results[level.value] = text
+            await mark_progress()
+            if self.cache is not None and cache_key is not None:
+                self.cache.put(cache_key, text)
+
+        return {
+            "section_id": node.id,
+            "section_hash": sh,
+            "gist": results.get(SummaryLevel.GIST.value, ""),
+            "synopsis": results.get(SummaryLevel.SYNOPSIS.value, ""),
+            "digest": results.get(SummaryLevel.DIGEST.value),
+            "model": self.summarizer.name,
+            "generated_at": now.isoformat(),
+        }
+
+
+class Summaries:
+    """Loaded summaries index. Read-only by section id."""
+
+    def __init__(
+        self,
+        sets: tuple[SummarySet, ...],
+        *,
+        doc_id: str,
+        model: str,
+    ) -> None:
+        self._by_id: dict[str, SummarySet] = {s.section_id: s for s in sets}
+        self._all = sets
+        self.doc_id = doc_id
+        self.model = model
+
+    @classmethod
+    def load(cls, doc_dir: Path) -> Summaries:
+        """Load ``summaries.json`` from a document directory."""
+        path = doc_dir / SUMMARIES_FILENAME
+        if not path.exists():
+            msg = f"summaries.json not found in {doc_dir}"
+            raise IndexNotFoundError(msg, details={"path": str(path)})
+
+        with path.open("r", encoding="utf-8") as fh:
+            payload = json.load(fh)
+
+        version = payload.get("format_version")
+        if version != SUMMARIES_FORMAT_VERSION:
+            msg = (
+                f"unsupported summaries format version: {version!r} "
+                f"(expected {SUMMARIES_FORMAT_VERSION})"
+            )
+            raise IndexNotFoundError(msg, details={"path": str(path)})
+
+        sets = tuple(
+            SummarySet(
+                section_id=record["section_id"],
+                gist=record["gist"],
+                synopsis=record["synopsis"],
+                digest=record.get("digest"),
+                model=record["model"],
+                section_hash=record["section_hash"],
+                generated_at=datetime.fromisoformat(record["generated_at"]),
+            )
+            for record in payload["summaries"]
+        )
+        return cls(sets, doc_id=payload["doc_id"], model=payload["model"])
+
+    # -- queries -----------------------------------------------------------
+
+    def get(self, section_id: str) -> SummarySet | None:
+        return self._by_id.get(section_id)
+
+    def require(self, section_id: str) -> SummarySet:
+        s = self.get(section_id)
+        if s is None:
+            msg = f"summary not found for section: {section_id!r}"
+            raise IndexNotFoundError(msg, details={"section_id": section_id})
+        return s
+
+    def __contains__(self, section_id: object) -> bool:
+        return isinstance(section_id, str) and section_id in self._by_id
+
+    def __len__(self) -> int:
+        return len(self._all)
+
+    def __iter__(self) -> Iterator[SummarySet]:
+        return iter(self._all)
+
+
+# ---------------------------------------------------------------------------
+# helpers
+# ---------------------------------------------------------------------------
+
+
+def _dedupe_preserve_order(levels: Sequence[SummaryLevel]) -> list[SummaryLevel]:
+    seen: set[SummaryLevel] = set()
+    out: list[SummaryLevel] = []
+    for lvl in levels:
+        if lvl not in seen:
+            seen.add(lvl)
+            out.append(lvl)
+    return out