forkreduce 0.1.0__py3-none-any.whl

forkreduce/followup.py

@@ -0,0 +1,215 @@
+"""The hook between rounds, and the built-in gap-fill strategy.
+
+``followup`` receives the round's merged output and query results. It may call
+the model (another query on the same prefix) or stay a pure function. An empty
+list ends the run, as does reaching ``max_rounds``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from pydantic import BaseModel
+
+from forkreduce.coverage import ReduceReport
+from forkreduce.failures import Failure, FailureKind
+from forkreduce.policy import Accept, Action, Drop, Fail
+from forkreduce.query import Query, QueryResult
+
+FollowupFn = Callable[["RoundState", Any], Awaitable[list[Query[Any]]] | list[Query[Any]]]
+
+
+@dataclass
+class RoundState:
+    """What a followup hook can see after one map/check/reduce round.
+
+    ``side_results`` collects model calls the hook makes (gap-fill queries, for
+    example) so their usage is not dropped. ``failures`` is the same idea for
+    errors the hook observed.
+    """
+
+    round: int
+    results: list[QueryResult[Any]]
+    outputs: list[Any]
+    merged: Any
+    report: ReduceReport | None
+    failures: list[Failure] = field(default_factory=list)
+    side_results: list[QueryResult[Any]] = field(default_factory=list)
+
+    def add_result(self, result: QueryResult[Any]) -> None:
+        self.side_results.append(result)
+        if result.failure is not None:
+            self.failures.append(result.failure)
+
+
+@dataclass(frozen=True)
+class Scope:
+    """One slice of the document a gap-fill query is responsible for."""
+
+    name: str
+    responsibility: str
+
+
+@dataclass
+class LeadReport:
+    """Positions a gap-fill pass already covered, and the new leads it found."""
+
+    kind: str
+    new_positions: list[str]
+    duplicate_positions: list[str]
+    leads: list[Any]
+
+
+class LeadPolicy(Protocol):
+    """Decide which new leads become the next round's queries."""
+
+    def decide(self, report: LeadReport) -> Action:
+        """Return an action for ``report``."""
+        ...
+
+
+class AcceptLeads:
+    """Default lead policy: keep every lead whose position is not already covered."""
+
+    def decide(self, report: LeadReport) -> Action:
+        return Accept()
+
+
+@dataclass
+class GapFill:
+    """One gap-fill query per scope, then a query per new lead.
+
+    Leads are deduped by position (``block_id`` by default) against the merged
+    output and against each other. The gap-fill query itself runs inside the
+    hook, on the same session, so it shares the prefix cache. Its
+    :class:`QueryResult` is attached to the round state.
+
+    ``lead_output`` and ``item_output`` must both be registered on the session.
+    """
+
+    lead_output: type[BaseModel]
+    item_output: type[BaseModel]
+    scopes: Sequence[Scope] | None = None
+    instructions: str = (
+        "Look for requirements in your scope that the merged list does not already cover. "
+        "Return leads only for those gaps."
+    )
+    policy: LeadPolicy | None = None
+    lead_items: Callable[[Any], Sequence[Any]] | None = None
+    lead_position: Callable[[Any], str] | None = None
+    covered_positions: Callable[[Any], Sequence[str]] | None = None
+    query_for_lead: Callable[[Any, RoundState], Query[Any]] | None = None
+
+    async def __call__(self, state: RoundState, session: Any) -> list[Query[Any]]:
+        scopes = list(self.scopes) if self.scopes else [Scope("all", "the whole document")]
+        covered = set(self._covered(state.merged))
+        fresh: list[Any] = []
+        duplicates: list[str] = []
+        for scope in scopes:
+            query = Query(
+                tail=self._gap_tail(state.merged, scope),
+                output=self.lead_output,
+                id=f"gapfill-r{state.round}-{scope.name}",
+                round=state.round,
+            )
+            result = await session.ask(query)
+            state.add_result(result)
+            if result.failure is not None or result.output is None:
+                continue
+            for lead in self._leads(result.output):
+                position = self._position(lead)
+                if not position or position in covered:
+                    duplicates.append(position or "")
+                    continue
+                covered.add(position)
+                fresh.append(lead)
+        report = LeadReport(
+            kind="leads",
+            new_positions=[self._position(lead) for lead in fresh],
+            duplicate_positions=[item for item in duplicates if item],
+            leads=list(fresh),
+        )
+        action = (self.policy or AcceptLeads()).decide(report)
+        if isinstance(action, Fail):
+            state.failures.append(
+                Failure(
+                    kind=FailureKind.PERMANENT,
+                    message=action.reason,
+                    retryable=False,
+                    code="leads_rejected",
+                )
+            )
+            return []
+        if isinstance(action, Drop):
+            banned = set(action.items)
+            fresh = [lead for lead in fresh if self._position(lead) not in banned]
+        elif not isinstance(action, Accept):
+            state.failures.append(
+                Failure(
+                    kind=FailureKind.PERMANENT,
+                    message=f"lead policy returned {type(action).__name__}, which gap-fill cannot run",
+                    retryable=False,
+                    code="bad_action",
+                )
+            )
+            return []
+        return [self._to_query(lead, state) for lead in fresh]
+
+    def _leads(self, output: Any) -> Sequence[Any]:
+        if self.lead_items is not None:
+            return list(self.lead_items(output))
+        for name in ("leads", "items"):
+            value = getattr(output, name, None)
+            if isinstance(value, list):
+                return value
+        return []
+
+    def _position(self, lead: Any) -> str:
+        if self.lead_position is not None:
+            return self.lead_position(lead)
+        block_id = getattr(lead, "block_id", None)
+        if isinstance(block_id, str):
+            return block_id
+        return ""
+
+    def _covered(self, merged: Any) -> list[str]:
+        if self.covered_positions is not None:
+            return list(self.covered_positions(merged))
+        return sorted(_collect_block_ids(merged))
+
+    def _to_query(self, lead: Any, state: RoundState) -> Query[Any]:
+        if self.query_for_lead is not None:
+            return self.query_for_lead(lead, state)
+        question = getattr(lead, "question", None) or getattr(lead, "text", None) or str(lead)
+        position = self._position(lead) or "lead"
+        return Query(
+            id=f"gap-r{state.round}-{position}",
+            tail=str(question),
+            output=self.item_output,
+            meta={"block_id": position, "gap": True},
+            round=state.round + 1,
+        )
+
+    def _gap_tail(self, merged: Any, scope: Scope) -> str:
+        if isinstance(merged, BaseModel):
+            rendered = merged.model_dump_json()
+        else:
+            rendered = "null" if merged is None else str(merged)
+        return f"Scope {scope.name}: {scope.responsibility}\n{self.instructions.strip()}\n\nAlready merged:\n{rendered}"
+
+
+def _collect_block_ids(value: Any) -> set[str]:
+    found: set[str] = set()
+    if isinstance(value, BaseModel):
+        if "block_id" in type(value).model_fields:
+            block_id = getattr(value, "block_id")  # noqa: B009 - field is optional on BaseModel
+            if isinstance(block_id, str) and block_id:
+                found.add(block_id)
+        for name in type(value).model_fields:
+            found |= _collect_block_ids(getattr(value, name))
+    elif isinstance(value, list):
+        for item in value:
+            found |= _collect_block_ids(item)
+    return found

