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

aid_runtime/retrieval/metrics.py

@@ -0,0 +1,213 @@
+"""Metrics over a retrieved set.
+
+The shipped :class:`~aid_runtime.evals.metrics.Metric` scores ``(case, prediction:
+str)``. That signature cannot express these metrics: recall@k and MRR score a *ranked
+set of passages with provenance*, and a citation-grounding score has to look at where
+each passage came from. Encoding a retrieved set into one string would make every
+metric re-parse it, and the parse would become the thing under test.
+
+So the retrieval metrics get their own protocol here rather than widening the shipped
+one. Widening ``Metric`` would change the meaning of an interface the eval gate
+already threshold-checks, and would do so for every existing implementation. The new
+protocol is additive and reuses :class:`~aid_runtime.evals.metrics.MetricResult`.
+
+The two protocols share the part the gate actually consumes -- ``name`` and
+``deterministic`` -- which is what lets retrieval metrics be handed to the shipped
+gate's ``metrics=`` parameter. See :mod:`aid_runtime.retrieval.gate` for why that
+hand-off needs a cast and why the cast is sound.
+
+``deterministic`` is ``True`` for all three, and that is a claim the gate enforces:
+none of these metrics calls a model, and every tie is broken by chunk id, so the same
+retrieved set always scores the same.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from aid_runtime.evals.metrics import MetricResult
+from aid_runtime.retrieval.citations import RetrievalResult
+from aid_runtime.retrieval.types import RetrievalMetricSpec
+
+__all__ = [
+    "CitationGrounding",
+    "MeanReciprocalRank",
+    "RecallAtK",
+    "RetrievalCase",
+    "RetrievalMetric",
+    "retrieval_metrics",
+]
+
+
+class RetrievalCase(BaseModel):
+    """One retrieval question and the documents that should answer it.
+
+    This mirrors the IR's ``retrieval`` metric node field for field, and it exists
+    because :class:`~aid_runtime.evals.dataset.EvalCase` carries ``metadata:
+    dict[str, str]`` -- a string map, which cannot hold ``expectDocs`` without
+    inventing an encoding. The expected documents, and ``k``, are what a retrieval
+    metric needs, so they are typed here instead of smuggled through a string.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    id: str = Field(min_length=1)
+    query: str
+    expect_docs: tuple[str, ...] = Field(min_length=1)
+    k: int = Field(gt=0)
+
+    @classmethod
+    def from_spec(cls, *, case_id: str, query: str, spec: RetrievalMetricSpec) -> RetrievalCase:
+        """Build a case from the IR's ``{kind, expectDocs, k}`` node."""
+        return cls(id=case_id, query=query, expect_docs=spec.expect_docs, k=spec.k)
+
+
+@runtime_checkable
+class RetrievalMetric(Protocol):
+    """Scores a ranked, cited retrieved set.
+
+    ``name`` and ``deterministic`` are part of the contract because the shipped gate
+    reads exactly those two properties when it checks that a thresholded metric is
+    allowed to gate CI.
+    """
+
+    @property
+    def name(self) -> str: ...
+
+    @property
+    def deterministic(self) -> bool: ...
+
+    async def score(self, case: RetrievalCase, retrieved: RetrievalResult) -> MetricResult: ...
+
+
+class RecallAtK:
+    """Fraction of expected documents present in the top ``k``.
+
+    Duplicated passages do not inflate the score: :attr:`RetrievalResult.doc_ids` is
+    de-duplicated, so a document that contributed three passages counts once. Recall
+    is about which documents were found.
+    """
+
+    def __init__(self, *, name: str = "recall@k", min_score: float = 1.0) -> None:
+        self._name = name
+        self._min_score = min_score
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def deterministic(self) -> bool:
+        return True
+
+    async def score(self, case: RetrievalCase, retrieved: RetrievalResult) -> MetricResult:
+        expected = set(case.expect_docs)
+        found = expected & set(retrieved.doc_ids[: case.k])
+        score = len(found) / len(expected)
+        missing = sorted(expected - found)
+        return MetricResult(
+            metric=self.name,
+            score=score,
+            passed=score >= self._min_score,
+            detail="" if not missing else f"missing: {', '.join(missing)}",
+        )
+
+
+class MeanReciprocalRank:
+    """Reciprocal rank of the first expected document, averaged over cases.
+
+    Scored per case, so the "mean" is taken by the gate across cases and runs. One
+    number per case keeps the metric a drop-in for the gate's per-case arithmetic.
+    """
+
+    def __init__(self, *, name: str = "mrr", min_score: float = 0.5) -> None:
+        self._name = name
+        self._min_score = min_score
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def deterministic(self) -> bool:
+        return True
+
+    async def score(self, case: RetrievalCase, retrieved: RetrievalResult) -> MetricResult:
+        expected = set(case.expect_docs)
+        for rank, doc_id in enumerate(retrieved.doc_ids[: case.k], start=1):
+            if doc_id in expected:
+                score = 1.0 / rank
+                return MetricResult(
+                    metric=self.name,
+                    score=score,
+                    passed=score >= self._min_score,
+                    detail=f"first expected document {doc_id!r} at rank {rank}",
+                )
+        return MetricResult(
+            metric=self.name,
+            score=0.0,
+            passed=self._min_score <= 0.0,
+            detail=f"none of {', '.join(sorted(expected))} were retrieved within k={case.k}",
+        )
+
+
+class CitationGrounding:
+    """Fraction of expected documents that a returned item cites *correctly*.
+
+    Distinct from recall because it scores citations rather than document ids, and
+    because it checks integrity first: a citation whose ``chunk_id`` or ``doc_id``
+    disagrees with the passage it is attached to is worse than a missing citation,
+    since it sends a reader to the wrong place while looking authoritative. Any
+    such mismatch scores 0 regardless of what was retrieved.
+
+    This definition is authored here: the IR has no citation metric kind, so there is
+    no upstream semantics to match. It is stated rather than derived.
+    """
+
+    def __init__(self, *, name: str = "citation-grounding", min_score: float = 1.0) -> None:
+        self._name = name
+        self._min_score = min_score
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def deterministic(self) -> bool:
+        return True
+
+    async def score(self, case: RetrievalCase, retrieved: RetrievalResult) -> MetricResult:
+        broken = [
+            item
+            for item in retrieved.items
+            if item.citation.chunk_id != item.chunk.chunk_id
+            or item.citation.doc_id != item.chunk.doc_id
+        ]
+        if broken:
+            return MetricResult(
+                metric=self.name,
+                score=0.0,
+                passed=False,
+                detail=(
+                    f"{len(broken)} citation(s) disagree with the passage they name, "
+                    f"first: {broken[0].citation} vs chunk {broken[0].chunk.chunk_id}"
+                ),
+            )
+        expected = set(case.expect_docs)
+        cited = {item.citation.doc_id for item in retrieved.items}
+        grounded = expected & cited
+        score = len(grounded) / len(expected)
+        missing = sorted(expected - grounded)
+        return MetricResult(
+            metric=self.name,
+            score=score,
+            passed=score >= self._min_score,
+            detail="" if not missing else f"uncited: {', '.join(missing)}",
+        )
+
+
+def retrieval_metrics() -> tuple[RetrievalMetric, ...]:
+    """The retrieval metrics this runtime ships, in a stable order."""
+    return (RecallAtK(), MeanReciprocalRank(), CitationGrounding())

