forkreduce 0.1.0__py3-none-any.whl

forkreduce/documents.py

@@ -0,0 +1,173 @@
+"""Documents and blocks frozen into the shared prefix."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+from pydantic import BaseModel, Field
+
+from forkreduce.errors import ForkReduceError
+from forkreduce.textutil import strip_control_chars
+
+END_MARKER = "===END_FIXED_DOCUMENTS==="
+
+
+class Block(BaseModel):
+    """One addressable span of a document.
+
+    Citation checks look up ``block_id`` and search for a quote inside ``text``.
+    """
+
+    block_id: str
+    text: str
+    meta: dict[str, str] = Field(default_factory=dict)
+
+
+class Document(BaseModel):
+    """A document made of blocks. ``doc_id`` is stable and appears in the rendered prefix."""
+
+    doc_id: str
+    title: str = ""
+    blocks: list[Block]
+
+
+@dataclass
+class DocumentIndex:
+    """Cleaned documents plus the exact text that is frozen into the prefix."""
+
+    documents: list[Document]
+    blocks: dict[str, Block] = field(default_factory=dict)
+    rendered: str = ""
+
+    def block_text(self, block_id: str) -> str | None:
+        block = self.blocks.get(block_id)
+        if block is None:
+            return None
+        return block.text
+
+
+DocumentsInput = str | Document | Block | Sequence[str | Document | Block]
+
+
+def coerce_documents(documents: DocumentsInput) -> DocumentIndex:
+    """Normalize caller input into cleaned blocks and a stable rendering.
+
+    A bare string becomes one document with a single block ``b1``. A sequence of
+    strings becomes blocks ``b1``, ``b2``, ... inside one document. Control
+    characters are stripped before the text is hashed or shown to the model, so a
+    null byte cannot break JSON tool arguments later.
+    """
+    docs = _coerce_many(documents)
+    if not docs:
+        raise ForkReduceError("documents must contain at least one block of text")
+    blocks: dict[str, Block] = {}
+    for doc in docs:
+        if not doc.blocks:
+            raise ForkReduceError(f"document {doc.doc_id!r} has no blocks")
+        for block in doc.blocks:
+            if not block.block_id:
+                raise ForkReduceError("block_id must be a non-empty string")
+            if block.block_id in blocks:
+                raise ForkReduceError(f"duplicate block_id {block.block_id!r}")
+            if not block.text.strip():
+                raise ForkReduceError(f"block {block.block_id!r} is empty after cleaning")
+            blocks[block.block_id] = block
+    rendered = render_documents(docs)
+    return DocumentIndex(documents=docs, blocks=blocks, rendered=rendered)
+
+
+def render_documents(documents: Sequence[Document]) -> str:
+    """Render documents into the frozen prefix body, ending with :data:`END_MARKER`."""
+    lines: list[str] = []
+    for doc in documents:
+        title = f" title={doc.title}" if doc.title else ""
+        lines.append(f"[document id={doc.doc_id}{title}]")
+        for block in doc.blocks:
+            lines.append(f"[block {block.block_id}]")
+            lines.append(block.text)
+            lines.append("")
+    lines.append(END_MARKER)
+    return "\n".join(lines)
+
+
+def _coerce_many(documents: DocumentsInput) -> list[Document]:
+    if isinstance(documents, str):
+        return [_string_document("doc", [documents])]
+    if isinstance(documents, Document):
+        return [_clean_document(documents)]
+    if isinstance(documents, Block):
+        return [_blocks_document("doc", [documents])]
+    if isinstance(documents, Sequence):
+        # A homogeneous sequence keeps the historical single-document shape.
+        # A mixed sequence keeps the caller's order: consecutive loose blocks
+        # or strings become one generated document at that position.
+        if documents and all(isinstance(item, str) for item in documents):
+            return [_string_document("doc", [item for item in documents if isinstance(item, str)])]
+        if documents and all(isinstance(item, Block) for item in documents):
+            return [_blocks_document("doc", [item for item in documents if isinstance(item, Block)])]
+        docs: list[Document] = []
+        pending_blocks: list[Block] = []
+        pending_strings: list[str] = []
+
+        def flush() -> None:
+            nonlocal pending_blocks, pending_strings
+            if pending_blocks:
+                docs.append(_blocks_document(_generated_doc_id(docs), pending_blocks))
+                pending_blocks = []
+            if pending_strings:
+                docs.append(_string_document(_generated_doc_id(docs), pending_strings))
+                pending_strings = []
+
+        for item in documents:
+            if isinstance(item, str):
+                if pending_blocks:
+                    flush()
+                pending_strings.append(item)
+            elif isinstance(item, Block):
+                if pending_strings:
+                    flush()
+                pending_blocks.append(item)
+            elif isinstance(item, Document):
+                flush()
+                docs.append(_clean_document(item))
+            else:
+                raise ForkReduceError("documents must be a string, Block, Document, or a sequence of those")
+        flush()
+        return docs
+    raise ForkReduceError("documents must be a string, Block, Document, or a sequence of those")
+
+
+def _generated_doc_id(existing: list[Document]) -> str:
+    used = {doc.doc_id for doc in existing}
+    if "doc" not in used:
+        return "doc"
+    number = 2
+    while f"doc-{number}" in used:
+        number += 1
+    return f"doc-{number}"
+
+
+def _string_document(doc_id: str, parts: list[str]) -> Document:
+    blocks: list[Block] = []
+    for index, part in enumerate(parts, start=1):
+        blocks.append(Block(block_id=f"b{index}", text=strip_control_chars(part)))
+    return Document(doc_id=doc_id, blocks=blocks)
+
+
+def _blocks_document(doc_id: str, blocks: Sequence[Block]) -> Document:
+    cleaned = [
+        Block(block_id=block.block_id, text=strip_control_chars(block.text), meta=dict(block.meta)) for block in blocks
+    ]
+    return Document(doc_id=doc_id, blocks=cleaned)
+
+
+def _clean_document(document: Document) -> Document:
+    return Document(
+        doc_id=document.doc_id,
+        title=strip_control_chars(document.title),
+        blocks=[
+            Block(block_id=block.block_id, text=strip_control_chars(block.text), meta=dict(block.meta))
+            for block in document.blocks
+        ],
+    )

forkreduce/errors.py

@@ -0,0 +1,39 @@
+"""Errors raised for caller mistakes.
+
+Model, transport, and budget failures are returned as :class:`forkreduce.failures.Failure`
+on a result object. The exceptions in this module mean the caller asked for something the
+session cannot do, and they are raised before a model call when that is possible.
+"""
+
+from __future__ import annotations
+
+
+class ForkReduceError(Exception):
+    """Base class for forkreduce caller errors.
+
+    ``code`` is an optional machine-readable label. A provider context-length
+    error is a :class:`forkreduce.failures.Failure` with code ``context_overflow``,
+    not an exception from this class.
+    """
+
+    code: str
+
+    def __init__(self, message: str, *, code: str = "") -> None:
+        super().__init__(message)
+        self.code = code
+
+
+class UnregisteredOutputError(ForkReduceError):
+    """A query or reduce asked for an output type that was not registered at build time.
+
+    Registering it later would change the tool set and invalidate the shared prefix cache,
+    so the call fails immediately instead of silently sending a different schema.
+    """
+
+
+class PrefixMismatchWarning(UserWarning):
+    """The supplied message history does not start with this session's frozen prefix.
+
+    The call still runs. Provider prefix caches will miss, because the shared head of the
+    prompt is no longer identical to the other branches.
+    """

forkreduce/executor.py

@@ -0,0 +1,200 @@
+"""Single-attempt execution and the default local executor.
+
+:class:`LocalExecutor` is what :meth:`Session.ask` uses: a semaphore, bounded
+retries, and backoff for transport errors. One attempt lives on
+:meth:`Session.run_once` so a Temporal activity can do the same call and let
+Temporal own the retry.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Any, Protocol
+
+from pydantic_ai.usage import RunUsage
+
+from forkreduce.failures import Failure, budget_failure
+from forkreduce.policy import Accept, BudgetReport, Fail, Policy
+from forkreduce.query import Query, QueryResult
+from forkreduce.stats import add_usage
+
+if TYPE_CHECKING:
+    from forkreduce.session import Session
+
+
+class Executor(Protocol):
+    """Runs one query to a :class:`QueryResult`, including whatever retries it owns."""
+
+    async def run_query(self, session: Session, query: Query[Any]) -> QueryResult[Any]:
+        """Execute ``query`` against ``session``."""
+        ...
+
+
+@dataclass
+class BudgetLedger:
+    policy: Policy
+    requests: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    _lock: asyncio.Lock | None = None
+    _token_gate: asyncio.Lock | None = None
+
+    def _lock_for(self) -> asyncio.Lock:
+        if self._lock is None:
+            self._lock = asyncio.Lock()
+        return self._lock
+
+    async def try_reserve(self) -> BudgetReport | None:
+        budget = self.policy.budget
+        if budget is None:
+            return None
+        async with self._lock_for():
+            problem = self._problem()
+            if problem is not None:
+                return problem
+            self.requests += 1
+            return None
+
+    async def force_reserve(self) -> None:
+        async with self._lock_for():
+            self.requests += 1
+
+    async def add_usage(self, usage: RunUsage) -> None:
+        async with self._lock_for():
+            self.input_tokens += usage.input_tokens
+            self.output_tokens += usage.output_tokens
+
+    def limits_tokens(self) -> bool:
+        """Whether a token ceiling requires calls to run one at a time."""
+        budget = self.policy.budget
+        if budget is None:
+            return False
+        return any(
+            value is not None for value in (budget.max_input_tokens, budget.max_output_tokens, budget.max_total_tokens)
+        )
+
+    @asynccontextmanager
+    async def token_turn(self) -> AsyncIterator[None]:
+        """Hold the token gate for one call when a token ceiling is set.
+
+        Request-only budgets do not take this gate, so they stay concurrent.
+        """
+        if not self.limits_tokens():
+            yield
+            return
+        async with self._lock_for():
+            if self._token_gate is None:
+                self._token_gate = asyncio.Lock()
+            gate = self._token_gate
+        await gate.acquire()
+        try:
+            yield
+        finally:
+            gate.release()
+
+    def _problem(self) -> BudgetReport | None:
+        budget = self.policy.budget
+        if budget is None:
+            return None
+        message: str | None = None
+        if budget.max_requests is not None and self.requests >= budget.max_requests:
+            message = f"request budget exhausted ({self.requests}/{budget.max_requests})"
+        elif budget.max_input_tokens is not None and self.input_tokens >= budget.max_input_tokens:
+            message = f"input token budget exhausted ({self.input_tokens}/{budget.max_input_tokens})"
+        elif budget.max_output_tokens is not None and self.output_tokens >= budget.max_output_tokens:
+            message = f"output token budget exhausted ({self.output_tokens}/{budget.max_output_tokens})"
+        elif budget.max_total_tokens is not None and (
+            self.input_tokens + self.output_tokens >= budget.max_total_tokens
+        ):
+            total = self.input_tokens + self.output_tokens
+            message = f"total token budget exhausted ({total}/{budget.max_total_tokens})"
+        if message is None:
+            return None
+        return BudgetReport(
+            kind="budget",
+            message=message,
+            requests=self.requests,
+            input_tokens=self.input_tokens,
+            output_tokens=self.output_tokens,
+        )
+
+
+class LocalExecutor:
+    """asyncio executor with a concurrency semaphore and bounded retries.
+
+    Validation failures retry immediately. Transport and rate-limit failures
+    sleep with exponential backoff. Permanent failures and budget exhaustion
+    are returned on the result and are not retried.
+    """
+
+    def __init__(self, policy: Policy) -> None:
+        self.policy = policy
+        self._semaphore = asyncio.Semaphore(policy.concurrency)
+
+    async def run_query(self, session: Session, query: Query[Any]) -> QueryResult[Any]:
+        async with self._semaphore:
+            return await self._run_with_retries(session, query)
+
+    async def _run_with_retries(self, session: Session, query: Query[Any]) -> QueryResult[Any]:
+        usage = RunUsage()
+        cache = 0
+        elapsed = 0.0
+        last: QueryResult[object] | None = None
+        attempts = 0
+        for attempt in range(self.policy.retries + 1):
+            attempts = attempt + 1
+            result = await session.run_once(query)
+            usage = add_usage(usage, result.usage)
+            cache += result.cache_hit_tokens
+            elapsed += result.elapsed_ms
+            last = result
+            # A budget rejection happens before any model call.
+            if result.attempts == 0:
+                return self._folded(result, usage, cache, elapsed, 0)
+            if result.failure is None:
+                return self._folded(result, usage, cache, elapsed, attempts)
+            if not result.failure.retryable or attempt >= self.policy.retries:
+                return self._folded(result, usage, cache, elapsed, attempts)
+            if result.failure.kind.value == "transport":
+                await self._sleep(attempt + 1)
+        assert last is not None
+        return self._folded(last, usage, cache, elapsed, attempts)
+
+    def _folded(
+        self,
+        result: QueryResult[object],
+        usage: RunUsage,
+        cache: int,
+        elapsed: float,
+        attempts: int,
+    ) -> QueryResult[object]:
+        failure = result.failure.with_attempts(attempts) if result.failure is not None else None
+        return replace(
+            result,
+            usage=usage,
+            cache_hit_tokens=cache,
+            elapsed_ms=elapsed,
+            attempts=attempts,
+            failure=failure,
+        )
+
+    async def _sleep(self, attempt: int) -> None:
+        delay = min(self.policy.retry_backoff * (2 ** (attempt - 1)), self.policy.max_backoff)
+        if delay > 0:
+            await asyncio.sleep(delay)
+
+
+def budget_action_failure(report: BudgetReport, policy: Policy) -> Failure | None:
+    """Apply the budget policy. ``None`` means the call may proceed anyway."""
+    action: Accept | Fail
+    if policy.budget_policy is None:
+        action = Fail(report.message)
+    else:
+        decided = policy.budget_policy(report)
+        if isinstance(decided, Accept):
+            return None
+        action = decided if isinstance(decided, Fail) else Fail(report.message)
+    return budget_failure(action.reason)

