pinecall-runtime 0.1.6__py3-none-any.whl

pinecall/retrieval/_search.py

@@ -0,0 +1,186 @@
+"""The hybrid search over one table: nearest vectors and best BM25 words, fused by rank in SQL."""
+
+import dataclasses
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from statistics import fmean
+from typing import Literal
+
+from psycopg import sql
+from psycopg.rows import DictRow
+
+from pinecall.postgres.pool import Connection
+from pinecall.retrieval.embed import halfvec
+
+type Evidence = Literal["strong", "weak", "none"]
+
+
+# Cormack's constant: a rank weighs 1 / (60 + rank), so cosine and BM25 are never compared.
+RRF_K = 60
+
+
+# The cutoff of nDCG, past any k a turn is handed.
+NDCG_AT = 10
+
+
+CANDIDATES_PER_BRANCH = 30
+
+
+# Measured on a golden of real questions; data for the reader and the judge, never a gate.
+STRONG = 0.45
+
+
+NONE_BELOW = 0.30
+
+
+# `hnsw.iterative_scan` makes the index scan past its first 40 candidates when the scope filters
+# them out: one tenant's rows among every tenant's would otherwise come back short, or empty.
+ITERATIVE_SCAN = "SET LOCAL hnsw.iterative_scan = relaxed_order"
+
+
+# `<@>` is BM25 negated, lower is better, and 0 when no term matched: those rows are no match.
+HYBRID = """
+WITH dense AS (
+    SELECT id, RANK() OVER (ORDER BY distance) AS r FROM (
+        SELECT id, embedding <=> %(q)s::halfvec AS distance FROM {table}
+        WHERE ({scope}) AND ({same_space})
+        ORDER BY embedding <=> %(q)s::halfvec
+        LIMIT %(room)s
+    ) AS nearest
+),
+sparse AS (
+    SELECT id, RANK() OVER (ORDER BY score) AS r FROM (
+        SELECT id, score FROM (
+            SELECT id, text <@> to_bm25query(%(words)s, {text_index}) AS score FROM {table}
+            WHERE ({scope})
+        ) AS scored
+        WHERE score < 0
+        ORDER BY score
+        LIMIT %(room)s
+    ) AS worded
+),
+fusion AS (
+    SELECT coalesce(dense.id, sparse.id) AS id,
+        (coalesce(1.0 / ({k} + dense.r), 0) + coalesce(1.0 / ({k} + sparse.r), 0))::float8 AS fused
+    FROM dense FULL OUTER JOIN sparse ON dense.id = sparse.id
+)
+SELECT fusion.id::text AS hit_id, fusion.fused AS hit_fused,
+    1 - (found.embedding <=> %(q)s::halfvec) AS hit_cosine, {columns}
+FROM fusion JOIN {table} AS found ON found.id = fusion.id
+ORDER BY fusion.fused DESC, fusion.id
+"""
+
+
+@dataclass(frozen=True)
+class SearchedTable:
+    """What differs between the tables searched: which rows a query sees and what a hit carries."""
+
+    name: str
+    # to_bm25query reads its statistics by the index's name.
+    text_index: str
+    # A WHERE clause over the table, its placeholders named: `q`, `words` and `room` are taken.
+    scope: sql.Composable
+    params: Mapping[str, object]
+    columns: tuple[str, ...]
+    # Added to the scope of the vector branch alone: a row embedded by another model is found by
+    # its words only.
+    same_space: sql.Composable = field(default_factory=lambda: sql.SQL("true"))
+
+
+@dataclass(frozen=True)
+class Hit:
+    """One row found: its id, its fused score, its cosine to the query, and the columns asked."""
+
+    id: str
+    fused: float
+    cosine: float
+    row: DictRow
+
+
+@dataclass(frozen=True)
+class Figures:
+    """Recall at k, and nDCG at 10."""
+
+    recall_at_k: float
+    ndcg_at_10: float
+
+
+async def hybrid(
+    connection: Connection, table: SearchedTable, *, vector: list[float], words: str, room: int
+) -> list[Hit]:
+    """Up to `room` rows from each branch, fused by reciprocal rank, the best first."""
+    query = sql.SQL(HYBRID).format(
+        table=sql.Identifier(table.name),
+        text_index=sql.Literal(table.text_index),
+        scope=table.scope,
+        same_space=table.same_space,
+        k=sql.Literal(RRF_K),
+        columns=sql.SQL(", ").join(sql.Identifier("found", column) for column in table.columns),
+    )
+    params = {**table.params, "q": halfvec(vector), "words": words, "room": room}
+    async with connection.transaction():
+        await connection.execute(ITERATIVE_SCAN)
+        rows = await (await connection.execute(query, params)).fetchall()
+    return [
+        Hit(
+            id=row["hit_id"],
+            fused=row["hit_fused"],
+            cosine=row["hit_cosine"],
+            row={column: row[column] for column in table.columns},
+        )
+        for row in rows
+    ]
+
+
+# A fused score means nothing across queries: relative to the best, the top is always 1.0.
+def relative_to_the_best(hits: list[Hit]) -> list[Hit]:
+    """Each hit's fused score divided by the best, the best first and ties by id."""
+    best = max((hit.fused for hit in hits), default=0.0)
+    ordered = sorted(hits, key=lambda hit: (-hit.fused, hit.id))
+    return [dataclasses.replace(hit, fused=hit.fused / best if best else 0.0) for hit in ordered]
+
+
+def top_cosine(hits: list[Hit]) -> float:
+    """The best cosine among the hits; 0.0 when there are none."""
+    return max((hit.cosine for hit in hits), default=0.0)
+
+
+def evidence_of(cosine: float) -> Evidence:
+    """How strongly the best cosine says the question is answered here."""
+    if cosine >= STRONG:
+        return "strong"
+    if cosine < NONE_BELOW:
+        return "none"
+    return "weak"
+
+
+# nDCG is normalised by the ideal order for as many answers as the question wants.
+def figures(ranks: Sequence[Sequence[int | None]]) -> Figures:
+    """The two figures over the ranks of each question's wanted answers; None is not found."""
+    if not ranks:
+        return Figures(recall_at_k=0.0, ndcg_at_10=0.0)
+    return Figures(
+        recall_at_k=fmean(_share_found(question) for question in ranks),
+        ndcg_at_10=fmean(_normalised(question) for question in ranks),
+    )
+
+
+def _share_found(ranks: Sequence[int | None]) -> float:
+    if not ranks:
+        return 0.0
+    return sum(rank is not None for rank in ranks) / len(ranks)
+
+
+def _normalised(ranks: Sequence[int | None]) -> float:
+    if not ranks:
+        return 0.0
+    ideal = sum(_discounted(place) for place in range(1, min(len(ranks), NDCG_AT) + 1))
+    return sum(_discounted(rank) for rank in ranks) / ideal
+
+
+# Past the tenth the model never saw it.
+def _discounted(rank: int | None) -> float:
+    if rank is None or rank > NDCG_AT:
+        return 0.0
+    return 1.0 / math.log2(rank + 1)