aid_runtime/retrieval/pgvector.py

@@ -0,0 +1,258 @@
+"""A pgvector-backed store, provisioned locally and never touched by CI.
+
+The driver is an optional extra, so this module must be importable without it.
+Nothing at module scope imports ``psycopg``; the import happens inside
+:func:`_load_driver` on first use, and a missing driver becomes a
+:class:`~aid_runtime.errors.ConfigurationError` that names the extra to install.
+Importing :mod:`aid_runtime.retrieval` therefore cannot fail because a database
+extra is absent, which is what lets the default dependency set stay unchanged.
+
+Two consequences of that choice are deliberate and visible:
+
+* The driver is loaded as an untyped module, so ``Any`` reaches this file. It is
+  confined to this one adapter rather than spreading a suppressed import across the
+  retrieval package.
+* Every statement is a fixed string. The embedding column is declared
+  dimension-less rather than ``vector(1536)``, so no computed identifier is ever
+  interpolated into SQL, and the pinned width is enforced where the pin actually
+  lives: the caller's :class:`~aid_runtime.retrieval.types.PinnedEmbedding`.
+
+Correctness here is expected to match the in-memory adapter, but nothing checks
+that: this adapter is not gate-covered, and its SQL path runs only on a machine where
+someone brought the database up themselves. See ``compose.pgvector.yml``.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Sequence
+from typing import Any
+
+from aid_runtime.errors import ConfigurationError, ProviderError
+from aid_runtime.retrieval.types import (
+    Chunk,
+    DistanceMetric,
+    SearchHit,
+    StoreKind,
+    StoreRecord,
+)
+
+__all__ = ["PgVectorStore", "resolve_uri"]
+
+_ENV_PREFIX = "env:"
+
+_CREATE_EXTENSION = "CREATE EXTENSION IF NOT EXISTS vector"
+_CREATE_TABLE = """
+CREATE TABLE IF NOT EXISTS aid_chunks (
+    chunk_id text PRIMARY KEY,
+    doc_id text NOT NULL,
+    uri text NOT NULL,
+    chunk_index integer NOT NULL,
+    start_offset integer NOT NULL,
+    end_offset integer NOT NULL,
+    body text NOT NULL,
+    metadata jsonb NOT NULL DEFAULT '{}'::jsonb,
+    embedding vector NOT NULL
+)
+"""
+_UPSERT = """
+INSERT INTO aid_chunks (
+    chunk_id, doc_id, uri, chunk_index, start_offset, end_offset, body, metadata, embedding
+) VALUES (
+    %(chunk_id)s, %(doc_id)s, %(uri)s, %(chunk_index)s, %(start_offset)s, %(end_offset)s,
+    %(body)s, %(metadata)s::jsonb, %(embedding)s::vector
+)
+ON CONFLICT (chunk_id) DO UPDATE SET
+    doc_id = EXCLUDED.doc_id,
+    uri = EXCLUDED.uri,
+    chunk_index = EXCLUDED.chunk_index,
+    start_offset = EXCLUDED.start_offset,
+    end_offset = EXCLUDED.end_offset,
+    body = EXCLUDED.body,
+    metadata = EXCLUDED.metadata,
+    embedding = EXCLUDED.embedding
+"""
+_CORPUS = (
+    "SELECT chunk_id, doc_id, uri, chunk_index, start_offset, end_offset, body, metadata, "
+    "embedding FROM aid_chunks ORDER BY chunk_id"
+)
+# `<=>` is cosine distance, `<#>` is negative inner product and `<->` is L2 distance.
+# Each is rewritten into "higher is better" so callers see one orientation, matching
+# `stores.similarity`. The tie-break on `chunk_id` matches the in-memory adapter too:
+# equal scores must rank identically on both.
+#
+# Every statement is spelled out in full rather than assembled from a column list, so
+# that no statement is ever *constructed*: the only values that reach SQL are bound
+# parameters.
+_SEARCH: dict[DistanceMetric, str] = {
+    DistanceMetric.COSINE: (
+        "SELECT chunk_id, doc_id, uri, chunk_index, start_offset, end_offset, body, metadata, "
+        "embedding, 1 - (embedding <=> %(vector)s::vector) AS score FROM aid_chunks "
+        "ORDER BY embedding <=> %(vector)s::vector, chunk_id LIMIT %(top_k)s"
+    ),
+    DistanceMetric.DOT: (
+        "SELECT chunk_id, doc_id, uri, chunk_index, start_offset, end_offset, body, metadata, "
+        "embedding, -(embedding <#> %(vector)s::vector) AS score FROM aid_chunks "
+        "ORDER BY embedding <#> %(vector)s::vector, chunk_id LIMIT %(top_k)s"
+    ),
+    DistanceMetric.L2: (
+        "SELECT chunk_id, doc_id, uri, chunk_index, start_offset, end_offset, body, metadata, "
+        "embedding, -(embedding <-> %(vector)s::vector) AS score FROM aid_chunks "
+        "ORDER BY embedding <-> %(vector)s::vector, chunk_id LIMIT %(top_k)s"
+    ),
+}
+
+
+def _load_driver() -> Any:
+    """Import the optional driver, or explain how to install it."""
+    try:
+        import psycopg
+    except ImportError as error:
+        raise ConfigurationError(
+            "the pgvector store needs its driver; install it with "
+            "`pip install aid-runtime-python[pgvector]` (the default install omits it)"
+        ) from error
+    return psycopg
+
+
+def resolve_uri(uri: str | None) -> str:
+    """Resolve an ``env:NAME`` reference to the variable's value.
+
+    An ``env:`` reference names a variable, not an alias: nothing here reaches the
+    model catalog, so resolving it at run time cannot silently change which store a
+    run wrote to.
+    """
+    if uri is None:
+        raise ConfigurationError("the pgvector store requires a uri")
+    if not uri.startswith(_ENV_PREFIX):
+        return uri
+    name = uri[len(_ENV_PREFIX) :]
+    if not name:
+        raise ConfigurationError(f"store uri {uri!r} names an empty environment variable")
+    value = os.environ.get(name)
+    if value is None:
+        raise ConfigurationError(
+            f"store uri {uri!r} refers to environment variable {name!r}, which is unset"
+        )
+    return value
+
+
+def _literal(vector: Sequence[float]) -> str:
+    """pgvector's text input format, so its Python package is not a dependency."""
+    return "[" + ",".join(repr(float(value)) for value in vector) + "]"
+
+
+class PgVectorStore:
+    """A single-table store: one row per chunk, keyed by content hash."""
+
+    def __init__(self, *, uri: str | None, dimensions: int | None = None) -> None:
+        self._uri = resolve_uri(uri)
+        self._dimensions = dimensions
+        self._connection: Any = None
+        self._closed = False
+
+    @property
+    def kind(self) -> StoreKind:
+        return StoreKind.PGVECTOR
+
+    async def _connect(self) -> Any:
+        if self._closed:
+            raise ConfigurationError("this store has been closed")
+        if self._connection is None:
+            psycopg = _load_driver()
+            try:
+                self._connection = await psycopg.AsyncConnection.connect(self._uri)
+            except Exception as error:
+                raise ProviderError(f"cannot connect to the pgvector store: {error}") from error
+            await self._ensure_schema()
+        return self._connection
+
+    async def _ensure_schema(self) -> None:
+        async with self._connection.cursor() as cursor:
+            await cursor.execute(_CREATE_EXTENSION)
+            await cursor.execute(_CREATE_TABLE)
+        await self._connection.commit()
+
+    async def upsert(self, records: Sequence[StoreRecord]) -> int:
+        if self._dimensions is not None:
+            for record in records:
+                if len(record.vector) != self._dimensions:
+                    raise ConfigurationError(
+                        f"store is pinned to {self._dimensions} dimensions but a record carries "
+                        f"{len(record.vector)}"
+                    )
+        if not records:
+            return 0
+        connection = await self._connect()
+        async with connection.cursor() as cursor:
+            for record in records:
+                await cursor.execute(_UPSERT, self._params(record))
+        await connection.commit()
+        return len(records)
+
+    async def search(
+        self, vector: Sequence[float], *, top_k: int, metric: DistanceMetric
+    ) -> tuple[SearchHit, ...]:
+        if top_k < 1:
+            raise ConfigurationError("top_k must be at least 1")
+        statement = _SEARCH.get(metric)
+        if statement is None:
+            raise ConfigurationError(f"pgvector cannot rank by {metric.value!r}")
+        connection = await self._connect()
+        async with connection.cursor() as cursor:
+            await cursor.execute(statement, {"vector": _literal(vector), "top_k": top_k})
+            rows = await cursor.fetchall()
+        return tuple(SearchHit(record=self._record(row), score=float(row[9])) for row in rows)
+
+    async def corpus(self) -> tuple[StoreRecord, ...]:
+        connection = await self._connect()
+        async with connection.cursor() as cursor:
+            await cursor.execute(_CORPUS)
+            rows = await cursor.fetchall()
+        return tuple(self._record(row) for row in rows)
+
+    async def aclose(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        if self._connection is not None:
+            await self._connection.close()
+            self._connection = None
+
+    def _params(self, record: StoreRecord) -> dict[str, Any]:
+        chunk = record.chunk
+        return {
+            "chunk_id": chunk.chunk_id,
+            "doc_id": chunk.doc_id,
+            "uri": chunk.uri,
+            "chunk_index": chunk.index,
+            "start_offset": chunk.start,
+            "end_offset": chunk.end,
+            "body": chunk.text,
+            "metadata": json.dumps(chunk.metadata, sort_keys=True),
+            "embedding": _literal(record.vector),
+        }
+
+    def _record(self, row: Sequence[Any]) -> StoreRecord:
+        return StoreRecord(
+            chunk=Chunk(
+                chunk_id=row[0],
+                doc_id=row[1],
+                uri=row[2],
+                index=row[3],
+                start=row[4],
+                end=row[5],
+                text=row[6],
+                metadata=row[7] or {},
+            ),
+            vector=tuple(float(value) for value in _parse_vector(row[8])),
+        )
+
+
+def _parse_vector(value: Any) -> list[float]:
+    """Accept either a driver-returned sequence or pgvector's text form."""
+    if isinstance(value, str):
+        stripped = value.strip().strip("[]")
+        return [float(part) for part in stripped.split(",")] if stripped else []
+    return [float(item) for item in value]

aid_runtime/retrieval/rerank.py

@@ -0,0 +1,133 @@
+"""Reordering a retrieved set before it is returned.
+
+Exactly one reranker exists in Phase 1, and it is lexical. That is a deliberate
+limitation rather than an unfinished one: the seam is what this layer owes, and the
+deterministic implementation is what lets a gate assert on ranking behaviour without
+an API key or a floating cross-encoder.
+
+Why rerank at all when hybrid search already blends two rankings? Because fusion
+answers "which candidates are worth considering" and reranking answers "which of
+these best answers the question", and the second question can look at the query and
+one passage together instead of at ranks alone. The lexical reranker does that with
+query-term coverage, which is cheap, explainable, and stable.
+
+The reranker **replaces** each item's score instead of adding a second number. After
+reranking, the number that ranks the list is the reranker's, and keeping a fused RRF
+value in the same field would put two incomparable scales in one place. Ranks are
+renumbered from 1 so ``RetrievalResult.context`` stays resolvable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol, runtime_checkable
+
+from aid_runtime.errors import ConfigurationError, FeatureDeferredError
+from aid_runtime.retrieval.citations import RetrievedItem, cite
+from aid_runtime.retrieval.hybrid import tokenize
+from aid_runtime.retrieval.types import RerankSpec
+
+__all__ = [
+    "LexicalReranker",
+    "Reranker",
+    "lexical_score",
+    "reranker_for",
+]
+
+_LEXICAL_MARKER = "lexical"
+
+
+def lexical_score(query: str, text: str) -> float:
+    """How much of ``query`` the passage actually covers, in ``[0, 1]``.
+
+    Query-term coverage, not term frequency: for "which passages answer this", a
+    passage containing each distinct query term once is as good a candidate as one
+    repeating a term, and rewarding repetition is how a keyword-stuffed chunk wins.
+    An exact phrase match breaks ties toward the passage quoting the query.
+    """
+    terms = set(tokenize(query))
+    if not terms:
+        return 0.0
+    present = set(tokenize(text))
+    coverage = len(terms & present) / len(terms)
+    if query.strip() and query.strip().lower() in text.lower():
+        # Small and bounded, so it can only order passages that already cover the
+        # query equally well; it can never promote an unrelated passage.
+        return min(1.0, coverage + 0.1)
+    return coverage
+
+
+@runtime_checkable
+class Reranker(Protocol):
+    """Reorder candidates for a query, deterministically."""
+
+    @property
+    def name(self) -> str: ...
+
+    async def rerank(
+        self,
+        query: str,
+        items: Sequence[RetrievedItem],
+        *,
+        top_n: int,
+    ) -> tuple[RetrievedItem, ...]: ...
+
+
+class LexicalReranker:
+    """Coverage-based reranking, with ties broken on chunk id."""
+
+    def __init__(self, *, name: str = "lexical") -> None:
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    async def rerank(
+        self,
+        query: str,
+        items: Sequence[RetrievedItem],
+        *,
+        top_n: int,
+    ) -> tuple[RetrievedItem, ...]:
+        if top_n < 1:
+            raise ConfigurationError("rerank top_n must be at least 1")
+        scored = [
+            (lexical_score(query, item.chunk.text), item.chunk.chunk_id, item) for item in items
+        ]
+        # A reranker that cannot read the query leaves the incoming order alone
+        # rather than shuffling it: every score is 0.0, so the tie-break would be
+        # chunk-id order, which is an arbitrary order dressed up as a decision.
+        if all(score == 0.0 for score, _, _ in scored):
+            return tuple(items[:top_n])
+        scored.sort(key=lambda entry: (-entry[0], entry[1]))
+        return tuple(
+            RetrievedItem(
+                chunk=item.chunk,
+                score=score,
+                rank=rank,
+                citation=cite(item.chunk),
+            )
+            for rank, (score, _, item) in enumerate(scored[:top_n], start=1)
+        )
+
+
+def reranker_for(spec: RerankSpec | None) -> Reranker | None:
+    """The reranker a retriever's ``rerank`` block asks for, if any.
+
+    ``rerank.model`` is free-form in the IR, and Phase 1 ships exactly one
+    implementation. A declared model outside the lexical family is refused rather
+    than quietly served by :class:`LexicalReranker`, for the same reason an
+    unimplemented source kind or store kind is refused: a retriever that appears to
+    rerank with a cross-encoder it never called is worse than one that says it
+    cannot. The guard matches on the name because the name is the only signal the
+    IR gives; the lexical family is what this runtime can honestly claim.
+    """
+    if spec is None:
+        return None
+    if _LEXICAL_MARKER not in spec.model.casefold():
+        raise FeatureDeferredError(
+            f"rerank model {spec.model!r} is declared but not implemented: Phase 1 ships one "
+            "deterministic lexical reranker, and a hosted cross-encoder is a later layer"
+        )
+    return LexicalReranker()