docsgraph 0.1.0a2__py3-none-any.whl

cairn/__init__.py

@@ -0,0 +1,5 @@
+"""Cairn — structure-aware, MCP-native retrieval for large documents."""
+
+__version__ = "0.1.0a2"
+
+__all__ = ["__version__"]

cairn/bench/__init__.py

@@ -0,0 +1,37 @@
+"""cairn-bench — evaluate Cairn against a naive vector-RAG baseline.
+
+A small framework for measuring retrieval recall and token cost on
+hand-curated question sets. The shipped starter dataset is small;
+production-grade evaluation requires more curation than fits in this
+codebase. The framework is the contribution; the dataset is a template.
+"""
+
+from cairn.bench.dataset import (
+    BenchDocument,
+    BenchQuestion,
+    BenchSuite,
+    load_suite,
+)
+from cairn.bench.judge import LLMJudge
+from cairn.bench.metrics import recall_at_k
+from cairn.bench.report import (
+    BenchSummary,
+    QuestionResult,
+    format_markdown_report,
+    write_json_report,
+)
+from cairn.bench.runner import BenchRunner
+
+__all__ = [
+    "BenchDocument",
+    "BenchQuestion",
+    "BenchRunner",
+    "BenchSuite",
+    "BenchSummary",
+    "LLMJudge",
+    "QuestionResult",
+    "format_markdown_report",
+    "load_suite",
+    "recall_at_k",
+    "write_json_report",
+]

cairn/bench/baseline.py