pinecall/retrieval/_sections.py

@@ -0,0 +1,70 @@
+"""A found chunk handed back as its section, and which sections a search hands back."""
+
+from collections.abc import Mapping, Sequence
+
+from pinecall.domain.scope import Scope
+from pinecall.postgres.pool import Connection
+
+# A found chunk comes back as its section: the chunks of its heading this many either side of it.
+SECTION_SPAN = 2
+
+
+# The best pages give every section that was found in them; every page after them gives one.
+PAGES_GIVING_ALL = 2
+
+
+# The chunks around each found one, under the same heading of the same file, in the file's order.
+SECTIONS = """
+SELECT wanted.at, chunk.text
+FROM unnest(%(ats)s::integer[], %(bases)s::text[], %(holders)s::text[], %(paths)s::text[],
+    %(headings)s::text[], %(firsts)s::integer[], %(lasts)s::integer[])
+    AS wanted (at, base, holder, path, heading, first, last)
+JOIN knowledge_chunks AS chunk ON chunk.org = %(org)s AND chunk.env = %(env)s
+    AND chunk.base = wanted.base AND chunk.holder = wanted.holder AND chunk.path = wanted.path
+    AND chunk.heading IS NOT DISTINCT FROM wanted.heading
+    AND chunk.ordinal BETWEEN wanted.first AND wanted.last
+ORDER BY wanted.at, chunk.ordinal
+"""
+
+
+def sections_chosen(rows: Sequence[Mapping[str, object]], k: int) -> list[int]:
+    """Which hits come back, best first: one per section, the best pages all theirs, k at most."""
+    chosen: list[int] = []
+    sections: set[tuple[object, object, object]] = set()
+    pages: list[tuple[object, object]] = []
+    for at, row in enumerate(rows):
+        page = (row["base"], row["path"])
+        section = (row["base"], row["path"], row["heading"])
+        if section in sections or (page in pages and pages.index(page) >= PAGES_GIVING_ALL):
+            continue
+        if page not in pages:
+            pages.append(page)
+        sections.add(section)
+        chosen.append(at)
+        if len(chosen) == k:
+            break
+    return chosen
+
+
+async def section_texts(
+    connection: Connection, scope: Scope, rows: Sequence[Mapping[str, object]]
+) -> dict[int, list[str]]:
+    """Each row's section as its chunks' indexed texts in the file's order, by the row's place."""
+    if not rows:
+        return {}
+    ordinals = [int(str(row["ordinal"])) for row in rows]
+    named = {
+        "org": scope.org,
+        "env": scope.env,
+        "ats": list(range(len(rows))),
+        "bases": [row["base"] for row in rows],
+        "holders": [row["holder"] for row in rows],
+        "paths": [row["path"] for row in rows],
+        "headings": [row["heading"] for row in rows],
+        "firsts": [ordinal - SECTION_SPAN for ordinal in ordinals],
+        "lasts": [ordinal + SECTION_SPAN for ordinal in ordinals],
+    }
+    texts: dict[int, list[str]] = {}
+    for row in await (await connection.execute(SECTIONS, named)).fetchall():
+        texts.setdefault(int(row["at"]), []).append(str(row["text"]))
+    return texts

