141 lines
5.5 KiB
Python
141 lines
5.5 KiB
Python
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()
|