@@ -0,0 +1,236 @@
+"""Naive RAG baseline.
+
+Splits the source text into fixed N-word chunks, embeds each chunk into a
+LanceDB table, and serves cosine-similarity retrieval. **Crucially, the
+chunker is structure-blind** — chunks straddle headings, paragraphs, and
+section boundaries. This is the failure mode Cairn is built to fix; the
+baseline lets us measure the cost of that failure.
+
+For comparable recall reporting against Cairn (which returns section ids),
+each chunk is mapped to "the section whose span contains the chunk's
+midpoint byte" at index time. Deepest section wins on overlap.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import re
+import shutil
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+import lancedb
+import pyarrow as pa
+from pydantic import BaseModel, ConfigDict
+
+from cairn.core.types import Document, SectionNode
+from cairn.embed.base import Embedder
+from cairn.index.vectors import l2_normalize
+
+_WORD = re.compile(r"\S+")
+
+
+class _Chunk(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    chunk_id: int
+    text: str
+    byte_start: int
+    byte_end: int
+    section_id: str | None
+
+
+class NaiveHit(BaseModel):
+    """One retrieval result from the naive baseline."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    chunk_id: int
+    section_id: str | None
+    text: str
+    score: float
+
+
+def chunk_text(text: str, *, chunk_size_words: int = 512) -> list[tuple[int, int, str]]:
+    """Split ``text`` into chunks of approximately ``chunk_size_words`` words.
+
+    Returns ``(byte_start, byte_end, chunk_text)`` for each chunk. Word
+    boundaries are whitespace-delimited.
+    """
+    if chunk_size_words < 1:
+        msg = f"chunk_size_words must be >= 1; got {chunk_size_words}"
+        raise ValueError(msg)
+    words = list(_WORD.finditer(text))
+    if not words:
+        return []
+    chunks: list[tuple[int, int, str]] = []
+    for i in range(0, len(words), chunk_size_words):
+        start = words[i].start()
+        end = words[min(i + chunk_size_words - 1, len(words) - 1)].end()
+        chunks.append((start, end, text[start:end]))
+    return chunks
+
+
+def assign_section(
+    chunk_start: int,
+    chunk_end: int,
+    sections: Sequence[SectionNode],
+) -> str | None:
+    """Return the deepest section whose ``span`` contains the chunk's midpoint."""
+    midpoint = (chunk_start + chunk_end) // 2
+    deepest: SectionNode | None = None
+    for section in sections:
+        if section.span.start <= midpoint < section.span.end and (
+            deepest is None or section.level > deepest.level
+        ):
+            deepest = section
+    return deepest.id if deepest is not None else None
+
+
+class NaiveRAG:
+    """Structure-blind chunk + vector search baseline."""
+
+    name = "naive-rag"
+    table_name = "chunks"
+
+    def __init__(
+        self,
+        embedder: Embedder,
+        *,
+        chunk_size_words: int = 512,
+        batch_size: int = 32,
+    ) -> None:
+        if batch_size < 1:
+            msg = f"batch_size must be >= 1; got {batch_size}"
+            raise ValueError(msg)
+        self.embedder = embedder
+        self.chunk_size_words = chunk_size_words
+        self.batch_size = batch_size
+
+    async def index(
+        self,
+        document: Document,
+        source_text: str,
+        *,
+        out_dir: Path,
+    ) -> None:
+        chunks = self._build_chunks(document, source_text)
+        if not chunks:
+            await asyncio.to_thread(self._write_empty_table, out_dir)
+            return
+
+        vectors: list[list[float]] = []
+        for i in range(0, len(chunks), self.batch_size):
+            batch = chunks[i : i + self.batch_size]
+            raw = await self.embedder.embed([c.text for c in batch])
+            vectors.extend(l2_normalize(v) for v in raw)
+
+        await asyncio.to_thread(self._write_table, out_dir, chunks, vectors)
+
+    def _build_chunks(self, document: Document, source_text: str) -> list[_Chunk]:
+        sections = document.sections
+        chunks: list[_Chunk] = []
+        for chunk_id, (start, end, text) in enumerate(
+            chunk_text(source_text, chunk_size_words=self.chunk_size_words)
+        ):
+            chunks.append(
+                _Chunk(
+                    chunk_id=chunk_id,
+                    text=text,
+                    byte_start=start,
+                    byte_end=end,
+                    section_id=assign_section(start, end, sections),
+                )
+            )
+        return chunks
+
+    def _write_table(
+        self,
+        out_dir: Path,
+        chunks: list[_Chunk],
+        vectors: list[list[float]],
+    ) -> None:
+        db_path = out_dir / "naive.lance"
+        if db_path.exists():
+            shutil.rmtree(db_path)
+        db = lancedb.connect(str(db_path))
+        schema = pa.schema(
+            [
+                pa.field("chunk_id", pa.int64()),
+                pa.field("section_id", pa.string()),
+                pa.field("text", pa.string()),
+                pa.field("vector", pa.list_(pa.float32(), self.embedder.dim)),
+            ]
+        )
+        table = db.create_table(self.table_name, schema=schema)
+        records: list[dict[str, Any]] = [
+            {
+                "chunk_id": c.chunk_id,
+                "section_id": c.section_id or "",
+                "text": c.text,
+                "vector": v,
+            }
+            for c, v in zip(chunks, vectors, strict=True)
+        ]
+        table.add(records)
+
+    def _write_empty_table(self, out_dir: Path) -> None:
+        db_path = out_dir / "naive.lance"
+        if db_path.exists():
+            shutil.rmtree(db_path)
+        db = lancedb.connect(str(db_path))
+        schema = pa.schema(
+            [
+                pa.field("chunk_id", pa.int64()),
+                pa.field("section_id", pa.string()),
+                pa.field("text", pa.string()),
+                pa.field("vector", pa.list_(pa.float32(), self.embedder.dim)),
+            ]
+        )
+        db.create_table(self.table_name, schema=schema)
+
+    async def retrieve(
+        self,
+        query: str,
+        *,
+        out_dir: Path,
+        k: int = 8,
+    ) -> list[NaiveHit]:
+        if k < 1:
+            msg = f"k must be >= 1; got {k}"
+            raise ValueError(msg)
+        embedded = await self.embedder.embed([query])
+        if not embedded:
+            return []
+        query_vec = l2_normalize(embedded[0])
+        return await asyncio.to_thread(self._sync_retrieve, out_dir, query_vec, k)
+
+    def _sync_retrieve(
+        self,
+        out_dir: Path,
+        query_vec: list[float],
+        k: int,
+    ) -> list[NaiveHit]:
+        db = lancedb.connect(str(out_dir / "naive.lance"))
+        table = db.open_table(self.table_name)
+        rows = (
+            table.search(query_vec)
+            .distance_type("cosine")
+            .limit(k)
+            .to_list()
+        )
+        hits: list[NaiveHit] = []
+        for row in rows:
+            distance = float(row["_distance"])
+            score = max(0.0, min(1.0, 1.0 - distance))
+            section_id = row.get("section_id") or None
+            hits.append(
+                NaiveHit(
+                    chunk_id=int(row["chunk_id"]),
+                    section_id=section_id,
+                    text=str(row["text"]),
+                    score=score,
+                )
+            )
+        return hits

