Artifex/control_plane/authoring/streaming.py

from __future__ import annotations

import os
import re
import tempfile
from dataclasses import dataclass
from pathlib import Path

from model_router.router import ModelRequestContract, ModelRouter


def word_count(text: str) -> int:
    return len(re.findall(r"\b\S+\b", text))


def atomic_write_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as handle:
            handle.write(text)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    except BaseException:
        try:
            os.unlink(temporary)
        except FileNotFoundError:
            pass
        raise


def merge_with_overlap(existing: str, continuation: str, max_overlap: int = 4000) -> str:
    existing = existing.rstrip()
    continuation = continuation.lstrip()
    if not existing:
        return continuation
    limit = min(len(existing), len(continuation), max_overlap)
    for size in range(limit, 39, -1):
        if existing[-size:] == continuation[:size]:
            return existing + continuation[size:]
    return existing + ("" if existing.endswith((" ", "\n")) else " ") + continuation


@dataclass(frozen=True)
class DraftResult:
    text: str
    attempts: int
    resumed: bool
    word_count: int


class ResumableDraftWriter:
    def __init__(self, router: ModelRouter) -> None:
        self.router = router

    def generate(
        self,
        *,
        request: ModelRequestContract,
        partial_path: Path,
        minimum_words: int = 3000,
        maximum_words: int = 15000,
        completion_marker: str = "[[END_OF_CHAPTER]]",
        max_attempts: int = 4,
    ) -> DraftResult:
        attempt_path = partial_path.with_name(partial_path.name + ".attempt")
        partial = self._reconcile(partial_path, attempt_path)
        if completion_marker in partial:
            completed = partial.partition(completion_marker)[0].rstrip()
            if word_count(completed) < minimum_words:
                partial = ""
                atomic_write_text(partial_path, partial)
        resumed = bool(partial)
        last_error = "generation did not complete"
        retry_feedback = ""
        short_completions = 0
        for attempt in range(1, max_attempts + 1):
            prompt = request.prompt + retry_feedback
            if partial:
                prompt += (
                    "\n\nContinue from the exact cutoff below. Return continuation prose only; do not restart "
                    "or summarize. Finish with the required completion marker.\n<saved-prose>\n"
                    + partial
                    + "\n</saved-prose>"
                )
            continued_request = ModelRequestContract(
                purpose=request.purpose,
                prompt=prompt,
                model_hint=request.model_hint,
                token_budget=request.token_budget,
                project=request.project,
                agent_version=request.agent_version,
            )
            attempt_path.parent.mkdir(parents=True, exist_ok=True)
            try:
                with attempt_path.open("w", encoding="utf-8", newline="\n") as handle:
                    for chunk in self.router.stream(continued_request):
                        handle.write(chunk.content)
                        handle.flush()
                        os.fsync(handle.fileno())
                partial = self._reconcile(partial_path, attempt_path)
                words = word_count(partial)
                if words > maximum_words:
                    raise RuntimeError(
                        f"generated prose exceeds maximum: {words} > {maximum_words} words"
                    )
                if completion_marker not in partial:
                    last_error = "provider completed without the chapter marker"
                    continue
                body = partial.partition(completion_marker)[0].rstrip()
                if word_count(body) < minimum_words:
                    last_error = f"completed chapter is shorter than {minimum_words} words"
                    short_completions += 1
                    partial = ""
                    atomic_write_text(partial_path, partial)
                    if short_completions >= 2:
                        break
                    retry_feedback = (
                        f"\n\nThe prior complete draft was too short. Write at least {minimum_words} words "
                        "and fully dramatize every planned scene without padding or repeating the chapter."
                    )
                    continue
                atomic_write_text(partial_path, body)
                return DraftResult(body, attempt, resumed, word_count(body))
            except Exception as exc:
                last_error = str(exc)
                partial = self._reconcile(partial_path, attempt_path)
                if word_count(partial) > maximum_words:
                    raise
        raise RuntimeError(
            f"chapter remains partial after {max_attempts} attempts at {partial_path}: {last_error}"
        )

    def _reconcile(self, partial_path: Path, attempt_path: Path) -> str:
        partial = partial_path.read_text(encoding="utf-8") if partial_path.exists() else ""
        if attempt_path.exists():
            partial = merge_with_overlap(partial, attempt_path.read_text(encoding="utf-8"))
            atomic_write_text(partial_path, partial)
            attempt_path.unlink()
        return partial.strip()