aid-runtime-python 0.1.39__py3-none-any.whl

aid_runtime/retrieval/chunking.py

@@ -0,0 +1,222 @@
+"""Turning documents into deterministic, size-bounded chunks.
+
+The IR declares a strategy, a size and an overlap; it does not declare the split
+algorithm. This module pins one, and the pin is the spec validator's own reading
+of the intent: *"spec v0.1 chunks by size (fixed, recursive, code)"*. ``semantic``
+is refused because it needs an embedding pass over the whole source at ingest --
+a different shape of job, not a different boundary rule.
+
+So all three implemented strategies are size-bounded splitters and the only thing
+that distinguishes them is **where a boundary is allowed to fall**:
+
+``fixed``
+    Anywhere. The window always fills to ``size``, cutting mid-word when that is
+    where the budget ends. This is the honest baseline other strategies improve on.
+``recursive``
+    Prefer the last paragraph break, then the last line break, then the last
+    sentence end, then the last word break inside the window.
+``code``
+    Prefer the last blank line, then the last line break, then the last statement
+    end. Code is line-structured, so a blank line is the closest thing it has to a
+    paragraph, and cutting mid-statement is the failure worth avoiding.
+
+Chunks are pure functions of the text, so the same input always produces the same
+chunks with the same ids in the same order -- which is what makes re-ingesting a
+corpus a no-op instead of a duplicate.
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Sequence
+from typing import Protocol, runtime_checkable
+
+from aid_runtime.errors import ConfigurationError, FeatureDeferredError
+from aid_runtime.retrieval.connectors import SourceDocument
+from aid_runtime.retrieval.types import Chunk, ChunkingSpec, ChunkingStrategy, chunk_id_for
+
+__all__ = [
+    "Chunker",
+    "CodeChunker",
+    "FixedChunker",
+    "RecursiveChunker",
+    "chunk_document",
+    "chunker_for",
+]
+
+#: A boundary found in the first half of a window is ignored and the window is
+#: filled to ``size`` instead. Without this, one early paragraph break near the end
+#: of a split would produce a stream of very short chunks, and a chunker that emits
+#: 5-character chunks is worse than one that cuts a sentence in half.
+_MIN_FILL_NUMERATOR = 1
+_MIN_FILL_DENOMINATOR = 2
+
+
+@runtime_checkable
+class Chunker(Protocol):
+    """Splits a document into chunks, deterministically."""
+
+    @property
+    def strategy(self) -> ChunkingStrategy: ...
+
+    def chunk(self, document: SourceDocument) -> tuple[Chunk, ...]: ...
+
+
+class _SizeChunker(abc.ABC):
+    """Shared windowing. Subclasses only choose where a boundary may fall."""
+
+    def __init__(self, *, size: int, overlap: int, keep_metadata: Sequence[str] = ()) -> None:
+        if size <= 0:
+            raise ConfigurationError("chunking size must be positive")
+        if overlap < 0:
+            raise ConfigurationError("chunking overlap must not be negative")
+        if overlap >= size:
+            raise ConfigurationError(
+                f"chunking overlap {overlap} must be smaller than size {size}; otherwise a "
+                "window would not advance and chunking would never terminate"
+            )
+        self._size = size
+        self._overlap = overlap
+        self._keep = tuple(keep_metadata)
+
+    @property
+    @abc.abstractmethod
+    def strategy(self) -> ChunkingStrategy: ...
+
+    @abc.abstractmethod
+    def _cut(self, text: str, start: int, limit: int) -> int:
+        """Where to end the window that starts at ``start`` and may reach ``limit``."""
+
+    def chunk(self, document: SourceDocument) -> tuple[Chunk, ...]:
+        text = document.text
+        chunks: list[Chunk] = []
+        start = 0
+        while start < len(text):
+            limit = min(start + self._size, len(text))
+            end = limit if limit >= len(text) else self._cut(text, start, limit)
+            # A boundary at or before `start` would stall the loop, so progress is
+            # guaranteed here rather than trusted from the strategy hook.
+            end = max(end, start + 1)
+            window = text[start:end]
+            if window.strip():
+                chunks.append(self._build(document, len(chunks), start, end, window))
+            if end >= len(text):
+                break
+            start = max(end - self._overlap, start + 1)
+        return tuple(chunks)
+
+    def _build(
+        self, document: SourceDocument, index: int, start: int, end: int, text: str
+    ) -> Chunk:
+        return Chunk(
+            chunk_id=chunk_id_for(document.doc_id, index, text),
+            doc_id=document.doc_id,
+            uri=document.uri,
+            index=index,
+            start=start,
+            end=end,
+            text=text,
+            metadata={
+                key: document.metadata[key] for key in self._keep if key in document.metadata
+            },
+        )
+
+    def _boundary(self, text: str, start: int, limit: int, markers: Sequence[str]) -> int | None:
+        """The end offset of the last ``marker`` in ``text[start:limit]``, if usable.
+
+        ``None`` means "no boundary worth taking", and the caller falls back to
+        ``limit``. The minimum-fill rule is applied here rather than by each
+        strategy so no strategy can forget it.
+        """
+        floor = start + (self._size * _MIN_FILL_NUMERATOR) // _MIN_FILL_DENOMINATOR
+        window = text[start:limit]
+        found: int | None = None
+        for marker in markers:
+            position = window.rfind(marker)
+            if position != -1:
+                end = start + position + len(marker)
+                if end >= max(floor, start + 1):
+                    found = end
+                    break
+        return found
+
+    def _fill(
+        self, text: str, start: int, limit: int, marker_groups: Sequence[Sequence[str]]
+    ) -> int:
+        """The first usable boundary from ``marker_groups``, in preference order."""
+        for markers in marker_groups:
+            found = self._boundary(text, start, limit, markers)
+            if found is not None:
+                return found
+        return limit
+
+
+class FixedChunker(_SizeChunker):
+    """Cuts on size alone."""
+
+    @property
+    def strategy(self) -> ChunkingStrategy:
+        return ChunkingStrategy.FIXED
+
+    def _cut(self, text: str, start: int, limit: int) -> int:
+        return limit
+
+
+class RecursiveChunker(_SizeChunker):
+    """Prefers prose boundaries: paragraph, then line, then sentence, then word."""
+
+    @property
+    def strategy(self) -> ChunkingStrategy:
+        return ChunkingStrategy.RECURSIVE
+
+    def _cut(self, text: str, start: int, limit: int) -> int:
+        return self._fill(
+            text,
+            start,
+            limit,
+            (("\n\n",), ("\n",), (". ", "! ", "? "), (" ",)),
+        )
+
+
+class CodeChunker(_SizeChunker):
+    """Prefers code boundaries: blank line, then line, then statement, then token."""
+
+    @property
+    def strategy(self) -> ChunkingStrategy:
+        return ChunkingStrategy.CODE
+
+    def _cut(self, text: str, start: int, limit: int) -> int:
+        return self._fill(
+            text,
+            start,
+            limit,
+            (("\n\n",), ("\n",), (";\n", ";\r\n", "; "), (" ",)),
+        )
+
+
+_STRATEGIES: dict[ChunkingStrategy, type[_SizeChunker]] = {
+    ChunkingStrategy.FIXED: FixedChunker,
+    ChunkingStrategy.RECURSIVE: RecursiveChunker,
+    ChunkingStrategy.CODE: CodeChunker,
+}
+
+
+def chunker_for(spec: ChunkingSpec) -> Chunker:
+    """Resolve a chunking block to an implementation.
+
+    ``semantic`` is refused here as well as in the spec validator, because the
+    runtime is reachable from generated code the validator never saw.
+    """
+    implementation = _STRATEGIES.get(spec.strategy)
+    if implementation is None:
+        raise FeatureDeferredError(
+            f"chunking strategy {spec.strategy.value!r} is declared but not implemented: spec "
+            "v0.1 chunks by size (fixed, recursive, code), and semantic chunking needs an "
+            "embedding pass over the whole source at ingest"
+        )
+    return implementation(size=spec.size, overlap=spec.overlap, keep_metadata=spec.keep_metadata)
+
+
+def chunk_document(document: SourceDocument, spec: ChunkingSpec) -> tuple[Chunk, ...]:
+    """Chunk one document with the strategy the spec names."""
+    return chunker_for(spec).chunk(document)

aid_runtime/retrieval/citations.py

@@ -0,0 +1,131 @@
+"""Retrieved passages, their provenance, and the context they assemble into.
+
+A retrieved set is only useful to a model if the passages arrive with provenance
+the model can repeat back, so ranking and citation are one type here rather than
+a search result plus a separate lookup table. :class:`RetrievalResult` is the
+runtime's answer type: it is what a ``retrieve`` step produces and what every
+retrieval metric scores.
+
+The assembled context is numbered by **rank**, not by position in the tuple. The
+two are the same in practice because ranks are assigned contiguously from 1, but
+naming the rank is what makes a model's ``[2]`` resolvable by the caller without
+re-running the search.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from aid_runtime.agents.guardrails import GuardrailRecord
+from aid_runtime.retrieval.types import Chunk
+
+__all__ = [
+    "Citation",
+    "RetrievalResult",
+    "RetrievedItem",
+    "assemble_context",
+    "cite",
+]
+
+_RECORD = ConfigDict(frozen=True, extra="forbid")
+
+
+class Citation(BaseModel):
+    """Where one retrieved passage came from.
+
+    ``doc_id`` is the document identity the IR's retrieval metric already uses
+    (``expectDocs: [getting-started.md]``), so a citation and a gate expectation
+    are the same vocabulary rather than two spellings of the same thing.
+    """
+
+    model_config = _RECORD
+
+    doc_id: str = Field(min_length=1)
+    chunk_id: str = Field(min_length=1)
+    uri: str = ""
+    index: int = Field(ge=0)
+
+    def __str__(self) -> str:
+        return f"{self.doc_id}#{self.index}"
+
+
+class RetrievedItem(BaseModel):
+    """One ranked passage: the chunk, its score, and its provenance."""
+
+    model_config = _RECORD
+
+    chunk: Chunk
+    score: float
+    rank: int = Field(ge=1)
+    citation: Citation
+
+
+class RetrievalResult(BaseModel):
+    """What a ``retrieve`` step returns.
+
+    ``items`` is rank-ordered and may be empty -- "the corpus has nothing for this
+    query" is a result, not an error, and a caller that needs a non-empty answer
+    says so with a metric rather than by crashing.
+    """
+
+    model_config = _RECORD
+
+    retriever: str = Field(min_length=1)
+    query: str
+    items: tuple[RetrievedItem, ...] = ()
+    #: What the ``retrieval``-stage guardrails recorded. A ``flag`` reaches a caller
+    #: here; a ``block`` never returns a result at all.
+    guardrails: tuple[GuardrailRecord, ...] = ()
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def top(self, k: int) -> tuple[RetrievedItem, ...]:
+        """The first ``k`` items, or all of them when the result is shorter."""
+        return self.items[:k]
+
+    @property
+    def context(self) -> str:
+        """The numbered excerpts a prompt binds."""
+        return assemble_context(self.items)
+
+    @property
+    def citations(self) -> tuple[Citation, ...]:
+        """One citation per item, in rank order."""
+        return tuple(item.citation for item in self.items)
+
+    @property
+    def doc_ids(self) -> tuple[str, ...]:
+        """Rank-ordered, de-duplicated document ids.
+
+        De-duplicated because recall is about *documents found*, and a document
+        that contributed three passages was found once. First occurrence wins, so
+        the order still reflects relevance.
+        """
+        return tuple(dict.fromkeys(item.citation.doc_id for item in self.items))
+
+
+def cite(chunk: Chunk) -> Citation:
+    """The citation for a chunk.
+
+    Derived from the chunk rather than passed alongside it, so a citation cannot
+    drift from the passage it names.
+    """
+    return Citation(
+        doc_id=chunk.doc_id,
+        chunk_id=chunk.chunk_id,
+        uri=chunk.uri,
+        index=chunk.index,
+    )
+
+
+def assemble_context(items: Sequence[RetrievedItem], *, separator: str = "\n\n") -> str:
+    """Render ranked items as numbered excerpts, each tagged with its citation.
+
+    An empty retrieval renders as the empty string rather than as a header with
+    nothing under it: a header alone reads to a model as "sources were found", and
+    a prompt that invites a model to invent a source is worse than an empty prompt.
+    """
+    return separator.join(f"[{item.rank}] {item.citation} :: {item.chunk.text}" for item in items)