cairn/bench/dataset.py

@@ -0,0 +1,109 @@
+"""Bench dataset format — Pydantic models + TOML loader.
+
+A suite lives in a single TOML file (loaded via stdlib ``tomllib``). The
+shape::
+
+    name = "Cairn Architecture v0.2.2"
+
+    [[documents]]
+    id = "architecture"
+    source = "ARCHITECTURE.md"
+
+    [[documents.questions]]
+    id = "vector-store"
+    question = "What is the default vector store?"
+    expected_anchors = ["2-5-vectors-v-semantic-overlay", "7-tech-stack"]
+    tags = ["definition", "tech-stack"]
+
+``expected_anchors`` use substring matching against retrieved section_ids
+so authors can use short suffixes instead of full hierarchical slugs.
+"""
+
+from __future__ import annotations
+
+import tomllib
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from cairn.core.errors import ConfigError
+
+
+class BenchQuestion(BaseModel):
+    """One bench question with ground-truth retrieval targets."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    id: str
+    question: str
+    expected_anchors: tuple[str, ...] = Field(
+        default=(),
+        description=(
+            "Substrings to match against retrieved section_ids. "
+            "Empty tuple = no recall computed for this question."
+        ),
+    )
+    tags: tuple[str, ...] = ()
+    reference: str | None = Field(
+        default=None,
+        description="Optional free-form reference answer for LLM-judged QA.",
+    )
+
+
+class BenchDocument(BaseModel):
+    """One source document and the questions asked of it."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    id: str
+    source: Path = Field(
+        description="Path to the source document, relative to the suite file."
+    )
+    questions: tuple[BenchQuestion, ...]
+
+
+class BenchSuite(BaseModel):
+    """A complete benchmark suite."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    name: str
+    documents: tuple[BenchDocument, ...]
+
+
+def load_suite(path: Path) -> BenchSuite:
+    """Parse a TOML bench file. Source paths are resolved relative to it."""
+    if not path.exists():
+        msg = f"bench suite not found: {path}"
+        raise ConfigError(msg, details={"path": str(path)})
+
+    with path.open("rb") as fh:
+        payload = tomllib.load(fh)
+
+    suite_dir = path.parent
+    documents_in = payload.get("documents", [])
+    documents: list[BenchDocument] = []
+    for doc in documents_in:
+        questions_in = doc.get("questions", [])
+        questions = [
+            BenchQuestion(
+                id=q["id"],
+                question=q["question"],
+                expected_anchors=tuple(q.get("expected_anchors", ())),
+                tags=tuple(q.get("tags", ())),
+                reference=q.get("reference"),
+            )
+            for q in questions_in
+        ]
+        documents.append(
+            BenchDocument(
+                id=doc["id"],
+                source=(suite_dir / doc["source"]).resolve(),
+                questions=tuple(questions),
+            )
+        )
+
+    return BenchSuite(
+        name=payload.get("name", path.stem),
+        documents=tuple(documents),
+    )

cairn/bench/judge.py

@@ -0,0 +1,126 @@
+"""LLM-as-judge: optional QA accuracy dimension for cairn-bench.
+
+For each question + retrieved context the judge:
+
+1. Asks an LLM to answer the question using only the provided context.
+2. Asks the same LLM to evaluate whether that answer matches the
+   reference answer the suite author supplied.
+
+Both calls go through an OpenAI-compatible ``/v1/chat/completions``
+endpoint, so the judge runs against Ollama (default), OpenAI, vLLM, etc.
+
+The judge is **optional**: when no ``LLMJudge`` is configured on the
+:class:`cairn.bench.runner.BenchRunner`, QA accuracy is simply not
+reported. Recall and token cost still come out either way.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from cairn.core.errors import IndexBuildError
+
+_ANSWER_SYSTEM = (
+    "You answer questions strictly from the provided context. "
+    'If the context is insufficient, reply with "I don\'t know."'
+)
+
+_JUDGE_SYSTEM = (
+    "You evaluate whether an AI assistant's answer is correct given a "
+    "reference answer. Reply with YES or NO on the first line, then one "
+    "sentence of justification."
+)
+
+
+class LLMJudge:
+    """Generates answers from retrieved context and judges them against a reference."""
+
+    def __init__(
+        self,
+        *,
+        base_url: str = "http://localhost:11434/v1",
+        model: str = "llama3.2:3b",
+        api_key: str | None = None,
+        timeout: float = 60.0,
+        temperature: float = 0.0,
+    ) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.model = model
+        self.api_key = api_key
+        self.timeout = timeout
+        self.temperature = temperature
+        self.name = f"openai-compat:{model}"
+
+    async def answer(self, question: str, context: str) -> str:
+        """Produce an answer to ``question`` using only ``context``."""
+        prompt = (
+            "Context:\n"
+            f"{context.strip() or '(no context provided)'}\n\n"
+            f"Question: {question}\n\nAnswer:"
+        )
+        return await self._chat(
+            [
+                {"role": "system", "content": _ANSWER_SYSTEM},
+                {"role": "user", "content": prompt},
+            ]
+        )
+
+    async def judge(
+        self,
+        question: str,
+        reference: str,
+        answer: str,
+    ) -> tuple[bool, str]:
+        """Return ``(is_correct, raw_response)`` for one (question, answer) pair."""
+        prompt = (
+            f"Question: {question}\n\n"
+            f"Reference answer: {reference}\n\n"
+            f"Assistant's answer: {answer}\n\n"
+            "Is the assistant's answer correct?"
+        )
+        raw = await self._chat(
+            [
+                {"role": "system", "content": _JUDGE_SYSTEM},
+                {"role": "user", "content": prompt},
+            ]
+        )
+        first_line = raw.strip().split("\n", 1)[0].strip().upper()
+        is_correct = first_line.startswith("YES")
+        return is_correct, raw
+
+    async def _chat(self, messages: list[dict[str, str]]) -> str:
+        headers = {"Content-Type": "application/json"}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        payload: dict[str, Any] = {
+            "model": self.model,
+            "temperature": self.temperature,
+            "messages": messages,
+        }
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.post(
+                f"{self.base_url}/chat/completions",
+                json=payload,
+                headers=headers,
+            )
+            if response.status_code >= 400:
+                msg = (
+                    f"judge endpoint returned HTTP {response.status_code}: "
+                    f"{response.text[:200]}"
+                )
+                raise IndexBuildError(
+                    msg,
+                    details={
+                        "status": response.status_code,
+                        "model": self.model,
+                        "base_url": self.base_url,
+                    },
+                )
+            data = response.json()
+        try:
+            return str(data["choices"][0]["message"]["content"]).strip()
+        except (KeyError, IndexError, TypeError) as exc:
+            msg = "judge response did not match OpenAI chat-completions shape"
+            raise IndexBuildError(msg, details={"response": data}) from exc