pinecall/retrieval/embed.py

@@ -0,0 +1,340 @@
+"""The box's embedder over HTTP, OpenAI-shaped or contextual: batched, retried, validated."""
+
+import asyncio
+import base64
+import logging
+import math
+import re
+import secrets
+import struct
+from array import array
+from collections.abc import Awaitable, Callable, Iterator, Mapping
+
+import httpx
+from pydantic import BaseModel, ConfigDict, ValidationError
+
+from pinecall.domain.errors import EmbedderUnreachable, NotAvailable, WrongWidth
+from pinecall.domain.names import Credentials, Json
+from pinecall.providers.catalog import Embedding
+
+logger = logging.getLogger(__name__)
+
+type Sleep = Callable[[float], Awaitable[None]]
+
+
+# The cutter measures a chunk with the same estimate, so a chunk that fits its cap fits a window.
+TOKENS_PER_WORD = 1.3
+
+
+# The contextual endpoint counts a window's chunks together against a 32 768-token context; the
+# margin is for the estimate.
+WINDOW_TOKENS = 24_000
+
+
+REQUEST_TOKENS = 100_000
+
+
+# Only a refusal about size is split: halving a credential refusal fails it twice.
+TOO_BIG = re.compile(
+    r"exceeds maximum|too many tokens|maximum context length|reduce the length|"
+    r"input is too large|max_tokens_per_request",
+    re.IGNORECASE,
+)
+
+
+MAX_SPLITS = 5
+
+
+RETRIES = 2
+
+
+# The request was fine and the server was not: a timeout, a lock, a rate limit, or a 5xx.
+RETRYABLE = frozenset({408, 409, 429})
+
+
+FIRST_WAIT_S = 0.5
+
+
+LONGEST_WAIT_S = 8.0
+
+
+# A server that asks for longer is misconfigured; its hour is not waited.
+LONGEST_ASKED_S = 60.0
+
+
+# A push waits for a contextual window of 24 000 tokens, which outlasts any client default; a
+# query keeps the client's own, and the turn's budget bounds it.
+PUSH_TIMEOUT = httpx.Timeout(connect=5.0, read=120.0, write=5.0, pool=5.0)
+
+
+class _Refusal(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    message: str = ""
+
+
+class _Answer(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    error: _Refusal | str | None = None
+
+
+class _Row(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    index: int | None = None
+    embedding: list[float] | str
+
+
+class _Rows(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    data: list[_Row]
+
+
+class _Documents(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+
+    data: list[_Rows]
+
+
+class Embedder:
+    """The embedder the box's `embedding` names, on the box's key, over one HTTP client."""
+
+    def __init__(
+        self,
+        embedding: Embedding,
+        key: Credentials,
+        client: httpx.AsyncClient,
+        *,
+        sleep: Sleep = asyncio.sleep,
+    ) -> None:
+        """Hold the client; nothing is asked of the vendor until the first text."""
+        self.embedding = embedding
+        path = "embeddings" if embedding.shape == "embeddings" else "contextualizedembeddings"
+        self._url = f"{embedding.url.rstrip('/')}/{path}"
+        self._headers = {"Authorization": f"Bearer {_key_of(embedding.vendor, key)}"}
+        self._client = client
+        self._sleep = sleep
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """One unit vector per text, in order; on `contextual` each text is a one-chunk document."""
+        waits = self._client.timeout
+        if self.embedding.shape == "contextual":
+            return [(await self._sent([text], waits))[0] for text in texts]
+        return [
+            vector
+            for batch in _windows(texts, REQUEST_TOKENS)
+            for vector in await self._sent(batch, waits)
+        ]
+
+    async def embed_documents(self, documents: list[list[str]]) -> list[list[list[float]]]:
+        """One list of unit vectors per document, one per chunk, in order."""
+        if self.embedding.shape == "contextual":
+            return [
+                [
+                    vector
+                    for window in _windows(chunks, WINDOW_TOKENS)
+                    for vector in await self._sent(window, PUSH_TIMEOUT)
+                ]
+                for chunks in documents
+            ]
+        every = [chunk for chunks in documents for chunk in chunks]
+        flat = [
+            vector
+            for batch in _windows(every, REQUEST_TOKENS)
+            for vector in await self._sent(batch, PUSH_TIMEOUT)
+        ]
+        cut: list[list[list[float]]] = []
+        for chunks in documents:
+            cut.append(flat[: len(chunks)])
+            flat = flat[len(chunks) :]
+        return cut
+
+    # A split keeps the order: the first half's vectors, then the second's.
+    async def _sent(
+        self, texts: list[str], waits: httpx.Timeout, *, depth: int = 0
+    ) -> list[list[float]]:
+        try:
+            return self._vectors(await self._posted(texts, waits), len(texts))
+        except EmbedderUnreachable as refused:
+            if depth >= MAX_SPLITS or len(texts) == 1 or not TOO_BIG.search(refused.words):
+                raise
+        middle = len(texts) // 2
+        first = await self._sent(texts[:middle], waits, depth=depth + 1)
+        return [*first, *await self._sent(texts[middle:], waits, depth=depth + 1)]
+
+    async def _posted(self, texts: list[str], waits: httpx.Timeout) -> bytes:
+        body = self._body(texts)
+        attempt = 0
+        while True:
+            try:
+                answer = await self._client.post(
+                    self._url, json=body, headers=self._headers, timeout=waits
+                )
+            except httpx.TransportError as failed:
+                if attempt == RETRIES:
+                    raise self._unreachable(str(failed) or type(failed).__name__) from failed
+                await self._sleep(_wait(attempt, {}))
+                attempt += 1
+                continue
+            retryable = answer.status_code in RETRYABLE or answer.is_server_error
+            if retryable and attempt < RETRIES:
+                await self._sleep(_wait(attempt, answer.headers))
+                attempt += 1
+                continue
+            words = _refused_by(answer)
+            if words is not None:
+                raise self._unreachable(words)
+            return answer.content
+
+    # The request names the encoding: the reply is decoded as asked, never guessed from its bytes.
+    def _body(self, texts: list[str]) -> dict[str, Json]:
+        wanted = self.embedding
+        if wanted.shape == "contextual":
+            return {
+                "model": wanted.model,
+                "input": [list[Json](texts)],
+                "encoding_format": "base64_int8",
+                "dimensions": wanted.dimensions,
+            }
+        return {
+            "model": wanted.model,
+            "input": list[Json](texts),
+            "encoding_format": "float",
+            "dimensions": wanted.dimensions,
+        }
+
+    def _vectors(self, content: bytes, sent: int) -> list[list[float]]:
+        try:
+            rows = self._rows(content)
+        except ValidationError as unread:
+            raise self._unreachable(f"a reply that is not its vectors: {unread}") from unread
+        if len(rows) != sent:
+            raise self._unreachable(f"{len(rows)} vectors for {sent} chunks")
+        vectors = [self._decoded(row.embedding) for row in self._in_order(rows)]
+        widths = sorted({len(vector) for vector in vectors})
+        if widths and widths != [self.embedding.dimensions]:
+            answered = " and ".join(str(width) for width in widths)
+            raise WrongWidth(
+                f"{self.embedding.model} answers {answered}-wide vectors; every vector column "
+                f"holds {self.embedding.dimensions}"
+            )
+        return [_unit(vector) for vector in vectors]
+
+    def _rows(self, content: bytes) -> list[_Row]:
+        if self.embedding.shape == "embeddings":
+            return _Rows.model_validate_json(content).data
+        documents = _Documents.model_validate_json(content).data
+        if len(documents) != 1:
+            raise self._unreachable(f"{len(documents)} documents for the one it was sent")
+        return documents[0].data
+
+    # Sorting alone is not enough: indexes 0, 0, 2 sort cleanly and give two chunks one vector.
+    def _in_order(self, rows: list[_Row]) -> list[_Row]:
+        if any(row.index is None for row in rows):
+            return rows
+        ordered = sorted(rows, key=lambda row: row.index or 0)
+        indexes = [row.index for row in ordered]
+        if indexes != list(range(len(rows))):
+            raise self._unreachable(f"a reply indexed {indexes}, not 0 to {len(rows) - 1}")
+        return ordered
+
+    def _decoded(self, embedding: list[float] | str) -> list[float]:
+        if self.embedding.shape == "embeddings":
+            if isinstance(embedding, str):
+                raise self._unreachable("a base64 vector where floats were asked for")
+            return embedding
+        if not isinstance(embedding, str):
+            raise self._unreachable("floats where base64_int8 was asked for")
+        # Signed: -1 travels as 0xff, and read unsigned it would be 255.
+        values = array("b")
+        values.frombytes(base64.b64decode(embedding))
+        return [float(value) for value in values]
+
+    def _unreachable(self, words: str) -> EmbedderUnreachable:
+        logger.warning("the embedder at %s did not answer: %s", self._url, words)
+        return EmbedderUnreachable(self.embedding.vendor, self._url, words)
+
+
+def estimated_tokens(text: str) -> int:
+    """The tokens a text costs, estimated as its words times 1.3; no tokenizer is loaded."""
+    return round(len(text.split()) * TOKENS_PER_WORD)
+
+
+# pgvector parses the text and the SQL casts it (`::halfvec`), so the driver needs no vector type.
+def halfvec(vector: list[float]) -> str:
+    """The vector as pgvector's text, each value rounded to IEEE half precision as the column is."""
+    return "[" + ",".join(str(_half(value)) for value in vector) + "]"
+
+
+def _half(value: float) -> float:
+    rounded: float = struct.unpack("<e", struct.pack("<e", value))[0]
+    return rounded
+
+
+# A text over the budget goes alone: the vendor refuses it or cuts it, and a refusal is split.
+def _windows(texts: list[str], budget: int) -> Iterator[list[str]]:
+    window: list[str] = []
+    tokens = 0
+    for text in texts:
+        cost = estimated_tokens(text)
+        if window and tokens + cost > budget:
+            yield window
+            window, tokens = [], 0
+        window.append(text)
+        tokens += cost
+    if window:
+        yield window
+
+
+def _key_of(vendor: str, key: Credentials) -> str:
+    if isinstance(key, str):
+        return key
+    named = key.get("api_key")
+    if isinstance(named, str):
+        return named
+    raise NotAvailable(f"the box's credentials for {vendor} hold no api_key")
+
+
+# A gateway may refuse with a 200 and an `error` body: the words are read whatever the status.
+def _refused_by(answer: httpx.Response) -> str | None:
+    try:
+        payload = _Answer.model_validate_json(answer.content).error
+    except ValidationError:
+        payload = None
+    if isinstance(payload, _Refusal) and payload.message:
+        return payload.message
+    if isinstance(payload, str) and payload:
+        return payload
+    if answer.is_error or payload is not None:
+        return f"HTTP {answer.status_code}"
+    return None
+
+
+def _wait(attempt: int, headers: Mapping[str, str]) -> float:
+    params = _retry_after(headers)
+    if params is not None and 0 < params <= LONGEST_ASKED_S:
+        return params
+    backoff = min(FIRST_WAIT_S * 2**attempt, LONGEST_WAIT_S)
+    # The jitter only shortens the wait, so the longest wait stays the longest.
+    return backoff * (1 - 0.25 * secrets.SystemRandom().random())
+
+
+# `retry-after-ms` first: a server that says 1500 ms is not rounded up to two seconds.
+def _retry_after(headers: Mapping[str, str]) -> float | None:
+    for name, scale in (("retry-after-ms", 0.001), ("retry-after", 1.0)):
+        raw = headers.get(name)
+        if raw is None:
+            continue
+        try:
+            return float(raw) * scale
+        except ValueError:
+            continue
+    return None
+
+
+# Both shapes answer unnormalised vectors; a zero vector has no direction to keep.
+def _unit(vector: list[float]) -> list[float]:
+    length = math.sqrt(sum(value * value for value in vector))
+    return vector if length == 0 else [value / length for value in vector]