aid_runtime/retrieval/connectors.py

@@ -0,0 +1,131 @@
+"""Reading a ``file`` source.
+
+The IR names five source kinds and Phase 1 implements one. That is the spec
+validator's own position ("spec v0.1 indexes a local ``file`` source; the remote
+loaders belong to the phase that ships them"), so the other four raise
+:class:`~aid_runtime.errors.FeatureDeferredError` here rather than indexing an
+empty corpus and looking healthy.
+
+**Document identity is the POSIX path relative to the source root.** That is the
+vocabulary the IR's retrieval metric already uses -- its ``expectDocs`` hold
+``getting-started.md``, not ``./docs/getting-started.md`` -- and it is
+machine-independent, so the same corpus indexed on Windows and on Linux yields the
+same ids and the same eval outcome. Absolute paths would make a baseline
+non-portable.
+
+The protocol is ``async`` even though the only implementation reads local files
+synchronously, because the kinds that follow (``s3``, ``http``, ``db``, ``mcp``) are
+genuinely concurrent. Making it sync now would be a breaking change to the seam
+exactly when the first real loader arrives.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from aid_runtime.errors import ConfigurationError, FeatureDeferredError
+from aid_runtime.retrieval.types import SourceKind, SourceSpec
+
+__all__ = [
+    "Connector",
+    "FileConnector",
+    "SourceDocument",
+    "open_connector",
+]
+
+_RECORD = ConfigDict(frozen=True, extra="forbid")
+
+
+class SourceDocument(BaseModel):
+    """One document read from a source, before chunking."""
+
+    model_config = _RECORD
+
+    #: POSIX path relative to the source root; the document's identity everywhere.
+    doc_id: str = Field(min_length=1)
+    #: The absolute resolved location, for a human reading a citation. Never matched on.
+    uri: str = Field(min_length=1)
+    text: str
+    metadata: dict[str, str] = Field(default_factory=dict)
+
+
+@runtime_checkable
+class Connector(Protocol):
+    """A source that yields documents in a deterministic order."""
+
+    @property
+    def kind(self) -> SourceKind: ...
+
+    async def documents(self) -> tuple[SourceDocument, ...]: ...
+
+    async def aclose(self) -> None: ...
+
+
+class FileConnector:
+    """Reads ``*.md``-style files under a local directory.
+
+    Files are read and returned in sorted ``doc_id`` order, never in directory
+    order: directory order is a filesystem accident, and a corpus whose chunk ids
+    depend on it would reindex itself differently on a different machine.
+    """
+
+    def __init__(self, root: str | Path, *, glob: str | None = None) -> None:
+        self._root = Path(root)
+        if not self._root.exists():
+            raise ConfigurationError(f"source path {str(self._root)!r} does not exist")
+        if not self._root.is_dir():
+            raise ConfigurationError(f"source path {str(self._root)!r} is not a directory")
+        self._glob = glob
+
+    @property
+    def kind(self) -> SourceKind:
+        return SourceKind.FILE
+
+    async def documents(self) -> tuple[SourceDocument, ...]:
+        return tuple(self._read(path) for path in self._paths())
+
+    async def aclose(self) -> None:
+        return None
+
+    def _paths(self) -> tuple[Path, ...]:
+        pattern = self._glob or "**/*"
+        found = [
+            path
+            for path in self._root.glob(pattern)
+            if path.is_file() and not path.name.startswith(".")
+        ]
+        return tuple(sorted(found, key=self._doc_id))
+
+    def _doc_id(self, path: Path) -> str:
+        return path.relative_to(self._root).as_posix()
+
+    def _read(self, path: Path) -> SourceDocument:
+        doc_id = self._doc_id(path)
+        return SourceDocument(
+            doc_id=doc_id,
+            uri=str(path.resolve()),
+            text=path.read_text(encoding="utf-8"),
+            metadata={"path": doc_id, "suffix": path.suffix},
+        )
+
+
+def open_connector(spec: SourceSpec, *, base_dir: str | Path | None = None) -> Connector:
+    """Open the connector for a source entry.
+
+    ``base_dir`` resolves a relative source ``uri``; it is explicit because the
+    runtime has no opinion about a working directory, and a corpus must not depend
+    on where a process happened to be started.
+    """
+    if spec.kind is not SourceKind.FILE:
+        raise FeatureDeferredError(
+            f"source kind {spec.kind.value!r} is declared but not implemented: spec v0.1 "
+            "indexes a local 'file' source, and the remote loaders (s3, http, db, mcp) "
+            "belong to the phase that ships them"
+        )
+    root = Path(spec.uri)
+    if not root.is_absolute() and base_dir is not None:
+        root = Path(base_dir) / root
+    return FileConnector(root, glob=spec.glob)