docsgraph 0.1.0a2__py3-none-any.whl

cairn/tools/search_semantic.py

@@ -0,0 +1,181 @@
+"""``search_semantic`` retrieval tool.
+
+Spec: ``docs/specs/mcp-tools.md`` §4.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+from typing import Any, Literal
+
+from cairn.core.errors import ToolError
+from cairn.embed.base import Embedder
+from cairn.tools.base import DocumentIndex, ToolResponse, estimate_tokens_of_payload
+
+IncludeField = Literal["synopsis", "head", "evidence"]
+
+_HEAD_CHARS: int = 200
+_EVIDENCE_CHARS: int = 360
+_MAX_EVIDENCE_TERMS: int = 40
+_VALID_INCLUDE: frozenset[str] = frozenset({"synopsis", "head", "evidence"})
+
+
+async def search_semantic(
+    index: DocumentIndex,
+    *,
+    embedder: Embedder,
+    query: str,
+    scope: str | None = None,
+    k: int = 8,
+    include: Sequence[IncludeField] = ("synopsis", "head", "evidence"),
+) -> ToolResponse:
+    """Dense vector search across the document.
+
+    The MCP server passes an :class:`Embedder` instance; tools receive it as
+    a typed dependency so they can be unit-tested with a fake.
+    """
+    if k < 1 or k > 32:
+        msg = f"k must be in [1, 32]; got {k}"
+        raise ToolError(msg, details={"k": k})
+    if not query.strip():
+        msg = "query must not be empty"
+        raise ToolError(msg)
+    bad = [x for x in include if x not in _VALID_INCLUDE]
+    if bad:
+        msg = f"invalid include values: {bad}"
+        raise ToolError(msg, details={"invalid": bad})
+
+    vectors = await embedder.embed([query])
+    if not vectors:
+        msg = "embedder returned no vector for query"
+        raise ToolError(msg)
+    query_vec = vectors[0]
+    if len(query_vec) != index.vectors.dim:
+        msg = (
+            f"query embedding dim {len(query_vec)} != "
+            f"index dim {index.vectors.dim}"
+        )
+        raise ToolError(msg)
+
+    hits = await index.vectors.search(query_vec, k=k, scope_prefix=scope)
+
+    include_set = set(include)
+    results: list[dict[str, Any]] = []
+    for hit in hits:
+        node = index.tree.get(hit.id)
+        if node is None:
+            # Stale vector index entry: skip rather than fail the whole call.
+            continue
+        summary = index.summaries.get(hit.id)
+        result: dict[str, Any] = {
+            "id": hit.id,
+            "title": node.title,
+            "score": hit.score,
+            "anchor": index.anchor(hit.id),
+        }
+        if "synopsis" in include_set and summary is not None and summary.synopsis:
+            result["synopsis"] = summary.synopsis
+        if "head" in include_set:
+            result["head"] = node.raw_text[:_HEAD_CHARS]
+        if "evidence" in include_set:
+            result["evidence"] = _evidence_snippet(node.raw_text, query)
+        results.append(result)
+
+    payload: dict[str, Any] = {
+        "query": query,
+        "scope": scope,
+        "hits": results,
+        "cursor": None,
+    }
+    return ToolResponse(
+        data=payload,
+        tokens_returned=estimate_tokens_of_payload(payload),
+    )
+
+
+def _evidence_snippet(text: str, query: str) -> dict[str, Any]:
+    """Return the strongest lexical evidence window for a semantic hit.
+
+    Semantic rank is vector-based, but a short lexical window gives humans and
+    agents a cheap explanation of what in the section may have caused the hit.
+    When there is no term overlap, fall back to the start of the section.
+    """
+    clean_text = text.strip()
+    if not clean_text:
+        return {"text": "", "matched_terms": [], "span": {"start": 0, "end": 0}}
+
+    terms = _query_terms(query)
+    best_start = 0
+    best_score = 0
+    lowered = clean_text.lower()
+    if terms:
+        candidate_starts: set[int] = {0}
+        for term in terms:
+            for match in re.finditer(re.escape(term), lowered):
+                candidate_starts.add(max(0, match.start() - _EVIDENCE_CHARS // 3))
+        for start in candidate_starts:
+            end = min(len(clean_text), start + _EVIDENCE_CHARS)
+            window = lowered[start:end]
+            score = sum(window.count(term) for term in terms)
+            if score > best_score:
+                best_score = score
+                best_start = start
+
+    start = best_start
+    end = min(len(clean_text), start + _EVIDENCE_CHARS)
+    snippet = clean_text[start:end]
+    if start > 0:
+        snippet = "..." + snippet.lstrip()
+    if end < len(clean_text):
+        snippet = snippet.rstrip() + "..."
+
+    matched = [term for term in terms if term in lowered[start:end]]
+    return {
+        "text": snippet,
+        "matched_terms": matched,
+        "span": {"start": start, "end": end},
+    }
+
+
+def _query_terms(query: str) -> list[str]:
+    query = query.lower()
+    stop = {
+        "a",
+        "an",
+        "and",
+        "are",
+        "for",
+        "how",
+        "is",
+        "of",
+        "or",
+        "the",
+        "to",
+        "what",
+        "where",
+    }
+    seen: set[str] = set()
+    out: list[str] = []
+
+    def add(term: str) -> None:
+        if term in seen or len(out) >= _MAX_EVIDENCE_TERMS:
+            return
+        seen.add(term)
+        out.append(term)
+
+    words = re.findall(r"[A-Za-z0-9_][A-Za-z0-9_-]*", query)
+    for word in words:
+        if len(word) < 3 or word in stop or word in seen:
+            continue
+        add(word)
+
+    for seq in re.findall(r"[\u3400-\u9fff]+", query):
+        if len(seq) >= 2:
+            add(seq)
+        # CJK queries often have no whitespace; bounded n-grams give the
+        # evidence window useful overlap without changing vector ranking.
+        for size in range(min(6, len(seq)), 1, -1):
+            for start in range(0, len(seq) - size + 1):
+                add(seq[start : start + size])
+    return out

cairn/xref/__init__.py

@@ -0,0 +1,24 @@
+"""Cross-reference extraction — directed edges between sections.
+
+Three kinds of edges (per ARCHITECTURE.md §2.4):
+
+- ``link``: explicit Markdown anchor links (``[text](#anchor)``)
+- ``textual``: numeric section references (``"§ 2.5"``, ``"Section 3"``)
+- ``entity``: sections that share a high-signal defined entity
+
+The :class:`XRefExtractor` Protocol is the seam. The default
+:class:`HeuristicXRefExtractor` produces all three kinds without any model
+dependency, and accepts an optional Entities reader for entity-mediated
+edges. LLM-verified textual references are a v0.2.3+ refinement.
+"""
+
+from cairn.xref.base import ExtractionEdge, XRefExtractor
+from cairn.xref.fake import FakeXRefExtractor
+from cairn.xref.heuristic import HeuristicXRefExtractor
+
+__all__ = [
+    "ExtractionEdge",
+    "FakeXRefExtractor",
+    "HeuristicXRefExtractor",
+    "XRefExtractor",
+]

cairn/xref/base.py

@@ -0,0 +1,50 @@
+"""XRefExtractor Protocol + intermediate ExtractionEdge type."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Iterable
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from cairn.core.types import Document, Span, XRefKind
+from cairn.index.entities import Entities
+
+
+class ExtractionEdge(BaseModel):
+    """One observed cross-reference, before deduplication.
+
+    Spans use the same convention as :class:`cairn.entity.base.ExtractionHit`:
+    offsets within the *source section's* ``raw_text``. Self-loops (``src ==
+    dst``) are dropped by the :class:`cairn.index.xrefs.XRefBuilder`; do not
+    emit them.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    src: str
+    dst: str
+    kind: XRefKind
+    confidence: float = Field(ge=0.0, le=1.0)
+    span: Span
+
+
+@runtime_checkable
+class XRefExtractor(Protocol):
+    """Pluggable cross-reference extractor.
+
+    The default :class:`cairn.xref.heuristic.HeuristicXRefExtractor` accepts
+    an optional ``Entities`` reader as a constructor argument; the Protocol
+    itself only mandates the ``extract`` shape.
+    """
+
+    name: str
+
+    def extract(
+        self,
+        document: Document,
+        *,
+        entities: Entities | None = None,
+    ) -> Awaitable[Iterable[ExtractionEdge]]:
+        """Return an iterable of cross-reference edges across ``document``."""
+        ...

cairn/xref/fake.py

@@ -0,0 +1,40 @@
+"""Deterministic cross-reference extractor for tests.
+
+Emits one edge between each consecutive pair of sections (in document
+order). Kind is always ``link``. Useful for tests that exercise builder /
+reader / tool behavior without coupling to extraction heuristics.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from cairn.core.types import Document, Span
+from cairn.index.entities import Entities
+from cairn.xref.base import ExtractionEdge
+
+
+class FakeXRefExtractor:
+    """Linear edge between consecutive sections."""
+
+    name = "fake:linear"
+
+    async def extract(
+        self,
+        document: Document,
+        *,
+        entities: Entities | None = None,
+    ) -> Iterable[ExtractionEdge]:
+        edges: list[ExtractionEdge] = []
+        sections = document.sections
+        for i in range(len(sections) - 1):
+            edges.append(
+                ExtractionEdge(
+                    src=sections[i].id,
+                    dst=sections[i + 1].id,
+                    kind="link",
+                    confidence=1.0,
+                    span=Span(start=0, end=0),
+                )
+            )
+        return edges

cairn/xref/heuristic.py

@@ -0,0 +1,217 @@
+"""HeuristicXRefExtractor — regex + entity-graph derivation, no model needed.
+
+Combines three sources into one extractor (per ARCHITECTURE.md §2.4):
+
+- **link**: explicit anchor links. ``[text](#anchor)`` resolves to a section
+  whose ``id`` ends in ``anchor``. Confidence 0.95 when unique, 0.75 when
+  multiple candidates exist.
+- **textual**: numeric references like ``§ 2.5`` or ``Section 3.1``. Mapped
+  to a section whose ``title`` starts with the same numeric prefix.
+  Confidence 0.7.
+- **entity**: pairs of sections that share a high-signal *defined* entity.
+  Confidence scales with the number of shared entities, capped at 0.8.
+  Computed only when an :class:`~cairn.index.entities.Entities` reader is
+  supplied.
+
+Self-loops are filtered. Duplicate ``(src, dst, kind)`` triples are
+deduplicated by the :class:`~cairn.index.xrefs.XRefBuilder`; this layer just
+emits ``ExtractionEdge`` records.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable, Iterator
+from typing import Final
+
+from cairn.core.types import Document, SectionNode, Span
+from cairn.index.entities import Entities
+from cairn.xref.base import ExtractionEdge
+
+# Markdown anchor link: [text](#anchor) — anchor uses kebab/slug form.
+_ANCHOR_LINK = re.compile(r"\[[^\]]+\]\(#([^)\s]+)\)")
+
+# Section reference: "§ 2.5", "Section 3.1", "Chapter 4.2". Captures the
+# numeric prefix only.
+_SECTION_REF = re.compile(
+    r"(?:§\s*|(?:Section|Chapter|§)\s+)(\d+(?:\.\d+)*)\b",
+    re.IGNORECASE,
+)
+
+
+# Confidence scores per ARCHITECTURE.md §2.4.
+_LINK_CONF_UNIQUE: Final = 0.95
+_LINK_CONF_AMBIGUOUS: Final = 0.75
+_TEXTUAL_CONF: Final = 0.7
+_ENTITY_CONF_BASE: Final = 0.3
+_ENTITY_CONF_STEP: Final = 0.2
+_ENTITY_CONF_CAP: Final = 0.8
+
+
+class HeuristicXRefExtractor:
+    """Regex + entity-graph cross-reference extractor."""
+
+    name = "heuristic:xref-v1"
+
+    async def extract(
+        self,
+        document: Document,
+        *,
+        entities: Entities | None = None,
+    ) -> Iterable[ExtractionEdge]:
+        sections = document.sections
+        if not sections:
+            return []
+
+        # Pre-build lookups used by link + textual extractors.
+        anchor_to_ids = _build_anchor_index(sections)
+        prefix_to_id = _build_prefix_index(sections)
+
+        edges: list[ExtractionEdge] = []
+        for section in sections:
+            edges.extend(_scan_links(section, anchor_to_ids))
+            edges.extend(_scan_textual(section, prefix_to_id))
+
+        if entities is not None:
+            edges.extend(_entity_mediated(sections, entities))
+
+        return edges
+
+
+# ---------------------------------------------------------------------------
+# Link extraction
+# ---------------------------------------------------------------------------
+
+
+def _build_anchor_index(sections: tuple[SectionNode, ...]) -> dict[str, list[str]]:
+    """Map heading-anchor → list of section_ids whose last slug equals it."""
+    index: dict[str, list[str]] = {}
+    for section in sections:
+        last = section.id.rsplit("/", 1)[-1]
+        index.setdefault(last, []).append(section.id)
+    return index
+
+
+def _scan_links(
+    section: SectionNode,
+    anchor_to_ids: dict[str, list[str]],
+) -> Iterator[ExtractionEdge]:
+    for match in _ANCHOR_LINK.finditer(section.raw_text):
+        anchor = match.group(1).lower()
+        candidates = anchor_to_ids.get(anchor, ())
+        if not candidates:
+            continue
+        unique = len(candidates) == 1
+        conf = _LINK_CONF_UNIQUE if unique else _LINK_CONF_AMBIGUOUS
+        for dst in candidates:
+            if dst == section.id:
+                continue
+            yield ExtractionEdge(
+                src=section.id,
+                dst=dst,
+                kind="link",
+                confidence=conf,
+                span=Span(start=match.start(1), end=match.end(1)),
+            )
+
+
+# ---------------------------------------------------------------------------
+# Textual extraction
+# ---------------------------------------------------------------------------
+
+
+def _build_prefix_index(sections: tuple[SectionNode, ...]) -> dict[str, str]:
+    """Map numeric title prefix (``"2.5"``) → section_id.
+
+    Only sections whose title starts with a digit-sequence are indexed.
+    First-seen wins on collision.
+    """
+    index: dict[str, str] = {}
+    for section in sections:
+        prefix = _numeric_prefix(section.title)
+        if prefix and prefix not in index:
+            index[prefix] = section.id
+    return index
+
+
+_TITLE_PREFIX = re.compile(r"^\s*(\d+(?:\.\d+)*)\b")
+
+
+def _numeric_prefix(title: str) -> str | None:
+    m = _TITLE_PREFIX.match(title)
+    return m.group(1) if m else None
+
+
+def _scan_textual(
+    section: SectionNode,
+    prefix_to_id: dict[str, str],
+) -> Iterator[ExtractionEdge]:
+    for match in _SECTION_REF.finditer(section.raw_text):
+        number = match.group(1)
+        dst = prefix_to_id.get(number)
+        if dst is None or dst == section.id:
+            continue
+        yield ExtractionEdge(
+            src=section.id,
+            dst=dst,
+            kind="textual",
+            confidence=_TEXTUAL_CONF,
+            span=Span(start=match.start(1), end=match.end(1)),
+        )
+
+
+# ---------------------------------------------------------------------------
+# Entity-mediated extraction
+# ---------------------------------------------------------------------------
+
+
+def _entity_mediated(
+    sections: tuple[SectionNode, ...],
+    entities: Entities,
+) -> Iterator[ExtractionEdge]:
+    """Emit edges between sections that share defined entities.
+
+    For each defined entity that appears in 2+ sections, every ordered
+    section pair contributes one (or one extra) shared count. Confidence
+    rises with the count: 1 → 0.5, 2 → 0.7, 3+ → 0.8 (cap).
+    """
+    # shared_counts[(src, dst)] = number of distinct defined entities in common
+    shared: dict[tuple[str, str], int] = {}
+    sample_span: dict[tuple[str, str], Span] = {}
+
+    for ent in entities.by_kind("defined"):
+        mention_sections = {m.section_id for m in ent.mentions}
+        if len(mention_sections) < 2:
+            continue
+        ordered = sorted(mention_sections)
+        for i, src in enumerate(ordered):
+            for dst in ordered[i + 1 :]:
+                shared[(src, dst)] = shared.get((src, dst), 0) + 1
+                shared[(dst, src)] = shared.get((dst, src), 0) + 1
+                # Pick *some* span — the first mention in the src section.
+                if (src, dst) not in sample_span:
+                    for m in ent.mentions:
+                        if m.section_id == src:
+                            sample_span[(src, dst)] = m.span
+                            break
+                if (dst, src) not in sample_span:
+                    for m in ent.mentions:
+                        if m.section_id == dst:
+                            sample_span[(dst, src)] = m.span
+                            break
+
+    for (src, dst), count in shared.items():
+        if src == dst:
+            continue
+        conf = min(
+            _ENTITY_CONF_CAP,
+            _ENTITY_CONF_BASE + _ENTITY_CONF_STEP * count,
+        )
+        span = sample_span.get((src, dst), Span(start=0, end=0))
+        yield ExtractionEdge(
+            src=src,
+            dst=dst,
+            kind="entity",
+            confidence=conf,
+            span=span,
+        )