forkreduce/failures.py

@@ -0,0 +1,239 @@
+"""Failure classification.
+
+Every model call ends as either a value or a :class:`Failure`. Nothing is dropped
+on the floor: transport and validation failures are retryable up to the policy
+bound, and permanent failures are returned to the caller.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import re
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any
+
+from pydantic import ValidationError
+from pydantic_ai.exceptions import (
+    ModelAPIError,
+    ModelHTTPError,
+    UnexpectedModelBehavior,
+    UsageLimitExceeded,
+    UserError,
+)
+
+try:
+    import httpx
+except ImportError:  # pragma: no cover - pydantic-ai always brings an http client
+    httpx = None  # type: ignore[assignment]
+
+
+class FailureKind(StrEnum):
+    """Coarse class used for retry decisions and Temporal non-retryable mapping."""
+
+    VALIDATION = "validation"
+    TRANSPORT = "transport"
+    PERMANENT = "permanent"
+    BUDGET = "budget"
+
+
+# 408/429 and the usual gateway/unavailable statuses are worth another attempt.
+# Other 4xx responses (auth, bad request, not found) will not succeed if repeated.
+_RETRYABLE_STATUS = frozenset({408, 409, 425, 429, 500, 502, 503, 504})
+
+
+@dataclass
+class Failure:
+    """A classified failure attached to a query or a run.
+
+    ``retryable`` is the executor's signal. Temporal maps it to
+    ``ApplicationError(non_retryable=not retryable)``.
+    """
+
+    kind: FailureKind
+    message: str
+    retryable: bool
+    code: str = ""
+    attempts: int = 1
+    detail: dict[str, Any] = field(default_factory=dict)
+
+    def with_attempts(self, attempts: int) -> Failure:
+        return Failure(
+            kind=self.kind,
+            message=self.message,
+            retryable=self.retryable,
+            code=self.code,
+            attempts=attempts,
+            detail=dict(self.detail),
+        )
+
+
+def classify_exception(exc: BaseException) -> Failure:
+    """Map an exception from a model call onto a :class:`Failure`."""
+    message = _short_message(exc)
+    if isinstance(exc, ModelHTTPError):
+        overflow = _context_overflow_from_provider(exc)
+        if overflow is not None:
+            return overflow
+        retryable = exc.status_code in _RETRYABLE_STATUS or exc.status_code >= 500
+        kind = FailureKind.TRANSPORT if retryable else FailureKind.PERMANENT
+        return Failure(
+            kind=kind,
+            message=message,
+            retryable=retryable,
+            code=f"http_{exc.status_code}",
+            detail={"status_code": exc.status_code},
+        )
+    if isinstance(exc, ModelAPIError):
+        return Failure(
+            kind=FailureKind.TRANSPORT,
+            message=message,
+            retryable=True,
+            code="transport",
+        )
+    if httpx is not None and isinstance(exc, (httpx.TimeoutException, httpx.TransportError)):
+        return Failure(
+            kind=FailureKind.TRANSPORT,
+            message=message,
+            retryable=True,
+            code="transport",
+        )
+    if isinstance(exc, (TimeoutError, asyncio.TimeoutError, ConnectionError)):
+        return Failure(
+            kind=FailureKind.TRANSPORT,
+            message=message,
+            retryable=True,
+            code="transport",
+        )
+    if isinstance(exc, (UnexpectedModelBehavior, ValidationError, UsageLimitExceeded)):
+        return Failure(
+            kind=FailureKind.VALIDATION,
+            message=message,
+            retryable=True,
+            code="validation",
+        )
+    if isinstance(exc, UserError):
+        return Failure(
+            kind=FailureKind.PERMANENT,
+            message=message,
+            retryable=False,
+            code="user_error",
+        )
+    return Failure(
+        kind=FailureKind.PERMANENT,
+        message=message,
+        retryable=False,
+        code="unexpected",
+        detail={"exception_type": type(exc).__name__},
+    )
+
+
+def wrong_output_failure(expected: str, actual: str) -> Failure:
+    """The model called a registered output tool other than the one this query asked for."""
+    return Failure(
+        kind=FailureKind.VALIDATION,
+        message=f"Expected output {expected}, got {actual}",
+        retryable=True,
+        code="wrong_output_type",
+        detail={"expected": expected, "actual": actual},
+    )
+
+
+def context_overflow_failure(
+    *,
+    estimated_tokens: int | None,
+    limit: int | None,
+    where: str,
+    reported_tokens: int | None = None,
+) -> Failure:
+    """Permanent failure when the provider reports that one prompt does not fit.
+
+    The local character estimate never produces this failure. ``where`` is
+    ``provider``. The library does not truncate or split the prompt.
+    """
+    size = reported_tokens if reported_tokens is not None else estimated_tokens
+    size_label = "reported" if reported_tokens is not None else "estimated"
+    size_text = "unknown" if size is None else str(size)
+    limit_text = "unknown" if limit is None else str(limit)
+    detail: dict[str, Any] = {"where": where}
+    if estimated_tokens is not None:
+        detail["estimated_tokens"] = estimated_tokens
+    if reported_tokens is not None:
+        detail["reported_tokens"] = reported_tokens
+    if limit is not None:
+        detail["limit"] = limit
+    return Failure(
+        kind=FailureKind.PERMANENT,
+        message=(f"context overflow: {size_label} {size_text} tokens exceeds context limit {limit_text} ({where})"),
+        retryable=False,
+        code="context_overflow",
+        detail=detail,
+    )
+
+
+def budget_failure(message: str) -> Failure:
+    return Failure(
+        kind=FailureKind.BUDGET,
+        message=message,
+        retryable=False,
+        code="budget_exceeded",
+    )
+
+
+_OVERFLOW_MARKERS = (
+    "context_length_exceeded",
+    "context_overflow",
+    "maximum context length",
+    "context window",
+    "too many tokens",
+    "prompt is too long",
+    "input is too long",
+    "exceeds the context",
+    "exceeds the model's maximum context",
+)
+
+
+def _context_overflow_from_provider(exc: ModelHTTPError) -> Failure | None:
+    text = _provider_error_text(exc)
+    if not any(marker in text.lower() for marker in _OVERFLOW_MARKERS):
+        return None
+    reported, reported_limit = _reported_context_sizes(text)
+    failure = context_overflow_failure(
+        estimated_tokens=None,
+        limit=reported_limit,
+        where="provider",
+        reported_tokens=reported,
+    )
+    failure.detail["status_code"] = exc.status_code
+    return failure
+
+
+def _provider_error_text(exc: ModelHTTPError) -> str:
+    body = exc.body
+    if isinstance(body, str):
+        rendered = body
+    elif body is None:
+        rendered = ""
+    else:
+        rendered = str(body)
+    return f"{exc} {rendered}"
+
+
+def _reported_context_sizes(text: str) -> tuple[int | None, int | None]:
+    """Pull ``(requested tokens, window)`` out of a provider error, when present."""
+    match = re.search(
+        r"maximum context length is (\d+).*?requested (\d+)",
+        text,
+        flags=re.IGNORECASE | re.DOTALL,
+    )
+    if match is None:
+        return None, None
+    return int(match.group(2)), int(match.group(1))
+
+
+def _short_message(exc: BaseException) -> str:
+    text = str(exc).strip() or type(exc).__name__
+    text = " ".join(text.split())
+    if len(text) > 500:
+        text = text[:500] + "..."
+    return f"{type(exc).__name__}: {text}"