smartpipe-cli 1.3.0__py3-none-any.whl

smartpipe/verbs/join.py

@@ -0,0 +1,413 @@
+"""The ``join`` verb (D21): match stdin against a second input, semantically.
+
+Embed → block → judge: the right side (``--right``) is finite — read whole,
+embedded once (chunked), held as an in-memory index; each left item embeds,
+blocks to its top-K nearest candidates, and only those pairs reach the chat
+model with the filter-style verdict schema. N·K calls, never N·M.
+
+Fail-before-spend order (D18): flags → grammar → right file exists/parses/
+non-empty → right side fully embedded → cost preview → the first judge call.
+A bad right side costs zero chat calls.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Protocol
+
+from smartpipe.core.errors import ExitCode, ItemError, TooManyFailures, UsageFault
+from smartpipe.engine.blocking import RightIndex, build_index, candidates
+from smartpipe.engine.chunking import estimate_tokens, mean_pool, split_text
+from smartpipe.engine.prompts import (
+    JUDGE_SCHEMA,
+    build_judge_request,
+    build_repair_request,
+    parse_join_predicate,
+)
+from smartpipe.engine.ranking import rank
+from smartpipe.engine.runner import (
+    Done,
+    FailurePolicy,
+    run_ordered,
+    should_halt,
+    should_halt_consecutive,
+)
+from smartpipe.engine.schema import validate_and_coerce
+from smartpipe.io import diagnostics, readers
+from smartpipe.io.inputs import STDIN
+from smartpipe.io.items import ItemSource, describe_source, item_from_line
+from smartpipe.io.progress import make_stderr_spinner
+from smartpipe.verbs.common import (
+    embed_budget,
+    embed_in_batches,
+    ensure_text,
+    interrupted_exit_code,
+    outcome_exit_code,
+)
+from smartpipe.verbs.convert import Converter, make_converter
+
+if TYPE_CHECKING:
+    from pathlib import Path
+    from typing import TextIO
+
+    from smartpipe.engine.prompts import Token
+    from smartpipe.io.inputs import InputSpec
+    from smartpipe.io.items import Item
+    from smartpipe.io.writers import OutputFormat, ResultWriter
+    from smartpipe.models.base import ChatModel, EmbeddingModel, ModelRef
+    from smartpipe.models.stt import RemoteTranscriber
+
+__all__ = ["JoinContext", "JoinRequest", "PairBook", "run_join"]
+
+_PREVIEW_THRESHOLD = 200  # estimated judge calls before the cost line appears (D21)
+
+
+@dataclass(frozen=True, slots=True)
+class JoinRequest:
+    predicate: str
+    right: Path
+    k: int
+    threshold: float | None
+    model_flag: str | None
+    embed_model_flag: str | None
+    concurrency_flag: int | None
+    output: OutputFormat
+    input: InputSpec = STDIN
+    fields: tuple[str, ...] | None = None
+    unmatched: Path | None = None  # write zero-match left items here, verbatim
+    allow_captions: bool = False  # cloud conversions opt-in (D33)
+    kind: str = "inner"  # inner | leftouter | anti (D38/11)
+
+
+class JoinContext(Protocol):
+    def remote_transcriber(self, chat_ref: ModelRef | None = None) -> RemoteTranscriber | None: ...
+
+    """The first verb that needs BOTH models — the container already has both."""
+
+    async def chat_model(self, flag: str | None = None) -> ChatModel: ...
+    async def embedding_model(self, flag: str | None = None) -> EmbeddingModel: ...
+    def concurrency(self, flag: int | None = None) -> int: ...
+    def writer(
+        self,
+        output_flag: OutputFormat,
+        *,
+        structured: bool,
+        stdout: TextIO,
+        fields: tuple[str, ...] | None = None,
+    ) -> ResultWriter: ...
+
+
+@dataclass(slots=True)
+class PairBook:
+    """Judge-call accounting (D21: the halt policies count judge calls)."""
+
+    policy: FailurePolicy
+    right_name: str
+    judged: int = 0
+    skipped: int = 0
+    consecutive: int = 0
+    succeeded: bool = False
+
+    def ok(self) -> None:
+        self.judged += 1
+        self.consecutive = 0
+        self.succeeded = True
+
+    def skip(self, left: Item, right_position: int, reason: str) -> None:
+        self.judged += 1
+        self.skipped += 1
+        self.consecutive += 1
+        diagnostics.warn(
+            f"skipped: {describe_source(left.source)} × "  # noqa: RUF001 — pinned pair mark
+            f"{self.right_name} line {right_position + 1} ({reason})"
+        )
+        if should_halt(self.policy, total=self.judged, skipped=self.skipped):
+            raise TooManyFailures(self.skipped, self.judged, reason)
+        if should_halt_consecutive(
+            self.policy, succeeded=self.succeeded, consecutive=self.consecutive
+        ):
+            raise TooManyFailures(self.skipped, self.judged, reason)
+
+
+async def run_join(
+    request: JoinRequest,
+    context: JoinContext,
+    *,
+    stdin: TextIO,
+    stdout: TextIO,
+    stop: asyncio.Event | None = None,
+) -> ExitCode:
+    tokens = parse_join_predicate(request.predicate)  # UsageFault on bad grammar
+    if request.k < 1:
+        raise UsageFault(f"--k must be >= 1, got {request.k}")
+    right_items = _load_right(request.right)
+    embed_model = await context.embedding_model(request.embed_model_flag)
+    log = diagnostics.DegradationLog()  # per-row conversion disclosure (D27)
+    kept_right, index, right_chunks = await _index_right(
+        embed_model, right_items, request.right.name, log
+    )
+    chat = await context.chat_model(request.model_flag)
+    concurrency = context.concurrency(request.concurrency_flag)
+    if request.kind not in ("inner", "leftouter", "anti"):
+        raise UsageFault("--kind takes inner, leftouter, or anti")
+    if request.kind == "anti" and request.unmatched is not None:
+        raise UsageFault(
+            "--unmatched with --kind anti is redundant — anti already puts unmatched rows on stdout"
+        )
+    writer = context.writer(
+        request.output,
+        structured=request.kind != "anti",  # anti emits left rows verbatim
+        stdout=stdout,
+        fields=request.fields,
+    )
+    items_iter, total = readers.resolve_items(request.input, stdin, stop=stop)
+    preview_cost(total, request.k, len(index))
+
+    converter = make_converter(
+        chat, allow_paid=request.allow_captions, log=log, stt=context.remote_transcriber(chat.ref)
+    )
+    book = PairBook(policy=FailurePolicy(), right_name=request.right.name)
+    spinner = make_stderr_spinner()
+    spinner.start(total=total)
+
+    async def worker(item: Item) -> tuple[Item, tuple[tuple[int, float], ...]]:
+        matches = await _join_one(
+            item,
+            log=log,
+            converter=converter,
+            embed_model=embed_model,
+            chat=chat,
+            tokens=tokens,
+            index=index,
+            kept_right=kept_right,
+            right_chunks=right_chunks,
+            request=request,
+            book=book,
+            stop=stop,
+        )
+        return item, matches
+
+    done = 0
+    skipped = 0
+    outcomes = run_ordered(
+        items_iter, worker, concurrency=concurrency, failure_policy=FailurePolicy(), stop=stop
+    )
+    matched_pairs = 0
+    unmatched_count = 0
+    unmatched_sink = (
+        request.unmatched.open("w", encoding="utf-8") if request.unmatched is not None else None
+    )
+    try:
+        async for outcome in outcomes:
+            if isinstance(outcome, Done):
+                left, matches = outcome.value
+                if request.kind != "anti":
+                    for position, score in matches:
+                        writer.write_record(
+                            {
+                                "left": _payload(left),
+                                "right": _payload(kept_right[position]),
+                                "_score": round(score, 4),
+                            }
+                        )
+                matched_pairs += len(matches)
+                if not matches:
+                    unmatched_count += 1
+                    match request.kind:
+                        case "anti":  # the unmatched row IS the finding — verbatim
+                            writer.write_text(left.raw)
+                        case "leftouter":  # every left row, match or not
+                            writer.write_record({"left": _payload(left), "right": None})
+                        case _:
+                            if unmatched_sink is not None:
+                                unmatched_sink.write(left.raw + "\n")
+                done += 1
+            else:  # Skipped — the left item itself failed (image, embed error, …)
+                diagnostics.warn(f"skipped: {describe_source(outcome.source)} ({outcome.reason})")
+                skipped += 1
+            spinner.advance()
+    finally:
+        spinner.finish()
+        writer.flush()
+        log.finish()
+        if unmatched_sink is not None:
+            unmatched_sink.close()
+    if request.unmatched is not None:
+        diagnostics.note(
+            f"join: {matched_pairs} matched · {unmatched_count} unmatched → "
+            f"{request.unmatched.name}"
+        )
+    elif request.kind != "inner":
+        diagnostics.note(f"join: {matched_pairs} matched · {unmatched_count} unmatched")
+    if stop is not None and stop.is_set():
+        diagnostics.interrupted_summary(processed=done, skipped=skipped + book.skipped)
+        return interrupted_exit_code(done=done, skipped=skipped + book.skipped)
+    return outcome_exit_code(done=done, skipped=skipped + book.skipped)
+
+
+async def _join_one(
+    item: Item,
+    *,
+    log: diagnostics.DegradationLog,
+    converter: Converter,
+    embed_model: EmbeddingModel,
+    chat: ChatModel,
+    tokens: tuple[Token, ...],
+    index: RightIndex,
+    kept_right: list[Item],
+    right_chunks: dict[int, ChunkedSide],
+    request: JoinRequest,
+    book: PairBook,
+    stop: asyncio.Event | None,
+) -> tuple[tuple[int, float], ...]:
+    item = await ensure_text(item, log=log, converter=converter)  # D33 ladder
+    budget = embed_budget(embed_model.ref.provider)
+    left_side: ChunkedSide | None = None
+    if estimate_tokens(item.text) > budget:
+        vector, left_side = await _chunked_side(embed_model, item.text, budget)
+        log.note(
+            describe_source(item.source),
+            "oversized → best-chunk judge",
+            f"{len(left_side.chunks)} chunks pooled for blocking",
+        )
+    else:
+        vector = (await embed_model.embed([item.text]))[0]
+    matches: list[tuple[int, float]] = []
+    for position, score in candidates(vector, index, k=request.k, threshold=request.threshold):
+        if stop is not None and stop.is_set():
+            break
+        right_item = kept_right[position]
+        left_for_judge = item
+        if left_side is not None:
+            left_for_judge = replace(item, text=left_side.best_text(index.vectors[position]))
+        chunked_right = right_chunks.get(position)
+        if chunked_right is not None:
+            right_item = replace(right_item, text=chunked_right.best_text(vector))
+        try:
+            verdict = await _judge(chat, tokens, left_for_judge, right_item)
+        except ItemError as exc:
+            book.skip(item, position, str(exc))
+            continue
+        book.ok()
+        if verdict:
+            matches.append((position, score))
+    return tuple(matches)
+
+
+def _payload(item: Item) -> dict[str, object]:
+    return dict(item.data) if item.data is not None else {"text": item.text}
+
+
+async def _judge(chat: ChatModel, tokens: tuple[Token, ...], left: Item, right: Item) -> bool:
+    """One verdict with the standard single repair; ItemError = the pair skips."""
+    request = build_judge_request(tokens, left, right)
+    reply = await chat.complete(request)
+    try:
+        verdict = validate_and_coerce(reply, JUDGE_SCHEMA)
+    except ItemError as first_error:
+        repair = build_repair_request(request, bad_reply=reply, error=str(first_error))
+        verdict = validate_and_coerce(await chat.complete(repair), JUDGE_SCHEMA)
+    return verdict.get("match") is True
+
+
+def _load_right(path: Path) -> list[Item]:
+    if str(path) == "-":
+        raise UsageFault(
+            "--right - reads nothing — stdin is join's left side\n"
+            "  The right side is a finite file smartpipe indexes up front.\n"
+            '  Example: cat stream.jsonl | smartpipe join "…" --right catalog.jsonl'
+        )
+    if not path.exists():
+        raise UsageFault(
+            f"no such file: {path}\n  --right needs a JSONL or plain-lines file to index."
+        )
+    lines = path.read_text(encoding="utf-8").splitlines(keepends=True)
+    items = [
+        replace(
+            item_from_line(line, index),
+            source=ItemSource("file", path.name, index),
+        )
+        for index, line in enumerate(lines)
+    ]
+    if not items:
+        raise UsageFault(
+            f"{path} is empty — a join against nothing is a mistake\n"
+            '  join needs a right side to match against. If you meant "keep nothing", '
+            "that's filter."
+        )
+    return items
+
+
+@dataclass(frozen=True, slots=True)
+class ChunkedSide:
+    """An oversized side's chunks + their vectors — the judge sees the best one."""
+
+    chunks: tuple[str, ...]
+    vectors: tuple[tuple[float, ...], ...]
+
+    def best_text(self, other_vector: tuple[float, ...]) -> str:
+        position, _score = rank(other_vector, self.vectors)[0]
+        return self.chunks[position]
+
+
+async def _chunked_side(
+    model: EmbeddingModel, text: str, budget: int
+) -> tuple[tuple[float, ...], ChunkedSide]:
+    """Chunk-embed one oversized text: pooled vector for blocking, chunks kept
+    so the judge reads the most-relevant one (D26/W3 — no more skipping)."""
+    chunks = split_text(text, budget)
+    vectors = await model.embed(list(chunks))
+    return mean_pool(vectors), ChunkedSide(tuple(chunks), tuple(vectors))
+
+
+async def _index_right(
+    model: EmbeddingModel,
+    items: list[Item],
+    right_name: str,
+    log: diagnostics.DegradationLog,
+) -> tuple[list[Item], RightIndex, dict[int, ChunkedSide]]:
+    """The build side, fully embedded before any chat spend (the preflight)."""
+    budget = embed_budget(model.ref.provider)
+    normal = [item for item in items if estimate_tokens(item.text) <= budget]
+    oversized = [item for item in items if estimate_tokens(item.text) > budget]
+    kept: list[Item] = []
+    vectors: list[tuple[float, ...]] = []
+    chunked: dict[int, ChunkedSide] = {}
+    async for outcome in embed_in_batches(model, normal, failure_policy=FailurePolicy()):
+        if isinstance(outcome, Done):
+            item, vector = outcome.value
+            kept.append(item)
+            vectors.append(vector)
+        else:
+            diagnostics.warn(f"skipped: {right_name} line {outcome.index + 1} ({outcome.reason})")
+    for item in oversized:
+        try:
+            pooled, side = await _chunked_side(model, item.text, budget)
+        except ItemError as exc:
+            diagnostics.warn(f"skipped: {right_name} line {item.source.index + 1} ({exc})")
+            continue
+        chunked[len(kept)] = side
+        kept.append(item)
+        vectors.append(pooled)
+        log.note(
+            f"{right_name} line {item.source.index + 1}",
+            "oversized → best-chunk judge",
+            f"{len(side.chunks)} chunks pooled for blocking",
+        )
+    return kept, build_index(vectors), chunked
+
+
+def preview_cost(total: int | None, k: int, index_size: int) -> None:
+    per_item = min(k, index_size)
+    if total is None:
+        diagnostics.preview(
+            f"join: up to {per_item} model calls per input line (cap with --max-calls)"
+        )
+        return
+    estimate = total * per_item
+    if estimate > _PREVIEW_THRESHOLD:
+        diagnostics.preview(
+            f"join: {total:,} left items · up to {per_item} candidates each = "
+            f"at most {estimate:,} model calls (cap with --max-calls)"
+        )