cairn/bench/metrics.py

@@ -0,0 +1,32 @@
+"""Recall@k computation for bench questions.
+
+An ``expected_anchor`` is considered matched if it appears as a substring
+of any retrieved section id within the top ``k`` results. This is a
+deliberate convenience: authoring full hierarchical slugs in YAML/TOML is
+brittle, so we accept short suffixes that uniquely identify a section.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+
+def recall_at_k(
+    retrieved_section_ids: Sequence[str],
+    expected_anchors: Sequence[str],
+    *,
+    k: int,
+) -> float:
+    """Return the fraction of expected anchors found in the top-k retrieval.
+
+    Returns ``1.0`` when ``expected_anchors`` is empty (vacuously true).
+    """
+    if not expected_anchors:
+        return 1.0
+
+    top_k = retrieved_section_ids[:k]
+    matched = 0
+    for expected in expected_anchors:
+        if any(expected in retrieved for retrieved in top_k):
+            matched += 1
+    return matched / len(expected_anchors)

cairn/bench/report.py

@@ -0,0 +1,143 @@
+"""Bench result types + JSON / markdown report writers."""
+
+from __future__ import annotations
+
+import json
+import statistics
+from collections.abc import Iterable
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+System = Literal["cairn", "naive"]
+
+
+class SystemResult(BaseModel):
+    """One system's result for one question."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    system: System
+    section_ids: tuple[str, ...]
+    recall_at_k: float = Field(ge=0.0, le=1.0)
+    tokens_returned: int = Field(ge=0)
+
+
+class QuestionResult(BaseModel):
+    """A question's results across both systems."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    document_id: str
+    question_id: str
+    question: str
+    expected_anchors: tuple[str, ...]
+    tags: tuple[str, ...]
+    cairn: SystemResult
+    naive: SystemResult
+
+
+class BenchSummary(BaseModel):
+    """Full bench output."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    suite_name: str
+    k: int
+    questions: tuple[QuestionResult, ...]
+
+    def cairn_mean_recall(self) -> float:
+        return _mean(q.cairn.recall_at_k for q in self.questions)
+
+    def naive_mean_recall(self) -> float:
+        return _mean(q.naive.recall_at_k for q in self.questions)
+
+    def cairn_mean_tokens(self) -> float:
+        return _mean(float(q.cairn.tokens_returned) for q in self.questions)
+
+    def naive_mean_tokens(self) -> float:
+        return _mean(float(q.naive.tokens_returned) for q in self.questions)
+
+
+def _mean(values: Iterable[float]) -> float:
+    materialized = list(values)
+    return statistics.mean(materialized) if materialized else 0.0
+
+
+# ---------------------------------------------------------------------------
+# Report writers
+# ---------------------------------------------------------------------------
+
+
+def write_json_report(summary: BenchSummary, path: Path) -> Path:
+    """Write a deterministic JSON report next to the bench's working dir."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "generated_at": datetime.now(UTC).isoformat(),
+        "suite": summary.suite_name,
+        "k": summary.k,
+        "summary": {
+            "cairn_mean_recall_at_k": summary.cairn_mean_recall(),
+            "naive_mean_recall_at_k": summary.naive_mean_recall(),
+            "cairn_mean_tokens": summary.cairn_mean_tokens(),
+            "naive_mean_tokens": summary.naive_mean_tokens(),
+        },
+        "questions": [q.model_dump(mode="json") for q in summary.questions],
+    }
+    with path.open("w", encoding="utf-8") as fh:
+        json.dump(payload, fh, ensure_ascii=False, indent=2)
+        fh.write("\n")
+    return path
+
+
+def format_markdown_report(summary: BenchSummary) -> str:
+    """Render a human-readable comparison table."""
+    n = len(summary.questions)
+    cairn_recall = summary.cairn_mean_recall()
+    naive_recall = summary.naive_mean_recall()
+    cairn_tokens = summary.cairn_mean_tokens()
+    naive_tokens = summary.naive_mean_tokens()
+
+    token_ratio = (
+        cairn_tokens / naive_tokens if naive_tokens > 0 else float("nan")
+    )
+
+    lines: list[str] = []
+    lines.append(f"# {summary.suite_name}")
+    lines.append("")
+    lines.append(f"Questions: **{n}** · k = **{summary.k}**")
+    lines.append("")
+    lines.append("## Headline")
+    lines.append("")
+    lines.append("| metric | naive vector RAG | Cairn |")
+    lines.append("|---|---:|---:|")
+    lines.append(
+        f"| mean recall@{summary.k} | {naive_recall:.2%} | **{cairn_recall:.2%}** |"
+    )
+    lines.append(
+        f"| mean tokens returned | {naive_tokens:,.0f} | **{cairn_tokens:,.0f}** "
+        f"({token_ratio:.1%} of naive) |"
+    )
+    lines.append("")
+
+    lines.append("## Per-question")
+    lines.append("")
+    lines.append("| question | recall (cairn) | recall (naive) | tokens (cairn / naive) |")
+    lines.append("|---|---:|---:|---:|")
+    for q in summary.questions:
+        lines.append(
+            f"| `{q.question_id}` {_truncate(q.question, 60)} "
+            f"| {q.cairn.recall_at_k:.0%} "
+            f"| {q.naive.recall_at_k:.0%} "
+            f"| {q.cairn.tokens_returned} / {q.naive.tokens_returned} |"
+        )
+
+    return "\n".join(lines) + "\n"
+
+
+def _truncate(text: str, width: int) -> str:
+    if len(text) <= width:
+        return text
+    return text[: width - 1] + "…"