docsgraph 0.1.0a2__py3-none-any.whl

cairn/index/tree.py

@@ -0,0 +1,274 @@
+"""Tree sub-index — persistence and queries for the structural backbone.
+
+`TreeBuilder` writes a deterministic ``tree.json`` from a parsed
+:class:`Document`. `Tree` loads and queries it.
+
+The tree is the primary navigation structure (ARCHITECTURE.md §2.1). All other
+sub-indexes key into the ``section_id`` namespace it defines.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Final
+
+from cairn.core.errors import IndexBuildError, IndexNotFoundError
+from cairn.core.types import Document, SectionNode, Span
+
+TREE_FILENAME: Final = "tree.json"
+TREE_FORMAT_VERSION: Final = 1
+
+
+class TreeBuilder:
+    """Writes the structural tree of a Document to ``tree.json``."""
+
+    def build(self, document: Document, *, out_dir: Path) -> Path:
+        """Serialize ``document.sections`` into ``out_dir/tree.json``.
+
+        Args:
+            document: The parsed document. Its `sections` must form a valid
+                forest (every non-root section's `parent` exists; every
+                referenced `child` exists).
+            out_dir: Directory to write into. Created if it does not exist.
+
+        Returns:
+            The path to the written ``tree.json``.
+        """
+        self._validate_tree(document)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        path = out_dir / TREE_FILENAME
+
+        payload: dict[str, Any] = {
+            "format_version": TREE_FORMAT_VERSION,
+            "doc_id": document.id,
+            "source_path": str(document.source_path),
+            "source_hash": document.source_hash,
+            "indexed_at": document.indexed_at.isoformat(),
+            "cairn_version": document.cairn_version,
+            "sections": [_section_to_dict(s) for s in document.sections],
+        }
+
+        with path.open("w", encoding="utf-8") as fh:
+            json.dump(payload, fh, ensure_ascii=False, indent=2, sort_keys=False)
+            fh.write("\n")
+        return path
+
+    @staticmethod
+    def _validate_tree(document: Document) -> None:
+        seen_ids: set[str] = set()
+        for section in document.sections:
+            if section.id in seen_ids:
+                msg = f"duplicate section id in document: {section.id!r}"
+                raise IndexBuildError(msg, details={"section_id": section.id})
+            seen_ids.add(section.id)
+
+        for section in document.sections:
+            if section.parent is not None and section.parent not in seen_ids:
+                msg = (
+                    f"section {section.id!r} references unknown parent "
+                    f"{section.parent!r}"
+                )
+                raise IndexBuildError(
+                    msg,
+                    details={"section_id": section.id, "parent": section.parent},
+                )
+            for child in section.children:
+                if child not in seen_ids:
+                    msg = (
+                        f"section {section.id!r} references unknown child "
+                        f"{child!r}"
+                    )
+                    raise IndexBuildError(
+                        msg,
+                        details={"section_id": section.id, "child": child},
+                    )
+
+
+class Tree:
+    """Loaded tree index. Read-only queries against the structural backbone."""
+
+    def __init__(
+        self,
+        sections: tuple[SectionNode, ...],
+        *,
+        doc_id: str,
+        source_hash: str,
+        indexed_at: datetime,
+    ) -> None:
+        self._sections = sections
+        self._by_id: dict[str, SectionNode] = {s.id: s for s in sections}
+        self._roots: tuple[SectionNode, ...] = tuple(
+            s for s in sections if s.parent is None
+        )
+        self.doc_id = doc_id
+        self.source_hash = source_hash
+        self.indexed_at = indexed_at
+
+    # -- construction --------------------------------------------------------
+
+    @classmethod
+    def load(cls, doc_dir: Path) -> Tree:
+        """Load ``tree.json`` from a document directory."""
+        path = doc_dir / TREE_FILENAME
+        if not path.exists():
+            msg = f"tree.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)
+
+        format_version = payload.get("format_version")
+        if format_version != TREE_FORMAT_VERSION:
+            msg = (
+                f"unsupported tree format version: {format_version!r} "
+                f"(expected {TREE_FORMAT_VERSION})"
+            )
+            raise IndexNotFoundError(msg, details={"path": str(path)})
+
+        sections = tuple(_section_from_dict(d) for d in payload["sections"])
+        return cls(
+            sections,
+            doc_id=payload["doc_id"],
+            source_hash=payload["source_hash"],
+            indexed_at=datetime.fromisoformat(payload["indexed_at"]),
+        )
+
+    # -- queries -------------------------------------------------------------
+
+    def get(self, section_id: str) -> SectionNode | None:
+        """Look up a section by id. Returns ``None`` if absent."""
+        return self._by_id.get(section_id)
+
+    def require(self, section_id: str) -> SectionNode:
+        """Look up a section by id, raising :class:`IndexNotFoundError`."""
+        node = self.get(section_id)
+        if node is None:
+            msg = f"section not found: {section_id!r}"
+            raise IndexNotFoundError(msg, details={"section_id": section_id})
+        return node
+
+    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._sections)
+
+    def __iter__(self) -> Iterator[SectionNode]:
+        """Yield every section in document order."""
+        return iter(self._sections)
+
+    def roots(self) -> tuple[SectionNode, ...]:
+        """Top-level sections (those with `parent is None`)."""
+        return self._roots
+
+    def children_of(self, section_id: str) -> tuple[SectionNode, ...]:
+        """Direct children of a section, in document order."""
+        node = self.require(section_id)
+        return tuple(self._by_id[cid] for cid in node.children)
+
+    def descendants_of(self, section_id: str) -> Iterator[SectionNode]:
+        """Depth-first traversal of a section's descendants (excluding self)."""
+        node = self.require(section_id)
+        stack: list[str] = list(reversed(node.children))
+        while stack:
+            current_id = stack.pop()
+            current = self._by_id[current_id]
+            yield current
+            stack.extend(reversed(current.children))
+
+    def ancestors_of(self, section_id: str) -> Iterator[SectionNode]:
+        """Walk parents from the section up to the root (excluding self)."""
+        node = self.require(section_id)
+        current = node.parent
+        while current is not None:
+            parent_node = self._by_id[current]
+            yield parent_node
+            current = parent_node.parent
+
+    def outline(
+        self,
+        *,
+        depth: int = 2,
+        focus: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Return a nested outline suitable for the ``outline`` MCP tool.
+
+        Each node has: ``id``, ``title``, ``level``, ``children`` (recursively),
+        plus ``truncated: True`` when descendants exist beyond ``depth``.
+        Summaries are **not** attached here — that is the MCP tool's job after
+        joining with the Summaries sub-index.
+        """
+        if depth < 1 or depth > 6:
+            msg = f"depth must be in [1, 6]; got {depth}"
+            raise IndexNotFoundError(msg)
+
+        if focus is None:
+            roots = self._roots
+            base_level = 0
+        else:
+            focused = self.require(focus)
+            roots = (focused,)
+            base_level = focused.level - 1
+
+        return [self._outline_node(s, depth, base_level) for s in roots]
+
+    def _outline_node(
+        self,
+        node: SectionNode,
+        depth: int,
+        base_level: int,
+    ) -> dict[str, Any]:
+        remaining = depth - (node.level - base_level)
+        children_payload: list[dict[str, Any]] = []
+        truncated = False
+        if remaining > 0 and node.children:
+            for child_id in node.children:
+                child = self._by_id[child_id]
+                children_payload.append(self._outline_node(child, depth, base_level))
+        elif node.children:
+            truncated = True
+
+        payload: dict[str, Any] = {
+            "id": node.id,
+            "title": node.title,
+            "level": node.level,
+            "children": children_payload,
+        }
+        if truncated:
+            payload["truncated"] = True
+        return payload
+
+
+# ---------------------------------------------------------------------------
+# (de)serialization
+# ---------------------------------------------------------------------------
+
+
+def _section_to_dict(s: SectionNode) -> dict[str, Any]:
+    return {
+        "id": s.id,
+        "title": s.title,
+        "level": s.level,
+        "parent": s.parent,
+        "children": list(s.children),
+        "span": {"start": s.span.start, "end": s.span.end},
+        "path": list(s.path),
+        "raw_text": s.raw_text,
+    }
+
+
+def _section_from_dict(d: dict[str, Any]) -> SectionNode:
+    span = d["span"]
+    return SectionNode(
+        id=d["id"],
+        title=d["title"],
+        level=d["level"],
+        parent=d["parent"],
+        children=tuple(d["children"]),
+        span=Span(start=span["start"], end=span["end"]),
+        path=tuple(d["path"]),
+        raw_text=d["raw_text"],
+    )

cairn/index/vectors.py

@@ -0,0 +1,287 @@
+"""Vectors sub-index — dense embeddings over LanceDB.
+
+Storage layout::
+
+    <doc_dir>/
+    ├── vectors.lance/            # LanceDB connect root
+    │   └── data.lance/           # table holding (id, vector)
+    └── vectors_manifest.json     # embedder name, dim, build metadata
+
+LanceDB is the v0.1 default per ARCHITECTURE.md §7. We use the sync API and
+wrap blocking calls in ``asyncio.to_thread`` to satisfy our async-by-default
+public surface without adopting LanceDB's still-evolving native async API.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import math
+import re
+import shutil
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, Final
+
+import lancedb
+import pyarrow as pa
+from pydantic import BaseModel, ConfigDict, Field
+
+from cairn.core.errors import IndexBuildError, IndexNotFoundError
+from cairn.core.types import Document, SectionNode
+from cairn.embed.base import Embedder
+
+VECTORS_DB_DIRNAME: Final = "vectors.lance"
+VECTORS_TABLE_NAME: Final = "data"
+VECTORS_MANIFEST_FILENAME: Final = "vectors_manifest.json"
+VECTORS_FORMAT_VERSION: Final = 1
+
+_SCOPE_PATTERN = re.compile(r"^[a-z0-9][a-z0-9_/-]*$")
+
+
+class VectorHit(BaseModel):
+    """One result row from a vector search."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    id: str
+    score: float = Field(ge=0.0, le=1.0)
+
+
+class VectorEntry(BaseModel):
+    """One persisted section vector."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    id: str
+    vector: list[float]
+
+
+def embedding_text(node: SectionNode) -> str:
+    """Compose the text we embed for a section.
+
+    Includes the title so heading information enters the embedding, and falls
+    back to title alone for sections with empty bodies.
+    """
+    body = node.raw_text.strip()
+    if not body:
+        return node.title
+    return f"{node.title}\n\n{body}"
+
+
+def l2_normalize(vec: list[float]) -> list[float]:
+    """Return the L2-normalized copy of ``vec``. Zero vectors are returned unchanged."""
+    norm = math.sqrt(sum(x * x for x in vec))
+    if norm == 0.0:
+        return list(vec)
+    return [x / norm for x in vec]
+
+
+class VectorBuilder:
+    """Embed and persist section-level vectors for a Document."""
+
+    def __init__(
+        self,
+        embedder: Embedder,
+        *,
+        batch_size: int = 32,
+    ) -> None:
+        if batch_size < 1:
+            msg = f"batch_size must be >= 1; got {batch_size}"
+            raise IndexBuildError(msg)
+        self.embedder = embedder
+        self.batch_size = batch_size
+
+    async def build(self, document: Document, *, out_dir: Path) -> Path:
+        """Embed every section and write ``vectors.lance/`` + manifest.
+
+        Returns the path to the manifest file.
+        """
+        out_dir.mkdir(parents=True, exist_ok=True)
+        db_dir = out_dir / VECTORS_DB_DIRNAME
+        manifest_path = out_dir / VECTORS_MANIFEST_FILENAME
+
+        ids = [s.id for s in document.sections]
+        texts = [embedding_text(s) for s in document.sections]
+
+        vectors: list[list[float]] = []
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+            raw = await self.embedder.embed(batch)
+            if len(raw) != len(batch):
+                msg = (
+                    f"embedder returned {len(raw)} vectors for batch of "
+                    f"{len(batch)}"
+                )
+                raise IndexBuildError(msg)
+            for vec in raw:
+                if len(vec) != self.embedder.dim:
+                    msg = (
+                        f"embedder returned dim={len(vec)} but expected "
+                        f"dim={self.embedder.dim}"
+                    )
+                    raise IndexBuildError(msg)
+                vectors.append(l2_normalize(vec))
+
+        await asyncio.to_thread(self._write_table, db_dir, ids, vectors)
+
+        now = datetime.now(UTC)
+        manifest = {
+            "format_version": VECTORS_FORMAT_VERSION,
+            "doc_id": document.id,
+            "embedder": self.embedder.name,
+            "dim": self.embedder.dim,
+            "section_count": len(ids),
+            "generated_at": now.isoformat(),
+        }
+        with manifest_path.open("w", encoding="utf-8") as fh:
+            json.dump(manifest, fh, ensure_ascii=False, indent=2)
+            fh.write("\n")
+        return manifest_path
+
+    def _write_table(
+        self,
+        db_dir: Path,
+        ids: list[str],
+        vectors: list[list[float]],
+    ) -> None:
+        # Full rebuild: clear any previous table data for a clean schema state.
+        if db_dir.exists():
+            shutil.rmtree(db_dir)
+
+        db = lancedb.connect(str(db_dir))
+        schema = pa.schema(
+            [
+                pa.field("id", pa.string()),
+                pa.field("vector", pa.list_(pa.float32(), self.embedder.dim)),
+            ]
+        )
+        table = db.create_table(VECTORS_TABLE_NAME, schema=schema)
+        if ids:
+            records = [
+                {"id": sid, "vector": vec} for sid, vec in zip(ids, vectors, strict=True)
+            ]
+            table.add(records)
+
+
+class Vectors:
+    """Loaded vectors index. Cosine-similarity search via LanceDB."""
+
+    def __init__(
+        self,
+        table: Any,
+        *,
+        doc_id: str,
+        embedder: str,
+        dim: int,
+    ) -> None:
+        self._table = table
+        self.doc_id = doc_id
+        self.embedder = embedder
+        self.dim = dim
+
+    @classmethod
+    def load(cls, doc_dir: Path) -> Vectors:
+        """Load vectors index from a document directory."""
+        manifest_path = doc_dir / VECTORS_MANIFEST_FILENAME
+        db_dir = doc_dir / VECTORS_DB_DIRNAME
+        if not manifest_path.exists():
+            msg = f"vectors manifest not found in {doc_dir}"
+            raise IndexNotFoundError(msg, details={"path": str(manifest_path)})
+        if not db_dir.exists():
+            msg = f"vectors.lance directory not found in {doc_dir}"
+            raise IndexNotFoundError(msg, details={"path": str(db_dir)})
+
+        with manifest_path.open("r", encoding="utf-8") as fh:
+            manifest = json.load(fh)
+
+        version = manifest.get("format_version")
+        if version != VECTORS_FORMAT_VERSION:
+            msg = (
+                f"unsupported vectors format version: {version!r} "
+                f"(expected {VECTORS_FORMAT_VERSION})"
+            )
+            raise IndexNotFoundError(msg, details={"path": str(manifest_path)})
+
+        db = lancedb.connect(str(db_dir))
+        table = db.open_table(VECTORS_TABLE_NAME)
+        return cls(
+            table,
+            doc_id=manifest["doc_id"],
+            embedder=manifest["embedder"],
+            dim=manifest["dim"],
+        )
+
+    async def search(
+        self,
+        query: list[float],
+        *,
+        k: int = 8,
+        scope_prefix: str | None = None,
+    ) -> list[VectorHit]:
+        """Return up to ``k`` nearest sections by cosine similarity.
+
+        When ``scope_prefix`` is given, results are restricted to sections
+        whose id equals the prefix or begins with ``f"{prefix}/"``.
+        """
+        if k < 1:
+            msg = f"k must be >= 1; got {k}"
+            raise IndexBuildError(msg)
+        if len(query) != self.dim:
+            msg = f"query dim {len(query)} != index dim {self.dim}"
+            raise IndexBuildError(msg)
+        if scope_prefix is not None and not _SCOPE_PATTERN.match(scope_prefix):
+            msg = (
+                f"invalid scope_prefix {scope_prefix!r}; only lowercase "
+                "alphanumeric, '-', '_', '/' allowed"
+            )
+            raise IndexBuildError(msg)
+
+        normalized = l2_normalize(query)
+        return await asyncio.to_thread(
+            self._sync_search, normalized, k, scope_prefix
+        )
+
+    def _sync_search(
+        self,
+        vec: list[float],
+        k: int,
+        scope_prefix: str | None,
+    ) -> list[VectorHit]:
+        q = self._table.search(vec).distance_type("cosine")
+        if scope_prefix is not None:
+            predicate = (
+                f"id = '{scope_prefix}' OR id LIKE '{scope_prefix}/%'"
+            )
+            q = q.where(predicate, prefilter=True)
+        rows = q.limit(k).to_list()
+
+        hits: list[VectorHit] = []
+        for row in rows:
+            distance = float(row["_distance"])
+            score = max(0.0, min(1.0, 1.0 - distance))
+            hits.append(VectorHit(id=str(row["id"]), score=score))
+        return hits
+
+    async def count(self) -> int:
+        """Total number of indexed sections."""
+        return await asyncio.to_thread(self._table.count_rows)
+
+    async def entries(self) -> list[VectorEntry]:
+        """Return every stored vector.
+
+        Repo-scoped search uses this to build a process-local flat index once,
+        then answers repeated MCP queries without reopening every per-document
+        LanceDB table.
+        """
+        return await asyncio.to_thread(self._sync_entries)
+
+    def _sync_entries(self) -> list[VectorEntry]:
+        rows = self._table.to_arrow().to_pylist()
+        return [
+            VectorEntry(
+                id=str(row["id"]),
+                vector=[float(value) for value in row["vector"]],
+            )
+            for row in rows
+        ]