Extract chapter draft content-rejection retry into _draft_with_content_retry helper

novelforge/agents/chapter/__init__.py

@@ -16,13 +16,15 @@
     PASS_OPTIONAL,
     PASS_REQUIRED,
     _CONTENT_RETRY_LIMIT,
+    _DRAFT_CONTENT_NOTE,
     _FORBIDDEN_WORDS,
     _HARD_BAN_THRESHOLD,
     _OVERUSED_PATTERNS,
     _PATTERN_THRESHOLD,
     _SOFT_LIMIT_PER_CHAPTER,
     _SOFT_LIMITED_WORDS,
     _call_with_content_retry,
+    _draft_with_content_retry,
     _format_anti_repetition_rules,
     _log_pass_failure,
     _sanitize_for_content_policy,

novelforge/agents/chapter/_helpers.py

@@ -160,6 +160,60 @@ def _call_with_content_retry(
     raise ContentRejectionError(f"Content retry limit exceeded for {action}")
 
 
+# Content-guidance note injected into special_instructions on draft retries.
+_DRAFT_CONTENT_NOTE = (
+    "CONTENT NOTE: A previous draft attempt was rejected by a content "
+    "filter. Handle all mature themes (violence, horror, psychological "
+    "distress, body horror, etc.) through implication, atmosphere, "
+    "tension, and literary restraint rather than graphic or explicit "
+    "description. Show emotional and psychological impact. The story's "
+    "dark tone must be preserved but conveyed through what is suggested "
+    "and felt, not what is shown in detail."
+)
+
+
+def _draft_with_content_retry(
+    build_prompt_fn: Callable[[str], list[dict]],
+    *,
+    action: str,
+    special_instructions: str,
+    chapter_num: int,
+    max_attempts: int = 3,
+) -> str:
+    """
+    Call the LLM to produce an initial chapter draft, with content-rejection retry.
+
+    On ``ContentRejectionError`` a content-guidance note is appended to
+    ``special_instructions`` and the prompt is rebuilt via ``build_prompt_fn``.
+    Up to ``max_attempts`` are made before the error is re-raised.
+
+    ``build_prompt_fn`` must accept a single ``instructions`` string and return
+    the message list to pass to :func:`call_llm`.
+    """
+    content_note = ""
+    for attempt in range(max_attempts):
+        try:
+            instructions = special_instructions
+            if content_note:
+                instructions = (
+                    f"{special_instructions}\n\n{content_note}"
+                    if special_instructions
+                    else content_note
+                )
+            return call_llm(build_prompt_fn(instructions), action=action)
+        except ContentRejectionError:
+            if attempt >= max_attempts - 1:
+                raise
+            logger.warning(
+                "Chapter %d draft rejected by content filter (attempt %d/%d), "
+                "adding content guidance and retrying",
+                chapter_num, attempt + 1, max_attempts,
+            )
+            content_note = _DRAFT_CONTENT_NOTE
+    # Unreachable, but keeps the type checker happy
+    raise ContentRejectionError(f"Draft content retry limit exceeded for {action}")
+
+
 # ---------------------------------------------------------------------------
 # Vocabulary watchlists and scanning
 # ---------------------------------------------------------------------------

novelforge/routes/generation/chapters.py

@@ -37,6 +37,7 @@
 )
 from novelforge.agents.chapter import (
     _format_characters,
+    _draft_with_content_retry,
     build_chapter_draft_prompt,
     run_continuity_gatekeeper, run_chapter_rhythm_classifier,
     run_character_state_updater, run_per_chapter_compression_check,
@@ -301,45 +302,23 @@ def _set_step(step_label: str) -> None:
 
             # 1. Draft (with content-rejection retry)
             _set_step(f"Chapter {chapter_num}: drafting")
-            _draft_content_note = ""
-            for _draft_attempt in range(3):
-                try:
-                    _draft_instructions = special_instructions
-                    if _draft_content_note:
-                        _draft_instructions = f"{special_instructions}\n\n{_draft_content_note}" if special_instructions else _draft_content_note
-                    text = call_llm(
-                        build_chapter_draft_prompt(
-                            premise, genre, title, chapter_num, chapter_title,
-                            chapter_outline_summary, characters_text,
-                            previous_summaries, target_per_chapter, _draft_instructions,
-                            chapter_architecture_context, chapter_timeline_context,
-                            chapter_fate_context, chapter_arc_context,
-                            chapter_antagonist_context, chapter_technology_context,
-                            chapter_theme_context, gatekeeper_brief,
-                            compression_guidance, chapter_rhythm_shape,
-                            chapter_rhythm_reason, chapter_pov_context,
-                            voice_prompt, perspective_prompt,
-                        ),
-                        action=f"Chapter {chapter_num}: drafting"
-                    )
-                    break
-                except ContentRejectionError as _draft_exc:
-                    if _draft_attempt >= 2:
-                        raise
-                    logger.warning(
-                        "Chapter %d draft rejected by content filter (attempt %d/2), "
-                        "adding content guidance and retrying",
-                        chapter_num, _draft_attempt + 1,
-                    )
-                    _draft_content_note = (
-                        "CONTENT NOTE: A previous draft attempt was rejected by a content "
-                        "filter. Handle all mature themes (violence, horror, psychological "
-                        "distress, body horror, etc.) through implication, atmosphere, "
-                        "tension, and literary restraint rather than graphic or explicit "
-                        "description. Show emotional and psychological impact. The story's "
-                        "dark tone must be preserved but conveyed through what is suggested "
-                        "and felt, not what is shown in detail."
-                    )
+            text = _draft_with_content_retry(
+                lambda instr: build_chapter_draft_prompt(
+                    premise, genre, title, chapter_num, chapter_title,
+                    chapter_outline_summary, characters_text,
+                    previous_summaries, target_per_chapter, instr,
+                    chapter_architecture_context, chapter_timeline_context,
+                    chapter_fate_context, chapter_arc_context,
+                    chapter_antagonist_context, chapter_technology_context,
+                    chapter_theme_context, gatekeeper_brief,
+                    compression_guidance, chapter_rhythm_shape,
+                    chapter_rhythm_reason, chapter_pov_context,
+                    voice_prompt, perspective_prompt,
+                ),
+                action=f"Chapter {chapter_num}: drafting",
+                special_instructions=special_instructions,
+                chapter_num=chapter_num,
+            )
 
             ch_ctx = ChapterContext(
                 architecture=chapter_architecture_context,

tests/test_draft_with_content_retry.py

@@ -0,0 +1,191 @@
+"""
+Tests for the _draft_with_content_retry helper.
+
+Verifies that:
+- A successful first attempt returns the LLM response immediately.
+- On ContentRejectionError the helper retries with a content-guidance note
+  appended to special_instructions.
+- The content note is combined with special_instructions correctly when
+  special_instructions is non-empty vs. empty.
+- After max_attempts the ContentRejectionError is re-raised.
+- The chapter_num is logged in the warning message on retry.
+"""
+
+from __future__ import annotations
+
+import logging
+from unittest.mock import call, MagicMock
+
+import pytest
+
+import novelforge.agents.chapter._helpers as chap_helpers
+from novelforge.agents.chapter import _draft_with_content_retry, _DRAFT_CONTENT_NOTE
+from novelforge.llm.client import ContentRejectionError
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_rejection(*args, **kwargs):
+    raise ContentRejectionError("content blocked")
+
+
+def _make_prompt_fn(recorded: list[str]):
+    """Return a prompt builder that records the instructions it receives."""
+    def _build(instructions: str) -> list[dict]:
+        recorded.append(instructions)
+        return [{"role": "user", "content": instructions}]
+    return _build
+
+
+# ---------------------------------------------------------------------------
+# Success on first attempt
+# ---------------------------------------------------------------------------
+
+
+class TestDraftWithContentRetrySuccess:
+    def test_returns_llm_response_on_first_attempt(self, monkeypatch):
+        monkeypatch.setattr(chap_helpers, "call_llm", lambda msgs, *, action: "draft text")
+        recorded: list[str] = []
+        result = _draft_with_content_retry(
+            _make_prompt_fn(recorded),
+            action="Chapter 1: drafting",
+            special_instructions="Write dark themes.",
+            chapter_num=1,
+        )
+        assert result == "draft text"
+        assert len(recorded) == 1
+        assert recorded[0] == "Write dark themes."
+
+    def test_prompt_fn_receives_empty_instructions(self, monkeypatch):
+        monkeypatch.setattr(chap_helpers, "call_llm", lambda msgs, *, action: "ok")
+        recorded: list[str] = []
+        _draft_with_content_retry(
+            _make_prompt_fn(recorded),
+            action="Ch 2: drafting",
+            special_instructions="",
+            chapter_num=2,
+        )
+        assert recorded[0] == ""
+
+
+# ---------------------------------------------------------------------------
+# Retry behaviour on ContentRejectionError
+# ---------------------------------------------------------------------------
+
+
+class TestDraftWithContentRetryOnRejection:
+    def test_retries_once_on_rejection_then_succeeds(self, monkeypatch):
+        """First call raises ContentRejectionError; second call succeeds."""
+        call_count = 0
+
+        def _selective_llm(msgs, *, action):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise ContentRejectionError("blocked")
+            return "retry draft"
+
+        monkeypatch.setattr(chap_helpers, "call_llm", _selective_llm)
+        recorded: list[str] = []
+        result = _draft_with_content_retry(
+            _make_prompt_fn(recorded),
+            action="Ch 3: drafting",
+            special_instructions="My instructions.",
+            chapter_num=3,
+        )
+        assert result == "retry draft"
+        assert call_count == 2
+        # First call: original instructions; second call: with content note appended
+        assert recorded[0] == "My instructions."
+        assert recorded[1] == f"My instructions.\n\n{_DRAFT_CONTENT_NOTE}"
+
+    def test_content_note_used_alone_when_no_special_instructions(self, monkeypatch):
+        """When special_instructions is empty, the retry uses only the content note."""
+        call_count = 0
+
+        def _selective_llm(msgs, *, action):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise ContentRejectionError("blocked")
+            return "ok"
+
+        monkeypatch.setattr(chap_helpers, "call_llm", _selective_llm)
+        recorded: list[str] = []
+        _draft_with_content_retry(
+            _make_prompt_fn(recorded),
+            action="Ch 4: drafting",
+            special_instructions="",
+            chapter_num=4,
+        )
+        assert recorded[1] == _DRAFT_CONTENT_NOTE
+
+    def test_raises_after_max_attempts_exhausted(self, monkeypatch):
+        """ContentRejectionError is re-raised once max_attempts is reached."""
+        monkeypatch.setattr(chap_helpers, "call_llm", _make_rejection)
+        with pytest.raises(ContentRejectionError):
+            _draft_with_content_retry(
+                _make_prompt_fn([]),
+                action="Ch 5: drafting",
+                special_instructions="instr",
+                chapter_num=5,
+                max_attempts=3,
+            )
+
+    def test_exact_attempt_count_matches_max_attempts(self, monkeypatch):
+        """call_llm is called exactly max_attempts times before re-raising."""
+        call_count = 0
+
+        def _always_reject(msgs, *, action):
+            nonlocal call_count
+            call_count += 1
+            raise ContentRejectionError("always blocked")
+
+        monkeypatch.setattr(chap_helpers, "call_llm", _always_reject)
+        with pytest.raises(ContentRejectionError):
+            _draft_with_content_retry(
+                _make_prompt_fn([]),
+                action="Ch 6: drafting",
+                special_instructions="",
+                chapter_num=6,
+                max_attempts=2,
+            )
+        assert call_count == 2
+
+    def test_warning_logged_on_retry(self, monkeypatch, caplog):
+        """A WARNING must be logged on each retry, including the chapter number."""
+        call_count = 0
+
+        def _fail_once(msgs, *, action):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise ContentRejectionError("blocked")
+            return "ok"
+
+        monkeypatch.setattr(chap_helpers, "call_llm", _fail_once)
+        with caplog.at_level(logging.WARNING, logger="novelforge.agents.chapter._helpers"):
+            _draft_with_content_retry(
+                _make_prompt_fn([]),
+                action="Ch 7: drafting",
+                special_instructions="",
+                chapter_num=7,
+            )
+        warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
+        assert warnings, "Expected at least one WARNING to be emitted on retry"
+        text = warnings[-1].getMessage()
+        assert "7" in text, f"Chapter number missing from warning: {text!r}"
+
+    def test_custom_max_attempts_respected(self, monkeypatch):
+        """max_attempts=1 means a single attempt with no retries."""
+        monkeypatch.setattr(chap_helpers, "call_llm", _make_rejection)
+        with pytest.raises(ContentRejectionError):
+            _draft_with_content_retry(
+                _make_prompt_fn([]),
+                action="Ch 8: drafting",
+                special_instructions="",
+                chapter_num=8,
+                max_attempts=1,
+            )

tests/test_progress_snapshot.py

@@ -298,7 +298,7 @@ def test_snapshot_written_on_runtime_error(self, mock_llm, mocker, monkeypatch,
         _create_progress(token)
         snap = _seed_snap()
 
-        mocker.patch("novelforge.routes.generation.chapters.call_llm", side_effect=RuntimeError("boom"))
+        mocker.patch("novelforge.agents.chapter._helpers.call_llm", side_effect=RuntimeError("boom"))
 
         written = self._collect_final_snapshot(monkeypatch, tmp_path, token)
         _run_chapter_generation_internal(token, snap, [], [], 0)
@@ -318,15 +318,16 @@ def test_snapshot_written_on_content_rejection(self, mock_llm, mocker, monkeypat
         _create_progress(token)
         snap = _seed_snap()
 
-        # ContentRejectionError has an inner retry loop that exhausts 3 attempts;
-        # patching with side_effect ensures all draft calls raise and the outer
-        # except ContentRejectionError handler is reached.
+        # ContentRejectionError has an inner retry loop (inside _draft_with_content_retry
+        # in novelforge.agents.chapter._helpers) that exhausts 3 attempts; patching
+        # _helpers.call_llm with side_effect ensures all draft calls raise and the outer
+        # except ContentRejectionError handler in the generation worker is reached.
         #
-        # Use gen_mod.ContentRejectionError (the class bound in generation.py's own
-        # namespace) so the raised instance matches what generation.py's except
-        # clauses check, regardless of any module-reload in other tests.
+        # Use ContentRejectionError imported at module level so the raised instance matches
+        # what generation.py's except clauses check, regardless of any module-reload in
+        # other tests.
         mocker.patch(
-            "novelforge.routes.generation.chapters.call_llm",
+            "novelforge.agents.chapter._helpers.call_llm",
             side_effect=ContentRejectionError("policy"),
         )
 
@@ -350,7 +351,7 @@ def test_snapshot_written_on_circuit_breaker(self, mock_llm, mocker, monkeypatch
         snap = _seed_snap()
 
         mocker.patch(
-            "novelforge.routes.generation.chapters.call_llm",
+            "novelforge.agents.chapter._helpers.call_llm",
             side_effect=CircuitBreakerError("cb"),
         )
 
@@ -374,7 +375,7 @@ def test_snapshot_written_on_all_providers_exhausted(self, mock_llm, mocker, mon
         snap = _seed_snap()
 
         mocker.patch(
-            "novelforge.routes.generation.chapters.call_llm",
+            "novelforge.agents.chapter._helpers.call_llm",
             side_effect=AllProvidersExhaustedError("all"),
         )