forkreduce/models.py

@@ -0,0 +1,48 @@
+"""OpenAI-compatible models.
+
+forkreduce itself accepts any pydantic-ai ``Model``. This helper is only for the
+common case of an OpenAI-style chat endpoint. Retries stay in forkreduce:
+the OpenAI client is constructed with ``max_retries=0``.
+"""
+
+from __future__ import annotations
+
+import os
+
+from pydantic_ai.models.openai import OpenAIChatModel
+
+
+def openai_compatible_model(
+    model_name: str | None = None,
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+) -> OpenAIChatModel:
+    """Build an ``OpenAIChatModel`` from arguments or environment variables.
+
+    Environment variables, when the matching argument is omitted:
+
+    * ``OPENAI_MODEL`` (or ``FORKREDUCE_MODEL``) for the model name
+    * ``OPENAI_API_KEY``
+    * ``OPENAI_BASE_URL``
+    """
+    from openai import AsyncOpenAI
+    from pydantic_ai.providers.openai import OpenAIProvider
+
+    from forkreduce.errors import ForkReduceError
+
+    name = model_name or os.environ.get("OPENAI_MODEL") or os.environ.get("FORKREDUCE_MODEL")
+    if not name:
+        raise ForkReduceError("Pass model_name or set OPENAI_MODEL")
+    key = api_key if api_key is not None else os.environ.get("OPENAI_API_KEY")
+    url = base_url if base_url is not None else os.environ.get("OPENAI_BASE_URL")
+    # A local OpenAI-compatible server often has no key. The client still wants
+    # a non-empty string when base_url is set.
+    resolved_key = key or ("api-key-not-set" if url else None)
+    if url and resolved_key is not None:
+        client = AsyncOpenAI(api_key=resolved_key, base_url=url, max_retries=0)
+    elif resolved_key is not None:
+        client = AsyncOpenAI(api_key=resolved_key, max_retries=0)
+    else:
+        client = AsyncOpenAI(max_retries=0)
+    return OpenAIChatModel(name, provider=OpenAIProvider(openai_client=client))

forkreduce/plan.py

@@ -0,0 +1,355 @@
+"""Chained map / check / filter / reduce / followup plan."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+from pydantic_ai.usage import RunUsage
+
+from forkreduce.checks import Check, CheckReport, Repair
+from forkreduce.context import ContextWarning
+from forkreduce.coverage import ReduceReport
+from forkreduce.failures import Failure, FailureKind
+from forkreduce.followup import FollowupFn, RoundState
+from forkreduce.query import Query, QueryResult
+from forkreduce.reduce import (
+    ReduceItem,
+    annotate_round,
+    flatten_merged,
+    reduce_items,
+)
+from forkreduce.session import Session
+from forkreduce.stats import UsageSummary, add_usage, summarize
+
+FilterFn = Callable[[Any], bool]
+
+
+@dataclass
+class RoundRecord:
+    """One map/check/reduce round."""
+
+    round: int
+    results: list[QueryResult[Any]]
+    outputs: list[Any]
+    merged: Any
+    report: ReduceReport | None
+    failures: list[Failure]
+    check_reports: list[CheckReport] = field(default_factory=list)
+
+
+@dataclass
+class RunResult:
+    """The full plan: merged output, per-round detail, usage, and every failure."""
+
+    merged: Any
+    rounds: list[RoundRecord]
+    results: list[QueryResult[Any]]
+    usage: UsageSummary
+    failures: list[Failure]
+    reports: list[Any]
+    context_warning: ContextWarning | None = None
+
+    @property
+    def cache_hit_rate(self) -> float:
+        return self.usage.cache_hit_rate
+
+    @property
+    def context_warnings(self) -> list[ContextWarning]:
+        """Prefix warning from the session, then each query warning, in run order."""
+        found: list[ContextWarning] = []
+        if self.context_warning is not None:
+            found.append(self.context_warning)
+        for result in self.results:
+            if result.context_warning is not None:
+                found.append(result.context_warning)
+        return found
+
+
+class Plan:
+    """A lazy pipeline. Nothing runs until :meth:`run`."""
+
+    def __init__(self, session: Session, queries: list[Query[Any]]) -> None:
+        self._session = session
+        self._queries = queries
+        self._check: Check | None = None
+        self._repair: Repair | None = None
+        self._failure_mode = session.policy.repair_failure
+        self._filters: list[FilterFn] = []
+        self._reduce_output: type[Any] | None = None
+        self._reduce_instructions = ""
+        self._reduce_check: Any = None
+        self._reduce_max_batch = 200
+        self._reduce_coverage: Any = None
+        self._reduce_max_rounds = 3
+        self._followup: FollowupFn | None = None
+        self._max_rounds = 1
+
+    def check(
+        self,
+        check: Check,
+        *,
+        repair: Repair | None = None,
+        on_failure: str | None = None,
+    ) -> Plan:
+        """Validate each output. ``repair`` may continue the branch to fix it.
+
+        ``on_failure`` overrides :attr:`Policy.repair_failure` for items that
+        are still failing when repair stops: ``needs_review``, ``drop``, or
+        ``fail``.
+        """
+        if on_failure is not None and on_failure not in {"needs_review", "drop", "fail"}:
+            raise ValueError("on_failure must be 'needs_review', 'drop', or 'fail'")
+        self._check = check
+        self._repair = repair
+        if on_failure is not None:
+            self._failure_mode = on_failure
+        return self
+
+    def filter(self, fn: FilterFn) -> Plan:
+        """Keep outputs for which ``fn`` returns true. Others stay on the run, not in reduce."""
+        self._filters.append(fn)
+        return self
+
+    def reduce(
+        self,
+        *,
+        output: type[Any],
+        instructions: str = "",
+        check: Any = None,
+        max_batch: int = 200,
+        coverage: Any = None,
+        max_rounds: int = 3,
+    ) -> Plan:
+        """Semantically merge the kept outputs. ``output`` must be registered."""
+        self._reduce_output = output
+        self._reduce_instructions = instructions
+        self._reduce_check = check
+        self._reduce_max_batch = max_batch
+        self._reduce_coverage = coverage
+        self._reduce_max_rounds = max_rounds
+        return self
+
+    def followup(self, fn: FollowupFn, *, max_rounds: int = 2) -> Plan:
+        """Call ``fn`` between rounds. ``max_rounds`` counts map rounds, including the first."""
+        if max_rounds < 1:
+            raise ValueError("max_rounds must be >= 1")
+        self._followup = fn
+        self._max_rounds = max_rounds
+        return self
+
+    async def run(self) -> RunResult:
+        """Execute the pipeline."""
+        rounds: list[RoundRecord] = []
+        all_results: list[QueryResult[Any]] = []
+        all_failures: list[Failure] = []
+        reports: list[Any] = []
+        merged: Any = None
+        carried: list[ReduceItem] = []
+        queries = list(self._queries)
+        fatal = False
+
+        for round_index in range(1, self._max_rounds + 1):
+            if not queries:
+                break
+            for query in queries:
+                query.round = round_index
+            results = await self._session.fork(queries)
+            checked, check_reports, round_failures, stop = await self._apply_checks(results)
+            all_failures.extend(round_failures)
+            reports.extend(check_reports)
+            kept: list[QueryResult[Any]] = []
+            new_items: list[ReduceItem] = []
+            for result in checked:
+                if round_index > 0 and result.output is not None:
+                    result.output = annotate_round(result.output, round_index)
+                if result.dropped or result.output is None:
+                    continue
+                if result.failure is not None and not result.needs_review:
+                    continue
+                if not self._keep(result.output):
+                    continue
+                kept.append(result)
+                new_items.append(ReduceItem(id=result.query_id, value=result.output, round=round_index))
+            report: ReduceReport | None = None
+            reduce_calls: list[QueryResult[Any]] = []
+            reduce_failure: Failure | None = None
+            if self._reduce_output is not None and (new_items or carried) and not stop:
+                outcome = await reduce_items(
+                    self._session,
+                    [*carried, *new_items],
+                    output=self._reduce_output,
+                    instructions=self._reduce_instructions,
+                    check=self._reduce_check,
+                    max_batch=self._reduce_max_batch,
+                    coverage=self._reduce_coverage,
+                    max_rounds=self._reduce_max_rounds,
+                    round_index=round_index,
+                )
+                merged = outcome.output
+                report = outcome.report
+                reports.append(report)
+                reduce_calls = list(outcome.calls)
+                if outcome.failure is not None:
+                    all_failures.append(outcome.failure)
+                    reduce_failure = outcome.failure
+                if outcome.fatal:
+                    fatal = True
+                carried = flatten_merged(merged, round_index=round_index)
+            round_results = [*checked, *reduce_calls]
+            all_results.extend(round_results)
+            rounds.append(
+                RoundRecord(
+                    round=round_index,
+                    results=round_results,
+                    outputs=[item.value for item in new_items],
+                    merged=merged,
+                    report=report,
+                    failures=[*round_failures, *([reduce_failure] if reduce_failure is not None else [])],
+                    check_reports=check_reports,
+                )
+            )
+            if fatal or stop or self._followup is None or round_index >= self._max_rounds:
+                break
+            state = RoundState(
+                round=round_index,
+                results=checked,
+                outputs=[item.value for item in new_items],
+                merged=merged,
+                report=report,
+            )
+            nxt = self._followup(state, self._session)
+            if inspect.isawaitable(nxt):
+                nxt = await nxt
+            all_results.extend(state.side_results)
+            all_failures.extend(state.failures)
+            # Side results were not part of this round record yet.
+            if rounds:
+                rounds[-1].results.extend(state.side_results)
+                rounds[-1].failures.extend(state.failures)
+            queries = list(nxt)
+            if not queries:
+                break
+
+        return RunResult(
+            merged=merged,
+            rounds=rounds,
+            results=all_results,
+            usage=summarize(all_results),
+            failures=all_failures,
+            reports=reports,
+            context_warning=self._session.context_warning,
+        )
+
+    async def _apply_checks(
+        self,
+        results: list[QueryResult[Any]],
+    ) -> tuple[list[QueryResult[Any]], list[CheckReport], list[Failure], bool]:
+        if self._check is None:
+            return results, [], [result.failure for result in results if result.failure], False
+        updated: list[QueryResult[Any]] = []
+        reports: list[CheckReport] = []
+        failures: list[Failure] = []
+        stop = False
+        for result in results:
+            if result.failure is not None or result.output is None:
+                if result.failure is not None:
+                    failures.append(result.failure)
+                updated.append(result)
+                continue
+            parts = [result]
+            output = result.output
+            messages = result.messages
+            problems = self._check(output, documents=self._session.documents)
+            repair_rounds = 0
+            while problems and self._repair is not None and repair_rounds < self._repair.max_rounds:
+                repair_rounds += 1
+                carrier = replace(result, output=output, messages=messages, query_id=result.query_id)
+                repaired = await self._session.ask(
+                    self._repair.build_query(carrier, problems, round_index=repair_rounds)
+                )
+                parts.append(repaired)
+                if repaired.failure is not None or repaired.output is None:
+                    if repaired.failure is not None:
+                        failures.append(repaired.failure)
+                    break
+                try:
+                    output = self._repair.apply(output, repaired.output, problems)
+                except Exception as exc:
+                    failures.append(
+                        Failure(
+                            kind=FailureKind.PERMANENT,
+                            message=f"repair apply failed: {exc}",
+                            retryable=False,
+                            code="repair_apply",
+                        )
+                    )
+                    break
+                messages = repaired.messages
+                problems = self._check(output, documents=self._session.documents)
+            usage, cache, attempts, elapsed = _fold_usage(parts)
+            notice = result.context_warning
+            if notice is None:
+                for part in parts:
+                    if part.context_warning is not None:
+                        notice = part.context_warning
+                        break
+            current = replace(
+                result,
+                output=output,
+                messages=messages,
+                usage=usage,
+                cache_hit_tokens=cache,
+                attempts=attempts,
+                elapsed_ms=elapsed,
+                failure=None,
+                check_failures=list(problems),
+                context_warning=notice,
+            )
+            if not problems:
+                action = "repaired" if repair_rounds else "ok"
+            else:
+                action = self._failure_mode
+                if self._failure_mode == "needs_review":
+                    current.needs_review = True
+                elif self._failure_mode == "drop":
+                    current.dropped = True
+                else:
+                    failure = Failure(
+                        kind=FailureKind.PERMANENT,
+                        message="; ".join(item.message for item in problems),
+                        retryable=False,
+                        code="check_failed",
+                    )
+                    current.failure = failure
+                    failures.append(failure)
+                    stop = True
+            reports.append(
+                CheckReport(
+                    query_id=result.query_id,
+                    round=result.round,
+                    failures=list(problems),
+                    repair_rounds=repair_rounds,
+                    action=action,
+                )
+            )
+            updated.append(current)
+        return updated, reports, failures, stop
+
+    def _keep(self, output: Any) -> bool:
+        return all(fn(output) for fn in self._filters)
+
+
+def _fold_usage(parts: Sequence[QueryResult[Any]]) -> tuple[RunUsage, int, int, float]:
+    usage = RunUsage()
+    cache = 0
+    attempts = 0
+    elapsed = 0.0
+    for part in parts:
+        usage = add_usage(usage, part.usage)
+        cache += part.cache_hit_tokens
+        attempts += part.attempts
+        elapsed += part.elapsed_ms
+    return usage, cache, attempts, elapsed